Zebra: In-Context and Generative Pretraining for Solving Parametric PDEs

2024-10-04

arXiv.org

AI Summary

To address the limitations of gradient-dependent inference, poor generalization, and low efficiency in neural surrogate models for time-varying parametric partial differential equations (PDEs), this paper proposes a gradient-free, generative autoregressive Transformer framework. Inspired by large language models, it pioneers the adaptation of in-context learning to PDE solving: states are serialized into tokenized sequences, context-conditioned modeling is employed, and pretraining enables zero-shot adaptation to unseen parameter configurations. The framework supports variable-length context inputs, zero-shot parameter generalization, and sampling-based uncertainty quantification. Evaluated on multiple challenging parametric PDE benchmarks, it significantly outperforms existing neural solvers, achieving simultaneous improvements in generalization, robustness, and inference efficiency.

Technology Category

• Machine Learning: Deep Generative Models & Autoencoders
• Computer Vision: Diffusion Models for Vision
• Natural Language Processing: Generation

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• User Modeling, Personalization and Recommendation: Explainable and interpretable methods for personalization
• Graph Algorithms and Modeling for the Web: Graph neural networks and deep learning approaches for Web-related graphs

Abstract

Solving time-dependent parametric partial differential equations (PDEs) is challenging, as models must adapt to variations in parameters such as coefficients, forcing terms, and boundary conditions. Data-driven neural solvers either train on data sampled from the PDE parameters distribution in the hope that the model generalizes to new instances or rely on gradient-based adaptation and meta-learning to implicitly encode the dynamics from observations. This often comes with increased inference complexity. Inspired by the in-context learning capabilities of large language models (LLMs), we introduce Zebra, a novel generative auto-regressive transformer designed to solve parametric PDEs without requiring gradient adaptation at inference. By leveraging in-context information during both pre-training and inference, Zebra dynamically adapts to new tasks by conditioning on input sequences that incorporate context trajectories or preceding states. This approach enables Zebra to flexibly handle arbitrarily sized context inputs and supports uncertainty quantification through the sampling of multiple solution trajectories. We evaluate Zebra across a variety of challenging PDE scenarios, demonstrating its adaptability, robustness, and superior performance compared to existing approaches.

Problem

• Solving time-dependent parametric PDEs with varying parameters
• Adapting to new PDE tasks without gradient-based optimization
• Generating trajectories and quantifying prediction uncertainty

Innovation

• Generative auto-regressive transformer for PDEs
• In-context learning without gradient adaptation
• Uncertainty quantification in trajectory generation