Training Flow Matching: The Role of Weighting and Parameterization

In plain words

Diffusion and flow-matching image generators are trained on noisy versions of real pictures. The network sees a picture with noise mixed in and has to predict something: the clean picture, the noise, or a direction of motion toward the clean picture. Each choice of target can be paired with a rule for how much each noise level counts during training. The field has used several pairs, and reports disagree about which pair works best.

This paper does not propose a new generator. It writes the common choices in one shared form, then trains and compares them on controlled synthetic images and on two small face and object image sets. Its main findings are that the standard flow-matching weighting and the signal-to-noise weighting do best among the options tried, and that the best prediction target depends on the network design and on how much training data there is. Convolutional U-Nets and transformers that cut the image into large patches behave differently.

TL;DR

• Clean-image, noise and velocity outputs, each with its time weight, can all be written as one weighted denoising loss §3 eq. 2 tab. 1.
• On CIFAR-10 and CelebA-64 with U-Nets, the SNR weight t²/(1−t)² and the flow-matching weight (1−t)⁻² are the two best weights; the authors give a heuristic argument for the (1−t)⁻² growth near t = 1 §4 tab. 3 tab. 4.
• With U-Nets and small-patch ViTs the velocity output is better at the tested time points; large-patch ViTs and an MLP on synthetic data, and training on 10k images, favour the clean-image output §5 fig. 3 fig. 6.
• No code is released. Every result is author-reported, from single runs, with no seeds or error bars given.

Problem

Denoising-based generators can be trained with different outputs (clean image, noise, velocity) and different time weights, and the usual pairings are chosen by convention. Prior work reports that predicting the clean image is better on high-resolution data with a large-patch transformer, and attributes this to a manifold assumption §1 §5. The authors say these choices are still heuristics without consensus §2. The paper asks which weights and outputs should be preferred, and how that depends on architecture and data §1.

Method

Setup §2 eq. 1. Noise x₀ ~ N(0, I) is transported to data x₁ along the straight line x_t = (1−t)x₀ + t·x₁ for t in [0, 1]. The network learns the velocity v by regression onto x₁ − x₀. The optimal velocity can also be written through the conditional mean of x₁ or of x₀ given x_t, which makes velocity, clean-image and noise prediction equivalent in theory §2.

One loss for all choices §3 eq. 2. The paper rewrites each training target as a weighted denoising loss on a denoiser D(x_t, t) with clean target x₁. Two things differ: the time weight w^t and the class 𝒞 of functions the network can represent after the output is mapped to D. Over all functions the minimiser is the same for every weight, so differences come from the restricted class and the weight §3. The three natural pairs are in tab. 1:

• Clean image: weight 1, D = N_θ (class 𝒞_den) eq. 3.
• Velocity: weight 1/(1−t)², D = x + (1−t)N_θ (class 𝒞_vel) eq. 4. This is standard flow matching §3.
• Noise: weight t²/(1−t)² (the SNR weight), D = (x − (1−t)N_θ)/t (class 𝒞_noise) eq. 5.

A fourth weight, "classic", is 1/t² on t in [(1+σ_max)⁻¹, 1] with σ_max = 19, which gives t_min = 0.05 §3, footnote tab. 1. The paper pairs it with no class.

Metrics §4 app. A. At a fixed time t, PSNR compares D(x_t, t) with x₁. The authors argue PSNR is cheap, does not need natural colour images, and can show where a model fails across noise levels; they also argue it can expose memorisation §4. Generation is scored by FID on 50k training images with the standard Inception backbone app. A.

Weighting experiments §4 tab. 3 tab. 4. A U-Net is trained with the parameterization fixed to 𝒞_vel and time sampled uniformly on [0, 1]. Clean image and noise weights and the classic weight are compared with the flow-matching and SNR weights. The exploding weights (1−t)⁻¹ and (1−t)⁻³ are also tested, with 𝒞_vel only §4 tab. 3.

Why (1−t)⁻² near t = 1 §4 eq. 6 eq. 7. This is a heuristic, not a proof. Write y_t = x_t / t = x₁ + ((1−t)/t)·x₀. Learning E[x₁ | y_t] is a heteroscedastic regression. As t → 1 the posterior is approximated by a Gaussian whose covariance Σ(t) goes to zero, which gives an inverse-covariance weighted loss eq. 6. In a Gaussian toy model with x₁ ~ N(0, τ²I), the posterior is given by eq. 7, and the paper states the optimal weight as 1/τ² + t²/(1−t)², which grows like (1−t)⁻² as t → 1 §4.

Parameterization experiments §5. The weight is fixed to the flow-matching weight and the three classes are compared tab. 3 fig. 2. Then the architecture varies: ViTs with patch sizes 4, 8 and 16 on CIFAR-10 fig. 3; a synthetic 32×32 dataset whose samples lie on an m-mode Fourier subspace with m in {4, 8, 16}, across four model types: a convolutional U-Net, two ViTs that cut the image into 4-pixel and 16-pixel patches, and an MLP §5 fig. 5; CelebA at 128×128 tab. 5; and the number of training images n for a grayscale CelebA-64 U-Net fig. 6 app. A.

Key equations, reproduced in our notation:

\min_{\theta}\ \mathbb{E}\big[\|v_\theta(x_t,t)-(x_1-x_0)\|^2\big],\quad x_t=(1-t)x_0+t\,x_1 \qquad [\text{eq. 1}]

\min_{D\in\mathcal{C}}\ \mathcal{L}(D)=\mathbb{E}\big[w^t\,\|D(x_t,t)-x_1\|^2\big] \qquad [\text{eq. 2}]

\mathcal{C}_{\text{den}}=\{N^\theta\} \qquad [\text{eq. 3}]

\mathcal{C}_{\text{vel}}=\{\mathrm{Id}+(1-t)\,N^\theta\} \qquad [\text{eq. 4}]

\mathcal{C}_{\text{noise}}=\{\tfrac{1}{t}(\mathrm{Id}-(1-t)\,N^\theta)\} \qquad [\text{eq. 5}]

\mathbb{E}\big[(x_1-f_\theta(y_t,t))^\top\Sigma(t)^{-1}(x_1-f_\theta(y_t,t))\big] \qquad [\text{eq. 6}]

x_1\mid(y_t=y,t)\sim\mathcal{N}\!\Big(\tfrac{\tau^2}{\tau^2+(1-t)^2/t^2}\,y,\ \sigma_{\text{post}}^2(t)\,I\Big),\quad \sigma_{\text{post}}(t)=\Big(\tfrac{1}{\tau^2}+\tfrac{t^2}{(1-t)^2}\Big)^{-1} \qquad [\text{eq. 7}]

The paper has no algorithm environments. The loss is given in eq. 2 and tab. 1, written out below as pseudocode. The reduction over pixels is not stated, so the sum here is a choice.

# Weighted denoising loss for one batch  [eq. 2]
# net(x, t) returns the raw network output N_theta  [§3]
# weighting: 'den', 'vel' or 'noise'  [tab. 1]
# parameterization: 'den', 'vel' or 'noise'  [tab. 1]
def weighted_denoising_loss(net, x1, weighting, parameterization):
    # Sample one time per example, uniform on [0, 1]  [§4] [eq. 2]
    t = uniform(0.0, 1.0, size=batch_size)     # (B,)
    # Standard Gaussian source sample  [§2]
    x0 = normal(0.0, 1.0, size=x1.shape)       # (B, C, H, W)
    # Linear interpolation between source and data  [§2] [eq. 1]
    x_t = (1 - t) * x0 + t * x1                # (B, C, H, W)
    # Raw output of the network  [§3]
    raw = net(x_t, t)                          # (B, C, H, W)
    # Map the raw output to a denoiser in the chosen class  [§3] [tab. 1]
    if parameterization == "den":
        denoiser = raw
    elif parameterization == "vel":
        denoiser = x_t + (1 - t) * raw
    else:
        denoiser = (x_t - (1 - t) * raw) / t
    # Time weight of the chosen weighting  [tab. 1]
    if weighting == "den":
        weight = 1.0
    elif weighting == "vel":
        weight = 1 / (1 - t) ** 2
    else:
        weight = t ** 2 / (1 - t) ** 2
    # Squared error against the clean sample, summed over pixels  [eq. 2]
    error = ((denoiser - x1) ** 2).sum(axis=(1, 2, 3))
    return (weight * error).mean()

The synthetic Fourier images are built as follows. Only the mode indices and the sampling of coefficients are specified; the scale and the nonlinearity constant are not stated.

# One Fourier-32 image with m active modes  [app. A]
# mode_list: the m selected (k, l) index pairs, fixed for a dataset  [app. A]
# scale: scale of the coefficient distribution (value not stated)  [app. A]
def fourier_image(mode_list, scale, N=32):
    # Spectrum with every frequency set to zero  [app. A]
    spectrum = zeros((N, N), dtype=complex)
    for (k, l) in mode_list:
        # Draw one coefficient per active mode (real draw assumed)  [app. A]
        a = normal(0.0, scale)
        spectrum[k, l] = a
        # Conjugate-symmetric partner so the image is real  [app. A]
        spectrum[(-k) % N, (-l) % N] = conj(a)
    # Inverse 2-D FFT with orthonormal scaling, real part  [app. A]
    image = real(ifft2(spectrum, norm="ortho"))
    # Bound the dynamic range with a per-sample tanh  [app. A]
    return tanh(alpha * image)                 # alpha: not stated

The lowest-frequency mode set ("lowfreq") is chosen by the periodic radius min(k, N−k)² + min(l, N−l)², with tie-breaking not stated app. A.

Hyperparameters

Most training settings are missing from the paper. Values marked "not stated" must be chosen by the implementer; the paper does not say they were copied from any library.

Results

Author-reported. Each row is one training run on the setup named in the caption. The paper reports no seeds, repetitions or error bars. PSNR is higher-is-better; FID is lower-is-better. The units of PSNR are not stated in the tables.

CIFAR-10, 32×32, U-Net, 1000 epochs, PSNR on 1000 images, FID on 50k training images tab. 3.

CelebA-64, 300 epochs, PSNR on 1000 test images, FID on 50k training images tab. 4.

CelebA-128, 300 epochs, PSNR on 1000 test images, FID on 50k training images tab. 5. Only two rows are reported.

ViT patch size, CIFAR-10, train FID on 50k images, 𝒞_den vs 𝒞_vel fig. 3. The PSNR-gap panel of this figure was not read; these are the FID values printed in its table.

Training set size, grayscale CelebA-64 U-Net, weight 1/(1−t)², 150000 iterations, train FID on 50k images fig. 6 app. A.

Spectral residual energy, Fourier data, m = 16, flow-matching loss, mean tab. 2 app. A. Lower is closer to the manifold. The text does not state a correlation coefficient between this energy and PSNR.

Headline findings the paper states §4 §5 §6.

• Lowest FID on both datasets is 𝒞_vel with the SNR weight (3.71 on CIFAR-10, 3.97 on CelebA-64), with 𝒞_vel and the flow-matching weight next (3.90 and 4.45) tab. 3 tab. 4.
• The authors conclude that the parameterization choice follows architecture locality: U-Nets and small-patch ViTs favour 𝒞_vel; large-patch ViTs and the MLP favour 𝒞_den on the Fourier data §5 fig. 3 fig. 5.
• With 10k training images, 𝒞_den has the lower FID; at 50k and 100k, 𝒞_vel is lower fig. 6.
• Dataset-level claims on the Fourier data (manifold dimension) are read from the PSNR-gap panels fig. 5, which were not read as figures for this card.

Paper vs code

No official code.

Searched, in order: the paper's LaTeX source and the arXiv abstract page list no code repository, project page or release link. The arXiv comments field says the work was published at the ICLR 2026 DeLTa workshop and names no code. GitHub search through the API for "flow matching weighting parameterization" and "flow matching denoiser velocity fourier manifold" returned no repositories. The authors' GitHub pages were checked (mathurinm, annegnx). One repository on annegnx, PnP-Flow, is plug-and-play image restoration code for a different paper (per its listing, from an ICLR 2025 paper), and it is not used here.

Pitfalls

• No code, no optimiser settings. The paper gives no learning rate, optimiser, schedule, network size or Fourier training settings. An implementation will differ from the paper wherever these are chosen. Do not fill them from another library's defaults without saying so.
• Small PSNR gaps. Under the 1/(1−t)² weight, PSNR differences between classes are at most about 0.1 dB on CIFAR-10 tab. 3 and at most about 0.5 dB on CelebA-64 tab. 4. The classic-weight rows differ much more (up to about 4.9 dB on CelebA-64 at t = 0.95), so the small gaps hold only for the flow-matching weight. The FID gaps are larger. The tables give single runs with no error bars, so these gaps are not established as significant.
• "Best weight" is not uniform. The text says the SNR weight gives the best results consistently §4. For 𝒞_den and 𝒞_vel this holds on FID. For 𝒞_noise the flow-matching weight has lower FID than the SNR weight on both datasets: 6.30 vs 7.15 on CIFAR-10 and 659.05 vs 681.05 on CelebA-64 tab. 3 tab. 4. Also, on CIFAR-10 𝒞_den the SNR weight beats the flow-matching weight on FID (4.60 vs 4.89) tab. 3.
• Classic weight: two readings. Section 3 derives the classic weight for classical denoisers trained on x_σ = x₁ + σx₀ with σ uniform on [0, σ_max] §3. The appendix and tab. 3 apply 1/t² on the flow-matching interpolation x_t. Section 4 samples t uniformly for all weighting runs §4, which suggests the second reading, but the paper does not state it for the classic runs.
• Singularity at t = 1. Weights 1/(1−t)² and t²/(1−t)² diverge as t → 1. The paper does not state clipping or a time cap. An implementation must choose one; the PSNR values near t = 0.95 are the closest the tables get.
• Noise output divides by t. The noise class has a 1/t factor, which grows at small t (high noise). The paper says the noise output fails at early times and attributes this to the 1/t factor, but the sentence is worded so that the factor both prevents and causes the failure §5. The card does not resolve this. Numerical handling of t near 0 is not discussed.
• Loss reduction not stated. The paper does not say whether the squared error is summed or averaged over pixels, or averaged over the batch. This changes the loss scale and therefore the effective learning rate.
• PSNR data range. The value range (for example [0, 1] or [−1, 1]) used for PSNR is not stated. PSNR values are not comparable across ranges.
• PSNR split. fig. 1 says the PSNR was computed on 1000 test images; tab. 3 says "1000 images" with no split. Treat the CIFAR-10 table's split as unknown.
• FID is against training images. FID is computed on 50k training images app. A tab. 3. For the 10k-image run the reference set is not described. The paper argues that a memorising model can have good FID with poor PSNR §4; it does not test memorisation directly.
• Possible duplicate value. The 100k-image 𝒞_vel FID in fig. 6 (4.45) equals the 𝒞_vel, 1/(1−t)² FID in tab. 4. The setups differ (grayscale and 150000 iterations vs RGB and 300 epochs), so this may be a coincidence or a transcription error. Check the source before relying on either value.
• Sign convention of the delta-PSNR plots. Captions of the delta-PSNR figures say "lower is better" for a difference from a reference, while the Fourier section defines ΔPSNR = PSNR(𝒞_den) − PSNR(𝒞_vel) and reads negative values as favouring 𝒞_vel §5. Check the plotted direction against the text before using those figures.
• Weight derivation notation. In eq. 7 σ_post(t) is written as the inverse of a precision, yet the distribution uses σ²_post. The stated optimal weight 1/τ² + t²/(1−t)² equals the precision, not the inverse square of σ_post §4. Read the derivation as the precision weight. Also, eq. 6 writes f_θ(y_t, t) where f_θ(y, t) is defined.
• Heuristic, not proof. The (1−t)⁻² argument uses a Gaussian approximation of the posterior as t → 1 and a toy Gaussian data model §4. The paper calls it a statistical justification; it does not prove optimality for real data. As an observation from eq. 7 (not stated in the paper): as τ → ∞ the precision in the derivation reduces to t²/(1−t)², the SNR weight.
• Architecture claims rest on few points. The weight and class tables sample five time values, and CelebA-128 tab. 5 has two rows. The claim that the velocity output wins "at every noise level" §5 is verified only at these five points, since the delta-PSNR curves were not read.
• Fourier dataset is underspecified. The coefficient distribution, its scale, the nonlinearity constant α and the training settings are not stated. The MLP and ViT sizes are "approximately" matched but not given. The claim that PSNR tracks distance to the manifold app. A is supported by one spectral-energy table at m = 16 with no correlation measure.
• Grayscale vs colour. The dataset-size runs and the Fourier data are grayscale. The Section 4 CelebA-64 runs have no stated channel count.
• Paper figures not read. Figure numbers here are inferred from LaTeX source order. The PDF was not read, so check figure numbers and the delta-PSNR plots against the published PDF before citing them.

Lineage

• Builds on: flow matching, Lipman et al. (ICLR 2023), "Flow matching for generative modeling" §2 [references.bib]; diffusion training with noise prediction, Ho et al. (2020) §2; weightings as ELBO terms, Kingma et al. (2023) §2; the equivalence of diffusion and Gaussian flow matching, Gao et al. (ICLR 2025 blog track) §2 [references.bib]. The paper's own unified denoising formulation §3 is its main framing.
• Compared against: Li and He (2025), "Back to basics: Let denoising generative models denoise", arXiv:2511.13720, which argues for clean-image prediction with a large-patch transformer (JiT) on high-resolution data §1 §5 [references.bib]. Jin and Wang (2026), "Revisiting Diffusion Model Predictions Through Dimensionality", arXiv:2601.21419, which shows the clean-image benefit in a toy linear setting §2 [references.bib].
• Changes: the paper proposes no new method. It separates weight from parameterization, and adds architectural locality and data size as factors §1 §6.
• Superseded by: not known from this wiki.