Rectified Flows for Fast Multiscale Fluid Flow Modeling

Published: 08 Jun 2026, Last Modified: 08 Jun 2026Accepted by TMLREveryoneRevisionsBibTeXCC BY 4.0
Abstract: We introduce \emph{ReFlow}, a conditional rectified-flow surrogate for PDE forecasting. Given an initial state $u_i$, ReFlow transports Gaussian noise $\xi$ to a sample from the conditional final-state law $p(u_f\mid u_i)$ by integrating a learned deterministic ODE. Unlike diffusion surrogates, which require many stochastic denoising steps, the rectified transport is close to straight in sampling time: on multiscale 2D flow benchmarks, ReFlow matches diffusion-level posterior statistics with as few as $8$ ODE steps, compared with $\ge 128$ network evaluations for score-based diffusion. We also give a law-level analysis for conditional PDE surrogates. We formulate the ideal conditional rectified velocity as a barycentric transport field and show that it pushes the reference law to the target conditional law. At fixed spatial resolution, we decompose the one-step law error into a coverage term, controlled by unresolved high-frequency content via structure functions or spectral tails, and a fit term measuring approximation of the ideal velocity field. We further show that ODE discretization error is governed by the variation of the learned velocity along sampled rectified trajectories. This motivates a curvature-aware sampler that uses an EMA proxy for trajectory-wise velocity variation to stabilize inference, especially out of distribution. Across incompressible and compressible 2D flows, ReFlow matches diffusion baselines in one-point Wasserstein statistics and energy spectra, preserves fine-scale structure missed by deterministic MSE models, and produces high-resolution conditional samples at substantially lower inference cost.
Submission Type: Long submission (more than 12 pages of main content)
Changes Since Last Submission: Implemented the changes requested by the Editor.
Assigned Action Editor: ~Markus_Lange-Hegermann1
Submission Number: 7351
Loading