只处理四分之一的气象数据点,却让天气预测更准更快。
Less Tokens, Better Forecasts: Sparse Residual Routing for Efficient Weather Prediction

- 只让25%的网格点经过复杂计算,其余保持原样
- 在0.25度分辨率下训练快3.18倍,内存降低2.2倍以上
- 适合需要高精度网格输出的实时天气预报系统
现有基于ViT的天气预报模型对所有空间网格点均匀计算,但相邻网格点常具相似值,大区域随时间平滑演变,导致大量中间计算冗余。传统降本方法如剪枝或融合会丢弃或合并网格点,但天气预报是时空密集预测任务,必须保留原始经纬网格上每个格点的物理意义,尤其在自回归推演中。本文提出Sparse-Reslim——一种无参数的即插即用路由模块,使稀疏计算与固定网格要求兼容。该模块仅将25%的空间网格点送入昂贵的中间Transformer层,并将其视为残差更新:仅计算被选点的变化量,再将该增量散射回全序列。未被选中的网格点完全保留路由前表示,不丢失、不替换为掩码,也不引入融合层或额外参数。在最高达0.25°的ERA5分辨率及两种模型架构(确定性Transformer与扩散模型)下,Sparse-Reslim在所有评估变量上均提升预测精度,同时显著降低开销:主设置下训练速度提升约2.5倍,0.25°时达3.18倍,峰值内存下降超2.2倍。受控分解表明,准确率提升主要来自稀疏路由本身,而随机选点则提供额外正则化效果且无选择器开销。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。