Mamba-Assisted Non-Markovian Closure for Reduced-Order Modeling
Authors: Zhi-Feng Wei, Saad Qadeer, Panos Stinis
Organizations: Advanced Computing, Mathematics, and Data Division, Pacific Northwest National Laboratory, Richland, WA 99354, USA · Department of Applied Mathematics, University of Washington, Seattle, WA 98195, USA · Division of Applied Mathematics, Brown University, Providence, RI 02912, USA
Abstract
Reduced-order modeling of high-dimensional dynamical systems is often hindered by closure effects arising from unresolved variables, which can introduce non-Markovian dependence into the resolved dynamics. Motivated by the history-dependent memory term arising in the Mori--Zwanzig formalism, we recast non-Markovian closure modeling as a sequence modeling problem and propose the Mamba-Assisted Closure (MAC) framework. MAC employs a Mamba-based sequence model to predict the closure from the resolved trajectory and couples the learned closure with the reduced-order governing equations through a numerical integrator to advance the resolved variables in time. During training, the selective scan mechanism in Mamba enables efficient parallel sequence processing with linear scaling in sequence length, while autoregressive inference proceeds through recurrent state updates at essentially constant per-step cost. We evaluate MAC on four benchmark systems with complementary characteristics: the viscous Burgers' equation, the chaotic two-scale Lorenz '96 system, the 3-bus DeMarco--Zheng power-grid system, and the dispersive Korteweg--de Vries equation. Across these benchmarks, MAC consistently improves predictive accuracy and long-time rollout stability relative to the comparison models, demonstrating an effective and computationally scalable approach to non-Markovian closure modeling.
Reduced-order models (ROMs) provide an efficient surrogate for complex multiscale systems, but their predictive accuracy is often compromised by truncation errors and the inadequate representation of interactions between resolved and unresolved scales. The missing effect of truncated (unresolved) scales on ROM (resolved) scales is often denoted as the closure problem. In this work, we formulate ROM closure modeling as a multi-fidelity (MF) learning problem and propose an uncertainty-aware MF framework based on conditional normalizing flow to enhance ROM predictive accuracy. The proposed approach learns a probabilistic mapping from low-fidelity (LF) ROM coefficients to high-fidelity (HF) coefficients, thereby improving predictive fidelity while quantifying the uncertainty associated with the learned closure. Two correction strategies are investigated: direct learning, in which HF coefficients are predicted directly from LF inputs, and residual learning, which learns the discrepancy between LF and HF coefficients and uses it to recover the corrected HF solution. The framework is demonstrated on a vortex merging problem governed by the two-dimensional Navier Stokes equations. Results show that both correction strategies improve ROM accuracy over uncorrected ROM, with residual learning achieving consistently better performance than direct learning. Moreover, the two proposed deep generative model-based strategies provide uncertainty quantification for the corrected ROM coefficients, which is critical for assessing prediction confidence and supporting the reliable use of ROMs in practical applications.
Jice Zeng, Shady E. Ahmed, David Barajas-Solano +1
We propose a latent score-based generative AI framework for learning stochastic, non-local closure models and constitutive laws in nonlinear dynamical systems of computational mechanics. This work addresses a key challenge of modeling complex multiscale dynamical systems without a clear scale separation, for which numerically resolving all scales is prohibitively expensive, e.g., for engineering turbulent flows. While classical closure modeling methods leverage domain knowledge to approximate subgrid-scale phenomena, their deterministic and local assumptions can be too restrictive in regimes lacking a clear scale separation. Recent developments of diffusion-based stochastic models have shown promise in the context of closure modeling, but their prohibitive computational inference cost limits practical applications in many real-world settings. This work addresses this limitation by jointly training convolutional autoencoders with conditional diffusion models in latent space, significantly reducing the dimensionality of the sampling process while preserving essential physical characteristics. Numerical results demonstrate that the joint training approach helps discover a proper latent space that not only guarantees small reconstruction errors but also ensures good performance of the diffusion model in the latent space. When integrated into numerical simulations, the proposed stochastic modeling framework via latent conditional diffusion models achieves significant computational acceleration while maintaining comparable predictive accuracy to standard diffusion models in physical space.
Turbulence is ubiquitous in engineering and science, yet direct simulation is prohibitively expensive. The Reynolds-averaged Navier-Stokes (RANS) equations provide savings exceeding ten orders of magnitude but introduce unclosed terms (the closure problem). Offline-trained machine-learning (ML) closures suffer distribution shift in predictive simulations, while ML methods that bypass the governing equations struggle to generalise from scarce high-fidelity data. We develop a physics-derived deep learning closure model for RANS, the Deep Algebraic Reynolds Stress Model (DARSM), which can be trained on small datasets and accurately generalise across Reynolds numbers, to unseen geometries, and to different flow regimes. A neural network maps flow invariants to empirical parameters in an implicit algebraic Reynolds stress equation, derived from the Reynolds stress transport equations under the weak-equilibrium assumption, imposing physics-based structure on the ML closure. End-to-end optimisation through the governing PDEs and the coupled implicit closure eliminates distribution shift, but both unrolled and implicit automatic differentiation fail on the stiff coupled solver. We derive adjoint equations that exploit the solver's implicit-explicit structure for efficient optimisation. On canonical square-duct and periodic-hill benchmarks, DARSM reduces average test velocity error over baseline RANS by 2-4× across Reynolds number, geometries, and flow regimes, with peak case-level reductions of 12×. The model trained on attached, anisotropy-dominated flows (square duct) accurately generalises without retraining to separated flows (periodic hills), a regime change in the underlying physics. DARSM also outperforms five established ML methods: offline training, tensor-basis neural networks, field-inversion machine learning, DeepONets, and physics-informed neural networks.
Daniel Dehtyriov, Jonathan F. MacArt, Justin Sirignano