🤖 machine learning

Less Tokens, Better Forecasts: Sparse Residual Routing for Efficient Weather Prediction

The paper introduces Sparse-Reslim, a parameter-free routing module that improves the efficiency and accuracy of ViT-based weather forecasting by processing only a sparse subset of tokens through expensive transformer blocks while preserving full-grid representations via residual delta updates, thereby achieving significant speedups and memory reductions without sacrificing forecast quality.

Original authors: Janet Wang, Yunbei Zhang, Lin Zhao, Xi Xiao, Jihun Hamm, Xiao Wang

Published 2026-07-07

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Janet Wang, Yunbei Zhang, Lin Zhao, Xi Xiao, Jihun Hamm, Xiao Wang

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 entire globe. To do this, computers break the Earth's atmosphere into a giant grid of millions of tiny squares (like a massive pixelated map). For every single square, the computer has to calculate how the wind, temperature, and pressure will change.

The Problem: The "Busy Bee" Bottleneck
Current weather models act like a very diligent, but slightly inefficient, bee. They visit every single flower (grid square) in the garden, no matter how boring or empty it is.

Some areas, like the calm center of a hurricane or a quiet ocean, don't change much.
Other areas, like the edge of a storm, are chaotic and change rapidly.

Yet, the computer spends the exact same amount of time and energy calculating the "boring" calm spots as it does the "exciting" storm spots. This is a waste of energy.

The Old Solutions: The "Cut and Paste" Mistake
Scientists tried to fix this by telling the computer to ignore the boring spots (pruning) or to combine them into one big blob (merging).

The Analogy: Imagine you are writing a story about a city. To save time, you decide to skip writing about the quiet suburbs and just write "suburbs exist" in a blank space.
The Issue: In weather forecasting, you can't just skip a spot. The weather in one spot affects the spot next to it. If you delete a spot or replace it with a "placeholder," the story breaks. The computer gets confused because the map no longer has a value for every single square, which ruins the prediction for the next day.

The New Solution: Sparse-Reslim (The "Delta Delivery" System)
The authors of this paper created a new method called Sparse-Reslim. Think of it as a smart delivery service that changes how the work gets done, without changing what gets delivered.

Here is how it works, using a simple analogy:

The Three-Stage Factory: Imagine the computer's brain is a factory with three rooms:
- Room 1 (Dense Early): Every single worker (grid square) gets a full briefing. Everyone is involved.
- Room 2 (Sparse Middle): This is the expensive, heavy-lifting room. Instead of sending all 10,000 workers in here, the manager randomly picks only 25% of them (the "selected tokens").
- Room 3 (Dense Late): Everyone comes back together for the final wrap-up.
The Magic Trick (The Residual Delta):
- In the middle room, only the 25% selected workers do the heavy calculation. They figure out how much the weather should change.
- The other 75% of workers sit in the hallway. They don't do any work, but they keep their original briefing exactly as it was.
- The Delivery: The 25% workers finish their math and hand over a small note called a "Delta" (which just means "the change").
- The computer takes this note and pastes the change back onto the specific spots where the workers were.
- Crucially: The 75% who sat in the hallway didn't get a "fake" note or a "blank" note. They just kept their original briefing. Nothing was deleted, nothing was replaced with a mask. The map remains 100% complete.

Why This is a Big Deal

Speed: Because the computer only does the hard math on 25% of the grid in the middle section, it runs 2.5 times faster. At the highest resolution, it's nearly 3 times faster.
Memory: It uses less than half the computer memory (VRAM) because it doesn't have to hold the heavy calculations for every single spot at once.
Better Accuracy: Surprisingly, the weather predictions are more accurate than the old, slow method.
- Why? The paper suggests that randomly picking which workers to send to the middle room acts like a "stochastic regularizer" (a fancy way of saying it prevents the computer from getting too focused on one specific pattern and helps it learn better). It's like a teacher randomly calling on students to answer questions; it keeps the whole class more alert and learning better than if the teacher only asked the same smart kids every time.

The Results
The team tested this on two types of weather models (one that predicts a single future, and one that generates many possible futures). In both cases:

The models trained much faster.
They used less memory.
They predicted wind, temperature, and pressure more accurately than the standard models.

In Summary
Sparse-Reslim is like a smart manager who realizes that not every part of the world needs a full-time team of experts working on it every second. By letting the quiet areas "rest" while the experts work on the active areas, and then simply adding the experts' findings back to the map, the system becomes faster, cheaper, and surprisingly, smarter.

Technical Summary: Sparse Residual Routing for Efficient Weather Prediction

Problem Statement

Data-driven weather forecasting has achieved significant success, with models like Pangu-Weather, ClimaX, and FuXi matching or exceeding traditional Numerical Weather Prediction (NWP) accuracy at a fraction of the inference cost. However, training these models remains a primary bottleneck. Most state-of-the-art models utilize Vision Transformers (ViTs), which apply uniform self-attention computation across all spatial tokens. This approach is inefficient because atmospheric fields exhibit high spatial autocorrelation; large regions often contain smooth, slowly varying conditions where intermediate token representations are redundant.

Existing token-efficiency methods, such as token pruning, merging, or random selection (e.g., Sprint for image generation), are ill-suited for weather forecasting. Weather prediction is a spatiotemporal dense prediction problem requiring a physically meaningful output for every grid cell on the original latitude-longitude grid, especially during autoregressive rollout. Standard methods that remove tokens, fuse them, or replace them with learned mask embeddings disrupt the fixed-grid structure and the temporal conditioning required for accurate forecasting.

Methodology: Sparse-Reslim

The authors propose Sparse-Reslim, a parameter-free, plug-in routing module designed to enable sparse token processing while preserving the dense grid requirement of weather forecasting.

Core Mechanism

Sparse-Reslim partitions the transformer blocks into three stages:

Dense Early Stage: The first $n_1$ blocks process all $L$ spatial tokens to build rich initial representations.
Sparse Middle Stage: The middle $n_2$ blocks process only a randomly selected subset of $K = \lfloor rL \rfloor$ tokens (where $r$ is the keep ratio, typically 0.25).
Dense Late Stage: The final $n_3$ blocks process the full sequence of $L$ tokens to refine spatial coherence before the prediction head.

Residual Delta Formulation

The key innovation lies in how the sparse stage interacts with the full sequence. Instead of replacing unselected tokens with placeholders or mask embeddings, Sparse-Reslim computes a residual delta:

The selected tokens pass through the sparse middle blocks.
The output of these blocks is compared to their input (pre-routing representation) to compute the change ( $\Delta$ ).
This $\Delta$ is scattered back to the original token positions.
Unselected tokens receive a zero delta, effectively retaining their pre-routing representations exactly via an implicit identity skip connection.

This approach ensures that no grid cell is dropped or replaced by a synthetic token, maintaining the physical integrity of the atmospheric state representation throughout the network.

Integration with Architectures

Deterministic Models (Res-Slim-ViT): The module integrates directly into the standard ViT block loop. The sparse middle blocks operate with reduced self-attention complexity $O(K^2D)$ instead of $O(L^2D)$ .
Generative Models (EDM/Diffusion): The authors employ an asymmetric attention design. While self-attention is sparsified over the target tokens ( $O(K^2D)$ ), cross-attention retains the full conditioning sequence ( $Z_c$ ) as keys and values ( $O(KL \cdot D)$ ). This ensures the model retains the full atmospheric context (history frames) necessary for denoising, even when the target sequence is sparsified.

Token Selection Strategy

The method utilizes random token selection rather than attention-score-based or learned selection. The authors argue this provides three benefits:

Zero scoring overhead.
Acts as a stochastic regularizer (analogous to dropout on the token dimension).
Avoids systematic spatial bias that might underrepresent quiescent regions (e.g., tropical areas) in favor of dynamically active regions.

Key Contributions

Sparse-Reslim Module: A parameter-free routing mechanism that enables sparse processing in ViT-based weather models without violating the dense grid constraint, using a residual delta formulation to preserve unselected token representations.
Empirical Validation: Comprehensive evaluation across two architecture families (deterministic Res-Slim-ViT and generative EDM) and three ERA5 resolutions (up to the operational 0.25° standard).
Performance Gains: Demonstration that sparse routing not only reduces computational cost but also improves forecast accuracy compared to dense baselines.
Orthogonality: Showing that the method is compatible with mixed-precision training and resolution downsampling.

Experimental Results

The authors evaluated Sparse-Reslim on the ERA5 reanalysis dataset using a deterministic Res-Slim-ViT and a diffusion-based generative model.

Efficiency

Training Speed: Sparse-Reslim achieves approximately 2.5× faster training in main settings, reaching a 3.18× speedup at the 0.25° resolution.
Memory: Peak VRAM usage is reduced by over 2.2× across all tested settings.
Inference: Inference time is reduced by roughly 1.8×.

Accuracy

Contrary to typical efficiency methods that trade accuracy for speed, Sparse-Reslim improves forecast accuracy across all evaluated variables (Z500, T850, T2m, U10) and resolutions.

At 1.0° resolution, Z500 RMSE decreased from 868.7 (dense) to 768.3 (Sparse-Reslim), and ACC increased from 0.551 to 0.651.
Similar improvements were observed for T2m, U10, and T850.
Ablation Studies: The accuracy gain is attributed to a combination of reduced redundant computation and stochastic regularization. Deterministic sparse routing captures most of the gain, while random selection provides an additional, consistent benefit without extra parameters.
Stability: The model remains stable during 10-day autoregressive rollouts, with skill decaying smoothly rather than compounding errors.

Scaling

The benefits of Sparse-Reslim scale with model size. Larger models (e.g., 48.6M parameters) achieve higher speedups (2.85×) and greater accuracy improvements (12.6% RMSE reduction) compared to smaller models, as deeper networks have a larger fraction of compute in the sparse middle blocks.

Significance and Claims

The paper claims that Sparse-Reslim is the first application of sparse token routing to the efficient training of ViT-based weather models in both deterministic and generative settings. Its primary significance lies in resolving the conflict between token efficiency and the dense prediction requirements of weather forecasting. By using a residual delta formulation, it avoids the "mismatch" of synthetic placeholder tokens found in other methods (like Sprint), allowing the model to maintain physically meaningful representations at every grid point.

The authors modestly note that while the method is highly effective, future work could explore adaptive token selection to dynamically allocate computation to meteorologically active regions (e.g., storm fronts) rather than relying on random selection. However, the current work establishes that random selection, combined with the residual routing architecture, is sufficient to achieve both significant speedups and accuracy improvements.