SlowGuess/NeurIPS_2023_OCR huggingface dataset:
We use the Conditioning dataset for Navier Stokes-2D proposed in [18], 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, a batch size of 64 and a learning rate of 3e-4. We use the same VICReg hyperparameters as for Burgers' Equation. We use the following augmentations and strengths (augmentations whose strength is not specified here are not used):
- Crop of size (16, 128, 128), respectively in temporal, x and y dimensions.
- Uniform sampling in  $[-1, 1]$  for the coefficients associated to  $g_{2}$  and  $g_{3}$  (applied respectively in x and y).
Uniform sampling in  $[-0.1, 0.1]$  for the coefficients associated to  $g_{5}$ .
Uniform sampling in  $[-0.01, 0.01]$  for the coefficients associated to  $g_{6}$  and  $g_{7}$  (applied respectively in x and y).
- Uniform sampling in  $[-0.01, 0.01]$  for the coefficients associated to  $g_{8}$  and  $g_{9}$  (applied respectively in x and y).
Buoyancy 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  $\left[\mathrm{Buoyancy}_{\min}, \mathrm{Buoyancy}_{\max}\right]$ . 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.
Time-stepping We mainly depart from [18] by using 20 epochs to learn from 1,664 trajectories as we observe the results to be similar, and allowing to explore more combinations of architectures and conditioning methods.
Time-stepping results In addition to results on 1,664 trajectories, we also perform experiments with bigger train dataset (6,656) as 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 11. As in [18], AdaGN outperforms Addition. Note that AdaGN is needed for our representation conditioning to significantly improve over no conditioning. Finally, we found a very small bottleneck in the MLP that process the representation to also be crucial for performance, with a size of 1 giving the best results.
Table 10: List of model hyperparameters and training details for the timestepping tasks.
Equation	Burgers'	KdV	KS	Navier Stokes
Neural Operator:
Model	CNN [12]	CNN [12]	CNN [12]	Modified U-Net-64 [18]
Optimization:
Optimizer	AdamW	AdamW	AdamW	Adam
Learning Rate	1e-4	1e-4	1e-4	2e-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	~ 1d	~ 2d	~ 2d	~ 1.5d
Table 11: One-step validation MSE  $\times 1e^{-3}$  (↓) for Navier-Stokes for different baselines and conditioning methods, with UNet<sup>mod64</sup> [18] as base model.
Dataset size	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 ± 0.02	1.17 ± 0.04
Time conditioned, AdaGN	2.37 ± 0.01	1.12 ± 0.02
Time + Rep. conditioned, AdaGN (ours)	2.35 ± 0.03	1.11 ± 0.01
Methods with ground truth buoyancy:
Time + Buoyancy conditioned, Addition	2.08 ± 0.02	1.10 ± 0.01
Time + Buoyancy conditioned, AdaGN	2.01 ± 0.02	1.06 ± 0.04