Masked Autoencoders are PDE Learners
Authors: Anthony Zhou
Amir Barati Farimani
Venue: Accepted by TMLR ()

Abstract: Neural solvers for partial differential equations (PDEs) have great potential to generate fast and accurate physics solutions, yet their practicality is currently limited by their generalizability. PDEs evolve over broad scales and exhibit diverse behaviors; predicting these phenomena will require learning representations across a wide variety of inputs which may encompass different coefficients, boundary conditions, resolutions, or even equations. As a step towards generalizable PDE modeling, we adapt masked pretraining for physics problems. Through self-supervised learning across PDEs, masked autoencoders can consolidate heterogeneous physics to learn rich latent representations. We show that learned representations can generalize to a limited set of unseen equations or parameters and are meaningful enough to regress PDE coefficients or the classify PDE features. Furthermore, conditioning neural solvers on learned latent representations can improve time-stepping and super-resolution performance across a variety of coefficients, discretizations, or boundary conditions, as well as on certain unseen PDEs. We hope that masked pretraining can emerge as a unifying method across large, unlabeled, and heterogeneous datasets to learn latent physics at scale.

Published in Transactions on Machine Learning Research (12/2024)
Masked Autoencoders are PDE Learners
Anthony Zhou
ayz29andrew. edu
Deportment of Mechanical Engineering Carnegie Mellon University
Amir Barati Farimani
barati@emu..cm edu
Department of Mechanical Engineering
Department of Machine Learning Carnegie Mellon University
Reviewed on OpenReview: /ttps://openreett opw.//vee.....ee..o net/ forun???? id= FdNNuuFFuUV
Abstract
Neural solvers for partial differential equations (PDEs) have great potential to generate fast and accurate physics solutions, yet their practicality currently limited by their generalizabil- ity. PDEs evolve over broad scales and exhibit diverse behaviors; predicting these phenomena will require learning representations across a wide variety of inputs which may encompass different coefficients, boundary conditions, resolutions, even equations. As a step towards generalizable PDE modeling, we adapt masked pretraining for physics problems. Through self-supervised learning across PDEs, masked autoencoders can consolidate heterogeneous physics to learn rich latent representations. We show that learned representations can generalize to limited set of unseen equations or parameters and are meaningful enough to regress PDE cocfficients or the classify PDE features. Furthermore, conditioning neural solvers on learned latent representations can improve time-stepping and super-resolution performance across variety of coefficients, discretizations, or boundary conditions, well as on certain unseen PDEs. We hope that masked pretraining can emerge as unifying method across large, unlabeled, and heterogeneous datasets to learn latent physics at scale.
1 Introduction
The physical world is incredibly complex; physical phenomena can be extremely diverse and span wide spatiotemporal scales-fron from neuron excitations to turbulent flow even global climate. Importantly, many of these phenomena can be mathematically modeled with time-dependent partial differential equations (PDEs) (FitzHugh 1961 Nagumo al. 1962 Lorenz 1963). These PDEs are generally analytically intractable and require the use of numerical solvers to obtain approximate solutions. For complex physics, these solutions can often be slow to obtain; furthermore, different PDEs often require careful design of tailored solvers.
Advances in deep learning have led to the design a new class of solvers for PDEs. These neural solvers can be extremely fast and display resolution invariance; however, neural networks introduce training difficulties and a lack of theoretical guarantces. Many important advances have been made address these challenges, with models becoming faster than numerical solvers within well-studied PDEs under certain setups, proposing
error bounds, and being extended to solve real-world problems. (Raissi et al. 2019 Lu et al. 2019 Li et al.
2020
Cao 2021 Brandstetter et al. 2022 Li et al. 2023 Kovachki et al. 2021 Li et al. 2024
A current frontier in neural PDE solvers lies in generalizing solvers to different parameters, conditions, or equations, thereby avoiding the need to collect new data and retrain networks when given unseen PDE dynamics. Prior work in this space has explored many methods to achieve this, from directly conditioning on
PDE coefficients TTakamoto al. 2023 Lorsung al. 2024 Shen et al. 2024 pretraining foundation
models across various equations (Subramanian et al. 2023 McCabe al.
2023 Hao al. 2024). Despite

---

Published in Transactions on Machine Learning Research (12/2024)
these advances, generalizable neural solvers remain significant challenge. PDEs can be incredibly diverse, and neural solvers must adapt to different coefficients, geometries, discretizations, or boundary conditions.
As a step towards addressing generalizability, we propose adapting masked pretraining methods to the PDE domain. This is motivated by the observation that masked pretraining can learn highly meaningful and broad knowledge in the computer vision and language domains (Devlin et al.] 2018). In addition, masked modeling approaches are known to scale well to large and diverse datasets (He et al. 2021). In practice, masked pretraining can also be easily implemented since it makes prior assumptions about the PDE data, and modern training pipelines and masked modeling architectures are efficient (Vaswani et al. 2023) Lastly, it is hypothesized that masked modeling approaches can generalize well duc to having very low inductive bias from masking large portions of the data (Feichtenhofer et al.).
As such, to study masked modeling within the PDE domain, we train masked autoencoders on a diverse set of 1D and 2D PDE data and evaluate their learned representations. We demonstrate that self-supervised masked pretraining can learn latent structure that can express different coefficients, discretizations, boundary conditions or PDEs under a common representation. Furthermore, we show that masked autoencoders (MAEs) can learn highly structured latent spaces through masking alone. MAE models can be used to improve downstream tasks such as predicting PDE features or guiding neural solvers in time-stepping or super-resolution through providing meaningful context. Our contributions suggest the possibility to transfer the scalability and flexibility of masked modeling from language and vision domains to physics creating rich, unified representations of diverse physics through self-supervised learning. We provide the code and datasets used in this study here: httpp://gitth: thub. comb saatthooyyhhou-mmm--ppoo
2 Related Work
Neural PDE Solvers The field of neural PDE solvers has grown rapidly and has shown great advceces in both the accuracy of solutions and the ability to adapt to PDE parameters. Infiniteddimensional neural operators (Li et al., 2020 Kovachki et al. 2023 Lu et al., 2019) have shown impressive accuracy in solving Lime-dependent PDEs by learning the mappings between initial conditions and solutions. However, these methods alone have shown brittleness with respect to changing PDE coefficients or boundary conditions (Gupta and Brandstetter 2022 Lu et al. 2021), prompting recent work to allow neural solvers to adapt to different PDE conditions.
A variety of approaches have considered adding PDE dynamics information or time-dependent trends to neural solvers. Common neural solvers can support conditional prediction through architecture choices (Gupta and Brandstetter 2022), and novel architectures can be designed to explicitly operate with PDE parameter knowledge (Brandstetter et al. 2022). Beyond directly conditioning on PDE dynamics, a class of neural PDE solvers has proposed the addition of an encoder or adaptive network to inform a forecaster network of different PDE coelficients (Wang et al.] 2021 Kirchmeyer el al. Takamoto et al. 2023 Lorsung et al.] (024).. At an even higher level, meta-learning approaches have been adapted PDE learning to maximize shared learning across different physics (Yin ct al. 2021 Zhang et al. 2023).
As an effort to work towards more generalizable PDE neural solvers, recent work has followed the success of pretraining and foundational models in the broader deep learning community.
Pretraining for PDEs
Based on contrastive pretraining methods in computer vision problems, (Chen et al. 2020 Schroff et al. 2015:
Zbontar et al., 2021 Bardes et al. 2022), ; contrastive PDE methods aim to leverage equation coefficients
Lorsung and Farimani| (024). physical invariances (Zhang et al.) 2023), or Lic point symmetrics MMialon et al. 2023 Brandstetter al. 2022) to define differences in PDE dynamics that can be organized in latent space. Another approach in PDE pretraining follows observed in-context learning and emergent behavior in
LLMs Wei et al. 2022 Brown et al. 2020; Radford et al.p to design neural PDE solvers that are capable of
following prompted PDE examples to forecast unseen dynamics (Yang et al. 2023 Chen et al. 2024h.
A more straightforward pretraining method focuses on directly training neural solvers to transfer to new PDE dynamics (Goswami et al. 2022 Chakraborty et al. 2022 Wang et al. 2022). This approach has also been scaled by training neural solvers with large and diverse training sets to characterize its transfer behavior (Subramanian et al. 2023), as well as shown to be generally more effective over other pretraining
2

---

Published in Transactions on Machine Learning Research (12/2024)
strategies (Zhou et al. 2024). As a step toward large-scale modeling, more principled training approaches have been proposed to learn PDE dynamics across diverse physics at scale. Recent work has proposed a combinatorial neural operator that learns different dynamics as separate modules (Tripura and Chakraborty
2023 , embedding separate PDEs to a common space to do multi-physics prediction (McCabe et al. 2023
incorporating denoising with scalable transformer architectures ures while training across diverse PDE da dalasets
(Hao al. 2024), and using a unified PDE embedding to align LLMs across PDE families (Shen al. 2024
Masked Pretraining Masked reconstruction is a popular technique popularized by the language pro- oessing (Devlin ct a 2018) and vision (Dosovitskly ct al. 2020 Xie et al. 2021 He al. 2021 domains pretrain models for downstream tasks, Masked modeling is a broad field that spans many masking strategies, architectures, and applications (Li et al., 2024); this ubiquity is attributed to scalability and architecture breakthroughs /aaswani et al. 2023) that allow meaningful context to be learned through masked reconstruction (Cao et al. 2022). In the field of neural PDE solvers, masked pretraining has initially been explored as a hethod to prefrain neural solvers directly (Chen et al., 2024). However, we take a separate approach by investigating if models can understand physics through masked self-supervised learning and how these latent representations are manifested. After validating this learning method, we seek to understand if this knowledge can be used to benefit common PDE tasks.
Situating our Contribution To frame our contribution within these past works, we draw comparisons with broader deep learning efforts to pretrain large encoders through self-supervision, such as BERT (Devlin let al.] 2018) or CLIP (Radford et al., 2021 Ramesh et al., 2002),, to be used in downstream tasks. For many reasons (e.g. data availability, architectures, pretraining strategies, compute resources), an equivalent work does not currently exist for PDEs. However, we hope to advance this research thrust to train general-purpor physics encoders, which could be of benefit to the PDE community. Specifically, pretrained encoders are often leveraged to accelerate model development and training by outsourcing certain architecture modules lo pretrained models. Furthermore, having standardized, encoded representations of diverse PDEs can help in interfacing with conventional ML pipclines, whether be allowing L.Ms to interface with PDE data or diffusion backbones to condition on PDE inputs.
3 Methods
We describe our main methods in Figure 1 We first pretrain an encoder and decoder to reconstruct masked inputs. This pretraining objective can have a few benefits. Masking destroys inherent biases present in the data; especially at high masking ratios, this forces models to learn very general latent representations. Lastly, masked modeling does not assume any prior physics or leverage features specific to a particular PDE, making the approach applicable to a wide varicty of physics.
A After pretraining, we evaluate using the learned representations for three downstream tasks after discarding the decoder. The first task is to regress or classify PDE features, such predicting coefficient values boundary conditions of a PDE sample. This can be done by directly appending a linear head to the pretrained encoder. The second task is to improve the performance of neural solvers by conditioning on an encoded representation of input physics. Both the encoded and original representations of the data are passed to the neural solver. Lastly, we consider improving super-resolution performance by passing an up-sampled input to a neural solver and conditioning on an encoded representation of the low-resolution PDE input.
3.1 Masked Pretraining for PDEs
We adapt the popular Masked Autoencoder (MAE) ct al.) to train VIT models (Dosovitskiy et al.) 2020 Arnab et al.) 2021 to reconstruct 1D and 2D PDE data. Data are partitioned into non-overlapping patches before a random subset of these patches is sampled to be masked. The masked patches are omitted from the encoder input; the encoder embeds only the unmasked patches through scries of transformer blocks, which allows for larger encoders and faster training at large masking ratios He et al. 2021). The encoded patches are then recombined with mask tokens according their position in the PDE solution. Positional embeddings are added again to preserve positional information
approach (He et al. 2021 Xie et al. 2021 Feichtenhofer
3

---

Published in Transactions on Machine Learning Research (12/2024)
Masked Protraining
Supervised Fine-tuning
Label v=l
Conditiana) Time-stepping
Cancitioml Super-resalution
FRD
Figure We investigate learning diverse PDE dynamics with masked autoencoders (MAE) and using learned representations benefit various downstream tasks. (Masked Pretraining) An encoder trained on unmaskod patches of spatiotemporal PDE data, while a docoder reconstructs true data from latent encodings and learned mask tokens. (Supervised Pine-tuning) Pretrained encoders can be used to quickly regress equation coefficients or predict key PDE features. (Conditional Time-stepping) Neural solvers achieve higher accuracy predictions through conditioning on MAE encodings. (Conditional Super-resolttion) SR models can also benefit from conditioning on MAE encodings, using a discretization inversion (D-1) and neural operator (Yang et al. 2023 to predict high-resolution physics.
before the latent encodings and mask tokens are decoded. The decoder can be shallower and narrower than the encoder because discarded in downstream tasks (He et al. 2021). which further reduces training costs. The decoded tokens are projected into the PDE space through a linear layer before reshaping the patches to the input dimension. Lastly, the output is compared ground truth PDE data through MSE loss.
Within this general framework, the primary considerations affecting performance and compute model size, masking ratio, and patch size. We study the effects of these parameters the resulting model's reconstruction performance. In general, we find that incrcasing model size requires more data to be effective. Furthermore, increased masking ratios decrease the reconstruction performance but may lead to better learned representations. Lastly, decreasing patch size improves reconstruction performance at the cost of increased compute. For additional details and hyperparameters we direct readers to Appendix
3.2 Lie Point Symmetry Data Augmentation
Table 1: Error at different
To emulate a larger pretraining dataset and aid in generalization, we inves- tigate using Lie angmentations during pretraining, an approach that follows the use of data augmentations vision or video pretraining (He et al.) 2021 Xie et al. 2021 Feichtenhofer et al.). Given a PDE, one can derive its Lie symmetries set of transformations (gi.....gh. 9.),, each with variable e; that modulates the magnitude of the transformation. At training time, apply sequentially, each with a randomly sampled to randomly augment PDE samples. This augmented PDE sample could represent a solution that has been shifted in space, time, or magnitude, among other transformations, but still propagates dynamics according the original PDE. For more detailed discussion of Lie point symmetries for PDEs, we refer the reader to Olver 1986) and Mialon al. (2023). To understand the effects of Lie data augmentation reconstruction accuracy, we apply randomly
augmentation probabilities
Aug. Prob.
Error
0.00
2.48o-03
0.25
1.24e-03
0.50
1.17e-03
0.75
1.27e-03
1.00
1.37c-03
4

---

Published in Transactions on Machine Learning Research (12/2024)
augment training samples with different probabilities, with 0 representing augmentation and representing always augmenting samples. The resulting validation MSE reconstruction losses are given in Table after training on ID KdV-Burgers data. We find that over-augmenting samples may shift the training distribution too much; furthermore, we generally find that the best results use smaller augmentation magnitudos c.
Multi-Rcsolution Pretraining PDE data are often discretized on a mesh, which can be both irregular and of different resolutions. This can present a challenge for vanilla transformers, which use the same positional encodings regardless of variations in imput sequence length. This is especially pronounced in 2D; flattening two spatial dimensions results in positional information needing to wrap around inputs, and as a result, the same positional embedding can represent very different points in space when given different discretizations. Indeed, many approaches have been proposed to adapl ViT models to inpuls of varying resolutions or patch sizes (Beyer et al. 2023 Tian et al.] 2023 Dehghani et al., 2023 Fan et al., 2024). However, to avoid changes to the ViT backbone, we consider different strategics to accommodate varying input lengths. In particular, we consider the setup where imputs of shape (z4) or (mt f2y myy) can vary along their spatial dimension (nz) by multiples of the patch size (Pa,Py). Snapshots in time, or time windows, are sampled with length ne, which discretizes the PDE trajectories to common time dimension.
The simplest strategy would be to pad the imputs to the maximum sequence length (Pad). However, this does not address the fact that positional information becomes overloaded and only provides information about input length. To address this, we consider interpolating positional embeddings from the maximum sequence length to variable sequence lengths (Interp)). This would allow tokens to receive approximate positions; however, this would not be accurate for irregular grids. Lastly, to adapt to arbitrary grids, consider using an embedder network, which could be a 1D or 2D CNN (depending on the spatial dimension), to project the spatial coordinate grid to a token (Token). We compare the validation MSE reconstruction loss of using different strategies Table 2 after training on multiresolution KdV-Burgers data, observing consistent improvement by using more sophisticated mot for embedding spatial discretizations. In ID, multiresolution pretraining is significantly easier to learn, so we use the Interp. strategy for simplicity. In 2D, the use of tokenized spatial grid information becomes more important, we use the Token strategy for 2D multiresolution experiments.
Table 2: Error w/ differ- ent resolution strategies
Strategy
Error
None
4.40e-03
Pad
3.81e-03
Interp.
3.23e-03
Token
1.89e-03
4
Experimental Setup
4.1 PDEs and Datasets
We describe a variety of PDEs used for masked pretraining and downstream evaluation. In 1D, we pretrain MAE models on the KdV-Burgers equation only, while in 2D we pretrain on the Heat, Advection, and Burgers equations simultaneously. In all PDEs, coefficients and forcing terms are randomly sampled to produce diverse dynamics within dataset.
ID KdV-Burgers Equation: The KdV-Burgers equation (KdV-B) contains the Heat, Burgers, Kdv equations corner cases modulated by coefficients (,, (Brandstetter al. 2022: Jeffrey and Mohamad 1991)
(KdV-B) Both initial conditions and forcing terms are generated from the periodic function 5(t,z), z)), where we uniformly sample Aj E [_0..50.55, E [-0.4,0.4] 0.4), E (1,,, (2,,) Φb E [0, 2m) while fixing J = 5, T. 16.
0,u + aui,,u 30 + 700zzu = 5(t,
J Σ + oz). 2(0, x) = 5(0,z)
8(t, x)
(1)
To generate diverse PDE samples, the coefficients are sampled uniformly from E 3), [0, 0., [0. Furthermore, samples are generated with discretization (ne, (z)) = (250, 100) on an interval x [0,16)
5

---

Published in Transactions on Machine Learning Research (12/2024)
from t 0 to t = 2. For 1D multi-resolution experiments, PDE data from the KdV-Burgers equation is downsampled to variable spatial resolutions.
ID Heat and Burgers Equations: are used to evaluate performance on PDEs that are a subset of pretraining samples. Furthermore, to evaluate extrapolation to unseen boundary conditions (BCs), samples of the Heat equation are also generated with Dirichlet and Neumann BCs in addition to periodic BCs.
The 1D Heat and inviscid Burgers (Bandstetter et al. 2022) equations
0,u vd = 8(1, x)
(Heat) (Burgers)
0, ud,u ((t,
We solve the equations with the same periodic BCs, initial conditions, and forcing function setup and as the 1D KdV-Burgers equation, but by selting the appropriale coefficient values. Specifically, we uniformly sample 3 E [0.1, 0.8] for the Heat equation and fix a 0.5, 3 to model the inviscid Burgers equation. These equations are also solved with discretization (ne, n22) = (250, 100) on an interval = [0, 16] from 0 to = 2. To generate data for the Heat equation that enforces Dirichlet Neumann boundary conditions, we write the Heat equation in its variational form after an implicit Fuler discretization:
J (uv + AvV Vojda= = [( (1**+ A1 A 1idt
(2)
This formulation can be solved using FEniCS (Alnaes et al. 2015 Logg al.
2012 To simplify the boundary
value problem, we set the forcing term ʃ f 0. Furthermore,
set [u(r 0) u(x L) = 0] to enforce
Dirichlet BCs, and [0,u(z= = 0) = 0zu(z = L) = 0] to enforce Neumann BCs; however, in both of these cases, the initial conditions in Equation need to be modified to respect the new BCs.
J A 2,sin(2xlxx/LI+
m(0,x) x)
Σ
Dirichlet ICs
(3)
j=1
x(0, = Σ A,cos(2ml,a//1)) 65), Neumann ICs
(4)
j=1
In both oquations, we uniformly sample A, E (2-.0.5..55, E (1,2,3) o, E (0, while fixing J = 5, L = 16.
1D Advection, Wave, and Srraamoto-Sivashnssky Equations: The Advection (Adv). Wave, and parameter- -ependett Saaamoto-Sivashnssye (KS) (Lippe et al. (2023) equations are considered to eval- uate downstream performance to new equations; the equations contain PDE terms that are unseen during pretraining. Additionally, the Wave equation is generated with Dirichlet and Neumann BCs (Brandstetter
et al. 2022 to evaluate unseen BCs novel PDE dynamics.
O.u cπ,u
(Adv)
Ouu 2² Ozz =0
(Wave)
0.u + uZ.u + vd..u Oxxzzu=
(KS)
For the ID Advection equation, initial conditions are generated according to Equation [ the wave speed is uniformly sampled from E [0.1, 3]. and periodic BCs are used. The solution domain and discretization are the same as previous cases, with (n4, nx)) (250, 100), = [0, 16), and time ranging from 1 = to 2.
For the 1D Wave equation, we solve with Dirichlet (u(x = 0) = u(a = L) = 0) and Neumann (Q.u(z = 0) = 0_u(r= L) 0) BCs, resulting in waves that either bounce or reflect off boundarics. The wave speed is fixed al = 2, and the initial condition is Gaussian pulse with unit amplitude and with its peak randomly sampled the spatial domain. Lastly, the equation is solved from t to 100 the interval x [-8, with discretization (ne, nz) (250, 100).
For the 1D KS equation, we use periodic BCs with initial conditions from Equation Following the data setup proposed by Lippe et al. 2023), additionally uniformly sample [0.75, 1.25 to vary the second-order

---

Published in Transactions on Machine Learning Research (12/2024)
term in the KS equation. Furthermore, due to the unique dynamics of the KS equation, solve the PDE from t= 0 to / = 100 on the interval [0, 64) with a discretization of (nt, nz) = (100, 100).
2D Heat, Advection, and Burgers Equations: The 2D Heat, Advection (Adv), and scalar Burgers (Rosofsky et al. 2023) equations are considered for both pretraining and downstream evaluation. For 2D) multi-resolution experiments, data from these equations are downsampled to variable spatial resolutions.
0,u - vv =0
(Heat)
0 Vu = 0
(Adv)
duu u(c Vu) ~V
(Burgers)
For the Heat equation, we uniformly sample the [2 10-3, 10-2, for the Advection equalion, uniformly sample
Cy] [0.1, and for the Burgers equation, uniformly sample e [7.5 x 10-3 1.5 x 10-2, and c c= [Ca, (..). For all equations, we use periodic BCs and solve on a grid (Mi, feg ny) = (100, 64, 64) on solution domain (æ, = [-1,1] from = to L = 2. Lastly, initial conditions are generated from:
w(0,x, x. v) = Σ A (ysin(222zzz - 2 xj)
(5)
j=1
Initial condition parameters are uniformly sampled from Ag E (-0.5,05),. E (-0.4,04)), E (1,2,3),4, lus E (1,2,3), Φ [0, 2x) while fixing 5, L =
2D Navier-Stokes Equations: Following the setup from Li et al. (2020), we consider the incompressible Navicr-Stokes (NS) equations in vorticity form, but randomly sample the viscosity v and forcing function (() amplitude. To ensure consistency with the pretraining dataset, our experiments model NS dynamics as scalar vorticity field; from this the velocity field can be derived from the Biot-Savart Law.
Oyw + u Vw wV V-u=0
(NS)
The solution is solved on grid (nt, nay nyy = (100, 64, on a solution domain (x,y) = 10, 1 from 25. PDE parameters are uniformly sampled from v {[1,2,3, 4,5, 7, x 10 (6.7,8.9) and
A E 2,, 4,5, 6, 7,,,8,9, 10) x 10-3, Lastly, initial conditions we sampled according to Ti et al. 2020 from Gaussian random field.
4.2 Data Augmentations
We implement Lie Point Symmetry Data Augmentations according to Braudstetter et al. (2022), including shifting and resampling PDE solutions with the Fourier shift theorem. Since we only augment PDE samples during pretraining, we consider symmetry groups for the 1D KdV-Burgers equation, as well as the 2D Heat, Advection, and Burgers equations. The 1D KdV-Burgers equation has the following Lie subalgebras [bragimov 1993):
a
a
X1
X2
X3 at dr
(6)
Dt
Or'
Tu
Taking the exponential map results in the following Lic groups [bragimov
1993
(i(e)(z,t, = (I,t+ed e. (,())(z,t,u) (z t,t, (2, (x el, L, +e), (Galilean Boost)
(Time Shift)
(7)
(Space Shift)
(8)
(9)
For the 2D Heat, Advection, and Burgers equations, there are many possible Lie subalgebras Tbragimov 1993). For simplicity, we only consider basic subset of these that apply all three equations, however, there is ample room to implement more symmetries:
7

---

Published in Transactions on Machine Learning Research (12/2024)
0
8
X1
X2
X3 = Oy
(10)
Ox:
These result in the following Lie groups in 2D:
9 (e))) (,, (,,,,te,, e, = (,,,
(Time Shift)
(11)
(X Shift)
(12)
(s(0)(,,u,1,)) (x,++at,u),
(Y Shift)
(13)
5 Results
In this section, we provide the results of different experiments designed to understand the capabilities of masked autoencoders. Firstly, we would like to understand the reconstruction capabilities of the masked autoencoder, or if the pretraining goal given by the masked objective is being met. Once MAE performance is validated on its pretraining objective, we seek to understand and visualize the representations learned by the MAE during pretraining. This can be done by projecting the latent embeddings to lower dimension to visualize qualitative trends. Although insightful, another more rigorous evaluation is the regression/clasts classificaisninni of physical variables. These are easy quantities to derive and are usually trivially known priori, however, they can serve as probe to quantitatively gauge model knowledge rather than rely on qualitative latent trends. This is analogous to a linear probe used to predict image rotations or grayscale learning for computer vision (Chen et al. 2020 ; indeed, coefficient regression has been used in prior PDE literature to gange model performance after self-supervised pretraining (Mialon et al. 2023
color in self-supervised
While these three tasks (masked reconstruction, latent visualization, variable pearession/clasificaaa provide knowledge, they are generally not useful on their own. To extend masked pretraining to practical tasks, we consider using a pretrained encoder to improve PDE time-stepping or super-resolution In particular, we are interested in whether the representations learned during masked pretraining can help in diverse physics scenarios by providing additional context to neural PDE solvers during time-stepping or super-resolution.
can
5.1 MAE Pretraining
MAE models are trained on 10000 samples of 1D KdV-Burgers PDE data in 1D and 12288 samples of 2D Heat, Advection, and Burgers PDE data in 2D. We display example results from masked pretraining in Figures 2 and A notable difference from vision and video domains is that physics does not follow human-recognizable structure or descriptions (c.g. backgrounds, actions, faces, shapes, etc.); furthermore, in addition to the overall meaning, the numerical accuracy of the reconstruction important. Despite this, MAE models are able to capture underlying physics and reconstruct PDFs within the pretraining set woth-both in 1D and 2D, and high masking ratios (75% and 90%). In general, for PDEs that are similar to those seen during pretraining, such as the 1D Heat inviscid Burgers equation, MAE models tend to interpolate well Furthermore, given information about the spatial discret ization, MAE models can adapt to different PDE resolutions when trained to reconstruct multi-resolution inputs. This is true in 2D as well, with example results shown in Appendix B.2.
For PDFs that contain novel equation terms or boundary conditions (BCs), the MAE extrapolation per- formance is limited. Zero-shot reconstruction of the 1D Advection and KS cquations shows mild trends, while reconstruction of the 2D Navier-Stokes equations is ineffective. To address this gap and investigate whether MAE models can perform complex high-resolution physics with multiple variables, MAE model to reconstruct pressure and velocity on 2D smoke buoyancy data with varying buoyancy factors (Gupta and Brandstetter 2022) and qualitatively show its reconstruction in Appendix B.3 In general, MAE models can adapt. to complex scenarios and multiple physical variables; however, many of the fine details (e.g. eddies, shedding) become lost at high masking ratios.
train an
Since the 1D KdV-Burge equation is solved with periodic BCs, we evaluate MAE extrapolation to Dirichlet and Neumann BCs when reconstructing the 1D Heat and Wave equations, shown in Appendix B.1 Overall
8

---

Published in Transactions on Machine Learning Research (12/2024)
Koy-Rurgers |Eretraining
Pulti-Resolutioo
Burgers
Advection
Figure 2: Example results after training on the 1D KdV-Burge equation with a masking ratio of 75%. For each triplet, we show the masked PDE (left), the MAE reconstruction (middle), and the ground-truth (right), and plot space and time on the x and y axes respectively. The MAE can reconstruct multiple resolutions of KdV-Burgers data and interpolate to the ID Heat and inviscid Burgers equations. For the 1D Advection and KS equations, which contain novel PDE terms (Ugg Urzwe), the extrapolation performance is limited.
Heat
Adsection
Burgers
X5
Figure Example results after training on combined set of the 2D Heat, Advection, and Burgers equations with masking ratio of 90%. We show the masked PDE (top), the MAE reconstruction (middle), and the ground truth (bottom), plotting space on the z-y axis at multipic snapshots in time. Despite good performance within the training set, the model is unable to extrapolate to the Navier-Stokes (NS) equations, which contain novel initial conditions, forcing terms, and dynamics.
trends remain consistent; the Heat equation, being more similar to the KuV-Burgers equation, shows limited roconstruction performance when extrapolating to novel BCs, while the Wave equation introduces novel PDE term (we)) and initial condition (Gaussian pulse), and result the zero-shot reconstruction poor.
5.2 Latent Space Evaluation
To better understand the latent representation learned by masked pretraining, use the MAE encoder to embed various PDE validation samples and visualize embeddings with t-SNE (van der Maaten and Hinton 2008) in Figure 4 Through self-supervised learning, MAE models can learn trends in PDEs without labeled data and with limited extrapolation abilities. For example, through masked reconstruction of the 1D KdV-Burgers equation, MAE models can learn coefficient-dependent trends (Figure AA). This can be applied PDEs not seen during training, as the same MAE model can distinguish samples from the Advection equation based wave speed (Figure 4B). This extrapolation is also observed when embedding samples from different PDEs; representations learned from pretraining the KdV-Burgers equations allow models to cluster PDE samples from the Heat, Burgers, Advection, and KS equations (Figure FCg Lastly, MAE models are able to distinguish varying resolutions of PDE data after multi-resolltion training, suggesting high model capacity across both diverse PDEs and discretizations (Figure D).
There are certainly limitations emergent trends learned by masked pretraining. Unscen boundary conditions can be challenging; when pretrained on periodic 1D KdV-Burgers data, MAE models can only distinguish between periodic and non-periodic Heat equation samples without understanding differences between Dirichlet.

---

Published in Transactions on Machine Learning Research (12/2024)
K0V-E
AN
1lam
Perseaic
Refore
NOT
125.
ehoricdic
After
Figure 4: (-SNE embeddings of various PDEs. Plots show embeddings before and after using the MAE to encode samples, shown on the top and bottom. The MAE latent space shows structure despite not seeing labels of coefficients, PDEs, or BCs. A: ID KdV-Burgers equation, colored by a. B: 1D Advection equation, colored by c. C: 1D Heat, Burgers, Advection, and KS equations, colored by PDE. D: 1D KdV-Burgers equation, colored by resolution. E: 1D Heat equation, colored by boundary condition. F: 2D Heat, Advection and Burgers equations, colored by and c.
and Neumann BCs (Figure ΕE). In addition, trends in 2D PDE data are more difficult to learn without labels. Despite separating between Heat, Adv, and Burgers samples, latent trends are only observed in 2D Advection samples based on wave speed e (Figure F
5.3 PDE Feature Prediction
To evaluate the latent representations learned from masked pretraining, we regress PDE coefficients and classify PDE boundary conditions, equation families, and spatial resolutions from an initial time window. Regression tasks separated by PDE (KaV-B, KS, Heat, Adu, Burgers, N8). Further, classification tasks are formulated as: (Heatpc): predicting Periodic, Dirichlet, Neumann BCs from the Heat equation, (Wavegc): predicting Dirichlet Neunann BCs from the Wave equation, (PDEs): predicting unseen PDEs from Heat, Adv, Burgers, and KS equation samples, and (Res.): predicting resolutions from n2 = (50, 60, 70, 80, 90, 100} in 1D or (nz; ny) = ((48, 48), (52, 52), (56, 56), (60, 60), (64, 64)} in 2D.
Several model variants are evaluated: a randomly initialized VIT baseline (MAF,), a pretrained, frozen MAE encoder with linear head (MAE)), and pretrained, fine-tuned MAE encoder (MAE). For these MAE modcls, regression and classification to the mumber of prediction features through simple MLP. The baseline encoder is VIT with the same model size and architecture as the pretrained MAE encoder, and when the MAE encoder is frozen, only the MLP head receives gradient updates. Additionally, we train MLP and CNN models to regress coefficients or classify PDE features in order to benchmark the difficulty of these tasks. Since these models cannot process variable sized inputs, the Res. experiment not performed for these simpler baselines. Results are shown Tables
performed using CLS token and projecting the CLS embodding
with full error bars in Appendix E.1
In ID. the frozen MAE, encoder is able to outperform supervised baseline the KdV-B, Heat, PDEs, and Res. tasks despite never seeing labels throughout training. Further performance gains be realized by allowing the MAE encoder to fino-tune on labeled data, and can outperform random initialization on unscen equations and BCs. An exception to this is the Advection and Wave equations. We hypothesize that these PDEs are heavily governed by the wave speed e and boundary conditions (bouncing vs. reflecting waves) and are simple trends that supervised models can learn quickly. Indoed, the simple MLP and CNN baselines are able to achieve perfect classification accuracy on the Wave equation, likely since the wave either reflects bounces off the boundaries according to the BC.
10

---

Published in Transactions on Machine Learning Research (12/2024)
Table ID PDE feature prediction after MAE pretraining on the KdV-Burgers equation. Models are fine-tuned on 2000 held-out, labeled samples for each task. We consider regressing coefficients across four PDEs as well as classifying BCs of the heat/wave equation, identifying equations from mixed set of PDEs, and sorting different spatial resolutions of PDE. Regression errors given RMSEx 10-2 and averaged over 5 seeds.
classification errors are given as X- X-Entx10-1
Model
KdV-B
Heat
Adv
KS
Heatyc
WaveBC
PDEs
Res.
MLP
6.641
1.290
2.435
0.821
1.733
0.000
27.24
CNN
1,835
1.250
1.470
0.342
0.195
0.000
0.427
MAR,
3.454
0.834
0.241
0.354
0.123
0.022
0.355
64.52
MAE,
1.334
0.677
0.551
0.368
1.164
1.816
0.174
63.34
MAE
0.905
0.505
0.244
0.156
0.018
0.053
0.025
28.32
Table 4: 2D PDE feature prediction after MAE pretraining on a combined set of 2D Heat, Advection, and Burgers equations. Models are fine-tuned on 1024 held-out, labeled samples for each task, or 3072 samples in the combined case. Regression errors are given RMSEx10-1 and classification errors are given X-Entx10-1 averaged over seeds.
Model
Heat
Adv
Burgers
Combined
NS
Res.
CNN
0.305
1.057
1.370
0.371
1.224
MAE,
0.084
0.506
0.682
0.320
0.748
0.694
MAE,
0.232
0.540
0.606
0.384
0.709
0.636
MAE
0.062
0.507
0.409
0.265
0.594
0.005
In 2D, the effects of fine-tuning are more pronounced. Within the pretraining set, only in the 2D Burgers task does freezing the MAE encoder outperform a supervised baseline; nevertheless, masked pretraining serves good initialization for supervised fine-tuning. In addition, despite differing physics, prior knowledge from simpler 2D PDEs seems to benefit regression on the Navier-Stokes equations. When classifying 2D Heat, Advection, and Burgers data based on their discrctization, MAE models greatly benefit from pretraining on multi-resolution data. We hypothesize that in 2D PDEs, variable spatial resolutions can be challenging to distinguish due flattening the spatial dimension when patchifying inputs, whereas in ID PDEs the data is already flattened.
5.4 Conditional Time-stepping and Super-resolution
1D Experiments To evaluate the use the MAE encoder for practical tasks, we train a neural solver on various ID PDEs predict or upsample physics. For prediction, or time-stepping, models are given solutions are time & and queried to predict solutions models are given low-rcsolution solution at time and queried to predict high-fidelity solution at time For our experiments in 1D, consider five PDEs (KdV-B, Heat, Burgers, Adv, KS) well as two PDEs under varying boundary conditions (Heatac, Waresc) and predicting physics various resolutions of the KdV-Burgers equation (Res.). Time-stepping is performed autoregressively by predicting multiple timesteps simultancously to reduce error accumulation (Brandstetter et al.] 2022). Furthermore, the pushforward trick (Brandsttter al. 2022) implemented. This adds model noise inputs during training by making a prodiction of future timestep and using that prediction as noised input the model; importantly gradients are not calculated for the initial pass. Lastly, we test on FNO (Li al., 2020 and Unet architectures (Gupta) and Brandstetter 2022) (Ronneberger et al. 2015), and add conditioning information to hidden states after
time t+ 1. For upsampling, or super-resolution,
convolutions (Ho et al. 2020 Nichol and Dhariwal 2021
For super-resolution (SR), we implement a pipeline in which a network encodes low-resolution physics before upsampling with discretization inversion operator D-1 (linear interpolation in 1D and bicubic interpolation in 2D) and mapping to an output function space with a neural operator (Yang al.) 2023).oolowwi Following this, we implement a Resnet encoder (Wang et al., 2021] Zhang al.) 2018)) followed by an interpolation scheme and FNO operator; both Resnet and FNO models provided conditioning information from MAE encodings of
11

---

Published in Transactions on Machine Learning Research (12/2024)
Table 5: Conditional ID PDE time-stepping and super-resolution. Models are trained on 2000 held-out fino-tuning samples to predict or upsample physics across several settings. Validation errors are reported normalized L2 loss (time-stepping) or RMSEx10-1 (SR) summed over all timesteps and averaged across five soeds.
Model
KdV-B
Heat
Burgers
Adv
KS
Heatnc
Wavenc
Res.
FNO
1.153
0.671
1.094
0.437
0.830
2.408
0.147
1.141
FNO-MAE,
1.043
0.655
1.121
0.431
0.821
1.747
0.135
1.018
FNO-MAE
1.037
0.643
0.952
0.294
0.812
1.846
0.148
1.07
Unet
0.823
0.420
0.649
0.194
1.333
4.249
0.747
0.766
Unt--MAEF
0.806
0.425
0.582
0.177
1.241
4.734
0.699
0.688
Unet-MAE
0.758
0.363
0.546
0.210
1.125
5.157
0.659
0.683
Interp.
0,540
0.345
1.357
0.231
3.599
0.225
0.347
SR
0.520
0.203
0.881
0.223
2.673
0.210
0.516
SR-MAE,
0.481
0.173
0.691
0.169
2.460
0.204
0.376
SR-MAE
0.475
0.151
0.676
0.194
2.422
0.170
0.349
low-resolution physics. The motivation behind using a two-step SR pipeline is to learn a vector embedding using the Resnet, then map from vector to function space with 7-1, and finally transform this function using an learnable operator the output. For additional details on hyperparameters, see Appendix C Additionally, evaluate a simple baseline using linear interpolation to upsample low-resolution inputs (Interp.).
After pretraining MAE encoder on the ID KdV-Burgers equation, we compare the base neural solvers (FNO, Unet, SR) to conditioning on a frozen MAE embedding (-MAE,) and allowing the MAE encoder to fine-tune when conditioning (-MAE). Results in 1D are presented in Table 5 Within the pretraining distribution (KdV-B) and certain PDEs, MAE conditioning consistently improves timo-stepping and super-resolution performance. In addition, allowing MAE encoders to fine-tune can further lower errors. However, there are various exceptions, in particular PDEs with unseen boundary conditions. Despite this, improvements are consistent across different neural solver architectures, suggesting that pretrained MAE models can be agnostic to downstream model choices. In addition, in 1D, SR results are less significant suggesting that simple interpolation schemes are often enough for these phenomena, especially for simple equations such the advection or wave PDEs.
2D Experiments Following the setup in ID, we repeat timeestepping/supe--essttttt experiments on 2D PDÉs (Heat, Ade, Burgers, NS) and a combined set of 2D Heat, Advection, and Burgers equations (Combined). Additionally, we evaluate timo-stepping performance on the combined 2D Heat, Advection, and Burgers equations discretized at variable resolutions (Res.). We follow the same conditioning and training strategies 1D experiments, but modify the architectures to support 2D inputs, and present results in Table G Additionally, the interpolation baseline implements bicubic upsampling (Interp.). After pretraining an MAE encoder on the 2D Heat, Advection, and Burgers equations, we observe improvements in conditional physics prediction and upsampling. Improvements tend to be more pronounced in 2D; we hypothesize that the increased difficulty of the task increases the importance of MAE encoder guidance in time-stepping and super-resolution. However, out-of--isttiibution datasets are still challenging: when extrapolating pretrained encoders to new PDEs, such as the Navier-Stokes equations, the performance is limited. Nevertheless, we observe similar trends whereby MAE conditioning is agnostic to downstream architectures.
Additional Benchmarks We consider two additional benchmarks: randomly initialized VIT encoder that embods PDE inputs to conditioning vector as well lincar model that encodes ground-truth cocfficient or boundary condition information to a conditioning vector. We present detailed results in Appendix E.2 and E.4 and discuss the overall results here, In general, we observe that the randomly initialized, fine-tuned encoder also improves PDE prediction and upsampling, and this improvement generally matches or outperforms the performance of the frozen MAE encoder. However, allowing the MAE encoder to fine-tune generally outperfomms this random initialization and approaches the linear benchmark.
12

---

Published in Transactions on Machine Learning Research (12/2024)
Table 6: Conditional 2D PDE time-stepping and super-resolution. Models are trained on 1024 held-out fine-tuning samples, or 3072 in the combined case, predict or upsample physics. Validation errors are reported as normalized L2 loss (time-stepping) or RMSEx10 (SR) summed over all timesteps and averaged across three seeds.
Model
Heat
Adv
Burgers
Combined
NS
Res.
FNO
0.427
2.301
0.417
0.978
0.466
1.006
FNO-MAE,
0.233
1.179
0.252
0.607
0.59
0.701
FNO-MAE
0.128
1.135
0.198
0.494
0.477
0.499
Unet
0.147
1.795
0.226
0.835
0.713
0.908
Unet-MAER
0.136
1.804
0.229
0.761
0.669
0.861
Unet-MAE
0.116
1.230
0.186
0.669
0.692
0.676
Interp.
0.492
0.937
0.488
0.673
0.367
SR
0.175
2.014
0.295
0.534
0.326
SR-MAE;
0.159
0.659
0.264
0.407
0.347
SR-MAE 0.152
0.639
0.253
0.472
0.337
Although the linear benchmark generally performs the best in 1D, there are certain exceptions to this. Equations dominated by the coefficient response generally suffer from coefficient conditioning; we observe that the model heavily overfits to the true coefficient and does not learn the underlying PDE dynamics. Furthermore, for PDEs that do not have coefficient information, such as the 1D inviscid Burgers equation, this linear benchmark cannot be applied. In these cases, MAE encodings can still improve performance, suggesting that there are latent PDE features beyond coefficient information that neural solvers can benefit. from. Lastly, in 2D the linear benchmark performs much worse, only achieving the lowest error in few sconarios, and in some cases harming performance of the base model. This could be because PDE dynamics becomes much more complex in 2D and relies less on coefficient information, which is a low-dimensional vector and provides sparse information.
Lastly, benchmark our method against other protraining methods, we consider pretraining an encoder using contrastive self-supervised technique proposed by Mialon et al. (2023) which relies on using Lie symmetries to cluster PDE data. in learned latent space. We contrastively pretrain encoder 1D KdV-Burgers data and the evaluate conditional timestepping performance these results in Appendix E.3 To summarize, our approach is on par with Lie contrastive pretraining for PDE samples within the pretraining distribution; however, when extrapolating to unseen PDFs, masked pretraining able to outperform contrastive methods.
various downstream ID PDEs. We present
6 Discussion
In this section, discuss results and provide additional insight. Although masking pretraining strategy is not physically valid, this requirement does not seem to be necessary for learning. Both in this study, and in related works (Hao et al.) 2024 Zhou et al., 2024 Rahman et al. 2024), noising or masking strategies are used to improve model performance, which both represent artificial phenomena. Within the context of masking, also natural to ask if a unique solution exists given a masked input. Certainly at the rete where all of the input masked the solution is not unique; however, it seems empirically that only small amount of information (25% in 1D and 10% in 2D) is needed, which corroborated by related work in the CV domain [Feichtenhofer et al). We can observe this through the validation error: if this quantity is small, unique solution is being regressed in the validation set.
While promising, the presented results in time-stepping and super-resolution would likely not outperform a foundation model trained specifically for these tasks (Hao et al. 2024 Herde et al. (024).. This is because direct transfer learning has boen shown to outperform surrogate objectives when pretraining for PDEs (Zhou) et al. 2024); one hypothesized reason for this the lack of abundant unlabeled data in the PDE domain (or equivalently, the downstream task uses also unlabeled data). However, encoder-style approaches such
13

---

Published in Transactions on Machine Learning Research (12/2024)
this work or contrastive PDE encoders (Mialon al. 2023 Zhang al 2023) are more flexible than foundation models, capable of being applied to arbitrary downstream architect ures and different fine-tuning tasks. Additionally, when using a surrogate objective to train an encoder, models can learn more general latent representations, compared to using a neural solver or foundation model to only predict the next step of PDE rollout. We show some preliminary results demonstrating this in Appendix F While this versatility is interesting, the practicality of this unclear, since time-stepping is so singularly important. However, this research direction still underexplored and perhaps future work will find an interesting set of uses for PDE encoders.
Within the context of masked protraining, there are few additional limitations to be recognized. It can be costly to fully fine-tune the pretrained encoder during time-stepping or super-resolution since it operates on all unmasked tokens, and as a result does not benefit from the speed gains during MAE pretraining. This makes full fine-tuning likely infeasible when compared to baselines. To address this, freezing the MAE encoder greatly improves the training speed but decreases performance, espocially on unseen PDEs. With enough PDE data during pretraining, freezing MAE encoders can be more practical since the pretraining distribution would cover most downstream cases. Lastly, for simpler ID dynamics when coefficient information is available, conditioning on ground-truth PDE parameters remains the best choice in most scenarios. However, these approaches are not exclusive; imitial work suggests that models provided with both ground-truts truth information and MAE embeddings can outperform models provided with just one of either.
7
Conclusion
We have presented a new method that extends masked pretraining from vision and video domains physics, and evaluated several PDE-specific modifications. We empirically validate MAE models through reconstructing diverse set of 1D and 2D PDEs and show limited generalization behavior to different spatial discretizations and unseen equations. Furthermore, MAE training and find structured trends that can be used to predict PDE features. In practice, MAE encoders can also be used to improve time-stepping and super-resolution tasks across diverse physics scenarios.
evaluate the latent representations learned during
A promising direction would be to scale MAF models to larger datasets, as the current approach exhibits the same scalability as the originally proposed MAE (He et al. 2021). Additionally, future work could explore masked modeling approaches in more complex 2D and 3D problems. Lastly, future work could explore manipulating latent physics to generate new solutions or performing arithmetic in an autoencoder latent space. We present a potential sctup for this in Appendix D which relies on encoding solutions from separate equations, adding them in latent space, and decoding this latent vector.
References
[1] Richard FitzHugh. Impulses and physiological states in theoretical models of nerve membrane. Biophysical Journal, 1(6):445 466, 1961 ISSN 0006-3495. doi: http:://00..arr///000000/0000000 URL http///wwwaaaooeeeeeetcoo///ttta
[2] J. Nagumo, S. Arimoto, and S. Yoshizawa. An active pulse transmission line simulating nerve axon. Proceedings of the IRF, 50(10):2061-207) 1962. doi: 11110///RPRO...16222222222
[3] Edward N. Lorenz. Deterministic nonperiodic flow. Journal of Atmospheric Sciences, 20(2):130 141, 1963. doi: 10. 10117775//Beuuuu0000000000000001 2.0. URL
www.sooo.s.amettc.ooooooo 1963_ 020 0130 dnf 2_0_co_________ 2m
ametsoc org/
[4] Maziar Raissi, Paris Perdikaris, and George Em Karniadakis. Physics-informet neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial dif- forential equations. Journal of Computational Physics, 378:686-707, 2019. ISSN 0021-9991. doi:
URL https: http://wwww.cinceddeeeeeee http:///rrrttccc///cccc/e//raaaa
p11/50021999118077225
5 Lu Lu, Pengzhan Jin, and George Em Karniadakis. Deeponet: identifying differential equations based on the universal approximation theorem of operators. 10 2019.
Learning nonlinear operators for
14

---

Published in Transactions on Machine Learning Research (12/2024)
doi: 10.1038/412256--22-00002225 URL http://p/w/atttt..////////////// 001-22226-021003-2--
[6] Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar. Fourier neural operator for parametric partial differential equations. 10 2020. URL http:://rxxv.orr/aas///000 08895
7 Shuhao Cao. Choose transformer: Fourier galerkin, 2021.
[8] Johannes Brandstetter, Daniel E. Worrall, and Max Welling. Message passing neural PDE solvers. CoRR, hbs/2202.03336, 2022. URL htttps://arxwv.rrg/ablll
[9] Zijie Li, Kazem Meidani, and Amir Barati Farimani. Transformer for partial differential equations' operator learning, 2023.
[10] Nikola Kovachki, Samuel Lanthaler, and Siddhartha Mishra. On universal approximation and error bounds for fourier neural operators. Journal of Mochine Learning Research, 22(290):1 76, 2021. URL http://jmlr http:://jalr.orr//ppprssss// pttp///ppppee////2220066666 html
[11] Zijie Li, Anthony Zhou, Saurabh Patil, and Amir Barati Farimani. Cafa: Global weather forecasting with factorized attention on sphere, 2024. URL http:////araa...rrraaa////....
[2] Makoto Takamoto, Francesco Alesiani, and Mathias Niepert. Learning neural pde solvers with parameter- guided channol attention. 2023. URL
[13] Cooper Lorsung, Zijie Li, and Amir Barati Farimani. Physics informed token transformer for solving partial diffcrential equations, 2024.
[14] Junhong Shen, Tanya Marwah, and Amoet Talwalkar. Ups: Towards foundation models for pdc solving via cross-modal adaptation, 2024.
[15] Shashank Subramanian, Peter Harrington, Kurt Keutzer, Wahid Bhimji, Dmitriy Morozov, Michael Mahoney, and Amir Gholami. Towards foundation models for scientific machine learning: Characterizing scaling and transfer behavior. 2023. URL httppp////rrvi.o/aabba/////s//// 2000.//0.00002
[16] Michael McCabe, Bruno Régaldo-Saint Blancard, Liam Holden Parker, Ruben Ohana, Miles Cranmer, Alberto Bietti, Michael Eickenberg, Siavash Golkar, Geraud Krawezik, Francois Lanusse, Mariel Pettee, Tiberiu Tesileamu, Kyunghyun Cho, and Shirley Ho. Multiple physics protraining for physical surrogate models, 2023.
[17] Zhongkai Hao, Chang Su, Songming Liu, Julius Berner, Chengyang Ying, Hang Su, Anima Anandkumar, Jian Song, and Jun Zhu. Dpot: Auto-regressive denoising operator transformer for large-scale pde pre-training, 2024.
[18] Jacob Devlin, Ming-Wer Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pro-training of deep bidirectional transformers for language understanding. 10 2018. URL http:///axiv..rgg// 1810 04805
[19] Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross B. Girshick. Masked autoencoders scalable vision learners. CoRR, 0t//2111106377 2021. URL https://arxlv.org. orw./bs/////////o//o/o.oooo.eb/e abs/
2111.06371
[20] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need, 2023.
[21] Christoph Feichtenhofer, Haoqi Fan, Yanghao Li, Kaiming He, and Meta Ai. Masked autocncoders as spatiotemporal learners. URL https//s/g.tttooocccc/ccocorrrrs
[22] Nikola Kovachki, Zongyi Li, Burigode Liu, Kamyar Azizzadenesheli, Kaushik Bhattacharya, Andrew Stuart, Anima Anandkumar, and Lorenzo Rosasco. Neural operator: Learning maps between function spaces with applications to pdes, 2023.
15

---

Published in Transactions on Machine Learning Research (12/2024)
[23] Jayesh K. Gupta and Johannes Brandstetter. Towards muttispatiioemporal-aalll generalized pde modeling, 9 2022. URL http://arxiv. httpaaaiiv..rrggbss...saa 15616
[24] Lu Lu, Xuhui Meng, Shengze Cai, Zhiping Mao, Somdatta Goswami, Zhongqiang Zhang, and George Em Karniadakis. A comprehensive and fair comparison of two neural operators (with practical extensions) based on fair data. 11 2021. doi: 10.1016/j.cma.22 20t...2222...7.... URL http://arxiv.. wwtpeiii..rrg//aaa org/abs/2111 05512http://dx.oo doi 10.
cma. 2022. 114778
[25] Rui Wang, Robin Walters, and Rose Yu. Meta-learning dynamics forecasting using task inference. 2 2021. URL http://arxii..ore hwtg//bb/22200 2010/22222212212771
[66] Matthicu Kirchmeyer, Yuan Yin, Jérémic Donå, Nicolas Baskiotis, Alain Rakotomamonjy, and Patrick Gallinari. Generalizing to new physical systems via context-informed dynamics model.
[27] Yuan Yin, Ibrahim Ayed, Emmanuel de Bézenac, Nicolas Baskiotis, and Patrick Gallinari. Leads: Learning dynamical systems that generalize across environments. 6 2021. URL http://arxiv.orgg abs/2106.04544
[28] Lu Zhang, Huaiqian You, Tian Gao, Mo Yu, Chung-Har Hao Lee, and Yue Yu. Metano: How transfer your knowledge on learning hidden physics. 2023. URL http://aaxiv..or//abb//////
[29] Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey F. Hinton. simple framework for contrastive learning of visual representations. CoRR, abs/2002.05709, 2020. URL https://arxiv.orgg
abs/2002.05709
[30] Florian Schroff, Dmitry Kalenichenko, and James Philbin. Facenet: unified embedding for face recognition and clustering. In 2015 IEEE Conference on Computer Vision and Pattern Recognition (CVPR). TEFE, June 2015. doi: /o0100//op:
URL. wtw..ddoi..rg//..11.11..
[31] Jure Zbontar, Li Jing, Ishan Misra, Yann LeCun, and Stéphane Deny, Barlow twins: Self-supervised learning via rodundancy reduction, 2021.
[32] Adrien Bardes, Jean Ponce, and Yann LeCun. Vicreg: Varianee-ivvaranccccvaaaan regularization for self-supervised learning, 2022.
[33] Cooper Lorsung and Amir Barati Farimani. Picl: Physics informed contrastive learning for partial differential equations, 2024.
[34] Rui Zhang, Qi Meng, and Zhi-Ming Ma. Deciphering and integrating invariants for neural operator learning with various physical mechanisms. National Science Review, 11(4). Decomber 2023. ISSN doi: (0.1093/nsr/nwadde
2053-714X.
URL httt::://xx.doi.orgg.. 1093/nst/nwad336
[35] Grégoire Mialon, Quentin Garrido, Hannah Lawrence, Danyal Rehman, Yann LeCun, and Bobak T. Kiani. Self-super learning with lie symmetries for partial differential equations. 2023. URL
http:///artw..orgabb////22 /tt//////////0055222222
366 Johannes Brandstetter, Max Welling, and Daniel E. Worrall. Lie point symmetry data augmentation for neural pde solvers. 2 2022. URL http://arxiv wwwpe/iivorr//bbbee
[37] Jason Wei, Yi Tay, Rishi Bommasani, Colin Raffel, Barret Zoph, Sebastian Borgeaud, Dani Yogatama, Maarten Bosma, Denny Zhou, Donald Metzler, Ed H. Chi, Tatsunori Hashimoto, Oriol Vinyals, Percy Liang, Jeff Dean, and William Fedus. Emergent abilitics of large language models, 2022.
[38] Tom B. Brown, Benjamin Mann, Nick Ryder, Melanic Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert- Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesso, Mark Chen, Eric Sigler, Matcusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are sho-shot learners, 2020.
16

---

Published in Transactions on Machine Learning Research (12/2024)
[39] Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners. URL
[40] Liu Yang, Siting Liu, Tingwei Meng, and Stanley J Osher. In-context operator learning with data prompts for differential equation problems. 2023. doi: 10.1073/pnas. URL httpe:///bi.orrg/100. pnas 2310142120
[41] Wuyang Chen, Jialin Song, Pu Ren, Shashank Subramanian, Dmitrry Morozov, and Michael W. Mahoney. Data-efficient operator learning via unsupervised pretraining and in-context learning. 2 2024. URL
http:://axiv..or/// wwwp://////aaba/aaae/oeoase http://raivv.rrr//b//20000......
[2] Somdatta Goswami, Katiana Kontolati, Michael D. Shiclds, and George Em Karniadakis. Deep transfer operator learning for partial differential equations under conditional shift. Nature Machine Intelligence, doi: 10.1038 -42222660002--00560-2- URL. http:
4(12):1155-1164 December 2022. ISSN 2522-5839.
11w../..rg//11118///222200000 1038/s42255 022-22566--22000---222
43] Ayan Chakraborty, Cosmin Anitescu, Xiaoying Zhuang, and Timon Rabczuk. Domain adaptation based transfer learning approach for solving pdes on complex geometries. Engineering with Computers, 38(5): 4569 4588, Oct 2022. ISSN 1435-5663. doi:
URL attpp://doil org/
1007/=00366 0000//800366-022-1666111 01661-2
[44] Hengje Wang, Robert Planas, Aparna Chandramowlishwaraa and Ramin Bostanabad. Mosaic flows: A transferable deep learning framework for solving pdes on unseen domains. Computer Methods in Applied Mechanics and Engineering, 389, 2022. ISSN 00457825. doi: (0100///.cma.222..114444....
[55] Anthony Zhou, Cooper Lorsung, AmirPouya Hemmasian, and Amir Barati Farimani. Strategios for pretraining neural operators, 2024. URL htttp:://axxiv.or//bb////// /
[46] Tapas Tripura and Souvik Chakraborty. A foundational neural operator that continuously learns without forgetting, 2023.
[47] Alexcy Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. 10 2020.
URL http://arxiv. http:///axxi..or/abb///////
11929
[48] Zhenda Xie, Zheng Zhang, Yue Cao, Yutong Lin, Jianmin Bao, Zhuliang Yao, Qi Dai, and Han Hu. Simmim: A simple framework for masked image modeling. CoRR, abs/2111.0988) 2021. URL
https://arxiv.orr wrg. 09888
[9) Siyuan Li, Luyuan Zhang, Zedong Wang, Di Wu, Lirong Wu, Zicheng Liu, Jun Xia, Cheng Tan, Yang Liu, Baigui Sun, and Stan Z. Li. Masked modeling for self-supervised representation learning on vision and beyond, 2024.
50] Shuhao Cao, Peng Xu, and David A. Clifton. How to understand masked autoencoders. 2 2022. URL http://aaxxi..ooraa////a/// 03670
[1] Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, Gretchen Krueger, and llya Sutskever. Learning transferable visual models from natural language supervision, 2021.
[2] Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen. Hicrarchical text-conditional image generation with clip latents, 2022.
[53] Qidong Yang, Alex Sinnande--Garcia. Paula Harder, Venkatesh Ramesh, Prasanna Sattegeri, Daniela Sxwareman, Campbell D. Watson, and David Rolnick. Fourier neural operators for arbitrary resolution climate data downscaling, 2023.
[54] Anurag Arnab, Mostafa Dehghani, Georg Heigold, Chen Sun, Mario Lučić, and Cordelia Schmid, Vivit: A video vision transformer, 2021. URL htl.stttp:: http:://aaavv/aaaaaaaaaeaae 15691
17

---

Published in Transactions on Machine Learning Research (12/2024)
[55] Peter Olver. Applications of Lie Groups to Differential Equations. Springer New York, NY, 1986.
[6] Lucas Beyer, Pavel Ixmailov, Alexander Kolesnikov, Mathilde Caron, Simon Kornblith, Xiaohua Zhai, Matthias Minderer, Michael Tschannen, Ibrahim Alabdulmohsin, and Filip Pavetic. Flexivit: One model for all patch sizes, 2023.
[57] Rui Tian, Zuxuan Wu, Qi Dai, Han Hu, Yu Qiao, and Yu-Gang Jiang. Resformer: Scaling vits with multi-resolution training, 2023.
[58] Mostafa Dehghani, Basil Mustafa, Josip Djolonga, Jonathan Heek, Matthias Minderer, Mathilde Caron, Andreas Steiner, Joan Puigcerver, Robert Geirhos, Ibrahim Alabdulmohsin, Avital Oliver, Piotr Padlowski, Alexey Gritsenko, Mario Lučié, and Neil Houlsby. Patch n' pack: Navit, vision transformer for any aspect ratio and resolution, 2023.
[59] Qihang Fan, Quanzeng You, Xiaotian Han, Yongfei Liu, Yunzhe Tao, Huaibo Huang, Ran He, and Hongxia Yang. Vitar: Vision transformer with any resolution, 2024.
[60] A. Jeffrey and M.N.B. Mohamad. Exact solutions to the kdv-burgers' equation. Wave Motion, 14 (4):369 375, 1991. ISSN 0165-2125. doi: https///dbi.rrg//1.006/ 10111//000/10111//////11111 (165-2225(91)9000))) URL https:
/www. www..sciencedirect.om com ott.://cw.....ooo........//aaaa.
[61] Martin S. Alnaes, Jan Blechta, Johan Hake, August Johansson, Benjamin Kehlet, Anders Logg, Chris N. Richardson, Johannes Ring, Marie E. Rognes, and Garth N. Wells. The FEniCS project version 1.5. Archive of Numerical Soflware, 3, 2015. doi: 10. 2015588/an.222551000
[62] Anders Logg, Kent-Andre Mardal, Garth N. Wells, et al. Automated Solution of Differential Equations by the Finite Element Method. Springer, 2012. doi: 100007997--3-42-2211
[63] Phillip Lippe, Bastiaan S. Veeling, Paris Perdikaris, Richard E. Turner, and Johannes Brandstetter. Pele-refine: Achieving accurale long rollouts with neural pde solvers, 2023.
[64] Shawn G. Rosofsky, Hani Al Majed, and E A Huerta. Applications of physics informed neural operators. Machine Learning: Science and Technology, 4, 2023.
[65] Nail H. Ibragimov. CRC Handbook of Lie Group Analysis of Differential Equations: Symmetries, Exact Solutions, and Conservation Laws. Number v. 1. Taylor & Francis, 1993. ISBN 9780849344886. URL
https://books. httpss////oooss.ggggle..ooo www.o///oooosiooooooooo (oookssiddmppm5yy1ge7777
[66] Laurens van der Maaten and Geoffrey Hinton. Visualizing data using f-sme. Journal of Machine Learning Research, ((86):2579220) 2605, 2008. URL
http://www...aaa//////////aaaaaa html
[67] Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation, 2015.
[68] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models, 2020.
[69] Alex Nichol and Prafulla Dhariwal. Improved denoising diffuusion probabilistic models, 2021.
[70] Xintao Wang, Liangbin Xie, Chao Dong, and Ying Shan. Real-esrgan: Training real-world blind super-resolttion with pure synthetic data, 2021.
[71] Yulun Zhang, Yapeng Tian, Yu Kong, Bineng Zhong, and Yun Fu. Residual dense network for image super-resolution, 2018
[72] Md Ashiqur Rahman, Robert Joseph George, Mogab Elleithy, Daniel Leibovici, Zongyi Li, Boris Bonev, Colin White, Julius Berner, Raymond A. Ych, Jean Kossaifi, Kamyar Azizzadenesheli, and Anima Anandkumar. Pretraining codomain attention neural operators for solving multiphysics pdes, 2024. URL https://arxic: http://aawwvv/..r//////////// 12553
18

---

Published in Transactions on Machine Learning Research (12/2024)
[73] Maximilian Herde, Bogdan Raonić, Tobias Rohner, Roger Käppeli, Roberto Molinaro, Emmanuel de Béxenac, and Siddhartha Mishra.
Poseidon: Efficient foundation models for pdes, 2024. URL
19101
[74] Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E. Hinton. Layer normalization, 2016.
19

---

Published in Transactions on Machine Learning Research (12/2024)
Table 7: MAE Hyperparameters during pretraining.
(a) 1D PDEs
(b) 2D PDEs
Parameters
Value
Parameters
Value
Batch Size
256
Batch Size
Epochs
Epochs
Encoder Dim
256
Encoder Dim
256
Decoder Dim
32
Decoder Dim
32
Patch Size
(5, 5)
Patch Size
(4, 4, 4)
Masking Ratio
0.75
Masking Ratio
0.90
Time Window
20
Time Window
16
Augmentation Ratio
0.5
Augmentation Ratio
0.5
Base LR
1e-3
Base LR
1e-3
Optimizer
AdamW
Optimizer
AdamW
Scheduler
OneCycleLR
Scheduler
OneCycleLR
A MAE Implementation
implement the MAE encoder and decoder, we use a ViT architecture (7), which uses a self-atenntion layer (20 and MLP, both with LayerNorms (74). We present hyperparameters in Table 7 To study the effects of various hyperparameters, including model size, masking ratio, and patch size,
ablation studies masked reconstruction of ID PDEs, and report reconstruction MSE errors validation set in Table Overall, we find that increasing model size in limited data regimes-onle 10000 KdV-Burgers samples were used in pretraining tends to contribute to overfitting and increases validation errors. Predictably, increasing the masking ratio increases reconstruction errors as a result of less information being provided to MAE models. Furthermore, decreasing the patch size reduces errors but requires higher compulational cost, which consistent with results in CV domains (56).
In 1D, the MAE is trained on single NVIDIA GeForce RTX 2080 Ti, and reaches convergence in about 6 hours. In 2D, the MAE is trained on single NVIDIA RTX A6000, and reaches convergence in about 24 hours. Model size, masking ratio, and patch size all affect the computational cost and can be used to tradcoff performance for compute and memory.
To motivate the VIT architecture, we investigate the effect different model choices and hyperparameters on performance. We find that FNO autoencoders tend have poor reconstruction capabilitios due introducing spurious high-frequency modes when masking spatially. Unet approaches fare better but still suffer sharp boundaries across masked and unmasked regions. Additionally, we evaluate different VIT variants, such as ViVIT (axial attention) and Swin Transformer (window attention). We find that restricting the attention mechanism reduces performance, and the additional speedups were not significant since masking out large portions of the input already reduces computation. Lastly, model performance tends to vary smoothly with changes in hyperparameters; for example, reducing patch size slightly increases performance across downstream tasks.
Table 8: MAE model ablation studies on the 1D KdV-Burgers equation.
(a) Model Size
(b) Masking Ratio
(c) Patch Size, in (pr, Pz)
# Params
Error
Masking Ratio
Error
Patch Size
Error
1M
2.37e-03
0.6
8.02e-04
(5, 5)
1.12e-03
5M
2.48e-03
0.75
1. 1.66e-03
(4, 4)
7.05e-04
25M
3.36c-03
0.90
6. 6.68c-03
(4, 2)
6.46c-04
(2, 4)
4.79c-04
20

---

Published in Transactions on Machine Learning Research (12/2024)
B
Additional MAE Examples
B.1 Additional 1D MAE Predictions
Salv-Burgers
Heaat
Burgers
Acvection
Heat Dirichlet
Mave Pirishler
Wave Ne umaann
Figure 5: Additional 1D MAF Reconstruction examples after pretraining on the 1D KdV-Burgers equation. Each triplet is shown with the masked sample (Left), MAE reconstruction (Middle), and ground truth PDE (right). We include additional reconstructions unseen boundary conditions for the Heat and Wave equations.
21

---

Published in Transactions on Machine Learning Research (12/2024)
B.2 Additional 2D MAE Predictions
Heat
Advection
Burgers
Stokes
Figure 6: Additional 2D MAE Reconstruction examples after pretraining on the 2D Hcat, Advection, and Burgers Equations. Each sample is shown with the masked sample (Top), MAE reconstruction (Middle), and ground truth PDE (Bottom). We include sample MAE predictions variable resolutions for the 2D Heat, Advection, and Burgers equations; the lowest resolution (top) is (48, 48), the medium resolution (middle) is (52, 52), and the high resolution (bottom) is (56, 56) We include additional reconstructions of the incompressible NS equations at the nalive resolution (64, 64).
22

---

Published in Transactions on Machine Learning Research (12/2024)
B.3 2D Smoke Buoyancy Predictions
Density
X Velocity
Velocity
Figure 7: MAE validation reconstructions after training on 2D Navier-Stokes data with variable buoyancy nz: nyy = (56, 128, 128) with three data channels (p,eay Uxy uy) and a masking ratio of 0.75. Triplets are shown with the masked input (top), MAE reconstruction (middle), and ground truth (bottom), with the top, middle, and bottom triplets displaying density, X velocity, and Y velocity. The complex dynamics is challenging: indeed, many of the fine details are lost in the MAE reconstruction. We train with larger model (45M params) and patch size (2, 8, 8), which takes around 9 hours on NVIDIA RTX A6000 GPU.
factors (23). The MAE model is trained on a resolution of (n
23

---

Published in Transactions on Machine Learning Research (12/2024)
Table 9: Hyperparameters for architectures used for time-stepping and super-resolution.
(a) FNO Hyperparameters
(b) Unet Hyperparameters
(c) Resnet Hyperparameters
Parameter
ID/2DD
Parameter
1D ID/2D
Parameter
ID/2D
Modes
24/12
Hidden channels
16
Hidden channels
Width
64/48
Channel mults.
(1, 2, 4)
# Blocks Cond. dim Conditioning
Layers Cond. dim
Cond. dim Conditioning
32
32
32
AdaGN
Add
Conditioning
Add
Init Ir
8e-4
Init Ir
8e-4
Init lr
8e-4
# Params
1M/2M
# Params
1M/3M
# Params
1M/5M
C Training Details
Hyperparameters used for the FNO and Unet models during time-stepping presented in Table Addi- tionally, for the SR pipeline, hyperparameters used for the Resnet encoder, which uses residual dense blocks 711, and FNO operator are reported in Table g We present a schematic of the SR pipeline in Figure
Low-res input
Resnet
D
Hi-res output
Encoder
Figure 8: Conditional super-resolution pipeline.
D Latent Arithmetic
MAE encoders have shown strong capabilities in extracting information from self-supervised PDE learning, creating opportunities to operate in this latent space. This could be beneficial since many PDEs are compositions of simpler phenomena, and recombining PDE solutions in latent space may result in obtaining novel PDE solutions for free. After pretraining on the 1D KdV-Burgers equation, we consider an arithmetic task where samples of the Heat and Burgers equation are embedded and added in the latent space before being decoded to a novel solution (Figure D. Concretely, we generate 1D Heat data from the PDE: duu vd = 0, 1D inviscid Burgers data from the PDE: 8,u + udzuu = 0, and 1D viscous Burgers data from adding the two PDEs: 0.u vdzzu ud,u O. When given identical initial conditions and coefficients, PDE reconstructions obtained from summing latent Heat and Burgers embeddings qualitatively resemble solutions of the viscous Burgers PDE. In addition, if different latent embeddings can be added, the weighting of each embedding can be varied to control the resemblance of the reconstruction to make interpolated samples that have more shock formation or diffusive behavior (e.g., more resemble the Burgers or Heat equation..
Hext Burgers Raatenti
Hest Murt-Burgers (Truet
Bungers
ancader
(fargeess)
Burgers
Figure 9: A proposed setup for operating on latent PDE vectors. PDE data is encoded after masked pretraining and summed in the latent space before being decoded. These reconstructions can approximate summed PDEs in physical space.
24

---

Published in Transactions on Machine Learning Research (12/2024)
E Additional Results and Statistical Significance
E.1
Feature Prediction
Table 10: ID PDE coefficient regression. Validation errors are given as RMSE x10 The mean error and standard deviation are calculated over five scods; error bars are reported as one standard deviation above or below the mean. Lowest errors that are statistically significant are bolded.
Model
KdV-B
Heat
Adv
KS
MAE,
3.454 + 0.131
0.834 + 0.041
0.241 + 0.051
0.354 + 0.104
MAE;
1.334 + 0.036
0.677 + 0.016
0.551 + 0.03
0.368 + 0.02
MAE
0.905 + 0.059
0.505 + 0.065
0.244 - 0.064
0.156 + 0.023
Table 11: ID PDE feature classification. Validation Errors are given as X-Ent x10-4, The mean error and standard deviation are calculated over five seeds; error bars are reported as one standard deviation above or below the mean. Lowest errors that statistically significant bolded.
Model
Wavegc
Heatbc
PDEs
Res.
MAF,
0.022 + 0.004
0.123 + 0.108
0.355 + 0.276
64.52 + 3.475
MAE 1.817 + 1.175
1.164 + 0.838 0.018 + 0.002
0.174 + 0.064
63.34 = 3.71
MAE
0.053 + 0.048
0.025 + 0.013
28.322 + 23.351
Detailed results for 1D PDE feature prediction tasks are reported in Tables 10 and 11 For 1D tasks, certain experiments have high variance; we hypothesize that this due to the fact that each seed samples random dataset of 2000 samples from a much larger dataset. This would make some seeds easier regress/ cllassfy than others, but within each seed the models follow trends consistent with the mean statistics. Furthermore, the magnitude of the X-Fnt, error is very small, leading to high variations after the model has learned most. of the relevant features.
Table 12: 2D PDE coefficient regression and fcature classification. Validation errors are reported as RMSE x10-1 for regression tasks and X-Ent x10-1 for classification tasks. The mean error and standard deviation are calculated over five seeds; error bars reported one standard deviation above below the mean. Lowest errors that are statistically significant are bolded.
Model
Heat
Adv
Burgers
MAE,
0.084 + 0.014
0.506 + 0.009
0.682 + 0.037
MAE;
0.232 0.01
0.54 + 0.015
0.606 0.012
MAE
0.062 0.003
0.507 + 0.006
0.409 + 0.008
Model
Combined
NS
Res.
MAE,
0.320 + 0.007
0.748 + 0.005
0.694 t 0.174
MAE,
0.384 + 0.009
0.709 +-t 0.01
0.636 + 0.061
MAE
0.265 + 0.007
0.594 + 0.038
0.005 + 0.002
Detailed results for 2D feature prediction tasks are reported in Table 12 The 2D results tend to be more consistent and have lower variance, since a fixed dataset was used for each seed and only the shuffling is changed.
25

---

Published in Transactions on Machine Learning Research (12/2024)
E.2 Time- sttepste
Table 13:
Conditional 1D PDE time-stepping. Validation errors are reported normalized L2 loss summed all PDE timesteps. The mean and standard deviation calculated over five seeds; error bars are reported one standard deviation above below the mean. Lowest errors that statistically significant are bolded.
Model
KdV-B
Heat
Burgers
Adv
FNO
1.132 + 0.037
0.671 + 0.039
1.094 0.060
0.437 + 0.052
FNO-EnC
1.041 + 0.029
0.644 + 0.038
1.129 0.062
0.347 = 0.074
FNO-MAF
1.077 + 0.060
0.655 + 0.008
1.121 0.051
0.431 + 0.054
FNO-MAE
1.060 + 0.032
0.643 = 0.035
0.952 + 0.038
0.294 + 0.033
FNO-Lin
0.936 +: 0.029
0.75 + 0.062
N/A
0.204 + 0.019
Unet
0.872 + 0.069
0.420 + 0.021
0.649 0.052
0.194 = 0.059
Unet-Ene
0.834 + 0.043
0.395 + 0.021
0.582 0.032
0.224 + 0.057
Unet-MAEr
0.833 = 0.038
0.425 + 0.012
0.582 = 0.017
0.177 + 0.012
Unet-MAE
0.795 + 0.040
0.363 + 0.010
0.546 + 0.024
0.21 + 0.047
Unet-Lin
0.659 + 0.045
0.445 + 0.008
N/A
0.166 + 0.032
Model
KS
Heatpc
Wavebo
Res.
FNO
0.83 + 0.028
0.147 + 0.015
2.408 + 0.848
1.141 + 0.021
FNO-EnC
0.82 - 0.082
0.133 + 0.019
2.012 1.194
1.038 + 0.037
FNO-MAE,
0.821 + 0.088
0.135 + 0.013
1.747 0.665
1.018 + 0.13
FNO-MAE
0.812 = 0.061
0.148 = 0.015
1.846 1.885
1.070 = 0.011
FNO-Lin
0.757 + 0.077
0.132 + 0.020
1.454 + 0.450
0.899 + 0.01
Unet
1.333 + 0.068
0.747 + 0.043
4.249 2.296
0.766 + 0.083
Unet-Ene
1.203 + 0.102
0.691 + 0.046
4.902 1.935
0.739 + 0.088
Unet-MAEr
1.241 + 0.055
0.699 = 0.023
4.734 2.135
0.688 = 0.077
Unet-MAB
1.125 + 0.029
0.659 + 0.047
5.157 1.760
0.683 + 0.087
Unet-Lin
1.172 + 0.039
0.717 = 0.027
4.727 = 2.093
0.573 + 0.095
Following the main paper, we introduce two conditional benchmarks. We evaluate a randomly initialized and fino-tuned ViT encoder with the same architocture as the MAE encoder (-Enc), as well as linear encoder that embeds the ground-trute PDE parameters as the conditioning information (-Lin).
In ID, time-stepping results tend to have high variance; however, overall trends are still consistent with those reported in the main body. The variance is likely attributed to variations in the dataset for each seed; each seed samples a different set of 2000 samples from a larger PDE dataset, and as result, some data splits may be casier than others. This results the variance being high across seeds, however, within a seed (i.e. within a dataset), model performance closely follows trends consistent with the mean statistics.
26

---

Published in Transactions on Machine Learning Research (12/2024)
Table 14: Conditional 2D PDE time-stepping. Validation errors are reported as normalized L2 loss summed over all PDE timesteps. The mean error and standard deviation are calculated over three scods; error bars are reported as one standard deviation above below the mean. Lowest errors that are statistically significant are bolded.
Model FNO
Heat
Adv
Burgers
0.427 + 0.006
2.301 + 0.094
0.417 0.063
FNO-Enc
0.152 + 0.013
1.909 + 0.399
0.241 = 0.032
FNO-MAE,
0.233 = 0.028
1.179 0.036
0.252 0.012
FNO-MAE
0.128 + 0.008
1.135 0.121
0.198 0.009
FNO-Lin
0.118 = 0.005
2.531 + 0.013
0.149 0.036
Unet
0.147 = 0.031
1.795 + 0.105
0.226 = 0.018
Unet-Ene
0.132 0.040
1.604 0.164
0.218 0.02
Unet-MAEr
0.136 + 0.009
1.804 0.066
0.229 0.017
Unet-MAE
0.116 = 0.031
1.23 = 0.161
0.186 0.011
Unet-Tin
0.153 + 0.043
2.571 + 0.011
0.215 0.006
Model
Combined
NS
Res.
FNO
0.978 + 0.055
0.466 + 0.014
1.006 + 0.02
FNO-EnC
0.767 = 0.028
0.514 + 0.123
0.709 = 0.055
FNO-MAF+
0.607 + 0.019
0.59 + 0.107
0.701 = 0.051
FNO-MAE
0.494 + 0.043
0.477 0.029
0.499 = 0.024
FNO-Lin
0.977 0.021
0.445 0.026
0.986 0.015
Unet
0.835 + 0.067
0.713 + 0.005
0.908 0.061
Unet-Enc
0.791 + 0.061
0.695 + 0.027
0.971 = 0.023
Unet-MAEE
0.761 + 0.051
0.669 + 0.031
0.861 = 0.028
Unet-MAE
0.669 0.015
0.692 0.039
0.676 0.064
Unet-Lin
1.013 0.03
0.635 0.002
1.098 0.026
In 2D, time-stepping results have much lower variance; this is likely due to the fact that each seed uses the same dataset, with only the shuffling changing. Furthermore, the linear benchmark is less effective; in most experiments learned encoding can outperform ground-truth PDE parameters, especially when predicting a combined multi-resolution dataset of PDEs.
E.3 Comparison to Contrastive Learning with Lie Augmentations
Table 15:
We compare our approach to a contrastive self-supervised approach. After training a masked and contrastive encoder on the KdV-B pretraining set, we compare conditioning an FNO backbone to time-step different downstream ID PDEs. The MAE encoder shows comparable performance within the pretraining set (KdV-B), but has better generalization behavior unseen PDEs. Validation errors are reported normalized T.2 loss summed over all PDE timesteps
Model
KdV-B
Heat
Burgers
Adv
FNO
1.506
0.827
1.386
0.567
-NO-Cnntra
1.171
0.918
0.916
0.555
FNO-MAE
1.183
0.721
0.831
0.299
27

---

Published in Transactions on Machine Learning Research (12/2024)
E.4
Super-resolution
Table 16:
Conditional ID super-resolttae Validation errors are reported as RMSE x10-! summed over all PDE timesteps. The mean error and standard deviation calculated over five seeds; error bars reported one standard deviation above below the mean. Lowest errors that are statistically significant are bolded.
Model
KdV-B
Heat
Burgers
Adv
SR
0.520 + 0.021
0.203 + 0.011
0.881 0.062
0.223 + 0.004
SR-Enc
0.489 + 0.022
0.166 = 0.021
0.642 0.074
0.202 = 0.003
SR-MAE,
0.481 0.039
0.173 = 0.015
0.691 0.027
0.169 = 0.032
SR-MAE
0.475 + 0.018
0.151 + 0.028
0.676 + 0.060
0.194 + 0.016
SR-Lin
0.484 + 0.017
0.131 + 0.017
N/A
0.133 = 0.013
Model
KS
He
Wavene
SR
2.673 + 0.101
0.210 + 0.016
0.516 + 0.015
SR-Enc
2.585 + 0.062
0.174 + 0.01
0.373 + 0.022
SR-MAE;
2.460 + 0.056
0.204 + 0.015
0.376 + 0.032
SR-MAE
2.422 = 0.052
0.170 = 0.019
0.349 0.022
SR-Lin
2.517 0.038
0.177 + 0.027
0.451 0.014
Differences between benchmarks for 1D super-resolution lend to be incremental. Despite this, using frozen MAE encoding remains simple method to improve performance with negligible training cost. In general, super-rosolution for 1D PDEs is a relatively casy task, and changes in model architecture do not significantly affect results.
Table 17:
Conditional 2D super-resolution. Validation errors reported RMSE x10-1 summed over all PDE timesteps. he mean error and standard deviation are calculated over three seeds; error bars are reported as one standard deviation above or below the mean. Lowest errors that are statistically significant are bolded.
Model
Heat
Adv
Burgers
SR
0.175 0.023
2.014 + 0.205
0.295 + 0.018
SR-Enc
0.152 + 0.01
0.804 = 0.052
0.252 = 0.047
SR-MAE
0.159 + 0.007
0.659 = 0.089
0.264 + 0.05
SR-MAE
0.152 + 0.004
0.639 + 0.125
0.253 + 0.017
SR-Lin
0.167 + 0.015
2.016 + 0.015
0.263 + 0.021
Model
Combined
NS
SR
0.534 0.037
0.326 + 0.039
SR-Enc
0.49 + 0.016
0.364 + 0.006
SR-MAE,
0.407 0.002
0.347 = 0.007
SR-MAE
0.472 + 0.005
0.337 + 0.012
SR-Lin
1.235 + 0.032
0.366 + 0.003
In 2D, the general trend remains similar to 1D results; clear model choices for super-resolution are not apparent. Despite this, using frozen MAE encoding often outperforms the lincar benchmark; this be good way to boost model performance without additional training cost when PDE samples are within the pretraining distribution.
28

---

Published in Transactions on Machine Learning Research (12/2024)
F Comparison of Latent Embeddings
To understand the versatility of masking pretraining, we would like to consider what it means to be PDE learner beyond predicting future physics. For example, when reasoning about physics, humans can add, subtract, and rearrange equations to derive new knowledge. Additionally, we are able to identify what similar or different physical phenomena are in order to reason about relationships between physical quantities. Lastly, when observing physics, we intuitively understand that solutions must evolve forward with time and be temporally coherent. In the context of machine learning, these behaviors would have to be manifested in latent representations of physical solutions, and to this end, we demonstrate that masked autoencoders learn expressive, general representations of PDEs.
To do this, we propose a set of experiments to compare the MAE model against other pretraining paradigms such as contrastive or transfer learning. Firstly, we pretrain and MAE or contrastive encoder (35) on the 1D KdV-Burgens pretraining set. Additionally, we pretrain FNO and Unet model to predict future timesteps on this dataset in order model transfer learning scenario. Given these pretrained models, we seek understand their latent representations. To do this, embed PDE samples from three different scenarios.
Firstly, embed equation samples from the Heat and Burgers equations and add these two embeddings in latent space; the average pairwise distance is calculated between this summed embedding and a corresponding embedding from the viscous Burgers equation, which is an experiment in latent arithmetic (Arithmetic). Next, the average pairwise distance between embeddings from the Heat equation is calculated; embedded samples have varying initial conditions but evolve with the same coefficients, which tests if models can embed similar dynamics similar representations (Similarity). Lastly, he average pairwise distance between embeddings of subsequent timesteps of the Heat equation is calculated, mcasuring temporal coherence of latent embedding (Temporal). To ensure fair comparison, embeddings PCA and normalized (v r/mar(1|v]]a) ([maa)(l Masked pretraining performs well across these experiments and learns general representations due to its minimal inductive bias, compared to contrastive learning or neural solvers which have specific objectives to either maximize similarity or predict next the timestep.
projected to a common dimension 16 with
Table 18: Models are pretrained on the same KdV-B dataset and used to produce latent embeddings across different scenarios. The average pairwise L2 distance reported in each case.
Model
Arithmetic |
Similarity
Temporal
Average |
MAE
0.5072
1.3674
0.5277
0.8007
Contrastive
1.4839
1.3131
0.8626
1.2199
FNO
0.9442
1.3562
0.7565
1.0190
Unet
0.9447
1.3562
0.7397
1.0135
29


Reviews:
[{"review_id": "d4JVRkCn6k", "reviewer": "Reviewer_rzub", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "### Strengths\n\n- Novel application of a successful self-supervised learning technique (ViT-based MAEs) to PDEs.\n\n- Positive empirical results showing benefits of MAE pretraining and finetuning for various PDE types and downstream tasks.\n\n### Weaknesses\n\n1. **Performance**: Modest performance gains compared to baselines.\n\n    - Does not consistently outperform simpler approaches like conditioning on ground truth coefficients (at least in 1d).\n\n    - Mostly requires retraining the full encoder for best performance on downstream tasks, incurring significant additional costs.\n\n    - Limited ability to extrapolate to PDEs very different from those seen during pretraining, not sufficiently supporting statements such as \"Through self-supervised learning, MAE models can learn trends in PDEs without labeled data and in equations outside the training set.\", \"We show that learned representations can generalize to unseen equations or parameters\".\n\n    - Not clear how sensitive the results are w.r.t. hyperparameter choices like architectures, training methods, masking ratio, etc. For instance, using a randomly initialized and fine-tuned ViT encoder almost performs as well as the finetuned MAE (in particular in 1d), indicating that the benefits might be due to a higher parameter count or architectural prior instead of the pretraining strategy.\n\n2. **Baselines**: Insufficient comparisons against other baselines.\n\n    - Another baseline should be to train different models (such as UNet, FNO, and Transformer NOs) directly for time-stepping on the same pretraining data used by the MAE. The trained models can then be used for the same downstream tasks, potentially after finetuning. \n\n    - Comparisons against other approaches of pretraining foundation models are missing, e.g., MPP (https://arxiv.org/abs/2310.02994), DPOT (https://arxiv.org/pdf/2403.03542), CoDA-NO (https://arxiv.org/abs/2403.12553), Poseidon (https://arxiv.org/pdf/2405.19101), Strategies for Pretraining Neural Operators (https://arxiv.org/abs/2406.08473), which all appeared at least 3 months before the submission. In particular, CoDA-NO explores masking strategies similar to MAEs. Also, only a small comparison against the contrastive SSL approach by Mialon et al. (2023) is presented.\n\n    - Why is only a ViT baseline considered for feature prediction tasks and there are no comparisons against other (e.g., simpler) models? Details on training the baseline and finetuning the MAE seem to be missing, however, are crucial for a fair comparison. \n\n    - Training and inference costs are not reported and compared. In particular, the higher parameter count and need for additional finetuning might make the method less competitive against baselines and numerical solvers.\n\n    - One should compare directly to FNOs native superresolutions capabilities. In particular, it is not clear why the ResNet and the interpolation are necessary. For reference, one should also show the performance of only using interpolation.\n\n\n3. **Datasets**: Limited evidence for scalability to larger and more complex PDE datasets.\n\n    - The paper focuses mainly on relatively simple 1D and 2D PDEs and the findings may not generalize to more complex systems. In particular, on more complex equations (such as NS) or out-of-distribution datasets performance seems to be limited. For instance, the t-SNE plots seem less structured for the 2d equations and it is also noted that \"trends in 2D PDE data are more difficult to learn without labels\".\n\n    - Previous models have considered pretraining on a much larger scale and with PDEs that are practically more relevant (see the datasets used in the pretraining foundation models mentioned above). Such scaling also seems necessary to verify the claim that \"With enough PDE data during pretraining, freezing MAE encoders can be more practical since the pretraining distribution would cover most downstream cases\".\n\n\n3. **Tasks:** Artificial downstream tasks and preliminary latent PDE arithmetic.\n\n    - Some downstream tasks seem quite artificial and not of practical interest, such as classifying boundary conditions or predicting resolutions (which are also both easily done from the input).\n\n    - Results on latent PDE arithmetic are very preliminary and it is unclear if it can be used for meaningful applications. There seems to be only one qualitative Figure with a single example to validate the latent PDE arithmetic. This is not sufficient to support statements like \"learned representations [...] are semantically meaningful by performing latent PDE arithmetic.\" or \"Additionally, we propose operating in this latent space to interpolate between encoded PDE vectors\".\n\n\n4. **Adaptation to PDEs**: Limited adaptation and theoretical analysis of MAEs in the context of PDEs.\n\n    - It seems like the paper is coming from a computer vision perspective without many PDE-specific adaptations. While this is generally not a shortcoming, it is unclear what is meant by the \"several PDE-specific modifications\" used for pretraining. Discretization-invariance is claimed as an important feature but ViTs are used as the backbone which despite the tweaks mentioned in \"Multi-Resolution pretraining\" cannot be consistently applied to arbitrary resolutions. In particular, there already exist more principled approaches to designing discretization-invariant ViTs (https://arxiv.org/abs/2406.06486).\n\n    - It is not investigated how physically valid masking is. For instance, is there even a unique solution for the considered PDEs or will the model regress the mean?\n\n    - It should be explained what is meant by \"labels\" in the context of PDEs, e.g., in \"unlabeled, and heterogeneous datasets to learn latent physics at scale\" and \"can learn highly structured latent spaces without the need for labels\". For instance, for a time-stepping objective, the same data as used for MAE is sufficient and no labels seem to be needed (see also \"Baselines\" above). Also, it is not clear what effort is meant by \"so no additional effort is needed to generate pretraining labels\". Typical labels, such as the coefficients, boundary conditions, etc., are anyways given when obtaining the data from a numerical solver.\n\n### Minor issues and questions\n\n1. It's not clear why the 1d Heat and Burgers eqs. (with Dirichlet boundary condition) are not \"explicitly seen\" in the 1d KdV-Burgers eq.? They seem to be at least in the training distribution.\n\n2. It should be explained what is meant by \"PDE dynamics information\"?\n\n3. What is an \"output functional space\"?\n\n4. \"We hypothesize that the increased difficulty of the task increases the importance of MAE encoder guidance in time-stepping and super-resolution. This increased difficulty also manifests itself when extrapolating pretrained encoders to new PDEs, as the performance on the Navier-Stokes equations is limited.\" This sounds like a contradiction, since the MAE encoder guidance does *not* help for NS?\n\n5. It is not clear what the purpose of the smoke buoyancy experiment is.\n\n6. Why would we even want different resolutions to be mapped to different parts of the latent space as shown in Figure 4D?\n\n7. Code is not yet available.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This work explores the use of masked autoencoders (MAEs) for learning representations of partial differential equations (PDEs). The key idea is to adapt the successful MAE approach from computer vision to the domain of PDEs, with the goal of learning generalizable representations that can be useful for various PDE-related downstream tasks. In particular, the authors pretrain MAE models on datasets of solutions to various PDEs (KdV-Burgers in 1d and Heat, Advection, and Burgers in 2d), and then evaluate the learned representations on tasks like PDE feature prediction, time-stepping, and super-resolution across different equations. In general, the experiments show a positive impact of pretraining, in particular in combination with additional fine-tuning, and a benefit over contrastive approaches in out-of-distribution settings.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1. *(critical)* How sensitive is the method to hyperparameter choices (see \"Performance\" above)?\n\n2. *(critical)* How does the performance compare to various other pretraining approaches beyond the contrastive learning baseline (see \"Baselines\" above)?\n\n3. How well does this approach scale to much larger and more complex PDE datasets (see \"Datasets\" above)?\n\n4. *(critical)* Can the latent PDE arithmetic be quantitatively evaluated (see \"Tasks\" above)?\n\n5. Can one investigate the theoretical properties of the proposed method (see \"Adaptation to PDEs\")?", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "Wk45M08jTn", "reviewer": "Reviewer_ejJu", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n\nNovelty: The adaptation of masked pretraining to the domain of PDEs is innovative, expanding the potential applications of this self-supervised learning method.\n\nComprehensive Evaluation: The authors evaluate the approach on a range of PDEs, including previously unseen equations, and cover diverse tasks such as time-stepping and super-resolution.\n\nPractical Implications: Demonstrating improvements in both time-stepping accuracy and super-resolution highlights the potential utility of MAEs in real-world physics applications.\n\n\nWeakness:\n\nLimited Extrapolation: Although the MAE model shows generalization to some unseen PDEs, performance declines significantly when applied to more complex equations like the Navier-Stokes equations.\n\nAmbiguity in Latent Arithmetic Results: While latent PDE arithmetic is an exciting idea, the paper does not clearly quantify how closely the latent solutions match true solutions.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper proposes an innovative approach using Masked Autoencoders (MAEs) for solving Partial Differential Equations (PDEs). The authors adapt self-supervised masked pretraining, a method typically used in natural language processing and computer vision, to physics problems. The primary contributions include demonstrating that MAEs can effectively generalize to unseen PDEs, improve performance in downstream tasks like PDE coefficient regression, and enhance the super-resolution of neural solvers. The authors present a framework for latent PDE arithmetic, showing that learned representations from MAEs are meaningful and can operate in a latent space.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Major concerns:\n\n1. Provide More movitation and Additional Quantitative Results for Latent Arithmetic: \n\nWhat is the rationale of doing latent arithmetic in the context of learning PDE solutions? Although viscous Burgers Eqn can be treated as the linear combination of invicid Burgers and Heat Eqn, it does not necessarily mean that the same holds in the latent space, mainly because of the nonlinearity of MAE and the PDE. The author should provide some justification/motivation of the latent arithmetics. \nPage 7 bottom: the author claim “MAE encoer/decoder has never seen the heat/burgers eqn” is this correct? Shouldn’t the MAE be pre-trained on kdv-burgers? In Appendix B.1, alpha, beta, and gamma are allowed to be zero so KDV-Burger actually contains heat and burgers equation.\n\nIncluding a quantitative comparison between reconstructed solutions from latent arithmetic and actual PDE solutions could strengthen the evidence for this method's effectiveness. For example, could you include viscous Burgers in Figure 4C and see if the cluster lie in the space spanned by the clusters of invicid Burgers and Heat Eqn?\n\n2. Justification for classification of BC: \n\nMy main concern is that classifying the Boundary condition from PDE solution data is not a well-posed problem. For a fixed PDE solution, one can calculate its Dirichlet boundary data and Neumann boundary data and in general Robin boundary data. What is point to infer boundary condition type for the PDE data? As the author pointed out in Figure 4 E, Solutions generated from Dirichlet and Neumann BC are mixed. This is quite natural because a Dirichlet BC f and a different Neumann BC g may generate the same PDE solution u.\n\n3. Discuss Computational Requirements: \n\nA section on the computational costs associated with training MAEs on multi-resolution and diverse PDE data would be helpful for readers considering implementation.\n\n4. Introduce Time-stepping and Super-Resolution: \n\nA brief introduction on the problem setup and main goals of time-stepping and super-resolution should be provided.\n\n\nMinor issues:\n\nNo year in reference [54]\n\nThe term Heat-Burgers PDE is very weird. It should be called viscous Burgers equation.\n\nThe last sentece in the paragraph above section 4.4 is not clear.\n\nSection 4.5, reference [64] was not shown correctly.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "R8KmcGeYda", "reviewer": "Reviewer_9TVg", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "__Strengths__ \n\n1. The effort to adapt pretraining methods to PDEs is timely and seems like a reasonable idea. \n2. The breadth of experiments are impressive and creative. \n3. Claims are fairly measured and do not overstate the results outside of the OOD claims discussed below. Areas where the method performed poorly are shown clearly. \n\n__Weaknesses__\n1. This is currently not written for a journal format. Understanding the specifics of the paper requires significant jumps between the main body and the appendix for details. A journal submission should contain all required information to understand precisely what was done and why. Particular technical details or extra experiments that are not core to the paper's goals can be left to the appendix, but simple information like patching approaches or PDE parameter ranges are core to the understanding the paper content. \n2. The paper currently lacks motivation. There is never any overarching goal mentioned or real world problem looking to be addressed besides that other ML researchers are working in the space yet comparisons to those other works are highly limited. This is sufficient for other members of that community who are already familiar with what their goals are, but a reader currently does not come away from this understanding why individuals are using ML to solve PDEs. Some specific examples:  \n      1. _\"any important advances have been made to address these challenges, with SOTA models outperforming numerical solvers on well-studied PDEs under certain setups, proposing error bounds, and being extended to solve real-world problems\"_- This is very vague. How are they outperforming numerical methods? Are they faster? Are they more accurate? Have they been able to be used on new problems that numerical methods didn't apply to? Note that a number of the citations here have been referenced as unfair numerical comparisons in [recent work](https://www.arxiv.org/pdf/2407.07218) discussing baselines used in this space. The stated goals of the papers can certainly be used to motivate more work in the space, but we should be cautious when stating that one method has outperformed others. \n      2. Many of the experiments come across as arbitrary. They do not seem tied to a real problem someone who might use your method might have. For instance, ablation results are typically listed in terms of reconstruction error - this is a synthetic task used to avoid requiring labels. The result that would be interesting to readers is downstream performance.    \n      3. However, downstream tasks also suffer from this issue. The addition of two latent states, for example, even with a background in this area, I cannot really see a case where this would be practically useful even if the the target task was not a subset of the training data (addressed below). This also applies to the resolution prediction tasks - users know the resolution of their data. It actually seems somewhat troubling that the model is able to recover this data and something that should discussed.\n3. The submission seems to use in and out of distribution labels inaccurately. For example, if we look at 1D data, the training set is KdV-B. This is a superset of the KdV, Burgers, and heat equation. Burgers is a non-linear extension of 1D advection. The gap between advection and the training data seems to be the rate of advective transport in the coefficient sampling rather than the equation itself. \n4. The resulting 1D visualizations all seem quite similar. Because the axes aren't labeled, it's difficult to determine the exact pattern, but there seems to be very little transport across the domain in any of the examples besides 1D advection and KS. \n5. Figures and tables are at times underlabeled. In figure 2, one can't see which direction is forward in time. The colors make two losses in one table readable, but splitting them out might make it clearer. However, table 3 just uses different losses in different rows. \n6. Terminology is sometimes inaccurate. For instance \"Heat-Burgers\" is just viscous Burgers - a term already used elsewhere. Dirichlet and Neumann BCs are entire families of BCs rather than specific BCs. The specific families are included in the appendix, but this should be made clear in the main text.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The submission describes the results of adapting Masked Autoencoding Pretraining to the PDE setting. They test the approach across a variety of basic 1D and 2D PDEs and find somewhat mixed performance. The pretrained encoder effectively distinguishes between PDEs, and as consequence also effectively predicts their properties. It also does well on in-distribution downstream tasks but does poorly OOD. The submission also includes useful ablations on choices made for the pretraining model.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Overall the breadth of the experiments is impressive and suggests there is content here meriting eventual acceptance, but the submission is not currently ready for publication. The main issues are the heavily truncated conference-style text and some factual errors around the differences between the PDEs. \n\n__Critical__\n\n1. Rewrite this into a journal paper - it's clear that the current submission is written to be a space-limited conference paper. Given the less restrictive page limits in TMLR, the authors should de-compress the submission significant. Architectural and patching decisions, what the setup is for super resolution or various tasks, and others should not require digging into the appendix. The final paper should be able to be read straight through with the user understanding all vital details. I'd really recommend using the space that's available. Define losses. Describe architectural choices. Include figures if it seems difficult to get points across in text. \n2. It is fine to include the current tests as \"in distribution\" but transfer from KdV-Burgers to subsets of the same equation is not generalization testing. Either update the text and explore around parameter ranges or focus on true OOD tasks like KS and NS. \n3. Label axes on figures. Especially given the limited diversity in the dynamics and initial conditions, it is honestly very hard to tell which direction is \"forward\" in time. \n4. Tie more results to downstream tasks - This connects to the motivation issues. For the paper to be interesting, the results need to be tied to tasks that researchers are actually trying to solve. The current task selection is diverse, but many of these are toy tasks that previous works have not tried to address and there's nothing explaining why we should care about these tasks. For the results to merit inclusion in the main text, they need to be connected to a concrete goal or tell us something interesting about how the model works. \n\n__Stronger paper__ \n\n1. The intro could use a pretty serious overhaul. Right now the motivation for this work is reads as \"other people are doing things in this area\". Strong motivation can make the paper significantly more interesting to readers who aren't finding this simply to look at MAE results for PDE data. \n2. Find ways to include ablation-like content in the main text (along with 4 in critical requests). These are some of the more interesting results to readers who DO work in this space. This is a largely empirical paper - it is valuable for people working in similar spaces to see how your choices impacted downstream tasks. \n3. Compare your results to other pretraining methods. While it's a new space, at this point there are 10+ \"foundation models\" for PDEs many of which have open source implementations and weights. For an empirical paper, it would be extremely helpful to readers to better understand how your approach compares to those on the tasks for which they're designed. It is not important for acceptance that the proposed method outperform them. \n4. Improve the discussion of the results - a lot of the discussion of results right now focuses on what is in and out of distribution, but based on the notes on KdV-B, it does not seem to be accurate. The paper currently does a good job of documenting experiments, but it doesn't provide insight into what works or why. \n5. The impact of training in \"latent space evaluation\" is currently unclear. There are no baselines. Does the raw \"pixel-space\" data also cluster in this way? Without such a baseline, we cannot tell if this clustering of the autoencoder result is useful. For instance, if we found that pixel-space results could separate Dirichlet and Neuman BCs then this would indicate that the model is discarding important information.\n6. Focus more on 2D and extend to 3D. 1D results lack impact. Numerical solvers for 1D PDEs are extremely fast to the point where solver speed for 1D PDEs is only a bottleneck for special cases like stiff systems. On top of this, very few real-world problems can be approximated accurately in 1D. There are even many documented cases where the same equations behave differently in different dimensions (turbulence in Navier-Stokes being the classic example here). 1D are often used for testing very specific phenomena - wave-front tracking for instance, but general prediction accuracy in 1D isn't that interesting unless you are specifically highlighting areas numerical methods do poorly.", "broader_impact_concerns": "No broader impact concerns.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: While there were 2 leaning reject recommendations, these appeared to be grounded in a feeling that the amount of novelty was lacking and that there was not enough extensive experimentation. I agree these are limitations, but they do not meet the criteria for rejection from TMLR. The third reviewer, who recommended leaning accept, noted that the paper satisfied both criteria of acceptance by TMLR. In going over the reviewers and authors discussion, I find myself agreeing that both evidence and audience requirements are satisfied. Therefore, I recommend accept. Given that third reviewer noted several additional points (see below) that could be addressed by the authors to further strengthen their paper, I am recommending accept with minor revision, so that the authors can address these points.
-----------------------------------------
(Additional points from Reviewer 3):
1. Still no direct comparisons to other pretraining approaches/foundation models so it's unclear how MAE compares to these in any task.
2. "However, encoder-style approaches such as this work or contrastive PDE encoders (Mialon et al., 2023; Zhang et al., 2023) are more flexible than foundation models, capable of being applied to arbitrary downstream architectures and different fine-tuning tasks." does not really have any backing evidence since there are not comparisons.
3. In general, statements about how MAE is applicable to all data types are undercut when the paper uses augmentation strategies that require knowledge of the specific equations and which Lie symmetries apply. I don't believe MAE requires these so the statement is still valid, but the particular approach used in this paper may not generalize blindly.

The first two are the more notable ones, but while ultimately addressing them would make the paper far stronger and significantly more interesting, what is here does currently meet the two major criteria.
Decision: Accept (poster)