Analyzing and Exploring Training Recipes for Large-Scale Transformer-Based Weather Prediction

Willard, Jared D., Harrington, Peter, Subramanian, Shashank, Mahesh, Ankur, O’Brien, Travis A. and Collins, William D. ORCID: https://orcid.org/0000-0002-4463-9848 (2025) Analyzing and Exploring Training Recipes for Large-Scale Transformer-Based Weather Prediction. Artificial Intelligence for the Earth Systems, 4 (2). ISSN 2769-7525

Preview

PDF (aies-AIES-D-24-0061.1) - Published Version Available under License Creative Commons Attribution. Download (23MB) | Preview

Abstract

The rapid rise of deep learning (DL) in numerical weather prediction (NWP) has led to a proliferation of models which forecast atmospheric variables with comparable or superior skill than traditional physics-based NWP. However, among these leading DL models, there is a wide variance in both the training settings and architecture used. Further, the lack of thorough ablation studies makes it hard to discern which components are most critical to success. In this work, we show that it is possible to attain high forecast skill even with relatively off-the-shelf architectures, simple training procedures, and moderate compute budgets. Specifically, we train a minimally modified Swin Transformer V2 (SwinV2) on ERA5 data and find that it attains superior skill in terms of mean-square errors of deterministic forecasts when compared against the European Centre for Medium-Range Weather Forecasts’ Integrated Forecasting System (IFS). Almost all DL–NWP systems share a core set of hyperparameters and design decisions. To aid and expedite future DL–NWP research, we present an in-depth, systematic exploration of different loss functions, model sizes and depths, patch sizes, and multistep training objectives. We also examine the model performance with metrics beyond the typical accuracy (ACC) and RMSE and investigate how the performance scales with model size. Through our open-source code, scoring pipelines, and models, we share our findings on key aspects of the training pipeline. These ablations reduce the necessity for expensive hyperparameter tuning and lower the barrier to entry for future DL–NWP research.

Item Type:	Article
Additional Information:	Data availability statement. The primary code used in this study is available at https://github.com/NERSC/swin_v2_weather/. The repository README also has links to ERA5 data, neural network model weights, and precomputed statistics for normalization. The version of Earth2Mip used that contains the SwinV2 transformer implementation can be found at https://github.com/jdwillard19/earth2mip-swin-fork.
Faculty \ School:	Faculty of Science > School of Environmental Sciences Faculty of Science > Tyndall Centre for Climate Change Research University of East Anglia Research Groups/Centres > Faculty of Science > Research Centres > Tyndall Centre for Climate Change Research
Related URLs:	https://journals.ametsoc.org/view/journa...
Depositing User:	LivePure Connector
Date Deposited:	22 Sep 2026 09:39
Last Modified:	30 Sep 2026 11:00
URI:	https://ueaeprints.uea.ac.uk/id/eprint/104610
DOI:	10.1175/AIES-D-24-0061.1

Downloads

Actions (login required)

View Item