Transition Flow Matching

1. Problem Statement

Current mainstream generative models, such as Flow Matching (FM) and Diffusion Models, rely on learning a local velocity field $v(x_t, t)$. To generate data, these models must numerically integrate this velocity field over multiple time steps (iterative sampling) to transport a noise distribution to the data distribution.

• Limitations: This iterative process is computationally expensive.
• Existing Solutions: Methods like Consistency Models and Distillation attempt to reduce inference steps (to one or few steps). However, Consistency Models often impose behavioral constraints on network outputs without a clear theoretical link to the ground-truth field, leading to unstable training that requires complex curricula. Distillation relies on pre-trained teacher models, preventing end-to-end training from scratch.
• Gap: There is a need for a principled, end-to-end trainable framework that directly models global state transitions to enable high-quality generation in a single step (or arbitrary steps) without relying on distillation or unstable consistency constraints.

2. Methodology: Transition Flow Matching (TFM)

The authors propose Transition Flow Matching, a paradigm that shifts focus from learning local velocities to learning the transition flow itself.

Core Concept

Instead of predicting the instantaneous velocity $v(x_t, t)$, TFM learns a mapping $X^\theta(x_t, t, r)$ that directly predicts the state at a future time $r$ given the current state at time $t$. This transition flow is a global quantity representing the integrated dynamics over the interval $[t, r]$.

Key Theoretical Derivations

1. Transition Flow Identity:
   The authors derive a fundamental identity connecting the transition flow $X(x_t, t, r)$ to the local velocity field. By defining the average velocity $u(x_t, t, r)$ over the interval $[t, r]$ and differentiating with respect to $t$, they prove:
   $$X(x_t, t, r) = x_{t \to r} + (r-t) \frac{d}{dt} X(x_t, t, r)$$
   Where $x_{t \to r}$ is the true marginal transition state. This identity serves as the theoretical foundation for the training objective.

2. Tractable Training Objective:
   Directly minimizing the marginal loss is intractable because the marginal transition state $x_{t \to r}$ and marginal velocity are unknown.

   • Solution: The authors introduce a Conditional Transition Flow Matching objective. By conditioning on the endpoint pair $Z = (X_0, X_1)$ and using a linear interpolant $X_t = (1-t)X_0 + tX_1$, the conditional transition state and velocity become tractable (constant).
   • Gradient Equivalence: They prove that the gradients of the intractable marginal loss and the tractable conditional loss coincide. Thus, minimizing the conditional loss effectively trains the model to satisfy the global Transition Flow Identity.

3. Training Procedure:

   • Input: Current state $x_t$, time $t$, and target time $r$ (where $t \le r$).
   • Target Construction: The target is constructed as $X_{\text{tgt}} = x_{t \to r}^Z + (r-t) \frac{d}{dt} X^\theta(x_t, t, r)$.
   • Time Derivative: The term $\frac{d}{dt} X^\theta$ is computed efficiently using Jacobian-Vector Products (JVP), leveraging the known conditional velocity $v(x_t, t | Z) = X_1 - X_0$.
   • Loss: The model minimizes the divergence (e.g., MSE) between the predicted transition flow and the constructed target.

4. Unified Perspective:
   The paper establishes a theoretical link between TFM and Mean Velocity Flow models, showing that TFM provides a unified framework that generalizes previous approaches while offering a clearer theoretical grounding for few-step generation.

3. Key Contributions

• New Paradigm: Proposes Transition Flow Matching, a framework that directly models the generation trajectory (global transition) rather than the local velocity field.
• Theoretical Foundation: Derives the Transition Flow Identity and proves the gradient equivalence between marginal and conditional objectives, enabling end-to-end training from scratch.
• Unified View: Clarifies the relationship between TFM and Mean Velocity models, offering a cohesive theoretical perspective.
• Efficiency: Enables one-step generation (and arbitrary step sizes) without requiring distillation from pre-trained teachers or complex consistency constraints.

4. Experimental Results

The authors evaluated TFM on synthetic data, CIFAR-10, and ImageNet (256x256).

• Synthetic Data (2D "M" Dataset): Visualizations show that TFM can generate high-quality samples in a single step, with trajectories that are coherent and consistent across different time grids.
• CIFAR-10:
  • TFM achieved the best Fréchet Inception Distance (FID) in one-step generation (NFE=1) compared to Flow Matching, Rectified Flow, and Consistency Models.
  • Unlike Consistency Models, which often degrade in performance as the number of steps increases, TFM's performance consistently improved or remained stable as NFE increased.
• ImageNet (256x256):
  • Using a Transformer backbone (SiT/MeanFlow architecture), TFM-XL achieved an FID of 3.02 for one-step generation, outperforming state-of-the-art methods like MeanFlow (3.43) and S-VFM (3.31).
  • The method demonstrated robust scalability and high fidelity even with minimal function evaluations (NFE=1).
• Ablation Studies:
  • Loss Metrics: Adaptive weighting (specifically $p=1$ in the weighted $L_2$ loss) yielded the best results.
  • Time Sampling: Logit-normal sampling for time steps $(t, r)$ proved superior to uniform sampling.
  • Conditioning: Conditioning on $(t, \Delta t)$ or $(t, r)$ both worked well, with relative time information being dominant.
  • CFG: The method supports Classifier-Free Guidance effectively, significantly improving sample quality without extra inference cost.

5. Significance

Transition Flow Matching represents a significant shift in generative modeling theory and practice:

1. End-to-End Efficiency: It solves the "few-step generation" problem without the instability of consistency training or the dependency on distillation. Models can be trained from scratch to generate high-quality images in a single step.
2. Theoretical Clarity: By deriving a rigorous identity for transition flows, it provides a mathematically grounded alternative to heuristic consistency constraints.
3. Flexibility: The framework naturally supports arbitrary step sizes and numbers of steps, making it highly versatile for applications requiring different trade-offs between speed and quality.
4. State-of-the-Art Performance: The empirical results on ImageNet demonstrate that modeling global transitions is a more effective strategy for few-step generation than learning local velocity fields, setting a new benchmark for one-step generative models.