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

arXiv: https://arxiv.org/abs/2410.03437 · HTML: https://mlsanity.com/paper/2410.03437/

## Key facts

- Cited by 16 papers in the map (data through 14 Oct 2026).
- Builds on 40 of its 59 references that are in the map.
- First posted on arXiv 4 Oct 2024.

## Builds on (most cited first, 30 of 40)

- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [LLaMA: Open and Efficient Foundation Language Models](https://mlsanity.com/paper/2302.13971/) (2023)
- [Scalable Diffusion Models with Transformers](https://mlsanity.com/paper/2212.09748/) (2022)
- [Model-Agnostic Meta-Learning for Fast Adaptation of Deep Networks](https://mlsanity.com/paper/1703.03400/) (2017)
- [Neural Discrete Representation Learning](https://mlsanity.com/paper/1711.00937/) (2017)
- [Neural Ordinary Differential Equations](https://mlsanity.com/paper/1806.07366/) (2018)
- [Taming Transformers for High-Resolution Image Synthesis](https://mlsanity.com/paper/2012.09841/) (2020)
- [Fourier Neural Operator for Parametric Partial Differential Equations](https://mlsanity.com/paper/2010.08895/) (2020)
- [Sinkhorn Distances: Lightspeed Computation of Optimal Transportation Distances](https://mlsanity.com/paper/1306.0895/) (2013)
- [Video Diffusion Models](https://mlsanity.com/paper/2204.03458/) (2022)
- [HuggingFace's Transformers: State-of-the-art Natural Language Processing](https://mlsanity.com/paper/1910.03771/) (2019)
- [MaskGIT: Masked Generative Image Transformer](https://mlsanity.com/paper/2202.04200/) (2022)
- [Visual Autoregressive Modeling: Scalable Image Generation via Next-Scale Prediction](https://mlsanity.com/paper/2404.02905/) (2024)
- [Deep Networks with Stochastic Depth](https://mlsanity.com/paper/1603.09382/) (2016)
- [Deep Anomaly Detection with Outlier Exposure](https://mlsanity.com/paper/1812.04606/) (2018)
- [An Explanation of In-context Learning as Implicit Bayesian Inference](https://mlsanity.com/paper/2111.02080/) (2021)
- [Autoregressive Image Generation without Vector Quantization](https://mlsanity.com/paper/2406.11838/) (2024)
- [Language Model Beats Diffusion -- Tokenizer is Key to Visual Generation](https://mlsanity.com/paper/2310.05737/) (2023)
- [Finite Scalar Quantization: VQ-VAE Made Simple](https://mlsanity.com/paper/2309.15505/) (2023)
- [Message Passing Neural PDE Solvers](https://mlsanity.com/paper/2202.03376/) (2022)
- [Axial Attention in Multidimensional Transformers](https://mlsanity.com/paper/1912.12180/) (2019)
- [MAGVIT: Masked Generative Video Transformer](https://mlsanity.com/paper/2212.05199/) (2022)
- [PDE-Refiner: Achieving Accurate Long Rollouts with Neural PDE Solvers](https://mlsanity.com/paper/2308.05732/) (2023)
- [Why Can GPT Learn In-Context? Language Models Implicitly Perform Gradient Descent as Meta-Optimizers](https://mlsanity.com/paper/2212.10559/) (2022)
- [Poseidon: Efficient Foundation Models for PDEs](https://mlsanity.com/paper/2405.19101/) (2024)
- [Towards Multi-spatiotemporal-scale Generalized PDE Modeling](https://mlsanity.com/paper/2209.15616/) (2022)
- [DPOT: Auto-Regressive Denoising Operator Transformer for Large-Scale PDE Pre-Training](https://mlsanity.com/paper/2403.03542/) (2024)
- [Towards Foundation Models for Scientific Machine Learning: Characterizing Scaling and Transfer Behavior](https://mlsanity.com/paper/2306.00258/) (2023)
- [DiffusionPDE: Generative PDE-Solving Under Partial Observation](https://mlsanity.com/paper/2406.17763/) (2024)

## Built on by (most cited first, 16 of 16)

- [Walrus: A Cross-Domain Foundation Model for Continuum Dynamics](https://mlsanity.com/paper/2511.15684/) (2025)
- [DISCO: learning to DISCover an evolution Operator for multi-physics-agnostic prediction](https://mlsanity.com/paper/2504.19496/) (2025)
- [Test-time Generalization for Physics through Neural Operator Splitting](https://mlsanity.com/paper/2602.00884/) (2026)
- [Towards Foundational Models for Dynamical System Reconstruction: Hierarchical Meta-Learning via Mixture of Experts](https://mlsanity.com/paper/2502.05335/) (2025)
- [Weight-Space Linear Recurrent Neural Networks](https://mlsanity.com/paper/2506.01153/) (2025)
- [Graph In-Context Operator Networks for Generalizable Spatiotemporal Prediction](https://mlsanity.com/paper/2603.12725/) (2026)
- [Efficient Generative Transformer Operators For Million-Point PDEs](https://mlsanity.com/paper/2512.04974/) (2025)
- [Chain of Operators: An Inference-Time Harness for In-Context Operator Learning](https://mlsanity.com/paper/2606.12318/) (2026)
- [Factorized Neural Implicit DMD for Parametric Dynamics](https://mlsanity.com/paper/2603.10995/) (2026)
- [A Theory of Diversity for Random Matrices with Applications to In-Context Learning of Schr"odinger Equations](https://mlsanity.com/paper/2601.12587/) (2026)
- [Posterior-First Neural PDE Simulation: Inferring Hidden Problem State from a Single Field](https://mlsanity.com/paper/2605.03247/) (2026)
- [ARC-STAR: Auditable Post-Hoc Correction for PDE Foundation Models](https://mlsanity.com/paper/2605.22222/) (2026)
- [A foundation model of numerical intelligence with cross-disciplinary generalization](https://mlsanity.com/paper/2607.28432/) (2026)
- [PDE-JEPA: Predictive Representation Learning of Latent Dynamics Modeling for Parametric PDEs](https://mlsanity.com/paper/2609.34715/) (2026)
- [Cross-Domain Pretraining for Steady-State Neural CFD Surrogates](https://mlsanity.com/paper/2610.10398/) (2026)
- [SPLICE: Sparse Physics Library Identification and Compositional Expansion](https://mlsanity.com/paper/2610.17029/) (2026)

## Abstract (arXiv)

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.