Scaling transformer neural networks for skillful and reliable medium-range weather forecasting

We thank the reviewer for the very detailed and constructive feedback, and for recognizing the technical contributions and good presentation of Stormer. We answer each of the reviewer's concerns below.

PanguWeather could then also be used to implement the randomized forecast strategy. It would be interesting to see this included in the model comparisons.

We agree that Pangu-Weather is capable of performing the randomized forecast strategy with the cost of training 3 separate models. We plan to include this experiment in the updated paper.

As a fairer comparison to Pangu, we have compared the non-ensemble version of Stormer with the baselines. Specifically, we performed the Pangu-style inference, where we only used the 24-hour interval forecasts to roll out into the future (i.e., 1-day=24, 2-day=24+24, 3-day=24+24+24, etc.), instead of combining different intervals.

Figure 1 in our PDF shows that non-ensemble Stormer outperforms Pangu and performs competitively with Graphcast. Moreover, we note that the ensembling technique in Stormer is much cheaper and easier to use than other methods such as training multiple networks, dropout, or IC perturbations, as we only have to train a single neural network and do not need extensive hyperparameter tuning. For better efficiency, one can always use the Homogeneous version of Stormer for inference, which only requires 3 forward passes and performs competitively to the Best m in n version, as shown in Figure 3 of our PDF.

How do you expect that the resolution difference between the Stormer and GraphCast/PanguWeather models affects the comparisons? Are you computing the losses at the native resolution of Stormer, downsampling the other models to it? Do you expect that it would change the results if you instead perform the comparison at the native resolution of GraphCast/PanguWeather and upsample the Stormer results to it?

We downsampled the forecasts of Graphcast and Pangu-Weather to 1.40625deg and compared all methods in this resolution. This is the same strategy that was used in WeatherBench 2 to compare different models at different resolutions. We expect the comparisons to change if we instead upsample the Stormer forecasts to 0.25deg. While this is technically possible, we (and Weatherbench 2) opted not to do so because upsampling introduces additional error due to the complexity of extrapolating information from lower to higher resolution. In contrast, downsampling is straightforward, effectively averaging neighboring pixels without significant loss of information.

What is the computational complexity of training and inference

During training, complexity is the same as training a single-interval model because we train a single model for the same amount of time. In other words, there is no computation overhead of randomized iterative forecasting during training. During inference, since we average multiple forecasts with different interval combinations, complexity scales linearly with the number of combinations. If computation is a critical issue, one should use the homogeneous inference of Stormer, which only uses 3 homogeneous combinations, while achieving competitive results to the Best m in n inference, as Figure 3 in our PDF shows.

How do you counter for boundary conditions? Do the variable and patch embeddings respect boundary conditions?

The model learns everything from data, and we do not enforce any special constraints on the model. This is similar to most deep learning methods like Pangu, Graphcast, etc.

Can the model accommodate different lead-time resolution query points (t = 1hr, 19hr, 91hr, etc.) that differ from the dataset

If we train the model on the 1-hour interval then it can produce forecasts at any lead time that is a multiplier of 1 hour. However, due to the huge size of 1-hourly ERA5, we subsampled the data to 6-hourly only, so the model can make forecasts at lead times that are multipliers of 6 hours. We do not expect a model to generalize well to an interval unseen during training.

Does the method use standard 2D ViT or a 3D transformer as the stormer block?

Stormer uses a standard transformer 2D backbone. The model relies on the weather-specific embedding module to aggregate information across different variables and pressure heights.

We thank the reviewer for the constructive feedback and the appreciation of the good presentation and good performance of Stormer. We answer each of the reviewer's concerns below.

The paper has very limited novelty and is incremental.

We acknowledge that some components in Stormer are similar to those of prior works. We also discussed the relation to them in the paper: Lines 204-205 discussed the variable embedding in ClimaX, Line 298 mentioned the pressure-weighted loss in Graphcast, and Lines 100-101 mentioned multi-step finetuning in previous works. We will update the paper to bring this discussion to the methodology section for better clarity.

However, we would like to emphasize the differences and contributions of Stormer:

• Stormer has a similar architecture to climax, but we use adaptive layer norm for time-conditioning, which we show is important to the performance in Figure 3c. Even with a fairly similar architecture, Stormer significantly outperforms ClimaX, showing the superiority of randomized iterative forecasting vs continuous pretraining + direct finetuning.
• We introduce randomized iterative forecasting, a paradigm not explored in previous works. This allows any architecture (GNN, transformers, etc.) to gain performance improvements w.r.t deterministic metrics with only minor computation overhead (during inference only) compared to single-interval models.
• Unlike previous works, we carefully ablate each component in Stormer to understand their importance in obtaining a good forecast performance. Via this paper, we show that it may not be necessary for a specialized neural network architecture, and a standard model like transformers can achieve state-of-the-art performance with careful design choices. We believe this extends the current understanding of data-driven weather forecasting, and is a valuable contribution to the community.

The paper misses out on a ton of related works

Thank you for pointing us to these works, we will discuss them in detail in the updated paper. Both Gencast and NeuralGCM were concurrent works and used much larger compute resources to run on higher-resolution data. Moreover, our goal in this paper is not to show SOTA (even though we could potentially achieve better performance with more compute and higher resolution data), but rather to show the power of simple and scalable approaches based on the transformers architecture, and what contributes to the performance of an ML weather forecasting model.

Figure 4 in our PDF compares Stormer with NeuralGCM and ClimODE. We additionally include Stormer (IC noise), a version of Stormer with initial condition perturbations. Specifically, for each combination in the “Best 32 in 128” strategy, we sample 4 different noises from a Gaussian distribution with a standard deviation of 0.1 and add to the input, resulting in a total number of 128 ensemble members. The figure shows that Stormer outperforms deterministic NeuralGCM and performs competitively with NeuralGCM ENS. With IC perturbations, the gap between Stormer and NeuralGCM is negligible. ClimODE performs significantly worse than the other methods. While ClimeODE may improve by training on higher-resolution data, we do not believe it can close this huge gap. We will add this comparison to the updated paper.

Recent works have shown that continuous time methods with physical inductive biases and hybrid ML modeling have surpassed transformer-based methods such as ClimODE and Neural GCM, providing uncertainty and interpretability. The proposed method does not offer any benefits to any of them.

We respectfully disagree. NeuralGCM, ClimODE, and Stormer are different approaches and each has its pros and cons:

• NeuralGCM is a hybrid model that combines a differentiable dynamical core with ML components for end-to-end training. While the dynamical core allows the method to leverage powerful general circulation models, it also has various drawbacks. First, to make predictions, the dynamical core in NeuralGCM has to solve discretized dynamical equations, which are more computationally expensive than forward-passing a neural network. Second, the performance of NeuralGCM is upper-bounded by the accuracy of the fixed dynamical core, while fully learnable models like Stormer continue improving with more data. This is a desirable property, given the scaling properties we have seen in Stormer and ClimaX.
• ClimODE introduces physical inductive biases into deep learning models to improve the interpretability of their method. However, in terms of forecasting accuracy, as Figure 4 in our PDF shows, ClimODE is quite inferior to Stormer or other state-of-the-art methods.

They still restrict the whole method to output only point estimates rather than giving uncertainties.

This work focuses on deterministic forecasting and how to achieve state-of-the-art performance in deterministic metrics with a simple but scalable architecture like a transformer. However, we note that we can make Stormer a probabilistic model with IC perturbations. Figure shows the probabilistic performance of Stormer with different noise levels (standard deviations of the Gaussian noise distribution). IC perturbations significantly improve CRPS and SSR metrics of Stormer as well as deterministic performance at long lead times, but may hurt the accuracy at short lead times. Moreover, it is difficult to find an optimal noise level for the spread-skill ratio across different variables and lead times. We can further improve this by using a better noise distribution or variable-dependent and lead-time-dependent noise scheduling, which we defer to future works.

We thank the reviewer for the very detailed and constructive feedback, and for recognizing the strong results and technical contributions of Stormer. We answer each of the reviewer's concerns below.

The paper is essentially using an ensembling technique but comparing against deterministic models.

We would like to clarify that even though Stormer uses an ensembling technique, we consider it a deterministic model, and that’s why we compare it with deterministic baselines. The reason is that while Stormer can produce multiple forecasts for a given lead time at inference, we found these forecasts not diverse enough and should not be used for uncertainty estimation. This is shown by the under-dispersion of Stormer w.r.t to the spread-skill ratio in Figure 2 of our PDF. We will add this discussion to the updated paper.

Thus, a fairer comparison would be to either 1) ensemble the deterministic ML-based baselines (e.g. through input perturbations or with lagged ensembles as in [1]) or 2) show results of Stormer without ensembling (even a "m=1 in n" forecast would be much fairer than the way it is now).

We agree these comparisons would provide more insights into the performance of Stormer. To do this, we have compared the non-ensemble version of Stormer with the baselines. Specifically, we performed the Pangu-style inference, where we only used the 24-hour interval forecasts to roll out into the future (i.e., 1-day=24, 2-day=24+24, 3-day=24+24+24, etc.), instead of combining different intervals.

Figure 1 in our PDF shows that non-ensemble Stormer outperforms Pangu and performs competitively with Graphcast. Moreover, we note that the ensembling technique in Stormer is much cheaper and easier to use than other methods such as training multiple networks, dropout, or IC perturbations, as we only have to train a single neural network and do not need extensive hyperparameter tuning. For better efficiency, one can always use the Homogeneous version of Stormer for inference, which only requires 3 forward passes and performs competitively to the Best m in n version, as shown in Figure 3 of our PDF. We will add this result and discussion to the updated paper.

Additionally, the proper way to evaluate an ensemble weather forecast is via probabilistic metrics such as the CRPS and spread-skill ratios. This should be included to properly assess the quality of the homogenous or "best m in n" ensembles.

To make Stormer a probabilistic forecast system, we need to introduce more randomization to the forecasts via IC perturbations. To do this, for each combination of intervals during the Best m in n inference, we added 4 different noises sampled from a Gaussian distribution, resulting in a total number of 128 ensemble members. Figure 2 in our PDF shows the RMSE and probabilistic metrics of Stormer with different standard deviations of the noise distribution.

The result shows that IC perturbations improve the probabilistic metric significantly, but may hurt the deterministic performance at short lead times. Moreover, it is difficult to find an optimal noise level for the spread-skill ratio across different variables and lead times. We can further improve this by using a better noise distribution or variable-dependent and lead-time-dependent noise scheduling, which we defer to future works. We will add these results and discussions to the updated paper.

On top of probabilistic metrics, it would be instructive to see an analysis of the generated spectra of Stormer.

We thank the reviewer for the suggestion and we will include this in the updated version of the paper. Due to the time constraint, we will defer this experiment to later. If the reviewer thinks this is crucial to assessing the paper, we will conduct this experiment during the discussion phase.

The authors seem to have missed section "7.3.1. Multi-mesh ablation" in the GraphCast paper.

We thank the reviewer for pointing this out. Our main message here is questioning whether we need a specialized neural network architecture for weather forecasting, or whether a standard architecture like transformers can work equally well. In this paper, we consider the Transformer architecture as a special reference due to its low inductive bias, good scaling properties, and great performance across different data modalities (text, image, audio, etc.). We will reword this part for better clarity in the updated version.

I would like to see a more transparent discussion of the exact contributions of this work and a more careful contextualization with prior work.

We initially wanted to describe our method first before connecting it with the related work, but we agree with the reviewer that it'd be better and more transparent to mention the prior works as we discuss each component of Stormer. We will update the paper accordingly. We answer the reviewer’s specific concerns as follows: 1) This is correct, but this seemingly small difference can make a huge difference in performance, as shown by the big gap between Stormer and ClimaX, as shown in Figure 4 of our PDF, 2) They are similar, but there are some nuanced differences. As the reviewer may have noticed, we train a single model for all intervals, while Pangu trains a separate model for each lead time. During inference, Pangu uses a single combination of intervals for each lead time that minimizes the rollout step, while we use a combination of them, 3) and 4) We agree with the reviewer and will add this discussion to the appropriate part of the main text.

Relatedly, I think that the paper could benefit from discussing the related work more in-depth.

In terms of architecture, Stormer and ClimaX both use a standard transformer backbone, and the only difference is Stormer uses adaptive layer norm for time-conditioning while ClimaX uses a simple additive embedding. On the other hand, Pangu-Weather and FengWu use a Swin-transformer backbone.

In terms of model training, ClimaX is pretrained with continuous forecasting but finetuned for direct forecasting. Pangu, Fuxi, and Fengwu are iterative models but with slightly different designs. Pangu trains a separate model for each lead time, Fengwu performs multi-step finetuning with a replay buffer, and Fuxi finetunes a separate model for each time range (short, medium, and long). We will include this discussion in the updated version.

Please include results on other variables than just T2M for your ablations.

Table 1 in our PDF shows the performance of the weighted and unweighted loss on 6 different variables, ranging from low to high-pressure levels. As expected, the weighted loss model achieves better accuracy for high-pressure variables while underperforming the unweighted version for low-pressure variables. We were aware of this trade-off and proposed to use the weighted loss in Stormer to focus on variables that are more important to forecasting and/or human activities. We will add these results to the updated paper for better clarity.

To me it seems that at inference time there does exist a computational overhead due to the larger size of the ensemble

We wrote this with regard to the training process, which is the major computation overhead of deep learning models. At inference, Stormer does require more computation compared to single-interval models, and computation scales with how many combinations of intervals we use during inference. And as we showed above, using 3 ensemble members of the homogeneous inference is enough to achieve a good performance. We will reword this part to avoid confusion in the updated manuscript.

It would be very insightful to see an ablation of the different prediction approaches in weather forecasting. That is, continous, iterative, and randomized iterative forecasting.

We expand the differences between these approaches below:

• Continuous vs randomized iterative forecasting: The difference can be seen in the performance gap between Climax and Stormer. Continuous requires a single model to be able to forecast with a wide range of lead times, e.g., 6 hours to 2 weeks, which is a challenging (and sometimes confusing) learning task for any model. In contrast, randomized iterative only trains the model with a small set of intervals (6, 12, 24), mitigating this problem. Moreover, continuous models do not generalize beyond lead times in the training range, as shown in Climax, while randomized iterative does.
• Iterative vs randomized iterative: The latter offers two advantages, data augmentation and the ability to combine the intervals to produce multiple forecasts. Empirically, the randomized iterative approach achieves significantly better accuracy than the iterative approach, as shown in the performance gap between Stormer and non-ensemble Stormer, while only incurring a slight computation overhead.

Can you expand on how the "best m" combinations are chosen?

By combinations, we mean from the set {6,12,24} pick any ordered combination that sums up to a certain lead time. For example, for a lead time of 1 day (24 hours), we can either do 6+6+6+6, 6+12+6, 24, 12+12, 12+6+6, 24, etc. We chose the best m based on the validation loss, and we picked n to be a reasonably large value (128) without hyperparameter tuning. Based on our preliminary results, the final performance is not sensitive to this value (we tried 32, 64, 128).

Indeed, there won't be many combinations for very small lead times, but these lead times also do not require ensembling multiple forecasts because individual forecasts are accurate enough. We will discuss this component in more detail in the updated paper.

Why does the time embedding module predict the second scale parameter? Why no second shift? Did you ablate this choice? What if you remove this second time embedding part (or the first)?

We did not ablate this choice. We adopted the adaptive layer normalization (adaLN) from the computer vision literature [1, 2], which is a common technique to condition a neural network on additional information like time.

[1] Perez, Ethan, et al. "Film: Visual reasoning with a general conditioning layer." Proceedings of the AAAI conference on artificial intelligence. Vol. 32. No. 1. 2018.

[2] Peebles, William, and Saining Xie. "Scalable diffusion models with transformers." Proceedings of the IEEE/CVF International Conference on Computer Vision. 2023.

Thank you for the quick response. I think that the revised paper will be a great addition to the literature. I look forward to reading it.