🤖 machine learning

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

该论文介绍了 Sparse-Reslim，一种无参数的路由模块，它通过仅让稀疏的 token 子集通过昂贵的 Transformer 块，同时通过残差增量更新保留全网格表示，从而在不牺牲预报质量的前提下，提高了基于 ViT 的天气预报的效率和准确性，并实现了显著的加速和内存减少。

原作者： Janet Wang, Yunbei Zhang, Lin Zhao, Xi Xiao, Jihun Hamm, Xiao Wang

发布于 2026-07-07

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Janet Wang, Yunbei Zhang, Lin Zhao, Xi Xiao, Jihun Hamm, Xiao Wang

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下，你正试图预测全球的天气。为了实现这一目标，计算机将地球的大气层分解成一个巨大的网格，由数百万个微小的方块组成（就像一张巨大的像素化地图）。对于每一个方块，计算机都必须计算风速、温度和气压将如何变化。

问题所在：“忙碌的小蜜蜂”瓶颈
目前的天气模型表现得像一只勤奋但效率略低的蜜蜂。它会访问花园里的每一朵花（网格方块），无论这些地方多么无聊或空旷。

有些区域，比如飓风中心或平静的海洋，变化并不大。
而其他区域，比如风暴的边缘，则非常混乱且变化剧烈。

然而，计算机在处理那些“无聊”的平静区域时所花费的时间和精力，与处理那些“兴奋”的风暴区域时是完全一样的。这是一种能量的浪费。

旧的解决方案：“剪切与粘贴”的错误
科学家们曾尝试通过让计算机忽略无聊的区域（剪枝）或将它们合并成一个大斑块（合并）来解决这个问题。

类比： 想象你正在写一个关于城市的故事。为了节省时间，你决定跳过对安静郊区的描写，只是在空白处写上“郊区存在”。
问题： 在天气预报中，你不能仅仅跳过某个地点。一个地点的情况会影响到相邻的地点。如果你删除了一个点或者用一个“占位符”替换了它，故事就会断掉。计算机会感到困惑，因为它在地图上不再拥有每个方块的数值，这会破坏对第二天的预测。

新的解决方案：Sparse-Reslim（“增量交付”系统）
该论文的作者创造了一种名为 Sparse-Reslim 的新方法。把它想象成一种聪明的快递服务，它改变的是“如何完成工作”，而不是改变“交付的内容”。

它是这样运作的，使用一个简单的类比：

三阶段工厂： 想象计算机的大脑是一个拥有三个房间的工厂：
- 第一间房（稠密早期）： 每一位工作人员（网格方块）都会得到一份完整的简报。每个人都参与其中。
- 第二间房（稀疏中期）： 这是昂贵的、重体力活的房间。经理并没有把所有 10,000 名工人送进去，而是随机挑选了仅 25% 的人（即“被选中的标记/tokens”）。
- 第三间房（稠密后期）： 所有人重新聚在一起进行最后的总结。
神奇的技巧（残差增量/Residual Delta）：
- 在中间的房间里，只有那 25% 被选中的工人进行繁重的计算。他们计算出天气应该如何变化。
- 其余 75% 的工人则坐在走廊里。他们不干活，但他们完整地保留了最初的简报，原封不动。
- 交付： 这 25% 的工人完成了数学计算，并递交了一张被称为“增量”（Delta，意为“变化”）的小纸条。
- 计算机将这张纸条上的“变化”重新粘贴回那些工人所在的特定位置。
- 至关重要的一点是： 那 75% 坐在走廊里的工人并没有得到一张“假”纸条或“空白”纸条。他们只是保留了原始的简报。地图仍然是 100% 完整的，没有任何东西被删除，也没有任何东西被替换成掩码。

为什么这意义重大

速度： 因为计算机在中间部分只对 25% 的网格进行复杂的数学运算，所以它的运行速度快了 2.5 倍。在最高分辨率下，它几乎快了 3 倍。
内存： 它使用的计算机内存（VRAM）不到原来的一半，因为它不需要同时为每一个点都持有沉重的计算数据。
更好的准确性： 出人意料的是，天气预测比旧的、缓慢的方法更准确。
- 原因： 论文指出，随机选择哪些工人进入中间房间起到了“随机正则化器”（stochastic regularizer）的作用（这是一种高级说法，指防止计算机过度专注于某种特定的模式，从而帮助它更好地学习）。这就像一位老师随机点名学生回答问题；比起只询问那些固定的优等生，这种方式能让全班同学都保持警觉并学得更好。

结果
团队在两种类型的天气模型（一种预测单一未来，另一种生成多种可能的未来）上进行了测试。在这两种情况下：

模型的训练速度大大加快。
它们使用了更少的内存。
它们对风速、温度和气压的预测比标准模型更加准确。

总结
Sparse-Reslim 就像一位聪明的管理者，他意识到并非世界的每一个部分都需要一支专家团队在每一秒钟都全力投入工作。通过让安静的区域“休息”，同时让专家们处理活跃区域，并将专家的发现重新添加回地图，该系统变得更快、更便宜，而且出人意料地更聪明。

技术摘要：用于高效天气预报的稀疏残差路由技术

问题陈述

数据驱动的天气预报已取得显著成功，诸如 Pangu-Weather、ClimaX 和 FuXi 等模型在匹配或超越传统数值天气预报（NWP）准确度的同时，大幅降低了推理成本。然而，训练这些模型仍然是一个主要的瓶颈。大多数最先进的模型使用视觉 Transformer（ViT），它们在所有空间标记（tokens）上应用统一的自注意力计算。这种方法是低效的，因为大气场具有高度的空间自相关性；大范围区域通常包含平滑且变化缓慢的状态，其中中间标记的表示是冗余的。

现有的标记效率方法，如标记剪枝、合并或随机选择（例如用于图像生成的 Sprint），并不适用于天气预报。天气预报是一个时空密集预测问题，要求在原始经纬度网格上的每个网格单元都必须有一个具有物理意义的输出，尤其是在自回归展开过程中。标准方法通过移除标记、融合它们或用学习到的掩码嵌入替换它们，会破坏固定的网格结构以及对准确预报至关重要的时序条件约束。

方法论：Sparse-Reslim

作者提出了 Sparse-Reslim，这是一种无参数、插件式的路由模块，旨在实现稀疏标记处理，同时保留天气预报对密集网格的要求。

核心机制

Sparse-Reslim 将 Transformer 块划分为三个阶段：

密集早期阶段（Dense Early Stage）： 前 $n_1$ 个块处理所有 $L$ 个空间标记，以构建丰富的初始表示。
稀疏中期阶段（Sparse Middle Stage）： 中间的 $n_2$ 个块仅处理随机选择的 $K = \lfloor rL \rfloor$ 个标记子集（其中 $r$ 是保留比例，通常为 0.25）。
密集后期阶段（Dense Late Stage）： 最后的 $n_3$ 个块处理完整的 $L$ 个标记序列，以在进入预测头之前优化空间相干性。

残差增量公式化（Residual Delta Formulation）

该研究的核心创新在于稀疏阶段如何与完整序列进行交互。Sparse-Reslim 并没有使用占位符或掩码嵌入来替换未被选中的标记，而是计算了一个残差增量（residual delta）：

选中的标记通过稀疏中间块进行处理。
将这些块的输出与其输入（路由前的表示）进行比较，以计算变化量（ $\Delta$ ）。
该 $\Delta$ 被散射回原始标记的位置。
未被选中的标记接收一个零增量，从而通过隐式恒等跳跃连接（identity skip connection）精确地保留其路由前的表示。

这种方法确保了没有网格单元被丢弃或被合成标记替换，从而在整个网络中维持了大气状态表示的物理完整性。

与架构的集成

确定性模型 (Res-Slim-ViT)： 该模块直接集成到标准的 ViT 块循环中。稀疏中间块以降低后的自注意力复杂度 $O(K^2D)$ 运行，而非 $O(L^2D)$ 。
生成式模型 (EDM/Diffusion)： 作者采用了非对称注意力设计。虽然自注意力在目标标记上是稀疏化的（ $O(K^2D)$ ），但交叉注意力保留了完整的条件序列 ( $Z_c$ ) 作为键（Keys）和值（Values）（ $O(KL \cdot D)$ ）。这确保了即使在目标序列被稀疏化时，模型仍能保留进行去噪所需的完整大气背景信息（历史帧）。

标记选择策略

该方法利用随机标记选择，而非基于注意力分数或学习到的选择。作者认为这提供了三个好处：

零评分开销。
经过类似于在标记维度上的 Dropout 作用，起到随机正则化的作用。
避免了系统性的空间偏差，即不会为了偏向动态活跃区域而系统性地低估静止区域（如热带地区）。

核心贡献

Sparse-Reslim 模块： 一种无参数的路由机制，能够在不违反天气预报密集网格约束的情况下，实现 ViT 天气模型的稀疏处理，并利用残差增量公式化来保留未选中的标记表示。
实证验证： 在两种架构家族（确定性 Res-Slim-ViT 和生成式 EDM）以及三种 ERA5 分辨率（最高达 0.25° 标准分辨率）上进行了全面的评估。
性能提升： 证明了稀疏路由不仅能降低计算成本，还能提高预报准确度（相比于密集基准模型）。
একই 正交性： 展示了该方法与混合精度训练及分辨率下采样具有兼容性。

实验结果

作者在 ERA5 再分析数据集上使用确定性 Res-Slim-ViT 和基于扩散的生成模型对 Sparse-Reslim 进行了评估。

效率

训练速度： 在主要设置下，Sparse-Reslim 实现了约 2.5 倍的训练加速，在 0.25° 分辨率下达到了 3.18 倍的加速。
显存： 在所有测试设置中，峰值显存占用降低了超过 2.2 倍。
推理： 推理时间减少了约 1.8 倍。

准确度

与典型的以牺牲准确度换取速度的效率方法不同，Sparse-Reslim 在所有评估的变量（Z500, T850, T2m, U10）和分辨率下都提高了预报准确度。

在 1.0° 分辨率下，Z500 的 RMSE 从 868.7（密集型）降至 768.3（Sparse-Reslim），ACC 从 0.551 升至 0.651。
在 T2m、U10 和 T850 上也观察到了类似的改进。
消融实验： 准确度的提升归功于减少冗余计算和随机正则化的结合。确定性稀疏路由捕获了大部分增益，而随机选择则在不增加额外参数的情况下提供了持续的额外收益。
稳定性： 模型在 10 天自回归展开过程中保持稳定，其技能衰减过程平滑，而非误差复合。

扩展性

Sparse-Reslim 的收益随模型规模扩大而增加。较大的模型（例如 48.6M 参数）比较小的模型获得更高的加速比（2.85×）和更大的准确度提升（RMSE 降低 12.6%），因为更深的网络有更大比例的计算量分布在稀疏中间块中。

重要性与主张

本文声称 Sparse-Reslim 是首次将稀疏标记路由应用于提高基于 ViT 的天气模型（包括确定性和生成式设置）的训练效率。其核心意义在于解决了效率与天气预报密集预测需求之间的冲突。通过使用残差增量公式，它避免了其他方法（如 Sprint）中出现的合成占位标记所导致的“失配”问题，从而使模型能够在每个网格点上维持具有物理意义的表示。

作者谦虚地指出，尽管该方法非常有效，但未来的工作可以探索自适应标记选择，以动态地将计算分配给气象活跃区域（如风暴锋面），而非仅仅依赖随机选择。然而，目前的工作已经证明，随机选择结合残差路由架构，足以同时实现显著的加速和准确度提升。