🤖 machine learning

Sonny: Breaking the Compute Wall in Medium-Range Weather Forecasting

The paper introduces Sonny, an efficient hierarchical transformer model that achieves competitive medium-range weather forecasting performance on a single GPU by utilizing a novel two-stage StepsNet architecture and exponential moving average training to overcome the high compute barriers of existing deep learning weather models.

Original authors: Minjong Cheon

Published 2026-03-24

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Minjong Cheon

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to predict the weather for the next week. For decades, scientists have used massive, super-computer-powered systems (like a giant, complex orchestra) to simulate the atmosphere. These systems are incredibly accurate, but they are so expensive and energy-hungry that only a few elite organizations can afford to run them.

Recently, AI models have tried to learn from past weather data instead of simulating physics. They are fast and cheap, but to get them to work well, they usually require "super-computing" power (hundreds of graphics cards) to train. This creates a barrier: if you are a university researcher with just a few computers, you can't build a world-class weather model.

Enter "Sonny."

Think of Sonny as the "smart, efficient student" who can ace the weather exam using just a single laptop, while the other students need a whole server farm. Here is how it works, broken down with simple analogies:

1. The "Two-Track" Brain (The StepsNet)

Most AI weather models try to learn everything at once: the wind, the temperature, the humidity, and the pressure. It's like trying to listen to a full orchestra while simultaneously solving a math problem. It's overwhelming and inefficient.

Sonny uses a Two-Track System:

The Slow Track (The Backbone): First, Sonny looks at the "big picture" drivers of the weather, like the wind and pressure fields. It treats these like the conductor of an orchestra. It figures out the general rhythm and structure of the weather system. Because it only focuses on these few key variables, it can be very deep and detailed without getting bogged down.
The Fast Track (The Details): Once the "conductor" has set the rhythm, Sonny brings in the rest of the musicians (temperature, humidity, etc.). It layers these details on top of the structure it already built.
The Analogy: Imagine building a house. The Slow Track is laying the foundation and framing the walls (the structure). The Fast Track is putting up the drywall, painting, and decorating (the details). By separating the heavy lifting from the finishing touches, Sonny saves massive amounts of energy.

2. The "Randomized Practice" (Randomized Dynamics)

Usually, weather models are trained to predict exactly 6 hours into the future, then 12, then 24. It's like a musician only practicing for exactly 5 minutes at a time.

Sonny is trained differently. It is told, "Predict the weather in 6 hours, or maybe 12, or maybe 24—just pick a random time!"

The Analogy: This is like a student studying for a test by practicing with random questions from the whole book, rather than just memorizing the answers to questions 1 through 10. This makes Sonny much more flexible and better at handling different timeframes without needing extra training.

3. The "Steady Hand" (EMA)

One of the biggest problems with AI weather models is that if you ask them to predict 10 days out, they often start to hallucinate or drift off course (like a drunk person walking in a straight line who eventually wanders into a bush).

To fix this, Sonny uses a technique called Exponential Moving Average (EMA).

The Analogy: Imagine you are learning to ride a bike. If you only look at your very last wobble, you might panic and fall. But if you look at your average balance over the last few seconds, you stay steady. Sonny keeps a "mental average" of its own weights during training. This smooths out the bumps and prevents the model from going crazy when predicting far into the future. It's like having a steady hand guiding the AI so it doesn't overreact.

4. The Results: Small but Mighty

The paper shows that Sonny can be trained on a single graphics card (an NVIDIA A40) in about 5.5 days.

The Comparison: Other top models might need 32 to 192 powerful cards and take weeks to train. Sonny does it on one card in less than a week.
The Performance: Even though it's small, Sonny is surprisingly good.
- It predicts the path of Typhoon Nanmadol almost as well as the massive super-computer models.
- It predicts the intensity of Winter Storm Elliott with high accuracy.
- It is particularly good at predicting weather in the Tropics (like hurricanes) over long periods, where other efficient models tend to fail.

Why Does This Matter?

Before Sonny, if you wanted to build a new, better weather model, you needed a bank account full of money and a warehouse full of computers. Sonny proves that you don't need a supercomputer to get great results; you just need a smart design.

It breaks down the "Compute Wall," allowing universities and smaller research groups to participate in the future of weather forecasting. It's like proving you can build a Ferrari engine in a garage, rather than needing a massive factory.

In short: Sonny is a clever, efficient weather forecaster that separates the "big picture" from the "details," practices with random timeframes, and keeps a steady hand to avoid mistakes—all while running on a single computer.

1. Problem Statement

Weather forecasting is critical for protecting lives and infrastructure, yet a significant barrier exists in the development of next-generation data-driven models. While recent deep learning models (CNNs, GNNs, Transformers) have matched or exceeded traditional Numerical Weather Prediction (NWP) systems, they typically rely on Large Weather Models (LWMs) requiring massive computational resources (e.g., hundreds of TPUs/GPUs) and extensive training times.

The Gap: This "compute wall" limits participation from academic groups with limited budgets and hinders the creation of reproducible, accessible baselines.
The Goal: Develop a Small Weather Model (SWM) that achieves competitive medium-range (up to 10 days) global forecast accuracy while being trainable on a single consumer-grade GPU within a reasonable timeframe.

2. Methodology: The Sonny Architecture

Sonny is an efficient hierarchical Transformer based on the StepsNet architecture, designed to balance physical fidelity with computational efficiency.

A. Data and Input

Dataset: WeatherBench2 (ERA5 reanalysis) with 6-hourly temporal resolution and 1.5° spatial resolution (121 × 240 grid).
Variables: 9 total variables including 4 surface variables (2m Temp, MSLP, Surface Air Pressure, 10m winds) and 5 atmospheric variables (Zonal/Meridional wind, Temp, Specific Humidity, Geopotential) across 13 pressure levels.
Split: 1979–2018 (Train), 2019 (Validation), 2022 (Test).

B. Training Objectives & Loss

Randomized Dynamics Forecasting: Instead of predicting absolute states or fixed intervals, Sonny predicts the change ( $\Delta_{\delta t}$ ) over a randomized time interval $\delta t$ sampled from $\{6, 12, 24\}$ hours. This improves generalization across different lead times.
Physics-Based Weighted Loss: The loss function incorporates:
1. Pressure Weighting ( $w(v)$ ): Assigns higher weights to variables near the Earth's surface (higher density) to prioritize societal impact.
2. Latitude Weighting ( $L(i)$ ): Corrects for the non-uniformity of the spherical grid.

C. Core Architecture: Two-Stage StepsNet

The model employs a Variable-Aware Embedding that splits inputs into two physically distinct groups:

Dynamics Group: Zonal/Meridional wind (U, V), Geopotential (Z), Pressure (P).
Thermodynamics Group: Temperature (T), Specific Humidity (Q).

The processing pipeline consists of two stages:

Step 1 (Slow Path): Processes the Dynamics group in a narrow hidden dimension ( $d_1$ ) using $N_1$ Transformer blocks. This stage captures large-scale atmospheric evolution and structural dependencies without the quadratic memory cost of processing all variables.
Step 2 (Fast Path): Concatenates the refined dynamic features from Step 1 with raw Thermodynamics embeddings ( $d_2$ ). The combined features are processed through $N_2$ full-width blocks ( $d = d_1 + d_2$ ). This stage models complex non-linear interactions (e.g., moisture advection, latent heat) guided by the pre-processed dynamic backbone.
Conditioning: A Randomized Dynamics Conditioning mechanism injects the time interval $\delta t$ into every Transformer block via adaptive Layer Normalization (adaLN-Zero), allowing flexible inference at various lead times.

D. Training Stability Strategy (EMA)

To avoid the computationally expensive fine-tuning phase common in other weather models, Sonny utilizes Exponential Moving Average (EMA) of model weights during the primary training stage (decay factor $\alpha = 0.9$ ). This smooths the optimization landscape, stabilizes long-term iterative rollouts, and eliminates the need for a separate fine-tuning stage.

3. Key Contributions

Sonny Model: An efficient SWM (20.5M parameters) capable of competitive medium-range forecasting under strict compute constraints.
Reproducible Training Recipe: A methodology using EMA and randomized dynamics that enables stable training and rollout without expensive infrastructure.
Accessibility Benchmark: Demonstrates that high-quality global forecasting is achievable on a single NVIDIA A40 GPU in approximately 5.5 days, democratizing access to advanced weather AI.

4. Experimental Results

Evaluated on WeatherBench2 against operational baselines (HRES, Met Office GM) and AI models (FastNet, Fengwu, Pangu-Weather).

Impact of EMA: Applying EMA reduced the Root Mean Square Error (RMSE) by an average of 3.34% across all variables and lead times, with a peak improvement of 4.75% on Day 3.
Competitive Performance:
- Extratropics: Sonny matches the accuracy of the Met Office GM and FastNet O96 for Z500 and MSLP.
- Tropics: Sonny significantly outperforms FastNet O96 in extended lead times (beyond Day 4) for Z500 and T850, maintaining high Anomaly Correlation Coefficient (ACC) scores where other models degrade rapidly.
Efficiency:
- Training: 5.5 days on 1x A40 (vs. weeks/months for LWMs).
- Inference: 125.48 ms latency (comparable to other efficient models).
- Compute: 96.81 G MACs (significantly lower than Fengwu or GraphCast).
Case Studies:
- Typhoon Nanmadol (2022): Accurately predicted track (137.5 km error) and intensity (977 hPa vs. 979 hPa reference), avoiding the over-smoothing typical of deep learning models.
- Winter Storm Elliott (2022): Successfully captured rapid intensification of a bomb cyclone at 1.5° resolution, with a central pressure error of only 4.9 hPa.

5. Significance and Future Work

Significance: Sonny proves that compute efficiency does not require a sacrifice in forecast skill. By leveraging physical priors (variable splitting) and training stability techniques (EMA), it provides a viable, reproducible baseline for the academic community, breaking the dependency on massive supercomputing clusters.
Future Directions:
- Leveraging randomized intervals for principled forecast uncertainty estimation.
- Integrating land-surface and oceanic variables for coupled modeling.
- Scaling to higher spatial resolutions and larger model sizes, given the favorable efficiency-accuracy trade-off demonstrated.

In summary, Sonny represents a paradigm shift toward accessible, high-fidelity AI weather forecasting, making state-of-the-art medium-range prediction feasible for researchers with limited hardware resources.