ENMA: Tokenwise Autoregression for Generative Neural PDE Operators (2506.06158v1)

Abstract: Solving time-dependent parametric partial differential equations (PDEs) remains a fundamental challenge for neural solvers, particularly when generalizing across a wide range of physical parameters and dynamics. When data is uncertain or incomplete-as is often the case-a natural approach is to turn to generative models. We introduce ENMA, a generative neural operator designed to model spatio-temporal dynamics arising from physical phenomena. ENMA predicts future dynamics in a compressed latent space using a generative masked autoregressive transformer trained with flow matching loss, enabling tokenwise generation. Irregularly sampled spatial observations are encoded into uniform latent representations via attention mechanisms and further compressed through a spatio-temporal convolutional encoder. This allows ENMA to perform in-context learning at inference time by conditioning on either past states of the target trajectory or auxiliary context trajectories with similar dynamics. The result is a robust and adaptable framework that generalizes to new PDE regimes and supports one-shot surrogate modeling of time-dependent parametric PDEs.

• The paper introduces a novel tokenwise autoregressive generative framework that accurately predicts spatio-temporal dynamics in parametric PDEs.
• By encoding irregular data into a structured latent space and utilizing a masked autoregressive Transformer, ENMA enhances scalability and inference speed.
• Extensive experiments demonstrate that ENMA outperforms benchmark models in accuracy and uncertainty quantification, paving the way for real-world applications.

Overview of ENMA: Tokenwise Autoregression for Generative Neural PDE Operators

This paper introduces ENMA, a novel approach in the computational modeling of parametric Partial Differential Equations (PDEs) that leverages continuous autoregressive neural operators. The ENMA framework addresses the fundamental challenge of predicting spatio-temporal dynamics from PDEs under uncertainty and varying physical regimes. The model employs a novel generative format to efficiently and accurately predict dynamic systems by autoregressing over latent representations rather than discrete tokens or full-frame data, enhancing both scalability and fidelity in predictions.

Key Components of ENMA

ENMA unfolds in an encode-generate-decode cycle:

1. Encoding: The framework begins with a sophisticated encoder that maps irregular and incomplete spatio-temporal data into a structured latent space using attention mechanisms and causal convolutions. This operation results in uniform latent tokens that maintain essential context from the physical phenomena despite data irregularities.
2. Generation: The generative segment utilizes a masked autoregressive Transformer architecture for latent space prediction. By leveraging flow matching, ENMA dynamically models conditional probability distributions for each token—supporting fine granularity and efficient sampling compared to traditional diffusion models.
3. Decoding: The process concludes with decoding generated latents back into the physical domain using a continuous variational autoencoder (VAE) structure, thus producing predictions that adhere closely to the original PDE dynamics.

Numerical and Theoretical Results

ENMA demonstrates notable efficacy through extensive experiments on various dynamical systems, including 1D and 2D PDEs like the Advection, Wave, Combined, Gray-Scott, and Vorticity equations. It consistently surpasses benchmark models in terms of:

• Accuracy: In standard forecasting tasks, ENMA achieves lower relative MSE compared to alternatives, particularly under conditions of sparse or incomplete input data and out-of-distribution regimes.
• Scalability: By operating predominantly in latent space, ENMA reduces computational overheads inherent in pixel-based or discrete token models, providing faster inference speeds while retaining precise predictive quality.
• Generative Capabilities: ENMA supports robust probabilistic forecasting, enabling uncertainty quantification and innovative trajectory generation in scenarios void of explicit initial conditions. These features are pivotal in chaotic systems and real-world applications like weather and turbulence modeling.

Implications and Future Directions

The implications of ENMA are profound, offering improvements in how AI systems model physical systems governed by PDEs. Its continuous generative approach paves the way for faster, more reliable surrogate models in engineering and scientific domains. It also suggests potential integrations with other AI subfields, such as reinforcement learning, to enhance adaptive control systems in dynamic environments. For future research, adaptations of ENMA could explore its applications in high-dimensional spaces or multi-physics coupling scenarios, further extending AI's role in simulating complex systems with near real-time computational efficiency.

In conclusion, ENMA represents a significant step forward in generative modeling for PDEs, seamlessly blending the precision of scientific computation with the adaptability of AI methodologies. Its introduction heralds an era of enriched modeling capabilities with broad applicability across multiple domains reliant on accurate PDE solvers.