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

Fuente: arXiv
Enregistré dans:
Détails bibliographiques
Auteurs principaux: Willard, Jared D., Harrington, Peter, Subramanian, Shashank, Mahesh, Ankur, O'Brien, Travis A., Collins, William D.
Format: Preprint
Publié: 2024
Sujets:
Accès en ligne:
Tags: Ajouter un tag
Pas de tags, Soyez le premier à ajouter un tag!
_version_ 1866914777779929088
author Willard, Jared D.
Harrington, Peter
Subramanian, Shashank
Mahesh, Ankur
O'Brien, Travis A.
Collins, William D.
author_facet Willard, Jared D.
Harrington, Peter
Subramanian, Shashank
Mahesh, Ankur
O'Brien, Travis A.
Collins, William D.
contents 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 SwinV2 transformer on ERA5 data, and find that it attains superior forecast skill when compared against IFS. We present some ablations on key aspects of the training pipeline, exploring different loss functions, model sizes and depths, and multi-step fine-tuning to investigate their effect. We also examine the model performance with metrics beyond the typical ACC and RMSE, and investigate how the performance scales with model size.
format Preprint
id arxiv_https___arxiv_org_abs_2404_19630
institution arXiv
publishDate 2024
record_format arxiv
spellingShingle 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.
Collins, William D.
Machine Learning
68T07, 86A10
J.2; I.2.6
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 SwinV2 transformer on ERA5 data, and find that it attains superior forecast skill when compared against IFS. We present some ablations on key aspects of the training pipeline, exploring different loss functions, model sizes and depths, and multi-step fine-tuning to investigate their effect. We also examine the model performance with metrics beyond the typical ACC and RMSE, and investigate how the performance scales with model size.
title Analyzing and Exploring Training Recipes for Large-Scale Transformer-Based Weather Prediction
topic Machine Learning
68T07, 86A10
J.2; I.2.6
url https://arxiv.org/abs/2404.19630