Technical Summary: Conservative Flows

Problem Statement

Modern generative modeling is predominantly defined by the paradigm of transport: learning a mapping from a simple reference distribution (typically Gaussian noise) to a complex data distribution. While score-based and flow-based diffusion models have achieved state-of-the-art results by reversing a forward noising process, this approach inherently requires learning a global transport map across a family of intermediate distributions.

The authors identify a fundamental limitation in the discretization of continuous-time dynamics used in these models. Specifically, while continuous-time Langevin dynamics theoretically preserve the target data distribution $p$ as an invariant measure, their explicit discrete-time approximations (e.g., the Unadjusted Langevin Algorithm, ULA) fail to do so exactly. The discretization error causes the generated samples to drift away from the true data manifold over time, breaking the requirement that the generation process remains within the data distribution at every step.

Furthermore, the authors argue that in many practical settings—such as augmented retrieval systems or latent generative pipelines—generation does not necessarily begin from an "information-less" noise state. Instead, valid initial states often exist in the form of retrieved data points, nearest neighbors, or latent codes that already lie near or on the target distribution. The paper posits that for such scenarios, a paradigm shift from transport (noise-to-data) to invariance (data-to-data) is more natural and efficient.

Methodology

The paper proposes Conservative Flows, a framework where generation is performed by a discrete stochastic dynamics that leaves the data distribution invariant. The system is initialized from data-supported states ($x_0 \sim p$) and evolves while preserving the probability law at every step ($x_j \sim p$ for all $j \ge 1$).

The framework introduces two distinct probability-preserving sampling mechanisms that can operate directly on top of existing pretrained flow-matching models without retraining:

1. Corrected Langevin Dynamics (dMALA)

This method adapts the Metropolis-Adjusted Langevin Algorithm (MALA) to generative modeling where the explicit density $p(x)$ is unavailable.

• Mechanism: It utilizes a learned score field $s_\theta$ or, more practically, a denoiser $D_\sigma$ trained to predict clean samples from noisy observations. The denoiser is derived directly from the bridge used in flow-matching training.
• Correction: To restore invariance, the method employs a Metropolis-Hastings correction. Instead of requiring the log-density difference $\log p(y) - \log p(x)$, which is unknown, the method uses line integration of the conservative score field (or the denoiser residual) along the path between the current state $x$ and the proposal $y$.
• Invariance Target: This mechanism preserves a Gaussian-smoothed distribution $p_\sigma = p * \mathcal{N}(0, \sigma^2 I)$, where $\sigma$ is determined by the noise level of the denoiser.
• Efficiency: By approximating the line integral with the trapezoidal rule, the acceptance ratio simplifies to a function of the squared residuals of the denoiser, requiring only two denoiser evaluations per step.

2. Predictor–Corrector Flow

This method targets the original data distribution $p$ directly, rather than a smoothed version.

• Mechanism: It operates as a two-step process:
  1. Predictor: A controlled amount of noise is added to the current sample $x \sim p$, moving it to an intermediate "bridge" distribution $p_\tau$ (where $\tau < 1$).
  2. Corrector: A learned flow map (or direct solution map) is applied to transport the noisy sample back to the data distribution $p$ (time $t=1$).
• Invariance: Under an exact flow model, this kernel is strictly probability-preserving. The method relies on the flow model's ability to accurately map samples from the bridge back to the data manifold.
• Flexibility: Unlike the local accept/reject nature of dMALA, this approach allows for larger, non-local moves by temporarily stepping off the manifold and returning via the learned flow.

Key Contributions

1. Paradigm Shift: The paper reframes generative modeling from "transport from noise" to "stationary probing of the data distribution." The central object is a stochastic dynamics that preserves the relevant law, initialized from data-supported states.
2. Two Probability-Preserving Mechanisms: The authors develop and formalize two specific samplers (dMALA and Predictor–Corrector) that correct the discretization errors of standard Langevin dynamics and flow-based generation.
3. Reuse of Pretrained Models: A critical practical contribution is that neither mechanism requires training a new generative model from scratch. Both can extract the necessary score, denoiser, or flow map directly from existing pretrained flow-matching checkpoints (e.g., SoFlow, SiT).
4. Theoretical Correction: The paper explicitly addresses the failure of discrete Langevin dynamics to preserve the target law and provides the mathematical derivation for correcting this via line integration and bridge flows.

Experimental Results

The framework was validated on a synthetic Swiss-roll target, ImageNet-256, and Oxford Flowers-102.

• Synthetic Target (Swiss-roll): Experiments demonstrated that while standard ULA drifts away from the target distribution, the corrected dynamics (dMALA) and Predictor–Corrector flow remain on the correct manifold, preserving the distribution over long chains.
• ImageNet-256:
  • Both methods consistently improved upon the baseline samplers of the pretrained backbones (SoFlow-XL/2-cond and SiT-REPA-XL).
  • dMALA achieved the lowest Fréchet Inception Distance (FID) scores, particularly with smaller smoothing levels ($\sigma=0.3$), indicating high fidelity to the data distribution.
  • Predictor–Corrector achieved higher Inception Scores (IS) and Precision, suggesting it generates more diverse and distinct samples, though with slightly higher FID than dMALA in some configurations.
  • Long-Chain Behavior: Both methods showed graceful degradation with increasing chain length ($K$), consistent with the accumulation of approximation errors in the pretrained backbone, but maintained performance significantly better than baselines.
• Computational Efficiency: The Predictor–Corrector method was shown to be competitive with direct ODE integration. For instance, 200 network function evaluations (NFEs) of the Predictor–Corrector method achieved a lower FID (4.91) than 800-step forward-Euler integration (6.51) of the same backbone.

Significance and Claims

The authors position this work as a new paradigm for generative modeling. They claim that by shifting the goal from transforming noise into data to maintaining the data distribution under learned dynamics, the framework enables:

• Faithful Sample Variation: By keeping the process on high-probability regions of the data manifold.
• Improved Data Fidelity: Through the correction of discretization errors that plague standard Langevin dynamics.
• Efficient Reuse: The ability to enhance existing pretrained models without retraining.

The paper modestly acknowledges limitations, noting that the effectiveness of the samplers is bound by the quality of the underlying pretrained flow or score-based model. Approximation errors in the score model can accumulate over repeated kernel applications. Additionally, the current validation is restricted to continuous domains (images), and extending the paradigm to discrete domains like text generation is identified as future work. The authors emphasize that their contribution is the unification of conservation principles into a generative framework organized around stationary probing rather than reverse noising.