技術要約：より大きなドメインへのトランスフォーマー・ニューラル・オペレータのゼロショット汎化

問題提起

トランスフォーマーベースのニューラル・オペレータは、複雑な幾何学的形状における偏微分方程式（PDE）の解オペレータを近似する上で、強力な性能を示してきた。しかし、既存のアプローチは通常、訓練中に固定されたドメインサイズを想定しており、これが推論時におけるドメインサイズの拡張能力を制限している。本論文は、ドメイン拡張（domain extension）、すなわち、訓練時に遭遇したよりも大幅に大きな空間ドメインに対してゼロショット推論を行うという課題に取り組んでいる。

超解像や部分観測タスク（これらは補間問題である）とは異なり、ドメイン拡張は根本的に**外挿問題（extrapolation problem）**である。現在のトランスフォーマーベースのオペレータは、有限の訓練ドメインによって暗黙的に制限された非局所的な積分カーネルに依存している。推論ドメインのサイズが変化すると、標準的な位置エンコーディング（例：Rotary Positional Embeddings, RoPE）の持つ暗黙的な周期性が、訓練中に一度も同時に観測されたことのない空間的に離れた点同士の間に「ゴースト」相互作用を引き起こす可能性がある。これにより、モデルが新しいドメインのスケールに対して物理的な局所性を尊重できなくなり、汎化能力が損なわれる。

手法

著者らは、基礎となるアーキテクチャを変更したり、最適化されたアテンション・カーネル（例：FlashAttention）との互換性を犠牲にしたりすることなく、トランスフォーマー・ニューラル・オペレータのアテンション機構内に明示的かつ制御可能な空間的局所性を強制する手法を提案している。

コアメカニズム：分解可能な局所性バイアス

このアプローチは、Rotary Positional Embeddings (RoPE) と、アテンションのロジットに加算される分解可能な距離バイアスを組み合わせたものである。

1. RoPEバックボーン: RoPEは移動等変性と表現力豊かな相対的位置エンコーディングを提供する。しかし、それは本質的に非局所的であり、準周期的である。
2. 局所性制約: 偽の長距離相互作用を防ぐため、アテンションのロジットからバイアス項 $\phi$ を減算する。このバイアスは幾何学的なポテンシャルとして機能し、アテンション・カーネルを特定の距離 $r$ 以内に急速に減衰させ、積分サポートを近傍 $B_r(c)$ に効果的に制限する。
3. 分解可能性: 高性能なGPUカーネルとの互換性を維持するため、バイアスはクエリの位置の関数とキーの位置の関数の関数（$z = q^\top R_\theta k - \phi$）へと分解可能でなければならない。

提案されるエンコーディング

論文では、このバイアスに関する2つの具体的な定式化を紹介している。

• Local Anisotropic Symmetric Positional Encoding (LASPE):
  異方的な $L_2$ 距離バイアスを使用する：
  $$z_{LASPE} = q^\top(c) R_\theta^\top(c-\xi) k(\xi) - \frac{\sqrt{d}}{2} (c-\xi)^\top A (c-\xi)$$
  ここで、$A$ は異方性を制御する対角行列である。この手法は対称なドメインには効果的であるが、二次項に依存しているため、座標値が非常に大きい場合に数値的不安定性を招く可能性がある。

• Local Anisotropic Asymmetric Positional Encoding (LAAPE):
  非対称性（例：風の流れにおいて、上流と下流で相互作用が異なる場合）に対処し、数値的安定性を向上させるために、著者らは双曲線関数（具体的には指数関数）に基づくバイアスを提案している：
  $$z_{LAAPE} = q^\top(c) R_\theta^\top(c-\xi) k(\xi) - \frac{\sqrt{d}}{2} \sum_i \left[ \exp(A_-^{1/2}(c-\xi)\cdot e_i) + \exp(A_+^{1/2}(\xi-c)\cdot e_i) \right]$$
  この定式化により、正の方向と負の方向で異なる減衰率（$A_+$ と $A_-$）を設定することが可能になり、カーネルのサポートを非対称かつ異方的にすることができる。これはLASPEの自然な拡張であるが、大きなドメインにおいても優れた安定性を提供する。

カーネル幅の選択

著者らは、カーネル幅（減衰パラメータ $\lambda$ によって制御される）は、訓練ドメインのサイズに基づいて選択されるべきであると主張している。具体的には、拡張されたテストドメインにおける偽の相関をモデルが学習することを防ぐため、アテンションの重みが訓練ドメインを超える距離において統計的に無視できるレベルになるように $\lambda$ を設定すべきである。

主な貢献

1. 問題の定式化: 本論文は、ドメイン拡張をニューラル・オペレータのゼロショット汎化課題として正式に定義し、現在の非局所的な仮定が外挿の設定において持つ限界を浮き彫りにした。
2. 新しい位置エンコーディング: 著者らは、RoPEと分解可能な距離バイアスを組み合わせたLAAPEおよびLASPEを提案している。これらのエンコーディングは以下を提供する：
   • 制御可能な局所性: 相互作用の範囲を明示的に制限する。
   • 異方性と非対称性: 方向依存の相互作用範囲をモデリングする能力（流体力学のような複雑な物理システムにおいて極めて重要）。
   • 効率性: 分解可能性を通じて、最適化されたアテンション・カーネル（FlashAttention）との完全な互換性を維持する。
3. 実証的検証: この手法は、2つの学術的なPDEベンチマーク（1次元の浅水波方程式および2次元のグレイ・スコット方程式）と、困難な3次元の産業応用（都市環境における微細規模の大気流）を用いて検証されている。

実験結果

モデルは小さなドメインで訓練され、段階的に大きくなるドメイン（スケーリング係数 $s$ が最大10倍）でテストされた。

• 学術的ベンチマーク (1D & 2D):

  • 標準的な RoPE モデルは、ドメインサイズが増大するにつれて顕著な誤差の悪化を示した（例：Gray-ScottのL1誤差は、$s=10$ において約30%から400%以上に増加した）。
  • LASPE および LAAPE モデルは、スケーリング係数に関わらず、ほぼ一定の誤差率を維持した。
  • 高いスケーリング係数（$s \ge 7$）において、LASPEは大きな座標値による数値的不安定性を示したが、LAAPE はLASPEを上回る性能を示した。

• 産業応用 (3D 大気流):

  • AB-SWIFT モデル（微細規模の大気流）を用いたタスクでは、強い異方性と非対称性（建物の配置が風上にあること）が含まれていた。
  • RoPE は、スケールされたドメイン（$5\times5$）において破綻し、不整合な流れの予測を生じさせた。
  • LAAPE は、全体的な流れの構造を維持し、後流領域を正確に予測した。誤差の増加はRoPEよりも大幅に低かった。著者らは、LAAPEに残る誤差は、オペレータ自体の失敗ではなく、より大きなドメインから生じる新しい流れのパターンに起因するものと考えている。

意義と主張

本論文は、構造化された空間的局所性を強制することが、ニューラル・オペレータを産業応用へとスケールさせるための重要な帰納バイアスであると主張している。相互作用の実効的なサポートを訓練ドメインと一致するように明示的に制約することで、モデルは、訓練データやモデル容量を比例して増やすことなく、より大きなドメインへと汎化できる。

著者らは、本研究を機械学習における物理学のための「適切なスケーリング則（appropriate scaling laws）」への一歩として位置づけている。これは、単にデータやモデルの規模に頼るのではなく、情報が空間をどのように伝播するかを反映するように帰納バイアスを調整することで、性能向上が得られることを示唆している。この手法は、堅牢なゼロショット・ドメイン拡張を可能にする。この能力は、これまでドメイン分解戦略やグラフベースのモデルに限定されていたが、本手法は高性能計算カーネルと互換性のある統一されたトランスフォーマー・アーキテクチャ内でこれを実現している。

著者らは限界についても認めており、局所性は汎化を向上させるものの、訓練ドメインよりも本質的に大きなスケールに存在する物理的な相関を捉えることはできないと述べている。そのようなケースについては、少数の大規模なケースに対するファインチューニングが依然として将来的な必要事項であるとしている。