Two step diffusion: fast sampling and reliable prediction of 3D Keller-Segel chemotaxis systems in fluid flows

Zhenda Shen, Zhongjian Wang, Jack Xin, Zhiwen Zhang

Keywords: 3D Keller–Segel chemotaxis systems, Two-Step Diffusion, Optimal Transport (OT), Wasserstein-2 distance (W2), GPU-friendly mini-batch OT training

Abstract: In this work, we study fast and reliable generative transport for the 3D Keller-Segel in different fluid flows, where the goal is to map initial particles $x_0$ to terminal states $x_1$ for a range of physical parameters $\sigma$. While the quadratic Wasserstein distance $W_2$ serves as a better metric for differences between distributions, optimizing $W_2$ directly is unstable and computationally expensive in high dimensions. 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 trained via a MeanFlow identity yields a deterministic, 1-NFE global transport that moves particles close to their terminal states without simulating forward diffusions. 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 EMD (Earth Mover's Distance)/OT (Optimal Transport) couplings computed on the MeanFlow outputs. Crucially, after the one-step transport (from Stage I) has concentrated mass on the correct support, the induced geometry stabilizes high-dimensional $W_2$ optimization. We validate our construction on the 3D Keller Segel setting with different flows and find that our method consistently reduces the empirical $W_2$ relative to one-step flows. Moreover, the two-stage refinement reduces MeanFlow's generalization error across $\sigma$, demonstrating improved robustness to parameter shift, which is evidenced by consistently lower $W_2$ over the $\sigma$ sweep. At the same time, the approach preserves the properties of 1-NFE and deterministic sampling, and yields clear qualitative gains in anisotropy and mass placement, as supported by our simulations and $W_2$-vs-$\sigma$ curves.

Primary Area: applications to physical sciences (physics, chemistry, biology, etc.)

Submission Number: 23927