Memory-Conditioned Flow-Matching for Stable Autoregressive PDE Rollouts

TL;DR

This paper introduces a memory-conditioned flow-matching approach for stable long-term autoregressive PDE simulations, addressing drift issues by incorporating structured memory and proving stability, with improved accuracy demonstrated on complex flow problems.

Contribution

The paper develops a novel memory-conditioned diffusion method for PDE rollouts, leveraging the Mori--Zwanzig formalism to improve stability and fidelity over long horizons.

Findings

Enhanced stability of long-horizon PDE rollouts.

Improved spectral and statistical fidelity in complex flows.

Theoretical proof of Wasserstein stability for the proposed method.

Abstract

Autoregressive generative PDE solvers can be accurate one step ahead yet drift over long rollouts, especially in coarse-to-fine regimes where each step must regenerate unresolved fine scales. This is the regime of diffusion and flow-matching generators: although their internal dynamics are Markovian, rollout stability is governed by per-step \emph{conditional law} errors. Using the Mori--Zwanzig projection formalism, we show that eliminating unresolved variables yields an exact resolved evolution with a Markov term, a memory term, and an orthogonal forcing, exposing a structural limitation of memoryless closures. Motivated by this, we introduce memory-conditioned diffusion/flow-matching with a compact online state injected into denoising via latent features. Via disintegration, memory induces a structured conditional tail prior for unresolved scales and reduces the transport needed to populate missing frequencies. We prove Wasserstein stability of the resulting conditional kernel. We then derive discrete Gr\"onwall rollout bounds that separate memory approximation from conditional generation error. Experiments on compressible flows with shocks and multiscale mixing show improved accuracy and markedly more stable long-horizon rollouts, with better fine-scale spectral and statistical fidelity.