SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Grégoire Mialon†  
Meta, FAIR
Quentin Garrido†  
Meta, FAIR  
Univ Gustave Eiffel, CNRS, LIGM
Hannah Lawrence Meta, FAIR MIT
Danyal Rehman  
MIT
Yann LeCun Meta, FAIR NYU
Bobak T. Kiani* MIT
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.
Dynamical systems governed by differential equations are ubiquitous in fluid dynamics, chemistry, astrophysics, and beyond. Accurately analyzing and predicting the evolution of such systems is 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 as an alternative method for analyzing differential equations at a fraction of the cost [1, 2, 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 to 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. In 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: gmiumon@meta.com, garridoq@meta.com, and bkiani@mit.edu, † Equal contribution
Figure 1: A high-level overview of the self-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, self-supervised learning uses augmentations (e.g. color-shifting for images, or Lie symmetries for PDEs) to train a network  $f_{\theta}$  to produce useful representations from input images. Given a smaller set of labeled data, these representations can then be used as inputs to a supervised learning pipeline, performing tasks such as predicting class labels (images) or regressing the kinematic viscosity  $\nu$  (Burgers' equation). Trainable steps are shown with red arrows; importantly, the representation function learned via SSL is not altered during application to downstream tasks.
To tackle these broader challenges, we take inspiration from the recent success of self-supervised learning (SSL) as a 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 [7]. 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 to solve downstream tasks with a limited amount of data (such as kinematic viscosity regression), as illustrated in Figure 1.
To do so, we leverage the joint embedding framework [8] for self-supervised learning, a popular paradigm for learning visual representations from unlabeled data [9, 10]. It consists of training an encoder to enforce similarity between embeddings of two augmented versions of a 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, SSL relies on human intuition to build hand-crafted augmentations (e.g. recoloring and cropping), whereas PDEs are endowed with a group of symmetries preserving the governing equations of the PDE [11, 12]. These symmetry groups are important because creating embeddings that are invariant under them would allow to 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 E). Symmetry groups are well-studied for common PDE families, and can be derived systematically or calculated from computer algebra systems via tools from Lie theory [11, 13, 14].
Contributions: We present a general framework for performing SSL for PDEs using their corresponding symmetry groups. In particular, we show that by exploiting the analytic group transformations from one PDE solution to 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-motivated 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
Figure 2: Pretraining and evaluation frameworks, illustrated on Burgers' equation. (Left) Self-supervised pretraining. We generate augmented solutions  $\pmb{x}$  and  $\pmb{x}'$  using Lie symmetries parametrized by  $g$  and  $g'$  before passing them through an encoder  $f_{\theta}$ , yielding representations  $\pmb{y}$ . The representations are then input to a projection head  $h_{\theta}$ , yielding embeddings  $z$ , on which the SSL loss is applied. (Right) Evaluation protocols for our pretrained representations  $\pmb{y}$ . On new data, we use the computed representations to either predict characteristics of interest, or to condition a neural network or operator to improve time-stepping performance.
adapting self-supervision to the physical sciences. We hope this work will serve as a starting point for developing foundation models on more complex dynamical systems using our framework.
We now describe our general framework for learning representations from and for diverse sources of PDE data, which can subsequently be used for a wide range of tasks, ranging from regressing characteristics of interest of a PDE sample to improving neural solvers. To this end, we adapt a popular paradigm for representation learning without labels: the joint-embedding 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,  $f_{\theta}$ , producing representations that can be used for downstream tasks. The SSL loss function is comprised of a similarity loss  $\mathcal{L}_{\mathrm{sim}}$  between projections (through a projector  $h_{\theta}$ , which helps generalization [15]) of the pairs of views, to make their representations invariant to augmentations, and a regularization loss  $\mathcal{L}_{\mathrm{reg}}$ , to avoid trivial solutions (such as mapping all inputs to the same representation). The regularization term can consist of a repulsive force between points, or regularization on the covariance matrix of the embeddings. Both functions similarly, as shown in [16]. This pretraining procedure is illustrated in Fig. 2 (left) in the context of Burgers' equation.
In this work, we choose variance-invariance-covariance regularization (VICReg) as our self-supervised loss function [9]. Concretely, let  $Z, Z' \in \mathbb{R}^{N \times D}$  contain the  $D$ -dimensional representations of two batches of  $N$  inputs with  $D \times D$  centered covariance matrices,  $\mathrm{Cov}(Z)$  and  $\mathrm{Cov}(Z')$ . Rows  $Z_{i,:}$  and  $Z_{i,:}'$  are two views of a shared input. The loss over this batch includes a term to enforce similarity  $(\mathcal{L}_{\mathrm{sim}})$  and a term to avoid collapse and regularize representations  $(\mathcal{L}_{\mathrm{reg}})$  by pushing elements of the encodings to be statistically identical:
$$
\mathcal {L} \left(\boldsymbol {Z}, \boldsymbol {Z} ^ {\prime}\right) \approx \frac {\lambda_ {i n v}}{N} \underbrace {\sum_ {i = 1} ^ {N} \left\| \boldsymbol {Z} _ {i , :} - \boldsymbol {Z} _ {i , :} ^ {\prime} \right\| _ {2} ^ {2}} _ {\mathcal {L} _ {\operatorname {s i m}} (\boldsymbol {Z}, \boldsymbol {Z} ^ {\prime})} + \frac {\lambda_ {r e g}}{D} \underbrace {\left(\left\| \operatorname {C o v} (\boldsymbol {Z}) - \boldsymbol {I} \right\| _ {F} ^ {2} + \left\| \operatorname {C o v} (\boldsymbol {Z} ^ {\prime}) - \boldsymbol {I} \right\| _ {F} ^ {2}\right)} _ {\mathcal {L} _ {\operatorname {r e g}} (\boldsymbol {Z}) + \mathcal {L} _ {\operatorname {r e g}} (\boldsymbol {Z} ^ {\prime})}, \tag {1}
$$
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_{1})$ , spatial shifts  $(g_{2})$ , and Galilean boosts  $(g_{3})$  with their corresponding infinitesimal transformations in the Lie algebra placed inside the figure. The shaded red square denotes the original  $(x, t)$ , while the dotted line represents the same points after the augmentation is applied.
where  $\| \cdot \| _F$  denotes the matrix Frobenius norm and  $\lambda_{inv},\lambda_{reg}\in \mathbb{R}^{+}$  are hyperparameters to weight the two terms. In practice, VICReg separates the regularization  $\mathcal{L}_{reg}(Z)$  into two components to handle diagonal and non-diagonal entries  $\operatorname {Cov}(Z)$  separately. For full details, see Appendix C.
Adapting VICReg to learn from PDE data: Numerical PDE solutions typically come in the form of a 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,x)$ : a vector of  $t$  time values, and a vector of  $x$  spatial values. We therefore process the sample to obtain a  $(3,t,x)$  tensor with the last two channels encoding spatial and temporal discretization, which can be naturally fed to neural networks tailored for images such as ResNets [17]. From these, we extract the representation before the classification layer (which is unused here). It is worth noting that convolutional neural networks have become ubiquitous in the literature [18, 12]. While the VICReg default hyper-parameters did not require substantial tuning, tuning was crucial to probe the quality of our learned representations to monitor the quality of the pre-training 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 a linear classifier on ImageNet. In our framework, we pick regression of a PDE coefficient, or regression of the initial conditions when there is no coefficient in the equation. The latter, commonly referred to as the inverse problem, has the advantage of being applicable to any PDE, and is often a challenging problem in the numerical methods community given the ill-posed nature of the problem [19]. Our approach for a particular task, kinematic viscosity regression, is schematically illustrated in Fig. 2 (top right). More details on evaluation tasks are provided in Section 4.
2.2 Augmentations and PDE Symmetry Groups
Background: PDEs formally define a systems of equations which depend on derivatives of input variables. Given input space  $\Omega$  and output space  $\mathcal{U}$ , a PDE  $\Delta$  is a system of equations in independent variables  $x\in \Omega$ , dependent variables  $u:\Omega \to \mathcal{U}$ , and derivatives  $(\pmb {u}_x,\pmb {u}_{xx},\dots)$  of  $\pmb{u}$  with respect to  $\pmb{x}$ . For example, the Kuramoto-Sivashinsky equation is given by
$$
\Delta (x, t, u) = u _ {t} + u u _ {x} + u _ {x x} + u _ {x x x x} = 0. \tag {2}
$$
Informally, a symmetry group of a PDE  $G^2$  acts on the total space via smooth maps  $G: \Omega \times \mathcal{U} \to \Omega \times \mathcal{U}$  taking solutions of  $\Delta$  to other solutions of  $\Delta$ . More explicitly,  $G$  is contained in the symmetry group of  $\Delta$  if outputs of group operations acting on solutions are still a solution of the PDE:
$$
\Delta (\boldsymbol {x}, \boldsymbol {u}) = 0 \Rightarrow \Delta [ g \cdot (\boldsymbol {x}, \boldsymbol {u}) ] = 0, \quad \forall g \in G. \tag {3}
$$
2A group  $G$  is a set closed under an associative binary operation containing an identity element  $e$  and inverses (i.e.,  $e \in G$  and  $\forall g \in G : g^{-1} \in G$ ).  $G : \mathcal{X} \to \mathcal{X}$  acts on a space  $\mathcal{X}$  if  $\forall x \in \mathcal{X}, \forall g, h \in G : ex = x$  and  $(gh)x = g(hx)$ .
For PDEs, these symmetry groups can be analytically derived [11] (see also Appendix A for more formal details). The types of symmetries we consider are so-called Lie point symmetries  $g: \Omega \times \mathcal{U} \to \Omega \times \mathcal{U}$ , which act smoothly at any given point in the total space  $\Omega \times \mathcal{U}$ . For the Kuramoto-Sivashinsky PDE, these symmetries take the form depicted in Fig. 3:
Temporal Shift:  $g_{1}(\epsilon):(x,t,u)\mapsto (x,t + \epsilon ,u)$
Spatial Shift:  $g_{2}(\epsilon):(x,t,u)\mapsto (x + \epsilon ,t,u)$  (4)
Galilean Boost:  $g_{3}(\epsilon):(x,t,u)\mapsto (x + \epsilon t,t,u + \epsilon)$
As in this example, every Lie point transformation can be written as a one parameter transform of  $\epsilon \in \mathbb{R}$  where the transformation at  $\epsilon = 0$  recovers the identity map and the magnitude of  $\epsilon$  corresponds to the "strength" of the corresponding augmentation. Taking the derivative of the transformation at  $\epsilon = 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 we 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 B.
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 as regressing viscosity), the Lie symmetries used should neither depend on viscosity nor change the viscosity of the output solution. Moreover, there is no guarantee as to which Lie symmetries are the most "natural", i.e. most likely 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.
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) [22, 23]. In our experiments, we find that Lie point augmentations applied at relatively small strengths perform the best (see Appendix E), 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 a regular grid of sampled points.
In this section, we provide a 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 [9, 20], 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, 13]. Most related to our 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
${}^{3}$  Technically,  $\epsilon$  is the magnitude and direction of the transformation vector for the basis element of the corresponding generator in the Lie algebra.
networks to satisfy a continuity equation [28], or enforcing physical constraints in dynamic mode decomposition [29]. For Hamiltonian systems, various works have designed algorithms that respect the symplectic structure or conservation laws of the Hamiltonian [30, 31].
Equations considered: We focus on flow-related equations here as a testing ground for our methodology. In our experiments, we consider the four equations below, which are 1D evolution equations apart from the Navier-Stokes equation, which we consider in its 2D spatial form. For the 1D flow-related equations, we impose periodic boundary conditions with  $\Omega = [0,L]\times [0,T]$ . For Navier-Stokes, boundary conditions are Dirichlet ( $v = 0$ ) as in [18]. Symmetries for all equations are listed in Appendix E.
1. The viscous Burgers' Equation, written in its "standard" form, is a nonlinear model of dissipative flow given by
$$
u _ {t} + u u _ {x} - \nu u _ {x x} = 0, \tag {5}
$$
where  $u(x,t)$  is the velocity and  $\nu \in \mathbb{R}^{+}$  is the kinematic viscosity.
2. The Korteweg-de Vries (KdV) equation models waves on shallow water surfaces as
$$
u _ {t} + u u _ {x} + u _ {x x x} = 0, \tag {6}
$$
where  $u(x,t)$  represents the wave amplitude.
3. The Kuramoto-Sivashinsky (KS) equation is a model of chaotic flow given by
$$
u _ {t} + u u _ {x} + u _ {x x} + u _ {x x x x} = 0, \tag {7}
$$
where  $u(x,t)$  is the dependent variable. The equation often shows up in reaction-diffusion systems, as well as flame propagation problems.
4. The incompressible Navier-Stokes equation in two spatial dimensions is given by
$$
\boldsymbol {u} _ {t} = - \boldsymbol {u} \cdot \nabla \boldsymbol {u} - \frac {1}{\rho} \nabla p + \nu \nabla^ {2} \boldsymbol {u} + \boldsymbol {f}, \quad \nabla \boldsymbol {u} = 0, \tag {8}
$$
where  $\pmb{u}(\pmb{x},t)$  is the velocity vector,  $p(\pmb{x},t)$  is the pressure,  $\rho$  is the fluid density,  $\nu$  is the kinematic viscosity, and  $\pmb{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 case of Burgers' equation or pseudo-spectral methods used to generate PDE learning benchmarking data (see Appendix F) [12, 18, 32]. Burgers', KdV and KS's solutions are 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 1.
Pretraining: For each equation, we pretrain a ResNet18 with our SSL framework for 100 epochs using AdamW [33], a batch size of 32 (64 for Navier-Stokes) and a learning rate of 3e-4. We then freeze its weights. To evaluate the resulting representation, we (i) train a linear head on top of our features and on a new set 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 a 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 problem 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, we train the same architecture, i.e. a ResNet18, using the MSE loss on downstream labels. Unless stated otherwise, we train the linear model for 30 epochs using Adam. Further details are in Appendix F.
Kinematic viscosity regression (Burgers): We pretrain a ResNet18 on 10,000 unlabeled realizations of Burgers' equation, and use the resulting features to train a linear model on a smaller, labeled