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

W4: Number of examples

We experimented with up to 5 context trajectories for Zebra since our focus is specifically on few-shot learning. As for the baselines, gradient based adaptation methods, directly comparable to ZEBRA are known to saturate fast with respect to the number of context examples, see [4] for example (Appendix D.2, Table 5): test errors quickly plateau as the number of context trajectories increases.

W5 : Number of frames

Our primary goal is to evaluate the generalization of neural solvers given minimal observations since our focus is on scarce data settings. We chose a constrained regimes: 2 timestamps for zero-shot and 1 trajectory + 1 condition for one-shot settings. These configurations form the core focus of our paper.

W6 : Datasets

The PDEBench datasets used in MPP exhibit limited parameter variation (around 10), making them unsuitable for our setting. In contrast, our datasets are specifically designed to represent a rich variety of phenomena. For one equation, we considered up to 1200 different parameter values when generating our datasets, compared to the O(10) variations in PDEBench.

W7 : ICON

We tried, but were unable to compare experimentally our method to ICON because it could not handle the spatio-temporal forecasting tasks addressed in our paper. While there are some similarities, the two models target different applications and are built on distinct principles. ICON is an encoder-decoder model relying on cross-attention conditioned on context embeddings, generating outputs pointwise without explicitly modeling process dynamics. In contrast, ZEBRA is a generative model (similar to GPT) employing an encode-process-decode framework. It models dynamics in latent space through its process module. A more precise claim would be that ZEBRA is the first in-context generative model designed for PDEs.

[1] Radford et al, 2018. Improving Language Understanding by Generative Pre-Training.

[2] Ohana et al, Neurips 2024. The Well: a Large-Scale Collection of Diverse Physics Simulations for Machine Learning.

[3] Bonnet et al., Neurips 2022. AirfRANS: High Fidelity Computational Fluid Dynamics Dataset for Approximating Reynolds-Averaged Navier-Stokes Solutions

[4] Kassaï Koupaï et al., Neurips 2024. GEPS: Boosting Generalization in Parametric PDE Neural Solvers through Adaptive Conditioning.

W0: Novelty

We are sorry if there is some confusion here, but MPP and Zebra significantly differ both in their objective and in their modeling framework. As for the objective, MPP adopts the conventional pretraining (on a set of PDEs) and fine-tuning paradigm. The latter requires in particular a significant number of examples.
On our side, we consider solving parametric PDEs with few-shot adaptation. Said otherwise the model is trained to adapt with only a few samples through the mechanism of in-context learning. Additionally, we focus on a scarce data regime for the adaptation. As for the models, MPP relies on a ViT architecture (Dosivitskiy 2020), while Zebra leverages a generative causal architecture inspired from Llama family models.

W1: Distinction with MPP

We believe there might be a misunderstanding: The novelty of our work lies in the in-context learning framework rather than the architecture itself. Furthermore, our model is fundamentally different from MPP. MPP employs a ViT (as noted in Appendix B of the MPP paper), whereas we use a causal architecture in a generative transformer, similar to a language model. Our model operates on discrete tokens, which is why we use a VQVAE. This enables Zebra to generate trajectory distributions—a capability not possible with MPP—and facilitates the use of in-context learning in a manner akin to language models.

W2 : MPP-in-context

Extending MPP for in-context learning is not straightforward. Video transformers have demonstrated limited empirical evidence of in-context learning capabilities compared to language transformers as Zebra. Following your suggestion, we however compared Zebra with MPP trained using in-context examples (referred to as MPP-in-context in the text) in both zero-shot (Table 4, see also below) and one-shot (Table 3, see also below) settings, as detailed in Sections 4.2 and 4.3 of the revised manuscript. Zebra consistently outperforms MPP-in-context in one-shot adaptation and surpasses it on 5 out of 7 datasets in the zero-shot setting using only two frames. Notably, training MPP directly with in-context examples systematically degrades its zero-shot performance compared to vanilla MPP. This suggests that generative pretraining better models dynamics under varying context sizes, aligning with findings in the literature [1].

W2.1 One-shot Adaptation

Conditioning from a similar trajectory.
Test results in relative L2 loss on the trajectory.

W2.2. Zero-shot Prediction from 2 Frames

Conditioning from a trajectory history with 2 frames as input.
Test results in relative L2 loss on the trajectory.

W3 : In-context learning efficiency

First, we emphasize that the zero-shot setting with two frames and the one-shot setting with one trajectory are not directly comparable. Performance varies according to the dataset: Zebra is better in the one-shot setting for Heat, Burgers and Wave boundary than in the zero-shot setting and it is the other way around for Advection, Combined, Wave 2D, and Vorticity 2D. The only comparable settings are presented in Figure 3 and Table 8 in Appendix D.1, where we see that including additional frames and examples consistently reduces error compared to pure zero-shot scenarios.

Regarding the potential of in-context learning for solving PDEs, this is one of the very first works exploring this direction. We agree that demonstrating its full efficiency will likely require additional efforts from the community, as was the case with language models, for example. Specifically, this will necessitate the development and use of appropriate large-scale datasets, which are currently being created by this community [2][3].

I thank the authors for such a great effort in addressing my concerns. I understand the difficulties in rerunning some of the experiments I demanded (eg. Q2), but I hope they can be done for the final version of the paper (please do the same for Q5 as well).

The explanations on superior 2D results are much appreciated, although I am not fully convinced. Thank you for other metrics as well, the quicker inference times boat well for the method's adoption by the ICLR community.

Concerning W3, I would argue that fine-tuning on a new task is a form of adaptation, even thought the model wasn't trained with that in mind. On the plus side, I think in-context learning offers a much broader form of adaptation. It is important that our field clearly separates these paradigms with precise definitions at some point.

Concerning W5, thank you for the added results. However, it is quite a shame that uncertainty quantification is now less emphasized. I really liked the parts about UQ, and I personally found this to be, more than any other contribution, the most useful for the community -- especially since it builds on techniques from LLM prone to hallucinations. I agree with Reviewer k4yp that $\pm 3\sigma*$ is quite arbitrary and that high CLs should not be anyone's objective, but I don't find the aleatoric/epistemic distinction to be particularly helpful, at least not to the point of removing that section entirely. Regardless, I hope a future work will fully investigate this topic.

All this taken into account, I've updated my various scores, particularly the overall score from 5 to 6.

W1: Contribution

We agree, but this is just to clearly indicate the validation setting of the paper.

W2: Referencing

Thank you. We have added the references in Table 1.

W3: Adaptive conditioning

There might be better wordings. By "adaptive conditioning" we want to highlight that the model is trained to adapt to new examples. In Zebra, the model is trained using in-context examples, in CODA and related gradient based adaptation method, the model is trained to adapt through gradient on a new example at inference step.

Fine tuning usually is used for pre-trained models in order to boost performance on a new task. In this setting the model is not trained to adapt.

You are correct, we will correct the sentence tying multi-environments and meta-learning.

W4: Analysis

Yes, this has not been addressed yet due to the page limit. We will revise this section to emphasize the strengths and limitations of each method in both settings.

Thank you, we have corrected the wording.

W5: Uncertainty quantification

We wished to focus here on generalization and in-context conditioning. However, as Zebra is a generative model, we also wanted to illustrate the possibility of quantifying uncertainty of the model's predictions. As this is not the main topic of the paper we have placed it in the appendix section. Addressing the UQ topic would require a dedicated paper in itself.

We thank the reviewer for these references, we have added them in the related work section.

Regarding the discretization from vector quantization, we will explicitly add that we refer to the discretization of values and not in space or time.

Q1 : Training size and OoD

We have conducted an analysis on advection and combined equation to understand the influence of the dataset training size, and show that generalizing even within the training range of parameters is challenging when the number of trajectories is reduced. As you suggested, we performed additional experiments in order to quantify this aspect. Please see Figure 14 in section D.3 for more detailed comments.

We have also performed experiments in an OoD context in order to test the capabilities of the models to extrapolate outside of the training parameters range. The results are now in the new Section 4.4. Overall, in this OoD setting, compared to the baselines, Zebra performs best in the one-shot and zero-shot settings. The gradient-based methods fail to generalize even in the one-shot scenario, for which they were designed. MPP on the other hand performs well in the zero-shot setting. Overall, the OoD setting is challenging for all the methods. They generalize to some extent when the distribution is close to the training one, but severe shifts would require that imply a change in the dynamics cannot be handled.

Q2: Periodic boundary

This is indeed a relevant remark. We initially thought that all models would be able to learn the differences between the Neumann and Dirichlet boundary conditions, however it appears that only the in-context methods are capable of doing so in our setting (MPP-in-context to some extent). We had no time during this period to perform the experiments, because this implies retraining a large number of models but we agree that this should induce better generalization.

W1: In-Context Learning

Sorry if there is a misunderstanding. In-context learning and learning using historical data are two clearly distinct settings. In the former, the model learns from examples provided in the context to perform a similar task with a new initial condition. The context includes multiple trajectories, each corresponding to a different initial condition. In the latter, only a single trajectory is considered, and the task is simply to forecast future states based on the past history of this single trajectory.

W2: Performance Concerns

We respectfully disagree. The datasets cover a wide range of parameter values, capturing diverse dynamics for each equation. This includes variations in coefficients, initial and boundary conditions, and physical variable values. Consequently, the training task is complex, as reflected in the baseline performance. Regarding the performance, for one-shot adaptation (Table 3), Zebra ranks first in 4 out of 7 datasets and second in the remaining 3. The zero-shot results (Table 4) exhibit a similar pattern. Furthermore, the additional out-of-distribution (OoD) experiments provided in the rebuttal show a consistent trend.

W3: Irregular data

This is a very relevant issue. In this study, we focus on generalization with respect to the parameters of the PDE and then considered regular grids for simplicity and for a better control of the experiments and of their conclusions. Our model could be naturally extended in many different ways to handle irregular meshes or even point sets. A good candidate could be using the encoder and decoder from [1] or [2] for example.

W4: OoD

As you suggested, we have provided out-of-distribution results on Heat, Vorticity and Wave2D in the new section 4.4 in Table 5 (see also table below). We tested all the methods - including Zebra - directly without finetuning, on the one-shot and zero-shot setting. In the one-shot setting, Zebra performs the best followed by MPP and then CODA. This suggests stronger generalization properties of in-context learning vs. gradient-based adaptation. Similarly, in the zero-shot setting, Zebra performs best on 2D while MPP achieves the lowest error on Heat equation. Again, CODA suffers from poor generalization in this setting.

---

Test results in relative L2 on the trajectory. -- indicates inference has diverged.
For each dataset:

• One-shot adaptation: Conditioning from a similar trajectory.
• Zero-shot prediction: Conditioning from a trajectory history with 2 frames as input.

[1] Serrano et al. Neurips 2024. AROMA: Preserving Spatial Structure for Latent PDE Modeling with Local Neural Fields

[2] Zhou et al. 2024. Text2PDE: Latent Diffusion Models for Accessible Physics Simulation

W0: Novelty

We would like to highlight that the originality of our framework does not lie in the proposed architecture, which leverages existing modules, although the combination is new for solving PDE problems. The originality of our framework is to propose in-context generative pretraining to allow adaptation without gradient steps for PDE solving. This is why we position our framework versus existing adaptation frameworks (adaptive conditioning and temporal conditioning) and show that our in-context model (i) outperforms SOTA baselines in this context, (ii) does not require gradient-based fine tuning at inference as these baselines do.

W1: Irregular meshes

The problem we address is generalization. We have then focused on regular grids for simplicity. We agree that the model could be extended in many different ways to handle irregular meshes or even point sets. A good candidate could be using the encoder and decoder from [1] or [2] for example.

W2 : Causal mask

We use a transformer based on Llama's architecture, which uses by default a causal attention mask (as in GPT, Gemini, etc.) where the token $s_k$ can only attend previous tokens $s_{< k}$. Note that we employ a causal scheme as we pretrain the transformer to predict the next token, however it could be beneficial to explore other pretraining strategies, such as the one proposed in [3].

W3: Misc.

• Yes, "the same underlying dynamics as the target" means the context trajectory and target trajectory share the same parameters with different initial conditions. We will clarify this wording in the final version. "Predictions are generated using a single sample" means that we predict the trajectory by generating a single sequence of tokens. We generate one trajectory per initial condition.

• Trends on Figure 2. You are right but, Figure 2 illustrates that the context size plays a role in the model performance. This is not the case for example for gradient-based adaptation where the performance saturates fast [4]. We think that this is a distinctive feature that we would like to highlight.

• Figure 3 with CAPE and CODA. We agree, we will try to finish the experiments for the final version.

• We will provide more visualizations in the final version of the manuscript.

• This section aims at a tentative explanation of the interpretability of the generative aspect of our model. To better understand the type of uncertainty of our model, we have conducted an analysis where we provide more or less information in the context at inference. Please see the additional Table 8 of Appendix D.1 in the revised version (see also table below). We have systematically observed a reduction of the standard deviation of the predictions when including more information, which experimentally suggests that some of the uncertainty is epistemic. Although there remains aleatoric uncertainty as our model is not perfect. The ability to predict trajectory distributions is a property of our generative model, but is shown here as an illustration. The topic would require a dedicated analysis which is beyond our scope. We have then moved this section in the appendix (Appendix D.1) with a mention in the main text.

---

Conditioning from a trajectory example and 1 frame or 2 frames as initial conditions.
Metrics include:

• Relative L2 loss (average accuracy).
• Relative standard deviation (average spread around the average prediction).
  The temperature is fixed at 0.1.

---

W4 (minor)

We thank the reviewer for spotting these errors, and have fixed them.

[1] Serrano et al. Neurips 2024. AROMA: Preserving Spatial Structure for Latent PDE Modeling with Local Neural Fields

[2] Zhou et al. 2024. Text2PDE: Latent Diffusion Models for Accessible Physics Simulation

[3] Tian et al. Neurips 2024. Visual Autoregressive Modeling: Scalable Image Generation via Next-Scale Prediction

[4] Kassaï Koupaï et al., Neurips 2024. GEPS: Boosting Generalization in Parametric PDE Neural Solvers through Adaptive Conditioning.

W1: OoD testing

We thank the reviewer for their suggestion on evaluating the models also on out-of-distribution settings. To this end we have included additional results concerning zebra and the main baselines on different scenarios. We have conducted several out-of-distribution tests: (1) forcing terms for heat equation, (2) smaller viscosity for viscosity 2D (3) larger celerity and damping terms for wave 2D. Overall the performance of all the methods degrade in OoD settings compared to in-distribution evaluations, and they extrapolate to parametric regimes that are close to the training distribution. For one-shot evaluation ZEBRA performs best, and for zero-shot it is best on 2 out of 3 evaluations. Please see table below. Note that interpolating in a large range of parameters as we did in our experiments is already a challenging task, that few models are able to handle.

---

Test results in relative L2 on the trajectory. -- indicates inference has diverged.
For each dataset:

• One-shot adaptation: Conditioning from a similar trajectory.
• Zero-shot prediction: Conditioning from a trajectory history with 2 frames as input.

---

W2 : Decoder reconstruction

In most cases, the decoder can reconstruct accurately the frames. The relative L2 reconstruction loss goes from 3e-4 on advection, to 1.7e-2 on vorticity (please see table below). We have provided an additional Table 12 in Appendix D.6, where we display the relative errors between the reconstruction and the ground truth for each dataset used in the paper.

Q1 : ICON

We were not able to compare our method to ICON, because we did not succeed to make it work on our examples. ICON cannot handle the spatio-temporal forecasting task dealt with in the paper. You are correct in pointing out similarities, however the two models target different applications and rely on different principles. ICON is an encoder-decoder model and not a generative mode. It generates outputs using cross-attention conditioned on the context embeddings. The output is generated pointwise. There is no explicit modeling of process dynamics there. ZEBRA is a generative model (like GPT) that relies on an encode-process-decode framework. It learns to model the dynamics in a latent space through the process module. A more precise claim could be that ZEBRA is the first in-context generative model for PDEs.

Q2 : Sobolev norm

We thank the reviewer for this suggestion. Indeed, the sobolev norm is prominent in PDE analysis, and has been proposed for training neural operators recently [1]. However, here we follow the current practice in the literature of reporting the metrics with the L2 norm.

Q3 : VQVAE influence

We thank the reviewer for the suggested analysis. We have conducted an ablation to understand the influence of the codebook size (which plays the role of the number of basis used for projection) on the reconstruction and prediction loss. We also observe a U shaped curve as a function of the size of the codebook. Please find the results in Appendix D.5 of the revised version, in Table 11 and Figure 15.

[1] Cho et al, 2024. Sobolev Training for Operator Learning