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

Existing ViT-based weather forecasting models apply uni-form computation across all spatial tokens, even though nearby atmo-spheric grid points often contain similar values and large regions evolvesmoothly over time. This makes much of the intermediate per-token com-putation redundant. Standard token-efficiency methods, such as pruningor merging, reduce cost by removing or fusing tokens. However, weatherforecasting is a spatiotemporal dense prediction problem in which a his-tory of atmospheric states must be mapped to future values on the orig-inal latitude-longitude grid. Thus, every grid cell must retain a physi-cally meaningful representation, especially under autoregressive rollout.We introduce Sparse-Reslim, a parameter-free plug-in routing modulethat makes sparse token processing compatible with this fixed-grid re-quirement. Sparse-Reslim routes only 25% of spatial tokens through theexpensive middle transformer blocks and treats those blocks as residualupdates: it computes the change produced for the routed tokens and scat-ters only this delta back to the full sequence. Unselected tokens keep theirpre-routing representations exactly, so no grid cell is dropped or replacedby a mask token, and no fusion layer or additional parameters are intro-duced. Across ERA5 resolutions up to the operational 0.25° standard andtwo model families, a deterministic Transformer and a diffusion model,Sparse-Reslim improves forecast accuracy on every evaluated variablewhile substantially reducing cost: training is about 2.5× faster in themain settings and reaches 3.18× speedup at 0.25°, with over 2.2× lowerpeak memory. A controlled decomposition shows that the accuracy gaincomes primarily from sparse routing itself, while random token selectionprovides an additional regularization benefit without selector overhead.Code is available at https://github.com/janet-sw/Sparse-Reslim.