Training Flow Matching: The Role of Weighting and Parameterization

This paper systematically investigates how loss weighting and output parameterization choices in flow matching models interact with data dimensionality, architecture, and dataset size to provide practical design insights for optimizing both denoising accuracy and generative quality.

The Gist

Imagine you are trying to teach a robot how to draw a perfect picture of a cat. You start with a canvas full of random static (noise) and want the robot to slowly transform that static into a clear image of a cat. This is how modern AI image generators, like Flow Matching and Diffusion models, work. They don't just "guess" the picture; they learn a step-by-step process of cleaning up the noise.

This paper is like a chef's guide to the perfect recipe. The authors aren't inventing a new cooking method; instead, they are testing different ingredients (mathematical settings) to see which combination makes the best cake. They focus on two main "ingredients":

1. The Weighting (How much attention to pay): Should the robot focus more on the very beginning of the process (when the image is just static) or the end (when it's almost a clear picture)?
2. The Parameterization (What the robot is asked to predict): Is it easier for the robot to guess "What is the final cat?" (Clean Image), "What is the static?" (Noise), or "Which direction should I move to get closer to the cat?" (Velocity)?

Here is the breakdown of their findings using simple analogies:

1. The Weighting: "The Spotlight"

Imagine the training process is a long journey from a dark cave (pure noise) to a sunny meadow (the clear image). The "weighting" is the spotlight the teacher shines on the robot.

• The Old Way: Some teachers shine the light equally everywhere.
• The Paper's Discovery: The best teachers shine the light much brighter near the end of the journey (when the image is almost clear).
• Why? Think of it like polishing a diamond. The rough shaping is important, but the final polishing (removing the last tiny scratches) requires the most attention to get a perfect shine. The paper proves mathematically that focusing on these "almost done" moments yields the best results. They found that a specific mathematical formula (called SNR weighting) acts like the perfect spotlight, making the robot learn faster and better.

2. The Parameterization: "The GPS vs. The Map"

This is the most interesting part. The robot needs to know what to predict at every step.

• Option A: Predict the Clean Image (The Map). The robot tries to guess the final picture right away.
• Option B: Predict the Noise (The Static). The robot tries to guess what the mess looks like so it can subtract it.
• Option C: Predict the Velocity (The GPS). The robot doesn't guess the destination or the mess; it just guesses "Which way should I walk?"

The Big Surprise:
For a long time, some researchers thought predicting the "Clean Image" (Option A) was best because real-world data (like photos) is simple and sits on a "low-dimensional manifold" (a fancy way of saying photos have patterns and aren't totally random). They thought, "If the data is simple, just guess the answer!"

The Paper's Verdict:
It depends entirely on what kind of brain (architecture) the robot has.

• The Local Brain (U-Net): Imagine a robot that looks at the picture one tiny tile at a time, like a person looking through a small tube. This robot works best when it follows the GPS (Velocity). It doesn't need to see the whole picture to know which way to step; it just needs local direction.
• The Global Brain (ViT): Imagine a robot that sees the whole picture at once, like a bird flying high above. This robot struggles with the GPS. It works better when it tries to predict the Clean Image (Map) directly.

The Patch Size Analogy:
The authors found that if you force the "Global Brain" to look at the picture in huge chunks (large patches), it gets confused by the GPS and fails. But if you break the picture into tiny pieces (small patches), the GPS works great again. It's not about the size of the picture, but how the robot sees it.

3. The Data Amount: "The Student's Library"

The paper also looked at how much data the robot has to study.

• Small Library (Few images): If the robot only has a few pictures to learn from, it's better off trying to memorize the "Clean Image" directly. It's like a student with a small textbook who should just memorize the answers.
• Huge Library (Many images): If the robot has millions of images, it can afford to learn the "GPS" (Velocity) rules, which helps it generalize better to new, unseen pictures.

The Takeaway

The paper concludes that there is no single "best" setting for everyone. It's like building a car:

• If you are driving on a bumpy, local road (using a U-Net), you want a GPS (Velocity) and a spotlight focused on the finish line.
• If you are flying a plane over a vast landscape (using a ViT with large patches), you might prefer a Map (Clean Image) and a different kind of spotlight.

In short: Don't just copy-paste settings from other AI models. You have to match your "brain" (architecture) and your "library" (data size) with the right "teaching style" (weighting and prediction target) to get the best results.

Technical Summary

1. Problem Statement

Flow Matching (FM) and Diffusion Models are state-of-the-art generative methods, yet fundamental questions remain regarding optimal training design choices. Specifically, practitioners face two critical decisions:

1. Loss Weighting: How should the training loss be weighted across different noise levels (time steps $t$)?
2. Parameterization: What should the neural network predict? Common choices include predicting the clean image ($x$-prediction/denoiser), the noise ($\epsilon$-prediction), or the velocity ($v$-prediction).

While empirical heuristics exist, there is a lack of principled understanding regarding how these choices interact with model architecture (e.g., U-Net vs. ViT), data dimensionality, and dataset size. Recent literature presents conflicting findings (e.g., Li & He, 2025 suggest $x$-prediction is superior due to manifold assumptions, while earlier FM works favor $v$-prediction). This paper aims to disentangle these factors to provide practical design guidelines.

2. Methodology

The authors propose a unified denoising framework to compare different training objectives.

• Unified Formulation: They express all training objectives as a weighted denoising problem:
  $$\min_{D \in \mathcal{C}} \mathbb{E}_{t, x_0, x_1} \left[ w_t \| D(x_t, t) - x_1 \|^2 \right]$$
  where $D$ is the denoiser function, $\mathcal{C}$ is the class of learnable functions determined by the parameterization, and $w_t$ is the time-dependent weighting.
• Parametrization Classes ($\mathcal{C}$):
  • $C_{den}$: Predicts clean image directly ($D = N_\theta$).
  • $C_{vel}$: Predicts velocity, reconstructing clean image as $D = x + (1-t)N_\theta$.
  • $C_{noise}$: Predicts noise, reconstructing clean image as $D = (x - (1-t)N_\theta)/t$.
• Weighting Schemes ($w_t$): They analyze standard weightings including Flow Matching ($1/(1-t)^2$), Signal-to-Noise Ratio (SNR,$t^2/(1-t)^2$), and classical denoising weights.
• Experimental Setup:
  • Datasets: Synthetic datasets with controlled geometry (Fourier modes to control intrinsic dimension), CIFAR-10, CelebA-64/128, and ImageNet (via literature comparison).
  • Architectures: U-Nets, Vision Transformers (ViT) with varying patch sizes, and MLPs.
  • Metrics:
    • PSNR (Peak Signal-to-Noise Ratio): Measures denoising accuracy at specific noise levels.
    • FID (Fréchet Inception Distance): Measures generative quality.
  • Key Innovation: The authors systematically decouple weightings and parameterizations, testing combinations that are not standard (e.g., using SNR weighting with velocity parameterization).

3. Key Contributions

A. Theoretical Insights on Weighting

• Optimal Weighting: The paper provides a statistical justification for why weightings diverging as $(1-t)^{-2}$ (specifically SNR and standard Flow Matching weights) perform best.
• Derivation: By framing the problem as heteroscedastic regression and applying maximum likelihood estimation, they show that as $t \to 1$ (low noise), the conditional variance of the data shrinks. To maximize likelihood, the loss must be weighted by the inverse of this variance, which scales as $(1-t)^{-2}$.
• Conclusion: Weightings that prioritize low-noise regimes (near $t=1$) are theoretically optimal and empirically yield the best PSNR and FID.

B. Re-evaluating Parameterization

• The Manifold Assumption Debate: The authors challenge the recent claim that $x$-prediction (denoiser) is universally superior due to data lying on a low-dimensional manifold.
• Architecture Dependency: They demonstrate that the optimal parameterization is not determined solely by data dimensionality but is critically dependent on architectural locality:
  • Local Architectures (U-Nets, ViTs with small patches): Strongly favor Velocity Parameterization ($C_{vel}$).
  • Global/Coarse Architectures (ViTs with large patches, MLPs): Favor Denoiser Parameterization ($C_{den}$).
• Mechanism: Large patch sizes in ViTs reduce the number of tokens and limit local inductive bias. In these "coarser" models, predicting the clean image directly is more stable. Conversely, models with strong local inductive biases (like U-Nets) benefit from the velocity formulation.

C. Data Regime and Sample Size

• Low Data Regime: In scenarios with limited training data (e.g., 10k images), Denoiser Parameterization ($C_{den}$) significantly outperforms velocity parameterization and leads to better generalization.
• High Data Regime: As dataset size increases (e.g., 100k+), the performance gap narrows, and velocity parameterization often regains superiority in local architectures.

4. Key Results

1. Weighting Performance:

   • SNR weighting ($w_t = t^2/(1-t)^2$) and Flow Matching weighting ($w_t = 1/(1-t)^2$) consistently achieve the highest PSNR and lowest FID.
   • Classical denoising weights (prioritizing high noise) are suboptimal for generative tasks.
   • There is a strong correlation between high PSNR (denoising accuracy) and low FID (generative quality).

2. Parameterization vs. Architecture:

   • U-Nets: $C_{vel}$ is superior regardless of data resolution (tested up to 128x128).
   • ViTs: Performance flips based on patch size.
     • Small patches (4x4): $C_{vel}$ wins.
     • Large patches (16x16): $C_{den}$ wins; $C_{vel}$ fails catastrophically at high noise levels.
   • This explains the discrepancy between Li & He (2025) (who used large-patch ViTs) and earlier FM works (which used U-Nets).

3. Manifold Dimension:

   • On synthetic data with controlled intrinsic dimension ($m$), the "manifold assumption" (that $C_{den}$ should win on low-dimensional data) only holds for coarse models (MLP, ViT-16).
   • U-Nets are largely insensitive to the intrinsic dimension of the data regarding the choice between $C_{den}$ and $C_{vel}$.

4. Decoupling Strategy:

   • The best results are achieved by decoupling the natural pairs. For instance, using SNR weighting with Velocity Parameterization often yields better results than the standard "Noise Prediction + SNR" or "Velocity + FM Weighting" pairs.

5. Significance

This paper provides a crucial "design manual" for training Flow Matching and Diffusion models:

• Resolves Conflicts: It reconciles conflicting literature by showing that architectural inductive bias (locality) is a stronger determinant of parameterization success than the manifold assumption alone.
• Theoretical Grounding: It moves loss weighting from empirical heuristics to a statistically grounded derivation based on inverse-variance weighting in heteroscedastic regression.
• Practical Guidance:
  • Use SNR or FM weighting ($\propto (1-t)^{-2}$) for almost all settings.
  • Use Velocity Parameterization if using U-Nets or ViTs with small patches.
  • Use Denoiser Parameterization if using ViTs with large patches, MLPs, or if training on very small datasets.
  • Decouple weighting and parameterization to optimize performance (e.g., SNR weight + Velocity param).

The work emphasizes that there is no "one-size-fits-all" solution; the optimal training objective must be tailored to the specific interplay between the model architecture and the data regime.