AROMA: Preserving Spatial Structure for Latent PDE Modeling with Local Neural Fields

TL;DR

AROMA (Attentive Reduced Order Model with Attention), a framework designed to enhance the modeling of partial differential equations (PDEs) using local neural fields, is presented, achieving greater stability and enable longer rollouts compared to conventional MSE training.

Abstract

We present AROMA (Attentive Reduced Order Model with Attention), a framework designed to enhance the modeling of partial differential equations (PDEs) using local neural fields. Our flexible encoder-decoder architecture can obtain smooth latent representations of spatial physical fields from a variety of data types, including irregular-grid inputs and point clouds. This versatility eliminates the need for patching and allows efficient processing of diverse geometries. The sequential nature of our latent representation can be interpreted spatially and permits the use of a conditional transformer for modeling the temporal dynamics of PDEs. By employing a diffusion-based formulation, we achieve greater stability and enable longer rollouts compared to conventional MSE training. AROMA's superior performance in simulating 1D and 2D equations underscores the efficacy of our approach in capturing complex dynamical behaviors.

Figures (26)

• Figure 1: AROMA inference: The discretization-free encoder compresses the information of a set of $N$ input values to a sequence of $M$ latent tokens, where $M < N$. The conditional diffusion transformer is used to model the dynamics, acting as a latent refiner. The continuous decoder leverages self-attentions (SA), cross-attention (CA) and a local INR to map back to the physical space. Learnable tokens are shared and encode spatial relations. Latent token $Z^t$ represents $u_t$ and $Z^{t+\Delta t}$ is the prediction corresponding to $u_{t+\Delta t}$.
• Figure 2: Spatial interpretation of the tokens through cross attention between $T^{geo}$ and $\bm{\gamma}(x)$ for each $x$ in the domain. Here we visualize the cross-attention of three different tokens for a given head. The cross attentions can have varying receptive fields depending on the geometries.
• Figure 3: Correlation over time for long rollouts with different methods on Burgers
• Figure 4: During training, we noise the next-step latent tokens ${\bm{Z}}^{t+\Delta t}$ and train the transformer to predict the "velocity" of the noise. Each DIT block is implemented as in peebles2023scalable.
• Figure 5: At inference, we start from $\tilde{{\bm{Z}}}^{t + \Delta t}_K \sim \mathcal{N}(0, I)$ and reverse the diffusion process to denoise our prediction. We set our prediction $\hat{{\bm{Z}}}^{t + \Delta t} = \tilde{{\bm{Z}}}^{t + \Delta t}_0$.