Reformulating Neural Operators in $d+1$ Dimensions for Embedding Evolution

本論文は、補助的な関数次元を介して埋め込みの進化をモデル化する、新しい$d+1$次元ニューラルオペレータフレームワークを導入しており、従来の埋め込みスケーリング手法に伴う計算コストを回避しつつ、多様な物理的ベンチマークにおいて最先端の精度と堅牢性を実現している。

要約

あなたは、コンピュータに、金属板を通る熱の伝わり方や嵐の中の水流のように、複雑な物理システムが時間の経過とともにどのように変化するかを予測させる方法を教えようとしていると想像してください。人工知能の世界では、これらの問題を解くために設計されたAIモデルは、**偏微分方程式（PDE）**と呼ばれる数学的な規則によって記述されることがよくあります。

長い間、これらの問題を解くために設計されたAIモデル（ニューラルオペレーターと呼ばれます）は、「力任せ」に近い戦略に依存してきました。もしモデルの精度が不十分であれば、エンジニアは単にモデルの内部チャネルやレイヤーを増やして、モデルをより「太らせる」ことで解決しようとしてきました。これは、バケツがすでに重くて扱いにくいとしても、より多くの水を運ぼうとして、バケツ自体を幅広にするようなものです。

この論文は、よりスマートな水の運び方を提案しています。単にバケツを幅広にするのではなく、バケツ自体に新しい次元を追加することを提案しているのです。

コアとなるアイデア：「影」の次元

物理的な世界（例えば、都市の2Dマップ）を、平らな紙のシートだと考えてみてください。従来のAIモデルは、そのシートを上から、一層ずつ眺めることで、そのパターンを理解しようとします。

著者であるHaoze Song氏とそのチームは、単に紙を見るのではなく、その紙に影やゴースト次元が付着していると想像すべきだと提案しています。彼らはこれを「補助的な次元（auxiliary dimension）」（これを「p次元」と呼びましょう）と呼んでいます。

• 従来の方法： 3Dの物体を理解しようとする際、2Dの写真を見ながら、細部が見えるように目を凝らす（ピクセルを増やす）だけの状態です。
• 新しい方法（SKNO）： 2Dの写真がある一方で、その写真の影を隣の壁に投影する特別な「影プロジェクター」を持っている状態を想像してください。写真と、その影の両方を同時に研究することで、より大きな写真を用意することなく、3Dの形状をより良く理解することができます。

この論文において、彼らはSKNO（Schrödingerised Kernel Neural Operator）と呼ばれるモデルを作成しました。これは、データが一つ余分な次元を持つ空間に存在するかのように扱います。モデルは、物理的なマップ上のデータを更新するだけでなく、そのマップとその影の両方のデータを同時に更新します。

その仕組み：「二つの視点」戦略

SKNOの魔法は、この余分な次元をどのように更新するかという点にあります。著者らは、量子物理学（具体的にはシュレディンガー方程式ですが、ここではシミュレーションとしてではなく、設計のブループリントとして使用しています）から着想を得た巧妙なトリックを使用しています。

彼らは「影」のデータを、以下の2つの異なる方法で同時に更新します。

1. 生の視点（The Raw View）： データをそのままの姿として見る（通常のテキストとして本を読むようなもの）。
2. フーリエの視点（The Fourier View）： データを波と周波数の混合物として見る（本を音波の楽譜として読むようなもの）。

これら2つの「影の次元」の視点を組み合わせることで、モデルは非常に効率的に複雑なパターンを捉えることができます。それは、まるで「普通の英語」と「詩的な英語」の両方を話す翻訳者がいるようなものです。その翻訳者は、一つの言語しか話せない人よりも、文のニュアンスをはるかに深く理解できるのです。

結果：より速く、より小さく、より正確に

チームは、単純な熱方程式から、高度に混沌とした3Dの流体爆発（レイリー・テイラー不安定性）に至るまで、10種類以上の困難な物理問題を用いてこの新モデルをテストしました。

判明したことは以下の通りです：

• 低い誤差： SKNOは、既存の最良のモデル（FNO、Transolver、DeepONetなど）よりも一貫して間違いが少ないという結果を出しました。
• 効率性： モデルを「太く」したり、計算コストを増やしたりすることなく、これらの結果を達成しました。実際、トレーニングはより速く、より少ない計算資源で済むことが多いという結果が出ました。
• 堅牢性（ロバスト性）： 未知のデータ（例えば、学習していない日の天候パターンや、より高い解像度のデータ）に対してテストされた際も、競合モデルよりも優れた性能を維持しました。データの「グリッド」のサイズが変わっても、混乱することがありませんでした。

まとめ

この論文は、難しい物理問題を解くためにAIモデルを単に大きく、重くするのではなく、データの「見方」を変えるべきであると主張しています。「影の次元」を追加し、データを2つの異なる数学的なレンズ（生および周波数ベース）を通じて更新することで、モデルは物理学の根本的なルールをより自然に学習します。

これは、「問題に対してより多くのリソースを投入する」ことから、「問題を見るためのより良い角度を見つける」ことへの転換です。その結果、より正確であるだけでなく、よりエレガントで効率的なモデルが誕生しました。

技術概要

技術要約：埋め込みの進化のための $d + 1$ 次元におけるニューラルオペレータの再定式化

問題提起

ニューラルオペレータ（NO）は、関数空間間のマッピングを学習するように設計されており、特に偏微分方程式（PDE）の解法に用いられる。近年の進展は、$d$ 次元の物理ドメインにおけるカーネルパラメータ化の洗練に焦点を当ててきたが、リフトされた埋め込み（embedding）の進化については十分に探索されていない。既存のアーキテクチャは、埋め込みの表現力不足を補うために、ブルートフォース的なスケーリング（埋め込み幅の拡大やヘッドの追加）を行う傾向がある。しかし、この戦略は高い計算コストを伴う。すなわち、密なチャネル混合は埋め込み幅に対して二次的にスケールし、ヘッドごとの因子分解も、ブロック対角構造を誘発することでクロスヘッド間の結合を弱めてしまうため、緩和効果は限定的である。本論文は、単に容量を増やすのではなく、「どのように」埋め込みを進化させるかを直接設計するという観点におけるギャップを特定している。

手法

著者らは、補助的な関数次元 $p$ を導入することで、ニューラルオペレータのパイプラインを $d + 1$ 次元へと再定式化することを提案する。提案されるフレームワークでは、埋め込みを物理ドメイン $D_x$ の上だけで進化させるのではなく、積ドメイン $D_x \times D_p$ 上の潜在的なスカラー関数として進化させる。

一般的なフレームワーク

1. リフティング（Lifting）: 入力場 $a(x)$ は、積ドメイン上のスカラー潜在関数 $v_0(x, p)$ へとリフトされる。これは、分離された線形写像 $v_0(x, p) = w^\top(p)a(x)$ を用いたリフティング演算子 $P$ によって実現される。
2. $(d+1)$ 次元の進化: 潜在関数は、学習可能な線形演算子 $\mathcal{L}$ と非線形写像 $\sigma$ のシーケンスを通じて進化する。核となるコンポーネントは、物理座標 $x$ と補助座標 $p$ の両方に対して作用するカーネル積分演算子 $\mathcal{K}$ である：
   $$\mathcal{K}_l[v_l](x, p) = \int_{D_x} \int_{D_p} \kappa_l(x, y, p, p') v_l(y, p') \, dp' \, dy$$
3. リカバリ（Recovery）: 進化した関数 $v_L(x, p)$ は、通常 $p$ に関する積分であるリカバリ演算子 $Q$ を通じて出力ドメインへと写される：$u_{pred}(x) = \int_{D_p} \chi(p) v_L(x, p) \, dp$。

Schrödingerised Kernel Neural Operator (SKNO)

本論文では、このフレームワークを SKNO と呼ばれるフーリエベースのモデルとして具体化している。主な設計上の選択肢は以下の通りである：

• 基底多様化された補助進化（Basis-Diversified Auxiliary Evolution）: 各空間位置において、補助次元 $p$ に沿った信号は、2つの異なる座標ビューを用いて更新される：
  1. 生の $p$ 座標混合: $p$ の空間ドメインにおける線形混合。
  2. フーリエ $p$ 座標混合: $p$ のフーリエ領域におけるスペクトル混合。
     このデュアルブランチ構造（$F_p^{-1} \tilde{A}_l F_p + B_l$）により、モデルは単に同じチャネル混合パスを複製することなく、両方のビューから特徴を捉えることが可能になる。
• 物理ドメインの伝播: SKNOは、$(L-1)$ 個のグローバル伝搬器（フーリエ領域の $x$ で対角化されたスペクトル畳み込み演算子を使用）と、グローバルなスペクトル手法によって失われる局所的な情報を捉えるための1つの最終的なローカル伝搬器を採用している。
• 残差接続: 学習の容易さと安定性を促進するため、線形ブロックには残差接続が含まれている。

主な貢献

1. オペレータレベルの再定式化: 著者らは、物理座標と補助座標の両方に関するカーネル積分を通じて潜在関数を進化させるよう、NOのパイプラインを再定式化し、埋め込み進化のための明示的なオペレータベースのメカニズムを確立した。
2. SKNO アーキテクチャ: 生の $p$ 座標とフーリエ $p$ 座標を混合する「基底多様化された補助進化」を利用して、ブルートフォース的なスケーリングなしに表現力を向上させる、Schrödingerised Kernel Neural Operatorを提案した。
3. 包括的な評価: 1次元の線形方程式から高度に非線形な3次元不安定性まで、10種類以上のベンチマークを用いてモデルを評価した。
4. 制御された分析: パフォーマンスの向上は、単なるパラメータ数の増加ではなく、アーキテクチャのデザイン（基底の多様性）に由来することを実証するために、厳密な比較とアブレーション研究を行っている。

実験結果

1D 熱伝導/移流方程式、1D Burgers、2D Darcy Flow、2D Gray-Scott、2D/3D Navier-Stokes、および3D Rayleigh-Taylor 不安定性を含むベンチマーク全体において、SKNOは評価されたベースライン（DeepONet, FNO, Transolver, CNO）の中で一貫して最も低い相対 $L_2$ 誤差を達成した。

• パフォーマンスの向上: 2D 非圧縮 Navier-Stokes ($\nu=10^{-5}$) において、SKNOは FNO と比較して相対 $L_2$ 誤差を約 37.1% 削減した。2D Gray-Scott では 42.1% の削減を実現した。3D Rayleigh-Taylor では、SKNO は 14.3% の誤差削減を達成した。
• 容量効率: 制御実験により、SKNO (A+B) は、より少ないパラメータ数と FLOPs で、系統的にスケールさせた FNO 変種や並列にスタックされた FNO を上回ることが示された。「B+B」変種（生の $p$ ブランチを複製したもの）は、基底多様化された「A+B」変種の性能に達しなかったことから、デュアル座標ビューの価値が確認された。
• 堅牢性: SKNO は優れた解像度不変性を示し、混合解像度での学習およびゼロショット超解像推論（例：128 グリッドで学習し、8192 でテスト）の下でも低い誤差を維持した。また、未知の時間領域に対する強力なゼロショット汎化能も示した。
• 効率性: 追加された次元にもかかわらず、SKNO は競争力のある学習時間を維持しており、埋め込みサイズに対して二次的な複雑さを持つ Transolver のような Transformer ベースのモデルをしばしば凌駕している。

意義と主張

本論文は、補助ドメインにおけるオペレータ進化が、ブルートフォース的な埋め込みスケーリングに代わる有望な選択肢であることを主張している。補助座標に沿ってオペレータ設計の原理を適用することで、モデルは埋め込みを広げるという法外な計算コストを伴わずに、表現力と近似能力を向上させている。

著者らは、「Schrödingerised」という命名は、補助座標に沿った構造化されたオペレータ進化のデザイン・インスピレーションとしての役割であり、PDE 自体に対する直接的な古典的数値加速メカニズムを主張するものではないことを強調している。結果は、提案された $d+1$ 次元の設計が、より低い誤差、優れた解像度堅牢性、および優れた容量効率によって裏付けられた、ニューラルオペレータの性能を向上させるためのより直接的かつ効率的な経路を提供することを示唆している。

最後に、今後の研究は、最終的なテスト誤差を超えてニューラルオペレータを比較するための定量的な基準を開発すること、具体的には、異なる集約デザインが最適化の軌跡や高次元誤差ランドスケープにおける局所解の選択にどのように影響するかを調査することに焦点を当てるべきであると述べている。