技術要約：Transformerのカーネル・アテンション学習のためのクーロン粒子モデル

問題提起

カーネル法は、非線形な類似性を符号化するための原理的なアプローチを提供するが、サンプル数に対するスケーラビリティの低さが課題となることが多い。ランダム特徴量（例：ランダムフーリエ特徴量）は、暗黙的なカーネルを明示的なランダム写像に置き換えることで、線形時間での推論を可能にし、この問題に対処している。しかし、決定的な限界が依然として存在する。それは、特徴量分布がラベルを観測する前に固定されているという点である。適切な類似性構造が未知である場合、この事前に決定された選択がダウンストリームの性能を大きく左右してしまう。既存のカーネル学習手法は、カーネル学習と特徴量推定を分離しているか、あるいは重要度サンプリングに依存しており、特定のタスクに対して特徴量分布を完全に最適化できていない可能性がある。

手法

著者らは、統計力学の枠組みを用いて、ランダム特徴量を相互作用する粒子として扱う、教師ありカーネル学習手順を提案している。コアとなる手法は、主に以下の3つのコンポーネントで構成される。

1. 粒子ベースの最適化

本手法は、ランダム特徴量のパラメータ $\Omega = \{\omega_1, \dots, \omega_N\}$ を粒子としてモデル化する。学習目的関数は、以下の2つの項からなる有限次元のハミルトニアン $H_N(\Omega_N)$ として定式化される：

• ターゲット・アライメント（標的適合性）： 実験的なカーネルが訓練ラベルと一致することを報酬とする項（カーネル・ターゲット・アライメントの最大化）。これは、データを説明する特徴量構成へと粒子を引き寄せる「引力」として機能する。
• 反発的正則化： 粒子の崩壊を防ぐためのリーシュ/クーロン反発ポテンシャル（$W_{N,s}$）。この項は、パラメータ空間 $\Omega$ 全体に粒子を分散させることで、学習される特徴量の多様性を促進する。

2. ランジュバン最適化と平均場極限

最適化は、射影ランジュバン動力学アルゴリズム（Algorithm 1）を用いて行われる。この確率的勾配降下法には以下が含まれる：

• ハミルトニアンの勾配から導出されるドリフト項。
• 逆温度 $\beta$ によってスケールされるノイズ項。
• 粒子がコンパクトかつ凸な領域内に留まることを保証するための射影ステップ。

理論的に、著者らはマッケーン–ヴラソフ（McKean–Vlasov）平均場極限を確立している。粒子数 $N \to \infty$ のとき、粒子の経験的測度は、偏微分方程式によって制御される確率密度 $\rho_t$ に収束する。この連続体極限は、データのポテンシャル（アライメント）と相互作用ポテンシャル（反発）のバランスを取る、自由エネルギー汎関数のワッサースタイン勾配流として学習ダイナミクスを記述する。

3. 線形化Transformerアテンションへの適用

本手法は、Transformerのアテンション機構において以下のように実装される：

• フェーズA（アライメント）： シーケンス表現に対するカーネル・ターゲット・アライメントを最適化することで、特徴量粒子 $\Omega_N$ を学習する。このフェーズでは、ランダム特徴量の分布を学習する。
• フェーズB（ファインチューニング）： 学習されたカーネル（粒子）を固定する。残りのTransformerパラメータ（例：投影行列）は、標準的なクロスエントロピー損失を用いて訓練される。
• 推論： 得られたアテンション機構は、学習された正のランダム特徴写像を使用し、線形アテンションの線形推論の計算量複雑性を維持したまま（シーケンス長への依存度を二次関数から線形へ削減）、動作する。

主な貢献

本論文は、主に4つの貢献を述べている：

1. 粒子ハミルトニアン： ターゲット・アライメントとリーシュ/クーロン反発正則化を組み合わせた、教師ありランダム特徴量学習のための新しい定式化。
2. 連続体極限を伴うランジュバン最適化： 平均場記述を可能にするマッケーン–ヴラソフ方程式を伴う具体的な最適化手順であり、学習ダイナミクスの統計力学的解釈を提供する。
3. 線形時間の実装： タスク適応型のカーネルを学習しながら、線形推論の複雑性を維持する、正のランダム特徴量アテンションへの実用的な適用。
4. 実証的検証： 学習されたカーネルが、固定されたランダムフーリエ特徴量や他のベースラインと比較して、合成データおよびNLPベンチマークにおいて精度、較正性（キャリブレーション）、および堅牢性を向上させることを示す証拠。

実験結果

著者らは、合成分類タスクおよび3つの文レベルのベンチマーク（SST-2、QQP、Rotten Tomatoes）を用いて手法を評価した。

• 合成結果： 最適化された粒子は、特に中程度のノイズ条件下において、固定されたランダムフーリエ特徴量や重要度サンプリングのベースラインを上回った。性能の差は特徴量の予算が増えるにつれて拡大した。最適化された粒子は、わずか64個の粒子を使用した場合でも、256個の粒子を用いた場合の性能の約90%を維持したが、固定特徴量はより急激に低下した。
• NLPベンチマーク：
  • SST-2： 「Kernel-PORF-softplus」バリアントが、最高の精度とマシューズ相関係数（MCC）を達成した。「Kernel-FAVOR」バリアントは、最も低い期待較正誤差（ECE）を記録した。
  • QQP： 「Kernel-softmaxfeat」モデル（$Q=K$ の場合）が最も強力な構成となり、バニラ版と比較してLogLossとBrierスコアを減少させた。
  • Rotten Tomatoes： 「Kernel-softmaxfeat」が、最高の精度、MCC、LogLoss、およびBrierスコアをもたらした。
• 効率性： 一回限りのアライメントフェーズ（フェーズA）により訓練時間は増加したが、推論時間はシーケンス長に対して線形に保たれ、対応するバニラ特徴写像の複雑性と一致した。

意義と主張

本論文は、ラベル・アライメント、反発的正則化、ランジュバン最適化、および平衡集中を結びつける統一されたフレームワークを提供することを主張している。その意義は、線形アテンションの計算効率を損なうことなく、実用的かつタスク適応的な方法で線形アテンションのカーネルを学習できる点にある。

著者らは、学習されたカーネルがタスク適応的な変分カーネルの周囲に集中することを強調している。大偏差原理（Large Deviation Principles）に基づく理論的結果は、関連するレート関数が一意の最小値を持つ場合、学習されたカーネルは粒子数または逆温度が増加するにつれて指数関数的に最適なカーネルへと収束することを示唆している。本研究は、固定されたランダム特徴量に代わるスケーラブルな選択肢として位置づけられ、特徴量分布をラベルから直接学習することが、識別能と較正指標の双方において測定可能な改善をもたらすことを実証している。