UtoMe: Observation-Uncertainty-Guided Token Merging for Weather Foundation Models

会议: NeurIPS 2026
arXiv: 2610.25001
代码: https://ornl-utome.github.io/
领域: 模型压缩
关键词: 气象基础模型、Token Merging、观测不确定性、集合数据同化、高效大模型

一句话总结

针对传统 Token Merging 仅依据隐层特征相似度容易受再分析输入观测噪声误导的问题，UtoMe 联合大气动力学显著性与集合数据同化离散度构建重要性评分，通过保护高分锚点和引入不确定性校准边惩罚实现高效气象 Transformer 压缩。

研究背景与动机

基于 Transformer 的数据驱动气象大模型（如 ClimaX、FourCastNet、Pangu-Weather 等）正在显著改变全球中期天气预报的效率与精度，但其计算开销依然极其庞大。在海量网格点输入下，注意力机制与前馈网络的反复计算构成了巨大的训练和推理瓶颈。简单的增大 Patch 尺寸属于静态下采样，会无差别地抹平所有空间细节；自适应序列缩减技术（如 Token Pruning 和基于二分图匹配的 Token Merging, ToMe）由此成为降低气象模型计算复杂度的关键方向。

然而，将计算机视觉中成熟的 ToMe 直接迁移至气象预报存在本质矛盾：气象模型的再分析输入（如 ERA5）并非无误差真值，而是融合异构观测与动力同化系统得出的隐状态估计，其观测不确定性具有强烈的非均匀空间结构。若仅依赖隐层特征的余弦相似度进行合并，算法无法甄别相似性究竟源于大气的真实物理冗余，还是由输入端的观测噪声和同化扰动所致，极易将高动态且对预报敏感的关键区域错误合并。

解决该矛盾的核心在于引入输入可靠性作为算力分配的物理引导信号。但若单纯剔除高不确定性区域，又会与急流、锋面等往往兼具高动态与高不确定性的天气系统产生冲突，必须将动力学显著性与输入可靠性联合建模。核心 idea：将大气物理动力学显著性与集合数据同化离散度联合构建 Patch 级重要性评分，前 25% 高分区域作为关键锚点严格免除合并，其余候选 Token 则通过不确定性校准的二分图匹配进行相似聚合。

方法详解

整体框架

UtoMe 的整体工作流涵盖输入特征提取、联合评分计算、锚点掩蔽保护以及不确定性校准二分图合并四个阶段。输入多通道大气状态网格后，系统一方面通过差分与异常振幅提取动力学梯度，另一方面对 ECMWF 集合数据同化（EDA）离散度做归一化与 Patch 均值池化，融合成标定输入可靠性的综合重要性得分。随后，系统冻结最高分锚点以保留关键物理结构，再对剩余 Token 施加重要性总和与差值惩罚完成相似合并，并在训练与推理全程激活合并操作。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["大气状态输入 x_t 与 EDA 离散度 u_t"] --> B["动力学显著性与观测离散度提取"]
    B --> C["联合重要性评分构建<br/>s_i = d_i - β q_i"]
    C --> D["锚点保护掩码筛选<br/>Top-ρ 高分 Token 冻结独立保留"]
    D --> E["不确定性校准二分图匹配<br/>m_ij = sim(z_i, z_j) - λ_s(s_i+s_j) - λ_Δ|s_i-s_j|"]
    E --> F["加权聚合与 Token 序列缩减"]
    F --> G["多层 Transformer 演化与天气预报输出"]

关键设计

1. 动力学显著性与观测离散度联合重要性评分：兼顾物理变化剧烈度与输入数据可信度

针对纯特征相似度对观测噪声盲目合并的痛点，UtoMe 在 Patch 级显式解耦并联合了“大气动态活跃度”与“输入不确定度”。对于关键预报变量集合 \(\mathcal{C}_{\mathrm{dyn}}\)，算法在输入网格上统计空间相邻差分与相对于通道空间均值的异常振幅，经过 Patch 均值池化和样本内标准化获得动力学显著性分数 \(d_i\)；同时，将 ECMWF 数据同化集合（EDA）离散度场沿通道做 Patch 归一化与标准化，得到观测不确定度得分 \(q_i\)。两者按可控权重融合生成最终 Patch 级重要性评分：

\[s_i = d_i - \beta q_i\]

在默认 \(\beta=1.0\) 下，动力学剧烈且观测离散度低（高信度）的 Patch 获得最高评分；即使某些区域因强对流导致观测扰动大，只要动力学梯度足够强，依然能保持高分，避免了粗暴过滤活跃天气的缺陷。

2. 核心动力学特征锚点保护机制：杜绝关键天气系统信息的渐进损失

标准 ToMe 允许全图任意 Token 参与二分图划分和配对合并，导致多次迭代后台风眼、高空急流和锋面过渡带等高频特征被平滑稀释。UtoMe 在每一层合并前，依据重要性评分 \(s_i\) 选取前 \(\rho\)（默认 \(\rho=0.25\)）比例的 Token 划入锚点集合 \(\mathcal{A}\)。锚点 Token 既不作为合并源节点，也不作为被合并的目标宿主节点，在整个 Transformer 前向过程中保留独立的隐层表示，将计算压缩的负荷完全转移到背景冗余区域。

3. 不确定性校准的二分图匹配：基于重要性一致性的注意力约束

在剔除保护锚点后的候选集合上，两两 Token 之间的潜在合并价值不仅取决于隐层特征的余弦相似度 \(\mathrm{sim}(z_i, z_j)\)，还受到其重要性属性的调控。UtoMe 构造了引入惩罚项的匹配亲和度矩阵：

\[m_{ij} = \mathrm{sim}(z_i, z_j) - \lambda_s (s_i + s_j) - \lambda_{\Delta} |s_i - s_j|\]

其中 \(\lambda_s\) 为总和惩罚系数，用于抑制两个均具备中等重要性 Token 的合并几率；\(\lambda_{\Delta}\) 为差值惩罚系数，避免将相对关键的 Token 强制吸附到平坦无变化的低价值背景 Token 上。各源节点以贪心方式选择匹配得分最高的目标节点完成均值池化合并。

4. 三重不确定性可用性模式适配：解除对外部同化集合产品的硬性依赖

在实际业务与离线预报中，高质量的外部实时 EDA 不确定性产品并不总是可得。UtoMe 设计了三种互补的部署模式：完全依赖外部 EDA 输入的 UtoMe-E（训练与推理双阶段均使用外部 UQ）；仅在训练期利用 EDA 离散度监督、在 Transformer 浅层（Layer 3）引出轻量卷积探针预测重要性分数的 UtoMe-P（推理期摆脱外部场）；以及完全不依赖外部 UQ、纯粹利用局部 \(3\times 3\) 空间不一致性作为伪不确定性标签自监督训练探针的 UtoMe-S，使模型能够灵活落地于各类缺少同化集合的预报场景。

损失函数 / 训练策略

模型预报主干采用纬度加权的 L1/L2 预报损失函数。对于 UtoMe-P 和 UtoMe-S 两种依赖浅层重要性预测器 \(g_\phi\) 的变体，由于 Token 合并决策的核心在于排序（Top-\(\rho\) 锚点选择及匹配排序），设计了辅助性成对排序损失（Pairwise-ranking loss），其损失权重设为 \(0.2\)：

\[\mathcal{L}_{\mathrm{rank}} = \sum_{i, j} \max(0, -(s_i - s_j)(\hat{s}_i - \hat{s}_j) + \gamma)\]

训练全程采用 AdamW 优化器，学习率设为 \(5\times 10^{-4}\)，并引入 1000 步 Warmup。在 UtoMe-P/S 中，前 4 层保持全 Token 运行以供探针抽取稳定表征，随后层开始执行阶段性合并。

实验关键数据

主实验

在 ERA5 72 小时全球确定性天气预报任务中，覆盖两档空间分辨率（\(5.625^\circ\) 粗网格与 \(1.40625^\circ\) 细网格），评价指标为 Z500、T850、T2m、U10、V10 的纬度加权均方根误差（wRMSE）、综合平均误差（Overall）及异常相关系数（ACC）。

在 \(1.40625^\circ\) 高分辨率网格下的主实验对比数据如下：

数据集 / 分辨率	压缩率 / 配置	Overall wRMSE (↓)	ACC (↑)	Z500 wRMSE (↓)	T850 wRMSE (↓)	编码器相对 FLOPs (↓)
ERA5 \(1.40625^\circ\)	无合并 (Full Token)	32.41 ± 0.28	0.964 ± 0.004	156.76 ± 1.34	1.299 ± 0.007	1.00×
ERA5 \(1.40625^\circ\)	ToMe (25% 减缩)	42.45 ± 0.24	0.949 ± 0.003	206.02 ± 1.17	1.452 ± 0.006	0.86×
ERA5 \(1.40625^\circ\)	UtoMe-E (25% 减缩)	39.59 ± 0.21	0.952 ± 0.003	191.93 ± 1.05	1.413 ± 0.005	0.86×
ERA5 \(1.40625^\circ\)	UtoMe-P (25% 减缩)	39.61 ± 0.01	0.952 ± 0.000	191.98 ± 0.03	1.425 ± 0.001	0.87×
ERA5 \(1.40625^\circ\)	UtoMe-S (25% 减缩)	39.84 ± 0.27	0.951 ± 0.004	193.14 ± 1.31	1.416 ± 0.007	0.87×
ERA5 \(1.40625^\circ\)	ToMe (50% 减缩)	67.32 ± 0.43	0.905 ± 0.005	328.21 ± 2.11	1.846 ± 0.011	0.55×
ERA5 \(1.40625^\circ\)	UtoMe-E (50% 减缩)	53.00 ± 0.36	0.913 ± 0.004	256.72 ± 1.76	1.759 ± 0.009	0.55×
ERA5 \(1.40625^\circ\)	UtoMe-P (50% 减缩)	53.14 ± 0.01	0.907 ± 0.000	257.37 ± 0.05	1.762 ± 0.001	0.57×
ERA5 \(1.40625^\circ\)	UtoMe-S (50% 减缩)	57.01 ± 0.04	0.908 ± 0.000	276.99 ± 0.18	1.752 ± 0.001	0.57×

消融实验

在 \(1.40625^\circ\) 分辨率下针对各设计组件（动力学、UQ、锚点、校准匹配）的逐项消融验证结果如下表：

配置	25% 减缩 Overall (↓)	50% 减缩 Overall (↓)	说明
ToMe（无 UQ、无锚点、无校准）	42.45	67.32	原始特征相似度基线
仅动力学显著性 (\(\beta=0\))	40.95	58.60	缺失观测不确定度校正
仅 UQ 不确定度 (\(s=-q\))	41.40	61.70	缺失动力学梯度，易误删锋面
联合评分，去除锚点保护	40.31	61.40	激进减缩下关键结构退化严重
联合评分 + 锚点，去除匹配惩罚	40.68	55.90	未约束剩余非锚点之间的跨类别合并
完整 UtoMe-E	39.59	53.00	全模块协同最优

关键发现

• 锚点保护在深度压缩时起决定性作用：当 Token 压缩率达到 50% 时，剔除锚点保护会导致 Overall 误差从 53.00 急剧恶化到 61.40（+8.40），证明对极端梯度天气特征施加不可合并硬约束是维持大模型宏观环流刻画能力的核心防线。
• 空间对齐比不确定性绝对标定更关键：对 UQ 场进行空间洗牌（Shuffled UQ）会使 50% 减缩下的误差恶化到 62.30（接近盲目合并），而对 UQ 施加平方根、平方或百分位排序变换时，误差变动在 0.28 以内，证实只要相对空间排布一致，模型对不确定性的绝对尺度具有强鲁棒性。
• 高动态高不确定性极端区域收益最明显：细分分析显示，在中高层环流（Z500）的高动态活跃区，UtoMe 相比 ToMe 带来的误差降低达 18.0%，同时在上/下尾部极端值预报上误差分别降低 22.4% 与 22.7%。

亮点与洞察

• 将数据源固有物理误差转化为算力调度信号：巧妙地将同化系统固有的 Ensemble Spread 从单纯的后验检验指标，升格为引导网络内部注意力与 Token 聚合的前验控制变量。
• 软硬兼施的双重序列压缩架构：以 Top-\(\rho\) 锚点构建“不可侵犯”的硬性拓扑骨架，辅以考虑可靠度相似性的软性二分图惩罚，兼顾了局部高频信息保真度与全局背景冗余压缩率。
• 解耦推理期外部产品依赖：通过设计基于浅层表征与空间差分自监督的预测头（UtoMe-P/S），成功消除了实际业务部署时必须同步加载多源庞大同化集合的落地壁垒。

局限与展望

• 低动态背景区域的次优表现：将过多注意力配额向动态剧烈区倾斜后，平稳区域（如副热带高压中心等）的 Token 受到更激进的合并，导致低动态分区的微小误差有所上升。
• 预报误差未达全 Token 无损水平：在 50% 激进压缩下虽然显著好于 ToMe，但相比 Full Token 模型仍有约 63.5% 的综合指标退化，高算力预算追求极致精度时仍需谨慎选择压缩比（5%-10% 减缩可在精度几乎无损下节省计算）。
• 长时程自回归累计稳定性待验证：当前评测主要基于 72 小时单步直出预测，在 14 天或亚季节尺度的多次 Rollout 迭代中，Token 合并引起的微小相空间扰动是否会快速发散仍有待探索。

相关工作与启发

• vs ToMe (Bolya et al.): ToMe 依赖网络浅层至深层的隐特征余弦距离做二分图配对，在大气科学输入自带观测噪声时容易误合并；UtoMe 融合外部 EDA 物理离散度与动力学显著性，引入硬锚点与不确定性距离惩罚，显著提升了大气环流的物理保真度。
• vs Sparse-Reslim (Wang et al.): Sparse-Reslim 通过稀疏残差路由仅激活部分 Token 的中间块计算并还原全序列；UtoMe 则是物理级压缩序列长度并在训练全程激活，实现了更紧凑的端到端表征与更直接的 FLOPs 缩减。

评分

• 新颖性: ⭐⭐⭐⭐⭐ 首次将数据同化集合离散度（物理不确定性）作为 Token 压缩调度的核心先验。
• 实验充分度: ⭐⭐⭐⭐⭐ 覆盖两大空间分辨率、全组件细致消融、空间洗牌对比及极端天气切片分析。
• 写作质量: ⭐⭐⭐⭐⭐ 理论推导清晰（三方方差分解），动机阐述扎实，图表详实自洽。
• 价值: ⭐⭐⭐⭐☆ 为高分辨率天气/气候大模型的低成本微调、快速推演及星载/边缘部署提供了行之有效的途径。

相关论文

• [ICML 2026] Quantifying the Uncertainty of Foundation Models with Singular Value Ensembles
• [NeurIPS 2026] PGSB: Pretrained-Guided Shared Basis for LoRA Model Merging
• [NeurIPS 2026] Distillation of Foundation Models for Time-dependent PDEs
• [NeurIPS 2026] FeatCal: Feature Calibration for Post-Merging Models
• [NeurIPS 2026] Compact SO(3) Equivariant Atomistic Foundation Models via Structural Pruning