SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Hyperparameters: The hyperparameters for CimaX were selected based on the original implementation and the author's adaptation to ClimateBench. The chosen values include a patch size of 16, embedding dimension of 1024, encoder depth of 8 with 16 attention heads, and a non-linear decoder comprising two MLP layers with GELU activation, followed by a linear prediction head layer. For training, we use an Adam optimizer with weight decay with a learning rate of 1e-3 and a weight decay of 1e-5 paired with a linear warm-up cosine annealing learning rate scheduler, with 5 warm-up epochs and a maximum of 50 epochs to keep it consistent with the training of other models. The warm-up start learning rate is set to 1e-8, with a minimum eta of 1e-8.
K.2 U-Net
The U-Net architecture, originally proposed by Ronneberger et al. (2015), has established itself as a highly regarded convolutional neural network (CNN) model for image segmentation tasks. Its effectiveness in capturing intricate details while preserving the global context has made it a preferred choice for pixel-level classification and localization in various domains, including medical imaging and computer vision.
In a recent study conducted by Orlova et al. (2022), the potential of ML models for sub-seasonal forecasting of crucial climate variables was explored. The research encompassed an ensemble of models, incorporating linear models, regression models, random forests, and a U-Net-based convolutional network. The results demonstrated the superior performance of these ML methods over traditional non-ML baselines. Notably, the U-Net architecture exhibited exceptional capabilities in handling spatial variability within climate forecasts.
Architecture: The U-Net structure comprises an encoder path, responsible for capturing contextual information and extracting high-level features from the input image, and a decoder path, tasked with reconstructing the segmented image using the learned features. The decoder is followed by an output layer, typically employing additional convolutional operations to reduce the number of channels to match the desired number of output features.
Distinguishing Features from the Original Implementation: Following the implementation by Orlova et al. (2022), we adopt a U-Net architecture with a pre-trained backbone for the task of climate projection. Orlova et al. (2022) employed an off-the-shelf U-Net model available in the segmentation models PyTorch library. This pre-existing U-Net implementation utilized a pre-trained VGG11 (Simonyan and Zisserman, 2015) encoder backbone which was trained on the ImageNet classification challenge (Deng et al., 2009).
To adapt the U-Net model to our specific requirements, we encountered the challenge of dealing with input grids of varying resolutions. Unlike the original approach, we devised a solution by introducing zero-padding to the nearest image size divisible by 32 for both the longitude and latitude dimensions as this is required for the pre-trained encoder. Hence, we incorporated an adapted average pooling mechanism, initially introduced by Liu et al. (2018), to restore the output to its original grid size after the U-Net decoder process.
Furthermore, as our objective involved modelling sequence-to-sequence relationships rather than sequence-to-one, we employed a Time Distributed Layer encapsulating the pre-trained segmentation model. This approach, commonly used to handle sequential image data with convolutions, allows us to construct a sequential model with one layer per time-step in the input. Each layer consists of the same U-Net architecture, enabling the model to be applied to every temporal slice of the input. The weights are shared and trained simultaneously during the backward step.
Hyperparameters: Consistent with the original implementation (Orlova et al., 2022), the hyperparameters adopted include a VGG11 encoder, a linear readout layer, the Adam optimizer with a learning rate of 2e-4, and a weight decay of 1e-6. Additionally, an exponential learning rate scheduler with a gamma value of 0.98 is utilized for training.
K.3 Convolutional LSTM
The Convolutional LSTM (CNN-LSTM) is a model that combines the sequential application of convolutional neural networks (CNNs) (LeCun et al., 1989) and long short-term memory (LSTM) (Hochreiter and Schmidhuber, 1997) networks. It is important to note that the CNN-LSTM should not be confused with an LSTM that employs convolutional operations for its gates. The motivation behind
the CNN-LSTM is to effectively process image-like data over time by extracting spatial dependencies using the CNN component and temporal dependencies using the LSTM component. Thus, the model can be used for climate projection tasks and has been recognized as the best performing baseline in terms of root mean squared error (RMSE) scores in ClimateBench (Watson-Parris et al., 2022), albeit with some spatial biases in its predictions.
Architecture: The architecture of the CNN-LSTM involves several key components. Firstly, the input data undergoes feature extraction through a CNN layer. To ensure the extraction of features at each time step while sharing the same module, a Time-Distributed Layer is applied. Different weights are assigned to the filters in this layer. Following the feature extraction, an average pooling operation is performed to reduce the spatial dimensionality of the extracted features. This is also done time-step wise. Subsequently, the features are fed into a single LSTM layer, which captures the temporal dependencies in the data. Finally, a linear readout layer is applied to obtain the desired output.
Distinguishing Features from the Original Implementation: In comparison to the original implementation in ClimateBench, several modifications have been made to fit our training pipeline and our task setting. The ClimateBench implementation did not employ a sequence-to-sequence approach, only considering the last output of the LSTM and also did only model one output variable at a time. Contrary to that, in our modified version, all outputs of the LSTM are taken into account, allowing for modelling of the entire sequence. Thus, the readout layer has been altered to enable the modelling of the complete sequence and also to enable modelling multiple output variables simultaneously. Furthermore, the implementation has been converted from TensorFlow to PyTorch, which may result in slight changes due to the differences in internal workings of the layers between the two different libraries.
Hyperparameters: Regarding the chosen hyperparameters, we have adhered to the settings selected in ClimateBench (Watson-Parris et al., 2022). The CNN component consists of a single convolutional layer with 20 filters, each with a kernel size of 3, and employs ReLU activation function. The pooling stage involves two 2D average pooling layers, one with a kernel size of (2,2) and the other with a kernel size of (longitude/2, latitude/2). The LSTM layer comprises a single layer with 25 units and uses ReLU as the activation function. The model is trained using the Adam optimizer with a learning rate of 2e-4, a weight decay of 1e-6, and an epsilon value of 1e-8. Additionally, an Exponential Decay Learning Rate Scheduler with a gamma of 0.98 is applied during training.
K.4 Gaussian Process
Gaussian process regression (Williams and Rasmussen, 2006) is a nonparametric and Bayesian regression method that has the advantage of providing uncertainty measures over predictions. Since we deal with large datasets, we use a stochastic variational variant of the Gaussian process for regression (SVGP; (Hensman et al., 2015)) that supports training with minibatches and scales well with the number of samples. To deal with the multi-output nature of the task, we use the Linear Model of Coregionalization (LMC) with 100 latent Gaussian processes. The number of latents used is a hyperparameter that controls the capacity of the model. We use Matern1.5 kernels and train with the Adam optimizer with a learning rate of 0.1 and batch size of 64.
As in ClimateBench (Watson-Parris et al., 2022), we use Empirical orthogonal functions (EOF) as a dimensionality reduction technique on the aerosols input (BC and SO2) and keep only the 5 first modes. While we use SVGP, the ClimateBench uses exact Gaussian Processes which can be problematic for large datasets since it scales cubically with respect to the number of examples.
Finally, to implement our SVGP method, we used the library GPytorch (Gardner et al., 2018) (https://gpytorch.ai/), a Pytorch implementation of Gaussian processes.
We showcase our results with the Gaussian Process on six climate models in Table 4. Gaussian Process is the best-performing model when predicting TAS (surface air temperature) for NorESM2-LM. Generally, we observe that the Gaussian Process performs really well for predicting TAS but when predicting PR (precipitation), it underperforms the models shown in Table 5 by a huge margin. For this reason, we chose to run the Gaussian Process on a smaller subset of the climate models and omitted it for the experiments run on 15 climate models.
GP
TAS	PR
AWI-CM-1-1-MR	0.264	0.538
BCC-CSM2-MR	0.247	0.536
EC-Earth3	0.279	0.580
FGOALS-f3-L	0.280	0.553
MPI-ESM1-2-HR	0.257	0.544
NorESM2-LM	0.248	0.553
Table 4: RMSE of Gaussian Process single-emulating temperature (TAS) and precipitation (PR) on six climate models. Emboldened values are showing when the GP performs best across all ML models (GP, U-Net, ConvLstm, CimaX, and CimaX $_{\text{Frozen}}$ ). Italic values show when the GP performs worst across all ML models.
L Experimental Setup
L.1 Single-Emulator
For all our models, we use a similar training procedure where we predict the outputs for entire input sequence in one go. Our inputs are of shape  $<$  batch, sequence_length, num_vars, lon, lat > and outputs are of size  $<$  batch, sequence_length, 2, lon, lat > where the 2 in output dimension denotes TAS and PR. For our experiments we have a sequence length of 12, meaning that we use data from 12 months and predict the output for those 12 months. More details about the input and output data, along with the variables used are in Section 2.1. Our training methodology differs slightly from those used by Nguyen et al. (2023) and Watson-Parris et al. (2022) in the sense that we don't use a lead time for out models.
We train all our models on one Nvidia-RTX8000 GPU with a batch size of 4 and make use of 32GB of RAM. The U-Net and ConvLSTM are trained for 50 epochs with an initial learning rate of  $2e - 4$  with an exponential decay scheduler. To keep the training consistent for CimaX models, we train the models for 50 epochs with an initial warm-up for 5 epochs. The learning rate for the warm-up is set to 1e-8 and then to 5e-4 for training. Other training details are kept the same as in the original implementation. Fig. 6 reports results from experiments without warm-up, while all tables in the Appendix report only results from experiments with warm-up. For the loss, we use the latitude-longitude weighted mean squared error (LLMSE) as implemented in Nguyen et al. (2023). We report the latitude-longitude weighted root mean squared error (RMSE) as our metric to evaluate the models.
L.2 Super-Emulator
We use the term super-emulator to refer to a single ML model trained on the entire data of all climate models, such that it is capable of projecting a unique output of every specific climate model, according to a user requirement. Training a single ML model (as opposed to many ML models each trained on one climate model data) may benefit from rich and diverse data from the collection of climate models (as opposed to only one). This joint knowledge results in a super-emulator model with parameter sharing and increased capacity. For this goal, the super-emulator should be able to distinguish between data inputs provided to it from different climate models. A naive super-emulator that is trained on multiple climate models as targets without any contextual information would be mapping the same inputs to multiple targets. This lack of context restricts the emulator's ability to accurately represent the different behaviours of climate models. In order to prevent the "one-to-many" mapping, a super-emulator needs to encode the climate models in the input data. One approach is to provide which climate model is associated with each sample, and another approach is to provide the climate model's internal dynamics as additional input. Both methods are outlined in the following:
Multi-head decoder: This method is centred around the idea of indexing each sample with its respective climate model. By associating each sample with a climate model, the mapping becomes a
"one-to-one" mapping. Here, we implement this by using a multi-head decoder on top of a neural network model. We are using the exact same NN models as in the single-emulator experiments as basemodels. They use the same inputs, outputs, and shapes as in the single-emulator experiments. On top of such a basemodel (U-Net, CimaX, ConvLSTM), a multi-head decoder is installed. Each climate model that participates in the super-emulator receives its own head. A head consists of several convolutional layers which can be determined by the user. Each head receives the gridded output of the basemodel and the climate model index for each sample. The training batches can contain samples from different climate models since each sample can be identified by its climate model index. Each sample is forwarded only through its matching head, and back-propagated through this head, while all other heads stay frozen. In our experiments, we used the ConvLSTM, U-Net, CimaX, CimaX $_{\text{Frozen}}$  as base models with the same training regimes as outlined in the single-emulator experiments. Due to memory requirements, we were only able to train the super-emulator on the smaller set of six climate models. The overall super-emulator was trained for a maximum of 100 epochs with a head for each climate model and each head had 2 convolutional layers with 32 units. We noticed that the performance of CimaX-based models plummeted and other models stagnated after  $\sim 20$  epochs with predictions converging to zero. This lead us to include a standard Early Stopping Regularization paradigm that stopped the training process around that threshold. We used the inbuilt pytorch Early Stopping class with patience set to 3 and a min delta of 0. Refer to Appendix M.2 for the results of those experiments. Further experiments are necessary to be able to analyze the potential of the multi-head decoder for super-emulation fully.
Sea-level pressure maps: The current approach of training a super-emulator solely on greenhouse gas, aerosols and aerosol precursor emission data from Input4MIPs does not sufficiently capture the intrinsic dynamics of the climate models. Climate models are not solely influenced by external factors but also by intrinsic dynamics, which can exhibit chaotic behaviour. Consequently, even numerical models generate different trajectories when provided with the same input data but different initializations. This internal variability is captured by the multiple ensemble members that are also included in ClimateSet. However, the current machine learning emulators fail to account for and explain this internal variability, resulting in fluctuations in the predicted targets. To address these challenges, one option is to incorporate one of the intrinsic dynamics present in climate models as an additional input alongside the forcing data. For instance, concatenating sea level pressure maps specific to the climate model and/or ensemble member with the input GHG and aerosol data would enable the required one-to-one mapping and partially explain the variability observed in the targets. This approach has the potential to enhance the emulator's overall performance. It would also enable querying for average predictions or interpolation between existing scenarios during inference time. By interpolating or averaging sea level pressure maps from existing runs, it becomes possible to generate new input features that incorporate both emissions data and intrinsic dynamics, even though the intrinsic dynamics have not been simulated by a climate model yet. Consequently, the machine learning emulator - despite using the output of a climate model - can still predict scenarios that the climate model itself has not run.
To summarize, the possibilities to design a super-emulator are plentiful and are subject to the intended purpose determined by the user. As such, they remain to be investigated thoroughly.
L.3 Generalization Capabilities of ML Models
In order to investigate the generalization capabilities of ML models on different climate models, we decided to choose a finetuning approach.
L.3.1 Finetuning Single-Emulators
For the single-emulator finetuning experiments, each ML model is first pre-trained on a specific climate model, and subsequently finetuned and evaluated on NorESM2-LM (Seland et al., 2020). The best model to be chosen for subsequent tuning is selected based on the validation metric. We chose NorESM2-LM as a generalization test case since it is the model that has been most widely used in previous work. We compare the performance of the finetuned ML model with the NorESM2-LM single-emulator to learn A) how well the ML model can generalize from one model to another, and B) how much transferable information each climate model has for NorESM2-LM. Similar to (Nguyen
<sup>7</sup>We plan to address this limitation in future work and will update our GitHub repo continuously.
et al., 2023), we hope that pre-training on climate datasets might help the ML model learn some patterns that are applicable to other datasets & climate models.
The single-emulator finetuning experiments follow the training procedure described in Section L.1. Those setups are used for pre-training. The pre-trained models are then loaded and afterwards finetuned on NorESM2-LM for 50 epochs. We report the latitude-longitude weighted root mean squared error (RMSE) as our evaluation metric.
L.3.2 Finetuning Super-Emulators
The multi-head decoder used for super-emulation can also be used to examine the generalization capability of a super-emulator. For pre-training a multi-head decoder super-emulator is used with heads for all participating climate models. The only climate model that has no header is the one on which the super-emulator is evaluated. During finetuning another header is added and only that head is finetuned. Subsequently, the performance of the super-emulator (e.g. trained on a set of climate models but without NorESM2-LM) can be compared with the performance of the finetuned super-emulator (e.g. finetuned on an added head for NorESM2-LM). We leave those experiments and analyses for future work.
M Experimental Results
M.1 Single-Emulator
The results of our experiments for training ML models on emulating single climate models are shown in Table 5. In these experiments, we train our models individually on 15 different climate models and report the RMSE for our two output variables, TAS (surface air temperature) and PR (precipitation).
Our experiments indicate a trend of larger ML-models outperforming the smaller ones on single-emulation. For most climate datasets, ClimaX performs much better than other models. This is not unexpected as ClimaX is a much larger model being finetuned on suitable pre-trained weights (stemming from ERA5 and two climate models) for this task.  $\mathrm{ClimaX}_{F}$ rozen follows usually with slightly worse RMSE values. Thus, finetuning the weights is indeed helpful to increase the performance of ClimaX. With some exceptions, U-Net is performing best after  $\mathrm{ClimaX}_{F}$ rozen, and ConvLSTM, the simplest and smallest model, shows often the worst performance across all ML-models. We conclude, that there is a general trend among the ML-models performance across several climate models.
However, there are exceptions to this general "leaderboard". We can identify three strong outliers, namely GFDL-ESM4 (Dunne et al., 2020), NorESM2-LM (Seland et al., 2020) and NorESM2-MM (Seland et al., 2020), where U-Net outperforms both CimaX models by a pretty large margin. Also on TaiESM1 (Wang et al., 2021), U-Net is performing exceptionally well and comes in second after CimaX. These results show, that it is not sufficient to evaluate an ML-emulator on a single climate model. When trained only on e.g. NorESM2 models, one might conclude that U-Net is the best performing model. Only when taking a set of climate model datasets intro consideration, we can evaluate how well a specific ML-model can emulate climate models.
Also note, that the performance differences between the ML models is not huge. For example, U-Net performs in a very similar range like  $\mathrm{ClimaX}_{f}$  rozen, outperforming it from time to time. This indicates that simple baselines can still keep up with larger models. We hope that other simple baselines will be investigated by other researchers in the future. Also further work on large models is needed to achieve performance gains that justify using such models over simple baselines.
M.2 Super-Emulator
Our super-emulation experiments showed that the ranking of the best performing ML models changes between single and super-emulation. While U-Net and ClimaX have been the best performing models when the goal is to emulate only one climate model (see Table 5), ConvLSTM is performing best when a set of climate models needs to be emulated (see Table 6). The ConvLSTM achieves the highest performance among all ML models for all six climate models on both TAS and PR. However, all super-emulator performance values are significantly worse than the single-emulator results. This
U-Net	ConvLSTM	ClimaX	ClimaXFrozen
TAS	PR	TAS	PR	TAS	PR	TAS	PR
AWI-CM-1-1-MR	0.321	0.326	0.445	0.438	0.263	0.267	0.316	0.321
BCC-CSM2-MR	0.232	0.258	0.295	0.333	0.205	0.229	0.232	0.258
CAS-ESM2-0	0.317	0.315	0.393	0.394	0.226	0.226	0.272	0.276
CNRM-CM6-1-HR	0.315	0.309	0.379	0.370	0.250	0.245	0.289	0.286
EC-Earth3	0.292	0.295	0.397	0.398	0.240	0.241	0.287	0.292
EC-Earth3-Veg-LR	0.299	0.300	0.410	0.410	0.248	0.250	0.290	0.290
FGOALS-f3-L	0.286	0.305	0.362	0.381	0.260	0.274	0.294	0.313
GFDL-ESM4	0.481	0.482	0.505	0.503	0.542	0.544	0.630	0.636
INM-CM4-8	0.232	0.235	0.335	0.346	0.212	0.214	0.257	0.261
INM-CM5-0	0.258	0.264	0.329	0.345	0.207	0.209	0.247	0.252
MPI-ESM1-2-HR	0.31	0.317	0.406	0.415	0.225	0.229	0.258	0.265
MRI-ESM2-0	0.292	0.292	0.357	0.360	0.221	0.219	0.262	0.262
NorESM2-LM	0.271	0.291	0.299	0.315	0.369	0.381	0.386	0.413
NorESM2-MM	0.291	0.313	0.317	0.333	0.376	0.395	0.420	0.435
TaiESM1	0.196	0.197	0.333	0.346	0.194	0.195	0.232	0.237
Table 5: RMSE values for each ML model single-emulating surface air temperature (TAS) and precipitation (PR) across 15 climate models. All reported values are averages over three seeds. The highest RMSE among the ML-models is emboldened for each climate model dataset. For most of the climate model datasets, CimaX seems to outperform other ML models. U-Net is the best performing model on some of the climate model datasets.
U-Net	ConvLSTM	ClimaX	ClimaXFrozen
TAS	PR	TAS	PR	TAS	PR	TAS	PR
AWI-CM-1-1-MR	0.579	0.571	0.396	0.398	0.908	1.172	0.647	0.634
BCC-CSM2-MR	0.334	0.299	0.270	0.238	0.495	0.88	0.366	0.338
EC-Earth3	0.508	0.481	0.287	0.286	0.616	1.287	0.491	0.480
FGOALS-f3-L	0.811	0.745	0.314	0.287	0.665	1.622	0.464	0.439
MPI-ESM1-2-HR	0.8	0.756	0.295	0.289	0.903	0.979	0.531	0.512
NorESM2-LM	0.566	0.532	0.261	0.249	0.686	1.457	0.455	0.426
Table 6: RMSE values for each ML model super-emulating surface air temperature (TAS) and precipitation (PR) on a set of six climate models. All reported values are averages over three seeds. The highest RMSE among the ML-models is emboldened for each climate model dataset. The performance of super-emulator models seems to be worse as compared to the single-emulator models as seen in Table 5. The smallest model, ConvLSTM, is performing best across all climate models. Smaller ML models seem to be the better super-emulators in our experiments.
can be attributed to the fact that the task of learning to emulate several climate models is harder than learning to emulate only one climate model.