技術要約：効率的な気象予測のための疎な残差ルーティング（Sparse Residual Routing）

問題提起

データ駆動型の気象予測は、Pangu-Weather、ClimaX、FuXiといったモデルにより、従来の数値予報（NWP）と同等またはそれを上回る精度を、極めて低い推論コストで実現し、大きな成功を収めてきました。しかし、これらのモデルの学習は依然として主要なボトルネックとなっています。最先端のモデルの多くはVision Transformer（ViT）を利用していますが、これらはすべての空間トークンに対して一様に自己注意（self-attention）計算を適用します。このアプローチは非効率的です。なぜなら、大気場には高い空間自己相関が存在し、広大な領域は滑らかで緩やかに変化する条件を含んでいるため、中間段階のトークン表現が冗長になることが多いからです。

トークン削減、マージ、あるいはランダム選択（画像生成におけるSprintなど）といった既存のトークン効率化手法は、気象予測には適していません。気象予測は時空間的な高密度予測問題であり、オートレグレッシブなロールアウト（逐次計算）において、元の緯度・経度格子上のすべてのグリッドセルに対して物理的に意味のある出力が求められます。トークンを削除したり、融合したり、学習されたマスク埋め込みに置き換えたりする標準的な手法は、固定された格子構造と、正確な予測に不可避な時間的コンディショニングを崩してしまいます。

手法：Sparse-Reslim

著者らは、気象予測の持つ高密度な格子要件を維持しながら、疎なトークン処理を可能にする、パラメータフリーのプラグイン・ルーティング・モジュールであるSparse-Reslimを提案しています。

コアメカニズム

Sparse-Reslimは、Transformerブロックを以下の3つのステージに分割します。

1. 高密度初期ステージ（Dense Early Stage）: 最初の $n_1$ ブロックが、豊かな初期表現を構築するために全 $L$ 個の空間トークンを処理します。
2. 疎な中間ステージ（Sparse Middle Stage）: 中間の $n_2$ ブロックは、ランダムに選択されたサブセットである $K = \lfloor rL \rfloor$ 個のトークンのみを処理します（ここで $r$ は保持率であり、通常は0.25です）。
3. 高密度後期ステージ（Dense Late Stage）: 最終的な予測ヘッドの前に空間的な一貫性を精緻化するため、最後の $n_3$ ブロックが全シーケンス $L$ 個のトークンを処理します。

残差デルタ定式化（Residual Delta Formulation）

本手法の核心的な革新は、疎なステージが全シーケンスとどのように相互作用するかという点にあります。選択されなかったトークンをプレースホルダーやマスク埋め込みで置き換える代わりに、Sparse-Reslimは**残差デルタ（residual delta）**を計算します。

1. 選択されたトークンが疎な中間ブロックを通過します。
2. これらのブロックの出力は、入力（ルーティング前の表現）と比較され、変化量（$\Delta$）が計算されます。
3. この $\Delta$ が元のトークンの位置へとスキャッター（散布）されます。
4. 選択されなかったトークンはゼロのデルタを受け取り、暗黙的なアイデンティティ・スキップ接続を介して、ルーティング前の表現を正確に保持します。

このアプローチにより、グリッドセルが脱落したり、合成されたトークンに置き換えられたりすることがなくなり、ネットワーク全体を通じて大気状態の物理的な整合性が維持されます。

アーキテクチャへの統合

• 決定論的モデル (Res-Slim-ViT): このモジュールは、標準的なViTブロックのループに直接統合されます。疎な中間ブロックは、自己注意の計算複雑性を $O(L^2D)$ から $O(K^2D)$ へと低減させます。
• 生成モデル (EDM/Diffusion): 著者らは非対称アテンション設計を採用しています。自己注意はターゲットとなるトークンに対して疎化（$O(K^2D)$）されますが、クロスアテンションは、キーおよび値としてフルコンディショニング・シーケンス（$Z_c$）を保持します（$O(KL \cdot D)$）。これにより、ターゲットシーケンスが疎化されている場合でも、デノイジングに必要な完全な大気コンテキスト（履歴フレーム）をモデルが保持できることが保証されます。

トークン選択戦略

本手法は、アテンション・スコアに基づく選択や学習された選択ではなく、ランダム・トークン選択を利用しています。著者らは、これが3つの利点を提供すると主張しています。

1. スコアリングのオーバーヘッドがゼロであること。
2. اینکه、ドロップアウトのトークン次元版のような、確率的な正則化として機能すること。
3. 動的に活発な領域（例：嵐の領域）を優先して、静穏な領域（例：熱帯地域）を過小評価するといった、系統的な空間バイアスを回避できること。

主な貢献

1. Sparse-Reslimモジュール: 残差デルタ定式化を用いて、高密度な予測要件を損なうことなく、ViTベースの気象モデルの効率的な学習を可能にする、パラメータフリーのルーティング・メカニズム。
2. 実証的検証: 2つのアーキテクチャ・ファミリー（決定論的Res-Slim-ViTおよび生成型EDM）と、3つのERA5解像度（運用標準である0.25°を含む）にわたる包括的な評価。
3. 性能向上: 疎なルーティングが計算コストを削減するだけでなく、高密度なベースラインと比較して予測精度をも向上させることを実証。
4. 直交性: 本手法が混合精度訓練および解像度のダウンサンプリングと互換性があることを示す。

実験結果

著者らは、決定論的なRes-Slim-ViTと拡散ベースの生成モデルを用いて、ERA5再解析データセット上でSparse-Reslimを評価しました。

効率性

• 学習速度: Sparse-Reslimは、主要な設定において約2.5倍高速な学習を実現し、0.25°解像度では最大3.18倍の高速化に達しました。
• メモリ: テストされたすべての設定において、ピークVRAM使用量が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など）に見られる合成プレースホルダー・トークンの「不一致」を回避し、すべてのグリッドポイントにおいて物理的に意味のある表現を維持することを可能にしています。

著者らは、本手法が非常に効果的である一方で、将来の研究として、ランダム選択ではなく、気象学的に活動的な領域（例：嵐のフロント）に計算を動的に割り当てるための適応的トークン選択を探索できる可能性があると控えめに述べています。しかし、現在の研究は、ランダム選択と残差ルーティング・アーキテクチャの組み合わせが、大幅な高速化と精度向上を両立させるために十分であることを確立しています。