Zebra: In-Context Generative Pretraining for Solving Parametric PDEs
Authors: Louis Serrano
Armand Kassaï Koupaï
Thomas X Wang
Pierre ERBACHER
Patrick Gallinari
Venue: ICML 2025 poster (2025)
Keywords: PDE
Domain adaptation
In-Context Learning
Generative Pre-training
Autoregressive Transformer
Primary area: applications->chemistry_physics_and_earth_sciences
TL;DR: Zebra is a generative auto-regressive transformer for solving parametric PDEs, leveraging in-context information to adapt to new dynamics.

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.

Zebra: In-Context Generative Pretraining for Solving Parametric PDEs
Louis Serrano' Armand Kassaï Koupai' Thomas X Wang' Pierre Erbacher 2 Patrick Gallinari 13
Abstract
Li et al., 2021; Koupai et al., 2024).
Solving time-dependent parametric partial differ- ential equations (PDEs) challenging for data- driven methods, as these models must adapt to
We address the parametric PDE problem (Cohen & Devore, 2015), where the goal is to train models on a limited set of scenarios representing a given physical phenomenon so that they can generalize across a wide range of new con- texts, including different PDE parameters. These parameters may encompass initial and boundary conditions, physical coefficients, and forcing terms. In this work, focus on purely data-driven approaches that do not incorporate prior knowledge of the underlying equations.
variations in parameters such as coefficients, fore- ing 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 in- ference complexity. Inspired by the in-context learning capabilities of large language models (LLMs), we introduce Zebra, a novel genera- tive auto-regressive transformer designed to solve parametric PDEs without requiring gradient adap- tation at inference. By leveraging in-context in- formation during both pre-training and inference, Zebra dynamically adapts to new tasks by con- ditioning input sequences that incorporate con- text example trajectories. As generative model, Zebra can be used to generate new trajecto- ries and allows quantifying the uncertainty of the predictions. We evaluate Zebra across a vari- ety of challenging PDE scenarios, demonstrating its adaptability, robustness, and superior perfor- mance compared to existing approaches.
A basic approach to this problem is to sample from the distribution of physical parameters, i.e., to train on differ- ent instances of PDE characterized by varying parameter values, with the goal of generalizing to unseen instances. i.i.d. assumption and requires
This approach relies training set that adequately represents the distribution of the underlying dynamical system-a a condition that is of- ten difficult to satisfy in practice due to the complexity of physical phenomena. Other approaches explicitly condi- tion on specific PDE parameters (Brandstetter et al., 2022b; Takamoto et al., 2023), relying on the availability of such prior knowledge. This assumes that a physical model of the observed system is known, making the incorporation of PDE parameters into neural solvers challenging beyond physical quantities. Moreover, in many cases, this prior knowledge is incomplete or entirely unknown.
An alternative approach involves adaptation to new PDE in- stances by leveraging observations from novel environments. Here we consider that an environment is characterized by set of parameters. For data-driven models, adaptation is often performed through fine tuning, which usually requires a significant amount of examples for the new environment. This is for example the setting adopted in many recent de- velopment of foundation models (Subramanian et al., 2023; Herde al., 2024; McCabe et al., 2023). This involves train- ing large model on variety of physics bbsied numerical simulations, each requiring a large amount of simulations with the expectation that it will generalize.
1. Introduction
major challenge for training neural solvers for time depen- dent partial differential equations (PDEs) or more generally for modeling spatio-temporal dynamics is to capture the variety of behaviors arising from complex physical phenom- ena. In particular, neural solvers, trained from a limited number of situations often fail to generalize to new physical contexts and situations (Chen et al., 2018; Raissi et al., 2019;
Our code is available on GitHub. 'Sorbonne Université, CNRS, ISIR, 75005 Paris, France "Naver Labs Europe, France 'Critéo AI Lab, Paris, France. Conresoondence to: Louis Ser- rano <louis. serrano @ werboonn-nniversieffrr... Patrick Gallinari <patrick.gallinariie @ sorbonne-une subbonne-nniversieefft.
More principled frameworks for adaptation leverage meta- learning, where the model is trained on simulations cor- responding to different environments-iie. varying PDE parameter values-so so that it can quickly adapt new and unseen PDE simulation instances using few trajectory ex-
Proceedings of the 42nd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Zebra: In-Context Generative Pretraining for Solving Parametric PDEs
amples (Park et al., 2023; Kirchmeyer et al., 2022; Yin et al., 2022). These flexible methods rely on gradient updates for adaptation, adding computational overhead.
next token objective. At inference, Zebra can handle varying context sizes for conditioning and supports uncer- tainty quantification, enabling generalization to unseen PDE parameters without gradient updates.
We explore a new direction for adaptation inspired by the successes of in-context learning (ICL) in natural language processing (NLP) and its demonstrated ability to generalize to downstream tasks without retraining (Brown al., 2020; Touvron al., 2023). We propose a framework, denoted Zebra, relying ICL for solving parametric PDEs with new parameter values, without any additional update of the model parameters.
Our main contributions include:
We introduce a generative autoregressive transformer for modeling physical dynamics. It operates on com- pact discretized representations of physical state obser- vations. This framework represents the first successful application of causal generative modeling using quan- tized representations of physical systems.
As for ICL in NLP, the model is trained to generate appro- priate responses given context examples and a query. The context examples will be trajectories from the same dynam- ies starting from different initial conditions. The query will consist for example of new initial state condition, that will serve as inference starting point for the new forecast. The proposed model is inspired from NLP approaches: it is a causal generative model that processes discrete token sequences encoding observations. It is trained to model the trajectory distributions of parameteric PDEs. This approach offers key advantages and greater flexibility compared to existing methods. It can leverage contexts of different types and sizes, requires only a few context examples to adapt to new dynamics, and allows us to cover a wide range of situations. It provides enhanced capabilities compared to more classical deterministic forecasting models. Notably, generative probabilistic models have been developed for physical problems such as weather forecasting (Price et al., 2023; Couairon et al., 2024) and even PDE solving (Lippe et al., 2023), demonstrating superior performance and capa- bilities over their deterministic counterparts, However, their setting is different, as they rely on diffusion models and are neither designed for adaptation nor intended to address the parametric PDE problem.
To hamness the in-context learning strengths of autore- gressive transformers, we develop a new pretraining strategy that conditions the model on example trajec- tories with similar underlying dynamics but different initial conditions.
Our generative model predicts trajectory distributions. This provides a richer information than deterministic auto-regressive models. This comes with enhanced capabilities including more accurate predictions, uncer- tainty measures, the ability to sample and generate new trajectories conditioned on some examples.
We propose accelerated inference procedure that is orders of magnitude faster than existing adaptation methods.
We evaluate Zebra in one-shot adaptation setting, where it must infer the dynamics from single context trajectory and unroll the dynamics from a new query initial condition (i.e., a single snapshot). The predicted trajectory is compared to a target trajectory governed by the same underlying dynamics as the context but starting from the query initial condition. Zebra's 's per- formance is benchmarked against domain adaptation baselines specifically trained for this task, and it consis- tently outperforms gradient-base -based adaptation methods on 2D datasets.
Some recent works have also explored adaptation through in-context learning for dynamics modeling. The closest to ours is probably (Yang et al., 2023), which also targets adap- tation to multiple environments of an underlying physical dynamics through prompting with examples. Their model employs a specific deterministic encoder-decoder architec- ture and limited to 1D ODEs or sparse 2D data due to scalability issues. More details and further references are provided in Appendix A.
2. Problem setting
2.1. Solving parametric PDEs
We aim to solve parametric time-dependent PDEs beyond the typical variation in initial conditions. Our goal is to train models capable of generalizing across wide range of PDE parameters. To this end, we consider time-dependent PDEs with different initial conditions, and with additional degrees of freedom, namely: (1) coefficient parameters such as fluid viscosity or advection speed - denoted by vector P (2) boundary conditions B, e.g. Neumann or Dirichlet; (3) forcing terms 6, including damping param- eter or sinusoidal forcing with different frequencies. We
On the technical side, Zebra introduces a novel gener- ative autoregressive solver for parametric PDEs. It em- ploys an encode-gend encode-generate- weeegeeee-eeoodeeeeeeeeee decode framework: first, a vector- quantized variational auto-encoder (VQ-VAE) (Oord et al., 2017) is learned to compress physical states into discrete tokens and decode back to the original physical space. Next, a generative autoregressive transformer is pre-trained with arbitrary size context examples of trajectories using
2

---

Zebra: In-Context Generative Pretraining for Solving Parametric PDEs
denote ç : B,6 and we define F&g as the set of PDE solutions corresponding to the PDE parameters 14, bound- ary conditions B and forcing term 6, and refer to Feg as PDE environment. Formally, a solution u(x, t) within F&g satisfies:
and Og are respectively the shared and specific parameters. At inference, for a new environment, only small number of parameters 0g is adapted from a small sample of observa- tions. This family of method will be our reference baseline in the following.
du
(6, u, du du dr' 0x2 .....
3. Zebra Framework
Ōt
We introduce Zebra, novel framework designed to solve parametric PDEs through in-context learning and flexible conditioning. Zebra utilizes generative autoregressive transformer to model partial differential equations (PDEs) within a compact, discrete latent space. A spatial CNN encoder is employed to map physical spatial observations into these latent representations, while CNN decoder accu- rately reconstructs them. Our pretraining pipeline consists of two key stages: 1) Learning finite vocabulary of phys- ical phenomena, and 2) Training the transformer using an in-context pretraining strategy, enabling the model to effec- tively condition on contextual information. At inference, Zebra allows to perform in-context learning from context trajectories as illustrated in Figure 1.
V Ω, VI E (0, (a,) t) = 0. E 02, Vt (0. (0 2) = u., A
(I)
where F is function of the solution u and its spatial deriva- tives on the domain Ω2 and also includes the forcing term 5 ; is the boundary condition constraint (e.g., spatial peri- odicity, Dirichlet, or Neumann) that must be satisfied at the boundary of the domain a2; and u is the initial condition sampled with probability measure u² ~ p.).
2.2. Adapaation for parametric PDE
Solving time-dependent parametric PDEs requires devel- oping neural solvers capable of generalizing to whole distribution of PDE parameters. In practice, changes in the PDE parameters often lead to distribution shifts in the tra- jectories which makes the problem challenging. Different directions are currently being explored and are briefly re- viewed below. We focus on pure ddiven approaches that do not make use of any prior knowledge on the equations. We make the assumption that the models are learned from numerical simulations so that it is possible to generate from multiple parameters. This emulates real situations where for example, a physical phenomenon is observed in different contexts.
3.1. Learning a finite vocabulary of physical phenomena
In order to leverage the auto-regressive transformer archi- tecture and adopt a next-token generative pretraining, we need to convert physical observations into discrete represen- tations. We do not quantize the observations directly but rather quantize compressed latent representations by em- ploying VQVAE (Oord et al., 2017), an encoder-decoder architecture with a quantizer component. Our encoder spa- tially compresses the input function u' through a convo- lutional model E which maps the input to continuous latent variable z' (u'). The latent variables are then quantized a vector of discrete codes zi using codebook Z of size K = |Z| through the quantization component q. For each spatial code 21,) in z the nearest codebook entry Z is selected. The decoder Dy reconstructs the signal ù' from the quantized latent codes ét. Both models are jointly trained to minimize the reconstruction error between the function u and its reconstruction n' Di o q e (u')).
Fine tuning pre-trained models The classical strategy for adapting to new settings is to fine tune models that have been pretrained on a distribution of the PDE parameters. This approach often relies on large fine tuning samples and involves updating all or a subset of parameters (Subrama- nian et al., 2023; Herde et al., 2024). We do not consider this option that has been shown to underperform SOTA adaptation approaches (Koupai et al., 2024).
Once this training step is done, we can tokenize a trajec- tory u tmt+ by applying our encoder in parallel on each timestamp to obtain vectors of discrete codes z tft++mr and retrieve the corresponding index entries t.+++ma from the codebook Z. Similarly, we detokenize discrete indices with the decoder. We provide a brief description of the VQVAE model and details on its architecture Appendix C.
Gradient-base based adaptation A more flexible approach relies on adaptation at inference time through meta-learning. posits that set of environments are available from which trajectories are sampled, each environment being defined by specific PDE parameter values (Zintgraf et al., 2019a; Kirchmeyer et al., 2022). The model is trained from a sampling from the environments distribution to adapt fast to a new environment. The usual formulation is to learn shared and specific environment parameters Go where 0
3.2. In-context modeling
We design sequences that enable Zebra to perform in- context learning on trajectories that share underlying dy-
3

---

Zebra: In-Context Generative Pretraining for Solving Parametric PDEs
Query Initial
Example Trajectory
Example Trajectory 2
Condition (IC)
Generated Trajectory for IC
De Tokennize
R
Jm
Time
genwaled
Tokenize
Tokenize
Tokenize
Auto regressive generation
Figure Zebra's inference pipeline from context trajectories. The model is provided with a set of example trajectories (u₃ and t2) that share the same underlying dynamics (e.g., identical PDE parameters) but differ in initial conditions. Given a new query initial condition u.. Zebra aims to generate corresponding trajectory that follows the same physical behavior. Each trajectory and initial condition first tokenized into flattened index sequences then autoregressively generates remaining tokens to obtain a plausible prediction. Finally, the generated indices are detokenized reconstruct the trajectory in physical space.
and sQ, which concatenated using specific formatting scheme. The transformer
namics with different initial states. To incorporate varying amounts of contextual information, draw number n E {1, Pmmx }. then sample trajectories sharing the same dy- namics, each with m snapshots starting from time f, denoted as (u f:t+mb ..... u (f)) These trajectories are tok- enized into index representations (s 1ttt++Ar
et al., 2018). Given a sequence S of discrete tokens of length N, the model is optimized to minimize the negative log-likelihood (cross-entropy loss):
Lrransformer
Es
Σ log (Sgg((s 1=1
A()) which are flattened into sequences 81,..., 8u: maintaining the temporal order from left right. In practice, we fix 7mmx and m 9.
where the model learns predict each token Syg condi- tioned on all previous tokens Spreq- Due to the trans- former's causal structure, earlier tokens in the sequence are not influenced by later ones, while later tokens benefit from more context, allowing for more accurate predictions. This structure naturally supports both generation in one- shot and few-shot setting within a unified framework. Our transformer implementation is based on the Llama archi- tecture (Touvron et al., 2023). Additional details can be found in Appendix C. Up our knowledge, this is the first adaptation of generative auto-regressive transformers to the modeling of physical dynamics.
Since our model operates on tokens from codebook, we found it advantageous to introduce special tokens to struc- ture the sequences. The tokens <bot>> (beginning of trajec- tory) and <eot> (end of trajectory) define the boundaries of each trajectory within the sequence. Furthermore, as we sample sequences with varying context sizes, we maximize the utilization of the transformer's is context window by stack- ing sequences that could also represent different dynamics. To signal that these sequences should not influence each other, we use the special tokens <bos> (beginning of se- quence) and <eos> (end of sequence). The final sequence design is:
3.4. Flexible inference: prompting and sampling
= <bot <bot>>ssk<eot>>>>>>>>>>>>>>>>>>> <s. <eot>
In this section, we outline the inference pipeline for Zebra across various scenarios. For simplicity, we assume that observations have already been tokenized and omit the detokenization process. Let s, represent the target token sequence be predicted.
And our pretraining dataset structured as follows:
<bos>[{||<eo>> <bos> (S)s>ss
3.3. Next-token pretraining
Prompting with n examples and an initial condition: The prompt structured as = <bos><bot> ( [5 _ <eot><bot>(!))
The transformer is trained using self-supervised learning on a next-token prediction task with teacher forcing (Radford

---

Zebra: In-Context Generative Pretraining for Solving Parametric PDES
allowing the model adapt based on the provided examples and initial condition.
exceeding those in previous studies and available simula- tion datasets (Yin al. (2022), Kirchmeyer et al. (2022), Blanke & Lelarge (2023), Nzoyem et al. (2024)). We con- duct experiments across seven datasets: five ID Advection, Heat, Burgers, Wave-b, Combined-and two 2D Wave 21), Vorticity 2D. These datasets were selected to encompass different physical phenomena and test generalization under changes various types of PDE parameters, as described below. The spatial resolution is set to 256 for 1D datasets and 64 64 for 2D datasets. We subsample each trajectory to keep 10 snapshots, except for the Wave-b dataset, where we retain 15,
Prompting with n examples and É frames: This setup combines context from multiple trajectories with the initial timestamps, structured
S = <bos>>>>> <bot> [s] [s0:mAr]
[s (m <eot>
<bot> [
At inference, we adjust the temperature parameter T of the classifier layer to calibrate the level of diversity of the next- token distributions. The temperature 7 scales the logits yi before the softmax function:
Varying PDE coefficients The changing factor is the set of coefficients μ in Equation (1). For Burgers. Heat, and Vorticity 2D equations, the viscosity coefficient v/ varies across environments. For Advection, the advection speed B3 changes. In Wave-c and Wave-2D, the wave's 's celerity e is unique to each environment, and the damping coefficient l varies across environments in Wave-2D. In the Combined equation, three coefficients (a, B, 7) vary, each influencing different derivative terms respectively: On." 0 8¹ on the right-hand side of Equation (1).
Uk
exp
Uk
p(S = k| = softmax
Σ, exp
V1
When 5 > 1. the distribution becomes more uniform, en- couraging exploration, whereas т < 1 sharpens the dis- tribution, favoring more deterministic predictions. During training, is kept fixed 7 =
dx + 02 0g
4. Experiments
In this section, we experimentally validate that our frame- work enables one-shot adaptation at inference. We follow the pretraining procedure outlined in Section 3 for each dataset described in Section 4.1 and evaluate Zebra across distinet scenarios. We first assess its performance in the one-shot setting for in-distribution parameters, comparing it to adaptation-basse based baselines (Section 4.2). We then ex- amine its generalization in out-o--distribution settings in Section 4.3. Next we illustrate and analyze the generative abilities of Zebra through two example tasks: uncertainty quantification and new trajectory generation in Section 4.4, with further analysis in Appendix D.2 and Appendix D.3. Fi- nally, in Section 4.5, we show how we can drastically accel- erate the adaptation at inference compared to gradient-basee based methods. More results are provided in Appendix D, includ- ing: deterministic pretraining (Appendix D.1), scaling be- havior with the number of training samples (Appendix D.4), the effect of codebook size (Appendix D.6), reconstruction capabilities (Appendix D.7), and the impact of the number of context examples (Appendix D.8).
Varying boundary conditions In this case, the vary- ing parameter is the boundary condition from Equation (I). For Wave-b, we explore two types of boundary condi- tions-Dirichlet and Deumann-applicn independently to each boundary, resulting four distinct environments.
Varying forcing term The varying parameter is the forc- ing term ō in Equation (1). In Burgers and Heat, the forcing terms vary by the amplitude, frequency, and shift coeffi-
cients: 0(1, = Ei=
1 A, sin (wyl+ + 2n'4 + +3))
A detailed description of the datasets is provided in Ap- pendix B, while Table 6 summarizes the number of envi- ronments used during training, the number of trajectories sharing the same dynamics, and the varying PDE parameters across environments. For testing, all methods are evaluated trajectories with new initial conditions in previously un- seen environments. These unseen environments include trajectories with both novel initial conditions and varying parameters, which remain within the training distribution for in-distribution evaluation and extend beyond for out- of-distribution testing. For each testing, we use 120 unseen environments for the 2D datasets and 12 for the ID datasets, with each environment containing 10 trajectories.
4.1. Datasets details
As in Kirchmeyer et al. (2022), we generate data in batches, where each batch of trajectories corresponds to single envi- ronment and shares the same PDE parameters while having different initial conditions. We consider various factors of variation across multiple datasets. To assess the generaliza- tion ability of our model across a wide range of scenarios, we use a significantly larger number of environments far
Regarding computational resources, training the VQ-VAE in ID takes approximately 4 hours on an RTX 24 GB GPU, while the transformer component requires around 15 hours. In the 2D setting, both training times increase to approxi- mately 20 hours each on a single A A100 80 GB GPU.
5

---

Zebra: In-Context Generative Pretraining for Solving Parametric PDEs
Table 1. One-shot adaptation. Conditioning from similar trajectory. Test results in relative 1.2 the trajectory. * indicates inference
has diverged.
Advection
Heut
Burgers
Wave b
Combined
Wave 2D
Vorticity 2D
CAPR
0.00941
0.223
0.213
0.978
0.00857
CODA
0.00687
0.546
0.767
1.020
0.0120
0.777
0.678
[CLS] VIT
0.140
0.136
0.116
0.971
0.0446
0.271
0.972
vin-in-context
0.0902
0.472
0.582
0.472
0.0885
0,390
0.173
Zebra
0.00794
0.154
0.115
0.245
0.00965
0.207
0.119
4.2. In-distribution generalization
updates, following the same approach used in CODA and CAPE. context is a transformer with separate temporal and spatial attention (Ho et al., 2019), where we stack context examples and preceding target frames in the temporal axis to provide in-context examples. Note that all these baselines are deterministic.
Setting We evaluate Zebra's ability to perform in- context learning by leveraging example trajectories that follow the same underlying dynamics as the target. For- mally, in the n-shot adaptation setting, we assume access to a set of n context trajectories (uf:mar ...., u (m) at inference time, all of which belong to the same dynamical system Fe.. The goal of the adaptation task is to accurately predict a future trajectory u. Atmmtt from new initial con- dition u.. knowing that the underlying target dynamics is shared with the provided context example trajectories.
Metries We evaluate the performance using the Rel- ative 1 norm between the predicted rollout trajec- tory u,
tnjectoryy and the ground truth u. !njectory, Lis 1idl 20ndayy la
Non Ejere
I 1 1|2
Sampling For Zebra, we use here a random sampling procedure at inference for generating the next tokens for all datasets, setting a low temperature (7 0.1) to prioritize accuracy over diversity. Predictions are generated using single sample under this configuration.
Results As evidenced in Table 5, Zebra demonstrates strong overall performance in the one-shot adaptation set- ting, often surpassing gradient-based adaptation methods. For the more challenging datasets, such as Burgers, Wave- b, and the 2D cases, Zebra consistently achieves lower relative 1.2 errors, highlighting its capacity to model com- plex dynamics effectively. Notably, Zebra excels in 2D environments, outperforming both CODA and [CLS] ViT and avoiding the divergence issues encountered by CAPE. While Zebra performs comparably to CODA on simpler datasets like Advection and Combined, its overall stability and versatility across range of scenarios, particularly in 2D settings, highlight its competitiveness. Overall, Zebra stands out as a reliable and scalable solution for adaptation for solving parametric PDEs, demonstrating that in-context learning offers robust alternative to existing gradient-based adaptation methods. The ablation study (see Appendix D.1) highlights an important advantage of Zebra's 's generative capabilities. When trained deterministically to predict the conditional expectation of the next token, the model accumu- lates significant error during the autoregressive rollout. In contrast, Zebra, trained to model trajectory distributions, can sample from this distribution at inference time, resulting in predictions that are more robust to error accumulation.
Baselines We evaluate Zebra against 4 baselines, CODA (Kirchmeyer et al., 2022) and CAPE (Takamoto et al., 2023), two SOTA adaptation methods. We also compare to two specifically designed ViT architectures: [CLS] ] ViT that performs adaptation by learning [CI.S] embedding and
designed for in-context training. CODA is a meta-learning framework designed for learning parametric PDEs. It leverages common knowledge from multiple environments where trajectories from a same en- vironment e share the same PDE parameter values. CODA training performs adaptation in the parameter space by learn ing shared parameters across all environments and a con- text vector ef specific to each environment. At inference, CODA adapts to new environment by tuning e" with sev- eral gradient steps. CAPE was not designed to perform adaptation via extra-trajectories, but instead needs the cor- rect parameter values as input to condition a neural solver. We adapt it to our setting, by learning a context c instead of using the real parameter values. During adaptation, we only tune this context cf via gradient updates. [CLS] J ViT is specifically designed baseline based on a vision trans- former (Peebles & Xie, 2023), integrating a [CLS] token that serves as a learned parameter for each environment. This token lets the model handle different dynamics, and during inference, we adapt the [CLS] vector via gradient
4.3. Out-of-distribution generalization
Datasets We evaluate our models on new PDE instances under the following distribution shifts: (i) Hear: We mod- ify the forcing term parameterization. Wide Forcing: The forcing coefficients are varied from Appendix B.3, with
6

---

Zebra:
In-Context Generative Pretraining for Solving Parametric PDEs
Figure 2. One-shot prediction Vorticity in the turbulent OoD regime E [le 5, le 4]. Top-row is the example given in context, mid-row is the ground truth trajectory, and the bottom-row is the generation with Zebra.
parameters sampled as A; E [-1.0, wi E 1-0.8, 0.8]
in Figure 2, the predictions are not as sharp as the ground truths, as the VQVAE was not explicitly trained to capture the part of the spectrum present in turbulent trajectories.
compared to the training ranges [-0.5, 0.5] and 00.4, 0.4]. Gaussian: ((t, = Σニ A, exp
(==ي) sin (w,t), 20
introducing smooth, localized excitations. This results in solutions similar the homogeneous case away from peaks. Square: x(t,z)
Table 2. -fit-o--iittiiuttio results. Test results in relative L2 on the trajectory. indicates inference has diverged. n/a indicates that the model was not evaluated this setting.
2 1- &5), which preserves the parameterization but introduces dis- continuities and gigh-frequeccy content, resulting in a more challenging OoD scenario. (ii) Vorticity: In the close setting, viscosity is sampled from [5 x 10 10-4, 10 3I. while the train- ing range is [10-3, 100-0 10. The far setting corresponds to a turbulent regime with viscosity in [10-5 110-5110-1) .100"); see Fig- ure 2 for a qualitative example. (iii) Wave 2/D: We increase both the wave celerity and damping beyond the training distribution. The wave celerity e is sampled from [500, 550] vs. [100, 500] during training, and the damping term is from [50, 60] vs. [0, 50] in training.
Σj_ A sin (wyt+ 155++2
Heat
Wave 2D
Vorticity 2D
Wide
Gaussian
Square
close
for
CAPE CODA
0.47
n/a
n/a
1.03
n/a
n/a
1.51
1.71
VIT WT-Ln--ontees
0.52
0.40
0.66
0.68
0.30
0.37
Febra
0.15
0.32
0.36
0.68
0.24
0.32
4.4. Generative ability of the model
The evaluation in Section 4.2 already shows that as gener- ative model, Zebra is less prone to error accumulation that deterministic auto-regressive models. We illustrate here ad- ditional benefits from the generative capabilities of Zebra a through two example tasks: uncertainty quantification and new trajectory generation. Further analysis of the behavior of zebra is provided in Appendix D.2 and Appendix D.3.
Setting We evaluate all the models in a one-shot setting on trajectories with out-of-distributin PDE parameters new initial conditions, making this particularly challenging test of generalization.
Uncertainty quantification Given a context example and an initial condition, Zebra a can generate multiple trajecto- ries thanks to the sampling operation at the classifier level. Statistics can then be derived from this sample of the trajec- tory distribution in order to assess for example the uncer- tainty associated to a prediction. An illustration is provided in Figure 3, the red curve represents the ground truth, the blue curve is the predicted mean and the blue shading indi- cates the empirical confidence interval (3 x standard devia- tion). Mean and standard deviation are calculated pointwise. We provide a more detailed analysis in Appendix D.2. In particular, it shows as expected that (i) uncertainty can be calibrated via the temperature parameter 7 (Figure 17), and (ii) it decreases with additional context (Table 9).
Results We report the scores in Table 2. Overall, all meth- ods experience performance degradation due to the distri- bution shift, with Zebra achieving the best results in three out of four experiments, while CODA and CAPE perform the worst. This poor performance for CAPE and CODA is expected the 2D datasets, as they already struggled to generalize within the training distribution. However, for the Heat equation, errors for CAPE and CODA double, whereas zebra maintains similar accuracy, demonstrating greater robustness to distribution shifts. Comparing Zebra and ViT-in-context to CAPE and CODA, it is remarkable that adaptation through in-context learning appears to be a more effective alternative than gradient- based adaptation for out-of-distribution generalization.
New trajectory generation As second illustration, we assess Zebra's ability to generate relevant new trajecto- ries conditioned on in-context examples alone, i.e. without prompting with an initial state query. This is similar to conditional image or text generation in LLMs. The gen-
Out-of-distribution generalization remains a challenging task, particularly under strong shifts. On the Vorticity dataset, Zebra adapts to large shifts in viscosity and pre- dicts the large-scale component of the dynamics. As shown
7

---

Zebra: In-Context Generative Pretraining for Solving Parametric PDEs
Mean prediction and confid ence Interval
Table 3. Comparison of distrihutions using the Wasserstein dis-
2.0
Aur Jin
tance between generated trajectories and real trajectories.
1.5
Distance Metric
Advection
Combined Equation
Gaussian noise vs. real data
18.22
16.15
Validation data vs. test data
5.11
1.87
0.5
zeb a-generated data vs, real data
5.57
2.21
3.0
compared to direct surrogate modeling, making inference costly. To address this, we propose a fast inference method that accelcrates inference by orders of magnitude relative to both the original model and gradient- based adaptation (Table 4). Instead of token-wise autoregressive generation, we predict entire frames once. This is achieved by re- placing the token-wise autoregressive generation process with trame-wise autoregressive surrogate, implemented as a UNet. The UNet, conditioned on a context embedding output by the transformer, takes frame w' as input and pre- diets &r A crucial component is the context embedding, which captures underlying dynamics from example trajec- tories. This is learned by introducing a new token, [DYN], at the transformer's 's input, analogous to [CLS] in BERT, allowing attention to encode context dynamics effectively. The implementation is detailed in Appendix C and in Fig- ure 7. Table shows that this reduces inference time by one to two orders of magnitude, making Zebra highly efficient.
1.0
50
100
150
200
250
Figure 3. Uncertainty quantification with Zebra in one-shot setting on Heat equation.
erated trajectories are sampled conditioned on context trajectory from new unseen environment. 'he key finding here is that Zebra effectively generates faithful trajectory distributions that closely match the real simulated trajectory distribution. Hence, given some examples from an unknown environment, Zebra could be used to generate trajecto- ries that comply with the distribution this environment. Qualitatively, Figure illustrates how the real (from a held out sample) and the generated distributions match at two different time steps (t 0, 9): the PCA projection, indicates strong alignment. Quantitatively, Table shows that the Wasserstein distance of the generated trajectories is comparable to the Wasserstein distance between valida- tion and test samples. As a calibration measure, we also provide in Table the Wasserstein distance between the real distribution and a Gaussian distribution. Further details are provided in Appendix D.3.
Table 4, Inference times for one-shot adaptation. Average time in seconds to predict a single trajectory given a context trajectory and initial condition. Times include adaptation and forecast for CODA and CAPE, while for Zebra, they include encoding. autoregressive prediction, and decoding.
Advection
Vorticity 2D
CAPE CODA
18s 31s
23s
28s
Zebra
3s
21s
7ehra + UNet
0.
0.14s
As shown in Table 5, this framework matches or outperforms pretrained Zebra in most cases. The dynamics embedding captures meaningful context, enabling efficient UNet train- ing. In contrast, methods like CAPE and CODA must learn both model weights and environment embeddings simulta- neously, making training less efficient.
L = 0.
(b) L 9.
Figure 4, PCA Visualization of generated (blue) vs. real (orange)
trajectories on Combined Equation
4.5. Accelerating inference
Table 5. Zebra vs Xebra + UNet, for the in-distribution one-shot setting. Test results in relative 1.2 error.
Since Zebra requires no gradient steps at inference, it is already faster than gradient-based adaptation (see Table 4). However, its autoregressive nature introduces significant overhead at inference: generating trajectories token by to- ken increases solver calls by one to two orders of magnitude
Advection
Combined
Wave 2D
Vorticity 2D
Zebra Zebra+UNet
0.00794 0.0072
0.00965 0.0138
0.207
0.119
0.150
0.0869
8

---

Zebra: In-Context Generative Pretraining for Solving Parametric PDEs
5. Limitations
Acknowledgements
Reconstruction quality The fidelity of the generated tra- jectories is constrained by the decoder's ability to recon- struct fine details from the quantized latent space. While the reconstruction quality is sufficient for the applications considered in this study, it may be inadequate for systems characterized by high-frequency physical phenomena. To address this limitation, scaling the codebook size (Yu et al., 2023a; Mentzer et al., 2023) or exploring alternatives to vector quantization (Li et al., 2024) could improve recon- struction accuracy-provided that the model's in-context leamning capabiliitiss are preserved.
We acknowledge the financial support provided by DLACLIM
DEEP- NUM (ANR- 21 21-CE23-0017-02), PHLUSIM (ANR-23-CE23-0025-00) and PEPR Sharp (ANR- 23-PEIA-0008, ANR, FRANCE 2030). We also thank Emmanuel de Bézenac, Jean- vves Franceschi, and Etienne Le Naour for insightful and stimu- lating discussions.
Impact Statement
This paper presents work whose goal is to advance the field of Machine Learning. There many potential societal consequences of our work, none which we feel must be specifically highlighted here.
Irregular grids Our current encoder and decoder rely on convolutional blocks, which limits the architecture to data defined on regular grids. Adopting more flexible encoding and decoding schemes-such
References
those proposed in Serrano al. (2024) could enable support for irregularly sampled inputs more complex geometric domains.
Agarwal, N., Ali, A., Bala, M., Balaji, Y., Barker, E., Cai, T., Chattopadhyay, P., Chen, Y., Cui, Y., Ding, Y,, et al. Cosmos world foundation model platform for physical ai. arXiv preprint arXiv:2501.03575 2025.
Data requirements Our data scaling analysis indicates that the framework requires a large number of training trajec- tories generalize effectively. As such, the current method is not well suited for data-scarce regimes and depends significant diversity in the parameter space. A more thor- ough investigation is needed to determine the minimal level of parameter variability required to ensure generalization in or close to the training distribution of PDE parameters.
Arakawa, A. Computational design for long-term numerical integration of the equations of fluid motion: Two-dimensional incompressible flow. part Journal of Computational Physics, I (1):119-143, 143, 1966.
ISSN 0021-9991. doi:
http://ddaaa/////000000019999991 URL science sitence/articcee//////2299966969
Blanke, M. and Lelarge, M. Interpretable meta-learning of physical systems. arXiv preprint arXiv:2312.00477, 2023.
Scalability and complexity The transformer in our chitecture applies causal attention across the full se- quence-comprising comprising the number of context examples N, time steps T, and spatial tokens per frame h x w. Extend- ing to 3D data with a third spatial dimension d increases the complexity to ONNhwwd)². This quadratic growth can lead to prohibitively long sequences in higher-dimensional settings, and may necessitate architectural changes such axial or factorized attention mechanisms.
Brandstetter, J., Welling, M., and Worrall, D. E. Lie point symmetry data augmentation for neural pde solvers. In International Conference on Machine Learning, pp. 2241- 2256. PMLR, 2022a.
Brandstetter, J., Worrall, D. E., and Welling, M. Message passing neural pde solvers. International Conference on Learning Representations, 2022b.
6. Conclusion
Brown, T. B., Mann, B., Ryder, N., Subbiah, M., Kaplan, J., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G.,
This study introduces Zebra, a novel generative model that adapts language model pretraining techniques for solving parametrie PDEs. We propose a pretraining strategy that enables Zebra to develop in-context learning capabilities. Our experiments demonstrate that the pretrained model per- forms competitively against specialized baselines across var- ious scenarios. Additionally, as a generative model, Zebra facilitates uncertainty quantification and can generate new trajectories, providing valuable flexibility in applications.
Askell, A., Agarwal, S., Herbert- Voss, A., Krueger, G., Henighan, T., Child, R., Ramesh, A., Ziegler, D. M., Wu, J., Winter, C., Hesse, C., Chen, M., Sigler, E., Litwin, M., Gray, S., Chess, B., Clark, J., Berner, C., McCandlish, S., Radford, A., Sutskever, I., and Amodei, D. Language models are few-shot learners, 2020.
Cao, Y., Liu, Y., Yang, L., Yu, R., Schaeffer, H., and Osher, S. Vicon: Vision in-context operator networks for multi-

---

Zebra: In-Context Generative Pretraining for Solving Parametric PDES
physics fluid dynamics prediction. 11 2024. URL http: //aaxiv org/ abs /211 6003.
image recognition scale. International Conference Learning Representations. 10 2021. URL http:
Chang, H., Zhang, H., Jiang, I.,, Liu, C., and Freeman, W. T. Maskgit: Masked generative image transformer. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pp. 11315-1133 2022.
Esser, P., Rombach, R., and Ommer, B. Taming transformers for high-resolution image synthesis. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 12873-12883, 2021.
Chen, R. T., Rubanova, Y., Bettencourt, J., and Duvenaud, D. K. Neural ordinary differential equations. Advances in neural information processing systems, 31, 2018.
Finn, C., Abbeel, P., and Levine, S. Model-agnosiee meta- learning for fast adaptation of deep networks. In Interna- tional conference on machine learning, pp. 1126-1135. PMLR, 2017.
Chen, W. Song, J., Ren, P., Subramanian, S., and Mahoney, M. W. Data-efficient operator learning via unsupervised pretraining and in-context learning. In NewrIPS, 2024. URL. https://gitte thub. com/ delta-lab-an aalabbaa/aee /att//aaaaaeeeeeeeeaaaaeeaaaa/ae wew__ffcfieennnggpeeeeeee
Gneiting, T. and Raftery, A. E. Strictly proper scor- ing rules, prediction, and estimation. Journal of the American Statistical Association, 102:359 - 378, 2007. URL nittps://ag semant cscholar aeg/CorpusID:1878355
Cohen, A. and Devore, R. Approximation of high- dimensional parametric pdes, Acta Numerica, 2015.
Cole, F., Lu, Y., O'Neill, R., and Zhang, T. Provable in-context learning of linear systems and linear elliptic pdes with transformers. arXiv preprint 24Xiv:2409.12299 2024.
Gupta, J. K. and Brandstetter, J. Towards multi- spatiotemporal-scat generalized pde modeling. Transac- tions on Machine Learning (TMLR) 07/2023, 2022.
Hao, Z., Su, C., Liu, S., Berner, J., Ying. C.,, Su, H., Anand- kumar, A., Song, J., and Zhu, J. Dpot: Auto-regressive denoising operator transformer for large-scale pde pre- training. 41th International Conference on Machine Learning (ICML 2024), 2024.
Couairon, G., Singh, R., Charantonis, A.. Lessig, C., and Monteleoni, C. Archesweather & archesweathergen: de- terministic and generative model for efficient ml weather forecasting. 12 2024. URL httt:://arxiv.or. abs/241 129771.
Hendrycks, D., Mazeika, M., and Dietterich, T. Deep anomaly detection with outlier exposure. arXiv preprint 1812 2018
Cuturi, M. Sinkhorn distances: Lightspeed computation of optimal transport. Advances in neural information processing systems, 26, 2013.
Dai, D., Sun, Y., Dong, L., Hao, Y., Ma, S., Sui, Z., and Wei, F. Why can gpt learn in-context? language models im- plicitly perform gradient descent as meta- meta-optimizerr. In Proceedings of the Annual Meeting of the Association for Computational Linguistics, pp. 4005-4019, 2023. ISBN 9781959429623. doi: 10. 10.18653//1/20223.nndll 247.
Herde, M., Raonić, B., Rohner, T., Käppeli, R., Molinaro, R., de Bézenac, E., and Mishra, S. Poseidon: Efficient foundation models for pdes. 2024.
Ho, J., Kalchbrenner, N., Weissenborn, D., and Salimans, Axial attention in multidimensional transformers. arXiv preprint
2019.
Dong, Q., Li, L., Dai, D., Zheng, C., Ma, J., Li, R., Xia, H., Xu, J., Wu, Z., Chang, B., Sun, X., Li, L., and Sui, Z. A survey on in-context learning. In Proceed- ings of the 2024 Conference Empirical Methods in Natural Language Processing, pp. 1107-1128, Miami, Florida, USA, November 2024. Association for Compu- tational Linguistics. doi: 10. 10.18665//1122222cccmmmm..m
Ho, J., Salimans, T., Gritsenko, A., Chan. W., Norouzi, M., and Fleet, D. J. Video diffusion models. Advances in Neural Information Processing Systems, 35:8633-8646 2022.
Hu, E. J., Shen, Y., Wallis, P.. Allen-Zhu, Z., Li, Y.. Wang, S., Wang, I.., and Chen, W. Lora: Low-rank adaptation of large language models. arXiv preprint 2021.
64. URL https: / aa http///lannhholoo.orr////000 emnlp main. 64/.
Dosovitskiy, A., Beyer, L., Kolesnikov, A., Weissenborn, D., Zhai, X., Unterthiner. T., Dehghani, M., Minderer, M.. Heigold, G., Gelly. S., Uszkoreit, J., and Houlsby,
Hu, P., Wang, R., Zheng, X., Zhang, T., Feng, H., Feng, R., Wei, L., Wang, Y., Ma, Z.-.., and Wu, T. Wavelet diffu- sion neural operator. arXiv preprint 24l::24122048333 2024.
N. image worth 16x 16 words: Transformers for
10

---

Zebra: In-Context Generative Pretraining for Solving Parametric PDEs
Huang, G., Sun, Y., Liu, Z., Sedra, D., and Weinberger, K. Q- Deep networks with stochastic depth. In Computer Vision- ECCV 2016: 14th European Conference, Amsterdam, The Netherlands, October 11 11-14, 2016. Proceedings, Part IV 14, pp. 646-661. Springer, 2016.
Mentzer, F,, Minnen, D., Agustsson, E., and Tschannen, M. Finite scalar quantization: Vq-vae made simple. Interna- tional Conference Learning Representations, 2023.
Nzoyem, D., Barton, A., and Deakin, Neural context flows for meta-leaming of dynamical systems. arXiv preprint arXiv:2405.00 154, 2024.
Huang, J., Yang, G., Wang. Z., and Park, J. Diffusionpde: Generative pde-solving under partial observation. arXiv preprint ar arXiv:2406.177 17763, 2024.
Oord, A. v. d., Vinyals, O., and Kavukcuoglu, K. Neu- ral discrete representation learning. arXiv preprint 2017.
Jiaqi, L., Cui, J., Yang. J., and Yang, B. Stochastic neural simulator for generalizing dynamical systems across environments. pp. 5909-5917, 08 2024. doi: 00.24963/jjeai.2224/555
Park, J., Berto, F., Jamgochian, A., Kochenderfer, M., and Park, J. First-order context-basse adaptation for gener- alizing to new dynamical systems. 2023. URL https: / oopnrevieew net/ forum? 2id= i 101 OhzqJ.
Kirchmeyer, M., Yin, Y., Donå, J., Baskiotis, N., Rako- tomamonjy, A., and Gallinari, P. Generalizing new physical systems via context-informed dynamiis model. In International Conference on Machine Learning, pp. 11283-11300 13301 PMLR, 2022.
Peebles, W. and Xie, S. Scalable diffusion models with transformers. In Proceedings of the IEEE/CVE Interna- tional Conference on Computer Vision, pp. -195-4205, 2023.
Koupai, A. K., Benet, J. M., Yin, Y., Vittaut, J.-N., and Gal- linari, P. Geps: Boosting generalization in parametric pde neural solvers through adaptive conditioning. Advances in Neural Information Processing Systems, 38, 2024.
Price, I., Sanchez-Gonzalez, A., Alet, F., Ewalds, T., El- Kadi, A., Stott, J., Mohamed, S., Battaglia, P., Lam, R., and Willson, M. Gencast: Diffusion-
ensemble forecasting for medium-range weather. pp. 1-46, 2023. URL http: /ttp:://aaxiiii org/ abs////b//aa 200os/2311.2.2 15796.
Li, T., Tian, Y., Li, H., Deng, M., and He, K. Autoregres- sive image generation without vector quantization. arXiv preprint arXiv:2406.11 18838, 2024.
Radford, A., Narasimhan, K., Salimans, T., and Sutskever, I. Improving language understanding by generative pre- training. 2018
Li, Z., Kovachki, N., Azizzadenesheli, K., Liu, B., Bhat- tacharya, K., Stuart, A., and Anandkumar, A. Fourier neural operator for parametric partial differential equa- tions. arXiv preprint urXiv:2010.08895. 2020.
Raissi, M., Perdikaris, P., and Karniadakis, G. E. Physics- informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations. Journal of Computational physics, 378:686-700 7076 2019.
Li, Z., Kovachki, N., Azizzadenesheli, K., Liu, B., Bhat- tacharya, K., Stuart, A., and Anandkumar, A. Fourier neural operator for parametric partial differential equa- tions. International Conference on Learning Represen- tations., 10 2021. URL http: /ttp:///axxiii aaxivv org/ abs/ 2010. 08895.
Ronneberger, O., Fischer, P., and Brox, T. U-net: Convolu- tional networks for biomedical image segmentation. vol- ume 9351, 2015. doi: 10.10077978-33311-247777 2457--4 28.
Serrano, L., Wang, T. X., Naour, E. L., Vittaut, J.-N., and Gallinari, P. Aroma: Preserving spatial structure for latent pde modeling with local neural fields. arXiv preprint arXiv:2406.0m 2.XY:22006.02176, 2024.
Lino, M., Pfaff, T., and Thuerey, N. Learning distribu- tions of complex fluid simulations with diffusion graph networks. arXiv preprint 2088022504.02843, 2025.
Subramanian, S., Harrington, P., Keutzer, K., Bhimji, W., Morozov, D., Mahoney, M. W., and Gholami, A. To- wards foundation models for scientific machine learning: Characterizing scaling and transfer behavior. Advances Neural Information Processing Systems 37 (NeurIPS 2023), 2023.
Lippe, P., Veeling, B. S., Perdikaris, P., Turner, R. E., and Brandstetter, Pde-refiner: Achieving accurate long roll- outs with neural pde solvers. In Neurips, pp. 1-36, 2023. URL http: / /awviv.ivvvv.////////...//aaavvv org/ abs 05732.
McCabe, M., Blancard, B. R.-S., Parker, L. H., Ohana, R., Cranmer, M., Bietti, A., Eickenberg, M., Golkar, S., Krawezik, G., Lanusse, F., et al. Multiple physics pre- training for physical surrogate models. arXiv preprint OrXiv:2310.000 10.02994, 2023.
Takamoto, M., Alesiani. F., and Niepert, M. Learning neu- ral pde solvers with crrameter-guider channel attention. International Conference on Machine Learning (ICML), 2023.

---

Zebra: In-Context Generative Pretraining for Solving Parametric PDEs
Tian, K., Jiang, Y., Yuan, Z., Peng, B., and Wang, L. Visual autoregressive modeling: Scalable image generation via next-scale prediction. arXiv preprint arXiv:2404.02905. 2024.
Touvron, H., Lavril, T., Izacard, G., Martinet, X., Lachaux, M.-A., Lacroix, T., Rozière, B., Goyal, N., Hambro, E., Azhar, F., Rodriguez, A., Joulin, A., Grave, E., and Lam- ple, G. Llama: Open and efficient foundation language models. CoRR, 2023.
Wolf, T. Huggingface's transformers: State-of-the- art natural language processing. arXiv preprint arXiv: 2019.
Xie, S. M., Raghunathan, A., Liang, P., and Ma, An explanation of in-context learning as implicit bayesian inference. In The Tenth International Conference Learning Representations, ICLR 2022, Virtual Event, April 25-29, 2022. OpenReview.net, 2022. URL ht htsps: //openreview net idddnndddddRRCCC-- JUMI.
Yang, L., Liu, S., Meng, T., and Osher, S. In-context operator learning with data prompts for differential equa- tion problems. Proceedings of the National Academy of Sciences, 120139):23100122100, 2023.
Yin, Y., Ayed, I., de Bézenac, E., Baskiotis, N., and Galli- nari, P. Leads: Learning dynamical systems that general- ize across environments. Neural Information Processing Systems, 2022.
Yu, L.,, Cheng, Y., Sohn, K., Lezama, J., Zhang, H., Chang, H., Hauptmann, A. G., Yang, M.-H., Hao, Y., Essa, I., et al. Magvit: Masked generative video transformer. In Proceedings of the IEFF/CVF Conference on Computer Vision and Pattern Recognition, pp. 10459, 2023a.
Yu, L., Lezama, J., Gundavarapu, N. B., Versari, L., Sohn, K., Minnen, D., Cheng, Y., Birodkar, V., Gupta, A., Gu, X., et al. Language model beats diffussion-toenizzrrii is key to visual generation. arXiv preprint ar0iv.2310.05737, 2023b.
Zintgraf, L., Shiarli, K., Kurin, V., Hofmann, K., and White- son, S. Fast context adaptation via meta-learning. In International Conference on Machine Learning, pp. p 7693- 7702. PMLR, 2019a.
Zintgraf, I.., Shiarlis, K., Kurin, V., Hofmann, K., and Whiteson, S. Fast context adaptation via meta-learning. 36th International Conference on Machine Learning. ICML 2019, 2019-June, 2019b
12

---

Zebra: In-Context Generative Pretraining for Solving Parametric PDEs
A. Related Work
A.1. Learning parametric PDEs
The classical ML paradigm The classical ML paradigm for solving parametric PDEs consists in sampling from the PDE parameter distribution trajectories to generalize to new PDE parameter values. It is the classical ERM approach. The natural way for generalizing to new PDE parameters is to explicitly embed them in the neural network (Brandstetter al., 2022b). (Takamoto et al., 2023) proposed a channel- mechanism guide neural solvers with the physical coefficients given as input; it requires complete knowledge of the physical system and are not designed for other PDE parameter values, e.g., boundary conditions. It is commonly assumed that prior knowledge are not available, but instead rely on past states of trajectories for inferring the dynamics. Neural solvers and operators learn parametric PDEs by stacking the past states as channel information as done in (Li et al., 2021), or by creating additional temporal dimension as done in video prediction contexts (Ho et al., 2022; McCabe et al., 2023). Their performance drops when shifts occur in the data distribution, which is often met with parametrie PDEs, as small changes in the PDE parameters can lead to various dynamics. To better generalize to new PDE parameter values, (Subramanian et al., 2023) instead leverages fine-tuning from pretrained models generalize to new PDE parameters. It however often necessitates a relatively large number of fine tuning samples to effectively adapt to new PDE parameter values, by updating all or a subset of parameters (Herde et al., 2024; Hao et al., 2024).
Gradient-based adaptation To better adapt to new PDE parameters values at inference, several works have explored learning on multiple environments. During training, a limited number of environments are available, each corresponding to a specific PDE instance. (Yin et al., 2022) introduced LEADS, a multi-task framework for learning parametric PDEs, where a shared model from all environments and a model specific to each environment are learned jointly. At inference, for a new PDE instance, the shared model remain frozen and only a model specific to that environment is learned. (Kirchmeyer et al., 2022) proposed to perform adaptive conditioning in the parameter space; the framework adapts the weights of model to each environment via a hyper-network conditioned by context vector cf specific to each environment. At inference, the model adapts to new environment by only tuning ca. (Park et al., 2023) bridged the gap from the classical gradient-based meta-learning approaches by addressing the limitations of second-order optimization of MAML and its variants (Finn et al., 2017; Zintgraf et al., 2019b). Other works have also extended these frameworks to quantify uncertainty of the predictions: (Jiaqi et al., 2024) proposed a conditional neural process to capture uncertainty in the context of multiple environments with sparse trajectories, while (Nzoyem et al., 2024) leveraged information from multiple environments to enable more robust predictions and uncertainty quantification.
In-context learning for PDE Inspired by the in-context learning (ICL) paradigm in large language models (LLMs), recent works have explored adapting this approach for solving PDEs and modeling dynamical systems. One of the earliest efforts in this direction is Yang et al. (2023), which aims learn operators capable of adapting to different physical scenarios by leveraging in-context examples. Their approach utilizes an encoder-decoder transformer, where the transformer encodes the context prompt. This prompt, together with a query, is then passed to the decoder, which predicts the corresponding output values of the state vector. However, since functions are represented as scattered point tokens, the model encounters computational complexity challenges and is primarily limited to ID ODEs or sparse 2D data. VICON (Cao et al., 2024) extends this framework by leveraging vision transformers (ViTs) that operate on image patches. However, the two approaches further differ in their problem settings: while our method focuses on adapting a surrogate model using a small number of context trajectories and then unrolling from an initial condition,as traditional PDE solvers do, VICON is designed for future-state forecasting based on access to past trajectory history. Our ViT In-Context baseline already captures an architecture that is structurally similar to VICON, but tailored to our adaptation formulation, thereby offering an implicit point of comparison.
Chen et al. (2024) takes a different approach, focusing on unsupervised pretraining for operator learning. Their method involves pretraining the encoder-decoder of neural operators on proxy tasks (such as masked prediction or super-resolution) that require only snapshots of the dynamics rather than full simulation data, followed by fine-tuning on target dynamics. In this case, ICL is used only at inference time, where context examples similar to query input are retrieved from the training set, and their solutions are aggregated and averaged to form the final prediction. This setting differs significantly from ours. Notably, all these approaches rely on deterministic ViT-like architectures, whereas our method employs generative stochastic model.
Note that in-context learning is still a not well understood phenomenon and that different hypotheses are being explored which attempt to fill this gap (Dong et al., 2024). Two prevalent explanations come from Bayesian perspective on ICL as
13

---

Zebra: In-Context Generative Pretraining for Solving Parametric PDEs
introduced in popular paper (Xie et al., 2022) and the gradient descent view as introduced e-g. in Dai et al. (2023) that identied dual form between transformer attention and gradient descent highlighting relations between GPTT-ased ICL and expicit fine tuning.
A.2. Generative models
Auto-regressive Transformers for Images and Videos Recent works have explored combining language modeling techniques with image and video generation, typically using VAE (Oord et al., 2017) paired with causal transformer (Esser et al., 2021) or bidirectional transformer (Chang et al., 2022). VQGAN (Esser et al., 2021) has become the leading framework by incorporating pereeptual and adversarial losses to improve the visual realism of decoder outputs from quantized latent representations. However, while these methods succeed in generating visually plausible images, they introduce bias-driven by perceptual and adversarial losses-thar that leads the network to prioritize perceptual similarity and realism, often causing reconstructions to deviate from the true input. In contrast, Zebra focuses on maximizing reconstruction accuracy, and we did not observe benefits from using adversarial or perceptual losses during training.
In video generation, models like Magvit (Yu et al., 2023a) and Magvit2 (Yu et al., 2023b) adopt similar strategies, using 3D CNN encoders to compress sequences of video frames into spatiotemporal latent representations by exploiting the structural similarities between successive frames in a video. However, such temporal compression is unsuitable for modeling partial differential equations (PDEs), where temporal dynamics can vary significantly between frames depending on the temporal resolution. With Zebra, we spatially compress observations using the encoder and learn the temporal dynamics with an auto-regressive transformer, avoiding temporal compression.
Generative PDE surrogate models Recent advances have leveraged diffusion-based generative models for surrogate modeling of PDEs. Huang et al. (2024) introduce DiffusionPDE, which employs diffusion model to complete partial observations and solve PDEs under limited supervision. Hu et al. (2024) propose the Wavelet Diffusion Neural Operator (WDNO), which performs diffusion in the wavelet domain and uses multi-resolltion training to capture sharp features and generalize across spatial scales, with a focus on long-term simulation and control. Lino et al. (2025) develop a diffusion graph network that learns full solution distributions of complex fluids; enables direct sampling of equilibrium states on unstructured meshes, facilitating efficient estimation of turbulent statistics. Finally, Hao al. (2024) present DPOT, an autoregressive transformer pre-trained with noise- peitteperturbe inputs large collection of PDE simulations.
B. Dataset details
Table 6. Dataset Summary
Dataset Name
Number of env.
Trajectorics per env.
Main parameters
Advection
1200
10
Advection speed Diffusion and forcing Diffusion and forcing Boundary conditions
Heat
1200
10
Burgers
1200
10
Wave boundary
4
3000
Combined equation
1200
10
a, 3,y
Wave 2D
1200
10
Wave celerity and damping Diffusion
Vorticity 2D
1200
10
B.1. Advection
We consider a ID advection equation with advection speed parameter 3::
0.u + 30,u 0
For each environment, we sample 3 with a uniform distribution in [0, 4]. We sample 1200 parameters, and 10 trajectories per parameter, constituting a training set of 12000 trajectories. At test time, we draw 12 new parameters and evaluate the performance on 10 trajectories each.
We fix the size of the domain I. 128 and draw initial conditions as described in Equation (3) in appendix B.5 and generate
14

---

Zebra: In-Context Generative Pretraining for Solving Parametric PDEs
solutions with the method of lines and the pseudo-spectral solver described in (Brandstetter et al., 2022b). We take 140 snapshots along 100s long simulations, which we downsample 14 timestamps for training. We used spatial resolution of 256.
B.2. Burgers
We consider the Burgers equation special case of the combined equation described in Appendix B.5 and initially in (Brandstetter et al., 2022b), with fixed Y 0 and 0.5. However, in this setting, we include a forcing term 5(t,z) = مj= Aj sin(w,t+ 2 L 0;) that can vary across different environments. We fix J = 5, L = 16. We draw initial conditions as described in Equation (3).
For each environment, we sample 3 with log-uniform distribution in [Le - 3, 5). and sample the forcing term coefficients uniformly: A, E [-0.5, 0.5], W/ E [-0.4, -0.4], é E (1,2.3), oj E [0, 2n]. We create a dataset of 1200 environments with 10 trajectories for training, and 12 environments with 10 trajectories for testing.
We use the solver from (Brandstetter et al., 2022b), and take 250 snapshots along the 4s of the generations. We employ spatial resolution of 256 and downsample the temporal resolution to 25 frames.
B.3. Heat
We consider the heat equation special case of the combined equation described in Appendix B.5 and initially in (Brandstetter et al., 2022b), with fixed y = 0 and = O. However, in this setting, we include forcing term 8(t, x) Σÿ=1 Aj sin(w,t + 2 [ + Φb)) that can vary across different environments. We fix J 5, L 16. We draw initial conditions as described in Equation (3).
For each environment, we sample 8 with uniform distribution in [le - 3, 5]. and sample the forcing term coefficients uniformly: A E [-0.5, 0.5), Wj E [-0.4, -0.4), l E (1,2,3). 0j E [0, 2T].. We create a dataset of 1200 environments with 10 trajectories for training, and 12 environments with 10 trajectories for testing.
We use the solver from (Brandstetter et al., 2022b), and take 250 snapshots along the 4s of the generations. We employ a spatial resolution of 256 and downsample the temporal resolution to 25 frames.
B.4. Wave boundary
We consider a ID wave equation as in (Brandstetter et al., 2022b).
Onu c²O = 0, E 8,
where is the wave velocity (e = our experiments). We consider Dirichlet B[u] = u = and Neumann B[u] = - = boundary conditions.
We consider 4 different environments each boundary can either respect Neumann Dirichlet conditions, and sample 3000 trajectories for each environment. This results in 12000 trajectories for training. For the test set, sample 30 new trajectories from these 4 environments resulting in 120 test trajectories.
The initial condition is a Gaussian pulse with a peak at a random location. Numerical ground truth is generated with the solver proposed in (Brandstetter et al., 2022b). We obtain ground truth trajectories with resolution (nzz, n2) = (256, 250), and downsample the temporal resolution to obtain trajectories of shape (256, 60).
B.5. Combined equation
We used the setting introduced in (Brandstetter et al., 2022b), but with the exception that we do not include a forcing term. The combined equation thus described by the following PDE:
10 2 30 + 70 (1,, 8(t,2).
(2)
J
=0,
uo (x(x)
Σ
A, (inn t,)-
(3)
j=1
15

---

Zebra: In-Context Generative Pretraining for Solving Parametric PDEs
For training, we sampled 1200 triplets of parameters uniformly within the ranges [0, 1]. 8 [0, 0.4), and 7 E [0, 1]. For each parameter instance, we sample 10 trajectories, resulting in 12000 trajectories for training and 120 trajectories for testing. We used the solver proposed in (Brandstetter al., 2022a) to generate the solutions. The trajectories were generated with spatial resolution of 256 for 10 seconds, along which 140 snapshots are taken. We downsample the temporal resolution to obtain trajectories with shape (256, 14).
B.6. Vorticity
We propose a 2D turbulence equation. We focus analyzing the dynamics of the vorticity variable, The vorticity, denoted by w. is a vector field that characterizes the local rotation of fluid elements, defined as = V x u. The vorticity equation is expressed as:
dw a
(u VWw - ~ =0
(4)
Here, u represents the fluid velocity field, // is the kinematic viscosity with = 1/ Re. For the vorticity equation, the parametric problem consists in learning dynamical systems with changes in the viscosity term.
For training, sampled 1200 PDE parameter values in the range [le 3, le 2). For test, evaluate model on 120 new parameters not seen during training in the same paramter range. For each parameter instance, sample 10 trajectory, resulting in 12000 trajectories for training and 1200 for test.
Data generation For the data generation, we use a 5 point stencil for the classical central difference scheme of the Laplacian operator. For the Jacobian, we use a second order accurate scheme proposed by Arakawa that preserves the energy, enstrophy and skew symmetry (Arakawa, 1966). Finally for solving the Poisson equation, we use Fast Fourier Transform based solver. We discretize periodic domain into 512 512 points for the DNS and uses RK4 solver with At le on temporal horizon [0, 2]. We then perform a temporal and spatial down-sample operation, thus obtaining trajectories composed of 10 states 64 x 64 grid.
We consider the following initial conditions:
k
k Roo
E(k)
V
exp
(5)
Kn
ko
Vorticity is linked to energy by the following equation:
E(k)
(k:
(6)
B.7. Wave 2D
We propose a damped wave equation, defined by
Ow Ot
e² Aw +
(7)
0t
where e is the wave speed and k is the damping coefficient. We are only interested in learning To tackle the parametric problem, we sample 1200 parameters in the range = [0, 50) and k = [100, 500]. For validation, we evaluate our model on 120 new parameters not seen during training in the same paramter range. For each parameter instance, we sample 10 trajectory, resulting 12000 trajectories for training and 1200 for validation.
Data generation For the data generation, we consider a compact spatial domain Ω represented 64 64 grid and discretize the Laplacian operator similarly. A is implemented using 5 discrete Laplace operator in simulation. For boundary conditions, null neumann boundary conditions imposed. We set At 6.25c and generate trajectories the temporal horizon [0, 5e 3). The simulation was integrated using fourth order runge-kutta kutta schema from initial condition corresponding to sum of gaussians:
(x x.,) + (y - y:)
wo(x, v) = 5 exp
(8)
20
1m
16

---

Zebra: In-Context Generative Pretraining for Solving Parametric PDEs
where we choose p 5 gaussians with o, 0.0.025.0011 4y Mo.1). ys ~ M0. We fixed C' to here. Thus, all initial conditions correspond to sum of gaussians of varying amplitudes.
C. Architecture details
C.1. Baseline implementations
For all baselines, we followed the recommendations given by the authors. We report here the architectures used for each baseline:
CODA: For CODA, we implemented U-Net (Ronneberger al., 2015) and a FNO (Li al., 2020) as the neural network decoder. For all the different experiments, reported in the results the best score among the two backbones used. We trained the different models in the same manner Zebra, i.e. via teacher forcing (Radford et al., 2018). The model is adapted to each environment using a context vector specific to each environment. For the size of the context vector, we followed the authors recommendation and chose a context size equals to the number of degrees of freedom used to define each environment for each dataset. At inference, we adapt to a new environment using 250 gradient steps.
CAPE: For CAPE (Takamoto et al., 2023), we adapted the method to an adaptation setting. Instead of giving true physical coefficients as input, we learn to auto-decode a context vector e as in CODA, which implicitly embeds the specific characteristics of each environment. During inference, we only adapt c with 250 gradient steps. For the architectures, UNET and FNO as the backbones, and reported the best results among the two architectures for all settings.
[CLS] VIT: For the ViT, meta-learning setting where the CLS token encodes the specific variations of each environment, At inference, the CLS token is adapted to new environment with 100 gradient steps.
simple vision transformer architecture (Dosovitskiy al., 2021), but adapt it to
Wi-in-context: We implement ViT-in-context using standard transformer architecture with separate temporal and spatial attention mechanisms, following Ho et al. (2019). During both training and inference, context examples are stacked along the temporal dimension. The model is trained to predict the next frame in the target sequence, conditioned on both the context examples and the preceding frames of the target sequence.
C.2. Zebra additional details
Zebra We describe the pretraining strategy in Section 3, and provide details on the architecture and its hyperparameters in Appendix C. The datasets used are described in Appendix B. We plan release the code, the weights of the models, and the datasets used in this study upon acceptance.
For clarity, we outline the pretraining steps of Zebra in Figure 5 illustrated with the vorticity 2D dataset.
We also provide illustrations of our inference pipeline in Figure We finally include a schematic view of the different generation possibilities with Zebra in Figure 8, using the sequence design adopted during pretraining.
Zebra + UNet Zebra is competitive both in-distribution and out-of- distribution, while also enabling uncertainty quantifi- cation due to its generative nature. Additionally, Zebra can generate novel trajectories and initial conditions, providing a way to sample complex initial states. As such Zebra is already faster than gradient- based adaptation methods, but since the model generates trajectory solutions token by token, the number of calls to the transformer increases by one or two orders of magnitude compared direct surrogate modeling, making the process costly.
In this section, we propose hybrid approach that leverages Zebra's pretrained knowledge in combination with conventional neural surrogate model. The objective is to develop framework that encodes context trajectories into an embedding vector, which then conditions a neural surrogate, similar to classical adaptation methods such as CODA and CAPE.
To achieve this, we finetune Zebra as an encoder to adapt a conditional surrogate model, such as a UNet (Ronneberger al., 2015). We introduce [DYN] token (short for dynamics), which is appended the right of the context sequence
17

---

Zebra: In-Conter, Generative Pretraining for Solving Parametric PDEs
VQVAE
Input: physical field
Output: reconstruction
codes
-
quantized codes
Codebook Z
quantization
Encoder CNN
Decoder CNN
min
22
2
In-context Pretraining
Trajectories sharing same dynamics
Trajectory 1
Trajectory 2
Output: Next token probabillies
p(5²) p(Sa) p(S. S
p(Ss)
v(Sx).
Transtormer
8
Sx-1 SN
Input: Sequence of indices
Time
Time
Special tokens
begilt Cevournce
Tokenize
Tokenize
dboss
<tots
<oots <bott
<bots <0003
Figure 5. Zebra's pretraining includes two stages. 1) A finite vocabulary of physical phenomena is learned by training VQ-VAE VAE on spatial representations. 2) During the pretraining, multiple trajectories sharing the same dynamics are tokenized and concatenated into common sequence S. The transformer is trained to predict the next-token by minimizing the cross-entropy loss.
18

---

Zebra: In-Context Generative Pretraining for Solving Parametric PDEs
Generated Trajectory
Context trajectory
Initial condition
DeTokenize
Time
8.
1am
genersted tokens
Time
Transformer
Tokenize
Tokenize
Auto-regressive géneration
81
dboss
dbob
81
1000 cbots
0
Figure Zebra's inference pipeline from context trajectory. The context trajectory and initial conditions are tokenized into index sequences that are concatenated according to the sequence design adopted during pretraining, The transformer then generates the next tokens to complete the sequence. detokenize these indices to get back the physical space.
during both training and inference. This allows Zebra to extract a dynamics embedding from the transformer's output, defined as:
Es Transformer(S,, [DYN])
where S represents a sequence tokens encoding the context trajectories. The embedding &s captures key properties the dynamics and is mapped to the conditioning space of UNet, which adapts the model each specific dynamics. Following Gupta & Brandstetter (2022), the UNet conditioning modifies the network's biases. Once the dynamics embedding is extracted, the UNet can directly predict the next state from the current state
U ++ UNet(u',
With this architecture, we effectively extract the key dynamics from example trajectories and use this information to adapt neural surrogate, significantly accelerating inference. The complete pipeline at inference is illustrated in Figure 7.
To efficiently finetune Zebra. we apply LoRA (Hu et al., 2021), keeping the transformer's weights frozen while learning low-rank updates. This setup enables the UNet to leverage Zebra' 's pre-learned representations while achieving substantial speedup. Compared to standalone Zebra, integrating UNet improves inference speed by factor of x30 in ID and 150 in 2D. The method is also considerably faster than CODA and CAPE while maintaining competitive performance across tasks. The inference times are summarized in Table 4.
Finally we illustrate the inference pipeline for accelerating the inference of Zebra with the UNet in Figure 7.
C.3. Auto-regressive transformer
Zebra's transformer is based on Llama's architecture, which we describe informally in Figure 9. We use the implementation provided by HuggingFace (Wolf, 2019) and the hyperparameters from Table 7 in our experiments. For training the transformer, we used a single NVIDIA TITAN RTX for the ID experiments and used single A 100 for training the model on the 2D datasets. Training the transformer on 2D datasets took 20h on a single A.00 and took 15h on a single RTX for the 1d dataset.
C.4. VQVAE
The quantizer used at the token level is a VQVAE model (Oord et al., 2017). As illustrated in Figure 10 this is an encoder-decoder architecture with an intermediate quantizer component.
19

---

Zebra: In-Context Generative Pretraining for Solving Parametric PDEs
Example Trajectory 1
Example Trajectory
UNet
dynamics embedding
Tokenize
Tokenize
Transformer
81
$2
sboss
<bot>
[DYN]
Figure 7. Zebra + UNet inference pipeline. Zebra serves as an encoder, utilizing the special token DYN to generate dynamics embedding from the context example trajectories. Once this emhedding obtained, the UNet autoregressively forecast the sequence significantly faster than a next-token transformer.
Table Hyperparameters for Zebra's Transformer
Hyperparameters
Advection
Heat
Burgers
Wave b Combined
Vorticity 2D
Wave 2D
sax
2048
2048
2048
2048
2048
8192
8192
batch_size
4
2
2
num num-gradien wgw.ggrdddnnntctnnen accuullltts
1
4
hidden_size
256
256
256
256
256
384
512
mlp ratio
4.0
4.0
4.0
4.0
4.0
4.0
4.0
depth
8
8
8
8
8
8
num heads
8
8
8
8
8
8
vocabulary -sze
264
264
264
264
264
2056
2056
start learning rate
le-4
le-4
le-4
le-4
le-4
le-4
le-4
beight_deeaay
le-4
le-4
le-4
lc44
le-4
le-4
le-4
scheduler
Cosine
Cosine
Cosine
Cosine
Cosine
Cosine
Cosine
num num_epoct epochs
100
100
100
100
100
30
30
20

---

Zebra: In-Context Generative Pretraining for Solving Parametric PDEs
a)
Conditional generation from similar trajectories
Auto-regressive generation
Transformer
<bot>
8
Dum 1
<eot>
<bot>
Ocm sis
<eot>
<bob>
generated prediction
Prompt examples
initial condition
b) Unconditional generation from similar trajectories
Auto-regressive generation
Transformer
0:m $ i
Orm $ is
<bos>
<bot>
<B0t>
<bot>
<eot>
<bot>
generation
Prompt examples
c) Unconditional generation
Auto-regressive generation
Transformer
<bos>>
<bot>
Orm
generation
<bos>>
begin of sequence
D:m $1
context trajectory
<bot>
begin of new trajectory
initial condition of target trajectory
<eot>
end of trajectory
Pigure Generation possibilities with Zebra.
21

---

Zebra: In-Context Generative Pretraining for Solving Parametric PDEs
Output: Next-token probabilties
Llama block
Softmax
LayerNorm
A
A
Linear
Multi-Head Attention
A
A
A
Llama block x 8
LayerNorm
Embedding layer
FFN
<bos>>
<bot>
$11
<eot>
<bot>
15
<eat>
<bob>
ig
<Bot>
<8OS>
input: Sequence of indices
<boss>
begin of sequence
Ouj Ai
Context trajectory
<bot>
begin of new trajectory
Initial condition target trajectory
<eot>
end of trajectory
Figure 9, Zebra's transformer architecture is based on Llama (Touvron et al., 2023).
22

---

Zebra: In-Context Generative Pretraining for Solving Parametric PDES
The encoder spatially compresses the input function et' by reducing its spatial resolution H W lower resolution while increasing the channel dimension d. This is achieved through convolutional model &w, which maps the input to continuous latent variable z' = Ew (u where z E Rixxxxd The latent variables are then quantized to discrete codes z using a codebook 2 of size K = |2| and through the quantization step q. For each spatial code [gg) the nearest codebook entry Z& is selected:
2 4.1bl = ( = arg min ||]])) - xx| F
The decoder De: reconstructs the signal ú' from the quantized latent codes z Both models jointly trained to minimize the reconstruction error between the function u' and its reconstruction u' v o (w').. The codebook 2 is updated using an exponential moving average (EMA) strategy, which stabilizes training and ensures high codebook occupancy. The training objective is:
|/ee' ú' 1|2
Lvo
| al|ss(z]]]
& (w)
where the first term is the Relative L2 loss commonly used in PDE modeling, and the second term is the commitment loss, ensuring encoder outputs are close to the codebook entries. The parameter o, set to 0.25, balances the two components. Here, sg denotes the hoopggradient operation that detaches a tensor from the computational graph.
We provide schematic view of the VQVAE framework in Figure 10 and detail the architectures used for the encoder and decoder on the ID and 2D datasets respectively Figure 11 and Figure 12. As detailed, we use residual blocks to process latent representations, and downsampling and upsampling block for decreasing increasing the spatial resolutions. We provide the full details of the hyperparameters used during the experiments in Table 8. For training the VQVAE, we used a single NVIDIA TITAN RTX for the ID experiments and used single V 1000 for training the model on the 2D datasets. Training the incoder-decode on 2D datasets took 20h on a single V100 and it took 4h on a single RTX for ID dataset.
Input: physical field
Output: reconstruction
codes
quantized codes
quantization
Encoder CNN
Decoder CNN
arg min ||= m/|l Aye
2
233
Codebook
Figure 10. Rebra's VQVAE used to obtain compressed and discretized latent representation. By retrieving the codebok index for each discrete representation, obtain discrete tokens encoding physical observations that can be mapped back to the physical space with high fidelity.
23

---

Zebra: In-Contex Generative Pretraining for Solving Parametric PDEs
Encoder
Decoder
GroupNorm + SiLU +
Conv, C->64
Conv, 64 ->
ResBlock 64
ResBlock 64
x4
DownBlock
x4
UpBlock
GroupNorm + SiLU Conv, 256 64
Conv, 64->256
ResBlock
VQ, bodebook_size-255
GroupNorm + SiLU Conv c in, o_out
DownBlock
UpBlock
Conv, stride=2 c, h -> C, h/2
Depth to space:
GroupNorm + SiLU + Conv c_out, c_out
h -> h-2
ResBlock -> 2'c
ResBlock c > c/2
Figure 11. Architecture of Zebra's VQVAE for ID datasets. Each convolution acts only on the spatial dimension and uses a kernel of size 3. The Residual Blocks are used process information and increase or decrease the channel dimensions, while the Up and Down blocks respectively up-sample and down-sample the resolution of the inputs. In ID, we used a spatial compression factor of 16 on all datasets. Every downsampling results in doubling of the number of channels, and likewise, every upsampling is followed by reduction of the number of channels by 2. We choose a maximum number of channels of 256.
24

---

Zebra: In-Contex Generative Pretraining for Solving Parametric PDEs
Encoder
Decoder
GroupNorm +SLUU+
Conv, C->128
Conv, 128 c
x2
ResBlock 128
x2
ResBlock 128
x2
DownBlock
x2
UpBlock
GroupNorm + SiLU
Conv, 16->512
Conv, 512 - 16
ResBlock
VQ, bodebook_size-2000
GroupNorm + SiLU Conv c_in, c_out
DownBlock
UpBlock
Conv, stride=2 c, h - c, h/2
Depth to space: c, h > c, h
GroupNorm SiLU Conv c out, c out
x2
ResBlock c - 2'c
x2
ResBlock C > c/2
Figure 12. Architecture of Zebra's VQVAE for 2D datasets. Each convolution acts only on the spatial dimensions and uses kernel of size 3. The Residual Blocks are used to process information and increase or decrease the channel dimensions, while the Up and Down blocks respectively sa-sam and down-sample the resolution of the inputs. In 2D, used spatial compression factor of for Vorticity, and 8 for Wave2D. Every downsampling results a doubling of the number of channels, and likewise, every upsampling is followed by a reduction of the number of channels by 2. We choose maximum number of channels of 1024.
25

---

Zebra: In-Context Generative Pretraining for Solving Parametric PDEs
Table 8. Hyperparameters for Zehra's VQVAE
Hyperparameters
Advection
Heat
Burgers
Wave b
Combined
Vorticity
Wave 2D
start hiddeeddddeeee size
64
64
64
64
64
128
128
max_hidden_siae
256
256
256
256
256
1024
1024
num sin__dow__bloo. blocks
4
4
4
2
codebook.size
256
256
256
256
256
2048
2048
code_die dii num_codebooks
64
64
64
64
64
16
16
2
2
2
1
2
True
True
True
True
True
True
True
tokens-pex frame
32
32
32
32
32
256
128
start learning_ran
3e-4
3e-4
3e-4
3e-4
3e-4
3e-4
3e-4
bright_decay
lc-4
le-4
le-4
le-4
le-4
le-4
lc-4
scheduler
Cosine
Cosine
Cosine
Cosine 1000
Cosine
Cosine
Cosine
num.epochs
1000
1000
1000
1000
300
300
D. Additional Quantitative results
D.1. Alternative pretrainings
We experimented with different pretraining strategies before settling on the pretraining approach proposed in Zebra. The most intuitive way to adapt the next-token prediction objective for dynamics modeling is to operate in a continuous latent space, omitting the quantization step used in Zebra and therefore using deterministic transformer instead of generative model. As shown in previous studies (Li al., 2024; Agarwal et al., 2025), we obtained better reconstruction results using autoencoder instead of VQVAF.
However, we encountered two critical challenges with this approach. First, the training loss plateaued quickly, illustrated in Figure 13. Compared to training generative transformer (with the negative ligliiklihhod) as in Zebra, the deterministic variant trained with MSE loss exhibited instability and failed to improve over training steps.
3x100
NLL
MSE
2x100
Train Loss
Train Loss
10a
1000
2000
3300 4000
1000
2000
5000
4000
Step
Step
Training loss (negative log-likelihood) with Zebra
(b) Training loss (MSE) with deterministic transformer
Figure 13. Comparison of the training process between the Zebra transformer and deterministic transformer Advection. In both cases, the model trained to predict next token. Zebra utilizes a discrete vocabulary and learns a probability distribution over the next token, whereas the deterministic transformer optimized predict the using MSE loss.
Second, while the model was able to predict the next token at inference, it could not generate an entire trajectory. It performed particularly poorly in the one-shot adaptation setting. As demonstrated in Figure 14, errors accumulated quickly during inference, leading to rapid divergence from the ground truth. This ultimately resulted in poor reconstructions when feeding the predicted tokens into the decoder, as seen in Figure
26

---

Zebra: In-Context Generative Pretraining for Solving Parametric PDEs
Ground Fuuhh Preticsoo
28
12
Chanral
22-
30
Takon 20
25 30
Figure 14. When trained with MSE for next-token prediction, inference suffers from instabilities, causing errors to grow exponentially. Here, show evolution over the sequence tokens of particular channel illustrate phenomenon. The dashed lines show the temporal transitions between two subsequent frames, as here each frame is represented by 16 tokens.
Eormale
0.0
0:25
-0.0
1.50
3.0
359
1.0
021
-252
Figure 15. The deterministic transformer, trained with the next-token objective, performs poorly in the one-shot adaptation task.
27

---

Zebra: In-Context Generative Pretraining for Solving Parametric PDES
These observations highlight the importance of the generative aspect when adopting a next-token objective. For this reason, opted for quantized representation combined with transformer modeling discrete distribution, standard approach in image and video generation, but which has never been explored for modeling physical phenomena. While alternative strategies exist (Tian et al., 2024), they involve non-trivial extensions and are left for future works.
That said, next-token prediction pretraining may not be the only viable framework for developing in-context capabilities. To explore this, we experimented with a direct next-frame prediction approach using deterministic setup trained with relative L2 loss. This method, which we called VIT-in-conter is based on video transformer operating on patches with bidirectional attention. While its training behavior was more stable, inference results remained unsatisfactory.
serves as baseline for evaluating other in-context pretrainings.
D.2. Uncertainty quantification
Mean prediction and confidence interval
Moan,
2.0
Mean,
Mean,
Mean. Time
Mean. Time 4
1.5
Hean, Time 6
Hean, Time 5
Mean, Time 7
Mean, Time
1.0
Mean,
BxStd,
0.5
0.0
-0.5
-1.0
50
100
150
200
250
Figure Uncertainty quantification with Zebra a one-shot setting on Heat equation
Setting Since Zebra a generative model, it allows us to sample multiple plausible trajectories for the same conditioning input, enabling the computation key statistics across different generations. By calculating the pointwise mean and standard deviation, effectively visualize the model's uncertainty in its predictions. In Figure 16, the red curve represents the ground truth, the blue curve is the predicted mean and the blue shading indicates the empirical confidence interval (3 x standard deviation).
Metrics Motivated by this observation, we investigate how varying the model's temperature parameter affects its predictions, specifically in the one-shot adaptation setting described in Section 4.2. By adjusting we aim to assess its impact on both the accuracy and variability of the predictions. We generate different trajectories and use them to compute several metrics. We employ four metrics to evaluate the model's uncertainty:
1. Relative 2 loss:
This assesses the accuracy of the generated trajectories by measuring the bias of the predictions relative the ground truth.
2. Relative standard deviation: We estimate the variability of the predictions using the formula: Relative Std ||+. lla
where m. and o. represent the empirical mean and standard deviation of the predictions, computed pointwise across 10 generations.
28

---

Zebra: In-Context Generative Pretraining for Solving Parametric PDEs
3. Confidence level:
We create pointwise empirical confidence intervals CT(x) (m, - 3Ø, (1), (s), 36, (r)
and compute the confidence level as: Confidence level Ri Ex (+)) This score indicates how often the ground truth falls within the empirical confidence interval generated from sampling multiple trajectories.
4. Continuous Ranked Probability Score (CRPS, Gneiting & Raftery (2007)) ) is a proper scoring rule that measures the accuracy of a probabilistic forecast by quantifying the difference between the predicted cumulative distribution function (CDF) and the empirical CDF of the observed value, with lower values indicating better calibration and sharpness.
5. Root Mean Squared Calibration Error (RMSCE, Hendrycks et al. (2018)) quantifies the discrepancy between predicted confidence and actual accuracy at given confidence level.
When modeling uncertainty, the model achieves a tradeoff between the quality of the mean prediction approxima- tion and the guarantee for this prediction to be in the corresponding confidence interval. Figure 17 illustrates the trade-off between mean prediction accuracy and uncertainty calibration. At lower temperatures, we achieve the most accurate predictions, but with lower variance, i.e. with no guarantee that the target value within the confidence interval around the predicted mean. Across most datasets, the confidence level then remains low (less than 80% for 7 0.25), indicating that the true solutions are not reliably captured within the empirical confidence intervals. Conversely, increasing the temperature results in less accurate mean predictions and higher relative standard deviations, but the confidence intervals become more reliable, with levels exceeding 95% for 7 > 0.5. Therefore, the temperature can be calibrated depending on whether the focus is on accurate point estimates or reliable uncertainty bounds.
Results
To better calibrate this temperature, we can therefore use a proper scroring metric such as the CRPS, and we can pick the temperature parameter that has the lowest CRPS value for given value (see Figure 17). We can see that Zebra models really well the distribution for Combined and Advection and not so much for Burgers and Heat somehow.
Finally, we examine how the model's uncertainty evolves additional information is provided as input. Specifically, compare Zebra's s average error and relative uncertainty when conditioned on one example trajectory, with one or two frames as initial conditions. Table 9 reports the relative L2 loss and relative standard deviation for both scenarios. The results clearly show that including the first two frames as initial conditions reduces both the error and the relative standard deviation consistently. This indicates that, while some of the uncertainty remains aleatoric, the epistemic uncertainty is reduced as more input information becomes available.
Table Uncertainty quantification in the onc-shot setting. Conditioning from a trajectory example and frame or 2 frames initial conditions. Metries include relative 1,2 loss (average accuracy) and relative standard deviation (average spread around the average prediction). The temperature fixed at 0.1.
Advection
Heat
Burgers
Wave
Combined
Rel. L
1 frame 0.006
0.156
0.115
0.154
0,008
Rel. f.?
2 frames
0.004
0.047
0.052
0.075
0.005
Rel. Std.
frame
0.003
0.062
0.048
0.074
0.005
Rel. Std.
frames
0.002
0.019
0.018
0.040
0.003
Comparison with ViT-in-Context We quantitatively evaluate the uncertainty using CRPS and RMSCE in the one-shot prediction task, as reported in Table 10, comparing Zebra with ViT-in-Context. Since model, we introduce stochasticity by: (1) adding Gaussian noise with small standard deviation (a 0.1) to the input, thereby perturbing the input; and (2) enabling stochastic depth at inference time by keeping the DropPath mechanism (Huang et al., 2016). Across all settings, Zebra consistently outperforms these simpler baselines, often by an order of magnitude in both CRPS and RMSCE.
is deterministic
Uncertainty over time Finally, we illustrate in Table 11 how CRPS and RMSCE evolve over successive autoregressive steps in the predicted trajectory for Zebra on the Burgers equation. As expected, CRPS increases over time due to error accumulation during rollout. In contrast, RMSCE remains stable, indicating that the model maintains consistent calibration throughout the prediction.
29

---

Zebra: In-Context Generative Pretraining for Solving Parametric PDEs
Relative Error
Relative Standard Deviation
Relative Standard Deviation (log scale)
100
Rollout Loss (log scale)
1011
10-1
10-2
10-2
0.00
0.25
0.50
0.75
1.00
1.25
1.50
1.75
2.00
0.00
0.25
0.75
1.00
1.25
1.50
1.75
2.00
Temperature
Temperature
Confidence level
CRPS
10"
Confidence level (log scale)
CRPS (log scale)
6 60-1
Advection
10-2
Burgers Heat
40-1
Combined
Wave
3 -0-1
95% Confidence level
0.00
0.25
0.50
0.75
1.00
1.25
1.50
1.75
2.00
0.00
0.25
0.50
0.75
1.00
1.25
1.50
1.75
2.00
Temperature
Temperature
Figure 17. Uncertainty quantification with Zebra. The main parameter of this study is the temperature (x-axis). We then look from left to right at (1) The rollout loss, iee. the relative L loss between the predictions and the ground truth; (2) The relative standard deviation to quantify spread around mean; (3) The confidence level, that measures the frequency observations empirical confidence interval. (4) The CRPS that measures the quality of the uncertainty, can be used to pick the temperature with the most calibrated uncertaity.
lie within the
30

---

Zebra: In-Context Generative Pretraining for Solving Parametric PDEs
Table 10. CRPS and RMSCE results
PDE benchmarks. Lower is better.
Metrie
Model
Advection
Heat
Burgers
Wave
Combined
VIT noise
0.0705
0.176
0.227
0.093
0.098
CRPS
VIT Dropout
0,0363
0,213
0.196
0.024
0.024
Zebra
0.0026
0.043
0.020
0.0129
0.0018
VIT noise
0.132
0.241
0.265
0.249
0.045
RMSCE
ViT Dropout
0.386
0.547
0.529
0.340
0.064
Zebra
0.074
0.055
0.048
0.124
0.074
Table 11. CRPS and RMSCE over time for Zebra on the Burgers dataset.
Metric
t= 1
t= 2
==
f
== 5
t= 6
t=
t
8
t= 9
CRPS
0.0057
0.0106
0.0148
0.0184
0.0216
0.0244
0.0271
0.0296 0.0457
0.0320
RMSCE
0.0511
0.0513
0,0505
0.0487
0.0483
0.0489
0.0467
0.0457
D.3. Analysis of the generation
Zebra is capable of generating completely novel trajectories for new environments, including the initial conditions. An example of generated trajectory for Vorticity 2D is shown in Figure 18, where the top shows the context trajectory used to guide the generation, and the bottom row displays the model's 's generated trajectory, including the initial condition. In this section, we evaluate on Combined and Advection the quality of the generated trajectories with Zebra.
Figure 18. Unconditional generation Vorticity 2D. The top-row is the example used guide the generation, and the bottom-row is the generated example. The model also generates the initial condition.
Setting We evaluate whether our pretrained model can generate new samples conditioned on trajectory observed in previously unseen test environment. Unlike previous settings, the transformer is not explicitly conditioned on tokens derived from a real initial condition. Instead, we expect it to generate trajectories, including their initial conditions, that follow the same dynamics as the observed context.
Our evaluation focuses four key aspects. First, we assess whether the generated trajectories faithfully follow the dynamics of the context. Second, analyze the diversity of the generated trajectories to determine if they are significantly different from one another. Third, compare the generated samples with those produced by numerical solvers to evaluate whether their distributions align. Finally, we examine the types of initial conditions generated by Zebra.
Metries To assess fidelit,, we generate ground truth trajectories using the physical solver originally used to construct the dataset. These simulations start from the initial conditions generated by Zebra, using ground truth environment parameters that Zebra itself does not have access to. We then compute the L distance between the generated trajectories and those obtained from the physical solver.
31

---

Zebra: In-Context Generative Pretraining for Solving Parametric PDEs
For diversity, we measure the average pairwise L distance between different trajectories generated by Zebra. The results for both fidelity and diversity are reported in Table 12 for the Advection and Combined Equations.
To compare the distribution of generated trajectories with that of numerical solvers, we compute the Wasserstein distance using the Sinkhorn algorithm (Cuturi, 2013). As baselines, we compare against two references. First, we compute the Wasserstein distance with purely random samples drawn from Gaussian noise, providing an upper bound on the problem. Second, we measure the Wasserstein distance between two independent sets of numerical solver trajectories (the validation and test sets), allowing us quantify the variability inherent the dataset itself.
Finally, qualitative analysis, perform a principal component analysis on the trajectories generated by Zebra and visualize the first two principal components in Figure 19 for Combined.
Sampling We use a default temperature of = 1.0. For each context trajectory, we sample 10 new trajectories in parallel.
Results Table 12 shows that Zebra generates new initial conditions and trajectories that respect the same physical laws as the given context. The model seems to have learned the statistical relationships between initial conditions and later timestamps. The high average L distance between samples indicates that the generated trajectories are diverse. This can also be observed in Figure 19, where the generated samples effectively cover the distribution of real samples.
Table 12, Fidelity and diversity metrics. The L distance measures fidelity to the context dynamics, while the average L quantifies sample diversity.
Model
L
Average L between samples
Advection
0.0185
1.57
Combined Equation
0.0136
1.59
To further assess distribution alignment, we compute the Wasserstein distance between the generated trajectories and those obtained with numerical solvers. The results in Table 13 indicate that Zebra achieves lower Wasserstein distances than Gaussian noise but remains slightly above the cross-distribution baseline (which compares the validation and test distributions), suggesting reasonable alignment with the true data distribution.
Table 13. Comparison of distributions using the Wasserstein distance between Zebra-generater trajectories and numerical solver samples.
Distance Metric
Advection
Combined Equation
Gaussian noise real data
18.22
16.15
Validation data test data
5.11
1.87
Zebra -egeerate data vs. real data
5,57
2.21
D.4. Dataset scaling analysis
We investigate how the one-shor shot error on the test evolves as we vary the size of the training dataset. To this end, we train the auto-regressive transformer on datasets containing 10, 100, 1000, and 12,000 trajectories and evaluate Zebra's generations on the test set, starting with two frames as inputs. The training time is proportional to the dataset size: for example, the number of training steps for 1,000 trajectories is 10 times the number of steps for 100 trajectories. The results are presented in Figure 20.
First, we observe that Zebra requires a substantial amount of data to generalize effectively to different parameter values, even within the training distribution. This aligns with findings in the literature that transformers, especially auto-regressive transformers, excel scaling performing well on very large datasets and for larger model architectures. However, for smaller datasets, this approach may not be the most efficient. We believe that Zebra's 's potential resides when applied to large amounts of data, making it an ideal candidate for scenarios involving large-scale training.
Second, for the Combined equation, we notice that performance plateaus between 100 and 1,000 trajectories. This may be due to insufficient training or a lack of diverse examples, as the Combined equation more challenging compared to the Advection equation, whose performance follows a more log-linear trend. This suggests that additional data or targeted training strategies might be needed achieve better generalization for more complex equations.
32

---

Zebra: In-Context Generative Pretraining for Solving Parametric PDEs
PCA time
PCA time t=9
Cenerated Data
Real Data
Real Data
PCA Componentt
PCA Component
PCA Component 1
PCA Component 1
(a) Distribution of generated initial conditions (t = 0).
(b) Distribution of generated trajectories at (f = 9).
Figure 19. Qualitative analysis of generated trajectorics. Zebra generates new initial conditions and trajectories for unseen test
environments. PCA projections visualize both generated and truc trajectories in a lower-dimensional space at t 0 and 9.
Scaling Law: Training Size vs. Error
Combined
Advection
Emor Loss (log scale)
10-3
10
101
10
Training dataset size (log scale)
Figure 20. Dataset scaling analysis. One-shot prediction the test set the training dataset size.
33

---

Zebra: In-Context Generative Pretraining for Solving Parametric PDEs
D.5. Inference time comparison
Table 14 compares the inference time for one-shot adaptation across different methods when predicting single trajectory given context trajectory and initial condition. For Zebra, the inference process, which includes encoding, auto- regressive prediction, and decoding, is much faster in ID and slightly faster in 2D. With Zebra, the bottleneck at inference is the autoregressive generation of tokens, which speed is about 128 tokens per second on V100 for 2D and an RTX for ID. The decoding is fast and can be done in parallel for the trajectory in one forward pass. In contrast, for CODA and CAPE, the majority of the inference time is spent on adaptation and gradient-based steps. Here the times were reported with 100 gradient steps, note that we used 250 for the rest of the experiments. We believe Zebra's inference time could be further optimized by (1) improving the optimization code and leveraging specialized hardware such as H100 (for flash attention) and LPUs (which show significant speed-ups against GPUs), and (2) increasing the number of tokens sampled per step (as in e.g. next-scale prediction (Tian et al., 2024)).
However, we have shown that it is possible to greatly accelerate inference by employing a deterministic neural surrogate on top of Zebra, which acts as a dynamic encoder. This framework is order of magnitudes faster than gradient- basedd adaptation methods.
Table Inference times for one-shot adaptation. Average time seconds to predict single trajectory given context trajectory and an initial condition. Times include adaptation and forecast for CODA and CAPE, while it includes encoding, auto-regressive prediction and decoding Zebra.
Advection
Vorticity 2D
CAPE CODA
18s
23s
31s
28s
Zebra Zebra + UNet
3s
21s
0.10s
0.14s
D.6. Influence of the codebook size
The codebook size K' is crucial hyperparameter. It directly affects the quality of the reconstructions, since a larger codebook can improve the reconstructions quality. However, it also impacts the dynamics modeling stage: the smaller the codebook, the easier it for the transformer to learn the statistical correlations between similar trajectories. To have a sense of this trade-off, we report the relative reconstruction errors and the one-shot prediction errors in Table 15. The reconstruction error decreases when the codebook size increases. However, the one-shot prediction error decreases from 32 64 codes but then gradually increases from 64 512. We that it follows U-curve in Figure 21. This phenomenon was observed in different context in (Cole al., 2024).
Table 15. Influence of the codebook size. Reconstruction error and one-shot prediction error on Burgers for different codebook sizes. Errrr relative L2.
Codebook Size
Reconstruction
One-shot Prediction
32
0.0087
0.116
64
0.0043
0,097
128
0.0024
0.124
256
0.0019
0.163
512
0.0015
1.093
34

---

Zebra: In-Context Generative Pretraining for Solving Parametric PDEs
Prediction Loss vs Codebook Size (Log Scale)
109
Prediction Loss
Prediction (Log Scale)
10
Codebook Size
Figure 21. One-shot accuracy vs codebook size. One-shot prediction error on the test set for various codebook sizes. Error in relative L2.
D.7. Reconstruction errors
We report the accuracy of the reconstructions from our encoder-decoder for different datasets in Table 16. Here, no dynamics is involved, we simply evaluate the quality of the encoding and of the decoding. On ID and 2D datasets, the decoding errors are respectively of 0.1 % and 1% on the test set.
Table Reconstruction errors.
Test relative 1.2 loss between reconstructions from Zebra's VQVAE and the ground truths.
Advection
Heat
Burgers
Wave
Combined
Wave 2D
Vorticity 2D
VQVAE of Zebra
0.0003
0.0019
0.0016
0,0011
0.0022
0.010
0.017
To assess how faithful signal encoding and deconding influences prediction accuracy, we evaluate Zebra an out-of- distribution setting on the Vorticity 2D dataset across several viscosity ranges. As viscosity decreases and departs from the training distribution, both the context and target trajectories exhibit increasingly high-frequency components which seen by the VQVAE. We report reconstruction scores alongside one-shot prediction errors in Table 17.
Table 17. Impact of encoding and decoding quality on prediction accuracy. Reconstruction and one-shot prediction error with Zebra. Metric is relative L.2
Viscosity Range
Type
Reconstruction
One-shot Prediction
[10-3 10-21 5 x 10-4 10-4,10-)) 10-3]
In-distribution Close OoD
0.02
0.12
0.13
0.24
[10 5 50-5, 10=4]
Far OoD
0.22
0.32
As the viscosity decreases and deviates from the training distribution, the reconstruction error increases, which is expected since the VQVAE was not exposed to these regimes. However, the one-shot prediction error remains stable. This suggests that the encoder continues to capture the essential characteristics of the context dynamics in this one-shot OoD setting, enabling the transformer to leverage these sequence indices for accurate forecasting of the www-frequency components.
D.8. Influence of the number of context examples
While all models were evaluated on a one-shot adaptation task, Zebra a was trained with up to context examples as input. This allows us to analyze the impact of context size on prediction accuracy. As shown in Table 18, performance saturates after approximately in-context examples.
35

---

Zebra: In-Context Generative Pretraining for Solving Parametric PDEs
Table 18. Rffect of number of context example. Few-shot prediction error (relative 1.2) on ID datasets as function of the numher of context examples.
Examples Advection
Heat
Burgers
Combined
0.0074
0.1563
0.1150
0.0095
0.0077
0.1310
0.1020
0.0074
0,0072
0,1272
0.1000
0,0078
0.0071
0.1272
0.0990
0.0075
0.0071
0.1310
0.1000
0.0073
E. Qualitative results
We provide visualizations of the trajectories generated with Zebra under different settings in the following figures. One- shot prediction: Figure 22, Figure 24, Figure 26, Figure 28, Figure 30, Figure 35, Figure 38. Uncertainty quantification: Figure 23, Figure 25, Figure 27, Figure 29, Figure 27.
36

---

Zebra: In-Context Generative Pretraining for Solving Parametric PDEs
E. Advection
Example
Initial condition
Tmm
Time
1
Tmm 2
m0 3
1.0
mm
1.0
0.5
Value
Value
0.0
0.5
-0.5
-1.0
-15
50
100
150
200
250
0
50
100
150
200
250
Ground Truth
Generation
Tim
Tim
1.0
Tine
Time
Time
0.5
0.5
Rine
Time
0.0
0.0
Value
Value
-0.5
-1.0
1.0
-1.5
-1.5
100
200
250
100
150
200
250
Figure 22, One-shot adaptation on Advection
37

---

Zebra: In-Context Generative Pretraining for Solving Parametric PDEs
Example
Initial condition
Timz
1.0
mmc1
7m
1.0
mm
002 6
0.5
Value
Value
0.0
-0.5
0.5
-1.0
-1.0
50
100
150
200
250
0
50
100
150
200
250
Ground Truth
Mean prediction and confidence interval
Tim
0
Mean,
Mean, Time
1.0
1.0
Mean, Time
Mean, Time
0.5
Time
0.5
Gt. time
Ieen, Time
Moin dxSio. Time
Value
Value
-0.5
-1.0
-1.0
Figure 23, Uncertainty quantification on Advection
38

---

Zebra: In-Context Generative Pretraining for Solving Parametric PDEs
E.2. Burgers
Example 1
Initial condition
T Tmme 1
Tim
0.6
Tim 2
mm
1.00
mm
0.75
0.4
0.50
0.25
Value
Value
0.0
0.00
-022
0.25
0.4
0.50
-0.6
-0.75
50
100
150
200
250
50
100
150
200
250
Ground Truth
Generation
Tim 0
Time
1.00
Iime
0.75
0.75
Time
Time
Rine
line
0.50
0.50
0.25
0.25
Value
Value
0.00
0.00
0.25
0.25
-0.50
0,50
-0.75
-0.75
150
250
100
200
250
Figure 24, One-shot adaptation on Burgers
39

---

Zebra: In-Context Generative Pretraining for Solving Parametric PDEs
Example
Initial condition
m
Time
Tmm mmm3
0.6
Time
0.6
0.2
0.2
Value
Value
0.0
0.0
-0.4
0.4
-0.6
0.6
-0.8
50
100
150
200
250
0
50
100
150
200
250
Ground Truth
Mean prediction and confidence interval
Tim
0
Mean,
Mean, Time
0.6
0.4
Marn, Time
Mean Time
0.2
St. time
Hean, Time
Moan SaSte. Time
0.0
0.0
Value
Value
-0.2
-02
-0.4
0.6
008
-0.8
Figure 25. Uncertainty quantification on Burgers
40

---

Zebra: In-Context Generative Pretraining for Solving Parametric PDEs
Example 1
Initial condition
0.8
Time
0.6
0.4
0.0
0.2
Value
Value
-0.5
002
-0.4
-1.0
-0.6
150
250
100
200
250
Ground Truth
Generation
1.00
1.00
Tim
Tmc1
Iin
Tim 2
Iin
m
Tim
0.75
Tmm 4
0.75
Tim
Ino
Dmm
0.50
0.50
0.25
Value
Value
0.00
0.25
-0.25
-0.50
-0.50
-0.75
-0.75
50
100
150
200
250
50
100
150
200
250
Figure 26. One-shot adaptation on Heat

---

Zebra: In-Context Generative Pretraining for Solving Parametric PDEs
Example
Initial condition
1.0
0.8
Tmme
0.6
0.5
Tim 9
0.4
Value
Value
0.0
0.0
-0.
-0.2
-0.4
-1.0
-0.6
100
200
250
100
150
250
Ground Truth
Mean prediction and confidence interval
Time
Tmc1
Mean line
mic
Moan line
m83
1.0
Mean, Time
ms
Mean: Time
m3
Hoin, Time
Idean,
0.5
Time
0.5
0.0
0.0
Value
Value
-0.5
1.0
-1.0
-1.5
-2.0
50
100
150
200
250
0
50
100
150
200
250
Figure Uncertainty quantification on Heat
42

---

Zebra: In-Context Generative Pretraining for Solving Parametric PDEs
E.3. Heat
E.4. Wave boundary
Example 1
Initial condition
1.2
Tim
7mm 1
Tmm 2
mm 3
1.0
m9c
0.8
Value
0.6
Value
0.4
0.2
50
100
150
200
250
50
100
150
200
250
Ground Truth
Generation
2.00
Time 0
1.75
1.75
Time
1.50
1.50
mi
Tie
Tim
Time
Tie 12
1.25
Tie 11
Value
1.00
Value
1.00
0.75
0.75
0.50
0.50
0.25
0.25
0.00
0.00
50
100
150
200
250
50
100
150
200
250
Figure 28, One-shot adaptation on Wave b
43

---

Zebra: In-Context Generative Pretraining for Solving Parametric PDEs
Example
Initial condition
1.00
mic lime
Time
Time
Time
0.8
0.75
0.50
Value
Value
0.00
0.4
-.25
0.50
0.2
0.75
-1.00
0
50
100
150
200
250
50
100
150
200
250
Ground Truth
Mean prediction and confidence interval
Vean
Vean Dms
0.75
1.5
Vean, Time
0.50
Vean, Time
Time 6
Ween, Time
1.0
Wean Time 10
Irann lin 21
0.25
Vean. Hin 22
Time
Value
Value
0.00
0.5
0.25
0.0
0.50
00.
-0.55
-1.00
-1.0
Figure 29, Uncertainty quantification Wave
44

---

Zebra: In-Context Generative Pretraining for Solving Parametric PDEs
E.5. Combined equation
Example
Initial condition
0.6
T
0.8
Time
Tmm1 1
Tmm
mm 3
0.6
mm
0.4
0.4
0.2
0.2
Value
Value
0.0
0.0
-02
-0.2
0.4
-0.4
0.6
-0.8
50
100
150
200
250
0
50
100
150
200
250
Ground Truth
Generation
T 0
Tim
0.6
Tmm
Time
Time
0.4
0.4
Tine
Time
0.2
0.2
Value
Value
0.0
0.0
-022
-02
-0.4
0.4
-0.6
-0.6
100
200
250
100
150
250
Figure 30. One-shot adaptation on Combined
45

---

Zebra: In-Context Generative Pretraining for Solving Parametric PDEs
Example
Initial condition
1.00
1me mmc Tim
1im
Tim
0.75
0.4
0.50
0.25
0.2
Value
Value
0.00
0.25
-02
-0.50
-0.75
0
50
100
150
200
250
50
100
150
200
250
Ground Truth
Mean prediction and confidence interval
0.8
Mean Time
Mean Time
0.6
0.6
Mean, Time
Mean Time
0.4
Time
0.4
St. time
Meen, Time
Moan 20Ste. Time
0.2
Value
Value
0.0
0.0
0.2
-0.4
-0.4
0.6
0.6
Figure 31. Uncertainty quantification on Combined equation
46

---

Zebra: In-Context Generative Pretraining for Solving Parametric PDEs
E.6. Vorticity 2D
Example bwo
twd
Exammle
boample bnd
Example Leb
Examate
Example Eed
Examp
GT
Gr
GTT
Pred bw2
Prod
Pedte5
bed
Prod
--81
Figure 32, One-shot adaptation on Vorticity 2D. Example 1.
Exumple 1-3
Eramsle 1-6
Exemate
Example
-01
Gr
GTte1
GTtes
GT
GFts?
GTte8
G7tw9
Intordayy
Figure 33. One-shot adaptation on Vorticity 2D. Example 2.
Fosmple
Exumale
Doumple
Exandle Enf.
Eeumale
Fommpete
t-0
tw]
GT
G7 3m
bwd
GT
671w6
GT
twb
Figure 34. One-shot adaptation on Vorticity 2D. Example 3.
47

---

Zebra: In-Context Generative Pretraining for Solving Parametric PDEs
E. OUT-OF-DISTRRIBUTION
Exemple
Domple
Exampll
Examppe
Ezamp
GT1e6
GTte?
Prod tw1
Med be2
Prod
ProStee
Pcd te5
Redted
Mc3t
Prod bed
Prod b=0
Figure 35. One-shot OoD adaptation on Vorticity 2D, Example
Exomple 2-1
Towmp 111
Esample 2n2
Diomple 2n3
Trmmlee tnd
Example
Ezample 6
Exemale 1
Esompll ted
Tanmp Ind
Intosaay
Prod tw1
Red 2w2
Prod tw3
Pro3 tm4
Pod
tedd
Prdd tw?
Prod ted
Prod Tw0
Figure 36. One-shot OoD adaptation on Vorticity 2D. Example 2.
Fowmpitst
Exumale
Doumple
Esample
Examale
Exumale
tw]
GT
GTte5
GT
tT
tT
G7
103
Figure 37, One-shot OoD adaptation on Vorticity 2D. Example 3.
48

---

Zebra: In-Contert Generative Pretraining for Solving Parametric PDEs
E.7. Wave 2D
Example
Exampitel
Exampicte
Gl ox3
011
Gf te]
Q11+8
Pred
Pred
Figure 38. One-shot adaptation Wave 2/). Example 1.
Example
Fompettsl
Foumale
Diumplee
Eoample
Fowmple twis
Examale
Exomppe
twp
Prod 2-1
Irod 0w2
Prod D+5
Pred tw4
Pidd be5
bed
Prod tw?
Prod bed
Prod be9
Figure 39. One-shot adaptation on Wave 2D, Example 2.
Example L-O
Exampetei
Example b=2
Example bed
Example Led
Esample
Example Lwd
Examale bwi
Example bwd
Examp b-g
O
Incond
Fred 1wl
Pred 2u2
Pred te3
Predited
Pred tm5
Predteff
Pred
Pred toll
Pred to
Figure 40. One-shot adaptation on Wave 2D. Example 3.


Reviews:
[{"review_id": "B1GniDgHSA", "reviewer": "Reviewer_6LCA", "summary": "The paper addresses the problem of predicting PDE solutions without explicit knowledge of the underlying dynamics using a generative auto-regressive transformer. The model is trained using contexts composed of different trajectories. During inference, the model is given a context containing trajectories governed by the same dynamics and the initial condition to evolve.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "### 1. On the uncertainty\n\nThank you for your comments. We agree that the main contribution lies in adaptation through in-context learning, with uncertainty emerging as a byproduct. We aim to clarify this distinction and welcome suggestions for improving it.\n\nThe generative capacity stems from the transformer, not the VQ-VAE. Our decoder-only transformer, as in language modeling, predicts the next token from past tokens and samples from this distribution. Since each sampled token depends on previous ones, uncertainty propagates through time.\n\nThe VQ-VAE encodes continuous frames into discrete tokens and compresses the input, making transformer training feasible. It strikes a balance between fidelity and compression. In contrast, deterministic vision transformers (e.g., ViT, MPP, Poseidon) lack generative ability and do not benefit from in-context examples as our model does (see Table 1).\n\nAlthough uncertainty modeling is not our focus, we found it relevant to highlight Zebra’s ability to sample diverse trajectories (e.g., Fig. 3). Below, we evaluate uncertainty quality using CRPS and RMSCE. CRPS measures accuracy of probabilistic predictions; RMSCE evaluates calibration when ground-truth uncertainties are unknown. Lower is better. Zebra consistently outperforms baselines, often by an order of magnitude:\n\n* CRPS and RMSCE Results\n\n| Metric | Model        | Advection | Heat  | Burgers | Wave b | Combined |\n|--------|--------------|-----------|-------|---------|--------|----------|\n| CRPS   | ViT + noise  | 0.0705    | 0.176 | 0.227   | 0.093  | 0.098    |\n|        | ViT Dropout  | 0.0363    | 0.213 | 0.196   | 0.024  | 0.024    |\n|        | Zebra        | **0.0026**| **0.043** | **0.020** | **0.0129** | **0.0018** |\n| RMSCE  | ViT + noise  | 0.132     | 0.241 | 0.265   | 0.249  | **0.045** |\n|        | ViT Dropout  | 0.386     | 0.547 | 0.529   | 0.340  | 0.064    |\n|        | Zebra        | **0.074** | **0.055** | **0.048** | **0.124** | 0.074    |\n\nBelow, CRPS and RMSCE over time for Zebra on Burgers. As expected, CRPS grows with error accumulation. RMSCE remains stable, indicating consistent calibration.\n\n| Dataset   | Metric | t=1    | t=2    | t=3    | t=4    | t=5    | t=6    | t=7    | t=8    | t=9    |\n|-----------|--------|--------|--------|--------|--------|--------|--------|--------|--------|--------|\n| **Burgers**   | CRPS | 0.0057 | 0.0106 | 0.0148 | 0.0184 | 0.0216 | 0.0244 | 0.0271 | 0.0296 | 0.0320 |\n|            | RMSCE | 0.0511 | 0.0513 | 0.0505 | 0.0487 | 0.0483 | 0.0489 | 0.0467 | 0.0457 | 0.0457 |\n\n### 2. New trajectory generation\n\nThis experiment is exploratory. The model generates a full trajectory, including its initial state, when prompted with examples (no initial state given). Multiple diverse outputs can be sampled from the same context. See Appendix D.3, Table 10.\n\n### 3. Clarifying the scope\n\nWe will clarify this in the introduction. Our focus on forecasting from an initial state highlights the role of context.\n\n### 4. \"provides richer information...\"\n\nAgreed. We will remove the sentence.\n\n### 5. On scaling gradient-based adaptation\n\nOur concern is the high computational cost of meta-learning with double-loop optimization. Alternatives like CODA avoid this but drastically increase parameters count. Our approach avoids both issues.\n\n### 6. Index representations\n\nWe agree this needs clearer explanation. Lines 146–152 mention that each frame is encoded into tokens via VQ-VAE. These tokens come from a finite vocabulary and are replaced by indices $s$ for processing in the transformer—similar to language modeling.\n\n### 7. On the main figure\n\nWe only show final timesteps, which may seem unrelated. This will be clarified in the caption. See Appendix E.6 for better illustrations.\n\n### 8. Caption revision\n\nThe model takes as input example trajectories $u_1$, $u_2$ and a new initial condition $u^0_*$. These are tokenized into $s_1$, $s_2$, $s^0_*$, concatenated, and fed into a transformer that predicts next tokens. The output is then detokenized into a physical trajectory.\n\n### 9. \"Sequences\"\n\nWe will clarify this. Each $s_i$ is a sequence of VQ-VAE indices.\n\n### 10. <eot> token\n\nWe will try to include it, possibly with smaller font.\n\n### 11. Resolution\n\nWe used 64×64 in 2D and 256 in 1D.\n\n### 12. \"Blurry\" expectations\n\nWe will remove this term if it’s unclear. We meant deterministic models tend to predict conditional expectations, which average over possibilities.\n\n### 13. Fig. 2 issue\n\nWe will fix it.\n\n### 14. Different dynamics\n\nEach context uses a single dynamic. During training, we sample 1–5 examples per context. The dataset has ~1200 dynamics, each with 10 trajectories. For better GPU usage, we stack sequences from different dynamics, but this is marginal.\n\n### 15. ICL training\n\nWe train across a variety of dynamics. This differs from MPP and Poseidon, which train across physics types. In our experiments (e.g., Table 1), all methods use the same data, and Zebra outperforms ViTs.", "claims_and_evidence": "The main claims in the paper are well-supported by clear and convincing evidence. In particular, the pretraining strategy is novel in its application to PDEs and demonstrates strong results in one-shot adaptation compared to alternative adaptation methods.\n\nHowever, the claim regarding the benefits of using a VQ-VAE to provide the model with generative capabilities could be better substantiated (see the primary weakness below).", "methods_and_evaluation_criteria": "The methods and evaluation criteria are appropriate for supporting the paper’s main claims.", "theoretical_claims": "The paper does not introduce theoretical claims in the sense of proving a new theorem.", "experimental_designs_or_analyses": "The main figures and tables of the paper are sound and correctly support the main claims of the paper. \n\n- Table 1 and Figure 2 demonstrate the zero-shot capabilities of the model\n- Table 2 presents compelling out-of-distribution results\n\nThe experiments assessing the generative aspect of the model, particularly in uncertainty quantification (Figure 3) and distribution quality (Figure 4, Table 3), are a bit weak.\n\n- Figure 3: It is unclear what this figure illustrates beyond showing that the generative model produces a reasonable mean prediction, with the ground truth contained within three standard deviations. A more convincing validation would involve a scenario where the true uncertainty is known and verifying whether the model accurately reproduces that uncertainty. However, the dataset used likely does not provide ground-truth uncertainty levels.\n\n- Figure 4 and Table 3: Their purpose is unclear. If they aim to demonstrate that the model produces states that fall within the training distribution, this is a relatively weak result. Any model achieving reasonable validation accuracy should satisfy this criterion, provided the validation set is not significantly different from the training set.", "relation_to_broader_scientific_literature": "Compared to gradient-based adaptation techniques such as (e.g. CODA paper, Kirchmeyer et al. 2022), which require retraining for every new environment, the paper introduces an approach capable of tackling zero-shot problems through in-context learning. The proposed model builds upon the transformer introduced in (McCabe et al. 2023) by enhancing its in-context capabilities, enabling it to predict from a single initial condition, which this model cannot achieve.", "essential_references_not_discussed": "All essential references have been discussed.", "other_strengths_and_weaknesses": "Strengths:\n\n- The in-context pretraining strategy is quite relevant to the ML & PDE community, offering the potential for training larger models with enhanced in-context capabilities.\n\n- The reduction in inference costs proposed by the model is valuable to the field.\n\nWeaknesses:\n\n- The justification for incorporating a variational autoencoder in the architecture is unclear, as uncertainty quantification is not a central focus. Moreover, the uncertainty appears to apply only to individual states rather than the overall prediction. If each frame is treated independently, the model merely generates independent random perturbations per frame, rather than meaningful uncertainty estimates over entire trajectories.", "other_comments_or_suggestions": "- The paper could emphasize more clearly that it tackles PDE prediction from an initial state only. At inference, even though the context contains trajectories with the correct dynamics, the model is given only an initial state from the trajectory to evolve. Since the model is presented with only one state to evolve, this task is technically more challenging than predicting a trajectory with multiple past states (which explains I think why you did not compare to (McCabe et al. 2023)). However, this also makes it a more restricted task, and this should be explicitly stated. For example, if a single frame does not fully capture the system’s state (e.g., missing pixels, incorrect high frequencies), having more past frames would reduce this uncertainty, but that falls outside your proposed task.\n- l074: \"provides richer information than deterministic...\". This holds only if the conditional distributions are correctly calibrated, which the paper does not adequately demonstrate. Otherwise, generating a conditional distribution is trivial, but proper calibration is key.\n- l079: \"scales better than gradient-based adaptation approaches\"—vague. What aspect of scaling are you referring to?\n- l163: Unclear what you mean by \"into index representations.\"\n- l181: The figure is confusing—why does the generated trajectory seem unrelated to the query? Is there an issue?\n- l182: The caption is poorly written and should be clarified. Example: \"Sequences have similar dynamics but different initial conditions, generating trajectories that follow the same underlying physical behavior from a new initial condition (query)\", \"according to the sequence design used.\"\n- l189: What does \"sequence\" refer to here? A sequence of image patches? A sequence of states? The explanation of the pretraining strategy should be clearer.\n- l202: Double \"((\" and \"))\".\n- l222: Why is there no <eot> token before the last <bot> token?\n- l257: The main text should specify the resolution of the states used.\n- l298: \"Deterministic models tend to predict conditional 'blurry' expectations...\"—this is not a valid argument. Blurriness in deterministic models can result from poor high-frequency reconstruction by the autoencoder, an issue that a VQ-VAE may also face. If you mean that a VQ-VAE reconstructs high frequencies faster, this needs to be demonstrated.\n- l338: Figure 2 appears to have an issue—captions or y-labels seem misaligned in the first column.", "questions_for_authors": "- What proportion of the contexts in the training dataset contain trajectories from different dynamics?\n- Your model should enable strong one-shot performance without the need for retraining (unlike CODA, Poseidon, MPP). However, this obscures the fact that your model is trained on a much broader variety of contexts, including multiple trajectories, sometimes from different dynamics. Could you comment on how much additional training is required for your method to develop in-context capabilities, compared to (McCabe et al. 2023) for example?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ypZ1Olf3pJ", "reviewer": "Reviewer_qxAk", "summary": "This paper proposes an in-context generative auto-regressive transformer for solving parametric PDEs and outperforms gradient adaptation methods.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "### 1. Continuous vs Discrete\n\nThanks for the comment. We agree that encoder quality is crucial for model performance. However, the answer is not trivial and we will try to answer in different points.\n\n* Our goal was to build an autoregressive probabilistic predictor for parametric PDE dynamics that adapts to new dynamics (new values of the PDE parameters). We took inspiration from in-context learning in language modeling, which relies on discrete tokens from a finite vocabulary.\n* *Discrete representations*: To evaluate encoder quality, we tested different codebook sizes. Larger codebooks preserve more information (closer to continuous encoding), as shown in Fig. 21 and Table 13. There is a tradeoff between reconstruction quality (improved with larger codebooks) and forecasting performance. We used a codebook size of 256 as a compromise across 1D datasets.\n* *Continuous representations*: A generative model on continuous representations would require predicting distributions over continuous states at each time step. We initially explored this, but it remains an open problem. Recent attempts for image generation [1] do not yet translate to the more complex problem of modeling PDE dynamics. Thus, a direct comparison with continuous autoregressive probabilistic models is not feasible.\n* It is however possible to compare with deterministic approaches operating on continuous representations, and this is what we did with the vision transformers (ViT) in tables 1 and 2. Additionally, Appendix D.1 (Fig. 15) shows that using a deterministic transformer to predict the next token over continuous latent representations performs poorly in the one-shot adaptation setting.\n\n[1] Autoregressive Image Generation without Vector Quantization, Li et al. 2024.\n\n### 2. New forcing terms for Heat\n\nGreat point. We added experiments with out-of-distribution (OoD) forcing terms: **Gaussian** and **Square** patterns.\n\n* **Gaussian:**  \nWe define forcing as a sum of spatial Gaussian pulses modulated in time:\n$$\nf(x, t) = \\sum_{i=1}^n A_i \\exp\\left(-\\frac{(x - \\mu_i)^2}{2\\sigma_i^2} \\right) \\sin(\\omega_i t)\n$$\nThis introduces smooth but localized excitation. Solutions resemble the homogeneous case in regions far from peaks, making this a \"close\" OoD test.\n\n* **Square:**  \nWe define forcing as a sum of square waves:\n$$\nf(x, t) = \\sum_{i=1}^n A_i \\text{sign} \\left( \\sin\\left( \\omega_i t + \\frac{2\\pi l_i x}{L} + \\phi_i \\right) \\right)\n$$\nThis preserves parameterization but introduces discontinuities and high-frequency content, making it a more challenging OoD case.\n\nWe evaluate both without retraining, using context examples from these new trajectories. Results below compare ViT-in-context and Zebra:\n\n| Model           | Gaussian Forcing | Square Forcing |\n|----------------|------------------|----------------|\n| ViT-in-context | 0.4023           | 0.6584         |\n| Zebra          | **0.3263**       | **0.3648**     |\n\nWhile both models show performance degradation, Zebra remains significantly more robust than the baseline.\n\n### 3. New trajectory generation\n\nOur apologies if this was unclear. This actually corresponds to the fidelity metric described in Appendix D.3 and we will indicate it more clearly in the main text. Table 10 in the appendix reports the average L2 distance between the generated trajectory and the ground truth trajectory, which is simulated using the numerical solver with the true physical parameters (inaccessible to Zebra) and the initial condition generated by Zebra. These results show that our model can produce consistent trajectories for a given set of physical parameters.\n\n### 4. Experiments from Appendix D.\n\nThanks for the suggestion. We agree and will reference these findings directly in the main text.\n\n### 5. Generative models\n\nThanks, we will include these references in a section on SOTA generative PDE simulators.\n\n### 6. Encoder-decoder analysis\n\nWe agree that encoding quality is key. As discussed in point 1, we analyzed the impact of codebook size on reconstruction and forecasting.  We found out that although the quality of the encoding benefits from larger codebooks, there is a compromise between the codebook size and the quality of the forecast. VQ-VAE offers a strong balance between performance and complexity. Let us know if any further analysis would be valuable for the camera-ready version.\n\n### 7. Plot differences\n\nAgreed, we will add additional plots to show the differences between ground truths and predictions.\n\n### Q1. # of environments.\n\nActually, this was arbitrary. We can align all test settings to 120 environments for the camera-ready version; this will not significantly change the results.", "claims_and_evidence": "Given that the authors claim this is the first successful application of discretized representations in physical systems, they should further analyze the impact of discretized representations on modeling physical systems. For example, an ablation study could be conducted to compare it with continuous representation.", "methods_and_evaluation_criteria": "For the experiment in Section 4.3 ‘Out-of-distribution generalization,’ could the prediction accuracy be explored when the external force term differs in form from the external force terms in the training dataset? In other words, when the external force term is no longer a sum of sine functions?", "theoretical_claims": "The paper does not have theoretical claims.", "experimental_designs_or_analyses": "For the ‘New trajectory generation’ experiment in Section 4.4, a more accurate evaluation should involve inputting the initial state and other parameters of the generated sequence into the solver to obtain the true trajectory, and then comparing it with the generated trajectory.", "relation_to_broader_scientific_literature": "This paper explores a new direction in PDE simulation, namely in-context learning, which differs from previous works, and they also use discretized representations.", "essential_references_not_discussed": "L84. There are other state-of-the-art generative models for PDE simulation. [1, 2, 3]\n\n[1] Huang J, Yang G, Wang Z, et al. DiffusionPDE: Generative PDE-Solving under Partial Observation. The Thirty-eighth Annual Conference on Neural Information Processing Systems.\n\n[2] Peiyan Hu, Rui Wang, et al. Wavelet Diffusion Neural Operator. The Thirteenth International Conference on Learning Representations.\n\n[3] Mario Lino Valencia, Tobias Pfaff, Nils Thuerey. Learning Distributions of Complex Fluid Simulations with Diffusion Graph Networks. The Thirteenth International Conference on Learning Representations.", "other_strengths_and_weaknesses": "Strengths:\n\nThe paper is overall novel as it explores a new direction. Moreover, its experiments are comprehensive and thorough, and the writing is clear.\n\nWeaknesses: \n\nMore analysis of the encoder-decoder architecture would be better, as the current analysis is quite limited, yet it is an important part of the algorithm.", "other_comments_or_suggestions": "It would be clearer to plot differences between the ground truths and predictions in Figure 2.", "questions_for_authors": "Sec 4.1. Why use a lot more unseen environments for the 2D datasets (120) than for the 1D datasets for testing (12)? There should be an explanation.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "RUVc0M3DoD", "reviewer": "Reviewer_ikHe", "summary": "This paper introduces Zebra, a novel generative autoregressive transformer designed to solve parametric partial differential equations (PDEs) without requiring gradient adaptation at inference. The key innovation is leveraging in-context learning to dynamically adapt to new PDE parameters. Zebra employs a two-stage framework: first, a vector-quantized variational auto-encoder (VQ-VAE) compresses physical states into discrete tokens; second, a generative autoregressive transformer is pre-trained on these tokens with a next-token prediction objective. During inference, Zebra conditions on example trajectories with similar dynamics but different initial conditions to predict new trajectories. As a generative model, Zebra enables uncertainty quantification and trajectory sampling. The authors demonstrate Zebra's effectiveness across various PDE scenarios, showing it outperforms gradient-based adaptation methods while being computationally more efficient. They also propose an accelerated inference procedure using a UNet conditioned on a dynamics embedding, further improving speed by orders of magnitude.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "### 1. DPOT  \nThank you for mentioning this interesting work. We will include this reference, along with those suggested by other reviewers, in the final version and discuss how it differs from our approach and problem setting.\n\n### 2. Explanations and Limitations  \nThanks for the comment. Clear explanations are indeed crucial. We will provide a concise overview of our model and its differences from prior work before the detailed description.\n\nAs for limitations, our scaling analysis (Appendix D.4, Fig. 20) suggests the framework suits equations with diverse physical parameters and many training trajectories. In data-scarce regimes, this approach may be less effective—a general issue for data-driven models in complex physics. Additionally, discrete tokenization may be a bottleneck for simulations with significant high-frequency content. Finally, we have not tested extensions to more demanding scenarios (e.g., 3D) due to compute constraints. We will highlight these limitations more clearly in the appropriate section.\n\n### 3. Computational Resources  \nThank you for pointing this out. In 1D, VQVAE training takes ~4 hours on an RTX 24 GB GPU, and the transformer takes ~15 hours. In 2D, these increase to ~20 h each on a single A100 80 GB. Overall, our method requires more training than vanilla ViT in-context but is comparable to meta-learning baselines.\n\n### 4. Notations  \nWe agree. In particular, distinguishing between flattened and non-flattened indices can be confusing. We will revise and simplify the notations accordingly.\n\n### 5. Complexity  \nThis is a great point. Our method applies attention over the full sequence—i.e., context examples $N$, timesteps $T$, and tokens per frame $h \\times w$. Adding a third spatial dimension $d$ yields a complexity of $\\mathcal{O}((NThwd)^2)$. This can result in long sequences, and more demanding cases may require architectural changes, e.g., axial transformers. We will discuss this further.\n\n### Q1: VQVAE  \nWe agree. Fidelity depends on codebook size: larger sizes reduce discretization error. Appendix Fig. 21 and Table 13 show how reconstruction and prediction errors evolve with codebook size. While reconstruction improves with larger codebooks, one-shot prediction shows a trade-off—e.g., for Burgers' equation, the best size is 64. Larger codebooks increase transformer classifier complexity, as the number of parameters scales with the number of classes.\n\nAs for the behavior for turbulent systems, we performed OoD experiments for the vorticity equation. Here are the results:\n\n| **Viscosity Range**              | **Type**             | **Reconstruction** | **Prediction** |\n|----------------------------------|----------------------|--------------------|----------------|\n| $[10^{-2},\\ 10^{-3}]$            | In-distribution      | 0.02               | 0.12           |\n| $[5 \\times 10^{-4},\\ 10^{-3}]$ | Close OOD            | 0.13               | 0.24           |\n| $[10^{-5},\\ 10^{-4}]$          | Far OOD              | 0.22               | 0.32           |\n\nAs the viscosity decreases and moves out of the training distribution, the reconstruction error increases, which is expected since the VQ-VAE was not exposed to these regimes. However, the one-shot prediction error increases less sharply. This suggests that the encoder still captures the main structures of the flow that are important for prediction. The transformer can rely on these latent representations to make accurate forecasts, even if the reconstruction is less precise. Overall, the model shows pretty good robustness to distribution shifts.\n\n### Q2: Memory Requirements  \n\nThe memory requirements remain manageable for 1D experiments. However, in 2D, they become more demanding, as training sequences can reach lengths of up to 8192 tokens. To handle this, we use a single A100 GPU with 80 GB of memory, which is sufficient to train the model effectively. For reference, this sequence length is comparable to that used during the pretraining phase of LLaMA 3.\n\n### Q3: Number of Examples  \nGreat remark. We ran additional experiments varying the number of context examples on 1D datasets. We can see that the performance reaches a plateau after about 3 in-context examples:\n\n| # Examples | Advection | Heat   | Burgers | Combined |\n|------------|-----------|--------|---------|----------|\n| 1          | 0.0074    | 0.1563 | 0.1150  | 0.0095   |\n| 2          | 0.0077    | 0.1310 | 0.1020  | 0.0074   |\n| 3          | 0.0072    | 0.1272 | 0.1000  | 0.0078   |\n| 4          | 0.0071    | 0.1272 | 0.0990  | 0.0075   |\n| 5          | 0.0071    | 0.1310 | 0.1000  | 0.0073   |", "claims_and_evidence": "The claims in the paper are generally well-supported by convincing evidence. The authors provide comprehensive experimental results across seven different PDE datasets, comparing Zebra against state-of-the-art adaptation methods (CODA, CAPE, and other baselines). The empirical evidence demonstrates Zebra's superior performance, particularly in challenging scenarios like 2D environments and out-of-distribution settings. The authors substantiate their claims about Zebra's generative capabilities through quantitative metrics (Wasserstein distance between generated and real distributions) and qualitative visualizations. The paper also provides clear evidence for the efficiency gains of their proposed accelerated inference method, with specific timing measurements showing improvements of one to two orders of magnitude.", "methods_and_evaluation_criteria": "The proposed methods and evaluation criteria in this paper are well-suited for the parametric PDE problem. The authors carefully designed their approach to address the specific challenges of adapting to various PDE parameters without gradient updates at inference time. The evaluation uses seven diverse PDE datasets (five 1D and two 2D) that encompass various physical phenomena and parameter variations, providing a comprehensive testbed. The authors appropriately evaluate adaptation capabilities in both in-distribution and out-of-distribution settings, which is crucial for real-world applications where parameters may shift.", "theoretical_claims": "The paper does not contain formal mathematical proofs for theoretical claims. It primarily focuses on empirical validation of the proposed Zebra framework through extensive experimentation", "experimental_designs_or_analyses": "I reviewed the experimental design and analyses in the paper, finding them largely sound and well-executed. The authors' one-shot adaptation experiments across seven datasets provide a comprehensive evaluation of Zebra's capabilities. Their comparison against multiple baselines (CAPE, CODA, [CLS]ViT, and ViT-in-context) establishes a fair benchmark for assessing performance.", "relation_to_broader_scientific_literature": "The paper also connects to recent foundation models for physics, such as those proposed by Subramanian et al. (2023) and Herde et al. (2024), but distinguishes itself by focusing on efficient adaptation rather than comprehensive pretraining. Additionally, it relates to specific work on using discrete representations for physical systems, building upon the VQ-VAE framework introduced by van den Oord et al. (2017).\nMost directly, Zebra extends the in-context operator learning approach proposed by Yang et al. (2023), but overcomes its limitations by incorporating a generative framework and developing a scalable architecture that can handle complex 2D dynamics.", "essential_references_not_discussed": "The paper does not cite recent advances in generative state-space models for physical systems, such as Hao et al. (2024) \"DPOT: Auto-regressive Denoising Operator Transformer for Large-scale PDE Pretraining.\" This work approaches similar problems using denoising techniques and offers an alternative perspective on generative pretraining for PDEs.", "other_strengths_and_weaknesses": "Pros:\nThe paper demonstrates several notable strengths. Its most significant contribution is the innovative application of in-context learning principles from language models to physical systems, creating a bridge between these domains. This cross-pollination of ideas shows considerable originality and opens new avenues for solving complex physical problems without gradient-based adaptation.\n\nZebra's generative nature provides meaningful uncertainty quantification, which is crucial for high-stakes scientific and engineering applications. This capability differentiates it from deterministic approaches and enhances its utility for decision-making under uncertainty.\n\nCons:\n\nCould be improved with more concise explanations of the transformer architecture. The technical details, while comprehensive, sometimes obscure the core conceptual innovations. Additionally, the paper would benefit from more explicit discussion of limitations, particularly regarding the types of PDEs where Zebra might underperform.\n\nThe paper lacks discussion of computational resources required for training, which is important for understanding the practical accessibility of this approach to researchers with varying computational capabilities.", "other_comments_or_suggestions": "When describing the sequence design, the notation could be clarified to better distinguish between different trajectories and their corresponding token representations. The transition between mathematical notation and practical implementation is sometimes difficult to follow.\n\nThe paper would benefit from a more detailed discussion of computational complexity as a function of physical domain size, particularly for scaling to higher-dimensional PDEs beyond the 2D cases presented.", "questions_for_authors": "1. Given that the discretization of physical fields through VQ-VAE is fundamental to your approach, how does Zebra's performance depend on the fidelity of this discretization? Specifically, how would performance degrade for highly turbulent or chaotic systems where fine-scale features are critical?\n\n2. While you demonstrate Zebra's computational efficiency compared to gradient-based methods, could you elaborate on the memory requirements and training costs?\n\n3. The paper shows impressive performance in one-shot adaptation, but how does Zebra's performance scale with additional context examples? Is there a point of diminishing returns, and does the benefit of additional context vary across different PDE types?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "5rtPhzYxF2", "reviewer": "Reviewer_tD1F", "summary": "This paper proposes a model for modeling physical dynamics by leveraging the power of large language models (LLMs). The model handles data given in the physical domain through quantized representations, utilizing in-context learning. However, the advantages of this approach over other LLM-based foundation models are unclear, and the specific contribution to addressing dynamics governed by partial differential equations (PDEs) remains ambiguous.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "### 1. On the originality\nThere seems to be a misunderstanding, and we apologize if our explanations were not sufficiently clear. We appreciate the opportunity to clarify the originality of our contribution.\nObjective: Our goal is to develop a family of models capable of solving parametric PDEs. We argue that classical empirical risk minimization strategies are not well-suited for this task and instead advocate for a learning-to-adapt approach. To this end, we introduce a novel method and demonstrate, across a broad range of scenarios, that it performs comparably to—and often surpasses—state-of-the-art models.\nModels: To the best of our knowledge, this is the first instance of an LLM-based approach for PDE modeling. By LLM, we refer to generative probabilistic models, akin to the decoders used in language modeling. Due to their probabilistic nature, these models enable the generation of trajectory distributions. In contrast, the models you referenced (S. Subramanian, 2023; M. Herde, 2024), as well as those listed in Appendix A.1 under \"The classical ML paradigm\", are deterministic vision transformers (ViTs) rather than generative models. These models rely on fine-tuning when exposed to new data, making them ill-suited for scenarios with limited data, and they cannot be used to generate trajectory distributions. Additionally, in our experiments, ViTs performed poorly for in-context learning.\nKey Benefits and Novelty: Our formulation draws inspiration from the in-context learning capabilities of generative LLMs and explores their potential application to solving parametric PDEs. In this context, in-context learning means that, when provided with examples of a new task, the model can generate an appropriate response to a query. In our case, the examples consist of trajectories from the same dynamical system, and the queries correspond to new initial conditions. This approach enables the model to generalize to new dynamics using only a few examples, without requiring retraining. Achieving this level of generalization with minimal data and computational cost, compared to existing baselines, is precisely our objective.\nWe acknowledge that solving downstream tasks may be important in certain contexts, but this falls outside the scope of our current work.\n\n### 2. On the training data\nWe respectfully disagree. Our datasets encompass seven distinct equations with a wide range of parameter values, leading to diverse and complex dynamics. These variations include differences in PDE coefficients, initial and boundary conditions, and forcing terms. As a result, the learning task is inherently challenging, as evidenced by the comparatively weaker performance of baseline models.\nOur model, consistently outperforms these baselines in both in-distribution and OoD settings. Additionally, our experimental setup considers a broader variety of parametric scenarios than the referenced papers. Even in the latest large-scale PDE benchmark, The Well [1], the distribution of PDE parameters is more limited compared to our experiments. In this sense, our evaluation setting presents a more rigorous challenge.\n\n[1] The well: a large-scale collection of diverse physics simulations for machine learning. Ohana et al. 2024. \n\n### 3. On VICON\nWe thank the reviewer for bringing this reference to our attention. We were not aware of this paper prior to submission, and we will include it in the related work section.\nHowever, a direct comparison is not feasible, as the referenced model addresses a different problem. While our approach focuses on adapting to a new parametric setting based solely on an initial condition—akin to traditional PDE solvers—the referenced work leverages past trajectory history for future-state forecasting. Additionally, no code is available for this unpublished paper, making direct empirical comparison difficult.\nThat said, our ViT In-Context baseline already captures a setup that is closely aligned with the referenced work but adapted to our problem setting. This serves as an implicit comparison between the two strategies.\nFinally, we note that this paper was publicly released only in late November 2024. Under ICML guidelines, it therefore qualifies as concurrent work (see below).\nExcerpt from the reviewer guidelines at https://icml.cc/Conferences/2025/ReviewerInstructions\n“Concurrent Works\nAuthors cannot expect to discuss other papers that have only been made publicly available within four months of the submission deadline. (This cut-off is adopted from the AISTATS and ICLR reviewing instructions.) Such recent papers should be considered as concurrent and simultaneous. Good judgement is necessary to decide whether a paper that has not yet been peer-reviewed should be discussed. The guideline is to follow the best practices of the specific subfield; the Area Chair can help in each case with these.”\n\n### 4. Relation To Broader Scientific Literature\nWe hope that the clarifications provided help to reassess this statement.", "claims_and_evidence": "Recent advancements in large language models LLMs have led to the proposal of several PDE-based foundation models using LLMs. However, it is unclear what specific advantages the proposed approach offers compared to existing models (S. Subramanian, 2023; M. Herde, 2024). The method seems to primarily incorporate widely used techniques like VQVAE, which have been extensively explored in deep learning literature, yet there is no discussion on how these techniques offer advantages for handling sequential data governed by PDE dynamics. It would be beneficial to include an ablation study on the VQVAE and special tokens to assess their impact.\n\nThe ability to handle downstream tasks is just as important as solving forward PDEs in foundation models. For this, the generalization ability is crucial. Where does the generalization ability of the proposed method come from? Does the proposed quantized representation offer inherent advantages in enhancing the model's generalization?\n\n\n**References**\n\nS. Subramanian et al., Towards Foundation Models for Scientific Machine Learning, 2023.\n\nM. Herde et al., Poseidon: Efficient Foundation Models for PDEs, 2024.", "methods_and_evaluation_criteria": "The training data used in this paper appears simplistic and does not reflect practical scenarios. The PDEs addressed are basic, and even the example of the Burgers' equation includes viscosity, which prevents shock formation, making it a challenging scenario. While these benchmarks are commonly used in the literature, it is important to acknowledge that the problems tackled in this work are not difficult. Additionally, practical data is rarely clean and often contains noise or corruption, as well as irregularly structured spatial data. I recommend testing the robustness of the proposed method in such more realistic and challenging scenarios.", "theoretical_claims": "This paper has no theoretical content.", "experimental_designs_or_analyses": "The paper conducts in-context training based on a ViT architecture, yet there is an existing an in-context model based on ViT (Y. Can et al, VICON: Vision In-Context Operator Networks for Multi-Physics Fluid Dynamics Prediction, 2024.) I recommend comparing performance, memory, and time efficiency between the proposed model and this existing model in a 2D setting.", "relation_to_broader_scientific_literature": "It is unclear how this paper contributes to advancements in both deep learning technology and PDE solvers.", "essential_references_not_discussed": "There are no such essential related works that are missing from the paper.", "other_strengths_and_weaknesses": "Please refer to the above responses.", "other_comments_or_suggestions": "There are instances in the main text where citations have two sets of parentheses, which should be corrected.", "questions_for_authors": "Please refer to the above responses.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper introduces a novel type of transformer for in-context inference of parametric PDEs using  specific pre-training strategies and providing uncertainty estimates. 

While there were some concerns initially about the evaluations and baselines used, most of these it seems were resolved during the rebuttal and discussion period, although it appeared to me not all of them. For instance, ref. 6LCA raised issues about the material presented in Figs. 3 & 4 which I agree with, but the authors didn’t really comment much on this in their rebuttal. Especially Fig. 4, which only shows PCA projections and apparently broad distributions (hard to judge, given lack of axes scales and a clear reference), I don't find convincing.

Also, in my mind some of the baselines used (CODA, Kirchmeyer et al., 2022; CAPE, Takamoto et al. 2023) are somewhat older, given the fast development in this field, and not quite SOTA anymore. I would like to encourage the authors to carefully check the most recent batches of NeurIPS and ICLR papers on meta- & few-shot learning of dynamical systems, and potentially include some of this work at least in the discussion. Finally, one referee stressed the paper may also need a better contextualization within the broader lit. on PDE solvers and better discussion of limitations (e.g., lack of interpretability or error bounds).

In general, however, while one referee remained skeptical and also expressed some concerns about novelty, three referees clearly advocated acceptance.
I will go with the referees’ majority vote and suggest acceptance, but would appreciate if the points above were acknowledged or addressed in the revision.
Decision: Accept (poster)