PaperlayerView exact revision on arXiv

Source-linked AI summary

Building Normalizing Flows with Stochastic Interpolants

Michael S. Albergo, Eric Vanden-Eijnden

arXiv:2209.15571v3cs.LGstat.ML

TL;DR

The paper addresses how to learn probability flows between base and target densities without coupling path selection to optimization or backpropagating through ODE solves. It introduces stochastic interpolants and a quadratic velocity objective, yielding explicit probability flows that compare well with diffusion models and extend ODE flows to 128×128 images.

  • Problem

    Learning a tractable probability flow between base and target densities requires choosing transport structure while estimating the associated velocity field.

  • Method

    The method specifies a time-differentiable stochastic interpolant whose density induces a velocity field learned by minimizing a simple quadratic objective, then solves the resulting probability flow equation.

  • Results

    The learned flow compares well with diffusions for image generation and demonstrates ab-initio ODE generation at 128×128 resolution.

  • Takeaways & Limitations

    Separating interpolant selection from objective minimization provides explicit control of the transport path and can support optimization toward optimal transport.

  • Takeaways & Limitations

    The associated score-based SDE has drift and diffusion coefficients that are singular at t = 0 and t = 1, although the velocity remains well-behaved.

Abstract

from arXiv · show

A generative model based on a continuous-time normalizing flow between any pair of base and target probability densities is proposed. The velocity field of this flow is inferred from the probability current of a time-dependent density that interpolates between the base and the target in finite time. Unlike conventional normalizing flow inference methods based the maximum likelihood principle, which require costly backpropagation through ODE solvers, our interpolant approach leads to a simple quadratic loss for the velocity itself which is expressed in terms of expectations that are readily amenable to empirical estimation. The flow can be used to generate samples from either the base or target, and to estimate the likelihood at any time along the interpolant. In addition, the flow can be optimized to minimize the path length of the interpolant density, thereby paving the way for building optimal transport maps. In situations where the base is a Gaussian density, we also show that the velocity of our normalizing flow can also be used to construct a diffusion model to sample the target as well as estimate its score. However, our approach shows that we can bypass this diffusion completely and work at the level of the probability flow with greater simplicity, opening an avenue for methods based solely on ordinary differential equations as an alternative to those based on stochastic differential equations. Benchmarking on density estimation tasks illustrates that the learned flow can match and surpass conventional continuous flows at a fraction of the cost, and compares well with diffusions on image generation on CIFAR-10 and ImageNet $32\times32$. The method scales ab-initio ODE flows to previously unreachable image resolutions, demonstrated up to $128\times128$.

1 INTRODUCTION

The paper formulates continuous-time transport between base and target densities through a probability-flow ODE built from a stochastic interpolant. Its key contribution is learning the velocity with a quadratic objective while separating optimization from path design and relating the learned flow to optimal transport and diffusion methods.

  • Transport formulation: Continuous-time generative modeling seeks a map whose flow transports a base density ρ0 to a target density ρ1 through intermediate densities ρt.The map is represented by an ODE with velocity field vt, equivalently satisfying the continuity equation.
  • Stochastic interpolants: The method introduces a time-differentiable interpolant It that equals x0 at t=0 and x1 at t=1, then samples xt by applying it to independent draws from ρ0 and ρ1.The resulting process xt is called a stochastic interpolant.
  • Learning the flow: The interpolant density has a velocity field that uniquely minimizes a simple quadratic objective, enabling a neural estimator and the InterFlow probability-flow ODE.This replaces direct maximum-likelihood optimization tied to choosing a transport path.
  • Learning the flow: Choosing the interpolant separates optimizing the objective from discovering a path between the base and target densities, unlike conventional maximum-likelihood flow training.The framework applies to any interpolant satisfying the endpoint conditions under additional regularity assumptions.
  • Guarantees and extensions: The learned velocity objective controls the Wasserstein-2 distance between the transported approximation and the target, while its training value provides a convergence diagnostic.The paper also proposes optimizing the stochastic interpolant to obtain optimal-transport solutions.
  • Related approaches: Unlike diffusion-based related methods, the approach directly constructs a finite-time probability-flow ODE and can be extended by optimizing the interpolant for optimal transport.The paper situates this construction against score-based flows, bridge methods, and prior interpolant approaches.

2 STOCHASTIC INTERPOLANTS AND ASSOCIATED FLOWS

Stochastic interpolants define finite-time density paths whose probability-flow velocity is learned through a unique quadratic minimization. The framework supports empirical estimation, transport optimization, flow-error bounds, and a Gaussian-base connection to score-based diffusion.

  • Proposition 1 identifies the interpolant velocity as the unique minimizer of a quadratic objective subject to the continuity equation.The objective is derived from the interpolant density and probability current.
  • The velocity objective is expressible as an expectation over sampled times and independently sampled endpoints, making it amenable to empirical estimation.This enables numerical minimization using samples from the base, target, and uniform time distribution.
  • The objective’s minimum can be optimized over interpolants to shorten the Wasserstein path, and under stated assumptions this procedure achieves optimal transport.The result relies on interpolability of the optimal density and existence of a classical solution.
  • The induced path has finite length even when it is not the shortest possible, while the quadratic objective can remain strictly positive at the exact velocity and therefore lacks a convergence baseline.These statements distinguish finiteness of the path from optimality and from direct convergence assessment.
  • A velocity approximation yields a pushforward density whose Wasserstein error is bounded under differentiability and uniform Lipschitz assumptions.The bound relates the target-density error to the objective H(ˆv).
  • With a standard Gaussian base and the trigonometric interpolant, velocity determines the interpolant score, but the associated SDE has singular drift and diffusion coefficients at both endpoints.The score relation permits, in principle, resampling intermediate densities or constructing score-based diffusion models.

3 PRACTICAL IMPLEMENTATION AND NUMERICAL EXPERIMENTS

InterFlow is trained by empirical optimization of a velocity-field objective and evaluated on two-dimensional, tabular, and image-generation tasks. The experiments show accurate multimodal transport, competitive likelihoods, and scaling to 128×128 images.

  • 2D density estimation: The velocity field is parameterized by a three-layer ReLU network with 256 hidden units per layer and optimized on G(v̂) for 10k epochs.The two-dimensional experiments use kernel density estimates over 80k generated samples.
  • 2D density estimation: InterFlow captures all modes of the target density without artificial stretching or smearing, indicating a smooth map on the toy distributions.The method connects the anisotropic swirls distribution to the checkerboard distribution and can sample intermediate times.
  • Dataset interpolation: The interpolant formulation learns velocity fields from arbitrary base and target densities and requires only samples, not query access to either density.This differs from the score-based diffusion perspective, where tractable training requires one distribution to be Gaussian.
  • Tabular data for higher dimensional testing: On tabular datasets, InterFlow achieves better or equivalent held-out likelihoods than all ODE-based models except BSDS300, where FFJORD is about 0.6% better.The experiments report improvements of up to 30% over baselines without directly optimizing likelihood.
  • Unconditional image generation: Image-generation likelihoods are competitive with diffusions on CIFAR-10 and ImageNet-32×32, with NLL values of 2.99 and 3.45.FID scores are close to diffusion results but slightly behind the best results; the model is also demonstrated on 128×128 Oxford flowers.

4 DISCUSSION, CHALLENGES, AND FUTURE WORK

The discussion emphasizes efficient training, separation of interpolation from transport-path selection, and direct finite-time probability-flow modeling. It also identifies links to optimal transport and score-based diffusion for future investigation.

  • Discussion: The method avoids backpropagation through ODE solves and provides a stable quadratic objective with a diagnostic for validating minimizers.The objective also controls the Wasserstein-2 distance between the model and target.
  • Discussion: Choosing the interpolant decouples optimization from transport-path selection, allowing finite-time interpolation between any two densities and direct likelihood calculation.The interpolant can additionally be optimized to achieve optimal transport and reduce the cost of solving the ODE for sampling.
  • Future work: The probability-flow ODE is distinguished from bridge-based and score-based approaches that produce diffusion-based models.The paper presents the relation to score-based diffusion as a basis for studying SDE–ODE trade-offs.
  • Technical conditions: The probability-flow construction relies on regularity conditions such as continuously differentiable densities and interpolants, with globally Lipschitz velocity assumptions used for flow existence and uniqueness.The appendix explains likelihood evaluation by evolving the probability flow backward to an earlier time.

B PROOF OF PROPOSITION 1

The proof establishes that the stochastic interpolant has a density and probability current satisfying the continuity equation, and that the associated velocity is characterized by quadratic minimization.

  • Density and current: Under Assumption B.1, the stochastic interpolant has a density and probability current, and its density satisfies the continuity equation with the prescribed endpoint densities.The current is obtained from the time derivative of the interpolant.
  • Velocity minimization: The velocity field defined from the probability current is the unique minimizer of the quadratic objective G(v̂) under the stated assumptions.The proof uses the identity v_t(x)ρ_t(x)=j_t(x) and establishes that the objective is well-defined.
  • Gradient-field formulation: The construction remains valid when optimization is restricted to gradient fields, whose minimizer is the gradient of a potential solving a Poisson equation.The potential is unique up to an additive constant.
  • Objective limitations: Even for gradient fields, the quadratic inequality is not generally saturated, so the objective’s minimum need not provide a direct convergence baseline.A Gaussian example demonstrates that the relevant inequality is not saturated in general.
  • Objective construction: The original objective avoids derivatives in x and t and permits direct optimization of the velocity field, whereas an alternative potential-based objective requires integration by parts.The alternative objective is related to the KILBO objective up to a sign and factor of two.

C THE CASE OF GAUSSIAN MIXTURE DENSITIES

For Gaussian mixture endpoints, the paper gives explicit interpolant densities and probability currents under linear interpolation, and characterizes the resulting velocity field.

  • Gaussian mixture setup: The Gaussian-mixture analysis assumes Gaussian components with positive mixture weights and uses a linear interpolant between endpoint densities.The interpolant density is given explicitly for this setting.
  • Gaussian mixture setup: When all covariance matrices are equal, the trigonometric interpolant preserves the common covariance matrix throughout the interpolation.This covariance preservation motivates the interpolant choice in the equal-covariance case.
  • Density and current: The interpolant density and probability current satisfy the continuity equation, with the current specified explicitly for Gaussian mixture densities.The proposition also yields consequences for the corresponding velocity field.
  • Velocity behavior: The resulting velocity grows at most linearly in x and approximately reduces to a mode-wise expression when Gaussian modes are well separated.This characterizes the velocity’s behavior in the separated-mode regime.

D OPTIMIZING TRANSPORT THROUGH STOCHASTIC INTERPOLANTS

The stochastic-interpolant velocity yields an exact transport map, but the resulting map is generally not optimal because the chosen interpolant density need not follow the Benamou–Brenier optimal density.

  • The velocity minimizing the stochastic-interpolant objective produces an exact transport map from ρ0 to ρ1.
  • The Benamou–Brenier problem seeks a transport minimizing the integrated kinetic energy over admissible density and velocity paths.
  • The stochastic-interpolant density is fixed by the selected interpolant and generally differs from the Benamou–Brenier optimal density.
  • Restricting velocities to gradient fields lowers the objective for the fixed interpolant density but does not generally produce an optimal transport map.
  • The paper therefore considers optimizing the interpolant, and possibly the base density, to improve transport optimality.

D.1 OPTIMAL TRANSPORT WITH OPTIMAL INTERPOLANTS

The paper formulates optimal transport as a max–min optimization over stochastic interpolants. Under interpolability and regularity assumptions, an optimal interpolant reproduces the Benamou–Brenier density and velocity.

  • Maximizing the minimum velocity objective over admissible interpolants targets the lower transport cost associated with the Benamou–Brenier problem.
  • An interpolable density is the pushforward of a standard Gaussian through a time-dependent invertible map that is continuously differentiable in time and space.
  • Any interpolable density can be represented as a stochastic-interpolant density, so the construction covers the optimal density when it is interpolable.
  • Under the proposition’s assumptions, a maximizing interpolant has probability density ρ*_t and minimizing velocity v*_t satisfying the Benamou–Brenier equations.
  • The max–min reformulation is equivalent to the Benamou–Brenier minimization through convexity, concavity, and the minimax theorem.
  • In practice, restricting the interpolant class makes optimization feasible but means the resulting map is no longer guaranteed to be optimal.
  • Rectification methods can approach optimality through iteration, but imperfect learned maps introduce bias that grows with iterations and can cause instability.
  • The base density can also be optimized, although optimizing it alone is insufficient when the interpolant remains fixed.

E PROOF OF PROPOSITION 3

The proof establishes a stability bound for the transport construction using differential inequalities and Gronwall’s inequality, while noting an optional regularization that was not used numerically.

  • Differential inequalities bound the evolution of the discrepancy between trajectories and velocity fields.
  • Gronwall’s inequality converts these differential bounds into a stability estimate from the zero initial discrepancy.
  • The objective can be regularized with a small λ > 0, but the paper’s numerical results omit this regularization.

F PROOF OF PROPOSITION 4 AND LINK WITH SCORE-BASED DIFFUSION

The paper connects stochastic-interpolant probability flows to score-based diffusion and extends the construction with latent mixture components. The ODE velocity remains finite at endpoints where the associated diffusion becomes singular.

  • F PROOF OF PROPOSITION 4 AND LINK WITH SCORE-BASED DIFFUSION: For interpolants of the specified form, the probability current can be manipulated to recover the score ∇log ρt(x) and establish the diffusion connection.
  • F PROOF OF PROPOSITION 4 AND LINK WITH SCORE-BASED DIFFUSION: With a Gaussian base, the continuity equation can be rewritten as a diffusion equation whose diffusion coefficient is negative under the stated coefficient conditions.
  • F PROOF OF PROPOSITION 4 AND LINK WITH SCORE-BASED DIFFUSION: The resulting equation is well-posed backward in time, and reversing it gives an SDE transforming the base density into the target.
  • F PROOF OF PROPOSITION 4 AND LINK WITH SCORE-BASED DIFFUSION: The diffusion drift and coefficient are singular at t = 1, whereas the probability-flow velocity remains finite at both endpoints.
  • Generalizations: Additional latent variables can enlarge the class of interpolant densities and potentially yield simpler or more optimal velocity fields.
  • Generalizations: A K-component construction softly clusters source and target samples, interpolating corresponding clusters through separate interpolants.
  • Generalizations: For mixture bases, the objective can be optimized over component means and covariances when target cluster masses are known.

H EXPERIMENTS FOR OPTIMAL TRANSPORT, PARAMETERIZING It, PARAMETERIZING ρ0

The experiments optimize the interpolant and, separately, the base density to shorten transport paths and improve likelihood estimation. Learned interpolants reduce path length and make the resulting velocity field easier to learn.

  • Parameterizing It: Optimizing a Fourier-parameterized interpolant reduces path length and improves likelihood estimation for checkerboard density estimation.The experiment uses M = 7 Fourier coefficients and compares learned with nonparametric interpolants.
  • Parameterizing It: The learned interpolant changes the interpolant terms ˆat and ˆbt relative to their initializations during training.
  • Parameterizing It: Minimizing transport cost can produce models that are practically easier to learn.
  • Parameterizing ρ0: Parameterizing the base density as a Gaussian with learned mean and covariance further reduces transport cost and improves final log likelihood.The learned base density brings the model closer to optimal transport.

I IMPLEMENTATION DETAILS FOR NUMERICAL EXPERIMENTS

The numerical experiments use empirical objective estimates, tabular and image-model configurations, and implementation choices for time sampling and neural architectures. These details include parallelizable objective evaluation and U-Net-based image models trained on a single A100 GPU.

  • The empirical estimate of the objective is computed from samples drawn from the uniform density on [0, 1].
  • The objective calculation is parallelizable.
  • Tabular experiments use dataset-specific architectures and hyperparameters summarized in Table 3.ReLU activations are used throughout except for BSDS300, which uses ELU activations.
  • Time sampling is reweighted with a Beta distribution under the heuristic that the flow should be well trained near the target.The reweighting preserves the same objective minimizer for any weight w(t).
  • Image models use a DDPM-style U-Net with sinusoidal time embeddings and are implemented on a single A100 GPU.The U-Net dimension multipliers are changed according to Table 4.

I.1 DETAILS ON COMPUTATIONAL EFFICIENCY AND DEMONSTRATION OF CONVERGENCE GUARANTEE

The experiments assess computational efficiency and convergence using training-speed comparisons and a shifted objective diagnostic. The diagnostic reaches zero on POWER, while image demonstrations include CIFAR-10, ImageNet 32×32, and Oxford Flowers at 128×128.

  • Computational efficiency: 400x speedup over maximum-likelihood ODE flows is reported on MiniBooNE, alongside more efficient likelihood ascent for InterFlow.
  • Convergence guarantee: The shifted objective tracks whether the model velocity meets the minimizer requirement tied to the squared path length.
  • Convergence guarantee: The shifted objective ˜G(ˆv) converges to ˜G(v) = 0 and remains there during POWER training.This behavior is used as a convergence signature, though the diagnostic is necessary but not sufficient for convergence.
  • Image demonstrations: Image-generation demonstrations include uncurated samples from CIFAR-10, ImageNet 32×32, and Oxford Flowers at 128×128.
Loading 2209.15571v3…