PaperlayerView exact revision on arXiv

Source-linked AI summary

OT-Flow: Fast and Accurate Continuous Normalizing Flows via Optimal Transport

Derek Onken, Samy Wu Fung, Xingjian Li, Lars Ruthotto

arXiv:2006.00104v5cs.LGstat.ML

TL;DR

CNFs offer flexible invertible mappings but are computationally costly because ODE integration may require many evaluations and exact Jacobian-trace computation is challenging. OT-Flow regularizes CNFs with optimal transport for straighter trajectories and provides exact trace computation, achieving similar quality with average 8x training and 24x inference speedups.

  • Problem

    CNFs are computationally expensive because ODE integration can require many evaluations and computing the Jacobian trace without forming the full Jacobian is challenging.

  • Method

    OT-Flow formulates CNFs as an optimal transport problem, uses HJB-based regularization to reduce time steps, and analytically computes the exact Jacobian trace with O(d) FLOPS.

  • Results

    OT-Flow produces results of similar quality to state-of-the-art CNFs while achieving average 8x training and 24x inference speedups.

  • Takeaways & Limitations

    OT-Flow combines OT-influenced regularization, reduced parameterization, discretize-then-optimize training, and exact trace computation for fast and accurate CNF training and testing.

  • Takeaways & Limitations

    Testing loss can be low even when the learned forward map differs substantially from the target distribution, and inverse-flow error can remain.

Abstract

from arXiv · show

A normalizing flow is an invertible mapping between an arbitrary probability distribution and a standard normal distribution; it can be used for density estimation and statistical inference. Computing the flow follows the change of variables formula and thus requires invertibility of the mapping and an efficient way to compute the determinant of its Jacobian. To satisfy these requirements, normalizing flows typically consist of carefully chosen components. Continuous normalizing flows (CNFs) are mappings obtained by solving a neural ordinary differential equation (ODE). The neural ODE's dynamics can be chosen almost arbitrarily while ensuring invertibility. Moreover, the log-determinant of the flow's Jacobian can be obtained by integrating the trace of the dynamics' Jacobian along the flow. Our proposed OT-Flow approach tackles two critical computational challenges that limit a more widespread use of CNFs. First, OT-Flow leverages optimal transport (OT) theory to regularize the CNF and enforce straight trajectories that are easier to integrate. Second, OT-Flow features exact trace computation with time complexity equal to trace estimators used in existing CNFs. On five high-dimensional density estimation and generative modeling tasks, OT-Flow performs competitively to state-of-the-art CNFs while on average requiring one-fourth of the number of weights with an 8x speedup in training time and 24x speedup in inference.

1 Introduction

Normalizing flows require invertible mappings and efficient Jacobian-determinant computation, while CNFs can incur substantial ODE and trace-computation costs. OT-Flow addresses these challenges by regularizing trajectories and computing traces exactly with competitive complexity.

  • A normalizing flow maps an arbitrary distribution to a standard normal through an invertible function and uses the change-of-variables formula.
  • O(d3) log-determinant computation motivates flow layers with exploitable Jacobian structure.
  • CNFs obtain invertible mappings by solving neural ODEs, with log-determinants accumulated through the dynamics’ Jacobian trace.
  • CNFs can be expensive because ODE solvers may require many dynamics evaluations and trace computation without explicitly building the Jacobian is challenging.
  • OT-Flow uses optimal transport regularization to encourage straight trajectories, reducing the time steps needed for ODE integration.
  • O(d) exact trace computation matches the complexity of one-vector Hutchinson estimation, while OT-Flow achieves similar quality with 8x faster training and 24x faster inference on average.

2 Mathematical Formulation of OT-Flow

OT-Flow reformulates CNF training as a regularized optimal transport problem and uses potential-flow structure to encourage straight, smooth trajectories. An HJB penalty further guides the learned dynamics and reduces practical integration cost.

  • OT-Flow adds an OT transport cost and HJB-violation penalty to the CNF objective, encouraging straight trajectories.
  • The L2 transport cost penalizes squared trajectory arc-length and can be accumulated within the ODE solver.
  • The relaxed OT formulation is convex before neural parameterization and yields straight, non-intersecting trajectories while preserving smoothness and invertibility.
  • OT theory supplies a potential function whose gradient defines the dynamics and whose HJB equations characterize the optimality structure.
  • The HJB regularizer favors plausible potentials without changing the solution of the regularized OT optimization problem.
  • Without HJB regularization, few-step models can have insufficiently penalized transport cost, poor performance, and unstraight trajectories.

3 Implementation of OT-Flow

OT-Flow implements a potential-based CNF with efficient ODE optimization and exact trace computation. Its design supports competitive computation while avoiding stochastic trace-estimation error.

  • OT-Flow defines the model, derives analytic formulas for fast exact trace computation, and describes an efficient ODE solver.
  • The implementation compares methods by potential Φ, regularizers, optimization approach, ODE solver, and trace computation.
  • Network: The potential combines a neural network with quadratic terms, using A, b, and c for linear dynamics and N for nonlinear dynamics.
  • ResNet: The two-layer ResNet maps space-time inputs into hidden features and applies a residual update with element-wise activation functions.
  • Trace Computation: The exact trace uses O(m · d) FLOPS for the first layer and O(m2 · d) FLOPS for each additional layer, matching the complexity of trace estimation.It avoids building the full Hessian by exploiting the identity structure of matrix E.
  • Results: Exact trace computation converges more quickly, reaches lower validation loss, and produces less variable training loss than a single-vector Hutchinson estimator.
  • Optimization: Automatic differentiation supports the discretize-then-optimize approach for backpropagation through the neural ODE.

4 Related Works

Related flow models use discrete architectures, neural ODEs, stochastic trace estimation, or optimal-transport regularization. OT-Flow differs through its potential-flow formulation and scalable exact trace computation.

  • Finite Flows: Finite normalizing flows use autoregressive or coupling architectures whose triangular Jacobians make determinant computation tractable.
  • Infinitesimal Flows: CNFs model flows with differential equations, providing an alternative to compositions of discrete transformations.
  • Trace Estimation: FFJORD reduces training trace cost from O(d2) to O(d) with Hutchinson estimation, but accurate inference requires exact trace computation.
  • Experimental Comparisons: Table 2 compares training iterations, forward-ODE function evaluations, and time per iteration across real-data density-estimation methods.
  • Flows Influenced by Optimal Transport: RNODE combines FFJORD with transport and Frobenius-norm regularization and reports 2.8x speedup.
  • Flows Influenced by Optimal Transport: OT-Flow shares RNODE’s L2 cost but uses a potential-flow approach and exact trace computation, avoiding the need for Frobenius-norm regularization in its stable training.

5 Numerical Experiments

The experiments evaluate OT-Flow on toy, real-data density estimation, and MNIST generation tasks, using MMD, loss, runtime, visualizations, and inverse error. OT-Flow achieves comparable real-data performance with substantially reduced training time, while supporting conditioned generation and latent interpolation.

  • OT-Flow is evaluated on seven two-dimensional toy problems, five real-data problems, and MNIST generation.
  • MMD compares generated samples with held-out data, and lower MMD indicates greater distributional agreement.MMD is not used during training, providing an external evaluation metric.
  • Density Estimation on Real Data Sets: OT-Flow reduces training time while maintaining competitive real-data density-estimation performance relative to FFJORD and RNODE.The comparisons use MMD, loss, training speed, testing speed, and visual samples.
  • Density Estimation on Real Data Sets: OT-Flow achieves comparable MMD to state-of-the-art methods on real data except GAS, where it learns a slightly smoothed representation.
  • OT-Flow's exact trace computation contributes to faster testing, while inverse error indicates numerical invertibility and approximation of the ODE solution.
  • MNIST: MNIST experiments generate class-conditioned images with a 64-dimensional flow and latent interpolations with an unconditioned 128-dimensional flow.Both procedures invert the flow and decode latent points into pixel-space images.

6 Discussion

The discussion presents OT-Flow as a fast CNF approach that addresses costly ODE integration and trace computation. OT regularization enables few time steps, while exact traces improve convergence at estimator-comparable complexity.

  • OT-Flow targets two computational challenges in CNFs: many ODE time steps and expensive trace computation.
  • Adding transport cost and an HJB regularizer allows OT-Flow to use few time steps without sacrificing performance.
  • Exact trace computation has complexity and cost comparable to state-of-the-art trace estimators, while providing better convergence.
  • The analytic gradient and trace approach is not limited to ResNet architectures, but other architectures require further derivation.

A Derivation of Loss C

This derivation connects CNF training to an optimal-transport formulation and uses HJB optimality conditions to motivate a trajectory regularizer. Because the initial density is unknown, the objective drops its associated constant term.

  • The flow trajectories map samples from ρ0 to ρ1, and their density evolution follows the change-of-variables formula.
  • The CNF objective is related to KL divergence between the final flowed distribution and the target normal distribution.
  • The unknown initial density makes log(ρ0(x)) a constant that can be dropped without changing the minimizer.
  • The OT formulation introduces a transport cost and a potential whose dynamics satisfy Hamilton-Jacobi-Bellman equations.
  • The HJB regularizer penalizes equation violations along trajectories because solving the high-dimensional terminal condition is difficult.
  • The HJB regularizer enables a two-step RK4 flow to produce straighter trajectories than an unregularized two-step model in the toy demonstration.

C Error Bounds

The timing and density-estimation results use repeated runs and summarize variation differently. Trace-runtime uncertainty is estimated by bootstrap confidence intervals, while real-data results report means and standard deviations across instances.

  • Trace timing uses 20 replications and 4,000 bootstrap resamplings of size 16 to form 99% confidence intervals.
  • Real-data density-estimation results average three trained instances for each model and dataset.Standard deviations for the three instances are reported separately.

D Implementation Details

The implementation combines residual-network optimization, exact trace computation, and empirical comparisons across density-estimation models. OT-Flow's parameterization and training choices support reduced computational cost while retaining expressibility.

  • Implementation: OT-Flow optimizes weights parameterizing a potential, with hyperparameters controlling regularization, hidden-space size, solver steps, and optimization settings.The experiments use a Runge-Kutta 4 solver and two regularization hyperparameters.
  • Implementation: The residual network propagates space-time inputs through an initial nonlinear layer followed by residual layers with fixed step size h > 0.The network uses K0, K1, ..., KM and biases b0, ..., bM as trainable layer parameters.
  • Exact trace computation: Exact trace computation is performed in one forward pass, using O(m · d) FLOPS for the first layer and continued updates for subsequent layers.The method exploits the identity structure of matrix E rather than constructing the full Hessian.
  • Model comparisons: OT-Flow requires fewer parameters than many discrete normalizing flows, while reduced-parameter FFJORD and RNODE models lacked sufficient expressibility.Those reduced-parameter models converged to poor MMD and loss values in the reported experiments.
  • Evaluation: Testing-loss comparisons include discrete-flow and other-model baselines, while evaluation also reports inverse error and MMD.The POWER example shows that a competitive testing loss can coexist with poor visual quality, MMD, and inverse error.

G Loss Metric

Testing loss can be misleading because it depends on trace integration accuracy and may not reflect invertibility or generation quality. The paper therefore emphasizes MMD and visual inspection for evaluation.

  • Metric limitations: Testing loss is unreliable for comparing flows because coarse trace integration can produce a low value while sacrificing invertibility.The metric depends on integrating the trace along the computed trajectory.
  • Metric limitations: Similar testing losses can conceal poor generation, as FFJORD pushes too many points to the origin and maps poorly to a Gaussian.The comparison is made using visualizations of 10^5 samples.
  • Metric limitations: A POWER model achieved testing loss −0.64 while its visualization, MMD, and inverse error indicated flawed density estimation.The high inverse error suggests that the integration was not trustworthy.

H Visualizations of High-Dimensional Data Sets

The appendix provides additional visualizations for model performance on POWER, GAS, HEPMASS, and MINIBOONE datasets. These figures extend the paper's inspection of high-dimensional density-estimation behavior.

  • POWER: Figure A4 visualizes model performance on POWER test data.Use it to inspect the reported POWER test-data behavior alongside other evaluation metrics.
  • GAS: Figure A5 visualizes model performance on GAS test data.Use it to inspect the reported GAS test-data behavior.
  • HEPMASS: Figure A6 visualizes model performance on HEPMASS test data.Use it to inspect the reported HEPMASS test-data behavior.
  • MINIBOONE: Figure A7 adds two-dimensional slices of the MINIBOONE density estimation to supplement Figure 5.The slices provide additional views of the high-dimensional point cloud.
Loading 2006.00104v5…