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

TL;DR

This paper introduces a two-step diffusion method combining a deterministic initial transport with a learned correction to efficiently and reliably approximate the evolution of distributions in 3D Keller-Segel and KPP equations under fluid flows, using Wasserstein metrics.

Contribution

It proposes a novel two-stage pipeline that enhances Wasserstein-based transport approximation by combining a global deterministic step with a learned correction, improving efficiency and accuracy.

Findings

Effective in 3D KS and KPP equations with fluid flows

Reduces Wasserstein approximation errors significantly

Applicable to ordered and chaotic flow regimes

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 $\sigma$ 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.