SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Table 1: Downstream evaluation of our learned representations for four classical PDEs (averaged over three runs, the lower the better  $(\downarrow)$ ). The normalized mean squared error (NMSE) over a batch of  $N$  outputs  $\widehat{\pmb{u}}_k$  and targets  $\pmb{u}_k$  is equal to  $\mathrm{NMSE} = \frac{1}{N}\sum_{k=1}^{N}\|\widehat{\pmb{u}}_k - \pmb{u}_k\|_2^2/\|\widehat{\pmb{u}}_k\|_2^2$ . Relative error is similarly defined as  $\mathrm{RE} = \frac{1}{N}\sum_{k=1}^{N}\|\widehat{\pmb{u}}_k - \pmb{u}_k\|_1/\|\widehat{\pmb{u}}_k\|_1$ . For regression tasks, the reported errors 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 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 (t,x,(y))	256×128	256×128	448×224	56×128×128
Characteristic of interest	Init. coeffs	Init. coeffs	Kinematic viscosity	Buoyancy
Regression metric	NMSE (↓)	NMSE (↓)	Relative error %(↓)	MSE (↓)
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 (↓)	NMSE (↓)	NMSE (↓)	MSE ×10-3(↓)
Baseline	0.508 ± 0.102	0.549 ± 0.095	0.110 ± 0.008	2.37 ± 0.01
+ SSL 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 Table 1 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-supervised approach to 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 to solve the inverse problem by regressing initial condition parameters from a snapshot of future time evolutions of the solution. Following [34, 12], for a domain  $\Omega = [0, L]$ , a truncated Fourier series, parameterized by  $A_{k}, \omega_{k}, \phi_{k}$ , is used to generate initial conditions:
$$
u _ {0} (x) = \sum_ {k = 1} ^ {N} A _ {k} \sin \left(\frac {2 \pi \omega_ {k} x}{L} + \phi_ {k}\right). \tag {9}
$$
Our task is to regress the set of  $2N$  coefficients  $\{A_k, \omega_k : k \in \{1, \dots, N\}\}$  from a snapshot of the solution starting at  $t = 20$  to  $t = T$ . This way, the initial conditions and first-time steps are never seen during training, making the problem non-trivial. For all conducted tests,  $N = 10$ ,  $A_k \sim \mathcal{U}(-0.5, 0.5)$ , and  $\omega_k \sim \mathcal{U}(-0.4, 0.4)$ . By neglecting phase shifts,  $\phi_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 1, 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,  $\boldsymbol{f} = (c_x, c_y)^\top$ , is constant in the two spatial directions over the course of a given evolution, and our aim is to regress the magnitude of this force  $\sqrt{c_x^2 + c_y^2}$  given a solution to the PDE. We reuse the dataset generated in [18], where  $c_x = 0$  and  $c_y \sim \mathcal{U}(0.2, 0.5)$ . 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 Table 1, 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.
Table 2: One-step validation MSE (rescaled by  $1e3$ ) for time-stepping on Navier-Stokes with varying buoyancies for different combinations of architectures and conditioning methods. Architectures are taken from [18] with the same choice of hyper-parameters. Results with ground truth buoyancies are an upper-bound on the performance a representation containing information on the buoyancy.
Architecture	UNetmod64	UNetmod64	FNO128modes16	UF1Netmodes16
Conditioning method	Addition [18]	AdaGN [35]	Spatial-Spectral [18]	Addition [18]
Time conditioning only	2.60 ± 0.05	2.37 ± 0.01	13.4 ± 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 a sequence of time steps (the "history") of a PDE to predict a future sequence of steps. For each equation we consider different conditioning schemes, to 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 as a new channel. The resulting input is fed to a CNN as in [12] to predict the next 20 time steps (illustrated in Fig. 4 (bottom right) in the context of Burgers' equation). As shown in Table 1, 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 across 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 hidden activations (denoted in [18] 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 to the conditioning embedding as in [18]. 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 [18], 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 a 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 [18]. Conditioning on our representation outperforms the baseline without conditioning.
We also report results for different architectures and conditioning methods for Navier-Stokes in Table 2 and Burgers in Table 8 (Appendix F.1) 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 Navier-Stokes, 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.
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],
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 unlabeled 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.
the former outperforms the latter in the PDE domain we consider. A possible explanation is that enforcing similar representations for two different views of 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. 4 illustrates how more pretraining data benefits our SSL setup, whereas in our experiments it did not help the supervised baselines.
Cropping: Cropping is a natural, effective, and popular augmentation in computer vision [21, 39, 40]. 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 to learn from the PDE data. Cropping also offers a 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-Stokes; here, crops must contain as much information as possible while making sure that pairs of crops have as little overlap as possible (to discourage the network from relying on spurious correlations). This explains the two modes appearing in Fig. 5. We make a similar observation for Burgers, while KdV and KS are less sensitive. Finally, crops help bias the network to learn features that are invariant to whether the input was taken near a boundary or not, thus alleviating the issue of boundary condition preservation during augmentations.
Selecting Lie point augmentations: Whereas cropping alone yields satisfactory representations, Lie point augmentations can enhance performance but require careful tuning. In order to choose which symmetries to include in our SSL pipeline and at what strengths to apply them, we study the effectiveness of each Lie augmentation separately. More precisely, given an equation and each possible Lie point augmentation, we train a SSL representation using this augmentation only and cropping. Then, we couple all Lie augmentations improving the representation over simply using crops. In order for this composition to stay in the stability/convergence radius of the Lie Symmetries, we reduce each augmentation's optimal strength by an order of magnitude. Fig. 5 illustrates this process in the context of Navier-Stokes.
This work leverages Lie point symmetries for self-supervised representation learning from PDE data. Our preliminary experiments with the Burgers', KdV, KS, and Navier-Stokes equations demonstrate the usefulness of the resulting representation for sample or compute efficient estimation of characteristics and time-stepping. Nevertheless, a 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 a particular set of PDEs, but we believe they can be expanded beyond our setting.
Augmentation	Best strength	Buoyancy MSE
Crop	N.A	0.0051 ± 0.0001
single Lie transform
+ t translate g1	0.1	0.0052 ± 0.0001
+ x translate g2	10.0	0.0041 ± 0.0002
+ scaling g4	1.0	0.0050 ± 0.0003
+ rotation g5	1.0	0.0049 ± 0.0001
+ boost g6*	0.1	0.0047 ± 0.0002
+ boost g8**	0.1	0.0046 ± 0.0001
combined
+ {g2, g5, g6, g8}	best / 10	0.0038 ± 0.0001
* linear boost applied in  $x$  direction (see Table 7)
**quadratic boost applied in  $x$  direction (see Table 7)
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.
Learning equivariant representations: Another interesting direction is to expand our SSL framework to learning explicitly equivariant features [41, 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 are 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 are only able to 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 a given PDE are not solely comprised of Lie point symmetries. Other types of symmetries, such as nonlocal symmetries or approximate symmetries like Lie-Backlund symmetries, may also be implemented as potential augmentations [13].
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" SSL augmentations for representation learning has shown significant efficacy over supervised methods. This approach's potential extends beyond the PDEs we study as many problems in mathematics, physics, and chemistry present inherent symmetries that can be harnessed for SSL. Future directions could include implementations of SSL for learning stochastic PDEs, or Hamiltonian systems. In the latter, the rich study of Noether's symmetries in relation to Poisson brackets could be a useful setting to study [11]. Real-world data, as opposed to simulated data, may offer a 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.
The authors thank Aaron Lou, Johannes Brandstetter, and Daniel Worrall for helpful feedback and discussions. HL is supported by the Fannie and John Hertz Foundation and the NSF Graduate Fellowship under Grant No. 1745302.
[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 Ronneberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Žídek, Anna Potapenko, et al. Highly accurate protein structure prediction with alphafold. Nature, 596(7873):583-589, 2021.
[92] Nathaniel Thomas, Tess Smidt, Steven Kearnes, Lusann Yang, Li Li, Kai Kohlhoff, and Patrick Riley. Tensor field networks: Rotation-and translation-equivariant neural networks for 3d point clouds. arXiv preprint arXiv:1802.08219, 2018.
[93] James Kirkpatrick, Brendan McMorrow, David HP Turban, Alexander L Gaunt, James S 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):1385-1389, 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. AI open, 1:57–81, 2020.
[95] Jonathan Masci, Davide Boscaini, Michael Bronstein, and Pierre Vandergheynst. Geodesic convolutional neural networks on riemannian manifolds. In Proceedings of the IEEE international conference on computer vision workshops, pages 37-45, 2015.
[96] Jun Li, Li Fuxin, and Sinisa Todorovic. Efficient riemannian optimization on the stiefel manifold via the Cayley transform. arXiv preprint arXiv:2002.01113, 2020.
[97] Bobak Kiani, Randall Balestriero, Yann Lecun, and Seth Lloyd. projunn: 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 preprint arXiv:1909.13334, 2019.
[99] Martin Arjovsky, Amar Shah, and Yoshua Bengio. Unitary evolution recurrent neural networks. In International conference on machine learning, pages 1120-1128. PMLR, 2016.
[100] Turgut Özis 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:1708.03888, 2017.
A PDE Symmetry Groups and Deriving Generators 18
A.1 Symmetry Groups and Infinitesimal Invariance 19
A.2 Deriving Generators of the Symmetry Group of a PDE 20
A.3 Example: Burgers' Equation 21
B Exponential map and its approximations 22
B.1 Approximations to the exponential map 23
C VICReg Loss 24
D Expanded related work 25
E Details on Augmentations 26
E.1 Burgers' equation 26
E.2 KdV 27
E.3 KS 28
E.4 Navier Stokes 28
F Experimental details 28
F.1 Experiments on Burgers' Equation 28
F.2 Experiments on KdV and KS 30
F.3 Experiments on Navier-Stokes 31
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 a priori and correspond to flips, resizing, or 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 [11]. 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 of
$$
\frac {\partial y}{\partial t} = \alpha x \quad \frac {\partial x}{\partial t} = - \alpha y
$$
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 to 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  $\alpha$  in the PDE), which is fixed for any solution to 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 a 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 [11].
In what follows, we give an informal overview to the derivation of Lie point symmetries. Full details and formal rigor can be obtained in Olver [11], Ibragimov [13], among others.
In the setting we consider, a differential equation has a set of  $p$  independent variables  $\boldsymbol{x} = (x^{1}, x^{2}, \ldots, x^{p}) \in \mathbb{R}^{p}$  and  $q$  dependent variables  $\boldsymbol{u} = (u^{1}, u^{2}, \ldots, u^{q}) \in \mathbb{R}^{q}$ . The solutions take the form  $\boldsymbol{u} = f(\boldsymbol{x})$ , where  $u^{\alpha} = f^{\alpha}(\boldsymbol{x})$  for  $\alpha \in \{1, \ldots, q\}$ . Solutions form a graph over a domain  $\Omega \subset \mathbb{R}^{p}$ :
$$
\Gamma_ {f} = \left\{\left(\boldsymbol {x}, f (\boldsymbol {x})\right): \boldsymbol {x} \in \Omega \right\} \subset \mathbb {R} ^ {p} \times \mathbb {R} ^ {q}. \tag {10}
$$
In other words, a given solution  $\Gamma_f$  forms a  $p$ -dimensional submanifold of the space  $\mathbb{R}^p \times \mathbb{R}^q$ .
The  $n$ -th prolongation of a given smooth function  $\Gamma_f$  expands or "prolongs" the graph of the solution into a larger space to include derivatives up to the  $n$ -th order. More precisely, if  $\mathcal{U} = \mathbb{R}^q$  is the solution space of a given function and  $f: \mathbb{R}^p \to \mathcal{U}$ , then we introduce the Cartesian product space of the prolongation:
$$
\mathcal {U} ^ {(n)} = \mathcal {U} \times \mathcal {U} _ {1} \times \mathcal {U} _ {2} \times \dots \times \mathcal {U} _ {n}, \tag {11}
$$
where  $\mathcal{U}_k = \mathbb{R}^{\dim(k)}$  and  $\dim(k) = \binom{p+k-1}{k}$  is the dimension of the so-called jet space consisting of all  $k$ -th order derivatives. Given any solution  $f: \mathbb{R}^p \to \mathcal{U}$ , the prolongation can be calculated by simply calculating the corresponding derivatives up to order  $n$  (e.g., via a Taylor expansion at each point). For a given function  $\boldsymbol{u} = f(\boldsymbol{x})$ , the  $n$ -th prolongation is denoted as  $\boldsymbol{u}^{(n)} = \operatorname{pr}^{(n)} f(\boldsymbol{x})$ . As a simple example, for the case of  $p=2$  with independent variables  $x$  and  $y$  and  $q=1$  with a single
dependent variable  $f$ , the second prolongation is
$$
\begin{array}{l} \pmb {u} ^ {(2)} = \mathrm {p r} ^ {(2)} f (x, y) = (u; u _ {x}, u _ {y}; u _ {x x}, u _ {x y}, u _ {y y}) \\ = \left(f; \frac {\partial f}{\partial x}, \frac {\partial f}{\partial y}; \frac {\partial^ {2} f}{\partial x ^ {2}}, \frac {\partial^ {2} f}{\partial x \partial y}, \frac {\partial^ {2} f}{\partial y ^ {2}}\right) \in \mathbb {R} ^ {1} \times \mathbb {R} ^ {2} \times \mathbb {R} ^ {3}, \tag {12} \\ \end{array}
$$