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

TL;DR

OT-NFM introduces an ODE-free, one-step generative model that learns direct transport maps using neural flows, reducing inference to a single forward pass while maintaining competitive quality.

Contribution

It proposes a novel ODE-free framework for generative modeling that directly learns transport maps, addressing mean collapse with optimal transport pairings.

Findings

Achieves competitive sample quality on MNIST and CIFAR-10.

Enables true one-step generation with a single network evaluation.

Addresses mean collapse through optimal transport coupling strategies.

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.