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

TL;DR

AROMA introduces a flexible neural framework that preserves spatial structure for PDE modeling, enabling efficient handling of diverse geometries and stable long-term simulations through diffusion-based training.

Contribution

It presents a novel encoder-decoder architecture with attention and diffusion techniques for improved PDE modeling from irregular and diverse data sources.

Findings

Outperforms existing methods in 1D and 2D PDE simulations.

Provides stable long-term predictions with diffusion-based training.

Effectively handles irregular grids and point clouds.

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.