Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Memory-Conditioned Flow-Matching for Stable Autoregressive PDE Rollouts

Victor Armegioiu

TL;DR

This work addresses the instability of autoregressive PDE rollouts in coarse-to-fine multiscale regimes by introducing memory-conditioned flow matching. Grounded in the Mori–Zwanzig decomposition, it shows that a compact memory state is essential to approximate history-dependent tail priors and stabilize long-horizon generation, while retaining a Markovian internal-time transport. The Memory-Conditioned Rectified Flow (MCRF) framework augments a convolutional backbone with a memory pathway, and provides a formal rollout bound that decomposes errors into memory-approximation and conditional-generation components. Empirically, on 2D compressible Euler benchmarks with shocks and mixing, the method delivers substantial stability gains and improved multi-scale fidelity, especially in the low- and mid-frequency bands, with additional robustness from degraded-conditioning training.

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önwall 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.

Memory-Conditioned Flow-Matching for Stable Autoregressive PDE Rollouts

TL;DR

This work addresses the instability of autoregressive PDE rollouts in coarse-to-fine multiscale regimes by introducing memory-conditioned flow matching. Grounded in the Mori–Zwanzig decomposition, it shows that a compact memory state is essential to approximate history-dependent tail priors and stabilize long-horizon generation, while retaining a Markovian internal-time transport. The Memory-Conditioned Rectified Flow (MCRF) framework augments a convolutional backbone with a memory pathway, and provides a formal rollout bound that decomposes errors into memory-approximation and conditional-generation components. Empirically, on 2D compressible Euler benchmarks with shocks and mixing, the method delivers substantial stability gains and improved multi-scale fidelity, especially in the low- and mid-frequency bands, with additional robustness from degraded-conditioning training.

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önwall 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.
Paper Structure (126 sections, 30 theorems, 132 equations, 9 tables, 7 algorithms)

This paper contains 126 sections, 30 theorems, 132 equations, 9 tables, 7 algorithms.

Table of Contents

  1. Introduction
  2. Contributions.
  3. Organization.
  4. Background and setup
  5. PDE setting, resolved observations, and distributional rollouts
  6. Mori--Zwanzig: exact memory and orthogonal forcing after elimination
  7. Conditional flow matching / rectified flows as stochastic transitions
  8. Why internal time matters for PDE rollouts.
  9. Littlewood--Paley cutoff and the "transport difficulty" viewpoint
  10. Why memory changes the difficulty of populating missing shells.
  11. Theory: memory-shaped priors and stable autoregressive rollouts
  12. Internal time and Mori--Zwanzig: prior-driven shells require a history convolution
  13. Why this forces memory.
  14. Rollouts: stability reduces everything to a one-step law defect
  15. Disintegration: with low modes fixed, mismatch localizes to the tail
  16. ...and 111 more sections

Key Result

Theorem 3.1

Let $\mathcal{F}_n := \sigma(E_0,\dots,E_n)$. Assume that, conditionally on $\mathcal{F}_n$, each shell covariance is isotropic: Define the history term (conditional mean) Fix $\tau\in[0,1]$ and set $X:=X^{(\tau)}_{n+1}$. Then for every $j>J$, where In particular, in low-SNR shells (i.e. $s_{j,n+1}^2\ll \sigma(\tau)^2$), the denoiser is prior-driven and governed by $\Delta_j M^\star_{n+1}$.

Theorems & Definitions (59)

  • Theorem 3.1: Shell-wise history prior in the optimal denoiser
  • Lemma 3.2: Error recursion
  • Remark 3.3: Relation to invariant-measure regularization
  • Lemma 3.4: $W_2$ reduction under pinned low modes
  • Lemma 3.5: Conditioning reduces $L^2$ projection error
  • Proposition 3.6: Irreducible mismatch without memory
  • Proposition 3.7: Conditional defect decomposition
  • Theorem 3.8: Rollout bound
  • Theorem 1.3: Internal-time shellwise posterior mean
  • Lemma 1.4: Conditioning reduces $L^2$ projection error
  • ...and 49 more