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

2026-07-02

AI Summary

This work addresses the computational redundancy in existing Vision Transformer (ViT)-based weather forecasting models, which uniformly process all spatial tokens despite much of the atmospheric field exhibiting sparsity or smoothness. Conventional token compression techniques often compromise the physical consistency of fixed latitude–longitude grids. To overcome this, the authors propose Sparse-Reslim, a parameter-free module that retains the full-grid representation while selectively applying intermediate Transformer computations to only 25% of spatial tokens, randomly sampled at each layer. These sparse activations are then scattered back into the full sequence as residual updates, preserving grid structure without token merging or discarding. The stochastic selection also provides implicit regularization. Evaluated on multi-resolution ERA5 data, Sparse-Reslim consistently improves forecast accuracy, accelerates training by up to 3.18×, and reduces peak memory consumption by over 2.2×.

Abstract

Existing ViT-based weather forecasting models apply uniform computation across all spatial tokens, even though nearby atmospheric grid points often contain similar values and large regions evolve smoothly over time. This makes much of the intermediate per-token computation redundant. Standard token-efficiency methods, such as pruning or merging, reduce cost by removing or fusing tokens. However, weather forecasting is a spatiotemporal dense prediction problem in which a history of atmospheric states must be mapped to future values on the original latitude-longitude grid. Thus, every grid cell must retain a physically meaningful representation, especially under autoregressive rollout. We introduce Sparse-Reslim, a parameter-free plug-in routing module that makes sparse token processing compatible with this fixed-grid requirement. Sparse-Reslim routes only 25% of spatial tokens through the expensive middle transformer blocks and treats those blocks as residual updates: it computes the change produced for the routed tokens and scatters only this delta back to the full sequence. Unselected tokens keep their pre-routing representations exactly, so no grid cell is dropped or replaced by a mask token, and no fusion layer or additional parameters are introduced. Across ERA5 resolutions up to the operational 0.25\textdegree{} standard and two model families, a deterministic Transformer and a diffusion model, Sparse-Reslim improves forecast accuracy on every evaluated variable while substantially reducing cost: training is about 2.5x faster in the main settings and reaches 3.18x speedup at 0.25\textdegree{}, with over 2.2x lower peak memory. A controlled decomposition shows that the accuracy gain comes primarily from sparse routing itself, while random token selection provides an additional regularization benefit without selector overhead.

Problem

• weather forecasting
• token efficiency
• spatiotemporal prediction
• dense prediction
• redundant computation

Innovation

• Sparse Residual Routing
• Token Efficiency
• Weather Forecasting
• Vision Transformer
• Dense Spatiotemporal Prediction