A Lagrangian View of Flow Matching

Peyman Milanfar

Source-linked AI summary

A Lagrangian View of Flow Matching

Peyman Milanfar

arXiv:2609.00198v1Aug 31, 2026cs.CVphysics.flu-dyn

TL;DR

The paper addresses the difficulty of discrete solvers facing a moving denoising target by developing a particle-centric derivation based on target-identity invariance. It formulates this condition as a quasi-linear advection PDE and shows that, under the residual ansatz, its characteristics are the straight-line trajectories of Flow Matching, while curvature is tied to Jacobian effects and intersecting characteristics motivate distillation.

Problem
Discrete solvers face a continuously shifting denoising target, which forces conservative step sizes, while standard Eulerian derivations do not expose the particle-level mechanics.
Method
The paper uses a Lagrangian invariance principle, derives a quasi-linear advection PDE, and solves it with the Method of Characteristics under a residual velocity ansatz.
Results
Under the residual ansatz, the self-consistent characteristics are exactly the straight lines of Flow Matching and Rectified Flow.
Takeaways & Limitations
Straight characteristics keep the target static and neutralize the Jacobian penalty, enabling massive solver step sizes; monitoring target drift can identify regions needing smaller steps or distillation.
Takeaways & Limitations
The derivation does not prove the residual ansatz is the unique velocity field solving the invariance PDE, and its diagnostic example is deliberately limited to two modes in one dimension.

Abstract

from arXiv · show

Modern explicit-time generative models, such as Flow Matching [Lipman et al., 2023] and Rectified Flow [Liu et al., 2023], are typically derived top-down via Optimal Transport and the continuity equation. This standard Eulerian approach focuses on the macroscopic transport of probability mass. In this paper, we present an alternative, bottom-up mechanical derivation grounded in a Lagrangian (particle-centric) perspective. By analyzing the local Taylor expansion of a continuous denoiser, we motivate a strict invariance condition required for optimal, singlestep generation: the conservation of target identity. Enforcing this condition yields a governing quasi-linear advection Partial Differential Equation (PDE). We demonstrate that solving this PDE via the Method of Characteristics analytically yields the straight-line trajectories of Flow Matching. This geometric perspective isolates the Jacobian of the denoiser as the primary source of trajectory curvature, providing a direct mathematical explanation for why straight-line flows enable massive step sizes, and why empirical models require distillation to flatten intersecting characteristics.

1 The Predictor-Corrector and the Moving Target

A continuous denoiser’s target changes as the state advances, so discrete predictor-corrector solvers must repeatedly update their destination and use small steps to avoid overshooting it.

The denoiser f(x,t) predicts the clean data state x0 from the current trajectory state x(t).
The target’s continuous change contains spatial Jacobian and temporal-drift terms arising from movement through the denoising field.
A predictor uses f(x,t) to compute the residual x − f(x,t) and determine the discrete state increment Δx.
A corrector reapplies the denoiser at x + Δx and the advanced time to reevaluate the target.
Because the destination shifts with local geometry and the receding timeline, standard diffusion solvers require extremely small Δt to avoid overshooting.

2 The Invariance Principle: From Continuity to Advection

The paper proposes conserving each particle’s target identity along its trajectory. This Lagrangian invariance condition becomes a quasi-linear advection PDE that isolates how denoiser geometry governs the flow.

Single-step generation ideally follows trajectories where the denoiser identifies the same clean target from the initial noise level onward.
The denoiser’s output must have zero total rate of change along an ideal trajectory.
Unlike Eulerian formulations that conserve probability mass across space, the Lagrangian view rides a particle and conserves its target identity x0.
Applying the chain rule and defining dx/dt as v(x,t) yields the governing invariance equation for the generative field.
The resulting first-order quasi-linear advection PDE requires temporal target drift to be neutralized by advection through the denoiser’s local Jacobian.

3 Solving the PDE: Deriving Flow Matching

The paper tests a residual-based velocity ansatz against the target-invariance PDE using the Method of Characteristics. Under the boundary condition at t = 0, self-consistency fixes straight-line characteristics, while the derivation does not establish ansatz uniqueness.

3 Solving the PDE: Deriving Flow Matching: The PDE constrains denoiser and velocity jointly, so the paper posits a residual-based velocity field rather than deriving one uniquely from the PDE.The ansatz is motivated by Tweedie’s formula and residual-driven flow.
3 Solving the PDE: Deriving Flow Matching: The Method of Characteristics tests whether this specific residual ansatz can satisfy the target-invariance principle and then solves its trajectories.It is not used to derive the ansatz itself.
3 Solving the PDE: Deriving Flow Matching: Imposing invariance along the ansatz trajectory makes the denoiser constant there, and f(x, 0) = x0 fixes that constant.This converts the coupled PDE system into a closed ODE in X(t).
3 Solving the PDE: Deriving Flow Matching: Integrating the resulting separable ODE yields the Flow Matching and Rectified Flow formulation, with characteristics as straight lines from data x0 to noise z1.The noise schedule and initial noise state enter through the integrated solution.
3 Solving the PDE: Deriving Flow Matching: The derivation is conditional: it shows consistency and trajectory determination for the residual ansatz, not uniqueness of the velocity field solving the invariance PDE.The Method of Characteristics, together with the boundary condition, determines the trajectories after the ansatz narrows the PDE solution considered.

4 The Jacobian Penalty and Trajectory Curvature

The Jacobian of the denoiser links target drift to trajectory curvature and solver difficulty. Curved diffusion trajectories encounter an exploding Jacobian near t →0, whereas Flow Matching’s straight characteristics keep the target static and permit large steps.

4 The Jacobian Penalty and Trajectory Curvature: The advection PDE expresses temporal target drift through the denoiser’s Jacobian and velocity.Rearranging the PDE isolates ∂f/∂t, while off-characteristic motion makes this drift non-zero.
4 The Jacobian Penalty and Trajectory Curvature: For an optimal MMSE denoiser with symmetric Jacobian, Jf is proportional to the posterior covariance matrix Σpost.This connects local denoiser geometry to model uncertainty.
4 The Jacobian Penalty and Trajectory Curvature: The target drift is negatively correlated with velocity and follows the principal axes of posterior uncertainty.The directional relation follows from the positive-semidefinite covariance constraint.
4 The Jacobian Penalty and Trajectory Curvature: As t →0, the data manifold sharpens and Jf’s eigenvalues diverge, forcing curved-trajectory solvers toward infinitesimal stable step sizes.The exploding Jacobian is encountered when ∂f/∂t ≠ 0.
4 The Jacobian Penalty and Trajectory Curvature: Flow Matching’s straight characteristics keep the target static, mathematically neutralize the Jacobian penalty, and allow massive solver step sizes.The cited condition is ∂f/∂t = 0 along the analytically solved characteristics.

5 The Engineering Reality: Intersecting Characteristics

Empirical Flow Matching models violate the ideal straight-characteristic picture when independently paired paths intersect, producing ambiguity, curvature, and slower sampling. A two-mode example and diagnostics localize this failure near crossings and motivate Reflow as path untangling.

5 The Engineering Reality: Intersecting Characteristics: Neural-network training violates the analytical PDE assumptions even though the PDE solution yields straight characteristics.This violation arises in empirical models rather than in the analytical construction.
5 The Engineering Reality: Intersecting Characteristics: Intersecting paths demand conflicting target values at one point, so deterministic denoisers average them and warp the vector field.The resulting marginalization bends learned trajectories away from perfect characteristic curves.
5 The Engineering Reality: Intersecting Characteristics: Reflow distillation trains on pre-matched non-intersecting trajectories to draw straight lines, reducing conflicting targets and enabling single-step generation.The procedure first simulates valid trajectories, then trains a new model on the resulting pairs.
5.1 A Minimal Working Example: The Two-Mode Crossing: In a 1D two-mode interpolation, straight characteristic families overlap near x = 0 at high noise because both modes share the same noise distribution.This provides a closed-form setting in which crossing-induced ambiguity can be examined directly.
5.1 A Minimal Working Example: The Two-Mode Crossing: Near-crossing trajectories become curved S-shapes, whereas distant trajectories remain nearly straight; starting exactly at z1 = 0 never resolves.The trajectories are integrated backward from t = 1 toward t ≈0.
5.1 A Minimal Working Example: The Two-Mode Crossing: Near-crossing trajectories have roughly 3.5× greater peak drift than far trajectories: max |df/dt| ≈9.23 versus ≈2.65.The drift remains elevated longer near the crossing and peaks around t ≈0.2–0.5.
5.1 A Minimal Working Example: The Two-Mode Crossing: The drift diagnostic |df(X(t), t)/dt| can identify unresolved crossings without ground truth and indicate where solvers should slow down or distillation should focus.The same signal is proposed for per-sample inference monitoring.
5.1 A Minimal Working Example: The Two-Mode Crossing: The example uses two modes, one dimension, and a closed-form posterior, so real-data use would require evaluating the diagnostic over many samples.The paper states that the mechanism is expected to scale directly, but the example itself is deliberately minimal.

6 Discussion and Practical Implications

The paper’s Lagrangian perspective centers analysis on the denoiser output produced during sampling rather than on abstract probability masses. This makes invariance violations measurable and connects path intersections directly to curvature, blur, and the need for further distillation.

6 Discussion and Practical Implications: The Lagrangian view complements macroscopic Eulerian guarantees by exposing the discrete mechanics of generative algorithms running on hardware.Its practical focus is the particle trajectory rather than only the transported distribution.
6 Discussion and Practical Implications: Because neural networks output denoiser values f(x, t), the invariant dtf(x, t) = 0 can be monitored directly during inference.Fluctuations indicate PDE violation, trajectory curvature, and a need for smaller steps or further distillation.
6 Discussion and Practical Implications: Intersecting characteristics force velocity averaging, which bends the flow and shifts the target; blurry 1-step outputs therefore indicate paths that remain crossed.The paper frames further Reflow as the way to untangle those paths.