Zebra: a continuous generative transformer for solving parametric PDEs

Authors: louis  serrano
pierre  erbacher
jean-noël  vittaut
patrick  gallinari
Venue: ICLR_cc_2024_Workshop_AI4DiffEqtnsInSci
Published: 2024-02-10 01:45:50

Abstract: Foundation models have revolutionized deep learning, moving beyond task-specific architectures to versatile models pre-trained using self-supervised learning on extensive datasets. These models have set new benchmarks across domains, including natural language processing, computer vision, and biology, due to their adaptability and state-of-the-art performance on downstream tasks. Yet, for solving PDEs or modeling physical dynamics, the potential of foundation models remains untapped due to the limited scale of existing datasets. This study presents Zebra, a novel generative model that adapts language model techniques to the continuous domain of PDE solutions. Pre-trained on specific PDE families, Zebra excels in dynamic forecasting, surpassing existing neural operators and solvers, and establishes a promising path for foundation models extensively pre-trained on varied PDE scenarios to tackle PDE challenges with scarce data.

Main paper: 
 introduction: 
 Recent advancements in computational science have led to the development of Neural Operators (Lu et al., 2021; Li et al., 2021) and Auto-Regressive Solvers (Brandstetter et al., 2022; Gupta & Brandstetter, 2022; Stachenfeld et al., 2022) , which represent significant progress in creating surrogate models for systems governed by partial differential equations (PDEs). These data-driven models, capable of extrapolating system dynamics from novel initial conditions with limited training data, offer an appealing alternative to traditional computational fluid dynamics (CFD) solvers by reducing the dependency on extensive prior physical knowledge. Despite their promise, these models face a significant challenge: they require retraining to adapt to new or slightly modified PDE parameters, indicating a fundamental gap in their generalizability.One approach to address this challenge involves conditioning models on PDE parameters (Brandstetter et al., 2022; Takamoto et al., 2023; Subramanian et al., 2023) . However, this strategy assumes the availability of accurate PDE parameters for all training and test data. An alternative strategy proposes adapting network weights to each specific PDE "environment" (Kirchmeyer et al., 2022) , but this approach is hampered by complex meta-training optimizations, limiting scalability and necessitating labels for each environment.In a tentative to exploit foundation models, McCabe et al. (2023) explored the feasibility of pretraining large models on diverse PDE parameters using a self-supervised approach. They introduced a video-transformer architecture for multiphysics pretraining, demonstrating effective transferability to closely related PDE parameters. Their focus is on training from multiple heterogeneous dynamics instead of solving parametric PDE as we do here.Our research builds on the foundation laid by self-supervised pretraining successes in both language (Devlin et al., 2018; Radford et al., 2018; 2019; 2021) and vision domains (He et al., 2020; Caron et al., 2021) . Given the sequential nature of time-dependent PDE solutions, we investigate whether a generative pretraining approach, similar to that used in language models, could enhance transferability across dynamical systems. Language models, pre-trained on vast and heterogeneous datasets, utilize a probabilistic framework to manage diversity effectively. However, their application to PDEs input sequence prediction ...Figure 1 : Diagram of Zebra. An encoder compresses the normalized sequence of states, a language model predicts the dynamic and a decoder restores the spatial resolution.has been limited due to challenges in processing continuous data. Recent efforts to improve numerical tokenization in large language models (Golkar et al., 2023; Gruver et al., 2023) address some of these challenges, yet scalability for physical signals remains a concern. Our contributions are as follows:• We present Zebra, a novel generative model that adapts language model architectures to the continuous representations of physical states employing a gaussian mixture component, allowing the generation of continuous trajectory distributions.• We pretrain Zebra and a selection of neural operators/solvers on distinct families of PDEs, evaluating their proficiency in learning diverse dynamics.• Zebra outperforms all baseline models in auto-regressive rollouts for new trajectories within the distribution of PDE parameters. 
problem setting: 
 We explore the potential of neural solvers and neural operators in capturing the dynamics governed by a spectrum of time-dependent parametric PDEs. Our objective would be to train models that can generalize to multiple parameters including initial conditions, boundary conditions, coefficient parameters, forcing terms. However as a first approach, we focus here on solving PDEs with varying coefficient parameters -such as fluid viscosity or advection speed -broadly denoted by c and keep other aspects of the PDEs similar across dynamics. We define F c as the set of PDE solutions corresponding to parameter c. A solution u(x, t) within F c satisfies:EQUATIONB(u)(t, x) = 0 ∀x ∈ ∂Ω, ∀t ∈ (0, T ] (2) u(0, x) = u 0 ∀x ∈ Ω (3)Our goal is to assess the capability of neural architectures to approximate the evolution operator G c : u t → u t+∆t without knowledge of the parameter c. Leveraging a sequence of past states u t-k∆t:t := (u t-k∆t , . . . , u t ) with k ≥ 1, we aim to train models in a self-supervised manner to forecast u t+∆t , thereby pretraining across diverse dynamics. We do not give access to the parameters during training to force the models to learn from data the multiple dynamics as in (McCabe et al., 2023) . In this work, we will consider only 1D equations with one physical channel. 
model: 
 Our model, called Zebra adopts an encode-generate-decode paradigm, but unlike other models, Zebra makes predictions within a continuous probabilistic framework, as pictured in Figure 1 . The architecture is composed of an encoder, a language model and a decoder as described below.Encoder : u 0:k → z 0:k . The role of the encoder is to spatially compress the information by reducing the spatial resolution L of the input function into a latent spatial representation with resolution L ′ < L and channel dimension d. We adopt a variational formulation, therefore the outputs of the encoder are multivariate gaussian parameters E w (u t ) = (µ t , σ t ) whereµ t = (µ t 1 , • • • , µ t L ′ ), σ t = (σ t 1 , • • • , σ t L ′ ), and from which latent codes can be sampled z t ∼ N (µ t , (σ t ) 2 ). To simplify notations, we will omit t and denote u 0:k = u t:t+k∆t . To be able to efficiently represent a sequence of physical states u 0:k , which can be of varying amplitudes and include diverse spectrum evolution, into a sequence of codes z 0:k , we first normalize the input sequence prior to the encoder. We utilize the first observed state in the sequence u 0 to compute mean m 0 and standard deviation s 0 across the spatial dimension and obtain the normalized states as follows: ũi = u i -m 0 s 0 +ϵ , for 0 ≤ i ≤ k. We use a fully convolutional architecture for the encoder.Continuous Language Model: z 0:k → ẑk+1 . The role of the language model is to learn and generate the dynamics auto-regressively. It exhibits two key distinctions with standard language model architectures: (i) it takes continuous inputs and (ii) it outputs a continuous probability distribution thus enabling the modeling of continuous trajectory distributions. Our architecture follows the continuous formulation proposed in Tschannen et al. (2024) . Given a sequence of latent state representations z 0:k = (z 0 , . . . , z k ) with z i ∈ R L ′ ×d , the language model returns a continuous probability distribution over the next token p(z k+1 |z 0:k ). Our approach utilizes an auto-regressive transformer applying causal attention to process the sequence of states. The sequence z 0:k , derived from normalized states ũ0:k is linearly projected to the transformer high-dimensional space. We additionally incorporate a bidirectional attention prefix (Liu et al., 2018) containing the mean m 0 and scale s 0 tokens projected using two-layer MLPs, f m and f s to provide additional context. We use RoPE Su et al. (2021) to encode the relative positions in the sequence.Therefore, we parameterize the next-token probability distribution p(z k+1 |z 0:k , m 0 , s 0 ) with dindependent gaussian mixture models (GMM). By fixing the number of components in the mixture to K, the output of the transformer for each token outside the prefix is a Decoder: ẑk+1 → ûk+1 . The decoder maps the latent tokens ẑk+1 to the physical space ûk+1 = D ψ ( ẑk+1 ). We denormalize with m 0 and s 0 to get the reconstruction.d × K matrix of mix- ture probabilities α next ∈ [0, 1] d×K , mean parameters µ next ∈ ]-∞, +∞[ d×K ,Inference Given an encoded sequence z 0:k , predicting the latent representation z k+1 requires predicting L ′ tokens autoregressively. Each tokenz k+1 l with 1 < l < L ′ is sampled from the K-GMM distribution z k+1 l ∼ GMM(α k+1 l , µ k+1 l , σ k+1 l) and added to the sequence. To unroll the dynamics, we use a rolling window of constant size, meaning that the prediction ûk+1 is made using an initial sequence u 0:k , the following prediction ûk+2 is made using (u 1:k , ûk+1 ).Training We do a two-stage training. We first train the encoder-decoder as a variational autoencoder (VAE) Kingma & Welling (2014) and then train the language model on the latent representations obtained with the encoder. We provide further details on training in the sections below.The encoder and decoder are jointly trained as a VAE. We use a mixed loss objective:L = L recon + β • L KL (4) with β > 0; where L recon = i=k+1 i=0 ∥u i -ûi ∥2 ∥u i ∥2is the relative L2 loss between the input u i and its reconstruction ûi through the encoder-decoder.The KL termL KL = i=k+1 i=t D KL (N (µ i , (σ i ) 2 ) || N (0, I)) regularizes the network.Zebra is trained using self-supervised learning on next token prediction task (Radford et al., 2018) using teacher forcing. Given the sequence z 0:k+1 and normalization scalars m 0 , s 0 , the model is trained to minimize the negative log-likelihood (see Appendix A) expressed below:L nll = - k+1 i=1 log p(z i |z 0 , . . . , z i-1 , m 0 , s 0 ) (5) 
experiments: 
 Datasets We assess the model's performance on 1D data, leaving more complex tests to further investigations. We used the 1D Advection (Advection) and 1D Burgers (Burgers) equation datasets from PDEBench (Takamoto et al., 2022) , training across all available PDE parameters. For Advection, the coefficient parameter that varies across equations is the advection speed β ∈ {0.1, 0.2, 0.4, 0.7, 1.0, 2.0, 4.0, 7.0} and for Burgers it is the fluid viscosity ν ∈ {0.001, 0.002, 0.004, 0.01, 0.02, 0.04, 0.1, 0.2, 0.4, 1.0, 2.0, 4.0}. We evaluate all models on in-distribution parameters. Namely, for each parameter, we have access to 10,000 trajectories, generated with different initial conditions, that we split in 95% for training, and 5% for testing. The trajectories were generated with a spatial resolution of 1024 and with a temporal resolution of 200, which we downsample to 256 and 100 respectively.Baselines We compare our model to current state-of-the-art neural architectures: • Fourier neural operator (FNO) (Li et al., 2021 ) • UNet (Gupta & Brandstetter, 2022) • ResNet (Gupta & Brand- stetter, 2022).Training and evaluation To make a fair comparison, we pretrain all baselines with teacherforcing, i.e. with 1-step prediction by learning to map (u t , . . . , u t+k∆t ) → u t+(k+1)∆t . We fix k = 4 for all models and all datasets. We use the relative L2 loss for training. We evaluate the pretraining in a zero-shot setting on in-distribution new trajectories. We rollout the models for 1, 10, and 25 steps starting from solutions at different times (t = 0, t = 25, t = 50, t = 70) and compute the relative L2 loss with the ground truth trajectories.The results in Figure 2 show that all the models interpolate up to some extent to the PDE parameters distribution. Our model outperforms the baselines on Advection and Burgers. When evaluating the performance at 1, 10, and 25 rollouts, Zebra consistently demonstrates lower error rates, indicating its enhanced ability to accurately predict and model multiple dynamics. Our model has been designed to be easily adaptable to multiple settings such as variable spatial and temporal resolutions, diverse context sizes, etc. Likelihood of K-GMM In the following we write the likelihood of Gaussian mixture model with K components.p(X|θ) = N n=1 K k=1 α k 1 2ασ 2 k exp - (x n -µ k ) 2 2σ 2 k (6) L = log p(X|θ) = N n=1 log K k=1 α k 1 2ασ 2 k exp - (x n -µ k ) 2 2σ 2 k (7)with X being modeled as a mixture of K Gaussian distributions. Each Gaussian distribution k has a mean µ k , variance σ 2 k , and mixture coefficient α kVariational auto-encoder We provide a schematic view of the VAE framework in Figure 3 . We use a spatial encoder and decoder with convolutions only, and normalize the input sequence before the encoder and denormalize to get the final reconstructions. We use normalization statistics computed spatially with the first observation of the sequence. Transformer We illustrate the components of our continuous transformer in Figure 4 . We linearly project the codes obtained with the encoder, and add a prefix with the tokens derived from the normalization statistics. We discard the prefix from the outputs of the causal transformer and then apply a mean head, a scale head, and a mixture head to obtain the GMM coefficients.Language model for images and videos Several works have considered language modeling for image or video generation by combining a VQ-VAE (Oord et al., 2017) with a generative language transformer (Brown et al., 2020) . For images, VQGAN Esser et al. (2021) MaskGIT proposed a bi-directional transformer (Devlin et al. (2018) ) with a scheduling mechanism to accelerate the sampling procedure. Departing from this direction, Tschannen et al. ( 2024) recently introduced a continuous transformer for image generation and assessed its performance with the causal and masked architectures. In video generation, similar techniques relying on quantization have also been introduced, first with VideoGPT (Yan et al., 2021) , and more recently with magvit (Yu et al., 2023a) and magvit2 (Yu et al., 2023b) . The latter, which introduced a new quantization scheme, obtained results on-par with diffusion models. To the best of our knowledge, there is no continuous transformer for video generation.Multiple Physics Recent works have advanced neural solvers and surrogate models for PDEs within Scientific Machine Learning (SciML). Brandstetter et al. (2022) introduced a messagepassing neural solver that includes additional parameter inputs to improve its time-stepping scheme. Takamoto et al. ( 2023) developed a channel attention mechanism, using parameter embeddings of PDE coefficients to boost the generalization of neural surrogate models across different PDE parameters. Subramanian et al. (2023) investigated the application of pre-trained ML models to SciML via transfer learning, finding that fine-tuning these models can significantly reduce the necessity for extensive downstream examples for adaptation. Kirchmeyer et al. (2022) addressed the issue of unknown PDE parameters at the time of inference by adjusting network weights to better suit observed dynamics, using meta-learning to adapt the network across varied environments. Extending the concept of adaptability, McCabe et al. (2023) delved into multiphysics modeling with foundation models. They proposed an architecture to integrate diverse physical signals from different PDEs into a single high-dimensional applying a transformer architecture, similar to those in prediction, to forecast dynamics. Their method, especially effective for 2D equations, highlights the potential of using foundation models in for modeling PDEs and broadly for SciML applications.The Burgers' equation under consideration lacks a forcing term, leading to a rapid decrease in fluid velocity amplitude over time. This aspect makes the dynamics challenging to predict at the early stages of the trajectory for all models, a phenomenon highlighted in Figure 5 . For instance, FNO is not able to unroll the dynamics from t = 0, while it does not diverge when starting from a later timestamp. In contrast, our model is capable of handling this dynamics quite well, as it consistently outperforms other baselines, for all the considered starting times (Figure 6 ). In the experiments, we fixed k = 4, which means the baselines have been trained and tested with an input sequence of size 5. They cannot be used outside of this scope, with less or more context. In contrast, our architecture benefits from the flexibility of transformers and can work with different context lengths. In theory, it should be enough to infer the correct PDE coefficient parameter from the observation: (u t , u t+∆t ) and thus to unroll the correct dynamics. However, in practice this depends on the nature of the equation, and providing more context is useful to decrease the uncertainty over the estimation. Experimentally, we observed that our model performed at the same level for Advection with a context of size 2, better than all the baselines, and that its performance slightly deteriorated for Burgers, though still outperforming other baselines. We show the results in Figure 7 C.3 QUALITATIVE RESULTSWe provide a comparison of the predicted trajectories with Zebra and the ground truth for several test trajectories on Advection with β=2.0 (Figure 8 , Figure 9 , Figure 11 , Figure 10 , Figure 12 ) and Burgers with ν = 0.01 (Figure 13 , Figure 14 , Figure 15 , Figure 16 , Figure 17 ). In both cases, we start from the initial condition at t = 0 and unroll for 25 steps. For clarity, we represent only one timestamp out of 5 for advection.  
 
 Appendix:  


Decision: This study presents Zebra, a novel generative model that adapts language model techniques to the continuous domain of PDE solutions, which can mitigate the limitation of dataset. I believe this paper can bring new insight to the community.  It is expected that authors will be addressing comments by the reviewers in the final draft.

Reviews: review: title: title: A clean and well-motivated paper that would benefit from more extensive experiments and more challenging comparisons.
main_review: main_review: Summary: 
The authors propose to utilize a continuous formulation of generative pre-trained Transformers for PDE forecasting. The continuous formulation is achieved by predicting the parameterizing a Gaussian Mixture Model. Additionally, the computational complexity is reduced by operating in the latent space of a pretrained VAE.
Experimental evaluation is performed on the 1d Burgers and Advection families of PDEs and compares favorably to previously published baseline methods.

Pros:
- Experimental results support consistent improvements over the baselines.
- The framework appears to be of general interest and utility beyond what is shown in this work.
- The text is easy to follow and clearly written.
- The approach is well-motivated from the related work and context of the field.

Cons:
- Technical novelty appears to be low, the authors cite several studies that essentially showed the utility of all components used here.
- Experimental evaluation is performed only on rather simple case studies, precluding any conclusions about the applicability of the framework to more relevant settings.
- The model is only compared against (commonly used) baselines, but not against numerical solvers or more advanced DL methods.
- Performance is only evaluated in terms of L2 loss.

Minor:
- The term "foundation model" is used inflationary here, no scaling properties of the model are shown, nor zero-shot performance on other families of PDE.
- A few language hickups (e.g. 'In a tentative to exploit foundation models')
- Out of curiosity: where does the model name come from?
title: title: Zebra: a continuous generative transformer for solving parametric PDEs
main_review: main_review: Zebra introduces a probabilistic framework for modeling PDEs taking inspiration from the language model community in NLP.


**Strengths**:
- A new approach, from the prism of generative model, to train foundation model for PDEs.
- No explicit incorporation of PDE parameter values in the model.
- Yet another way to perform the *latent evolution* of PDE dynamics.
- Natural support for continuous (a.k.a. discretization-free) modeling.


**Weaknesses**:  
- No separate *validation* set. Hence, the *test* set is somewhat "seen" by the model.
- Generalization ability to unseen parameter distribution (OOD) is not tested.
- No experiments on 2D and 3D PDEs, making the generalizability of the proposed approach to higher dimensions uncertain.

**Questions**:
- For the benefit of the broader SciML community, will the codebase be made public upon acceptance?  

> We pretrain Zebra and a selection of neural operators/solvers on distinct families of PDEs, evaluating their proficiency in learning diverse dynamics.

- Were both the Advection and Burgers' datasets combined and a single Zebra model was trained? If yes, how was the specific PDE identified?

**Suggestions for improvement**:
- Make notations in Figure 1 consistent with the Problem Setting in Section 2 or vice-versa.