Rectified Flows for Fast Multiscale Fluid Flow Modeling

TL;DR

The paper tackles the challenge of efficiently sampling multiscale, chaotic fluid flows by learning Rectified Flows (RF), a deterministic ODE-based surrogate that transports input-output laws along near-straight trajectories. By matching a barycentric velocity $v^\star$ and employing curvature-aware inference with EMA-based blending, RF achieves diffusion-like law fidelity with far fewer steps (approximately 8–10 vs 128+). The authors develop a law-level framework linking structure functions to a coverage term and show that straightness of transport controls discretization error, enabling robust step-size adaptation. Empirical results on 2D incompressible and compressible flows demonstrate that RF matches or surpasses diffusion baselines in one-point $W_1$ and spectral metrics while offering substantial speedups and training efficiency, making multiscale CFD surrogates practical for uncertainty quantification and real-time prediction. The work provides a principled path from law-level theory to fast, accurate surrogates for chaotic PDEs with broad implications for scientific computing and AI-assisted modeling.

Abstract

Statistical surrogate modeling of fluid flows is hard because dynamics are multiscale and highly sensitive to initial conditions. Conditional diffusion surrogates can be accurate, but usually need hundreds of stochastic sampling steps. We propose a rectified-flow surrogate that learns a time-dependent conditional velocity field transporting input-to-output laws along near-straight trajectories. Inference is then a deterministic ODE solve, making each function evaluation more informative: on multiscale 2D benchmarks, we match diffusion-class posterior statistics with only (8) ODE steps versus (\ge 128) for score-based diffusion. Theoretically, we give a law-level analysis for conditional PDE forecasting. We (i) connect one-point Wasserstein field metrics to the (k=1) correlation-marginal perspective in statistical solutions, (ii) derive a one-step error split into a **coverage** term (high-frequency tail, controlled by structure functions/spectral decay) and a **fit** term (controlled by the training objective), and (iii) show that rectification-time **straightness** controls ODE local truncation error, yielding practical step-size/step-count guidance. Motivated by this, we introduce a curvature-aware sampler that uses an EMA straightness proxy to adapt blending and step sizes at inference. Across incompressible and compressible multiscale 2D flows, it matches diffusion baselines in Wasserstein statistics and spectra, preserves fine-scale structure beyond MSE surrogates, and significantly reduces inference cost.

Figures (29)

• Figure 1: Overview of ReFlow. (a): ReFlow closely matches high-order statistical moments (in the Wasserstein $W_1$-sense) of fluid flows by learning a neural surrogate to approximate the push-forward of the distribution $p(u_{\tau|\tau=0} \vert u_i)$, while minimizing inference-time computational costs. (b) Macro-architecture overview: A UViT based network predicts the velocity field $v_\theta$, conditioned on Gaussian noise, initial data $u_i$ and boundary conditions, and diffusion time $\tau$. Noise and initial / boundary conditions are concatenated as inputs. (c): Results for the cylindrical shear flow dataset. (d): The KDE plots monitor the evolution of a down-projected version of the push-forward measure $p_\tau := X_\theta(\tau)_\# p(u_{\tau=0} \vert u_i)$, where $X_\theta$ is the flow induced by the learned vector field $v_\theta$. Results for the Richtmyer Meshkov dataset showing flow matching between noise ($p_{\tau=0}$) and target ($p_{\tau=1}$) distributions. ReFlow solves an ODE, learning a constant velocity field that enables straight-line inference paths (visualized with principal component analysis and kernel density estimation at different diffusion times).
• Figure 2: Inference cost versus model size. Rectified flows need $4$–$8$ ODE steps vs. GenCFD’s $64$–$128$ SDE steps. U/̄Net depth varies from $4$ to $8$ blocks (parameters $5.5$ M to $10.2$ M).
• Figure 3: Energy-spectrum evolution for a random cloud–shock initial datum. ReFlow (blue) tracks high-wavenumber energy faster and more accurately than GenCFD (red).
• Figure 4: Density trajectories for Sample 1: ReFlow (top) versus GenCFD (bottom).
• Figure 5: UNet Backbone Architecture Used in ReFlow. This schematic illustrates the core UNet-based architecture used within the ReFlow framework, structured across three resolution levels. For clarity, the number of blocks per level is set to one in this illustration. In the actual GenCFD and ReFlow configurations, each block is repeated four times per level. The bottleneck shows an asymmetry between the encoder and decoder sides: the block on the encoder side includes a Convolution Block, Positional Embedding, and Attention Block, while the corresponding block on the decoder side omits the Positional Embedding.

Theorems & Definitions (50)

• Remark 4.1: Disintegration form
• Theorem 4.2: Terminal error decomposition
• Corollary 4.3: High-probability control of terminal error
• proof : Proof sketch
• Lemma 4.4: Local truncation error for explicit Euler on $\dot u=v(u,\tau)$
• Lemma 4.5: Bernstein upper bound for bandlimited fields
• proof : Proof sketch
• Lemma 4.6: Annulus lower bound
• proof : Proof sketch
• Corollary 4.7: Capacity requirement for typical strain levels