更少的令牌，更好的预测：用于高效天气预测的稀疏残差路由

# 更少的令牌，更准确的预测：面向高效天气预报的稀疏残差路由
来源：https://arxiv.org/html/2607.02829
11institutetext:杜兰大学 22institutetext:东北大学 33institutetext:阿拉巴马大学伯明翰分校 44institutetext:橡树岭国家实验室###### 摘要

现有的基于ViT的天气预报模型对所有空间令牌应用统一计算，尽管邻近的大气网格点通常包含相似的值，且大片区域随时间平滑演变。这使得中间逐令牌计算的大部分冗余。标准的令牌效率方法，如剪枝或合并，通过移除或融合令牌来降低成本。然而，天气预报是一个时空密集预测问题，必须将大气状态的历史映射到原始经纬度网格上的未来值。因此，每个网格单元必须保留物理上有意义的表示，尤其是在自回归滚动预测中。我们提出Sparse-Reslim，一个无参数的可插拔路由模块，使稀疏令牌处理与这种固定网格要求兼容。Sparse-Reslim仅将25%的空间令牌路由经过昂贵的中间 Transformer 块，并将这些块视为残差更新：它计算路由令牌产生的变化，并仅将此增量散射回完整的序列。未选中的令牌精确保留其路由前的表示，因此没有网格单元被丢弃或被掩码令牌替换，也没有引入融合层或额外参数。在ERA5分辨率高达业务化0.25°标准以及两个模型系列（确定性 Transformer 和扩散模型）上，Sparse-Reslim在每项评估变量上均提高了预测精度，同时大幅降低成本：主要设置下训练速度提升约2.5倍，在0.25°分辨率下达到3.18倍加速，峰值内存降低超过2.2倍。受控分解表明，精度提升主要来自稀疏路由本身，而随机令牌选择提供了额外的正则化收益，无需选择器开销。代码发布于https://github.com/janet-sw/Sparse-Reslim。

## 1 引言

数据驱动的天气预报已达到转折点[weyn2019can, scher2018toward, dueben2018challenges]。基于再分析数据训练的模型现在已匹配或超越经过数十年发展的业务化数值天气预报（NWP）系统的精度，且推理速度快数个数量级[bi2023accurate, lam2023learning, bodnar2025foundation]。这些模型在高维网格化大气场（0.25°分辨率的ERA5每个变量包含超过一百万个网格点）上运行，必须产生空间密集的输出，其中每个网格单元都带有物理上有意义的预测。随着社区向更高分辨率和更大模型容量推进，训练成本成为主要瓶颈，因为更精细的网格产生更多的空间令牌，而自注意力机制随令牌数量呈二次方扩展，同时更大的模型进一步增加了每个令牌的计算和激活内存。

在领先的天气预报模型（Pangu-Weather[bi2023accurate]、ClimaX[nguyen2023climax]、Stormer[nguyen2024scaling]、FuXi[chen2023fuxi]、Aurora[bodnar2025foundation]及其生成式对应模型GenCast[price2023gencast]和CorrDiff[mardani2025residual]）中，一个显著的共同点是它们都依赖于视觉变换器（ViT）或其变体。自注意力机制对所有空间令牌一视同仁，在平静开阔海洋上的令牌与在温带气旋核心的令牌上花费相同的O(L²D)计算。然而，大气场高度自相关：全球大片区域携带平滑、缓慢变化的条件，其中间表示很大程度上是冗余的。这表明，通过仅在最昂贵的层中处理代表性令牌子集，可以节省大量计算。

令牌剪枝[rao2021dynamicvit, liang2022not, kong2022spvit, xu2022evo]、令牌合并[bolya2022token]和基于窗口的注意力[liu2021swin]在视觉社区的ViT效率研究中已被广泛探索。最近，Sprint[park2025sprint]提出了一种用于扩散 Transformer 的密集-稀疏-密集块调度方案，仅将随机选择的令牌子集路由经过中间块，在图像生成中实现了巨大的训练加速。将这些方法直接应用于天气预报并非易事，因为天气预报是一个时空密集预测问题：模型基于过去的大气状态进行条件预测，并必须在每个网格点产生物理上有意义的未来值，通常用于自回归滚动预测。令牌剪枝和合并可能移除或改变空间位置，而Sprint用学习到的掩码嵌入替换跳过的令牌，随后通过线性投影进行融合。这对于静态图像生成是可接受的，但在天气预报中会产生不匹配：后续块操作的是真实大气表示和固定网格位置上合成占位符令牌的混合，而未明确保留预测问题的时间条件结构。

在这项工作中，我们提出Sparse-Reslim，一种用于基于ViT的天气预报器的无参数可插拔路由模块，使稀疏令牌处理与密集预测兼容。Sparse-Reslim将 Transformer 块划分为密集早期、稀疏中期和密集晚期三个阶段。在稀疏中期阶段的入口，它随机选择L个空间令牌中的一部分r，得到K = ⌊rL⌋个路由令牌，并仅通过昂贵的中间块处理这个子集。关键的是，Sparse-Reslim不对完整序列用处理后的令牌和占位符进行替换，而是计算稀疏中间块产生的*残差增量*，并仅将此增量散射回原始令牌位置。未选中的令牌接收零增量，因此精确保留其路由前的表示，确保所有L个空间令牌位置在整个模型中保持有效。该模块不需要掩码嵌入、融合层或额外的可学习参数，并可在现有ViT块循环中通过单个配置标志启用。

我们在多个分辨率下的ERA5[hersbach2020era5]以及两个架构系列上验证Sparse-Reslim：（1）从ORBIT-2百亿亿次系统[wang2025orbit]改编的确定性Res-Slim-ViT，以及（2）遵循EDM框架[karras2022elucidating]的基于扩散的生成模型。对于生成式变体，我们采用非对称注意力设计，其中自注意力被稀疏化，而交叉注意力保留完整的条件序列。在所有设置中，保持比率r=0.25产生约2.5倍的训练加速。值得注意的是，Sparse-Reslim在所有评估变量和指标上也提高了预测精度，优于密集基线。我们的消融实验表明，这种提升并非仅由随机正则化解释：确定性稀疏路由本身已捕获大部分改进，而随机令牌选择在不引入选择器参数或评分开销的情况下提供了额外的较小收益。我们的主要贡献是：

- •我们提出了Sparse-Reslim，一种用于时空天气预报的稀疏路由模块，它保留所有空间令牌位置，同时尊重时间条件。它不添加任何可学习组件，并可作为即插即用模块集成到基于ViT的天气预报器中。
- •我们提供了涵盖两个架构系列（确定性和生成式）、三种ERA5网格分辨率以及四个大气变量的实证研究。Sparse-Reslim在评估设置中持续提高精度，且其效益在更高分辨率下更为显著。
- •我们表明，稀疏路由与其他效率策略（包括混合精度训练和分辨率下采样）正交，并可与之结合以进一步获益。据我们所知，这是稀疏令牌路由首次应用于基于ViT的天气预报模型的训练效率提升，且同时涵盖确定性和生成式设置。

参考图注图1：Sparse-Reslim概览。(a) 三阶段令牌路由：密集早期（n1块）、稀疏中期（n2块，在K=⌊rL⌋个令牌上）和密集晚期（n3块），未选中的令牌走恒等路径。(b, c) 集成到确定性Res-Slim-ViT和生成式Res-Slim-EDM骨干网络中。
## 2 相关工作

### 2.1 AI天气预报模型

数据驱动的天气预报已成为传统NWP的引人注目的替代方案，以极低的计算成本达到可比或更好的精度。任务是从初始条件X₀预测未来大气状态X_T ∈ R^{V×H×W}，其中V表示物理变量，H×W表示空间网格，T表示目标提前时间。在确定性方法中，Pangu-Weather[bi2023accurate]引入了在气压层面体积上运行的三维地球特定 Transformer，而GraphCast[lam2023learning]采用基于网格的图神经网络在球面上进行消息传递。ClimaX[nguyen2023climax]提出了一种变量无关的ViT，在异构气候数据集上训练，Stormer[nguyen2024scaling]结合了移位窗口注意力与天气特定的归纳偏置。FengWu[chen2023fengwu]使用多模态、多任务跨模态融合 Transformer，将技巧性预测推进到超过10天提前时间，FuXi[chen2023fuxi]和Aurora[bodnar2025foundation]分别通过级联和基础模型设计扩展了前沿。

在概率方面，GenCast[price2023gencast]应用扩散模型生成校准的集合预报，而OmniCast[nguyen2025omnicast]和CorrDiff[mardani2025residual]探索了生成式方法来建模未来天气状态的完整分布，以解决影响确定性自回归滚动预测的模糊和误差累积问题。尽管架构多样，所有这些模型共享一个特征：它们对所有空间令牌应用统一计算，无论局部信息复杂度如何。这在更高分辨率下变得越来越浪费，因为二次注意力成本主导训练时间，而大部分网格是平滑且冗余的。作为ORBIT-2百亿亿次气候降尺度系统[wang2025orbit]的一部分引入的Res-Slim-ViT，通过结合ViT骨干网络与卷积残差路径的双路径设计展现了强可扩展性，使其成为高效令牌路由的自然测试平台。

### 2.2 高效视觉变换器

自注意力的二次成本激发了对降低ViT计算负担的大量研究，减少处理的令牌数量是最常研究的策略之一。基于注意力的令牌剪枝移除注意力得分低的令牌[liang2022not, kong2022spvit, xu2022evo, fayyaz2022adaptive]。几种方法进一步训练评分模块，以学习任务相关的剪枝标准[rao2021dynamicvit, rao2023dynamic]。令牌合并提供了互补方向：ToMe[bolya2022token]通过二分图匹配压缩相似令牌，减少序列长度同时保留表征内容。其他方法在注意力层面本身操作，包括基于窗口的注意力[liu2021swin]和学习稀疏注意力模式[wei2023sparsifiner, chen2021chasing]。

最近一系列工作专门针对生成模型训练。Sprint[park2025sprint]引入了用于扩散 Transformer 的密集-稀疏-密集块调度：只有一小部分随机选择的令牌通过中间块，而丢弃的令牌被可学习的掩码嵌入替换，并随后通过线性投影与浅层密集特征融合。虽然对于图像生成有效，但Sprint并不直接适用于天气预报。天气预报模型基于过去大气状态条件，并必须在原始网格上返回完整的未来状态。在自回归使用中，该输出成为下一个输入。因此，用掩码嵌入替换未处理的位置不仅仅是增加一个小投影层：它破坏了固定网格位置上大气状态表示的连续性。我们的残差增量公式通过保持未选中令牌为其实际的路由前表示，并仅将路由令牌计算的更新散射回去，解决了这一差距。在生成式设置中，我们通过稀疏化目标令牌上的自注意力，同时保留对条件序列的完整交叉注意力，进一步保持了时间条件。

## 3 方法

本节描述Sparse-Reslim，一种用于基于ViT的天气模型的即插即用稀疏令牌路由模块。图1 (https://arxiv.org/html/2607.02829#S1.F1)提供了Sparse-Reslim的视觉概览。我们首先介绍共享的令牌化方法和基础架构（第3.1节 (https://arxiv.org/html/2607.02829#S3.SS1)），然后介绍路由模块（第3.2节 (https://arxiv.org/html/2607.02829#S3.SS2)），及其在确定性（第3.3节 (https://arxiv.org/html/2607.02829#S3.SS3)）和生成式（第3.4节 (https://arxiv.org/html/2607.02829#S3.SS4)）预测管线中的集成。

### 3.1 预备知识

令牌化。两种架构共享一个每变量令牌化方案。每个变量v由专用卷积嵌入 \mathtt{PatchEmbed}_v 独立分块，产生维度为D的令牌。添加可学习的变量嵌入 e_v，并通过与共享查询 q ∈ R^D 的交叉注意力聚合变量：

z_i = CrossAttn( q, { p_{i,v} + e_v }_{v∈V} ) ∈ R^D, (1)

其中 p_{i,v} 是空间位置i处变量v的块令牌。这产生序列 Z ∈ R^{L×D}，其中 L = (h/p)(w/p) 个令牌，并附加正弦位置和空间分辨率嵌入。

确定性预测。确定性模型 f_θ 直接回归预测目标。给定输入历史 X，将历史维度展平到通道维度，模型预测：

Ŷ = f_θ(X) = CNN_skip(X) + Unpatchify( MLP_head( ViT(Z) ) ), (2)

其中 ViT(⋅) 表示 N 个 Transformer 块，CNN_skip 是绕过 Transformer 的卷积残差路径（Conv → PixelShuffle → Conv）。双路径设计让轻量级 CNN 捕捉密集局部模式，使 Transformer 专注于长距离依赖。训练最小化纬度加权均方误差：

L_det = (1/|V_out|) ∑_v ω_v · (∑_i cos(φ_i) (Ŷ_{v,i} - Y_{v,i})²) / (∑_i cos(φ_i)), (3)

其中 ω_v 是每变量权重，φ_i 是网格