ENMA: Tokenwise Autoregression for Generative Neural PDE Operators

2025-06-06

AI Summary

To address the weak generalization of neural operators for time-varying parametric PDEs under data sparsity, uncertainty, and cross-parameter-domain scenarios, this paper proposes ENMA—a generative neural operator. Methodologically, ENMA introduces a novel token-level autoregressive generation paradigm, integrating flow matching loss with context-aware latent-space modeling. It unifies a generative masked Transformer, multi-head attention, and spatiotemporal convolutional encoders to enable parameter-agnostic conditional generation in the latent space. ENMA supports single-shot surrogate modeling, irregular spatial sampling inputs, and in-context learning at inference time. Empirically, it achieves significantly improved spatiotemporal prediction accuracy and robustness on unseen physical parameters and dynamical regimes. By enabling high-fidelity, data-efficient modeling of complex physical systems with minimal supervision, ENMA establishes a new paradigm for generalizable scientific machine learning.

Technology Category

• Machine Learning: Deep Generative Models & Autoencoders
• Natural Language Processing: Generation
• Computer Vision: Generative Adversarial Networks (GANs) for Vision

Application Category

• Social Networks and Social Media: Generative AI / large language models and their impact on social systems
• Economics, Online Markets and Human Computation: Economic ramifications for generative AI infrastructure and applications
• Search and Retrieval-Augmented AI: Retrieval-Augmented Generation (RAG) and multi-modal RAG

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.

Problem

• Solving time-dependent parametric PDEs with neural solvers
• Handling uncertain or incomplete data using generative models
• Generalizing across diverse physical parameters and dynamics

Innovation

• Generative masked autoregressive transformer for tokenwise prediction
• Attention-based encoding of irregular spatial observations
• Spatio-temporal convolutional encoder for latent compression