Zebra: In-Context Generative Pretraining for Solving Parametric PDEs

TL;DR

Zebra is a novel transformer-based model that uses in-context learning to solve parametric PDEs efficiently without gradient updates, enabling dynamic adaptation, trajectory generation, and uncertainty quantification.

Contribution

Introducing Zebra, a generative auto-regressive transformer that leverages in-context learning for solving parametric PDEs without gradient-based adaptation.

Findings

Zebra outperforms existing methods in accuracy and robustness.

It can generate new solution trajectories effectively.

Zebra provides uncertainty quantification for predictions.

Abstract

Solving time-dependent parametric partial differential equations (PDEs) is challenging for data-driven methods, as these models must adapt to variations in parameters such as coefficients, forcing terms, and initial conditions. State-of-the-art neural surrogates perform adaptation through gradient-based optimization and meta-learning to implicitly encode the variety of 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 example trajectories. As a generative model, Zebra can be used to generate new trajectories and allows quantifying the uncertainty of the predictions. We evaluate Zebra across a variety of challenging PDE scenarios, demonstrating its adaptability, robustness, and superior performance compared to existing approaches.