Research map / Probabilistic and scientific modeling / Differential equations and PDEs
Zebra: In-Context Generative Pretraining for Solving Parametric PDEs
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Zebra is a generative autoregressive transformer over VQ-VAE tokens that solves parametric PDEs via in-context learning from example trajectories, with no gradient adaptation at inference.
Problem
Neural surrogates for time-dependent parametric PDEs must adapt to new coefficients, boundary conditions, and forcing terms; existing adaptation methods rely on gradient-based fine-tuning or meta-learning, which adds inference cost.
Approach
Pretrain a causal generative transformer on token sequences of concatenated trajectories sharing the same dynamics (different initial conditions), so that at inference it adapts in-context from a few example trajectories and unrolls from a new initial condition, without weight updates; a VQ-VAE compresses states into discrete tokens.
Key results
- In one-shot in-distribution evaluation, Zebra achieves the lowest relative L2 on Burgers (0.115), Wave b (0.245), Wave 2D (0.207) and Vorticity 2D (0.119), outperforming CODA and CAPE in 2D where CAPE diverges.
- On out-of-distribution Heat (Wide forcing) Zebra reaches 0.15 relative L2 versus 1.03 for CODA and 0.47 for CAPE.
- Zebra + UNet reduces inference time to 0.10s on Advection and 0.14s on Vorticity 2D versus 18-31s for CAPE and CODA.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use in-context Zebra instead of the strongest listed baseline | One-shot PDE adaptation · Burgers, Wave b, Wave 2D, Vorticity 2D | relative L2: 0.116 -> 0.115; 0.971 -> 0.245; 0.271 -> 0.207; 0.678 -> 0.119 | Helped |
| Use in-context Zebra instead of the strongest listed baseline | One-shot PDE adaptation · Heat | relative L2: 0.136 -> 0.154 | Hurt |
| Use in-context Zebra instead of the strongest listed baseline | One-shot PDE adaptation · Advection, Combined | relative L2: 0.00687 -> 0.00794; 0.00857 -> 0.00965 | Mixed |
| Use in-context Zebra instead of the strongest listed baseline | One-shot PDE adaptation · Heat Wide, Heat Gaussian, Heat Square, Vorticity 2D close, Vorticity 2D far | relative L2: 0.47 -> 0.15; 0.40 -> 0.32; 0.66 -> 0.36; 0.30 -> 0.24; 0.37 -> 0.32 | Helped |
| Use Zebra rather than ViT-in-context | Out-of-distribution one-shot adaptation · Wave 2D | relative L2: 0.68 -> 0.68 | No clear effect |
| Replace token-wise autoregressive decoding with a [DYN]-conditioned UNet (ablation) | Zebra · Advection | relative L2: 0.00794 -> 0.0072 | Helped |
| Replace token-wise autoregressive decoding with a [DYN]-conditioned UNet (ablation) | Zebra · Combined | relative L2: 0.00965 -> 0.0138 | Hurt |
| Replace token-wise autoregressive decoding with a [DYN]-conditioned UNet (ablation) | Zebra · Wave 2D, Vorticity 2D | relative L2: 0.207 -> 0.150; 0.119 -> 0.0869 | Helped |
| Replace token-wise autoregressive decoding with a [DYN]-conditioned UNet (ablation) | Zebra · Advection, Vorticity 2D | inference time: 3s -> 0.10s; 21s -> 0.14s | Helped |
| Train deterministically to predict the conditional expectation of the next token rather than model trajectory distributions (ablation) | Zebra · Autoregressive rollout | relative L2 error | Hurt |
| Increase the number of context examples | Zebra · Heat uncertainty quantification | uncertainty | Helped |
Surprising findings
- Gradient-based baselines CAPE and CODA diverge or degrade badly on 2D datasets and under out-of-distribution shifts, while in-context approaches are more robust.
- A deterministic variant accumulates significant error during autoregressive rollout, whereas sampling from the learned trajectory distribution is more robust.
Limitations
- Fidelity limited by VQ-VAE reconstruction of fine or high-frequency details; predictions for turbulent OoD are less sharp.
- Requires large numbers of training trajectories and parameter diversity; not suited to data-scarce regimes.
- Quadratic attention cost over context, time, and space tokens limits scaling to 3D; encoder/decoder restricted to regular grids.
Experimental setup
Datasets: Advection Heat Burgers Wave boundary Combined equation Wave 2D Vorticity 2D
Benchmarks: Advection Heat Burgers Wave b Combined Wave 2D Vorticity 2D
Metrics: relative L2 error Wasserstein distance inference time
Baselines: CODA CAPE [CLS] ViT ViT-in-context
Builds on: VQ-VAE Llama CODA CAPE VQGAN
Closest papers by meaning
- CALM-PDE: Continuous and Adaptive Convolutions for Latent Space Modeling of Time-dependent PDEs (NeurIPS 2025 · similarity 0.87)
- Boosting Generalization in Parametric PDE Neural Solvers through Adaptive Conditioning (NeurIPS 2024 · similarity 0.87)
- Text2PDE: Latent Diffusion Models for Accessible Physics Simulation (ICLR 2025 · similarity 0.86)
- Vectorized Conditional Neural Fields: A Framework for Solving Time-dependent Parametric Partial Differential Equations (ICML 2024 · similarity 0.86)
- DiffusionPDE: Generative PDE-Solving under Partial Observation (NeurIPS 2024 · similarity 0.85)
- Poseidon: Efficient Foundation Models for PDEs (NeurIPS 2024 · similarity 0.85)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.