pith. sign in

arxiv: 2604.06413 · v1 · submitted 2026-04-07 · 💻 cs.LG

ODE-free Neural Flow Matching for One-Step Generative Modeling

Xiao Shou This is my paper

Pith reviewed 2026-05-10 20:02 UTC · model grok-4.3

classification 💻 cs.LG
keywords neural flow matchingoptimal transportone-step generationgenerative modelingflow mapsdiffusion modelsimage generationmean collapse
0
0 comments X

The pith

Optimal transport pairings let neural networks learn direct one-step maps from noise to data.

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

The paper seeks to replace the repeated vector-field integrations of diffusion and flow-matching models with a single learned transport map that takes noise straight to data. A reader would care because this cuts inference from dozens of network calls down to one forward pass. The authors show that training such a map directly collapses to the data mean unless the noise-data pairs are consistent, prove that consistency is required for non-degenerate solutions, and enforce it by replacing random pairings with optimal-transport couplings computed on minibatches and online. Experiments on synthetic data, MNIST, and CIFAR-10 confirm that the resulting model matches the quality of multi-step baselines while using only a single evaluation.

Core claim

We propose Optimal Transport Neural Flow Matching (OT-NFM), an ODE-free generative framework that parameterizes the flow map with neural flows, enabling true one-step generation with a single forward pass. We show that naive flow-map training suffers from mean collapse, where inconsistent noise-data pairings drive all outputs toward the data mean. We prove that consistent coupling is necessary for non-degenerate learning and address this using optimal transport pairings with scalable minibatch and online coupling strategies.

What carries the argument

Optimal Transport Neural Flow Matching (OT-NFM), a direct parameterization of the transport map by a neural flow that is trained on consistent optimal-transport couplings instead of random noise-data pairs.

Load-bearing premise

That optimal transport can supply consistent, unbiased pairings at minibatch scale and online without introducing new bias, and that a neural network can accurately represent the resulting non-degenerate transport map.

What would settle it

Training OT-NFM with the proposed couplings and then observing that generated samples still concentrate near the data mean on a test distribution would show that consistent couplings are not sufficient for non-degenerate learning.

Figures

Figures reproduced from arXiv: 2604.06413 by Xiao Shou.

Figure 1
Figure 1. Figure 1: Synthetic transport results. Flow trajectories under four coupling strategies (columns) on four 2-D benchmarks (rows): Gauss → Checkerboard, Gauss → Spiral, Gauss →Crescent, and 8-GMM →2- Moons. Black dots: source samples (t = 0); blue dots: generated samples (t = 1); olive arrows: learned flow trajectories. Per-batch OT produces tangled, inconsistent trajectories across all tasks. Minibatch OT, LOOM, and … view at source ↗
Figure 2
Figure 2. Figure 2: MNIST coupling ablation. Each panel shows 100 generated samples (10×10 grid). NFM (no OT): complete mean collapse — every output is a blurred average with no digit identity. NFM (minibatch OT): precomputed pairing restores sharp, diverse digits in a single forward pass. CFM (no OT) provides an unguided flow matching baseline with independent coupling. CFM (minibatch OT) (100 NFE) and MeanFlow (1 NFE) produ… view at source ↗
Figure 4
Figure 4. Figure 4: CIFAR-10 samples (1-NFE). We train our neural flow model on CIFAR-10 (32×32, RGB) with precomputed minibatch OT (N=50,000, B=256, 5 sweep epochs) for 100,000 steps using the Adam optimizer with learn￾ing rate 2×10−4 and cosine annealing decay on an NVIDIA A100 GPU [PITH_FULL_IMAGE:figures/full_fig_p009_4.png] view at source ↗
Figure 3
Figure 3. Figure 3: Neural flow trajectories on MNIST. Each row shows Fθ(t, x0) for a fixed noise vector x0 ∼ N (0, I) evaluated at t ∈ {0.0, 0.1, . . . , 1.0}. At t=0 the identity condition gives pure noise; coarse digit structure emerges by t≈0.3–0.4 and class identity is committed by t≈0.5; background noise is suppressed and strokes sharpen monotonically toward the clean output at t=1.0. No ODE solver is used; generation i… view at source ↗
Figure 5
Figure 5. Figure 5: Trajectory ablation (global OT). Top row: 8-GMM → 2-moons. Bottom row: Gaus￾sian → Checkerboard. Columns show Cosine (left), Polynomial α=2 (center ), and Stochastic σ=0.5 (right). Linear interpolation ( [PITH_FULL_IMAGE:figures/full_fig_p011_5.png] view at source ↗
read the original abstract

Diffusion and flow matching models generate samples by learning time-dependent vector fields whose integration transports noise to data, requiring tens to hundreds of network evaluations at inference. We instead learn the transport map directly. We propose Optimal Transport Neural Flow Matching (OT-NFM), an ODE-free generative framework that parameterizes the flow map with neural flows, enabling true one-step generation with a single forward pass. We show that naive flow-map training suffers from mean collapse, where inconsistent noise-data pairings drive all outputs toward the data mean. We prove that consistent coupling is necessary for non-degenerate learning and address this using optimal transport pairings with scalable minibatch and online coupling strategies. Experiments on synthetic benchmarks and image generation tasks (MNIST and CIFAR-10) demonstrate competitive sample quality while reducing inference to a single network evaluation.

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, simulated authors' rebuttal, and a circularity audit. Tearing a paper down is the easy half of reading it; the pith above is the substance, this is the friction.

Referee Report

3 major / 2 minor

Summary. The paper proposes Optimal Transport Neural Flow Matching (OT-NFM), an ODE-free framework for one-step generative modeling. It parameterizes the flow map directly with neural flows, proves that consistent noise-data couplings are necessary to avoid mean collapse, and uses scalable minibatch and online optimal transport pairings to provide those couplings. Experiments on synthetic data, MNIST, and CIFAR-10 report competitive sample quality with inference reduced to a single network evaluation.

Significance. If the necessity proof holds and the OT approximations preserve sufficient consistency, the work offers a concrete path to single-pass generation that avoids the multi-step integration cost of diffusion and flow-matching models. The use of standard image benchmarks and the explicit identification of the mean-collapse failure mode are positive contributions that could influence efficient generative modeling research.

major comments (3)
  1. [Abstract / necessity proof] Abstract and the necessity proof section: the claim that minibatch and online OT strategies suffice to satisfy the consistent-coupling condition is load-bearing for the one-step guarantee. The manuscript must demonstrate (via bound or empirical diagnostic) that residual bias in these approximations does not drive the learned map toward degeneracy, as any such bias would directly contradict the necessity result invoked to justify the framework.
  2. [Experiments] Experiments section: reported results on MNIST and CIFAR-10 lack error bars, multiple random seeds, or ablation on pairing batch size. Without these, it is impossible to verify that the single-step samples are statistically competitive rather than artifacts of a single run or favorable pairing.
  3. [Method] Method section on neural flow parameterization: the exact loss formulation when the transport map is learned from OT pairings, and the architectural choices that prevent collapse even under approximate couplings, require additional equations and pseudocode to support reproducibility.
minor comments (2)
  1. [Abstract] Define 'neural flows' explicitly on first use and distinguish from flow matching or normalizing flows to avoid notation confusion.
  2. [Introduction] Add a short related-work paragraph contrasting OT-NFM with existing one-step methods (e.g., distilled diffusion, GANs) to clarify the incremental contribution.

Simulated Author's Rebuttal

3 responses · 0 unresolved

We thank the referee for the thoughtful and constructive review. The comments identify important areas for strengthening the rigor and reproducibility of the manuscript. We address each major comment below and commit to revisions that directly respond to the concerns raised.

read point-by-point responses

  1. Referee: [Abstract / necessity proof] Abstract and the necessity proof section: the claim that minibatch and online OT strategies suffice to satisfy the consistent-coupling condition is load-bearing for the one-step guarantee. The manuscript must demonstrate (via bound or empirical diagnostic) that residual bias in these approximations does not drive the learned map toward degeneracy, as any such bias would directly contradict the necessity result invoked to justify the framework.

    Authors: We agree that the sufficiency of the approximate couplings is central to the framework and that the necessity proof alone does not automatically guarantee non-degeneracy under approximation. In the revised manuscript we will add both an empirical diagnostic (measuring the effective coupling inconsistency via average transport cost deviation across training batches and correlating it with output variance to confirm absence of collapse) and a short theoretical remark bounding the propagation of residual bias into the learned map under the Lipschitz assumptions already used in the necessity proof. These additions will be placed in the necessity proof section and referenced from the abstract. revision: yes

  2. Referee: [Experiments] Experiments section: reported results on MNIST and CIFAR-10 lack error bars, multiple random seeds, or ablation on pairing batch size. Without these, it is impossible to verify that the single-step samples are statistically competitive rather than artifacts of a single run or favorable pairing.

    Authors: We acknowledge that the current experimental reporting is insufficient for statistical confidence. We will rerun the MNIST and CIFAR-10 experiments with at least five independent random seeds, report mean FID (and other metrics) together with standard error bars, and add an ablation table varying the OT pairing batch size over a range that includes the values used in the main results. The revised experiments section will present these new tables and figures. revision: yes

  3. Referee: [Method] Method section on neural flow parameterization: the exact loss formulation when the transport map is learned from OT pairings, and the architectural choices that prevent collapse even under approximate couplings, require additional equations and pseudocode to support reproducibility.

    Authors: We will expand the method section with the precise training objective that incorporates the OT-derived pairings (including the explicit expectation over the approximate coupling), the full set of architectural hyperparameters for the neural flow, and a pseudocode listing of the end-to-end training loop. We will also add a short paragraph explaining the architectural elements (e.g., residual connections and output scaling) that, in conjunction with the consistent-coupling condition, empirically stabilize training even when the OT approximation is imperfect. revision: yes

Circularity Check

0 steps flagged

No significant circularity; derivation relies on external OT theory and internal proof without reduction to inputs

full rationale

The paper's central derivation proceeds by identifying mean collapse in naive flow-map training, proving the necessity of consistent couplings for non-degenerate maps, and then applying optimal transport pairings (an external mathematical construct) via practical minibatch/online strategies to enable direct parameterization of the transport map with neural flows. This yields the one-step claim without any step that defines the output in terms of itself, renames a fitted quantity as a prediction, or reduces via self-citation chains. The proof and OT application are independent of the final generative performance metrics, keeping the framework self-contained against external benchmarks.

Axiom & Free-Parameter Ledger

0 free parameters · 0 axioms · 0 invented entities

Abstract-only review yields no explicit free parameters, axioms, or invented entities beyond standard neural network training assumptions and optimal transport properties from prior literature.

pith-pipeline@v0.9.0 · 5421 in / 1187 out tokens · 70346 ms · 2026-05-10T20:02:24.972898+00:00 · methodology

discussion (0)

Sign in with ORCID, Apple, or X to comment. Anyone can read and Pith papers without signing in.

Lean theorems connected to this paper

Citations machine-checked in the Pith Canon. Every link opens the source theorem in the public Lean library.

What do these tags mean?
matches
The paper's claim is directly supported by a theorem in the formal canon.
supports
The theorem supports part of the paper's argument, but the paper may add assumptions or extra steps.
extends
The paper goes beyond the formal theorem; the theorem is a base layer rather than the whole result.
uses
The paper appears to rely on the theorem as machinery.
contradicts
The paper's claim conflicts with a theorem or certificate in the canon.
unclear
Pith found a possible connection, but the passage is too broad, indirect, or ambiguous to say the theorem truly supports the claim.

Reference graph

Works this paper leans on

25 extracted references · 25 canonical work pages

  1. [1]

    In: ICLR (2023)

    Albergo, M.S., Vanden-Eijnden, E.: Building normalizing flows with stochastic interpolants. In: ICLR (2023)

    work page 2023

  2. [2]

    Advances in neural information processing systems34, 21325–21337 (2021)

    Biloˇ s, M., Sommer, J., Rangapuram, S.S., Januschowski, T., G¨ unnemann, S.: Neural flows: Efficient alternative to neural odes. Advances in neural information processing systems34, 21325–21337 (2021)

    work page 2021

  3. [3]

    Boffi, Michael S

    Boffi, N.M., Albergo, M.S., Vanden-Eijnden, E.: Flow map matching. arXiv preprint arXiv:2406.07507 (2024) 10 Figure 5:Trajectory ablation (global OT).Top row: 8-GMM→2-moons. Bottom row: Gaus- sian→Checkerboard. Columns show Cosine (left), Polynomialα=2 (center), and Stochasticσ=0.5 (right). Linear interpolation (Figure 1) produces the straightest trajecto...

    work page arXiv 2024

  4. [4]

    Advances in neural information processing systems31(2018)

    Chen, R.T., Rubanova, Y., Bettencourt, J., Duvenaud, D.K.: Neural ordinary differential equations. Advances in neural information processing systems31(2018)

    work page 2018

  5. [5]

    In: ICLR (2023)

    Chen, T.: On the importance of noise scheduling for diffusion models. In: ICLR (2023)

    work page 2023

  6. [6]

    In: The Thirteenth International Conference on Learning Representations (2025)

    Davtyan, A., Dadi, L.T., Cevher, V., Favaro, P.: Faster inference of flow-based generative models via improved data-noise coupling. In: The Thirteenth International Conference on Learning Representations (2025)

    work page 2025

  7. [7]

    In: NeurIPS (2021)

    Dhariwal, P., Nichol, A.: Diffusion models beat GANs on image synthesis. In: NeurIPS (2021)

    work page 2021

  8. [8]

    In: ICLR (2025)

    Frans, K., Hafner, D., Levine, S., Abbeel, P.: One step diffusion via shortcut models. In: ICLR (2025)

    work page 2025

  9. [9]

    NeurIPS 2025 (2025)

    Geng, Z., Deng, M., Bai, X., Kolter, J.Z., He, K.: Mean flows for one-step generative modeling. NeurIPS 2025 (2025)

    work page 2025

  10. [10]

    Machine Learning110(2), 393–416 (2021)

    Gouk, H., Frank, E., Pfahringer, B., Cree, M.J.: Regularisation of neural networks by enforcing lipschitz continuity. Machine Learning110(2), 393–416 (2021)

    work page 2021

  11. [11]

    In: NeurIPS (2020)

    Ho, J., Jain, A., Abbeel, P.: Denoising diffusion probabilistic models. In: NeurIPS (2020)

    work page 2020

  12. [12]

    Advances in neural information processing systems35, 26565–26577 (2022)

    Karras, T., Aittala, M., Aila, T., Laine, S.: Elucidating the design space of diffusion-based generative models. Advances in neural information processing systems35, 26565–26577 (2022)

    work page 2022

  13. [13]

    Advances in Neural Information Processing Systems37, 104180–104204 (2024) 11

    Kornilov, N., Mokrov, P., Gasnikov, A., Korotin, A.: Optimal flow matching: Learning straight trajec- tories in just one step. Advances in Neural Information Processing Systems37, 104180–104204 (2024) 11

    work page 2024

  14. [14]

    In: ICLR (2023)

    Lipman, Y., Chen, R., et al.: Flow matching for generative modeling. In: ICLR (2023)

    work page 2023

  15. [15]

    In: The Eleventh International Conference on Learning Representations (ICLR) (2023)

    Liu, X., Gong, C., Liu, Q.: Flow straight and fast: Learning to generate and transfer data with rectified flow. In: The Eleventh International Conference on Learning Representations (ICLR) (2023)

    work page 2023

  16. [16]

    In: ICLR (2018)

    Miyato, T., et al.: Spectral normalization for generative adversarial networks. In: ICLR (2018)

    work page 2018

  17. [17]

    In: Proceedings of the IEEE/CVF international conference on computer vision

    Peebles, W., Xie, S.: Scalable diffusion models with transformers. In: Proceedings of the IEEE/CVF international conference on computer vision. pp. 4195–4205 (2023)

    work page 2023

  18. [18]

    In: The Thirty-ninth Annual Con- ference on Neural Information Processing Systems (2025)

    Petrovi´ c, K., Atanackovic, L., Moro, V., Kapu´ sniak, K., Ceylan, I.I., Bronstein, M.M., Bose, J., Tong, A.: Curly flow matching for learning non-gradient field dynamics. In: The Thirty-ninth Annual Con- ference on Neural Information Processing Systems (2025)

    work page 2025

  19. [19]

    In: MICCAI (2015)

    Ronneberger, O., et al.: U-net: Convolutional networks for biomedical image segmentation. In: MICCAI (2015)

    work page 2015

  20. [20]

    In: ICML (2015)

    Sohl-Dickstein, J., Weiss, E., Maheswaranathan, N., Ganguli, S.: Deep unsupervised learning using nonequilibrium thermodynamics. In: ICML (2015)

    work page 2015

  21. [21]

    Song, Y., Dhariwal, P.: Improved techniques for training consistency models (2024)

    work page 2024

  22. [22]

    ICML (2023)

    Song, Y., Durkan, C., et al.: Consistency models. ICML (2023)

    work page 2023

  23. [23]

    In: International Conference on Learning Represen- tations (2021)

    Song, Y., Sohl-Dickstein, J., Kingma, D.P., Kumar, A., Ermon, S., Poole, B.: Score-based generative modeling through stochastic differential equations. In: International Conference on Learning Represen- tations (2021)

    work page 2021

  24. [24]

    Transactions on Machine Learning Research pp

    Tong, A., Fatras, K., Malkin, N., Huguet, G., Zhang, Y., Rector-Brooks, J., Wolf, G., Bengio, Y.: Im- proving and generalizing flow-based generative models with minibatch optimal transport. Transactions on Machine Learning Research pp. 1–34 (2024)

    work page 2024

  25. [25]

    Inductive moment matching.arXiv preprint arXiv:2503.07565, 2025

    Zhou, L., Ermon, S., Song, J.: Inductive moment matching. arXiv preprint arXiv:2503.07565 (2025) 12

    work page arXiv 2025