Self-Supervised Learning with Lie Symmetries for Partial Differential Equations

TL;DR

This work tackles learning robust representations of PDEs from heterogeneous, potentially unlabeled data by combining self-supervised learning with Lie point symmetries as principled data augmentations. It adapts the VICReg joint-embedding framework to PDE data by treating space-time PDE solutions as multi-channel inputs and pretraining an encoder on unlabeled realizations, then evaluating representations on downstream tasks like parameter regression and time-stepping. The authors demonstrate that symmetry-informed SSL representations can outperform supervised baselines on several tasks across Burgers', KdV, KS, and Navier–Stokes equations, and that larger unlabeled datasets further boost performance. They discuss limitations related to boundary conditions and cross-PDE transfer, and chart future directions toward explicit equivariance, foundation-model-style multi-PDE representations, and extensions to broader scientific data domains.

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. Code: https://github.com/facebookresearch/SSLForPDEs.

Figures (7)

• 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.
• Figure 2: Pretraining and evaluation frameworks, illustrated on Burgers' equation. (Left) Self-supervised pretraining. We generate augmented solutions ${\bm{x}}$ and ${\bm{x}}'$ using Lie symmetries parametrized by $g$ and $g'$ before passing them through an encoder $f_\theta$, yielding representations ${\bm{y}}$. The representations are then input to a projection head $h_\theta$, yielding embeddings ${\bm{z}}$, on which the SSL loss is applied. (Right) Evaluation protocols for our pretrained representations ${\bm{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.
• 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.
• 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 we use, the performance improves, further reinforcing the usefulness of self-supervised representations. (Right) Buoyancy regression on Navier-Stokes' equation. We notice a similar trend as 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.
• Figure 5: (Left) Isolating effective augmentations for Navier-Stokes. Note that we do not study $g_3$, $g_7$ and $g_9$, which are respectively counterparts of $g_2$, $g_6$ and $g_8$ applied in $y$ instead of $x$. (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.

Theorems & Definitions (2)

• Example B.1: Exponential map on symmetry generator of Burger's equation
• proof