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

Authors: tung  nguyen
rohan  shah
hritik  bansal
troy  arcomano
romit  maulik
veerabhadra  kotamarthi
ian  foster
sandeep  madireddy
aditya  grover
  cmu
argonne national laboratory
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 05:12:58

Abstract: Weather forecasting is a fundamental problem for anticipating and mitigating the impacts of climate change. Recently, data-driven approaches for weather forecasting based on deep learning have shown great promise, achieving accuracies that are competitive with operational systems. However, those methods often employ complex, customized architectures without sufficient ablation analysis, making it difficult to understand what truly contributes to their success. Here we introduce Stormer, a simple transformer model that achieves state-of-the art performance on weather forecasting with minimal changes to the standard transformer backbone. We identify the key components of Stormer through careful empirical analyses, including weather-specific embedding, randomized dynamics forecast, and pressure-weighted loss. At the core of Stormer is a randomized forecasting objective that trains the model to forecast the weather dynamics over varying time intervals. During inference, this allows us to produce multiple forecasts for a target lead time and combine them to obtain better forecast accuracy. On WeatherBench 2, Stormer performs competitively at short to medium-range forecasts and outperforms current methods beyond 7 days, while requiring orders-of-magnitude less training data and compute. Additionally, we demonstrate Stormer’s favorable scaling properties, showing consistent improvements in forecast accuracy with increases in model size and training tokens. Code and checkpoints are available at https://github.com/tung-nd/stormer.

Main paper: 
 introduction: 
 Weather forecasting is a fundamental problem for science and society. With increasing concerns about climate change, accurate weather forecasting helps prepare and recover from the effects of natural disasters and extreme weather events, while serving as an important tool for researchers to better understand the atmosphere. Traditionally, atmospheric scientists have relied on numerical weather prediction (NWP) models [2] . These models utilize systems of differential equations describing fluid flow and thermodynamics, which can be integrated over time to obtain future forecasts [29, 2] . Despite their widespread use, NWP models suffer from many challenges, such as parameterization errors of important small-scale physical processes, including cloud physics and radiation [44] . Numerical methods also incur high computation costs due to the complexity of integrating a large system of differential equations, especially when modeling at fine spatial and temporal resolutions. Furthermore, NWP forecast accuracy does not improve with more data, as the models rely on the expertise of climate scientists to refine equations, parameterizations, and algorithms [30] .To address the above challenges of NWP models, there has been an increasing interest in data-driven approaches based on deep learning for weather forecasting [11, 42, 48] . The key idea involves training deep neural networks to predict future weather conditions using historical data, such as the ERA5 reanalysis dataset [16, 17, 40, 41] . Once trained, these models can produce forecasts in a few seconds, as opposed to the hours required by typical NWP models. Because of the similar spatial structure between weather data and natural images, early works in this space attempted to adopt standard vision architectures such as ResNet [39, 8] and UNet [49] for weather forecasting, but their performances lagged behind those of numerical models. However, significant improvements have been made in recent years due to better model architectures and training recipes, and increasing data and compute [22, 35, 32, 3, 26, 5, 7] . Pangu-Weather [3] , a 3D Earth-Specific Transformer model trained on 0.25 • data (721×1440 grids), was the first model to outperform operational IFS [47] . Shortly after, GraphCast [26] scaled up the graph neural network architecture proposed by Keisler [22] to 0.25 • data and showed improvements over Pangu-Weather. Despite impressive forecast accuracy, existing methods often involve highly customized neural network architectures with minimal ablation studies, making it difficult to identify which components contribute to their success. For example, it is unclear what the benefits of 3D Earth-Specific Transformer over a standard Transformer are, and how critical the multi-mesh message-passing in GraphCast is to its performance. A deeper understanding, and ideally a simplification, of these existing approaches is essential for future progress in the field. Furthermore, establishing a common framework would facilitate the development of foundation models for weather and climate that extend beyond weather forecasting [32] .In this paper, we show that a simple architecture with a proper training recipe can achieve state-of-theart performance. We start with a standard vision transformer (ViT) architecture, and through extensive ablation studies, identify the three key components to the performance of the model: (1) a weatherspecific embedding layer that transforms the input data to a sequence of tokens by modeling the interactions among atmospheric variables; (2) a randomized dynamics forecasting objective that trains the model to predict the weather dynamics at random intervals; and (3) a pressure-weighted loss that weights variables at different pressure levels in the loss function to approximate the density at each pressure level. During inference, our proposed randomized dynamics forecasting objective allows a single model to produce multiple forecasts for a specified lead time by using different combinations of the intervals for which the model was trained. For example, one can obtain a 3-day forecast by rolling out the 6-hour predictions 12 times or 12-hour predictions 6 times. Combining these forecasts leads to significant accuracy improvements, especially for long lead times. We evaluate our proposed method, namely Scalable transformers for weather forecasting (Stormer), on WeatherBench 2 [41] , a widely used benchmark for data-driven weather forecasting. Stormer achieves competitive forecast accuracy of key atmospheric variables for 1-7 days and outperforms the state-of-the-art beyond 7 days. Notably, Stormer achieves this performance by training on more than 5× lower-resolution data and orders-of-magnitude fewer GPU hours compared to the baselines. Finally, our scaling analysis shows that the performance of Stormer improves consistently with increases in model capacity and data size, demonstrating the potential for further improvements. 
background and preliminaries: 
 Given a dataset D = {X i } N i=1 of historical weather data, the task of global weather forecasting is to forecast future weather conditions X T ∈ R V ×H×W given initial conditions X 0 ∈ R V ×H×W , Figure 2 : Different approaches to weather forecasting. Direct and continuous methods output forecasts directly, but continuous forecasting is adaptable to various lead times by conditioning on T . Iterative forecasting generates forecasts at small intervals δt, which are rolled out for the final forecast. Our proposed randomized iterative forecasting combines continuous and iterative methods. in which T is the target lead time, e.g., 7 days; V is the number of input and output atmospheric variables, such as temperature and humidity; and H × W is the spatial resolution of the data, which depends on how densely we grid the globe. This formulation is similar to many image-to-image tasks in computer vision such as segmentation or video frame prediction. However, unlike the RGB channels in natural images, weather data can contain up to 100s of channels. These channels represent actual physical variables that can be unbounded in values and follow complex laws governed by atmospheric physics. Therefore, the ability to model the spatial and temporal correlations between these variables is crucial to forecasting.There are three major approaches to data-driven weather forecasting. The first and simplest is direct forecasting, which trains the model to directly output future weather " X T = f θ (X 0 ) for each target lead time T . Most early works in the field adopt this approach [11, 42, 48, 39, 8, 49] . Since the weather is a chaotic system, forecasting the future directly for large T is challenging, which may explain the poor performances of these early models. Moreover, direct forecasting requires training one neural network for each lead time, which can be computationally expensive when the number of target lead times increases. To avoid the latter issue, continuous forecasting uses T as an additional input: " X T = f θ (X 0 , T ), allowing a single model to produce forecasts at any target lead time after training. MetNet [43, 12, 1] employed the continuous approach for nowcasting at different lead times up to 24 hours, WeatherBench [39] considered continuous forecasting as one of the baselines, and ClimaX [32] used this approach for pretraining. However, since this approach still attempts to forecast future weather directly, it suffers from the same challenging problem of forecasting the chaotic weather in one step. Finally, iterative forecasting trains the model to produce forecasts at a small interval " X δt = f θ (X 0 ), in which δt is typically from 6 to 24 hours. To produce longer-horizon forecasts, we roll out the model by iteratively feeding its predictions back in as input. This is a common paradigm in both traditional NWP systems and the two state-of-the-art deep learning methods, Pangu-Weather and GraphCast. One drawback of this approach is error accumulation when the number of rollout steps increases, which can be mitigated by a multi-step loss function [22, 26, 5, 7] . In iterative forecasting, one can forecast either the weather conditions X δt or the weather dynamics ∆ δt = X δt -X 0 , and X δt can be recovered by adding the predicted dynamics to the initial conditions. In this work, we adopt the latter approach, which we refer to as iterative dynamics forecasting. We show empirically that our approach achieves superior performance relative to the former approach in Section 4.2. Figure 2 summarizes these different approaches.Question: Does the paper discuss the limitations of the work performed by the authors?Answer: [Yes] Justification: Stormer currently operates on 1.40625 • data which is lower-resolution than other leading methods. We mention this limitation in Section 6.Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
methodology: 
 We introduce Stormer, a skillful method for weather forecasting, and show that a simple architecture can achieve competitive forecast performances with a well-designed framework. We first present the overall training and inference procedure of Stormer, then describe the model architecture we use in practice. Section 4.2 empirically demonstrates the importance of each component of Stormer. Lead time (days) We adopt the iterative approach for Stormer, and train the model to forecast the weather dynamics ∆ δt = X δt -X 0 , which is the difference between two consecutive weather conditions, X 0 and X δt , across the time interval δt. A common practice in previous works [22, 26] is to use a small fixed value of δt such as 6 hours. However, as we show in Figure 3a , while small intervals tend to work well for short lead times, larger intervals excel at longer lead times (beyond 7 days) due to less error accumulation. Therefore, having a model that can produce forecasts at different intervals and combine them in an effective manner has the potential to improve the performance of single-interval models. This motivates our randomized dynamics forecasting objective, which trains Stormer to forecast the dynamics at random intervals δt by conditioning on δt:EQUATIONin which P (δt) is the distribution of the random interval. In our experiments, unless otherwise specified, P (δt) is a uniform distribution over three values δt ∼ U{6, 12, 24}. These three time intervals play an important role in atmospheric dynamics. The 6 and 12-hour values help to encourage the model to learn and resolve the diurnal cycle (day-night cycle), one of the most important oscillations in the atmosphere driving short-term dynamics (e.g., temperature over the course of a day). The 24-hour value filters the effects of the diurnal cycle and allows the model to learn longer, synoptic-scale dynamics which are particularly important for medium-range weather forecasting [19] .From a practical standpoint, this randomized objective provides two benefits. First, randomizing δt enlarges the training data, serving as data augmentation. Second, it allows a single trained model to generate various forecasts for a specified lead time T by creating different combinations of intervals δt that sum to T . For example, to forecast 7 days ahead, one could use 12-hour forecasts 14 times or 24-hour forecasts 7 times. Our experiments show that combining these forecasts is crucial for achieving good accuracy, especially for longer lead times. While both our approach and the continuous approach use the time interval as an additional input, we perform iterative forecasting instead of direct forecasting. This avoids the challenge of directly modeling chaotic weather and offers more flexibility for combining different intervals at test time.Due to the large number of variables being predicted, we use a physics-based weighting function to weigh variables near the surface higher. Since each variable lies on a specific pressure level, we can use pressure as a proxy for the density of the atmosphere at each level. This weighting allows the model to prioritize near-surface variables, which are important for weather forecasting and have the most societal impact. The final objective function that we use for training is:EQUATION)The expectation is over δt, X 0 , and X δt which we omit for notational simplicity. In this equation, w(v) is the weight of variable v, and L(i) is the latitude-weighting factor commonly used in previous works to account for the non-uniformity when we grid the spherical globe [40, 22, 35, 32, 3, 26] . The pressure-weighted loss was first introduced by GraphCast [26] , and we show that it also helps with a different architecture.To produce forecasts at a lead time beyond the training intervals, we roll out the model several times. Since the model's forecasts are fed back as input, the forecast error accumulates as we roll out more steps. To alleviate this issue, we finetune the model on a multi-step loss function. Specifically, for each gradient step, we roll out the model K times, and average the objective (2) over the K steps:EQUATIONIn practice, we implement a three-phase training procedure for Stormer. In the first phase, we train the model to perform single-step forecasting, which is equivalent to optimizing the objective in (2).In the second and third phases, we finetune the trained model from the preceding phase with K = 4 and K = 8, respectively. We use the same sampled value of the interval δt for all K steps. We tried randomizing δt at each rollout step, but found that doing so destabilized training as the loss value at each step is of different magnitudes, hurting the final performance of the model. Multi-step finetuning was used in FourCastNet [35] and also adopted in more recent works [22, 26] .At test time, Stormer can produce forecasts at any time interval δt used during training. Thus the model can generate multiple forecasts for a target lead time T by creating different combinations of δt that sum to T . We consider two inference strategies for generating forecasts:Homogeneous In this strategy, we only consider homogeneous combinations of δt, i.e., combinations with just one value of δt. For example, for T = 24 we consider [6, 6, 6, 6] , [12, 12] , and [24] .Best m in n We generate n different, possibly heterogeneous combinations of δt, validate each combination, and pick m combinations with the lowest validation losses for testing.The two strategies offer a trade-off between efficiency and expressivity. The homogeneous strategy only requires running three combinations for each lead time T , while best m in n provides greater expressivity. Upon determining these combinations and executing the model rollouts, we obtain the final forecast by averaging the individual predictions. This approach achieves a similar effect to ensembling in NWP, where multiple forecasts are generated by running NWP models with different perturbed versions of the initial condition [27] . As target lead times extend beyond 5-7 days and individual forecasts begin to diverge due to the chaotic nature of the atmosphere, averaging these forecasts is a Monte Carlo integration approach to handle this sensitivity to initial conditions and the uncertainty in the analyses used as initial conditions [31] . We note that our inference strategy is distinguished from that used in Pangu-Weather: while Pangu-Weather trains a separate model for each time interval δt, we train a single model for all δt values by conditioning on δt. Additionally, while Pangu-Weather relies on a single combination of intervals to minimize rollout steps, our method improves forecast accuracy by averaging multiple forecasts derived from diverse combinations.We instantiate the framework in Section 3.1 with a simple Transformer [45] -based architecture. Due to the similarity of weather forecasting to various dense prediction tasks in computer vision, one might consider applying Vision Transformer (ViT) [10] for this task. However, weather data is distinct from natural images, primarily due to its significantly higher number of input channels, representing atmospheric variables with intricate physical relationships. For example, the wind fields are closely related to the gradient and shape of the geopotential field, and redistribute moisture and heat around the globe. Effectively modeling these interactions is critical to forecast accuracy.The standard patch embedding module in ViT, which uses a linear layer for embedding all input channels within a patch into a vector, may not sufficiently capture the complex interactions among input atmospheric variables. Therefore, we adopt for our architecture a weather-specific embedding module, consisting of two components, variable tokenization and variable aggregation.Variable tokenization Given an input of shape V × H × W , variable tokenization linearly embeds each variable independently to a sequence of shape (H/p) × (W/p) × D, in which p is the patch size and D is the hidden dimension. We then concatenate the output of all variables, resulting in a sequence of shape (H/p) × (W/p) × V × D.We employ a single-layer cross-attention mechanism with a learnable query vector to aggregate information across variables. This module operates over the variable dimension on the output of the tokenization stage to produce a sequence of shape (H/p)×(W/p)×D. This module offers two primary advantages. First, it reduces the sequence length by a factor of V , significantly alleviating the computational cost as we use transformer to process the sequence. Second, unlike standard patch embedding, the cross-attention layer allows the models to learn non-linear relationships among input variables, enhancing the model's capacity to capture complex physical interactions. We present the complete implementation details of the weather-specific embedding in Section B.Figure 3b shows the superior performance of weather-specific embedding to standard patch embedding at all lead times from 1 to 10 days. A similar weather-specific embedding module was introduced by ClimaX [32] to improve the model's flexibility when handling diverse data sources with heterogeneous input variables. We show that this specialized embedding module outperforms the standard patch embedding even when trained on a single dataset, due to its ability to model interactions between atmospheric variables through cross-attention effectively.Following weather-specific embedding, the tokens are processed by a stack of transformer blocks [45] .In addition to the input X 0 , the block also needs to process the time interval δt. We do this by replacing the standard layer normalization module used in transformer blocks with adaptive layer normalization (adaLN) [37] . In adaLN, instead of learning the scale and shift parameters γ and β as independent parameters of the network, we regress them with an one-layer MLP from the embedding of δt. Compared to ClimaX [32] which only adds the lead time embedding to the tokens before the first attention layer, adaLN is applied to every transformer block, thus amplifying the conditioning signal.Figure 3c shows the consistent improvement of adaLN over the additive lead time embedding used in ClimaX. Adaptive layer norm was widely used in both GANs [21, 4] and Diffusion [9, 36] to condition on additional inputs such as time steps or class labels. Figure 7 illustrates Stormer's architecture. We refer to [32] for illustrations of the weather-specific embedding block.Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?• The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors). • It should be clear whether the error bar is the standard deviation or the standard error of the mean. Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.• The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
experiments: 
 We compare Stormer with state-of-the-art weather forecasting methods, and conduct extensive ablation analyses to understand the importance of each component in Stormer. We also study Stormer scalability by varying model size and the number of training tokens. We conduct all experiments on WeatherBench 2 (WB2) [41] , a standard benchmark for data-driven weather forecasting.We train and evaluate Stormer on the ERA5 dataset from WB2, which is the curated version of the ERA5 reanalysis data provided by ECMWF [17] . In its raw form, ERA5 contains hourly data from 1979 to the current time at 0.25 • (721×1440 grids) resolution, with different atmospheric variables spanning 137 pressure levels plus the Earth's surface. WB2 downsamples this data to 6-hourly with 13 pressure levels and provides different spatial resolutions. In this work, we use the 1.40625 • (128×256 grids) data. We use four surface-level variables -2-meter temperature (T2m), 10-meter U and V components of wind (U10 and V10), and Mean sea-level pressure (MSLP), and five atmospheric variables -Geopotential (Z), Temperature (T), U and V components of wind (U and V), and Specific humidity (Q), each at 13 pressure levels {50, 100, 150, 200, 250, 300, 400, 500, 600, 700, 850, 925, 1000}. We use 1979 to 2018 for training, 2019 for validation, and 2020 for testing.Stormer architecture: For the main comparison in Section 4.1, we report the results of our largest Stormer model with 24 transformer blocks, 1024 hidden dimensions, and a patch size of 2, which is equivalent to ViT-L except for the smaller patch size. We vary the model size and patch size in the scaling analysis. For the remaining experiments, we report the performance of the same model as for the main result, but with a larger patch size of 4 for faster training. Training: For the main result in Section 4.1, we train Stormer in three phases, as described in Section 3.1.2. We train the model for 100 epochs for the first phase, 20 epochs for the second, and 20 epochs for the third. We perform early stopping on the validation loss aggregated across all variables, and evaluate the best checkpoint of the final phase on the test set. For the remaining experiments, we only train Stormer for the first phase due to computational constraints.Evaluation: We evaluate Stormer and two deep learning baselines on forecasting nine key variables: T2m, U10, V10, MSLP, Z500, T850, Q700, U850, and V850. These variables are also used to report the headline scores in WB2. For each variable, we evaluate the forecast accuracy at lead times from 1 to 14 days, using the latitude-weighted root-mean-square error (RMSE) metric. For the main results, we use best m in n inference for rolling out Stormer as it yields the best result, with m = 32 and n = 128 chosen randomly from all possible combinations. For the remaining experiments, we use homogeneous inference for faster evaluations. We provide results on the non-ensemble version of Stormer, probabilistic metrics with IC perturbations, a comparison between two inference strategies, and additional ablation studies in Appendix C.We compare the forecast performance of Stormer with Pangu-Weather [3] and GraphCast [26] , two leading deep learning methods for weather forecasting. Pangu-Weather employs a 3D Earth-Specific Transformer architecture trained on the same variables as Stormer, but with hourly data and a higher spatial resolution of 0.25 • . GraphCast is a graph neural network that was trained on 6-hourly ERA5 data at 0.25 • , using 37 pressure levels for the atmospheric variables, and two additional variables, Lead time (days) total precipitation and vertical wind speed. Both Pangu-Weather and GraphCast are iterative methods. GraphCast operates at 6-hour intervals, while Pangu-Weather uses four distinct models for 1-, 3-, 6-, and 24-hour intervals, and combines them to produce forecasts for specific lead times. We include Climatology as a simple baseline. We also compare Stormer with IFS HRES, the state-of-the-art numerical forecasting system, and IFS ENS (mean), which is the ensemble version of IFS. Since WB2 does not provide forecasts of these numerical models beyond 10 days, we defer the comparison against these models to Appendix C.1.Results: Figure 4 evaluates different methods on forecasting nine key weather variables at lead times from 1 to 14 days. For short-range, 1-5 day forecasts, Stormer's accuracy is on par with or exceeds that of Pangu-Weather, but lags slightly behind GraphCast. At longer lead times, Stormer excels, consistently outperforming both Pangu-Weather and GraphCast from day 6 onwards by a large margin. Moreover, the performance gap increases as we increase the lead time. At 14 day forecasts, Stormer performs better than GraphCast by 10% -20% across all 9 key variables. Stormer is also the only model in this comparison that performs better than Climatology at long lead times, while other methods approach or even do worse than this simple baseline. The model's superior performance at long lead times is attributed to the use of randomized dynamics training, which improves forecast accuracy by averaging out multiple forecasts, especially when individual forecasts begin to diverge.Moreover, we also note that Stormer achieves this performance with much less compute and training data compared to the two deep learning baselines. We train Stormer on 6-hourly data of 1.40625 • with 13 pressure levels, which is approximately 190× less data than Pangu-Weather's hourly data at 0.25 • and 90× less than that used for GraphCast, which also uses 6-hourly data but at a 0.25 • resolution with 37 pressure levels. The training of Stormer was completed in under 24 hours on 128 A100 GPUs. In contrast, Pangu-Weather took 60 days to train four models on 192 V100 GPUs, and GraphCast required 28 days on 32 TPUv4 devices. This training efficiency will facilitate future works that build upon our proposed framework.We analyze the significance of individual elements within Stormer by systematically omitting one component at a time and observing the difference in performance.We evaluate the effectiveness of our proposed randomized iterative forecasting approach. Figure 5a compares the forecast accuracy on surface temperature of Stormer and three models trained with different values of δt. Stormer consistently outperforms all singleinterval models at all lead times, and the performance gap widens as the lead time increases. We attribute this result to the ability of Stormer to produce multiple forecasts and combine them to improve accuracy. We note that Stormer achieves this improvement with no computational overhead compared to the single-interval models, as the different models share the same architecture and were trained for the same duration.Impact of pressure-weighted loss: Figure 5b shows the superior performance of Stormer when trained with the pressure-weighted loss. Intuitively, the weighting factor prioritizes variables that are nearer to the surface, as these variables are more important for weather forecasting and climate science. Dynamics vs. absolute forecasts: We justify our decision to forecast the dynamics ∆ δt by comparing with a counterpart that forecasts X δt . Figure 5c shows that forecasting the changes in weather conditions (dynamics) is consistently more accurate than predicting complete weather states. One possible explanation for this result is that it is simpler for the model to predict the changes between two consecutive weather conditions than the entire state of the weather; thus, the model can focus on learning the most significant signal, enhancing forecast accuracy.We examine Stormer's scalability in terms of model size and training tokens. We evaluate three variants -Stormer-S, Stormer-B, and Stormer-L, with parameter counts similar to ViT-S, ViT-B, and ViT-L, respectively. To understand the impact of training token count, we vary the patch size from 2 to 16, quadrupling the training tokens each time the patch size is halved. Figure 6 shows a significant improvement in forecast accuracy with larger models, and the performance gap widens with increased lead time. Since we do not perform multi-step fine-tuning for these models, minor performance differences at short intervals may magnify over time. While multi-step fine-tuning could potentially reduce this gap, it is unlikely to eliminate it entirely. Reducing the patch size also improves the performance of the model consistently. From a practical view, smaller patches mean more tokens and consequently more training data. From a climate perspective, smaller patches capture finer weather details and processes not evident in larger patches, allowing the model to more effectively capture physical dynamics that drive weather patterns. 
related work: 
 Deterministic weather forecasting Deep learning offers a promising approach to weather forecasting due to its fast inference and high expressivity. Early efforts [11, 42, 48] attempted training simple architectures on small weather datasets. To facilitate progress, WeatherBench [40] provided standard datasets and benchmarks, leading to subsequent works that trained Resnet [15] and UNet architectures [49] for weather forecasting. These works showed the potential of deep learning but still displayed inferior accuracy to numerical systems. However, significant improvements have been made in the last few years. Keisler [22] proposed a graph neural network (GNN) that performs iterative forecasting with 6-hour intervals, performing comparably with some NWP models. FourCastNet [35] trained an adaptive Fourier neural operator and was the first neural network to run on 0.25 • data. Pangu-Weather [3] , with its 3D Earth-Specific Transformer design, trained on high-resolution data, surpassed the benchmark IFS model. Following this, GraphCast [26] scaled up Keisler's GNN architecture to 0.25 • , achieving even better results than Pangu-Weather. FuXi [6] was a subsequent work that trained a SwinV2 [28] on 0.25 • data and showed improvements over GraphCast at long lead times. However, FuXi requires finetuning multiple models specialized for different time ranges, increasing model complexity and computation. FengWu [5] was a concurrent work with FuXi that also focused on improving long-horizon forecasts, but has not revealed complete model architecture and training details. ClimODE [46] introduced physical inductive biases to provide better interpretability but was empirically inferior to existing methods.Probabilistic weather forecasting In addition to high accuracy, a desired ability of a weather forecasting model is to quantify forecast uncertainty. One common approach to achieve this is to combine an existing architecture with a probabilistic loss function. Gencast [38] was one of the first works in this direction, combining the Graphcast architecture with a diffusion objective [18, 21] , followed by Graph-EFM [33] , which combined a hierarchical variant of Graphcast with the VAE objective [23] . This approach allows the model to generate multiple forecasts and estimate uncertainty after training. In an orthogonal approach, NeuralGCM [25] proposed a hybrid forecasting system that combined a differentiable dynamical core with ML components for end-to-end training. The dynamical core allows the method to leverage powerful general circulation models and generate forecast ensembles via IC perturbations similar to NWP. However, the dynamical core in NeuralGCM is more computationally expensive than forward-passing a neural network and can limit the method's performance with an imperfect circulation model. Weather and climate modeling is crucial for understanding and tackling climate change. Creating better models based on deep learning could offer faster and cheaper alternatives to expensive numerical simulations. These models could improve weather predictions, extreme event forecasts, and climate projections. They might also help reduce the carbon footprint, better prepare for natural disasters, and enhance our knowledge of the Earth. However, relying only on deep learning models requires careful checks and monitoring, especially when predicting new or uncertain scenarios. In all experiments, the variable tokenization module is a standard patch embedding layer usually used in ViT, and the aggregation module is a single-layer multi-head cross-attention. The first embedding of δt is a linear layer, and the adaLN module in each block employs a 2-layer MLP. For the main comparison with the current methods, we train a Stormer model with a patch size of 2, 1024 hidden dimensions, and 24 Stormer blocks. For the scaling experiments, we vary the hidden dimensions, number of blocks, and patch size. For the rest of the ablation studies, we use a patch size of 4, hidden dimension of 1024, and 24 blocks.Input normalization We compute the mean and standard deviation for each variable in the input across all spatial positions and all data points in the training set. This means each variable is associated with a scalar mean and scalar standard deviation. During training, we standardize each variable by subtracting it from the associated mean and dividing it by the standard deviation.Output normalization Unlike the input, the output that the model learns to predict is the difference between two consecutive steps. Therefore, for each variable, we compute the mean and standard deviation of the difference between two consecutive steps in the training set. What it means to be "consecutive" depends on the time interval δt. If δt = 6, we collect all pairs in training data that are 6-hour apart, compute the difference between two data points in each pair, and then compute the mean and standard deviation of these differences. Since we train Stormer with randomized δt, we repeat the same process for each value of δt.As mentioned in Section 3, we train Stormer in three phases with the following objective:EQUATIONwhere the number of rollout steps K is equal to 1, 4, and 8 in phase 1, 2, and 3, respectively. For phases 2 and 3, we finetune the best checkpoint from the preceding phase.For pressure-level variables, we assign weights proportionally to the pressure level of each variable. For 4 surface variables, we assign w = 1 for T2m and w = 0.1 for the remaining variables U10, V10, and MSLP. The surface weights were proposed by GraphCast [26] and we did not perform any additional hyperparameter tuning.For the 1st phase, we train the model for 100 epochs. We optimize the model using AdamW [24] with learning rate of 5e -4, parameters (β 1 1 = 0.9, β 2 = 0.95) and weight decay of 1e -5. We used a linear warmup schedule for 10 epochs, followed by a cosine schedule for 90 epochs.For the 2nd and 3rd phases, we train the model for 20 epochs with a learning rate of 5e -6 and 5e -7, respectively. We used a linear warmup schedule for 5 epochs, followed by a cosine schedule for 15 epochs. Other hyperparameters remain the same.We perform early stopping for all phases, where the criterion is the validation loss aggregated across all variables at lead times of 1 day, 3 days, and 5 days for phases 1, 2, and 3, respectively. We save the best checkpoint for each phase using the same criterion.We use PyTorch [34] , Pytorch Lightning [13] , timm [50] , numpy [14] and xarray [20] for data processing and model training. We trained Stormer on 128 40GB A100 devices. We leverage mixed-precision training, Fully Sharded Data Parallel, and gradient checkpointing to reduce memory.As different models are trained on different resolutions of data, we follow the practice in WB2 to regrid the forecasts of all models to the same resolution of 1.40625 • (128 × 256 grid points). We then calculate evaluation metrics on this shared resolution. Similarly to WB2, we evaluate forecasts with initial conditions at 00/12UTC for all days in 2020. U10m (m/s) Figure 8 compares Stormer with both deep learning and numerical methods. We take IFS and IFS ENS from WB2 which is only available until day 10. Similar to its deep learning counterparts, Stormer achieves lower RMSE compared to the IFS model for most variables, except for near-surface temperature (T2m) at initial lead times, and only performs slightly worse than IFS ENS. To the best of our knowledge, Stormer is the first model trained on 1.40625 • data to surpass IFS.Additionally, we compare Stormer and the baselines on latitude-weighted ACC, another common verification metric for weather forecast models. ACC represents the Pearson correlation coefficient between forecast anomalies relative to climatology and ground truth anomalies relative to climatology. ACC ranges from -1 to 1, where 1 indicates perfect correlation, and -1 indicates perfect anticorrelation. We refer to WB2 [41] for the formulation of ACC. Figure 9 shows that similarly to RMSE, Stormer achieves competitive performance from 1 to 5 days, and outperforms the baselines by a large margin beyond 6 days.We verify the importance of multi-step fine-tuning by comparing Stormer after the 1st phase (K = 1) and after the 3rd phase (K = 8). Figure 10 shows that multi-step fine-tuning significantly improves performance at long lead times. Even though our inference strategy can be considered ensemble forecasting, we note that it is much cheaper and more efficient than common techniques such as training multiple networks, dropout, or IC perturbations, as we only have to train a single neural network and do not need extensive Since Stormer can produce forecast ensembles after training, we can consider it a probabilistic forecast system. However, our preliminary results suggested that different forecasts from Stormer are underdispersive and should not be used for uncertainty estimation. 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 of 128 ensemble members.Figure 12 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.Our two inference strategies, Homogeneous and Best m in n, provide a tradeoff between efficiency and forecast accuracy. Figure 13 compares the performance of these two strategies across different variables at different lead times. The results show that Homogeneous performs competitively with Best m in n, while being much more efficient, requiring only 3 forward passes compared to n. Lead time (days) We visualize forecasts produced by Stormer at lead times from 1 days to 14 days for 9 key variables. All forecasts are initialized at 0UTC January 26th 2020. Each figure illustrates one lead time, where each row is for each variable. The first column shows the initial condition, the second column shows the ground truth at that lead time, the third column shows the forecast, and the last column shows the bias, which is the difference between the forecast and the ground truth. Justification: This paper introduces Stormer, a scalable and skillful method for mediumrange weather forecasting. Both the abstract and introduction reflect this contribution.Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in the paper. • The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers. • The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper. 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA]Justification: NA Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes] Justification: We used and cited Weatherbench 2, an open-source benchmark for data-driven weather forecasting.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.• If this information is not available online, the authors are encouraged to reach out to the asset's creators. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix:  


Decision: This paper is about weather forecasting, in particular the authors propose a new (vision) transformer architecture that has several advantages over the state of the art: weather-specific patch embedding, randomized dynamics forecast, and pressure-weighted loss.

The main contribution of this paper is to show that a simple transformer architecture can obtain state of the art performance with simple settings, in particular with a careful transformer setup that is specific to weather forecasting.

Reviewers were positive, in particular noting that the paper is very well written and understandable, the simple approach, the important problem setting, ablation results, and improvements to the state of the art in long term weather forecasting. During rebuttal the authors provided additional results requested by the reviewers, in particular comparison with other baselines and a fairer comparison as Stormer is basically an ensemble model, a non-ensemble Stormer still performs very well compared to the state of the art.

To improve the paper I recommend to consider the (many) reviewer comments, in particular clarifying contributions, (conceptual) comparisons of the proposed method to previous methods in the literature, and incorporating the results presented in the rebuttal phase into the main paper or appendix.

This area chair also suggest to consider uncertainty estimation, as was also mentioned by reviewers, as it is an important component of weather forecasting, and since the method is partially an ensemble, simple uncertainty estimates could be extracted from it.

Given the positive contribution to the state of the art, and good opinion from reviewers, this paper should be accepted.

Reviews: review: paper_summary: paper_summary: The authors introduce a vision transformer-based method, Stormer, designed for medium-range weather forecasting. Ablations identify multiple important components of the method, including a weather-specific patch embedding, "randomized" dynamics forecast, and a pressure-weighted loss. The randomized forecasting component is similar to prior continuous models but complemented with iterative rollouts and an ensemble technique "best m in n" that exploits the multiple discretizations possible to forecast a specific lead time iteratively.
The experiments show competitive results in terms of RMSE and ACC against deep learning and physics-based baselines.
strength_weakness: strength_weakness: - The evaluation is somewhat unfair. 

The paper is essentially using an ensembling technique but comparing against deterministic models. It is well known that ensembling improves ensemble-mean RMSE and ACC scores, especially for longer-range horizons (which is were Stormer, unsurprisingly, shines the most against the non-ensemble baselines). This can be actually clearly observed with the physics-based baselines IFS HRES (deterministic) and IFS ENS (ensemble mean) in Figures 8 and 9, where the ensemble mean shines the most on long-range horizons. 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). 

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. This probabilistic evaluation is actually supported by Weatherbench2, so it should be easy for the authors to extend their current evaluation with probabilistic metrics (+this will give you IFS ENS baseline up to 15 days ahead for free...).
- On top of probabilistic metrics, it would be instructive to see an analysis of the generated spectra of Stormer.
- The following is wrong: "it is unclear (...) how critical the multi-mesh message passing in GraphCast is to its performance". The authors seem to have missed section "7.3.1. Multi-mesh ablation" in the GraphCast paper.
- I would like to see a more transparent discussion of the exact contributions of this work and a more careful contextualization with prior work. For example, in section 3.1 it would be good to be more candid about that 1) the objective is exactly the same as for a continuous forecasting model. The only difference between the two seems to be during inference time (and the different lead times used for training); 2) the overall train+inference method is essentially the same (neural architectures aside) as for Pangu-weather (PW) but with Stormer training one model for all lead times, while PW trains one per lead time. Noting the pros and cons of each would be useful to. Even more useful would be to have a crisp ablation for Stormer, where you drop the time-conditioning and train 3 separate models for each lead time (6, 12, 24hours). This would control for the differences in architecture between Stormer and PW, giving valuable insights into the pros and cons of each; 3) The pressure-weighted loss is taken from GraphCast but this is not properly referenced in the manuscript (this is not mentioned in section 3.1.1 that introduces it but only much later); 4) Same for multi-step fine-tuning, which is very common in prior work.
- Relatedly, I think that the paper could benefit from discussing the related work more in-depth. For example, I seem to have missed any discussion that carefully compares Stormer with other (vision) transformer for weather forecasting methods such as Pangu-Weather, ClimaX, FengWu, etc.
- Please include results on other variables than just T2M for your ablations. Different variables might behave quite differently... Also, this is a particularly unfortunate choice for Fig. 5b because the fact that the pressure-weighted loss improves T2M RMSE is fairly trivial since the pressure-weighted loss is very strongly upweighting the influence of T2M on the total loss relative to other variables.
- Line 292: *"We note that Stormer achieves this improvement with no computational overhead compared to the single-interval models, as the different models share the same architecture and were trained for the same duration."*. Is this true? To me it seems that at inference time there does exist a computational overhead due to the larger size of the ensemble...

[1] A Practical Probabilistic Benchmark for AI Weather Models (https://arxiv.org/abs/2401.15305)
questions: questions: - 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. Exploring their (dis-)advantages would be beneficial and very interesting. This would also tie well with the paper's desire to carefully ablate important components of the design stack.
- Can you expand on how the "best m" combinations are chosen? By combinations you mean e.g. (3h+6h=9 and 6h+3h=9h, but not 9*1h?) Are you choosing them on the validation set? How do you choose the "n" combinations that you validate? Also, some lead times won't have as many combinations (e.g. 6h only has one, 12h two etc.), right? That would seem worth mentioning as a potential limitation of this ensembling technique.
- 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)?
- For completeness, consider specifying the parameter sizes of your models in section 4.3. explicitly. Also, do you discuss the exact hyperparameters used for the smaller model sizes?
- Is input normalization done per variable or per variable AND pressure level?
limitations: limitations: The authors have not adequately addressed the limitations, mentioning the spatial resolution as the only limitation of their method (but also claiming it to be an advantage in other parts of the paper). A more upfront discussion of (potential) limitations would be welcome. For example, there is no discussion on how the proposed cross-attention-based patch embedding might scale to higher spatial resolutions or how it might impact inference speed.
paper_summary: paper_summary: The paper proposes a transformer-based model for weather prediction. Experiments show improvement in downstream predictions.
strength_weakness: strength_weakness: - The paper has very limited novelty and is incremental. The transformer-based architecture is the same as ClimaX and other transformer-based methods; the variable embeddings are similar to what ClimaX proposes, multi-step fine-tuning is the same as FourCastNet, and pressure-weighted loss is by GraphCast. 
- The paper misses out on a ton of related works, such as  ClimODE (https://arxiv.org/abs/2404.10024), Neural GCM (https://arxiv.org/abs/2311.07222), GenCast (https://arxiv.org/abs/2312.15796), etc all seem to be missing and comparisons to them are missing.
- 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. 
- Although the authors utilize randomized dynamic forecasting mechanism to have stability over predictions, they still restrict the whole method to output only point estimates rather than giving uncertainties.
questions: questions: -  What is the computational complexity of training and inference? Are there any ablation studies regarding the compactness and inference time studies?
- How do you counter for boundary conditions? Do the variable and patch embeddings respect boundary conditions?
- Can the model accommodate different lead-time resolution query points (t = 1hr, 19hr, 91hr, etc.) that differ from the dataset, showcasing its generalizability and applicability in modeling weather?
- Does the method use standard 2D ViT or a 3D transformer as the stormer block?
limitations: limitations: See weaknesses and questions.
paper_summary: paper_summary: This paper introduces a deep learning weather prediction model called Stormer. Stormer is vision transformer type network that employs various techniques to improve the network's performance in weather forecasting applications, including a "weather-specific embedding" that first processes each variable separately, and a technique for producing multiple weather scenarios for each input by using the model's ability to be tuned to various lead times. By using a lower spatial resolution than competing models, Stormer is much faster to train and inference than current state of the art models, while still achieving competitive performance at short-to-medium time scales and superior performance at longer time scales.
strength_weakness: strength_weakness: I do not see serious weaknesses with the paper. However, one point that could be improved is with the randomized forecast strategy introduced in the paper. PanguWeather is also able to make weather forecasts with variable lead times (although it achieves this with separately trained models). If I'm not mistaken, PanguWeather could then also be used to implement the randomized forecast strategy. It would be interesting to see this included in the model comparisons.
questions: questions: 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?
limitations: limitations: The authors do briefly discuss the limitations of their approach and possible future directions. However, I think one point that gets glossed over by the authors is that due to its low resolution Stormer is unable to resolve weather features that models such as PanguWeather and GraphCast can resolve. Thus the performance gain of lower resolution comes at a cost of worse ability to resolve more localized weather features. It's also not clear how this impacts the model's representation of weather extremes, which are often found at small scales.

Comments: title: Response to the authors
Thank you
Releasing code and checkpoints
Checking in (2)
Checking in (2)
Checking in
Checking in
Rebuttal by Authors (2/2)
Rebuttal by Authors (3/3)
Rebuttal by Authors (2/3)
comment: 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.
We thank the reviewer for appreciating our efforts and raising the score. We strongly agree with the reviewer that in terms of architecture, ClimaX and Stormer are very similar. However, we would also like to emphasize that this is the exact message we wanted to convey in this paper: with a carefully designed training recipe, we can significantly boost the performance of an existing architecture for weather forecasting. We hope that the experimental results and ablation studies in this paper will shed light on the impact of different design choices for future works in the field.
We thank the reviewer for the constructive and thoughtful feedback. We strongly agree with the reviewer that there is a clear distinction between non-ensemble and ensemble forecasting, and each approach has its pros and cons. We will update the methodology and experiment sections to discuss this in detail as the reviewer suggested. We would also like to answer the remaining comments below.

> To highlight differences between models on short-range forecasts (~1-3 days ahead), consider showing relative plots (e.g. relative to HRES for non-ensemble models, and ENS for ensembles).

Thank you for the suggestion. We will include this plot in the revised paper.

> Given the architectural similarity of Climax and Stormer, can the underperformance of Climax (for weather forecasting) be mostly explained by its training data?

We note that the ERA5 data used to train ClimaX and Stormer are very similar, and ClimaX was even pretrained with CMIP6 before being finetuned on ERA5. We believe the difference in performance mainly comes from the training recipe for each method. ClimaX was trained to perform direct forecasting for each lead time, while Stormer performs iterative forecasting with additional multi-step finetuning and pressure-weighted loss. The empirical results suggested that the latter approach excels in weather forecasting, which agrees with the current literature (Pangu, Graphcast, etc., are all iterative methods). This is also the main message we want to convey in this paper: with a carefully designed training recipe, we can significantly boost the performance of an existing architecture for weather forecasting.
I appreciate the authors' comprehensive rebuttal and the additional insightful results provided in the PDF. The response has addressed most of my concerns, leaving only a key issue regarding the comparison of ensembles and non-ensembles (i.e. point predictions) that I touch on below. Based on this, I am raising my score to 6 and recommend accepting the paper. I apologize for the delayed reply, which was due to my being on holiday.

The non-ensemble Stormer results are particularly valuable. This version performs comparably or slightly below GraphCast and surpasses Pangu-Weather, which are impressive outcomes. I strongly recommend including this plot and its associated discussion in the main text. 

My primary concern, which I urge you to address in your revised version, is the need for extreme caution when comparing your ensemble version with non-ensemble baselines, especially for setting a correct example for follow-up work. This relates to your statement:

> 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.

While I understand your perspective, I have two reservations:

1. Ensembling typically improves results. Comparing point predictions to ensemble-mean results disadvantages the former. Although your ensembling technique is model-specific and interesting, a fair comparison would be against baseline ensembles (e.g., IFS ENS, stochastic models like NeuralGCM, diffusion-based models like GenCast, etc.) You do this in the rebuttal PDF and I would like to see the same level of rigor across your whole paper.
2. An ensemble-mean prediction, whether from Stormer or any other ensemble like IFS ENS, likely does not represent a physically realistic state. This is evident in the overly smooth predictions in Figures 12-16 and would likely be apparent in the spectra of these predictions compared to non-ensemble ones.

To address these concerns, I strongly recommend:

- Clearly distinguishing between ensemble and non-ensemble results throughout the paper.
- Providing separate comparisons: ensemble Stormer with other ensemble methods (e.g., IFS ENS) and non-ensemble Stormer with non-ensemble baselines. It's fair of course to show them all together in a plot (as in your rebuttal PDF), but I would be cautious about doing direct comparisons between ensemble predictions and point predictions which boil down to "the ensemble-mean is better".

These additions would significantly strengthen the paper and provide valuable context for future research in this field. I trust that the authors can revise their paper to reflect this. 
 
Some minor comments:
- To highlight differences between models on short-range forecasts (~1-3 days ahead), consider showing relative plots (e.g. relative to HRES for non-ensemble models, and ENS for ensembles).
- Given the architectural similarity of Climax and Stormer, can the underperformance of Climax (for weather forecasting) be mostly explained by its training data?
Thanks for the detailed response and clarifications. I am adjusting my score. However, the method still looks like a minor adaptation of ClimaX with an adaptive layer norm.
We thank the reviewer again for the constructive review and feedback. We will make sure to include these discussions into the paper.
As a side note, we promise to publicly release all the code and checkpoints to the community.
Thanks for the response, this clarifies my comments. I'll keep the original review score.
We thank you again for your constructive review and feedback. We sincerely hope the reviewer has had time to read our rebuttal and additional experiments, which we believe have addressed and answered the reviewer's concerns and questions. As the discussion ends today, please let us know any further questions or concerns you would like to discuss. We understand that you are busy, and we truly value the time and effort you put into this process.
We thank you again for your constructive review and feedback. We sincerely hope the reviewer has had time to read our rebuttal and additional experiments, which we believe have addressed and answered the reviewer's concerns and questions. As the discussion ends today, please let us know any further questions or concerns you would like to discuss. We understand that you are busy, and we truly value the time and effort you put into this process.
Thank you again for your review. We made significant efforts to address the reviewer's concerns and sincerely hope that our responses have adequately addressed the concerns you previously raised. Since the discussion period is short and drawing to a close soon, are there further questions or concerns we should discuss? We understand that you are busy, and we truly value the time and effort you put into this process. Thank you in advance for your continued support.
Thank you again for your review. We made significant efforts to address the reviewer's concerns and sincerely hope that our responses have adequately addressed the concerns you previously raised. Since the discussion period is short and drawing to a close soon, are there further questions or concerns we should discuss? We understand that you are busy, and we truly value the time and effort you put into this process. Thank you in advance for your continued support.
> 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.
> For completeness, consider specifying the parameter sizes of your models in section 4.3. explicitly. Also, do you discuss the exact hyperparameters used for the smaller model sizes?

The total parameter count of Stormer is 400 million. Different model sizes of Stormer correspond to the standard sizes of ViT models in the computer vision literature, defined at https://github.com/huggingface/pytorch-image-models/blob/main/timm/models/vision_transformer.py. We will add the parameter count and hyperparameters of all models to the updated version. 

> Is input normalization done per variable or per variable AND pressure level?

Input normalization is done per variable AND pressure level.
> 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.