[論文レビュー] Learning Operators with Coupled Attention

Georgios Kissas, J. S. Seidman+4 more|arXiv (Cornell University)|Jan 4, 2022

Hydrological Forecasting Using AI被引用数 56

ひとこと要約

LOCAは、関数空間間の写像を学習する新しい演算子学習フレームワークを提案します。Kernel-Coupled Attentionを用いて普遍近似保証と、PDE/ODEおよび気候関連タスクにおける高いデータ効率性、ロバスト性、一般化性能を実現します。

ABSTRACT

Supervised operator learning is an emerging machine learning paradigm with applications to modeling the evolution of spatio-temporal dynamical systems and approximating general black-box relationships between functional data. We propose a novel operator learning method, LOCA (Learning Operators with Coupled Attention), motivated from the recent success of the attention mechanism. In our architecture, the input functions are mapped to a finite set of features which are then averaged with attention weights that depend on the output query locations. By coupling these attention weights together with an integral transform, LOCA is able to explicitly learn correlations in the target output functions, enabling us to approximate nonlinear operators even when the number of output function in the training set measurements is very small. Our formulation is accompanied by rigorous approximation theoretic guarantees on the universal expressiveness of the proposed model. Empirically, we evaluate the performance of LOCA on several operator learning scenarios involving systems governed by ordinary and partial differential equations, as well as a black-box climate prediction problem. Through these scenarios we demonstrate state of the art accuracy, robustness with respect to noisy input data, and a consistently small spread of errors over testing data sets, even for out-of-distribution prediction tasks.

研究の動機と目的

• 時空系のための関数空間間の演算子学習を動機づける。
• 出力関数の相関を捉える新しいアテンションベースのアーキテクチャを導入する。
• 提案モデルの普遍近似保証を提供する。
• 実データと合成データの両方でデータ効率性、ノイズに対するロバスト性、良好な一般化を示す。

提案手法

• 入力関数を有限の特徴集合 v(u)へリフトし、クエリ位置 y での出力を E_phi(y)[v(u)] によって計算する。
• Kernel-Coupled Attention (KCA) を導入し、カーネル積分演算子を用いて異なる出力クエリ位置間でアテンション重みを結合する。
• カーネル kappa を、リフトされた入力 q_theta と普遍カーネルによって定義し、結合されたアテンション分布を生じるよう正規化する。
• 入力関数を D(u) 特徴写像を通して普遍関数 f によってエンコードし、v(u)=f(D(u)) を形成する。
• 変形とノイズに対する頑健性のため、オプショナルなスペクトルエンコーダ D としてウェーブレット散乱を用いる。
• 標準のL2損失を用いた経験的リスクによる訓練を提供し、カーネル積分のモンテカルロまたは区分推定近似を有効にする。
• 実用的な側面を議論する：y の位置エンコーディング、積分の離散化戦略、勾配法による最適化。

実験結果

リサーチクエスチョン

• RQ1 LOCA は関数空間間の任意の連続演算子を近似できるか（普遍近似性）？
• RQ2 出力クエリ位置全体でアテンションを結合することは、出力測定数が少ないときの精度とデータ効率性を向上させるか？
• RQ3 ノイズを含む入力と分布外/一般化シナリオで LOCA はどう機能するか？
• RQ4 ODE/PDEおよび気候データタスクにおいて、LOCAは既存の演算子学習法とどう比較されるか？
• RQ5 普遍性と性能を保つ実用的な実装選択は何か？

主な発見

• 適切な仮定の下で LOCA は普遍近似性を満たす。
• 入力あたりの出力評価数が少ないとき、結合されたアテンションは精度を向上させる。
• LOCA はノイズへのロバスト性を示し、外れ値を減らし、誤差は中央値付近に集中する。
• 合成データおよび実データ（地表の気温と気圧）で、最先端の精度とより良い一般化、訓練データを超えた外挿を含む、実証的結果を示す。
• モデルは高いデータ効率性を達成し、競合手法と比較してラベル付きデータのごく一部（6-12%）しか必要としない。