Two-Step Diffusion: Fast Sampling and Reliable Prediction for 3D Keller--Segel and KPP Equations in Fluid Flows

TL;DR

This work addresses the challenge of fast, distribution-level transport for 3D Keller–Segel and KPP systems in fluid flows by proposing a Two-Step Diffusion framework that first learns a fast, one-step Meanflow initializer and then applies a near-identity Deep Particle corrector trained with a mini-batch $W_2$ objective using warm-started OT couplings. The approach converts a high-dimensional OT problem into a tractable refinement, yielding substantial reductions in $W_2$ compared to Meanflow alone and exhibiting strong extrapolation and stability in stiff advection regimes. Empirically, the method delivers accurate invariant measures and eigenvalue convergence for KPP fronts, with qualitative improvements in KS density ridges under chaotic flows and robustness to 4D extensions, indicating practical applicability to high-dimensional scientific transport problems. Overall, the Two-Step Diffusion framework offers a principled, geometry-aware workflow that couples fast, deterministic transport with principled OT refinement to achieve reliable, scalable distribution matching in complex fluid-transport dynamics.

Abstract

We study fast and reliable generative transport for the 3D KS (Keller-Segel) and KPP (Kolmogorov-Petrovsky-Piskunov) equations in the presence of fluid flows with the goal to approximate the map between initial and terminal distributions for a range of physical parameters $σ$ under the Wasserstein metric. To minimize the inaccuracy of direct Wasserstein solver, we propose a two-stage pipeline that retains one-step efficiency while reinstating an explicit $W_2$ objective where it is tractable. In Stage I, a Meanflow-style regressor yields a deterministic, one-step global transport that moves particles close to their terminal states. In Stage II, we freeze this initializer and train a near-identity corrector (Deep Particle, DP) that directly minimizes a mini-batch $W_2$ objective using warm-started optimal transport couplings computed on the Meanflow outputs. Crucially, after the one-step transport (from Stage I) concentrating mass on the approximated correct support, the induced geometry stabilizes high-dimensional $W_2$ computation of the direct Wasserstein solver. We validate our construction in the 3D KS and KPP equations subject to fluid flows with ordered and chaotic streamlines.

Figures (17)

• Figure 1: Two-Step pipeline. Meanflow produces an initial estimator $\widehat{x}_0$ with flow-based supervision, the PDE solver provides $x_0$, and Deep particle refines to $\widetilde{x}_0$ via Wasserstein-2 Distance supervision.
• Figure 2: $W_2$ across $\sigma$ (left table) and $W_2$ vs. $\sigma$ (right). Superscripts on $\sigma$ denote usage: $\blacktriangle$—interpolation, $\bullet$—extrapolation. DP refinement consistently achieves the lowest cost and shows the largest gains in the singular perturbation regime ($\sigma\!\gtrsim\!150$). MF uses one-shot sampling; DP is the refinement in Algorithm \ref{['alg:dp_refine']}.
• Figure 3: Qualitative comparisons on x-y, x-z, y-z at $\sigma=160$ for the 3D Keller-Segel system in a laminar flow. (a) Reference solution projected to three coordinate planes. (b) Predicted solution projected to the three coordinate planes by Meanflow ($W_2=0.0403$) (c) DP Refinement solution projects to the three coordinate planes ($W_2=0.0082$).
• Figure 4: Qualitative comparisons of x-y, x-z, y-z planes at $\sigma=110$ for the 3D Keller-Segel system in Kolmogorov flow. (a) Reference solution projected to three coordinate planes. (b) Predicted solution projected to the three coordinate planes by Meanflow (c) DP Refinement solution projects to the three coordinate planes.
• Figure 5: Empirical invariant measures on $\mathbb{T}^2$ across diffusion constants. Rows correspond to $\sigma=2^{-3.75}$ (top, within the training range) and $\sigma=2^{-4}$ (bottom, extrapolation beyond training); columns show Meanflow (left), MF+DP (middle), and the resolved FK reference (right). Each panel is a $72{\times}72$ histogram on $[0,2\pi)^2$. The DP corrector contracts spurious mass and sharpens the anisotropic ridge, bringing the warm-start distribution visibly closer to the reference at both $\sigma$, including the extrapolation case.