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

2026-04-07

AI Summary

This work addresses the inefficiency of traditional diffusion and flow-matching models, which require costly iterative inference, and the tendency of direct transport map learning to suffer from mean collapse. The authors propose OT-NFM, a framework that parameterizes the optimal transport–induced flow map directly via neural flows, enabling single-step generation without solving ordinary differential equations. By introducing scalable minibatch and online optimal transport coupling strategies, the method enforces consistent pairing between noise and data samples within an end-to-end differentiable training pipeline, thereby avoiding degenerate solutions. Empirical results on MNIST, CIFAR-10, and synthetic benchmarks demonstrate that the model achieves competitive generation quality with only a single forward pass.

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.

Problem

• generative modeling
• flow matching
• mean collapse
• optimal transport
• one-step generation

Innovation

• Neural Flow Matching
• Optimal Transport
• One-Step Generation
• ODE-Free Generative Modeling
• Transport Map