Improving and generalizing flow-based generative models with minibatch optimal transport
Pith reviewed 2026-05-12 12:47 UTC · model grok-4.3
The pith
Conditional flow matching trains continuous normalizing flows without simulation or Gaussian assumptions.
A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.
Core claim
Generalized conditional flow matching (CFM) is a simulation-free training objective for continuous normalizing flows that regresses the vector field onto a target derived from a conditional path between source and target samples. The optimal transport CFM (OT-CFM) variant constructs these paths via minibatch optimal transport, yielding simpler flows that train stably, support faster inference, and approximate dynamic optimal transport when the exact plan is known.
What carries the argument
Conditional flow matching (CFM), a regression objective that learns the vector field of a continuous normalizing flow by matching to a conditional path, with optimal transport used to select straight paths between samples.
If this is right
- Continuous normalizing flows can be trained with a stable regression loss similar to diffusion models while keeping deterministic and fast inference.
- The source distribution can be arbitrary and its density does not need to be evaluated during training.
- OT-CFM produces simpler flows that require fewer integration steps at inference time.
- When the true optimal transport plan is known, the learned model approximates dynamic optimal transport.
- Performance gains appear on both unconditional generation and conditional tasks such as image translation and single-cell trajectory inference.
Where Pith is reading between the lines
- The ability to use non-Gaussian sources may allow flow models to incorporate domain-specific priors more directly than diffusion approaches.
- Straighter paths from OT-CFM could reduce the sensitivity of flow models to numerical integration errors in high dimensions.
- The link to dynamic optimal transport opens the possibility of using trained flows to solve transport problems in new domains.
Load-bearing premise
Minibatch optimal transport plans computed from finite samples are close enough to the true continuous optimal transport plan.
What would settle it
On a low-dimensional problem where the exact optimal transport plan between source and target can be computed analytically, measure whether OT-CFM using minibatch plans produces flows and samples nearly identical to those from the exact plan.
read the original abstract
Continuous normalizing flows (CNFs) are an attractive generative modeling technique, but they have been held back by limitations in their simulation-based maximum likelihood training. We introduce the generalized conditional flow matching (CFM) technique, a family of simulation-free training objectives for CNFs. CFM features a stable regression objective like that used to train the stochastic flow in diffusion models but enjoys the efficient inference of deterministic flow models. In contrast to both diffusion models and prior CNF training algorithms, CFM does not require the source distribution to be Gaussian or require evaluation of its density. A variant of our objective is optimal transport CFM (OT-CFM), which creates simpler flows that are more stable to train and lead to faster inference, as evaluated in our experiments. Furthermore, we show that when the true OT plan is available, our OT-CFM method approximates dynamic OT. Training CNFs with CFM improves results on a variety of conditional and unconditional generation tasks, such as inferring single cell dynamics, unsupervised image translation, and Schr\"odinger bridge inference.
Editorial analysis
A structured set of objections, weighed in public.
Referee Report
Summary. The manuscript introduces a generalized conditional flow matching (CFM) framework for simulation-free training of continuous normalizing flows (CNFs). It defines a family of regression objectives that avoid the need for Gaussian source distributions or source density evaluation. A prominent variant, optimal transport CFM (OT-CFM), constructs targets from minibatch optimal transport plans between source and target samples; the authors state that this yields simpler, more stable flows with faster inference. They further claim that OT-CFM approximates dynamic optimal transport when the true (population) OT plan is available, and report empirical gains on tasks including single-cell dynamics inference, unsupervised image translation, and Schrödinger bridge problems.
Significance. If the approximation result and empirical improvements hold under rigorous verification, the work supplies a practical, flexible alternative to both diffusion-style training and earlier CNF objectives. By removing the Gaussian-source restriction and incorporating minibatch OT for path simplification, it could streamline training of deterministic flows while retaining efficient inference. The connection to dynamic OT and the reported gains on scientific and vision tasks would be of interest to the generative modeling community.
major comments (2)
- [Abstract and theoretical analysis of OT-CFM] Abstract and the theoretical section on OT-CFM: the claim that OT-CFM approximates dynamic OT when the true plan is available is central to the paper's theoretical contribution. The implemented algorithm uses minibatch OT plans computed on finite samples, yet no error bound, convergence rate, or sensitivity analysis with respect to batch size is supplied to show that the minibatch coupling yields a velocity field sufficiently close to the dynamic OT solution. This gap directly affects whether the approximation statement can be transferred to the practical method.
- [Experiments] Experimental section (reported results on single-cell and image tasks): the soundness assessment notes that ablation details and statistical testing are not fully verifiable from the provided material. Without explicit variance estimates across multiple runs or controls isolating the effect of minibatch size on the learned vector field, it is difficult to confirm that the claimed stability and speed improvements survive rigorous evaluation.
minor comments (2)
- [Method] Notation for the generalized CFM objective could be clarified by explicitly distinguishing the regression target derived from the OT plan versus the data-only case.
- [Figures] Figure captions and axis labels in the inference-time and stability plots should include batch-size values used for the minibatch OT computations.
Simulated Author's Rebuttal
We thank the referee for their constructive and detailed comments. We address each major point below with clarifications and indicate where revisions have been or will be made to the manuscript.
read point-by-point responses
-
Referee: [Abstract and theoretical analysis of OT-CFM] Abstract and the theoretical section on OT-CFM: the claim that OT-CFM approximates dynamic OT when the true plan is available is central to the paper's theoretical contribution. The implemented algorithm uses minibatch OT plans computed on finite samples, yet no error bound, convergence rate, or sensitivity analysis with respect to batch size is supplied to show that the minibatch coupling yields a velocity field sufficiently close to the dynamic OT solution. This gap directly affects whether the approximation statement can be transferred to the practical method.
Authors: The theoretical result (Section 3.3) establishes that OT-CFM with the population-level OT plan yields a velocity field that approximates the dynamic OT solution; the proof proceeds by showing that the conditional OT interpolation produces the same marginal velocity as the dynamic formulation. The manuscript does not claim this equivalence holds exactly for the minibatch estimator used in practice. Minibatch OT is presented as a computationally tractable surrogate whose empirical behavior (simpler paths, stable training) is validated separately in the experiments. We agree that a formal sensitivity analysis or error bound relating minibatch size to the population solution would be desirable, but deriving such a bound requires additional regularity assumptions on the data distribution that lie outside the scope of this work. In the revision we have added explicit wording in the abstract and theory section distinguishing the population guarantee from the minibatch implementation, together with a short paragraph discussing the empirical justification for the minibatch approximation. revision: partial
-
Referee: [Experiments] Experimental section (reported results on single-cell and image tasks): the soundness assessment notes that ablation details and statistical testing are not fully verifiable from the provided material. Without explicit variance estimates across multiple runs or controls isolating the effect of minibatch size on the learned vector field, it is difficult to confirm that the claimed stability and speed improvements survive rigorous evaluation.
Authors: We acknowledge the need for greater statistical transparency. The revised manuscript now reports mean and standard deviation over five independent random seeds for all quantitative metrics on the single-cell and image-translation benchmarks. We have also inserted a new ablation subsection that varies minibatch size (32, 64, 128, 256) while holding all other hyperparameters fixed, and tabulates the resulting effects on (i) training-loss variance (as a proxy for stability) and (ii) wall-clock inference time. These controls directly isolate the contribution of the minibatch OT plan and support the claims of improved stability and faster inference. revision: yes
Circularity Check
No circularity: CFM objective and OT-CFM approximation are independently derived
full rationale
The paper introduces CFM as a regression-based training objective for CNFs that regresses to conditional vector fields constructed from data couplings (including OT plans for the OT-CFM variant). The statement that OT-CFM approximates dynamic OT when the true plan is available is presented as a mathematical result shown from the definitions of the objective and the Benamou-Brenier formulation, without reducing to a fitted parameter, self-citation, or renaming of inputs. No load-bearing step equates a claimed prediction to its own construction; the minibatch implementation is an approximation whose error is left unquantified but does not create definitional circularity in the core derivation.
Axiom & Free-Parameter Ledger
axioms (1)
- domain assumption The probability path admits a well-defined time-dependent vector field that can be regressed against a target derived from conditional or OT plans.
Forward citations
Cited by 60 Pith papers
-
What Time Is It? How Data Geometry Makes Time Conditioning Optional for Flow Matching
Data geometry makes time identifiable from noisy interpolants at rate O(1/sqrt(d-k)), rendering the time-blindness gap asymptotically negligible relative to coupling variance.
-
Generative Modeling with Flux Matching
Flux Matching generalizes score-based generative modeling by using a weaker objective that admits infinitely many non-conservative vector fields with the data as stationary distribution, enabling new design choices be...
-
Generative models on phase space
Generative diffusion and flow models are constructed to remain exactly on the Lorentz-invariant massless N-particle phase space manifold during sampling for particle physics applications.
-
FlowHijack: A Dynamics-Aware Backdoor Attack on Flow-Matching Vision-Language-Action Models
FlowHijack is the first dynamics-aware backdoor attack on flow-matching VLAs that achieves high success rates with stealthy triggers while preserving benign performance and making malicious actions kinematically indis...
-
Cross-Space Distillation: Teaching One-Step Students with Modern Diffusion Teachers
Introduces a Bridge latent interface that maps mismatched student latents into teacher space, enabling distillation from modern diffusion teachers to compact one-step students and raising SD 1.5 HPSv3 from 5.4 to 9.4 ...
-
A Distributionally Robust Framework for Learned Reconstructions in Inverse Problems
Introduces structured DRO for learned inverse problem reconstructions with ambiguity sets aligned to the forward operator, yielding explicit dual representations and a worst-case bound that induces Tikhonov regulariza...
-
Bridging Vision and Language Concepts through Optimal Transport Semantic Flow
OTF-CBM replaces static cosine similarity in vision-language CBMs with data-driven optimal transport flow to improve concept alignment, accuracy, and faithfulness.
-
Patched Flow Matching: Generative Wall-Pressure Reconstruction Beyond Training-Domain Scales from Sparse Sensors
Patched Flow Matching reconstructs full-resolution wall-pressure fields on domains four times larger than training data from 0.25% sensor coverage by fusing short-domain DNS patch priors with sparse measurements via t...
-
Intrinsic Flow Matching on Quantum Pure-State Manifolds with Phase-Aligned Transport
IFM learns deterministic tangent velocity fields on CP^{d-1} via Pancharatnam phase-aligned paths, recovering marginal transport with endpoint and stability guarantees while showing empirical gains over Euclidean flow...
-
TriFlow: Generating Artist-Like 3D Mesh Topology via Nearest-Vertex Vector Fields
TriFlow synthesizes nearest-vertex vector fields via flow-matching to generate artist-like 3D mesh topology, then extracts meshes via clustering and topology-aware QEM simplification.
-
Start Right, Arrive Right: Asynchronous Execution via Initial Noise Selection
PAINT reframes asynchronous flow-based action chunking as an initial noise selection problem solved via backward Euler inversion and a repainting rule.
-
Learning to Distort: Weakly-Supervised Image Quality Transfer for Prostate DWI Correction
A weakly-supervised image quality transfer method generates synthetic distorted DWI images from quality labels to train improved distortion correction models for prostate MRI.
-
Learning Individual Dynamics from Sparse Cross-Sectional Snapshots
CADENCE recovers individualized continuous-time trajectories from cross-sectional snapshots via context-anchored latent dynamics, a bijective score-based encoder, and SMoE routing, with claimed identifiability guarant...
-
Increasing the Precision of Surrogate Models for Weak Lensing Mass Maps with Flow Matching
A flow matching generative model produces weak lensing mass maps with fidelity improved to below 1% and 5% on basic and higher-order statistics relative to GAN benchmarks.
-
SURGE: Approximation and Training Free Particle Filter for Diffusion Surrogate
URGE performs unbiased path-wise importance reweighting via Girsanov estimation for derivative-free inference-time scaling in diffusion models, proving equivalence to particle-wise SMC and outperforming baselines empirically.
-
PACE: Geometry-Aware Bridge Transport for Single-Cell Trajectory Inference
PACE recovers geometry-consistent continuous transport dynamics from single-cell time-course snapshots via state-time dependent anisotropic Riemannian metrics, alternating cross-time couplings, and neural bridges, out...
-
Learning Unbiased Permutations via Flow Matching
PermFlow applies conditional flow matching on the affine subspace of doubly stochastic matrices with a closed-form tangent projector and nearest-target coupling to capture multimodal permutation distributions.
-
The Velocity Deficit: Initial Energy Injection for Flow Matching
Flow matching underestimates velocities due to MSE loss leading to integration lag; Initial Energy Injection corrects the start-end asymmetry, improving FID by 44.6% and achieving 5x speedup on ImageNet-1k.
-
Aligning Flow Map Policies with Optimal Q-Guidance
Flow map policies enable fast one-step inference for flow-based RL policies, and FMQ provides an optimal closed-form Q-guided target for offline-to-online adaptation under trust-region constraints, achieving SOTA performance.
-
Generative Transfer for Entropic Optimal Transport with Unknown Costs
A generative transfer framework using iterative path-wise tilting integrated with conditional flow matching recovers target entropic optimal transport couplings from reference samples, achieving O(δ) convergence in Wa...
-
A Call to Lagrangian Action: Learning Population Mechanics from Temporal Snapshots
Wasserstein Lagrangian Mechanics learns second-order population dynamics from observed marginal snapshots without specifying the Lagrangian and outperforms gradient flow methods on tasks like vortex dynamics and embry...
-
A Call to Lagrangian Action: Learning Population Mechanics from Temporal Snapshots
Wasserstein Lagrangian Mechanics learns second-order population dynamics from observed marginals without specifying the Lagrangian and outperforms gradient flow methods on periodic dynamics like vortex motion and flocking.
-
A Call to Lagrangian Action: Learning Population Mechanics from Temporal Snapshots
Wasserstein Lagrangian Mechanics formalizes second-order dynamics in Wasserstein space and provides an algorithm to learn them from observed marginals without specifying the Lagrangian, outperforming gradient flows on...
-
Curated Synthetic Data Doesn't Have to Collapse: A Theoretical Study of Generative Retraining with Pluralistic Preferences
Recursive generative retraining with pluralistic preferences converges to a stable diverse distribution that satisfies a weighted Nash bargaining solution.
-
Stochastic Transition-Map Distillation for Fast Probabilistic Inference
STMD distills the full transition map of diffusion sampling SDEs into a conditional Mean Flow model to enable fast one- or few-step stochastic sampling without teacher models or bi-level optimization.
-
SDFlow: Similarity-Driven Flow Matching for Time Series Generation
SDFlow uses similarity-driven flow matching with low-rank manifold decomposition and a categorical posterior to generate high-fidelity long time series in VQ space without step-wise error accumulation.
-
FluxFlow: Conservative Flow-Matching for Astronomical Image Super-Resolution
FluxFlow is a conservative pixel-space flow-matching framework for astronomical super-resolution that incorporates real atmospheric uncertainty and a training-free Wiener correction, outperforming baselines on a new 1...
-
Generative Modeling with Orbit-Space Particle Flow Matching
OGPP is a particle flow-matching method using orbit-space canonicalization and geometric paths that achieves lower error and fewer steps than prior approaches on 3D benchmarks.
-
Talker-T2AV: Joint Talking Audio-Video Generation with Autoregressive Diffusion Modeling
Talker-T2AV achieves better lip-sync accuracy, video quality, and audio quality than dual-branch baselines by separating high-level shared autoregressive modeling from modality-specific low-level diffusion refinement ...
-
Self-Improving Tabular Language Models via Iterative Reward-Guided Post-Training
TabGRAA enables self-improving tabular language models through iterative group-relative advantage alignment using modular automated quality signals like distinguishability classifiers.
-
FlowEqProp: Training Flow Matching Generative Models with Gradient Equilibrium Propagation
FlowEqProp trains flow matching generative models using gradient equilibrium propagation on a 25k-parameter MLP for digit generation without backpropagation, producing recognizable samples and allowing quality gains f...
-
Generative modeling of granular flow on inclined planes using conditional flow matching
A conditional flow matching model trained on DEM simulations reconstructs granular flow velocity fields from as little as 11-16% sparse boundary data, outperforming deterministic CNN baselines while providing uncertai...
-
Is Flow Matching Just Trajectory Replay for Sequential Data?
Flow matching on time series targets a closed-form nonparametric velocity field that is a similarity-weighted mixture of observed transition velocities, making neural models approximations to an ideal memory-augmented...
-
DisRFM: Polar Riemannian Flow Matching for Structure-Preserving Graph Domain Adaptation
DisRFM uses polar Riemannian flow matching on constant-curvature manifolds to align graph domains while preserving label-relevant topology via radial Wasserstein and angular confidence matching.
-
On The Hidden Biases of Flow Matching Samplers
Empirical flow matching introduces coupled biases from plug-in estimation, including altered statistical targets, non-gradient minimizers, and non-unique dynamics via flux-null fields, with base distribution controlli...
-
NEAT: Neighborhood-Guided, Efficient, Autoregressive Set Transformer for 3D Molecular Generation
NEAT achieves state-of-the-art 3D molecular generation on QM9 and GEOM-Drugs via a neighborhood-guided autoregressive set transformer that ensures atom-level permutation invariance and offers a significant speed advantage.
-
Sundial: A Family of Highly Capable Time Series Foundation Models
Sundial uses TimeFlow Loss for native pre-training of Transformers on continuous time series from TimeBench, achieving SOTA point and probabilistic forecasting with millisecond inference.
-
Stochastic Interpolants: A Unifying Framework for Flows and Diffusions
Stochastic interpolants unify flow-based and diffusion-based generative models by bridging target densities exactly via latent-variable processes whose drifts minimize quadratic objectives.
-
Generative Models on Analog Hardware with Dynamics
AIS framework with time-varying piecewise parameters and hidden states enables analog hardware to perform generative modeling, achieving FID 27.6 on MNIST and 80.8 on Fashion-MNIST with 23uJ per image using 4-bit spar...
-
CurvSegFlow: Time-Conditioned Flow Matching for Robust Segmentation of Curvilinear Structures in Noisy Biomedical Images
CurvSegFlow applies time-conditioned flow matching with a U-Net backbone and triple-term loss to progressively refine segmentations of thin structures in noisy images, reporting competitive performance on microtubule,...
-
Benchmarking Neural Speech Compression from a Rate-Distortion Perspective
ECC integrates hyperprior side information, channel-wise context, latent residual prediction, temporal modeling, and entropy skip into a learned entropy model, yielding 39.9% and 76.3% average BD-rate reductions on Vi...
-
GenPO++: Generative Policy Optimization with Jacobian-free Likelihood Ratios
GenPO++ achieves exact Jacobian-free likelihood ratio computation for generative flow policies by embedding history states as auxiliary memory in a high-order reversible ODE solver.
-
GenSBI: Generative Methods for Simulation-Based Inference in JAX
GenSBI delivers JAX-native implementations of generative SBI methods with transformer backbones and reports near-ideal calibration scores on standard benchmarks.
-
SURGE: Approximation and Training Free Particle Filter for Diffusion Surrogate
SURGE is an unbiased particle filter that fuses diffusion-model simulations with noisy observations via sequential Monte Carlo reweighting over diffusion trajectories.
-
PACE: Geometry-Aware Bridge Transport for Single-Cell Trajectory Inference
PACE recovers geometry-consistent continuous transport dynamics from destructive single-cell snapshots via anisotropic metrics and neural bridges, reducing reconstruction distances by 23.7% on average across seven datasets.
-
Towards Continuous Sign Language Conversation from Isolated Signs
Constructs continuous sign conversation data from isolated signs using retrieval and diffusion models to train a direct sign-to-sign conversational AI.
-
Intervention-Based Time Series Causal Discovery via Simulator-Generated Interventional Distributions
SVAR-FM uses simulator clamping to produce interventional distributions and flow matching to identify time series causal structures, with an error bound that predicts sign reversal of causal effects below a simulator ...
-
Curated Synthetic Data Doesn't Have to Collapse: A Theoretical Study of Generative Retraining with Pluralistic Preferences
Recursive generative retraining with heterogeneous rewards converges to a stable distribution satisfying a weighted Nash bargaining solution, preserving diversity under stated conditions.
-
Debiased Counterfactual Generation via Flow Matching from Observations
Observational and counterfactual distributions are linked by identical support and invariant features, enabling a flow-matching estimator with semiparametric efficiency correction to generate debiased counterfactuals ...
-
SDFlow: Similarity-Driven Flow Matching for Time Series Generation
SDFlow learns a global transport map via similarity-driven flow matching in VQ latent space, using low-rank manifold decomposition and a categorical posterior to handle discreteness, yielding SOTA long-horizon perform...
-
SixthSense: Task-Agnostic Proprioception-Only Whole-Body Wrench Estimation for Humanoids
SixthSense infers whole-body contact events and wrenches in humanoids from proprioception and IMU data alone by tokenizing histories and estimating a sparse contact-event flow with conditional flow matching.
-
FlowS: One-Step Motion Prediction via Local Transport Conditioning
FlowS achieves state-of-the-art single-step motion prediction on Waymo Open Motion Dataset by using scene-conditioned anchor trajectories and a step-consistent displacement field to make local transport accurate in on...
-
Fisher Decorator: Refining Flow Policy via a Local Transport Map
Fisher Decorator refines flow policies in offline RL via a local transport map and Fisher-matrix quadratic approximation of the KL constraint, yielding controllable error near the optimum and SOTA benchmark results.
-
PRiMeFlow: Capturing Complex Expression Heterogeneity in Perturbation Response Modelling
PRiMeFlow is a flow-matching model that approximates the full empirical distribution of single-cell gene expression after perturbations.
-
Monte Carlo Event Generation with Continuous Normalizing Flows
Continuous normalizing flows improve unweighting efficiency in Monte Carlo event generation for high-jet-multiplicity collider processes by factors up to 184, with wall-time gains of about ten when combined with coupl...
-
FluxMC: Rapid and High-Fidelity Inference for Space-Based Gravitational-Wave Observations
FluxMC integrates flow matching with parallel tempering MCMC to converge in under five hours on high-fidelity IMRPhenomHM waveforms for massive black hole binaries, where standard methods fail after hundreds of hours ...
-
MIOFlow 2.0: A unified framework for inferring cellular stochastic dynamics from single cell and spatial transcriptomics data
MIOFlow 2.0 learns stochastic cellular trajectories from transcriptomics data via neural SDEs, unbalanced optimal transport for growth, and a joint latent space unifying gene expression with spatial features.
-
Flow Matching is Adaptive to Manifold Structures
Flow matching achieves near-minimax optimal statistical consistency for manifold-supported distributions, with convergence rates governed by intrinsic dimension and smoothness rather than ambient dimension.
-
Flow marching for a generative PDE foundation model
Flow Marching jointly samples noise and physical time to learn a velocity field for generative PDE modeling, paired with a latent autoencoder and efficient transformer for large-scale pretraining on 2.5M trajectories.
-
Latent Stochastic Interpolants
Latent Stochastic Interpolants jointly optimize encoder-decoder and a latent-space stochastic interpolant using a continuous-time ELBO to transform arbitrary priors into aggregated posteriors.
discussion (0)
Sign in with ORCID, Apple, or X to comment. Anyone can read and Pith papers without signing in.