Self-Supervised Learning with Lie Symmetries for Partial Differential Equations
Authors: Grégoire Mialon
Quentin Garrido
Hannah Lawrence
Danyal Rehman
Yann LeCun
Bobak Kiani
Venue: NeurIPS 2023 poster (2023)
Keywords: Self-supervised learning
partial differential equations
Lie symmetries
data augmentation
TL;DR: We apply self-supervised learning to partial differential equations, using the equations' Lie point symmetries as augmentations.

Abstract: Machine learning for differential equations paves the way for computationally efficient alternatives to numerical solvers, with potentially broad impacts in science and engineering. Though current algorithms typically require simulated training data tailored to a given setting, one may instead wish to learn useful information from heterogeneous sources, or from real dynamical systems observations that are messy or incomplete. In this work, we learn general-purpose representations of PDEs from heterogeneous data by implementing joint embedding methods for self-supervised learning (SSL), a framework for unsupervised representation learning that has had notable success in computer vision. Our representation outperforms baseline approaches to invariant tasks, such as regressing the coefficients of a PDE, while also improving the time-stepping performance of neural solvers. We hope that our proposed methodology will prove useful in the eventual development of general-purpose foundation models for PDEs.

Self-Supervised Learning with Lie Symmetries for Partial Differential Equations
Grégoire Mialon
Quentin Garrido Meta, FAIR Univ Gustave Eiffel, CNRS, LIGM
Hannah Lawrence Meta, FAIR MIT
Meta, FAIR
Danyal Rehman MIT
Yann LeCun
Bobak T. Kiani
Meta, FAIR NYU
MIT
Abstract
Machine learning for differential equations paves the way for computationally efficient alternatives to numerical solvers, with potentially broad impacts in science and engineering. Though current algorithms typically require simulated training data tailored to given setting, from heterogeneous sources, or from real dynamical systems observations that are messy incomplete. In this work, we learn general- -prrpos representations of PDEs from heterogeneous data by implementing joint embedding methods for self- supervised learning (SSL), framework for unsupervised representation learning that has had notable success in computer vision. Our representation outperforms baseline approaches to invariant tasks, such regressing the coefficients of a PDE, while also improving the time-stepping performance of neural solvers. We hope that our proposed methodology will prove useful in the eventual development of general-purpose foundation models for PDEs.
may instead wish to learn useful information
Introduction
Dynamical systems governed by differential equations are ubiquitous in fluid dynamics. chemistry.
astrophysics, and beyond. Accurately analyzing and predicting the evolution of such systems of paramount importance, inspiring decades of innovation in algorithms for numerical methods. However, high-accuracy solvers are often computationally expensive. Machine learning has recently arisen an alternative method for analyzing differential equations fraction of the cost 22 3). Typically, the neural network for a given equation is trained on simulations of that same equation, generated by numerical solvers that are high-accuracy but comparatively slow 4 What if we instead wish learn from heterogeneous data, e.g., data with missing information, or gathered from actual observation of varied physical systems rather than clean simulations?
For example, we may have access to a dataset of instances of time-evolution, stemming from a family of partial differential equations (PDEs) for which important characteristics of the problem, such as viscosity or initial conditions, vary or are unknown. this case, representations learned from such a large, "unlabeled" dataset could still prove useful in learning to identify unknown characteristics, given only a small dataset "labeled" with viscosities or reaction constants. Alternatively, the "unlabeled" dataset may contain evolutions over very short periods of time, or with missing time intervals; possible goals are then to learn representations that could be useful in filling in these gaps, or regressing other quantities of interest.
"Correspondence to: gmialon @ meta.com garridog@meta.com and bkiani @ mit.edu
Equal contribution
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Unlabeled Data
Self-Supervised Learning
Learned Representation Map f
Application to Downstream Tasks
Dog
Bird
Augment
Augment
Fi
Classification "Dog"
Images
F2
F2
Detection
Fd
Fd
.001
.003
PDE
Augment:
Augment
'i
T
Regression 0.003
(Burgers' Equation)
F2
12
Time-stepping
Figure I: A high-level overview of the self- sulf-supervised learning pipeline, in the conventional setting of image data (top row) as well as our proposed setting of a PDE (bottom row). Given a large pool of unlabeled data, supervsue learning or Lie symmetries for PDEs) to train a network fa to produce useful representations from input images. Given smaller of labeled data, these representations can then be used as inputs supervised learning pipeline, performing tasks such as predicting class labels (images) or regressing the kinematic viscosity y (Burgers' equation). Trainable steps are shown with red arrows; importantly the representation function learned via SSL is not altered during application to downstream tasks.
augmentations (e.g. color-shifting for images,
To tackle these broader challenges, we take inspiration from the recent success of self-supervised learning (SSL) as tool for learning rich representations from large, unlabeled datasets of text and images [5 6) Building such representations from and for scientific data is a natural next step in the development of machine learning for science 77]. In the context of PDEs, this corresponds to learning representations from a large dataset of PDE realizations "unlabeled" with key information (such as kinematic viscosity for Burgers' equation), before applying these representations solve downstream tasks with a limited amount data (such as kinematic viscosity regression), as illustrated Figure [
To do so, we leverage the joint embedding framework [8] for self-supervised learning. popular paradigm for learning visual representations from unlabeled data [9] [0]. It consists of training an encoder to enforce similarity between embeddings of two augmented versions of given sample to form useful representations. This is guided by the principle that representations suited to downstream tasks (such as image classification) should preserve the common information between the two augmented views. For example, changing the color of an image of a dog still preserves its semantic meaning and we thus want similar embeddings under this augmentation. Hence, the choice of augmentations is crucial. For visual data, SSI. relies on human intuition to build hand-crafted augmentations (e.g. recoloring and cropping). whereas PDEs endowed with group of symmetries preserving the governing equations of the PDE 11 12]. These symmetry groups are important because creating embeddings that invariant under them would allow capture the underlying dynamics of the PDE. For example, solutions to certain PDEs with periodic boundary conditions remain valid solutions after translations in time and space. There exist more elaborate equation- specific transformations as well, such as Galilean boosts and dilations (see Appendix D Symmetry groups are sell-studied for common PDE families, and can be derived systematically or calculated from computer algebra systems via tools from Lie theory [1 14].
We present general framework for performing SSL for PDEs using their corre- sponding symmetry groups In particular, we show that by exploiting the analytic group transfor- mations from one PDE solution another, we can use joint embedding methods to generate useful representations from large, heterogeneous PDE datasets. We demonstrate the broad utility of these representations on downstream tasks, including regressing key parameters and time-stepping. on simulated physically-motivate datasets. Our approach is applicable to any family of PDEs, harnesses the well-understood mathematical structure of the equations governing PDE data - a luxury not typically available in non-scientific domains and demonstrates more broadly the promise of
Contributions:

---

Supervised downstream task
Self-supervised pretraining
fe
r
y
fo
he
Lreg
Frozen
Trained
Saim
Representation conditioned time-stepping
y
P
g
fo
ho
L
fo
Trained
Frozen
Figure 2: Pretraining and evaluation frameworks, illustrated on Burgers' equation. (Leftt Self- supervised pretraining. We generate augmented solutions æ and x' using Lie symmetries parametrized by and g before passing them through an encoder fe. yielding representations y. The represen- tations are then input to a projection head ha, yielding embeddings z, on which the SSL loss applied. (Right) Evaluation protocols for our pretrained representations y. On new data, we use the computed representations to either predict characteristics of interest, or to condition a neural network operator to improve time-stepping performance.
adapting self-supervision to the physical sciences. We hope this work will serve as starting point for developing foundation models more complex dynamical systems using our framework.
2 Methodology
We now describe our general framework for learning representations from and for diverse sources of PDE data, which can subsequently be used for wide range of tasks, ranging from regressing characteristics of interest of a PDE sample improving neural solvers. To this end, we adapt popular paradigm for representation learning without labels: the joint-emeedding self-supervised learning.
2.1 Self-Supervised Learning (SSL)
Background: In the joint-embedding framework, input data is transformed into two separate "views", using augmentations that preserve the underlying information in the data. The augmented views are then fed through a learnable encoder, fe. producing representations that can used for downstream tasks. The SSL loss function is comprised of similarity loss Ssi between projections (through a projector ha, which helps generalization (15)) of the pairs of views, to make their representations invariant augmentations, and regularization loss Lieg. to avoid trivial solutions (such as mapping all inputs to the same representation). The regularization term can consist of a repulsive force between points, regularization the covariance matrix of the embeddings. Both function similarly, as shown in [16]. This pretraining procedure is illustrated in Fig. (left) in the context of Burgers' equation.
In this work, we choose variance-invarianc--ovvaa supervised loss function [9]. Concretely, let Z, Z' NxD contain the D-dimensional repre- sentations of two batches of N inputs with x centered covariance matrices, Cov(2 Z) and Cov (zv. Rows Z, and Z are two views of shared input. The loss over this batch includes term to enforce similarity (Cam) and a term to avoid collapse and regularize representations (Creg) by pushing elements of the encodings to be statistically identical:
regularization (VICReg) as our self-
Xine
Xreg Clv -11 Cov(Z') I12:).
C(Z. Z') 2
Zt.:
Z/..122 +
(1)
Cra (Z) (Zn)
Com 22,
3

---

Temporal Shift, Si (f,
Spatial Shift, 22 (f,a,
Galilean Boost, da (f.
Time, 2
tū + Ou
Space,
Figure 3: One parameter Lie point symmetries for the Kuramoto-Sivashinsky (KS) PDE. The transformations (left to right) include the un-modified solution (u), temporal shifts (g)). spatial shifts (92), and Galilean boosts (ga) with their corresponding infinitesimal transformations in the Lie algebra placed inside the figure. The shaded red square denotes the original (ar, t), while the dotted line represents the same points after the augmentation is applied.
where ||r denotes the matrix Frobenius norm and Xino, Areg E IR are hyperparameters to weight the two terms. In practice, VICReg separates the regularization freg (Z) into two components to handle diagonal and non-diagonal entries Cov( Z)) separately. For full details, see Appendix C
Adapting VICReg learn from PDE data: Numerical PDE solutions typically come in the form of tensor of values, along with corresponding spatial and temporal grids. By treating the spatial and temporal information as supplementary channels, we can use existing methods developed for learning image representations. As an illustration, a numerical solution to Burgers consists of a velocity tensor with shape (t, r):: a vector of time values, and vector of spatial values. We therefore process the sample to obtain (3, f, x) tensor with the last two channels encoding spatial and temporal discretization, which can be naturally fed neural networks tailored for images such as ResNets [17]. From these, we extract the representation before the classification layer (which unused here). is worth noting that convolutional neural networks have become ubiquitous in the literature I& [2]. While the VICReg default hyper-parameterss did not require substantial tuning, tuning was crucial to probe the quality of our learned representations to monitor the quality of the pre treirrann step. Indeed, SSL loss values are generally not predictive of the quality of the representation, and thus must be complemented by an evaluation task. In computer vision, this is done by freezing the encoder, and using the features to train linear classifier framework, we pick regression of a PDE coefficient, or regression of the initial conditions when there is coefficient in the equation. The latter, commonly referred the advantage being applicable to any PDE, and is often a challenging problem the numerical methods community given the ill- ill-posed nature of the problem (19). Our approach for particular task, kinematic viscosity regression, is schematically illustrated in Fig. 2 (top right). More details on evaluation tasks are provided in Section
ImageNet. In
the inverse problem, has
2.2 Augmentations and PDE Symmetry Groups
Background: PDEs formally define systems of equations which depend on derivatives of input variables. Given input space Ω and output space 24 PDE A is system of equations in independent variables æ E 22 dependent variables u : - 24 and derivatives (Uay Ugaa.... u with respect to x. For example, the Srramoto-Sivashinse
equation is given by
A(z, ł, u) + uw, + Man
0.
Informally, a symmetry group of PDE G acts on the total space via smooth maps G Ω x U Ω 1 taking solutions of A other solutions of A. More explicitly, G is contained in the symmetry group of A if outputs of group operations acting on solutions are still a solution of the PDE:
A(z, u) = = A [g- (z, u)] = 0, Vg E G.
AA group G a set closed under an associative binary operation containing an identity element e and inverses (i.e. e E G and Vig E G - G). G X + X acts on space X if Vo E X, Vg. h E G er and (gh)æ g(hx).

---

For PDEs, these symmetry groups can be analytically derived [1]] (see also Appendix for more formal details). The types of symmetries we consider are so-called Lie point symmetries g : 92 x + Ω xU,, which act smoothly at any given point in the total space 52 x U. For the Kuramoto-Sivashinssy PDE, these symmetries take the form depicted in Fig. 3
Temporal Shift: (1(c) :(z, E, u) + (z, 1 + e, u)
Spatial Shift: 92(e) =(x,, 1, u) + (2 + 6,2,
Galilean Boost: 9s(c) :(2z, ł, u) + (z + ct, ut
As in this example, every Lie point transformation can be written as a one parameter transform of e E R where the transformation at c O recovers the identity map and the magnitude of corresponds to the "strength" " of the corresponding augmentation |
Taking the derivative of the transformation at = 0 with respect to the set of all group transformations recovers the Lie algebra of the group (see Appendix A). Lie algebras are vector spaces with elegant properties (e.g., smooth transformations can be uniquely and exhaustively implemented), so parameterize augmentations in the Lie algebra and implement the corresponding group operation via the exponential map from the algebra to the group. Details are contained in Appendix
PDE symmetry groups as SSL augmentations, and associated challenges: Symmetry groups of PDEs offer a technically sound basis for the implementation of augmentations: nevertheless, without proper considerations and careful tuning, SSL can fail to work successfully (20). Although we find the marriage of these PDE symmetries with SSL quite natural, there are several subtleties to the problem that make this task challenging. Consistent with the image setting, we find that, among the list of possible augmentations, crops are typically the most effective of the augmentations in building useful representations (21]. Selecting a sensible subset of PDE symmetries requires some care; for example, if one has a particular invariant task in mind (such regressing viscosity), the Lie symmetries used should neither depend on viscosity nor change the viscosity of the output solution. Morever, there is no guarantee to which Lie symmetries to produce solutions that are close to the original data distribution; this is also likely a confounding factor when evaluating their performance. Finally, precise derivations of Lie point symmetries require knowing the governing equation, though a subset of symmetries can usually be derived without knowing the exact form of the equation, as certain families of PDEs share Lie point symmetries and many symmetries arise from physical principles and conservation laws.
the most "natural", i.e. most likely
Sampling symmetries: We parameterize and sample from Lie point symmetries in the Lie algebra of the group, to ensure smoothness and universality of resulting maps in some small region around the identity. We use Trotter approximations of the exponential map, which are efficiently tunable to small errors, to apply the corresponding group operation to an element in the Lie algebra (see Appendix B [2 23). In our experiments, we find that Lic point augmentations applied at relatively small strengths perform the best (see Appendix H), as they are enough to create informative distortions of the input when combined. Finally, boundary conditions further complicate the simplified picture of PDE symmetries, and from a practical perspective, many of the symmetry groups (such as the Galilean Boost in Fig. 3 require a careful rediscretization back to regular grid of sampled points.
3 Related Work
In this section, we provide concise summary of research related to our work, reserving Appendix D for more details. Our study derives inspiration from applications of Self-Supervised Learning (SSL)
in building pre-trained foundational models 24.. For physical data, models pre-trained with SSL have been implemented in areas such as weather and climate prediction (7] and protein tasks 25 26 but none have previously used the Lie symmetries of the underlying system. The SSL techniques we use are inspired by similar techniques used in image and video analysis ( [2)]. with the hopes of learning rich representations that can be used for diverse downstream tasks.
Symmetry groups of PDEs have a rich history of study [11]3] Most related to work, 12 used Lie point symmetries of PDEs as a tool for augmenting PDE datasets in supervised tasks. For some PDEs, previous works have explicitly enforced symmetries or conservation laws by for example constructing networks equivariant to symmetries of the Navier Stokes equation (27), parameterizing
'Technically, is the magnitude and direction of the transformation vector for the basis element of the corresponding generator in the Lie algebra.
5

---

networks to satisfy a continuity equation (2), or enforcing physical constraints in dynamic mode decomposition [29]. For Hamiltomian systems, various works have designed algorithms that respect the symplectic structure or conservation laws of the Hamiltonian R0 31
4 Experiments
Equations considered: focus on flow-related equations here as testing ground for our methodology. In our experiments, we consider the four equations below, which are ID evolution equations apart from the Navier-Stokes equation, which we consider in its 2D spatial form. For the ID flow-related equations, we impose periodic boundary conditions with ! [0, L] x [O, For Navier-Stokes, boundary conditions are Dirichlet (v = 0) as in [8]. Symmetries for all equations are listed in Appendix
1. The viscous Burgers' Equation, written in its "standard" form, is nonlinear model of dissipative flow given by
U1 + utl Vule 0,
(5)
where u(x, t) is the velocity and E R+ the kinematic viscosity.
2. The Korteweg-de Vries (KdV) equation models waves on shallow water surfaces as
U4 WMr
Mxxx
(6)
where u(x, t) represents the wave amplitude. 3. The Kuramoto-Sivashinsky (KS) equation is a model of chaotic flow given by
Uy uu, + Usz Uyars
O,
where u(x, t) is the dependent variable. The equation often shows up in ceaction-diffusii systems, as well as flame propagation problems.
4. The incompressible Navier-Stokes equation in two spatial dimensions given by
U - u Vu - -Vp+uVut
-,
(8)
where u(x, is the velocity vector, p(z, is the pressure, p is the fluid density, the kinematic viscosity, and f is an external added force (buoyancy force) that we aim to regress in our experiments.
Solution realizations are generated from analytical solutions in the of Burgers' equation pseudo-spectral methods used to generate PDE learning benchmarking data (see Appendix 12 I8 (32. Burgers', KdV and KS's solutions generated following the process of [12] while for Navier Stokes we use the conditioning dataset from [18]. The respective characteristics of our datasets can be found in Table
Pretraining: For each equation, we pretrain a ResNer18 with our SSL framework for 100 epochs using AdamW [33) batch size of 32 (64 for Navier-Stokes) and learning rate of 3e-4. We then freeze its weights. To evaluate the resulting representation, we (i) train linear head on top of our features and on new of labeled realizations, and (ii) condition neural networks for time-stepping on our representation. Note that our encoder learns from heterogeneous data in the sense that for given equation, we grouped time evolutions with different parameters and initial conditions.
4.1 Equation parameter regression
We consider the task of regressing equation-related coefficients in Burgers' equation and the Navier- Stokes' equation from solutions to those PDEs. For KS and KdV we consider the inverse probem of regressing initial conditions. We train a linear model on top of the pretrained representation for the downstream regression task. For the baseline supervised model, ResNet18, using the MSE loss on downstream labels. Unless stated otherwise, we train the linear model for 30 epochs using Adam. Further details in Appendix F
train the same architecture, i.e.
Kinematic viscosity regression (Burgers): We pretrain ResNet18 on 10, 000 unlabeled realizations of Burgers' equation, and use the resulting features to train linear model smaller, labeled

---

Table 1: Downstream evaluation of our learned representations for four classical PDEs (averaged over three runs, the lower the better (U)). The normalized mean squared error (NMSE) over a batch of N outputs üg and targets UR is equal to NMSE * Σx= ق 1 1x. Relative error is similarly defined as RF = N Σk=1 |Í with supervised methods are the best performance across runs with Lie symmetry augmentations applied. For timestepping. we report NMSE for KdV, KS and Burgers as in 12, and MSE for
uA |I / | For regression tasks, the reported errors
Navier-Stokes for comparison with [18].
Equation
KdV
KS
Burgers
Navier-Stokes
SSL dataset size
10,000
10,000
10,000
26,624
Sample format (1, r, (v))
256x 128
256x 128
448x224
56x 128x 128
Characteristic of interest
Init. coeff's
Init. coeffs
Kinematic viscosity Relative error %6()
Buoyancy MSE (J)
Regression metric
NMSE (4)
NMSE (1)
Supervised
0.102 + 0,007
0.117 + 0.009
1.18 + 0.07
0,0078 + 0.0018
SSL repr. linear head
0.033 + 0.004
0.042 + 0.002
0.97 + 0.04
0.0038 + 0.0001
Timestepping metric
NMSE (4)
NMSE (U)
NMSE (L)
MSE x x10 (()
Baseline
0.508 + 0.102
0.549
0.095
0.110 0.008
2.37 + 0.01
+ repr. conditioning
0.330 + 0.081
0.381 + 0.097
0.108 + 0.011
2.35 + 0.03
dataset of only 2000 samples. We compare to the same supervised model (encoder and linear head) trained on the same labeled dataset. The viscosities used range between 0.001 and 0.007 and are sampled uniformly. We can see in that we are able to improve over the supervised baseline by leveraging our learned representations. This remains true even when also using Lie Point symmetries for the supervised baselines or when using comparable dataset sizes, as in Figure 4 We also clearly see the ability of our self- self-supervised approach leverage larger dataset sizes, whereas we did not see any gain when going to bigger datasets in the supervised setting.
Initial condition regression (inverse problem): For the KS and Kdv PDEs, we aim solve the inverse problem by regressing initial condition parameters from a snapshot of future time evolutions of the solution. Following [3 [1]] for domain f [0, L]. a truncated Fourier series, parameterized by As, Wk, Ok. is used to generate initial conditions:
N
(TTwxi L
uo (x)
Σ
As sin
+ Φz
Our task is regress the set of 2N coefficients {Ag,we: h {1...., N}} from snapshot of the solution starting at t = 20 to L = T. This way, the initial conditions and first-time steps never seen during training, making the problem non-trivial. For all conducted tests, N = 10, A& ~ ((00.5 0.5), and W/ ~ ~u(-0.4. 0.4). By neglecting phase shifts, Φk. the inverse problem is invariant to Galilean boosts and spatial translations, which we use as augmentations for training our SSL method (see Appendix [E. The datasets used for KdV and KS contains 10,000 training samples and 2,500 test samples. As shown in Table the SSL trained network reduces NMSE by a factor of almost three compared to the supervised baseline. This demonstrates how pre-training via SSL can help to extract the underlying dynamics from a snapshot of a solution.
Buoyancy magnitude regression: Following 18,, our dataset consists of solutions of Navier Stokes (Equation (8)) where the external buoyancy force, directions over the course of given evolution, and our aim is to regress the magnitude of this force
(c2. Gyy) T, is constant in the two spatial
V + cg given a solution to the PDE. We reuse the dataset generated in [8] where e, = 0 and Cy 0.... In practice this gives us 26,624 training samples that we used as our "unlabeled" dataset, 3,328 to train the downstream task on, and 6,592 to evaluate the models. As observed in T I the self-supervised approach is able to significantly outperform the supervised baseline. Even when looking at the best supervised performance (over 60 runs), or in similar data regimes as the supervised baseline illustrated in Fig. 4 the self- supervised baseline consistently performs better and improves further when given larger unlabeled datasets.
7

---

Table 2: One-step validation MSE (rescaled by 1c3) for time-stepping on Navier-Stokes with varying buoyancies for different combinations of architectures and conditioning methods. Architectures are taken from [8] with the same choice of hyper-paramenn an upper-bound on the performance a representation containing information on the buoyancy.
Results with ground truth buoyancies are
Architecture
UNe4mad64
UNe4mo064
FNO 128modes
Conditioning method
Addition [8
AdaGN 35
Spatial-Speetral
[B
Addition [8
Time conditioning only
2.60 +: 0.05
2.37 1: 0.01
13.4 1: 0.5
3.31 + 0.06
Time SSL repr. cond.
2.47 0.02
2.35 0.03
13.0 + 1.0
2.37 + 0.05
Time true buoyancy cond.
2.08 + 0.02
2.01 + 0.02
11.4 + 0.8
2.87 + 0.03
4.2 Time-stepping
To explore whether learned representations improve time-stepping, we study neural networks that use sequence of time steps (the "history") of a PDE to predict future sequence of steps. For each equation we consider different conditioning schemes, fit within the data modality and be comparable to previous work.
Burgers, Korteweg-de Vries, and Kuramoto-Sivashinsky We time-step on 2000 unseen samples for each PDE. To do so, we compute a representation of 20 first input time steps using our frozen encoder, and add it new channel. The resulting input is fed to CNN in [12 to predict the next 20 time steps (illustrated in Fig. 4 (bottom right) in the context of Burgers' equation). As shown in T conditioning the neural network or operator with pre-trained representations slightly reduces the error. Such conditioning noticeably improves performance for KdV and KS, while the results are mixed for Burgers'. A potential explanation is that KdV and KS feature more chaotic behavior than Burgers, leaving room for improvement.
Navier-Stokes' equation: As pointed out in [18], conditioning a neural network or neural operator on the buoyancy helps generalization accross different values of this parameter. This is done by embedding the buoyancy before mixing the resulting vector either via addition to the neural operator's 's hidden activations (denoted in [8] as "Addition"), or alternatively for UNets by affine transformation of group normalization layers (denoted as "AdaGN" and originally proposed in (35) ). For our main experiment, we use the same modified UNet with 64 channels as in [18] for our neural operator, since it yields the best performance on the Navier-Stokes dataset. To condition the UNet, we compute our representation on the 16 first frames (that are therefore excluded from the training), and pass the representation through a two layer MLP with a bottleneck of size 1, in order to exploit the ability of our representation to recover the buoyancy with only one linear layer. The resulting output is then added the conditioning embedding as in [8]. Finally, we choose AdaGN as our conditioning method, since it provides the best results in [18. We follow a similar training and evaluation protocol to [8]. except that we perform 20 epochs with cosine annealing schedule on 1.664 trajectories instead of 50 epochs, as we did not observe significant difference in terms of results, and this allowed to explore other architectures and conditioning methods. Additional details are provided in Appendix [F As baseline, we use the same model without buoyancy conditioning. Both models are conditioned on time. We report the one-step validation MSE on the same time horizons as [8] Conditioning on our representation outperforms the bascline without conditioning.
We also report results for different architectures and conditioning methods for Navier-Stokes in Table and Burgers in Table (8 (Appendix F.I) validating the potential of conditioning on SSL representations for different models. FNO 36 does not perform as well as other models, partly due to the relatively low number of samples used and the low-resolution nature of the benchmarks. For -tokkrsStokes, we also report results obtained when conditioning on both time and ground truth buoyancy, which serves as an upper-bound on the performance of our method. We conjecture these results can be improved by further increasing the quality of the learned representation, e.g by training on more samples or through further augmentation tuning. Indeed, the MSE on buoyancy regression obtained by SSL features, albeit significantly lower than the supervised baseline, is often still too imprecise to distinguish consecutive buoyancy values in our data.
4.3 Analysis
Self-supervised learning outperforms supervised learning for PDEs: While the superiority of self- supervised over supervised representation learning is still an open question in computer vision [37 38)
8

---

16-2
Viscosity regression
Buoyancy regression
1.20
Supervised SSL w/ LPS
1.15
Moan Squared Error
7.5
7.0
1.05
6.5
supervised
Relative
6.0
Best sup supervised
1.00
SSL LPS
8.5
0.95
5.0
0.90
4.5
0.85
1000
2000
3000 4000
5000
6000
7000 8000
9000
2000
4000
8000
8000 10000 12000 14000 16000 18000 Unlabeled dataset size
Unlabeled dataset size
Figure 4: Influence of dataset size on regression tasks. (Left) Kinematic regression on Burger's equation. When using Lie point symmetries (LPS) during pretraining, we are able to improve performance over the supervised baselines, even when using an unlabled dataset size that is half the size of the labeled one. As we increase the amount of unlabeled data that use, the performance improves, further reinforcing the usefulness of self-supervised representations. (Right) Buoyancy regression Navier-Stokes' equation. We notice similar trend in Burgers but found that the supervised approach was less stable than the self-supervised one. As such, SSL brings better performance as well as more stability here.
the former outperforms the latter in the PDE domain we consider. A possible explanation is that enforeing similar representations for two different views the same solution forces the network to learn the underlying dynamics, while the supervised objectives (such as regressing the buoyancy) may not be as informative of a signal to the network. Moreover, Fig. fillustrates how more pretraining data benefits SSL setup, whereas in our experiments it did not help the supervised baselines.
Cropping is natural, effective, and popular augmentation in computer vision 21 In the context of PDE samples, unless specified otherwise, we crop both in temporal and spatial domains finding such a procedure is necessary for the encoder learn from the PDE data. Cropping also offers typically weaker means of enforcing analogous space and time translation invariance. The exact size of the crops is generally domain dependent and requires tuning. We quantify its effect in Fig-5 in the context of Navier-Stoken, here, crops must contain as much information possible while making sure that pairs of crops have as little overlap as possible (to discourage the network from relying spurious correlations). This explains the two modes appearing in Fig, (5) We make similar observation for Burgers, while KdV and KS are less sensitive. Finally, crops help bias the network learn features that are invariant whether the input was taken near boundary not, thus alleviating the issue of boundary condition preservation during augmentations.
Cropping:
39
Selecting Lie point augmentations: Whereas cropping alone yields satisfactory representations, Lie point augmentations can enhance performance but require careful tuning. order to choose which symmetries to include in our SSL pipeline and at what strengths apply them, study effectiveness of each Lie augmentation separately. More precisely, given an equation and each possible Lie point augmentation, cropping. Then, we couple all Lie augmentations improving the representation over simply using crops. In order for this composition stay in the stability/convergene radius of the Lie Symmetries, we reduce each augmentation' 's optimal strength by an order of magnitude. Fig. 5 illlsttraees this process in the context of Navier-Stokes.
train SSL representation using this augmentation only and
Discussion
This work leverages Lie point symmetries for self-supervised representation learning from PDE data. Our preliminary experiments with the Burgers', KdV, KS, and Navier-Stoke equations demonstrate the usefulness of the resulting representation for sample or compute efficient estimation of characteristics and time-stepping. Nevertheless, number of limitations are present in this work, which we hope can be addressed in the future. The methodology and experiments in this study were confined to particular set of PDEs, but we believe they can be expanded beyond our setting.

---

Buoyancy regresssion MSE 102
Augmentation
Best strength
Buoyancy MSF
0.75
Crop
N.A
0.0051 + 0.0001
32
0.75
0.76
0.70
0.31
0.70
single Lie transform
+ + translate 91
0.1
0.0052 + 0.0001
0.65
Spatall 120
0.70
+ translate 92
10.0
0.0041 + 0.0002
0.60
+ scaling g4
1.0
0.0050 + 0.0003
+ rotation gs
1.0
0.0049 0.0001
0.49
0.69
0.73
0.75
+ boost ge
0.1
0.0047 +: 0.0002
+ boost ge
0.1
0.0046 + 0.0001
0.50
combined
128
0.62
0.61
0.75
0.45
+
[g2, 95, 96. 0s)
best / 10
0.0038 + 0.0001
linear boost applied in direction (see Table
quadratic boost applied direction (see Table
Temporal trup
Figure 5: (Left) Isolating effective augmentations for Navier-Stokes. Note that we do not study 93-
97 and 9o. which are respectively counterparts of 92- Is and Is applied in y instead of z. (Right) Influence of the crop size on performance. We see that performance is maximized when the crops are as large as possible with as little overlap as possible when generating pairs of them.
Learning equivariant representations: Another interesting direction is to expand our SSL frame- work to learning explicitly equivariant features [4] 42). Learning equivariant representations with SSL could be helpful for time-stepping, perhaps directly in the learned representation space.
Preserving boundary conditions and leveraging other symmetries: Theoretical insights can also help improve the results contained here. Symmetries generally derived with respect to systems with infinite domain or periodic boundaries. Since boundary conditions violate such symmetries, we observed in our work that we only able implement group operations with small strengths. Finding ways to preserve boundary conditions during augmentation, even approximately, would help expand the scope of symmetries available for learning tasks. Moreover, the available symmetry group operations of given PDE are not solely comprised of Lie point symmetries, Other types of symmetries, such nonlocal symmetries or approximate symmetries like Lie-Backlund symmetries, may also be implemented as potential augmentations [3]
Towards foundation models for PDEs:
A natural next step for our framework is to train a common representation on a mixture of data from different PDEs, such as Burgers, KdV and KS, that are all models of chaotic flow sharing many Lie point symmetries. Our preliminary experiments are encouraging yet suggest that work beyond the scope of this paper is needed to deal with the different time and length scales between PDEs.
Extension to other scientific data: In our study, utilizing the structure of PDE solutions as "exact" SSI. augmentations for representation learning has shown significant efficca over supervised methods. This approach's 's potential extends beyond the PDEs we study as many problems in mathematics, physics, and chemistry present inherent symmetries that can harnessed for SSL. Future directions could include implementations of SSL for learning stochastic PDEs, Hamiltonian systems. In latter, the rich study of Noether's symmetries in relation to Poisson brackets could a useful setting to study [II]. Real-world data, opposed to simulated data, may offer nice application to the SSL setting we study. Here, the exact form of the equation may not be known and symmetries of the equations would have to be garnered from basic physical principles (e.g., flow equations have translational symmetries), derived from conservation laws, or potentially learned from data.
Acknowledgements
The authors thank Aaron Lou, Johannes Brandstetter, and Daniel Worrall for helpful feedback and discussions. HI. is supported by the Fannie and John Hertz Foundation and the NSF Graduate Fellowship under Grant No. 1745302.
10

---

References
[1] Mazier Raissi, Paris Perdikaris, and George E. Karniadakis. Physics- informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations. Journal of Computational Physics, 778:686-707, 2019. ISSN 0021-9991. URL https://doi org/ 10. 11t./1116.. j.p. 2018. .45
[2] George E. Karniadakis, Ioannis G. Kevrekidis, Lu Lu, Paris Perdikaris, Sifan Wang, and Liu Yang. Physics-informed machine learning. Nature Reviews Physics, 3(6):422-440, 2021. URL
https://doi. 10.
[3] Samuel H Rudy, Steven I. Brunton, Joshua I. Proctor, and J Nathan Kutz. Data-driven discovery of partial differential equations. Science advances, 3(4):e1602614, 20177
[4] Steven L. Brunton, Joshua L. Proctor, and J. Nathan Kutz. Discovering governing equations from data by sparse identification of nonlinear dynamical systems. Proceedings of the National Academy of Sciences, 113(15)33932-3937 2016. URL Wtl.hhtps://doi.... 10. 1073/ pnas 1517384113
[5] Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark. et al. Learning transferable visual models from natural language supervision. In arXiv preprint arXiv: 103.00020, 2021.
[6] Mathilde Caron, Hugo Touvron, Ishan Misra, Hervé Jégou, Julien Mairal, Piotr Bojanowski, and Armand Joulin. Emerging properties in sulf--upe vision transformers. In Proceedings of the IEEE/CVF International Conference Computer Vision, pages 9650-9660, 2021.
[7] Tung Nguyen, Johannes Brandstetter, Ashish Kapoor, Jayesh K Gupta, and Aditya Grover. ClimaX: A foundation model for weather and climate. arXiv preprint arXiv:2301. 1003::330...0333, 2023.
[8] Jane Bromley, Isabelle Guyon, Yann LeCun, Eduard Säckinger, and Roopak Shah. Signature verification using a" siamese" time delay neural network. Advances in neural information processing systems, 6, 1993.
[9] Adrien Bardes, Jean Ponce, and Yann LeCun. Vicreg: Varianne-iivvriancc--vvaan regular- ization for self-supervised learning. arXiv preprint
2021.
[10] Jean-Bastien Grill, Florian Strub, Florent Altché, Corentin Tallec, Pierre Richemond, Elena Buchatskaya, Carl Doersch, Bernardo Avila Pires, Zhaohan Guo, Mohammad Gheshlaghi Azar, et al. Bootstrap your own latent-a new approach to self-supervised learning. NeurIPS, 2020.
[11] Peter J. Olver. Symmetry groups and group invariant solutions of partial differential equations. Journal of Differential Geometry, 14:497-542, 1979.
[12] Johannes Brandstetter, Max Welling, and Daniel E Worrall. Lie point symmetry data augmen- tation for neural pde solvers. arXiv preprint arXh::2202.07643, 2022.
[13] Nail H Ibragimov. CRC handbook of Lie group analysis of differential equations, volume 3. CRC press, 1995.
[14] Gerd Baumann. Symmetry analysis of differential equations with Mathematica Springer Science & Business Media, 2000.
[15] Florian Bordes, Randall Balcstriero, Quentin Garrido, Adrien Bardes, and Pascal Vincent. Guillotine regularization: Improving deep networks generalization by removing their head. arXiv preprint arXiv:2206. 13378, 2022.
[16] Quentin Garrido, Yubei Chen, Adrien Bardes, Laurent Najman, and Yann Lecun. On the duality between contrastive and non-contrastive self-supervised learning. arXiv preprint arXiv:2206.02577 2022.
[17] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of IEFE conference on computer vision and pattern recognition, pages 770-778, 2016.
11

---

[18] Jayesh K Gupta and Johannes Brandstetter. Towards multi saati-pppttotmmpoaasssalll generalized pde modeling. TMLR, 2022.
[19] Victor Isakov. Inverse problems for partial differential equations, volume 127. Springer, 2006.
[20] Randall Balestriero, Mark Ibrahim, Vlad Sobal, Ari Morcos, Shashank Shekhar, Tom Goldstein, Florian Bordes, Adrien Bardes, Gregoire Mialon, Yuandong Tian, et al. A cookbook of self- supervised learning. arXiv preprint arXiv:2304. 12210, 2023.
[21] Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In International conference on machine learning, pages 1597-1607. PMLR, 2020.
[22] Hale F Trotter. On the product of semi-groups of operators. Proceedings of the American Mathematical Sociery, (144):54--51,, 1959.
[23] Andrew M Childs, Yuan Su, Minh C Tran, Nathan Wiebe, and Shuchen Zhu. Theory of trotter enror with commutator scaling. Physical Review X, 11(1):011020, 2021.
[24] Rishi Bommasani, Drew A Hudson, Ehsan Adeli, Russ Altman, Simran Arora, Sydney von Arx, Michael S Bernstein, Jeannette Bohg, Antoine Bosselut, Emma Brunskill, al. On the opportunities and risks of foundation models. arXiv preprint arXiv:2108.07258 2021.
[25] Nadav Brandes, Dan Ofer, Yam Peleg, Nadav Rappoport, and Michal Linial. Proteinbert: a universal deep-learning model of protein sequence and function. Bioinformatics, 38(8): 2102-2110, 2022.
[26] Jianyi Yang. Ivan Anishchenko, Hahnbeom Park, Zhenling Peng, Sergey Ovchinnikov, and David Baker. Improved protein structure prediction using predicted interresidue orientations. Proceedings of the National Academy of Sciences, 117(3): 1060::199661100 1503, 2020.
[27] Rui Wang, Robin Walters, and Rose Yu. Incorporating symmetry into deep dynamics models for improved generalization. arXiv preprint arXiv:2002.03061,, 2020.
[28] Jack Richter-Powell, Yaron Lipman, and Ricky TQ Chen. Neural conservation laws: A divergence- free perspective. arXiv preprint arXiv:2210.01741, 2022.
[29] Peter J. Baddoo, Benjamin Herrmann, Beverley J. McKeon, J. Nathan Kutz, and Steven L.
Brunton. Physics-informed dynamic mode decomposition (pidmd), 2021. URL httpss
Sarxiv.org/abs/2111 013777
[30] Mare Finzi, Ke Alexander Wang, and Andrew G Wilson Simplifying hamiltonian and lagrangian neural networks via explicit constraints. Advances in neural information processing systems, 33: 13880-13888) 2020.
[31] Yuhan Chen, Takashi Matsubara, and Takaharu Yaguchi. Neural symplectic form: learn- ing hamiltonian equations on general coordinate systems. Advances in Neural Information Processing Systems, 34: 16659-1666667, 2021.
[32] Makoto Takamoto, Timothy Praditia, Raphael Leiteritz, Daniel MacKinlay, Francesco Alesiani, Dirk Pflüger, and Mathias Niepert. Pdebench: An extensive benchmark for scientific machine learning. Advances in Neural Information Processing Systems, 35:1596-1611, 2022.
[33] Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint arXiv:1711.051011 2017.
[34] Yohai Bar-Sinai, Stephan Hoyer, Jason Hickey, and Michael P. Brenner. driven discretizations for partial differential equations. Proceedings of the National Academy of Sciences, 116(3))15344-1554) 2019. URL https://doi. http://doooorr///0// www.oo../////01..111
Learning data-
pnas. 1014.1811401111
[35] Alexander Quinn Nichol and Prafulla Dhariwal. Improved denoising diffusion probabilistic models. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research,
pages 8162-8171. PMLR. 18-24 Jul 2021.
12

---

[36] Zongyi Li. Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar. Fourier neural operator for parametric partial differential equations, 2021.
[37] Mert Bulent Sariyildiz, Yannis Kalantidis, Karteek Alahari, and Diane Larlus, No reason for no supervision: Improved generalization in supervised models. In International Conference Learning Representations, 2023.
[38] Maxime Oquab, Timothée Darcet, Théo Moutakanni, Huy Vo, Mare Szafraniec, Vasil Khalidov, Pierre Fernandez, Daniel Haziza, Francisco Massa, Alaaeldin El-Nouby, et al. Dinov2: Learning robust visual features without supervision. arXiv preprint arXiv:2304.07199 2023.
[39] Adrien Bardes, Jean Ponce, and Yann LeCun. VICRegl: Self-super learning of local visual features. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors. Advances in Neural Information Processing Systems, 2022. URL pettps://openreii.e lee//forum?iraa22ssee333.
[40] Xinlei Chen and Kaiming He. Exploring simple siamese representation learning. In CVPR, 2020.
[41] Robin Winter, Marco Bertolini, Tuan Le, Frank Noe, and Djork-Aré Clevert. Unsuper- vised learning of group invariant and equivariant representations. In S. Koyejo, S. Mo- hamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh, editors, Advances in Neural In- formation Processing Systems, volume 35, pages 31942-31 31956. Curran Associates, Inc.. 2022. URL htttps:/pproceediniig newripesrpppp cc/ papera hhppprrflleee//////aaaaaaaaeeeea wwwes//ppeer///eeeeeeeeeeeeaeeet htlle//paper/022/fii22 Paper- Conf erence. eereenee.dff
[42] Quentin Garrido, Laurent Najman, and Yann Lecun. Self-supervised learning of split invariant equivariant representations. arXiv preprint arXiv:2302. 2023.
[43] Janpou Nee and Jinqiao Duan. Limit of trajectories of the coupled viscous burgers' equations. Applied mathematics letters, 11(1):77-1, 1998.
[44] Peter J Olver. Symmetry groups and group invariant solutions of partial differential equations. Journal of Differential Geometry, 14(4):497-542, 1979.
[45] Andrew Baker. Matrix groups: An introduction to Lie group theory. Springer Science & Business Media, 2003.
[46] John Dollard, Charles Friedman, and Pesi Rustom Masani. Product integration with applications to differential equations, volume 10. Westview Press, 1979.
[47] Masuo Suzuki. General theory of fractal path integrals with applications to many-body theories and statistical physics. Journal of Mathematical Physics, 32(2):400 407, 1991.
[48] Robert McLachlan and G Reinout W Quispel. Splitting methods. Acta Numerica, 11: 2002.
[49] Stéphane Descombes and Mechthild Thalhammer. An exact local error representation of exponential operator splitting methods for evolutionary problems and applications to lincar schrödinger equations in the semi-classical regime. BIT Numerical Mathematics, 50(4): 729-749, 2010.
[50] Klaus-Jochen Engel, Rainer Nagel, and Simon Brendle. One-parameter semigroups for linear evolution equations, volume 194. Springer, 2000.
[51] Claudia Canzi and Graziano Guerra. A simple counterexample related to the lic lic-trotter product formula. In Semigroup Forum, volume 84, pages 499-504. Springer, 2012.
[52] Mahdi Ramezanizadeh, Mohammad Hossein Ahmadi, Mohammad Alhuyi Nazari, Milad Sadeghzadeh, and Lingen Chen. A review on the utilized machine learning approaches for modeling the dynamic viscosity of nanofluids. Renewable and Sustainable Energy Reviews, 1 114:109345, 2019.
13

---

[53] William D Fries, Xiaolong He, and Youngsoo Choi. Lasdi: Parametric latent space dynamics identification. Computer Methods Applied Mechanics and Engineering. 399:115436, 2022.
[54] Xiaolong He, Youngsoo Choi, William Fries, Jon Belof, and Jiun-Shyan Chen. glasdi: Parametric physics greedy latent space dynamics identification. arXiv preprint wrXiv:224..12055, 2022.
[55] Rahmad Syah, Naeim Ahmadian, Marischa Elveny, SM Alizadeh, Meysam Hosseini, and Afrasyab Khan. Implementation of artificial intelligence and support vector machine learning to estimate the drilling fluid density in high-pressure high-tempprattre wells. Energy Reports, 141106-4113. 2021.
[56] Ricardo Vinuesa and Steven L Brunton. Enhancing computational fluid dynamics with machine learning. Nature Computational Science, 26)):58-3666, 2022.
[57] Maziar Raissi, Alireza Yazdani, and George Em Karniadakis. Hidden fluid mechanics: Learning velocity and pressure fields from flow visualizations. Science, 1676648)) 2020.
[58] Ehsan Adeli, Luning Sun, Jianxun Wang, and Alexandros Taflanidis. An advanced spatio- temporal convolutional recurrent neural network for storm surge predictions. arXiv preprint arXiv:2204.0990 095011, 2022.
[59] Pin Wu, Feng Qiu, Weibing Feng, Fangxing Fang, and Christopher Pain. A non-intrusive reduced order model with transformer neural network and its application. Physics of Fluids, 34(1 15:30, 2022.
[60] Léonard Equer, T. Konstantin Rusch, and Siddhartha Mishra. Multi-scale message passing neural pde solvers, 2023. URL httLthttp///// //arxiv. org/ 2302 03580
[61] Byungsoo Kim, Vinicius C Azevedo, Nils Thuerey, Theodore Kim, Markus Gross, and Barbara Solenthaler. Deep fluids: A generative network for parameterized fluid simulations. In Computer forum, forum, volume 38(2), pages 59-70. Wiley Online Library, 2019.
[62] Bethany Lusch, J Nathan Kutz, and Steven L Brunton. Deep learning for universal linear embeddings of nonlinear dynamics. Nature communications, (11):9950, 2018.
[63] Peter J Schmid. Dynamic mode decomposition of numerical and experimental data. Journal of fluid mechanics, 656:5-28, 2010.
[64] Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick. Momentum contrast for unsupervised visual representation learning. In Proceedings of the IEEEJCVE conference on computer vision and pattern recognition, pages
2020.
[65] Xinlei Chen, Haoqi Fan, Ross Girshick, and Kaiming He. Improved baselines with momentum contrastive learning. arXiv preprint arXiv:2003. 042977 2020.
[66] Xinlei Chen, Saining Xie, and Kaiming He. An empirical study of training self-supervised vision transformers. In ICCV, 2021.
[67] Chun-Hsiao Yeh, Cheng- Yao Hong, Yen-Chi Hsu, Tyng-Luh Liu, Yubei Chen, and Yann LeCun. Decoupled contrastive learning. arXiv preprint arXiv:2/ 2021.
[68] Aaron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding. arXiv preprint arXiv: ArXivv1887.0333 1807.03778.. 2018.
[69] Jeff Z HaoChen, Colin Wei, Adrien Gaidon, and Tengyu Ma. Provable guarantees for self- supervised deep learning with spectral contrastive loss. NeurIPS, 34, 2021.
[70] Mathilde Caron, Piotr Bojanowski. Armand Joulin, and Matthijs Douze. Deep clustering for unsupervised learning. In ECCV, 2018
[71] Mathilde Caron, Ishan Misra, Julien Mairal, Priya Goyal, Piotr Bojanowski, and Armand Joulin. Unsupervised learning of visual features by contrasting cluster assignments. In NeurIPS, 2020.
14

---

[72] Jure Zbontar, Li Jing, Ishan Misra, Yann LeCun, and Stéphane Deny. Barlow twins: Self- supervised learning via redundancy reduction. In ICML, pages 12310-12320. PMLR, 2021.
[73] Aleksandr Ermolov, Aliaksandr Siarohin, Enver Sangineto, and Nicu Sebe. Whitening for self-supervised representation learning, 2021.
[74] Zengyi Li, Yubei Chen, Yann LeCun, and Friedrich T Sommer. Neural manifold clustering and embedding. arXiv preprint ar arXiv:2201.10000, 2022.
[75] Vivien Cabannes, Bobak T Kiani, Randall Balestriero, Yann LeCun, and Alberto Bietti. The ssl interplay: Augmentations, inductive bias, and generalization. arXiv preprint Xii:22302.02774, 2023.
[76] Grégoire Mialon, Randall Balestriero, and Yann Lecun. Variance-covariance regulariza- tion enforces pairwise independence in self-supervised representations. arXiv preprint arXiv:2209. 14905, 2022.
[77] Lars Schmarje, Monty Santarossa, Aimon-Martin Schröder, and Reinhard Koch. A survey semi-, self-and unsupervised learning for image classification. IEEE Access, 9:82146-824 82168, 2021.
[78] Olmo Cerri, Thong Nguyen, Maurizio Pierini, Maria Spiropulu, and Jean-Roch Vlimant. Variational autoencoders for new physics mining at the large hadron collider. Journal of High Energy Physics, 2019(5):1-29, 2019.
[79] Carl Edward Rasmussen and Hannes Nickisch. Gaussian processes for machine learning (gpml) toolbox. The Journal of Machine Learning Research, 11:3011-3015, 2010.
[80] Mahmut Kaya and Hasan Şakir Bilge. Deep metric learning: A survey. Symmetry, 11(9))1066,
2019.
[81] Zichao Long, Yiping Lu, Xianzhong Ma, and Bin Dong. Pde-net: Learning pdes from data. In International conference on machine learning, pages 3208 3216. PMLR, 2018.
[82] M Giselle Fernández- Godino, Chanyoung Park, Nam-H Ho Kim, and Raphael T Haftka. Review of multi-fidelity models. arXiv preprint arXiv:1609.077
2016.
[83] Alexander IJ Forrester, András Sóbester, and Andy J Keane. Multi-fidelity optimization via surrogate modelling. Proceedings of the royal society a: mathematical, physical and engineering sciences, 463(2088):3551-333 2007.
[84] Leo Wai-Tsun Ng and Michael Eldred. Multifidelity uncertainty quantification using non- intrusive polynomial chaos and stochastic collocation. In 53rd aicalasme//aselaassaass struc- tures, structural dynamics and materials conference 20th aiaa/asme/ahs adaptive structures conference 14th aiaa, page 1852, 2012.
[85] Paris Perdikaris, Maziar Raissi, Andreas Damianou, Neil D Lawrence, and George Em Karni- adakis. Nonlinear information fusion algorithms for data-efficient multi- multti-iielity modelling. Proceedings of the Royal Society A: Mathematical, Physical and Engineering Sciences, 473 (2198):20160751 2017.
[86] Michael Bronstein, Joan Bruna, Taco Cohen, and Petar Veličković. Geometric deep learning: Grids, groups, graphs, geodesics, and gauges. arXiv preprint arXiv:2104.1 2021.
[87] Manzil Zaheer, Satwik Kottur, Siamak Ravanbakhsh, Barnabas Poczos, Russ R Salakhutdinov, and Alexander Smola. Deep sets. Advances in neural information processing systems, 30, 2017.
[88] Risi Kondor and Shubhendu Trivedi. On the generalization of equivariance and convolution in neural networks to the action of compact groups. In International Conference on Machine Learning, pages 2747-2755. PMLR, 2018.
[89] Taco Cohen and Max Welling. Group equivariant convolutional networks. In International conference machine learning, pages 2990-2999, PMLR, 2016.
15

---

[90] Victor Garcia Satorras, Emiel Hoogeboom, and Max Welling. E (n) equivariant graph neural networks. In International conference on machine learning, pages 9323-9332. PMLR, 2021.
[91] John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ron- neberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Židek, Anna Potapenko, al. Highly accurate protein structure prediction with alphafold. Nature, 596(7873)
2021.
[92] Nathaniel Thomas, Tess Smidt, Steven Kearnes, Lusann Yang, Li Li, Kai Kohlhoff, and Patrick Riley. Tensor field networks: Rotation-and cunslationeequivariinn neural networks for 3d point clouds. arXiv preprint 18X.00180...829,, 2018.
[93] James Kirkpatrick, Brendan McMorrow, David HP Turban, Alexander L Gaunt, James Spencer, Alexander GDG Matthews, Annette Obika, Louis Thiry, Meire Fortunato, David Pfau, et al. Pushing the frontiers of density functionals by solving the fractional electron problem. Science, 374(6573)11885-339) 2021.
[94] Jie Zhou, Ganqu Cui, Shengding Hu. Zhengyan Zhang, Cheng Yang, Zhiyuan Liu, Lifeng Wang, Changcheng Li, and Maosong Sun. Graph neural networks: A review of methods and applications. Al open, 1:57-81, 2020.
[95] Jonathan Masci, Davide Boscaini, Michael Bronstein, and Pierre Vandergheynst. Geodesic con- volutional neural networks on riemannian manifolds. In Proceedings of the IEEE international conference computer vision workshops, pages 37-45, 2015.
[96] Jun Li, Li Fuxin, and Sinisa Todorovic. Efficient riemannian optimization the stiefel manifold via the cayley transform. arXiv preprint arXiv:2002.01 2020.
[97] Bobak Kiani, Randall Balestriero, Yann Lecun, and Seth Lloyd. projumm: efficient method for training deep networks with unitary matrices. arXiv preprint arXiv:2203.05483, 2022.
[98] Zhengdao Chen, Jianyu Zhang, Martin Arjovsky, and Léon Bottou. Symplectic recurrent neural networks. arXiv preprini arXiv: 1909. 13334, 2019.
[99] Martin Arjovsky, Amar Shah, and Yoshua Bengio. Unitary evolution recurrent neural networks. International conference on machine learning, pages 1120-1128. PMLR, 2016.
[100] Turgut Ōziş and iSMAIL Aslan. Similarity solutions to burgers' equation in terms of special functions of mathematical physics. Acta Physica Polonica B, 2017.
[101] SP Lloyd. The infinitesimal group of the navier-stokes equations. Acta Mechanica, 38(1-2): 85-98, 1981.
[102] Yang You, Igor Gitman, and Boris Ginsburg. Large batch training of convolutional networks. arXiv preprint arXiv: 7788. 03888, 2017.
16

---

Appendix
Table of Contents
PDE Symmetry Groups and Deriving Generators
18
A.1 Symmetry Groups and Infinitesimal Invariance
19
A.2 Deriving Generators of the Symmetry Group of PDE
20
A.3 Example: Burgers' Equation
21
Exponential map and its approximations
22
B.I
Approximations the exponential map
23
VICReg Loss
24
D Expanded related work
25
Details on Augmentations
26
E.
Burgers' equation
26
2 Kdv
27
KS
28
Navier Stokes
28
Experimental details
28
Fl
Experiments on Burgers' Equation
28
F2
Experiments on KdV and KS
30
F3
Experiments Navier-Stokes
31
17

---

A PDE Symmetry Groups and Deriving Generators
Symmetry augmentations encourage invariance of the representations to known symmetry groups of the data. The guiding principle is that inputs that can be obtained from one another via transformations of the symmetry group should share a common representation. In images, such symmetries are known priori and correspond to flips, resizing, rotations of the input. In PDEs, these symmetry groups can be derived as Lie groups. commonly denoted as Lie point symmetries, and have been categorized for many common PDEs [T]]. An example of the form of such augmentations is given in Figure 6 for a simple PDE that rotates a point in 2-D space. In this example, the PDE exhibits both rotational symmetry and scaling symmetry of the radius of rotation. For arbitrary PDEs, such symmetries can be derived, as explained in more detail below.
Example: symmetries and invariances
dy
dx
āt ay
āt
rotational symmetry
rotation speed
scaling
t 2
(inveriant quantity)
synumelyy
t =
t=
Figure 6: Illustration of the PDE symmetry group and invariances of a simple PDE, which rotates a point in 2-D space. The PDE symmetry group here corresponds scalings of the radius of the rotation and fixed rotations of all the points over time. A sample invariant quantity is the rate of rotation (related to the parameter a in the PDE), which is fixed for any solution this PDE.
The Lie point symmetry groups of differential equations form a Lie group structure, where elements of the groups are smooth and differentiable transformations. It is typically easier to derive the symmetries of system of differential equations via the infinitesimal generators of the symmetries, (i.e., at the level of the derivatives of the one parameter transforms). By using these infinitesimal generators, one can replace nonlinear conditions for the invariance of a function under the group transformation, with an equivalent linear condition of infinitesimal invariance under the respective generator of the group action [1].
In what follows, we give an informal overview to the derivation of Lie point symmetries. Full details and formal rigor can be obtained Olver [11], Ibragimov 13 among others.
In the setting we consider, a differential equation has a set of p independent variables æ (xr1, _ ,..., xP) E RP and q dependent variables take the form u = f(z), where u² = fx) for a E [l,.... 4). Solutions form a graph over a domain Q IRP:
(u', u² ...., u²) R9. The solutions
r ((x,, f(x) E C R R2
(10)
In other words, given solution y forms pdmmmensoonal submanifold of the space Rp x R9.
The n-th prolongation of a given smooth function I'y expands or "prolongs" the graph of the solution into larger space include derivatives up to the n-th order. More precisely, if 2 H is the solution space of given function and f R - U. then we introduce the Cartesian product space of the prolongation:
(1) = 11 x U1 x U2 X x U...
(11)
where U₁ = dkim(k) and dim(k) = ( is the dimension of the so-called jet space consisting of all k-th order derivatives. Given any solution f Rp + 24, the prolongation simply calculating the corresponding derivatives up to order n (e.g., via a Taylor expansion at each
be calculated by
point). For a given function u = f(xz), the n-th prolongation denoted as (x) = pr( f(x)). As a simple example, for the case of p = 2 with independent variables and y and q = 1 with a single
18

---

dependent variable f, the second prolongation is u pp!!2) f(x = (w; Uxa May Mxar Mary Myyy)
[: 0f 0f 00 0
0 Or dy O= drdy' On²
(12)
ER' x x
which is evaluated at a given point (z: y) in the domain. The complete space R2 x u(n) often called the n-th order jet space [II].
system of differential equations a set of differential equations A Rp x u - Rd of the independent and dependent variables with dependence on the derivatives up to maximum order of n:
(x u(l)) 0, I....,
(13)
A smooth solution is thus function f such that for all points in the domain of æ:
(x, pr(*) [(x)) 0, ,,...,l.
(14)
In geometric terms, the system of differential equations states where the given map A vanishes on the jet space, and forms subvariety
S = ((x u (x)))= : (xx u (m)) = 0} C R x (()))
(15)
Therefore to check if solution is valid, one can check if the prolongation of the solution falls within the subvariety Za. As an example, consider the one dimensional heat equation
A U4
Clizz
(16)
We can check that f(z, t) sin(x)e -t is solution by forming its prolongation and checking if it falls withing the subvariety given by the above equation:
(r f(x t) = (sin(x)e= cos(sre at -csin(z)e si-(eee -ccos(x)e -e,,t, sin(()) sin(
A (x,t,u!") -esin(z)e esin(x)e 0.
(17)
A.1 Symmetry Groups and Infinitesimal Invariance
symmetry group G for system of differential equations is a set of local transformations to the function which transform one solution of the system of differential equations to another. The group takes the form of Lie group, where group operations can be expressed as a composition of one-parameter transforms. More rigorously, given the graph of solution I defined in Eq. (10). group operation E maps this graph to a new graph
, (( i g (æx, u.) (æ, u) E Ery,,
(18)
where (æ,, ú) label the new coordinates of the solution in the set ry, For example, if x = (x. u(2, t), and g acts on (z. u) via
(I, + (a tt,
then ū(z,1) u(x, = w(i - ol, 1) + where (7. = (x + el,
Note, that the set g- y may not necessarily be graph of new v-luueed function; however, since all transformations are local and smooth, one can ensure transformations are valid in some region near the identity of the group.
As an example, consider the following transformations which are members of the symmetry group of the differential equation uur 0. 91 (t) translates single spatial coordinate by amount t and Ø2 scales the output coordinate u by an amount e":
(1 (I,4) = (x + u). go (r) (z, u) = (x, u))
(19)
easy to verify that both of these operations are local and smooth around a region of the identity, as sending r, t 0 recovers the identity operation. Lie theory allows one to equivalently describe
19

---

the potentially nonlinear group operations above with corresponding infinitesimal generators of the group action, corresponding to the Lie algebra of the group. Infinitesimal generators form a vector field over the total space f x M, and the group operations correspond to integral flows over that vector field. To map from a single parameter Lie group operation to its corresponding infinitesimal generator, we take the derivative of the single parameter operation at the identity:
d Ug Vglex,,)
di gt) (x, u)
(20)
where g(0)- (z, u) = (I, u).
To map from the infinitesimal generator back to the corresponding group operation, one can apply the exponential map
exp(tu) (at, u) = g(t) (x, u),
(21)
where exp: g + G. Here, (-) maps from the Lie algebra, g. to the corresponding Lie group, G. This exponential map can evaluated using various methods, as detailed Appendix and Appendix
Returning the example earlier from Equation (19). the corresponding Lie algebra elements are
0g 02 < 9i (t) (x., u) (x: 1, Ug: = ud_ + 92( (z,, u) = (x, "...c"..).
(22)
Informally, Lie algebras help simplify notions of invariance as allows one to check whether functions differential equations are invariant to group by needing only to check it at the level of the derivative of that group. In other words, for any vector field corresponding to a Lie algebra element, given function is invariant to that vector field if the action of the vector field the given function evaluates to zero everywhere. Thus, given a symmetry group, one can determine a set of invariants using the vector fields corresponding the infinitesimal generators of the group. To determine whether a differential equation is in such a set of invariants, we extend the definition of prolongation to act on vector fields
d
(*(") (m)
(r(") [exp(ev)] (or, ()))
(23)
de
le=0 ===
A given vector field is therefore an infinitesimal generator of symmetry group G of system of differential equations A, indexed by u E (1,....I} if the prolonged vector field of any given solution is still solution:
(*)
(n) v (x, (")))] 0, I,....,, whenever A(x, u (")) 0.
(24)
For sake of convenience and brevity, we leave out many of the formal definitions behind these concepts and refer the reader to [1I] for complete details.
A.2 Deriving Generators of the Symmetry Group of PDE
Since symmetries of differential equations correspond to smooth maps, is typically easier to derive the particular symmetries of a differential equation via their infinitesimal generators. To derive such generators, we first show how to perform the prolongation of a vector field. As before, assume we have p independent variables x² ..... x" and dependent variables u ...., u', which are a function of the dependent variables. Note that we use superscripts to denote a particular variable. Derivatives with respect to given variable are denoted via subscripts corresponding the indices. For example, the variable u/12 denotes the third order derivative of ul taken twice with respect to the variable x. and once with respect to x?. As stated earlier, the prolongation of vector field is defined as the operation
d (r(") [exp(ev)] (z, (m)).
(r( v (m))()))
(25)
de c==
To calculate the above, we can evaluate the formula on a vector field written a generalized form. Le., any vector field corresponding the infinitesimal generator of symmetry takes the general form
(x, 4)
Sa (2, 4
(26)
Or
Ou
1=1
20

---

Throughout, we will use Greek letter indices for dependent variables and standard letter indices for independent variables. Then, have that
9 pr D-
dy (z, v(()) ΣΣ J
a
(27)
Oug
where is a tuple of dependent variables indicating which variables are in the derivative of Duyy"
Each oi (ac, u(() is calculated as
C m))
II D;
0a
Σ Stuf
+
(28)
SEJ
where "²,; =00g//00 /dx and Dj is the total derivative operator with respect to variable defined
D,P(z,ul")) = dr +
OP
ΣΣ
JR Juy
OP
(29)
After evaluating the coefficients, 0% (x, u (,), we can substitute these values into the definition of the vector field's prolongation in Equation (27). This fully describes the infinitesimal generator of the given PDE, which can be used to evaluate the necessary symmetries of the system of differential equations. An example for Burgers' equation, a canonical PDE, presented the following.
A.3 Example: Burgers' Equation
Burgers' equation is PDE used to describe cinvection-diffusion phenomena commonly observed in fluid mechanics, traffic flow, and acoustics 43). form or its "viscous" form. The potential form is
The PDE can be written in either its "potential"
Mt = Myz 2
(30)
Cautionary note: We derive here the symmetries of Burgers' equation in its potential form since this form is more convenient and simpler to study for the sake of an example. The equation we consider in our experiments is the more commonly studied Burgers' equation in its standard form which does not have the same Lie symmetry group (see Table 4T. Similar derivations for Burgers' equation in its standard form can be found in example 6.1 of
Following the notation from the previous section, = and q = 1. Consequently, the symmetry group of Burgers' equation will be generated by vector fields of the following form
Ô + T(x 1, at
i + olx L, u) Ou
à
= E(x, L u) Dz
(31)
where we wish to determine all possible coefficient functions, E(1, x, u), r(x, I, u), and 6(x, w) such that the resulting one-parameter sub-group exp (ev) is symmetry group of Burgers' equation.
To evaluate these coefficients, we need prolong the vector field up 2nd order, given that the highest-degree derivative present in the governing PDE is of order 2. The 2nd prolongation of the vector field can be expressed as
pr (2) x = v
+
+
dett
+
tt
(32)
Du.
Ou,
Juga
Dulzz
Applying this prolonged vector field to the differential equation in Equation imal symmetry criteria that
30 get the infinites-
(2) v ,,
L, [2)) =0² -0--0"00 -0xx 2u O.
(33)
To evaluate the individual coefficients, we apply Equation (8 Next, we substitute every instance of with u Uzzs and equate the coefficients of each monomial in the first and second-order
21

---

Table 3: Monomial coefficients in vector field prolongation for Burgers' equation.
Monomial
Coefficient
1
Φ Øez
Ug
20. 2(Oru Sex) = -5t
ug
2 Ex) Taz (ouu 22zu) Qu
ug
2T. 2 2Tzu Suu - 64
-2T, Tuu
Uxz
Taz + (du 2 Qu 78
Ug Uga
2T. TTzu 3ru 5u
2ra Tuu
2Tu
Tu -Tu
Mat
2T,
Ug
2TT, = 0
derivatives of to find the pertinent symmetry groups. Table below lists the relevant monomials as well as their respective coefficients.
Using these relations, we can solve for the coefficient functions. For the case of Burgers' equation, the most general infinitesimal symmetries have coefficient functions of the following form:
<(1,0)) kn + har + 2ki + 4kgxt
(34)
(1) = 2/2 + 2kal + 4hip(?
(35)
ott, x, u) (k3 -k5 2ket kg (x,
(36)
where ......y Ks IR and (xx is arbitrary solution to Burgers' equation. These coefficient functions can be used to generate the infinitesimal symmetries. These symmetries are spanned by the six vector fields below:
0
(37)
U2 0
(38)
13 0,
(39)
81 = x0g + 2t0,
(40)
2td, zū,
(41)
V6 4xt0, + 48 0: (x² + 21)a.
(42)
as well as the infiniteddimensional subalgebra: solution the heat equation. The relationship between the Heat equation and Burgers' equation can be seen, whereby if u is replaced by w = c", the Cole holf transformation is recovered.
rax "0, Here, >(x, f) is
any arbitrary
B Exponential map and its approximations
As observed in the previous section, symmetry groups are generally derived in the Lie algebra of the group. The exponential map can then be applied, taking elements of this Lie algebra to the corresponding group operations. Working within the Lie algebra a group provides several benefits. First, a Lie algebra to yield new elements of the Lie algebra (and the group, via the exponential map). Second, when generators of the Lie algebra are closed under the Lie bracket of the Lie algebra (i.e., the generators form a basis for the structure constants of the Lie algebra). any arbitrary Lie point symmetry can be obtained via element of the Lie algebra (i.e. the exponential map is surjective onto the connected component of the identity) [I]] In contrast, composing group operations in an arbitrary, fixed sequence not guaranteed to be able to generate any element of the group. Lastly, although not extensively detailed here, the "strength," or magnitude, of Lie algebra elements can be measured using an appropriately selected norm. For instance, the operator norm of matrix could be used for matrix Lie algebras.
vector space, elements of the Lie algebra can be added and subtracted
22

---

In certain cases, especially when the element v in the Lie algebra consists of single basis element, the exponential map exp(v) applied that element of the Lie algebra can be calculated explicitly. Here, applying the group operation to tuple of independent and dependent variables results in the so- called Lie point transformation, since is applied at a given point exp(en) ((x)) + , f())).
Consider the concrete example below from Burger's equation.
Example B.1 (Exponential map on symmetry generator of Burger's equation). contains the Lie point symmetry Vyy (x,,t)e "0,, with corresponding group transformation exp(ev,) (z, t, u) (z,t, log (e" + +7)).
The Burger's equation
Proof. This transformation only changes the component. Here, we have
exp (eye-"8 0..) u u+
Σ
1=
(43)
- =u + eye
2722
2272--11
2n
+
1 -
3u +...
2
3
Applying the series expansion log(1 + x) = x
x²
2
+
x¹
3
we get
exp (0,e"0,,) = u + log (1 + e
log (e") + log (1 + eye
(44)
= log (e" -
In general, the output of the exponential map cannot be easily calculated as we did above, especially if the vector field is weighted sum of various generators. In these cases, still apply the exponential map to a desired accuracy using efficient approximation methods, which we discuss next.
B.1 Approximations to the exponential map
For arbitrary Lie groups, computing the exact exponential map is often not feasible due to the complex nature of the group and its associated Lie algebra. Hence, it necessary to approximate the exponential map to obtain useful results. Two common methods for approximating the exponential map are the truncation of Taylor series and Lie- Lie-Trotter approximations.
Taylor series approximation Given a vector field v in the Lie algebra of the group, the exponential map can be approximated by truncating the Taylor series expansion of exp(v). The Taylor series expansion of the exponential map is given by:
1 2
exp(v) Id tu
(45)
n
To approximate the exponential map, we retain finite number of terms in the series:
k tn n! o((l)))))
exp(v)
(46)
where & is the order of the truncation. The accuracy of the approximation depends on the number of terms retained in the truncated series and the operator norm ||o||. For matrix Lie groups, where is also matrix, this operator norm is equivalent to the largest magnitude of the eigenvalues of the matrix [45]. The error associated with truncating the Taylor series after & terms thus decays exponentially with the order of the approximation.
Two drawbacks exist when using the Taylor approximation. First, for given vector field v. applying v. if to given function ff requires algebraic computation of derivatives. Alternatively, derivatives can also be approximated through finite difference schemes, but this would add an additional source of error. Second, when using the Taylor series to apply a symmetry transformation of a PDE to starting solution of that PDE, the Taylor series truncation will result in a new function, which not necessarily a solution of the PDE anymore (although it can be made arbitrarily close to a solution by increasing the truncation order). Tro-trr approximations, which we study next, approximate the exponential map by composition of symmetry operations, thus avoiding these two drawbacks.
23

---

Lie-Trotter series approximations The Trottro approximation is an alternative method for approximating the exponential map, particularly useful when one has access group elements directly, i.e. the closed-form output of the exponential map on each Lie algebra generator), but they are non-commutative. To provide motivation for this method, consider two elements X and Y the Lie algebra. The Lie-Trotter formula (or Lie product formula) approximates the exponential of their sum [22] 46).
X
X k
exp(X X +
exp
exp
exp
exp
(47)
7--00
where & is positive integer controlling the level of approximation.
The first-order approximation above can be extended higher orders, referred as the Lie-Trotter- Suzuki approximations, Though various different such approximations exist, we particularly use the following recursive approximation scheme
47
23
for given Lie algebra component Σx= Vs.
T2(v) exp
11 2
exp
exp
Up
exp
Up 2
exp
Pp-1
exp
11
2
2
2
(2)
T2x 2
Ug )²
T2x 200(2((11) 4uk)
T2x- 2u
(48)
4 41/(2k-1)
To apply the above formula, we tune the order parameter and split the time evolution into r segments to apply the approximation exp(v) = Tp (p(w/r). For the p-th order, the number of stages in the Suzuki formula above equal to 2 50/2-1 so the total number of stages applied is equal to 2r 2 5p/2-1
These methods are especially useful in the context of PDEs, as they allow for the approximation of the exponential map while preserving the structure of the Lie algebra and group. Similar techniques are used in the design of splitting methods for numerically solving PDEs 48 49]. Crucially, these approximations will always provide valid solutions the PDEs, since each individual group operation in the composition above is itself a symmetry of the PDE. This is in contrast with approximations via Taylor series truncation, which only provide approximate solutions.
As with the Taylor series approximation. the p-th order approximation above is accurate to (|)) with suitably selected values of and [23]. As cautionary note, the approximations here may fail to converge when applied to unbounded operators (50] [1). In practice, we tested a range of bounds to the augmentations and tuned augmentations accordingly (see Appendix E)
C VICReg Loss
In our implementations, we use the VICReg loss as our choice of SSL loss (9] This loss contains three different terms: variance term that ensures representations do not collapse to a single point, covariance term that ensures different dimensions of the representation encode different data, and an invariance term to enforce similarity of the representations for pairs of inputs related by an augmentation. We go through each term in more detail below. Given a distribution T from which to draw augmentations and set of inputs I; the precise algorithm to calculate the VICReg loss for batch of data is also given in Algorithm I
Formally, define our embedding matrices as Z, Z' E NND Next, we define the similarity criterion, f.im- as
Lim Sum v) u --li U
which we use to match our embeddings, and to make them invariant to the transformations. To avoid collapse of the representations, we use the original variance and covariance criteria to define our regularisation loss, Lreg, as
(Z) Aoov C(Z) + Avar V(Z), with
C(Z) D
Σ Cov(Z)) Z) and
itj
D
V(Z) Σ
max 0, /Var(Z.s))
j=1
24

---

Algorithm 1 VICReg Loss Evaluation
Hyperparameters: Aear, Acou, Xinn, 7 E R Input: N inputs in a batch {x; E RD
N
VICRegLoss(NV, Iis Avors Acoes Asnus 7):
1: Apply augmentations i, T to form embedding matrices Z, Z' E ERRND:
2,: = he (Je (x))) and Z = he (Jo (1' x.))
2: Form covariance matrices Cov( Z), Cov (Z)) E RDDD:
N Zi Σ
1
Cov ( = N
Σ
ZZ - Zi.) (Z, Zii)
1=1
i=1
3: Evaluate loss: L'Z,Z') = Aver (z,Z') + Noon Com Z Z') + Cine Zz 2')
N
Luar (0ar(Z,,2 D
max(0, 7 Cov (z);: max(00,7 Cov( Cov(Z")ii)
i=1
Loo0 (2, 2')
Σ
(Zov(Z)iy] + [Cov( (Z))
(Z,
1c 1.
1=1
4: Return: 2')
The variance criterion, V(Z), ensures that all dimensions in the representations are used, while also serving as a normalization of the dimensions. The goal of the covariance criterion to decorrelate the different dimensions, and thus, spread out information across the embeddings.
The final criterion is
LViCRez (z, Z')
Ainv
Lxim Zijiv Zliin) + (zeg(Z') + Greg ()).
N
Hyperparameters Avary Acous Aino, 7 E E weight the contributions of different terms in the loss. For all studies conducted this work, we use the default values of Avar = Ainv = 25 and Aooe = 1, unless specified. In our experience, these default settings perform generally well.
D Expanded related work
Recent work on machine learning for PDEs has considered both invariant prediction tasks 522 and time-series modelling (53 544. In the fluid mechanics setting, models learn dynamic viscosities. fluid densities, and/or pressure fields from both simulation and
Machine Learning for PDEs
real-world experimental data 55 56 57] For time-dependent PDEs, prior work has investigated the efficacy of convolutional neural networks (CNNs), recurrent neural networks (RNNs), graph neural networks (GNNs), and transformers in learning to evolve the PDE forward time (34 58 59 60).
This has invoked interest in the development of reduced order models and learned representations for time integration that decrease computational expense, while attempting maintain solution accuracy, Learning representations of the governing PDE can enable time-stepping in a latent space, where the computational expense is substantially reduced [61] Recently, for example, Lusch et al. have studied learning the infinite-die seiiitediimesional Koopman operator to globally linearize latent space dynamics [2] Kim et al. have employed the Sparse Identification parameterize latent space trajectories and combine them with classical ODE solvers to integrate latent space coordinates arbitrary points time (33] Nguyen et al. have looked the development of foundation models for climate sciences using transformers pre- pre-trained on well-established climate
Nonlinear Dynamics (SINDy) framework to
25

---

datasets [7]. Other methods like dynamic mode decomposition (DMD) are entirely data-driven and find the best operator to estimate temporal dynamics [63]. Recent extensions of this work have also considered learning equivalent operators, where physical constraints like energy conservation or the periodicity of the boundary conditions are enforced [29]
Self-supervised learning All joint embedding self-supervised learning methods have a similar
objective: forming representations across given domain of inputs that are invariant to a certain set of transformations. Contrastive and non-contrastive methods are both used, Contrastive methods [21 (4]
push away unrelated pairs of augmented datapoints, and frequently rely on the InfoNCE criterion 688 although in some cases, squared similarities between the embeddings have been employed 669 Clustering-basee
65
66
67
methods have also recently emerged [70] 7 6). of contrasting pairs of samples, samples are contrasted with cluster centroids. Non-contrastive methods [TC 40
where instead
72 73 39 aim bring together embeddings of positive samples. However, the primary difference between contrastive and non- non-contrastive methods lies in how they prevent representational collapse. In the former, contrasting pairs of examples are explicitly pushed away to avoid collapse. In the latter, the criterion considers the set of embeddings as a whole, encouraging information content maximization to avoid collapse. For example, this can achieved by regularizing the empirical covariance matrix of the embeddings. While there can be differences in practice, both families have been shown to lead to very similar representations [6 75). An intriguing feature in many SSL frameworks is the use of projector neural network after the encoder, on top of which the SSL loss is applied. The projector was introduced in (21]. Whereas the projector is not necessary for these methods to learn satisfactory representation, it is responsible for an important performance increase. Its exact role is an object of study [76 [5].
We should note that there exists a myriad of techniques. including metric learning, kernel design, autoencoders, and others 777 78 80 8 to build feature spaces and perform unsupervised learning. Many of these works share similar goal to ours, and we opted for SSL due to its proven efficacy in fields like computer vision and the direct analogy offered by data augmentations. One particular methodology that deserves mention is that of multi-fidelity modeling, which can reduce dependency on extensive training data for learning physical tasks [82 83 84). The goals of multi- fidelity modeling include training with data of different fidelity (2) or enhancing the accuracy of models by incorporating high quality data into models (85). In contrast, SSL aims to harness salient features from diverse data sources without being tailored to specific applications. The techniques we employ capitalize on the inherent structure in a dataset, especially through augmentations and invariances.
Equivariant networks and geometric deep learning In the past several years, an extensive set of literature has explored questions in the so-called realm of geometric deep learning tying together aspects of group theory, geometry, and deep learning [6). In one line of work, networks have been designed to explicitly encode symmetries into the network via equivariant layers or explicitly
symmetric parameterizations [87
88 89 90). These techniques have notably found particular
application in chemistry and biology related problems [9] 9 9 as well learning on graphs
(94). Another line of work considers optimization over layers or networks that are parameterized over Lie group 195 96 97 98| 199. Our work does not explicitly encode invariances or structurally parameterize Lie groups into architectures as in many of these works, but instead tries to learn representations that are approximately symmetric and invariant to these group structures via the SSL. As mentioned in the main text, perhaps more relevant for future work are techniques for learning equivariant features and maps [4 42
E Details on Augmentations
The generators of the Lie point symmetries of the various equations we study are listed below. For symmetry augmentations which distort the periodic grid in space and time, we provide inputs v and t to the network which contain the new spatial and time coordinates after augmentation.
E.1 Burgers' equation
As a reminder, the Burgers' equation takes the form
Ue
WMy Vuuz 0.
(49)
26

---

Lie point symmetries of the Burgers' equation are listed in Table 4 There are five generators. As we will see, the first three generators corresponding translations and Galilean boosts are consistent with the other equations we study (KS, KdV, and Navier Stokes) as these are all flow equations.
Table 4: Generators of the Lie point symmetry group of the Burgers' equation in its standard form 100]
Group operation
Lie algebra generator
(x,t,u) H
91 (space translation)
ed,
x 10+t,,,0
02 (time translation)
cd:
(z, f 1 c u)
93 (Galilean boost)
c(t0z + 0.)
x + et i, u e
94 (scaling)
c(x0, + 2t0; - ud,)
(c'r e 2t e e=
gs (projective)
e(x10, + 1=0, + (xr tu) Qw)
u +
c(x tu)
el
Comments regarding error in 12 As a cautionary note, the symmetry group given in Table of 12 for Burgers' equation is incorrectly labeled for Burgers' equation in its standard form. Instead, these augmentations are those for Burgers' equation its potential form, which is given as:
U2 U Vugz 0.
(50)
Burgers' equation in its standard form is + vo, VUge O, which can be obtained from the transformation v Ug. The Lie point symmetry group of the equation in its potential form contains more generators than that of the standard form. To apply these generators to the standard form of Burgers' equation, one can convert them via the Cole-Hopf transformation. but this conversion loses the smoothness and locality of some of these transformations (i.e., some are no longer Lie point transformations, although they do still describe valid transformations between solutions of the equation's 's corresponding form).
Note that this discrepancy does carry through in their experiments: 112 data as solutions to Heat equation, which they subsequently transform into solutions of Burgers' equation via a Cole-Hopf transform. Therefore, in their code, they apply augmentations using the Heat equation, for which they have the correct symmetry group. We opted only to work with solutions Burgers' equations itself for slightly fairer comparison to real-world settings, where convenient transform to linear PDE such as the Cole-Hopf transform is generally not available.
only consider input
E.2
Kdv
Lie point symmetries of the KdV equation are listed in Table Though all the operations listed are valid generators the symmetry group. only Φ1 and gs are invariant the downstream task of the inverse problem. (Notably, these parameters are independent of any spatial shift). Consequently, during SSL pre-training for the inverse problem, only 91 and 93 were used for learning representations. In contrast, for time-stepping, all listed symmetry groups were used.
Table 5: Generators of the Lie point symmetry group of the KdV equation. The only symmetries used in the inverse task of predicting initial conditions 91 and gs since the other two not invariant to the downstream task.
Group operation (xr, t, u) H
Lie algebra generator
91 (space translation)
e0,
x + e u)
92 (time translation)
c0,
(z, t + u)
93 (Galilean boost)
<(tO, + 0n)
x + et t, u c
94 (scaling)
((r,, 3tt; -2ud.)
E e 3t 2
27

---

E.3 KS
Lie point symmetries of the KS equation are listed in Table 6 All of these symmetry generators are shared with the KdV equation listed in Table 4 Similar to KdV, only 91 and Os are invariant to the downstream regression task of predicting the initial conditions. In addition, for time-stepping, all symmetry groups were used in learning meaningful representations.
Table 6: Generators of the Lie point symmetry group of the KS equation. The only symmetries used the inverse task of predicting initial conditions are 9i and ga since 92 is not invariant to the downstream task.
Group operation
Lie algebra generator
(x, u) H
91 (space translation)
e0,
[ e u)
92 (time translation)
e0,
(x. l 146
93 (Galilean boost)
(10= + O.)
x + d it u a+t
E.4 Navier Stokes
Lie point symmetries of the incompressible Navier Stokes equation are listed Table 7 J01. As pressure is not given an input to any of our networks, the symmetry In was not included in our implementations. For augmentations gE, and gE,, we restricted attention only to linear (t) = (t) = t or quadratic (x) = (t) = t² functions. This restriction was made to maintain invariance to the downstream task of buoyancy force prediction in the linear case easily calculable perturbations to the buoyancy by an amount 2c to the magnitude in the quadratic case. Finally, fix both order and steps parameters in our Lie-Trotter approximation implementation to for computationnal efficiency.
Experimental details
Whereas implemented own pretraining and evaluation (kinematic viscosity, initial conditions and buoyancy) pipelines, we used the data generation and time-stepping code provided on Github by [12) for Burgers', KS and KdV, and in [18] for Navier-Stokes (MIT License), with slight modifica- tion to condition the neural operators on our representation. All our code relies relies on Pytorch. Note that the time-stepping code for Navier-Stokes uses Pytorch Lightning. We report the details of the training cost and hyperparameters for pretraining and timestepping in Table and Table TO respectively.
F.1 Experiments on Burgers' Equation
Solutions realizations of Burgers' equation were generated using the analytical solution 33 obtained from the Heat equation and the Cole-Hopf transform. During generation, kinematic viscosities, and initial conditions were varied.
Representation pretraining We pretrain a representation on subsets of our full dataset containing 10, 000 1D time evolutions from Burgers equation with various kinematic viscosities, V, sampled uniformly in the range [0.001, 0.007]. and initial conditions using similar procedure to [12]. We generate solutions of size 224 x 448 in the spatial and temporal dimensions respectively, using the default parameters from (1). We train a ResNetl 88 [17 encoder using the VICReg [ approach to
joint embedding SSL, with a smaller projector (width 512) since we use a smaller ResNet than the
original paper. We keep the same variance, invariance and covariance parameters as in [9] We use
the following augmentations and strengths:
Crop of size (128, 256), respectively. in the spatial and temporal dimension.
Uniform sampling in [-2, 2] for the coefficient associated to g1-
Uniform sampling in [0, 2] for the coefficient associated to g2-
Uniform sampling in [-0.2, 0.2) for the coefficient associated gs-
28

---

Table Generators of the Lie point symmetry group of the incompressible Navier Stokes equation. Here, correspond to the velocity of the fluid in the r, y direction respectively and corresponds to the pressure. The last three augmentations correspond to infinite dimensional Lie subgroups with choice of functions (t), Fig a(4) that depend on l only. For invariant tasks, we only used settings where (t) (t), (t) = t (linear) or Ex (t), (t) = t² (quadratic) to ensure invariance to the downstream task or predictable changes in the outputs of the downstream task. These augmentations are listed as numbers 6 to 9.
Group operation
Lie algebra generator
(z.y, ł, u, U. H (z.y, t t C+c, u,
91 (time translation)
cb
92 (z translation)
+02
x [ y. 1.
9s (y translation)
c0
(I: y + c t, v.
< (2t0, + z0g + y0
94 (scaling)
(Exx e'y. c²at
u
0 c 22 p
ud, - vd, 2p0p)
cos y ysin
ssine E + y cOs t,
9s (rotation)
(xx0, n0 + ud, vd
c usine
usin + v cos .P)
9o (r linear boost)?
et0,, + 0.)
[ v,t, u U.p)
97 (y linear boost)'
c(tO, + 0,)
(z. y+ et, u, u
98 (a quadratic boost)²
<(P0, + 2t0, - -2x0p)
x (x+ et? .y.t, u + 2et 2 p 2x
99 quadratic boost)²
e(00, + 2t0, - -2yd,)
(I, y + et ,,u, U 2et
p 2y
< (1),n, + (1)O,
x EE (t) V.t,
gE. (a general boost) 3
- (1)p))
u + c F (1) v, P F. F="x(()a))
(t), + E' (t)j, gE, (y general boost)3
(x. y + c (F,(()
- (E)))))),))))
u, + cE' y(t)
p E" v(t)y
9q (additive p (t))
(s. U. E, u, Us p q(t)
case of gE, DE, where (1) = Fry (1) = (linear function of t)
2 case of gE, or gE, where (t) = Fy, (t) = 1² (quadratic function of t)
(t), (t).). qtt) can be any given smooth function that only depends on t
* Uniform sampling in [- for the coefficient associated to 94.
We pretrain for 100 epochs using AdamW (33 and a batch size of 32. Crucially, we assess the quality of the learned representation via linear probing for kinematic viscosity regression, which we detail below.
Kinematic viscosity regression We evaluate the learned representation as follows: the ResNet18 is frozen and used as an encoder to produce features from the training dataset. The features are passed through a linear layer, followed by a sigmoid to constrain the output within [M/in Mmma]. The learned model is evaluated against our validation dataset, which is comprised of 2, 000 samples.
Time-stepping We use a 1D CNN solver from 12 as our baseline. This neural solver takes Tр previous time steps as input, to predict the next T'y future ones. Each channel (or spatial axis, if we view the input as 2D image with one channel) is composed of the realization values, u, at T. times, with spatial step size dr, and time step size dt. The dimension of the input is therefore (Tp + 2, 224). where the extra two dimensions are simply to capture the scalars dæ and dt. We augment this input with our representation. More precisely, we select the encoder that allows for the most accurate linear regression of v with our validation dataset, feed it with the CNN operator input and reduce the resulting representation dimension to d with a learned projection before adding it as supplementary channels to the input, which is now (Tp + 2 + d, 224).
We set Tp = 20, Tý 20, and Rsamples 2, 000. We train both models for 20 epochs fol-
29

---

Table 8: One-step validation NMSE for time-stepping on Burgers for different architectures.
Architecture
ResNetld
FNOId
Baseline (no conditioning)
0.110 + 0.008
0.184 + 0.002
Representation conditioning
0.108 + 0.011
0.173 + 0.002
lowing the setup from 12]. In addition, we use AdamW with a decaying learning rate and different configurations of 3 runs each:
Batch size E (16, 64).
Learning rate E {0.0001, 0.00005).
F.2 Experiments on KdV and KS
To obtain realizations of both the KdV and KS PDEs, we apply method lines, and compute spatial derivatives using pseudo-spectral method, in line with the approach taken by [2
Representation pretraining To train on realizations of KdV, we use the following VICReg param- eters: Avar 25, Aine 25, and Acoe = 4. For the KS PDE, the Avar and Xine remain unchanged, with Acee = 6. The pre-training is performed on a dataset comprised of 10, 000 ID time evolutions of each PDE, each generated from initial conditions described in the main text. Generated solutions were of size 128 x 256 in the spatial and temporal dimensions, respectively. Similar to Burgers' equation, R encoder in conjunction with a projector of width 512 was used for SSL pre-training. The following augmentations and strengths were applied:
Crop of size (32, 256), respectively, in the spatial and temporal dimension.
Uniform sampling in [-0.2, 0.2] for the coefficient associated ga-
Initial condition regression The quality of the learned representations is evaluated by freezing the R8 encoder, training separate regression head to predict values of As and wis and comparing the NMSE to supervised baseline. The regression head was fully-connected network, where output dimension commensurate with number of initial conditions used. In addition, tinge-constrained sigmoid was added to bound the output between [-0.5, 2Tr), where the bounds were informed by the minimum and maximum range of the sampled initial conditions. Lastly, similar to Burgers' equation, the validation dataset comprised of 2, 000 labeled samples.
The same ID CNN solver used for Burgers' equation serves as the baseline for time-stepping the KdV and KS PDEs. We select the ResNet18 encoder based the one that provides the most accurate predictions of the initial conditions with our validation set. Here, the input dimension now (T, + 2, 128) to agree with the size of the generated input data. Similarly to Burgers' equation, Tp = 20, T'y 20, and Reamples = 2, 000. Lastly, Adam W with the same learning rate and batch size configurations as those seen for Burgers' equation were used across time-stepping runs each.
Time-stepping
A sample visualization with predicted instances of the KdV PDE is provided in Fig. be- low:
Ground Truth
Predicted (SSL pre-training)
Predicted (CNN baseline)
Figure 7: Illustration of the 20 predicted time steps for the KdV PDE. (Left) Ground truth data from PDE solver; (Middle) Predicted u(x; t) using learned representations; (Right) Predicted output from using the CNN baseline.
30

---

Table 9: List of model hyperparameters and training details for the invariant tasks. Training time includes periodic evaluations during the pretraining.
Equation
Burgers'
KdV
KS
Navier Stokes
Network:
Model
ResNet18
ResNet18
ResNet18 512
ResNet18
Embedding Dim.
512
512
512
Optimization:
Optimizer
LARS [[02]
AdamW
Adam
AdamW
Learning Rate
0.6
0.3
0.3
3e-4
Batch Size
32
64
64
64
Epochs
100
100
100
Nb of exps
300
30
~ 30
300
Hardware:
GPU used
Nvidia V100
Nvidia M4000
Nvidia M4000
Nvidia V100
Training time
5h
11/
12h
48h
E3 Experiments Navier-Stokes
We use the Conditioning dataset for Navier Stokes-2D proposed in [8] consisting of 26,624 2D time evolutions with 56 time steps and various buoyancies ranging approximately uniformly from 0.2 to 0.5.
Representation pretraining We train a ResNet18 for 100 epochs with AdamW, batch size of 64 and learning rate of 3e-4. We use the same VICReg hyperparameters as for Burgers' Equation. We the following augmentations and strengths (augmentations whose strength is not specified here are not used):
Crop of size (16, 128, 128), respectively in temporal, and y dimensions.
Uniform sampling in [-1, 1] for the coefficients associated to ø2 and ga (applied respectively in x and y).
Uniform sampling in [-0.1, for the coefficients associated to gs-
Uniform sampling in [-0.01, 0.01] for the coefficients associated Is and gr (applied respectively x and y).
Uniform sampling in [-0.01, for the coefficients associated to gs and go (applied respectively in and y).
Buoyancy regression We evaluate the learned representation as follows: the ResNet18 is frozen and used as encoder to produce features from the training dataset. The features passed through a linear layer, followed by a sigmoid to constrain the output within [Buoyancy min- Buoyancy Both the fully supervised baseline (ResNet18 linear head) and our (frozen ResNet18 + linear head) model are trained on 3, 328 unseen samples and evaluated against 6, 592 unseen samples.
We mainly depart from 18 by using 20 epochs lear from 1,664 trajectories as we observe the results to be similar, and allowing to explore more combinations architectures and conditioning methods.
Time-stepping
Time-stepping results In addition to results on 1,664 trajectories, we also perform experiments with bigger train dataset (6,656) in (18) using 20 epochs instead of 50 for computational reasons. We also report results for the two different conditioning methods described in (18), Addition and AdaGN. The results can be found in Table As in (18), AdaGN outperforms Addition. Note that AdaGN is needed for our representation conditioning to significantly improve over no conditioning. Finally, found very small bottleneck in the MIP that process the representation also be crucial for performance, with a size of giving the best results.
31

---

Table 10: List of model hyperparameters and training details for the timestepping tasks.
Equation
Burgers'
Kdv
KS
Navier Stokes
Neural Operator:
Model
CNN
P2
CNN
12
CNN 12
Modified U-Net-64 [8
Optimization:
Optimizer
A
AdamW
Adam
Adam 2c-4
Learning Rate
le-4
le-4
le-4
Batch Size
16
16
16
32
Epochs
20
20
20
20
Hardware:
GPU used
Nvidia V100
Nvidia M4000
Nvidia M4000
Nvidia V100 (16)
Training time
ld
2d
2d
1.5d
Table 11: One-step validation MSE -1e-3 (L) for Navier-Stokes for different baselines and condi- tioning methods, with UNetmed64 [18 as base model.
-
Dataset
1,664
6,656
Methods without ground truth buoyancy:
Time conditioned, Addition
2.60 + 0.05
1.18 + 0.03
Time Rep. conditioned, Addition (ours)
2.47 1: 0.02
1.17 0.04
Time conditioned, AdaGN
2.37
1.12 1 0.02
Time Rep, conditioned, AdaGN (ours)
2.35 1 0.03
1.11 1 0.01
Methods with ground truth buoyancy:
Time Buoyancy conditioned. Addition
2.08 1: 0.02
1.10 0.01
Time Buoyancy conditioned, AdaGN
2.01 0.02
1.06 1 0.04
32


Reviews:
[{"review_id": "YrDJJ4aI0P", "reviewer": "Reviewer_VTPv", "summary": "This paper presents a general framework for self-supervised learning in a PDE context. In a way that is principled and natural, PDE symmetry groups are used to make the requisite augmentations from which self-supervision will learn structure; the augmentations are selected carefully so as to keep the regressed quantity (for the downstream task) constant. The utility of this method is demonstrated by regressing quantities of interest on four downstream PDEs, and the results indicate that the improvement of predicting from self-supervised representations is considerable relative to straightforward supervised learning. This work paves the way for adapting self-supervision to the physical sciences by way of carefully considered symmetry group-oriented data augmentations.", "questions": "My questions to the authors are listed below:\n\n1. Have you put more thought into preserving boundary conditions during augmentation and/or rectifying the fact that the symmetry transformations applied are no longer the proper symmetries, since the system no longer has infinite domain?\n\n2. Have you done any further investigation into why self-supervised learning gives results that considerably exceed that of supervised learning alone?\n\n### Verdict\n\nIn this paper, the authors consider the novel idea of applying self-supervised learning to PDEs. Although they lay the basic groundwork (e.g. performing augmentation via Lie symmetries), I feel that more consideration is needed to find a way to preserve boundary conditions during augmentation, such that the work is theoretically principled. The results are promising, in that the proposed approach considerably outperforms supervised learning, but little investigation is done as to why (this is a surprising/interesting observation and I would like to see more of an explanation given). As such, I recommend a borderline reject rating for this paper.\n\n### Additional Comments and Minor Corrections\n\nThe writing and presentation was, for the most part, very good (as mentioned above). I would recommend clarifying some of the table headings in results figures, for example, changing \"Best strength\" in Figure 5 (left) to \"Best strength ($\\epsilon$).\" A few minor corrections are given below:\n\nL221: \"test samples.As shown in\" -> \"test samples. As shown in\"\n\nL229: \"for to evaluate the models\" -> \"to evaluate the models\"\n\nL253: \"buyoancy\" -> \"buoyancy\"\n\nL254: \"which the hardest evaluation setting -> \"which is the hardest evaluation setting\"\n\n## Post-rebuttal update\n\nAfter reading the authors' rebuttals, overall, I think the paper will make for a positive contribution to the conference. I have increased my score to a 5.", "limitations": "The authors have adequately addressed limitations of their work in the \"Discussion\" section (Section 5).", "rating": 5, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 2, "strengths": "1. The high-level idea is relatively novel, in that I do believe this is the first work that aims to make use of self-supervised learning for partial differential equations, with augmentations performed naturally according Lie symmetries. The most closely related work [12] uses Lie point symmetries of PDEs in order to augment PDE datasets, but this happens purely in the context of supervised tasks. \n\n2. On the four PDE datasets presented/four equations considered, the results provide a fairly convincing improvement over just supervised learning alone. The mean squared error of prediction is considerably reduced (Table 1).\n\n3. In terms of writing and presentation, the paper is quite good. The field is well introduced, with relevant literature cited. The idea is clearly presented and figures are given to illustrate the methodology (Figures 1, 2, and 3 were all quite instructive).", "weaknesses": "1. Although the idea of using Lie symmetry-based augmentation for self-supervised learning is novel, I was underwhelmed at the rigor with which these augmentations were applied. Symmetries of PDEs are with respect to infinite domain or periodic boundaries. The boundary conditions imposed in the paper violate such symmetries, and hence make the justification for this sort of operation theoretically dubious. The authors admit that, naturally, because they violate these symmetries, they can only implement the group operations with \"small strengths.\" I believe that the work would benefit from further investigating how to preserve these boundary conditions during augmentation, thereby providing a theoretically sound basis for the self-supervised learning proposal.\n\n2. One of the interesting discoveries of this paper is that learning on top of self-supervised representations is considerably better than just immediately employing supervised learning techniques (Table 1); this is in stark contrast to what we see in e.g. computer vision, where both exhibit roughly the same performance. Some cursory post hoc analysis is given to explain this observation, although no further investigation is done in terms of providing a good understanding (be it theoretical or empirical).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for taking the time to review our paper and for the valuable comments and suggestions. We are glad that you found our high-level idea novel and our work well-presented. We acknowledge the concerns raised and would like to address them as follows.\n\n___\n### On preserving boundary conditions during augmentation and/or rectifying the fact that the symmetry transformations applied are no longer the proper symmetries\n\nYour point about the use of Lie symmetries and their association with infinite domain or periodic boundaries is well-taken. As discussed in our draft, we agree that more rigorous treatment of this matter could fortify our approach. This is, in some ways, a rather fundamental challenge in SSL: a similar situation arises in the image setting, where symmetries such as resize or translation also incur boundary issues (for instance, if an important piece of the image is omitted by this augmentation).\n\nImportantly, our use of crops as an SSL augmentation mitigates this issue by biasing the network to learn local features that are in some sense robust to errors on the boundary. In essence, cropping ensures that learned features are invariant to whether the data was collected at a boundary or not. Both in the PDE and the image setting, this use of crops is crucial to the success of SSL.\n\nIn more formally dealing with boundary conditions, we explored a number of ideas which are listed below. All of these ideas either significantly increased the complexity of the task or lost sight of the generality of the SSL PDE approach, so we opted to leave them to future work.\n- Approximate symmetries and discrete symmetries offer a way to explore many more symmetries, some of which may preserve boundary conditions. However, deriving these is no longer systematic: in the Lie symmetry case and in the approximate case, they provide solutions that are accurate only up to the order of the magnitude of the symmetry operation. These are also typically derived in settings with infinite domains. So apart from offering more flexibility in the symmetries, it is not obvious these directly help the boundary issue.\n- Another class of symmetries commonly studied are Lie Bäcklund symmetries [1], which offer a means to transform solutions from one PDE to another, but these are also typically derived in infinite domains.\n- Symmetries can be enforced at an infinitesimal level by using the Lie derivative defined as \n\\begin{equation}\\lim_{t\\to 0} \\frac{f\\bigl(\\Phi^t_X(p)\\bigr) - f\\bigl(p\\bigr)}{t},\\end{equation}\nwhere $\\Phi^t_X$ is the exponential map for the Lie algebra vector field $X$ applied for an amount $t$. This Lie derivative can be enforced to be close to zero in the domain of the PDE outside of boundaries. Though this is nice in practice, the Lie derivative does not apply an augmentation, and we do not know how to apply this in SSL settings.\n\n\n\nTo summarize, at least in this first work introducing the method, we wanted to offer a general approach to SSL for learning PDE data that was relatively simple to implement. The above approaches would have greatly complicated our approach, for questionable gain and loss of generality. As evidenced both in existing self-supervised work for images and our experiments, the representations learned even by these imperfect  augmentations still contain  rich information for downstream tasks. We welcome more thoughts and discussion from the reviewer on this point.\n\n\n### Why SSL works better\n\nIn general, understanding when and why SSL works is an active area of research. We know from prior theoretical and experimental results that, when done right, SSL pre-training finds feature spaces that are suited to diverse downstream targets and less prone to fitting trivial features [2-4]. From analyzing loss plots, we have some evidence that this is the case here.\n\nWe have added loss plots for kinematic viscosity regression (Burgers') in Figure 1 (see attached PDF) explaining the observed gaps in Table 1 of our paper between “Supervised” and “SSL features”. Both methods rely on a ResNet18 architecture. In the “Supervised” setting, this network is trained from scratch, whereas for the “SSL features” method, the network was pre-trained with our SSL framework, subsequently frozen, and then a linear model was trained on top of the output features. “SSL features” (i) are rich yet can easily be leveraged with a linear layer, allowing fast convergence towards small test errors and (ii) are less prone to overfitting than supervised learning while using the same architecture. The gaps observed in Table 1 might be reduced via longer “Supervised” training or other efforts, but this would go out of the scope of a fair evaluation. SSL features are simply competitive and easier to use.\n\n\n**References:**\n\n[1] Ibragimov, Nail H. CRC handbook of Lie group analysis of differential equations. Vol. 3. CRC press, 1995.\n\n[2] Chen, Ting, et al. \"A simple framework for contrastive learning of visual representations.\" International conference on machine learning. PMLR, 2020.\n\n[3] HaoChen, Jeff Z., et al. \"Provable guarantees for self-supervised deep learning with spectral contrastive loss.\" Advances in Neural Information Processing Systems 34 (2021): 5000-5011.\n\n[4] Cabannes, Vivien, et al. \"The ssl interplay: Augmentations, inductive bias, and generalization.\" International Conference on Machine Learning. PMLR, 2023.\n"}, {"review_id": "TEz7MZpv2n", "reviewer": "Reviewer_swdq", "summary": "In this paper, the authors propose to use self-supervised learning for obtaining an embedding that can robustly be used for predicting some quantities of interest or for time stepping. Particularly, they use the joint-embedding framework for SSL. They use symmetry groups for training the embedding to be invariant to these symmetry groups and thus capture the underlying behavior of the PDEs. ", "questions": "See the weaknesses. Further, there are more questions that are not particularly weakening the paper but could make it clearer.\n\n1. Figure 5 and the corresponding text is very interesting to me. It seems that the results from the Lie transformation(0.0038 compared to 0.0078 in supervised) and Crop (0.0052 compared to 0.0078) are not that different. It seems to me that applying some kind of SSL training is more important than using a particular Lie transformation per se. Is it worth the effort to find out the best possible selection of Lie point augmentations? Further, is this selection task dependent? How does one train a generic model for several tasks then?\n2. It would be nice to see some additional information on the computational overhead from SSL compared to supervised training. (see weaknesses, ideally, I would like to also know the comparison with DeepONets)\n", "limitations": "N/A", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "+ novelty in using Lie symmetry groups for SSL training.\n+ Improvement in training physics-based models in much lesser time.\n+ covered a comprehensive list of symmetry groups for different PDEs. This would be very useful to the community.", "weaknesses": "- For a person from an engineering background, I do think SSL is quite similar to the ideas of Multi-fidelity modeling (https://arxiv.org/abs/1609.07196, https://arxiv.org/abs/2110.04170, https://arxiv.org/abs/1903.00104). Both help reduce the training data and both are very easy to implement. Unfortunately, SSL for the quit nh\n- The authors talk about taking steps for building foundation models, I would imagine some kind of amalgamation of all the PDEs while training the SSL. Like you use all the data from all the PDEs to build and train an SSL model. That seems closer to how foundation models would behave. I feel like the paper gives big hopes and doesn't deliver them in the implementation. If they were to not talk about foundation models and just say they build SSL frameworks for PDEs, I would have been very happy.\n- The authors make a comparison between supervised ones and SSL-trained ones. but, there are a whole host of PDE-residual-based models like DeepONets, FNOsetc. (which the authors acknowledge in the supplement). It would be nice to know if SSL is even needed for PDE-based foundation models. Can PDE-residuals sufficiently help us navigate through foundation models? Is there a need for SSL? Thats not established fairly. (note that I acknowledge the comparisons, but is that sufficient to establish the necessity of SSL for PDEs?)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for taking their time to review the paper, praising the novelty of our work, and providing valuable comments and feedback. We address the reviewer’s questions and comments below.\n\n***\n\n### SSL and Multi-fidelity modeling\n\nPart of the reviewer’s response was cut off. Could they please clarify what remained in this comment? \n\n### On foundation models for PDEs\n\nIndeed, we mention foundation models for PDEs as part of an aspirational research program, but only take the first step towards this goal. We will correct the phrasing accordingly. Moreover, to bring our work closer to this ambitious goal, we added new experiments training a common representation from the mixture of our Burgers, KdV and KS datasets (all models of chaotic flow that share many Lie point symmetries), using crops and experimenting with different sets of Lie augmentations that are common to all three.\n\nWe evaluate the representations on parameter regression tasks and report averaged results over 3 runs. Our preliminary results are encouraging yet show that mixing data from different equations is not straightforward:\n\n| Task             | Burgers kin. viscosity (%)   |\n|-----------------------------|--------------|\n| Supervised  | 1.55 ± 0.01 |\n| SSL features |  1.53 ± 0.02|\n\nOn Burgers we match the supervised baseline in a short training regimen (50 total epochs), showing that the model can learn good representations with heterogeneous data sources. However, in the short time frame of the rebuttal, we could not get conclusive results when evaluating on KS or KdV.\nMixing the equations requires addressing new challenges (e.g dealing with different chaotic behavior and different time and length scales between PDEs). It is a very interesting direction for future work.\n\n### Is SSL needed for foundation models? What about PDE-residuals?\n\nThe reviewer raises some good points about the necessity of SSL, and if we understand correctly, their inquiries cover two related issues: (1) the comparative effectiveness of SSL in constructing foundation models, and (2) the requirement of SSL in light of PDE-residual-based methods, which can operate in semi-supervised or unsupervised environments.\n\n\nAddressing (1), to clarify: we are not asserting that SSL is categorically the superior approach for building foundation models or representation learning methods. Nevertheless, based on very strong results in the realm of image processing [5], we are confident in SSL's potential when applied to PDEs.\n\nFor (2), we concur that many potentially effective PDE-residual-based methods exist. However, these methods seem to serve a different purpose, i.e learning an approximation to a differential operator, while SSL attempts to learn a rich yet easy to leverage feature space. Our added experiments with FNO (see summary response) show the latter can benefit the former. Could the reviewer elaborate on the comparison they expect?\n\n### Selection of augmentations\n\nUsing crops as the only augmentation provides a decent baseline. However, useful Lie point augmentations are required for best performance. We have evidence of this for both regression/classification tasks and time-stepping. We report two new experiments in this regard (see Table below):\nFor Burgers’, we trained a representation with the crop augmentation only. It does not outperform the supervised baseline in kinematic viscosity regression as opposed to the representation trained with crop and Lie point augmentations.\n\nFor Navier-Stokes, representations with MSE of 0.0052 (crop only) and 0.0038 (crop + Lie augmentations) reported in Fig. 5 for buoyancy regression are actually different. Our best time-stepping model (UNet + AdaGN conditioning) conditioned on the crop-only representation does not outperform the baseline, in contrast to conditioning on the representation trained with crop and Lie point augmentations.\n\n| Task             | Burgers kin. viscosity (%)   | Navier-Stokes time-stepping MSE (1e-3)   |\n|-----------------------------|--------------|-------------|\n| Supervised baseline  |  1.18 ± 0.07 | 2.37  ±  0.01 |\n| SSL features (Crop only) | 2.3 ± 0.2 | 2.9 ± 0.8 |\n| SSL features (Crop + Lie augmentations) | 0.97 ± 0.04  | 2.35 +- 0.03 |\n\nIt may be possible to improve on a given task by `\"overfitting” the augmentation parameters, but it is observed both in SSL for computer vision [5] and in our work that a good set of augmentations for SSL provides the best results across a wide range of tasks. In our experiments for example, the best representation evaluated during pre-training works out of the box for time-stepping (Table 1).\n\n### Computational overhead of SSL vs supervised\n\nSSL pre-training typically has higher training costs compared to just supervised methods. Crucially, pre-training is a fixed cost, and subsequent training costs for downstream tasks starting from SSL features are negligible: one linear layer on top of frozen SSL features is usually sufficient, whereas supervised learning requires tuning the whole network to get similar results [1-2]. Since SSL features are computationally cheap to use and transfer better [3-4], the SSL approach is advantageous when features are reused. Since the purpose of SSL and DeepONets seem different (see comment above), could the reviewer elaborate on the comparison they expect?\n\n**References:**\n\n[1] Chen, Ting, et al. \"A simple framework for contrastive learning of visual representations.\" International conference on machine learning. ICML 2020.\n\n[2] Bardes, Adrien et al. \"Vicreg: Variance-invariance-covariance regularization for self-supervised learning.\" ICLR 2022.\n\n[3] Ericsson, Linus et al. \"How well do self-supervised models transfer?.\" CVPR 2021.\n\n[4] Tian, Yonglong, et al. \"Rethinking few-shot image classification: a good embedding is all you need?.\"ECCV 2020\n\n[5] Oquab, Maxime, et al. “DINOv2: learning robust visual features without supervision”, arXiv preprint (2023)"}, {"review_id": "ehlNscdP5I", "reviewer": "Reviewer_xeoP", "summary": "This paper proposes to learn general-purpose representations of PDEs from heterogeneous  data by implementing joint embedding methods for self-supervised learning. Learned representation outperforms baseline approaches for invariant tasks such as regressing the coefficients of a PDE and improve the time-stepping performance of neural solvers.", "questions": "- Why do authors only consider Lie point symmetry for PDE? Is there any other symmetry group that applies to PDE?\n- If a PDE does not have periodic boundary conditions, does augmentation using symmetry group still valid?\n- Can elaborate more on why learned representation cannot improve time-stepping performance for the Burgers equation?\n", "limitations": "The limitations of this paper are discussed.", "rating": 3, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 1, "strengths": "- Propose a general framework for self-supervised learning in PDE by using symmetry transformation in PDE.\n- Learned representation can be used in several tasks, such as determining unknown parameters and time-stepping.\n- This paper is well-written and the presentation is clear.\n", "weaknesses": "- The novelty of this paper is quite marginal. The SSL framework adopted in the paper is well-known and not customized for the specific PDE problem. The only contribution is the augmentation of PDE solutions according to symmetry groups, which is also well-studied in previous literature, such as in [1].\n- The evaluation on time-stepping is obviously not enough. For example, this paper doesn’t compare with several important baselines, such as MPPDE [2], FNO [3], UNet, PINN, etc. I think the authors need to show the effectiveness of using learned representation to improve at least a few of these models on limited labeled training data.\n- The regression task of determining external buoyancy force (a constant value) in the NS equation is quite simple in practice. In practice, there are more complicated forcing terms. For example, in FNO [3], they use a forcing term containing cosine and sine waves. Can authors use the forcing term in FNO and regress the parameters in that forcing term?\n- The viscosity in the NS equation is also an important parameter. Can authors also provide regression results on determining the viscosity?\n\n[1] Brandstetter, Johannes, Max Welling, and Daniel E. Worrall. \"Lie point symmetry data augmentation for neural PDE solvers.\" International Conference on Machine Learning. PMLR, 2022.\n\n[2] Brandstetter, Johannes, Daniel Worrall, and Max Welling. \"Message passing neural PDE solvers.\" arXiv preprint arXiv:2202.03376 (2022).\n\n[3] Li, Zongyi, et al. \"Fourier neural operator for parametric partial differential equations.\" arXiv preprint arXiv:2010.08895 (2020).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their feedback and insightful comments. We address the reviewer’s questions below and at times, challenge their criticisms. We welcome further discussion.\n\n***\n### Novelty\n\nWe acknowledge that SSL for computer vision and the study of symmetry groups of PDEs are *separately* well-established. We respectfully disagree on novelty, and share examples that shed light on our contributions:\n- The role of data augmentation in supervised learning [1] and SSL varies fundamentally. In SSL, augmentations dictate properties preserved by the encoder, and no learning happens without them. As further illustration of this, our conclusions differ from [1] on the utility of augmentations. For instance, time translations often impede learning good representations for Burgers' or Navier-Stokes, contrasting with findings in [1] (Tables 3d and 4d).\n- The Lie-Trotter techniques for implementing augmentations (Appendix B) are new to the ML for PDE community as far as we are aware, and enable more smoothly and universally applying augmentations. This implementation may even be useful outside of SSL. For comparison Ref. [1] applied Lie symmetries by choosing an arbitrary order for the basis elements, which neither guarantee universality (there exist augmentations which cannot be performed) nor smoothness (the Lie algebra has no guarantees on norm of augmentation). \n- SSL comes with many moving parts to which subsequent representations may or may not be sensitive [2]. For example, it was unclear a priori that using the generic task of coefficient or initial condition regression as metrics for pre-training evaluation (replacing the ImageNet top-1 accuracy in computer vision) would select embeddings that transfer nicely to other tasks such as time-stepping.\n\nSSL algorithms based on joint embedding architectures have application in various regimes outside its original use in computer vision, including published works for graphs [3], language [4], point clouds [5], and more. These works elucidate where SSL is useful, how it can be adapted, and the extent to which SSL can be applied for learning real-world data in general. Overall, we believe it is very valuable to share our insights gained for PDE data.\n\n\n### Improving other models such as FNO\n\nWe acknowledge that using our representations to improve other baselines would make this work stronger. We refer the reviewer to our general comment, where we show the effectiveness of the SSL representations to condition FNOs, Fourier U-Nets, and U-Nets with a supplementary conditioning method, AdaGN.\n\n### Determining external buoyancy in the NS equation\n\nFrom a certain viewpoint, regressing the buoyancy force (a constant), within the NS equation might appear straightforward. One can simply calculate the derivatives and functions in the PDE and regress it. However, this is not the case here:\nCrucially, our network is agnostic to the specific form of the PDE. To regress the buoyancy force, it has to disentangle the parameter from the complex, nonlinear map of the PDE. \nEven in a supervised setting, the performance is far from the buoyancy's resolution, showing that it remains a hard task empirically. \n\nThis makes our task valid to gauge the quality of our representation.\n\n### Navier-Stokes regression results on viscosity\n\nThe NS benchmark we use (PDEArena) comes with a fixed viscosity. We did not have sufficient time to generate a new dataset with varying viscosities.  However, we will add this task in a revised version of the manuscript. \n\n### On the choice of Lie point symmetry for PDEs\nWe primarily focused on Lie point symmetries due to their systematic derivation and well-established applications in PDEs. We acknowledge that other forms of symmetries, such as approximate, Bäcklund, and discrete transformations, can also be considered. However, these are typically challenging to derive, may introduce other sources of errors and, like Lie symmetries, are typically derived in infinite domain settings (so do not address boundary issues). When learning data from multiple different types of PDEs, these may be useful but this would require different implementations and changes to our setup. \n\n### On augmentation using symmetry groups in non-periodic boundary conditions\nYou have touched upon an important limitation that we also point out in our work. Since Lie symmetries do not preserve boundaries, we limit and test the magnitude of transformations to avoid large errors (see discussion beginning on line 142). In short, addressing boundary conditions more directly would either escalate the complexity of the task or divert us from the generality we aimed to maintain. Please refer to the general response and response to reviewer VTPv for a more complete answer.\n\n### On conditioning time-stepping for Burgers equation\nIt is likely that the dynamics of Burgers’ equation are easy to predict with little room for improvement compared to KdV and KS given the same resources. A similar observation is made in [1] (Appendix D, Tables 3 and 4), where the normalized MSE for Burgers is far lower than that for KdV and KS, when provided with a few hundreds of samples.\n\n\n**References:**\n\n[1] Brandstetter, Johannes et al. \"Lie point symmetry data augmentation for neural PDE solvers.\" ICML 2022.\n\n[2] Balestriero, Randall, et al. \"A cookbook of self-supervised learning.\" arXiv preprint (2023).\n\n[3] Xie, Yaochen, et al. \"Self-supervised learning of graph neural networks: A unified review.\" IEEE TPAMI (2022).\n\n[4] Chuang, Ching-Yao, et al. \"Debiased contrastive learning.\" NeurIPS 2020.\n\n[5] Jiang, Li, et al. \"Guided point contrastive learning for semi-supervised point cloud semantic segmentation.\" CVPR 2021."}, {"review_id": "9wG6CNrFYm", "reviewer": "Reviewer_rnp9", "summary": "The paper presents an approach for self-supervised learning of PDEs. The main approach uses Lie symmetries (similar e.g. to translational symmetries for images) in the solutions of differential equations not only for data augmentation (as has been done before), but for representation learning for downsteam tasks.", "questions": "I guess for a paper of this level ob abstraction, you do not need to define what a group and what a group operation is? (l. 115)", "limitations": "The approach is certainly interesting. However, where I am working with differential equations, one only knows the symmetries from the differential equations. Hence, I rarely see a need for this approach.\n\nOtherwise, limitations and open questions are given both in the paper (e.g. l. 99ff, 130ff, etc.)\n", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "The paper is a pleasure to read. There are nice diagrams (perhaps even too fancy?) and a nice verbal presentation everywhere. The appendix comes with an introduction into the theory, which is also well written and correct. The literature is thoroughly cited (Of course, there is always more to cite, but the amount is huge and the choice are reasonable.) Many experiments are conducted. Even limitations and drawbacks are clearly stated, such that after reading the paper a really feel what works easily, what works with problems, and what is not yet known. I feel well-informed after reading the paper.\n\nFurthermore, the topic is clearly relevant for NeurIPS: it is a non-trivial machine learning question with relevant application domains.\n\nIn particular the correct usage of PDE-terms is refreshing to see when reviewing NeurIPS papers about PDEs. Here, I could find no problems. All terms were used correctly, even though more could and should have been said about PDEs, but the paper only has 9 pages; more is saind in the appendix. This quality in mathematical writing is not the norm at NeurIPS! And, at the same time, the paper is not technical.\n", "weaknesses": "I could follow the paper. However, I have a strong math background that includes symmetries of PDEs. The authors try really hard (and are quite good) in making everything clear to the audience. However, I guess that much of the NeurIPS audience might have problems with some techniques used in the paper.\n\nAfter reading the paper, it was not clear to me what kind of data is necessary. In one of your experiments, you used data on a grid and were thus able to apply ResNet18. More details in this regard (what is possible? what is reasonable?) would have been nice.\n\nQuite a lot of data is necessary for training. If I have 10000 samples of a PDE, I kind of know the PDE anyway. This, together with the high dependency on several hyperparameters and the additional limitation (see below), is the main problem of this paper which make is \"just a strong NeurIPS paper\" instead of an outstanding paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for taking their time to review the paper and providing valuable comments and feedback. We are glad the reviewer is happy with the quality of the experiments and the completeness of our manuscript. We address the reviewer’s questions and comments below.\n\n***\n\n### Question 1:\n>After reading the paper, it was not clear to me what kind of data is necessary. In one of your experiments, you used data on a grid and were thus able to apply ResNet18. More details in this regard (what is possible? what is reasonable?) would have been nice.\n\nThank you for raising this interesting point. Going further, real-world data could indeed come in the form of irregular grids making the use of architectures such as ResNet more challenging. There, one could leverage neural networks for solving PDEs with graph based methods that handle irregular grids [1-2] and SSL frameworks that have been extended to architectures more suited to such data [3]. Additionally, SSL for computer vision comes with recipes for using transformers, which could also be leveraged for irregular grids although such an approach may require more data than ResNets or GNNs. We will add discussion around these points to the updated draft of our paper.\n\nAs an aside, Lie symmetries are still implementable on irregular grids, since implementing Lie point symmetries only requires knowing the values of the solution at any given point (hence the name). This is not necessarily true of other types of symmetries such as Lie Bäcklund symmetries which require knowing the solution at a neighborhood around every point to implement them.\n\n### Question 2:\n>Quite a lot of data is necessary for training. If I have 10000 samples of a PDE, I kind of know the PDE anyway. This, together with the high dependency on several hyperparameters and the additional limitation (see below), is the main problem of this paper which make is \"just a strong NeurIPS paper\" instead of an outstanding paper.\n\nAlthough our current approach assumes we know the family of the PDE, our datasets mix realizations that have different initial conditions and equation parameters such as kinematic viscosity or buoyancy. In that regard, having models that, as those we propose, generalize to unseen initial conditions or equation parameters is already very valuable.\n\nWe also refer the reviewer to our answer to reviewer swdq, where we provide an experiment learning representations from a dataset mixing Burgers, KdV and KS realizations. Although very preliminary, this setup would remove the need to know the PDE family, while alleviating the data requirement for each equation.\n\nWe also show in figure 4 that we already have improved performance with our approach with only a few thousands data points. While this may already be a big requirement in certain scenarios, this indicates that it can still be applied in more data-constrained regimes.\n\n### Other relevant comment:\n\n> The approach is certainly interesting. However, where I am working with differential equations, one only knows the symmetries from the differential equations.\n\nThank you for pointing this out. Indeed, knowing the differential equation allows one to derive symmetries. This is the most ideal setting. Nevertheless, symmetries can also be deduced by knowing what class of differential equations one PDE is a part of or knowing properties of the PDE. For example, all flow related equations share common symmetries such as translations and Galilean boosts (e.g., see the added shared experiment in response to reviewer swdq). Furthermore, many symmetries can be derived from known conservation properties by Noether's theorems. So in general, knowing an equation precisely is best, but this is not always required in practice. Thank you again for raising this clarification. We will comment on this further in our updated draft.\n\n**References:**\n\n[1] Belbute-Peres, Filipe De Avila, Thomas Economon, and Zico Kolter. \"Combining differentiable PDE solvers and graph neural networks for fluid flow prediction.\" international conference on machine learning. PMLR, 2020.\n\n[2] Brandstetter, Johannes, Daniel Worrall, and Max Welling. \"Message passing neural PDE solvers.\" arXiv preprint arXiv:2202.03376 (2022).\n\n[3] You, Yuning et al., “Graph Contrastive Learning with Augmentations”, NeurIPS 2020"}]

Author rebuttal:
## Summary Response

We thank the reviewers for their insightful comments and many great questions. We have responded to each reviewer’s comments separately, and are sharing a summary response covering the common threads in the reviewers’ responses. To enhance our responses, we have added experiments to back up many of the points we make below. We truly value further input from the reviewers on these matters.

***

### Replicating experiments for different architectures:
Multiple reviewers pointed out that they would be more confident in our results if we replicated time-stepping for different architectures, especially those that operate in Fourier space. We agree with this perspective, and are happy to share new results for different architectures:
- For Burgers: FNO1d, conditioned as detailed in our work.
- For Navier-Stokes: Fourier Neural Operator [2] with SpaSpec conditioning [3], Fourier U-Nets with Addition conditioning [3], and an additional conditioning method of U-Nets, AdaGN [3].




| Burger's (NMSE)               | ResNet1d     | FNO1d       |
|-----------------------------|--------------|-------------|
| No conditioning (baseline)  | 0.110 ± 0.008| 0.184 ± 0.002 |
| Representation conditioning | 0.108 ± 0.011| 0.173 ± 0.002 |



| Navier-Stokes (MSE x 1e-3)                                              | U-Nets + Addition | U-Nets + AdaGN | FNO + SpaSpec  | Fourier U-Nets + Addition |
|-----------------------------------------------------------------|-------------------|----------------|----------------|--------------------------|
| Time conditioning only (baseline)                                | 2.60 ± 0.05       | 2.37 ± 0.01    | 13.4 ± 0.5    | 3.31 ± 0.06             |
| Time conditioning + representation conditioning (ours)           | 2.47 ± 0.02       | 2.35 ± 0.03    | 13 ± 1.0        | 2.37 ± 0.05             | 




### Handling boundary conditions:

As two reviewers have pointed out and as detailed in the limitations section of our work, boundary conditions can be a source of error when implementing Lie point symmetries. Since virtually all symmetry derivations are done assuming infinite or periodic domains, it is not necessarily obvious how to mathematically handle these boundaries in derivations of symmetries. In practice, this technical issue is not as major as one would expect:
- A similar situation arises in image settings where augmentations like resize or translation cause issues on the boundary. However, the goal is to learn global features that are invariant to these boundary issues. In pursuit of that higher level goal, SSL is still effective at learning with these augmentations included.
- Crops are an essential augmentation both here and in the standard image setting [1]. For PDEs, this is partly because crops help bias the network to learn features that are invariant to whether the input was taken near a boundary or not.
- There is no obvious technical solution to this problem as far as we are aware. Indeed, other forms of PDE symmetry groups like approximate symmetries, discrete symmetries, or Lie Bäcklund symmetries are also derived in infinite domains and do not address boundary issues. These symmetry groups are typically much harder to derive, more complicated, and do not necessarily come with the nice Lie algebraic structure associated with the Lie point symmetry group.

There is much more to say on this point and our more detailed thoughts are captured in our response to reviewer VTPv which we recommend that reviewers look at. We are open to more ideas and happy to continue the discussion on this point.   



**References:**

[1] Chen, Ting, et al. "A simple framework for contrastive learning of visual representations." International conference on machine learning. PMLR, 2020.

[2] Li, Zongyi, et al. "Fourier neural operator for parametric partial differential equations." arXiv preprint arXiv:2010.08895 (2020).

[3] Gupta, Jayesh K., and Johannes Brandstetter. "Towards multi-spatiotemporal-scale generalized pde modeling." arXiv preprint arXiv:2209.15616 (2022).



Meta-review: The paper proposes to use self-supervised learning for improving certain PDE-related problems. Many PDEs come with point Lie symmetries, so this can generate different 'views' of the dataset, thus standard techniques can be used. 
The numerical examples report significant improvement over baselines in some cases. 
The task considered in the paper is regression of the coefficients, i.e. efficient representation is indeed important. 

The reviewers agree that the paper is well-written. However, the problem itself is not very common (regression task) and is indeed a classical regression task. Thus, the paper combines two well-known facts: some PDE has symmetries; 
this can be used together with SSL, and there is no much specifics about that. Also, the question of boundary condition is valid, thus making the method even more empirical.

Overall, although the paper is a combination of two standard techniques, it has not been shown to be effective before in this particular domain (inverse problems).The more detailed studies will follow.
Decision: Accept (poster)