Zebra: In-Context Generative Pretraining for Solving Parametric PDEs

1. On the originality

There 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. Objective: 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. Models: 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. Key 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. We acknowledge that solving downstream tasks may be important in certain contexts, but this falls outside the scope of our current work.

2. On the training data

We 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. Our 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.

[1] The well: a large-scale collection of diverse physics simulations for machine learning. Ohana et al. 2024.

3. On VICON

We 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. However, 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. That 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. Finally, we note that this paper was publicly released only in late November 2024. Under ICML guidelines, it therefore qualifies as concurrent work (see below). Excerpt from the reviewer guidelines at https://icml.cc/Conferences/2025/ReviewerInstructions “Concurrent Works Authors 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.”

4. Relation To Broader Scientific Literature

We hope that the clarifications provided help to reassess this statement.

1. Continuous vs Discrete

Thanks 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.

• 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.
• 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.
• 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.
• 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.

[1] Autoregressive Image Generation without Vector Quantization, Li et al. 2024.

2. New forcing terms for Heat

Great point. We added experiments with out-of-distribution (OoD) forcing terms: Gaussian and Square patterns.

• Gaussian:
  We define forcing as a sum of spatial Gaussian pulses modulated in time:

$$f(x, t) = \sum_{i=1}^n A_i \exp\left(-\frac{(x - \mu_i)^2}{2\sigma_i^2} \right) \sin(\omega_i t)$$

This introduces smooth but localized excitation. Solutions resemble the homogeneous case in regions far from peaks, making this a "close" OoD test.

• Square:
  We define forcing as a sum of square waves:

$$f(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)$$

This preserves parameterization but introduces discontinuities and high-frequency content, making it a more challenging OoD case.

We evaluate both without retraining, using context examples from these new trajectories. Results below compare ViT-in-context and Zebra:

While both models show performance degradation, Zebra remains significantly more robust than the baseline.

3. New trajectory generation

Our 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.

4. Experiments from Appendix D.

Thanks for the suggestion. We agree and will reference these findings directly in the main text.

5. Generative models

Thanks, we will include these references in a section on SOTA generative PDE simulators.

6. Encoder-decoder analysis

We 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.

7. Plot differences

Agreed, we will add additional plots to show the differences between ground truths and predictions.

Q1. # of environments.

Actually, this was arbitrary. We can align all test settings to 120 environments for the camera-ready version; this will not significantly change the results.

1. DPOT

Thank 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.

2. Explanations and Limitations

Thanks 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.

As 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.

3. Computational Resources

Thank 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.

4. Notations

We agree. In particular, distinguishing between flattened and non-flattened indices can be confusing. We will revise and simplify the notations accordingly.

5. Complexity

This 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.

Q1: VQVAE

We 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.

As for the behavior for turbulent systems, we performed OoD experiments for the vorticity equation. Here are the results:

As 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.

Q2: Memory Requirements

The 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.

Q3: Number of Examples

Great 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:

1. On the uncertainty

Thank 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.

The 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.

The 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).

Although 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:

• CRPS and RMSCE Results

Below, CRPS and RMSCE over time for Zebra on Burgers. As expected, CRPS grows with error accumulation. RMSCE remains stable, indicating consistent calibration.

2. New trajectory generation

This 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.

3. Clarifying the scope

We will clarify this in the introduction. Our focus on forecasting from an initial state highlights the role of context.

4. "provides richer information..."

Agreed. We will remove the sentence.

5. On scaling gradient-based adaptation

Our 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.

6. Index representations

We 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.

7. On the main figure

We only show final timesteps, which may seem unrelated. This will be clarified in the caption. See Appendix E.6 for better illustrations.

8. Caption revision

The 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.

9. "Sequences"

We will clarify this. Each $s_i$ is a sequence of VQ-VAE indices.

10. <eot> token

We will try to include it, possibly with smaller font.

11. Resolution

We used 64×64 in 2D and 256 in 1D.

12. "Blurry" expectations

We will remove this term if it’s unclear. We meant deterministic models tend to predict conditional expectations, which average over possibilities.

13. Fig. 2 issue

We will fix it.

14. Different dynamics

Each 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.

15. ICL training

We 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.