技術要約: Flexformer

問題提起

Transformerモデルは、長距離の依存関係を捉えるためにドット積アテンション機構に依存しているが、この機構はシーケンス長 $N$ に対して二次的な時間および空間計算量 ($O(N^2)$) を招く。このボトルネックは、長いシーケンスに対するスケーラビリティを著しく制限している。カーネルベースのリニアアテンション手法は、特徴写像を用いてアテンション機構を近似することで、計算量を線形 ($O(N)$) に削減できるが、既存のアプローチは効率性と表現力の間のトレードオフに直面している：

• 固定カーネル: PerformerやRFAのような手法は、ソフトマックスカーネルを近似するためにランダムフーリエ特徴量（RFF）に依存している。しかし、これらはカーネルを固定かつ非学習的なものとして扱っており、ソフトマックスが最適であるという仮定に基づいているが、これはあらゆるシナリオにおいて必ずしも成立するとは限らない。
• 弱学習可能カーネル: 近年の研究（Hedgehog、Polaformerなど）は、ソフトマックスの低エントロピー特性を模倣するために、学習可能な写像と特定の関数（指数関数、冪関数）を導入している。しかし、これらの手法は、学習されたカーネルファミリーが広範なアテンションパターンを包含するのに十分な表現力を備えていることを保証しておらず、性能を損なう可能性がある。

核心となる課題は、線形計算量を維持しつつ、固定カーネルや弱学習可能なベースラインを凌駕する十分な表現力を提供できる、柔軟で学習可能なアテンションカーネルを構築することである。

手法

著者らは、Random Fourier Features (RFF) を用いて、データ駆動型でアテンションカーネルを学習する線形Transformerフレームワークである Flexformer を提案する。

理論的基礎

本手法は、スペクトル表現理論に基づいている。

1. 定常カーネル: ボルの定理（定理3.1）に基づき、定常カーネル $k(x-y)$ はスペクトル密度 $p(\omega)$ に対応する。標準的なRFFは、固定された分布（例：ガウス分布を用いたソフトマックス）から周波数 $\omega$ をサンプリングすることによってこれを近似する。
2. 非定常カーネル: ヤグロムの定理（定理3.2）に基づき、非定常カーネルは結合スペクトル密度を通じて表現できる。

Flexformer アーキテクチャ

Flexformerは、周波数を固定された分布からサンプリングする代わりに、スペクトル周波数をエンドツーエンドで最適化可能な学習可能なパラメータとして扱う。

• Flexformer$_s$ (定常): 周波数 $\{\omega_i\}$ は学習可能なパラメータである。これにより、モデルはデータから直接、最適なスペクトル密度 $p(\omega)$ を学習することができる。特徴写像は以下のように構成される：
  $$\tilde{\phi}_n(x) = \exp\left(\frac{\|x\|^2}{2\sqrt{d}}\right) \phi_n(x)$$
  ここで、$\phi_n(x)$ は学習可能な周波数のコサイン項とサイン項で構成される。この定式化は、ソフトマックスカーネルの不偏推定量を厳密に含んでおり、最適であればソフトマックスアテンションを復元できることを保証している。

• Flexformer$_n$ (非定常): 表現力を高めるため、著者らはヤグロムの定理を用いて非定常カーネルへとフレームワークを拡張した。これには、周波数のペア $\{(\omega_{1i}, \omega_{2i})\}$ と、固定の $2\sqrt{d}$ に代わるスケーリングパラメータ $\tau$ の学習が含まれる。特徴写像は以下のようになる：
  $$\tilde{\phi}_n(x) = \exp\left(\frac{\|x\|^2}{\exp(\tau)}\right) \phi_n(x)$$
  非定常バリアントは、理論的に定常バリアントよりも厳密に高い表現力を持つ。

計算量

両方のバリアントは、シーケンス長 $N$ に対して線形時間計算量 $O(Ndd')$および**線形空間計算量**$O(Nd + Nd' + dd')$を維持する。ここで、$d$は隠れ次元、$d'$ は特徴写像の次元である。

主な貢献

1. 学習可能なカーネルフレームワーク: ランダムフーリエ特徴量の枠組み内で学習可能なスペクトル周波数を利用するFlexformerを提案し、データ駆動型で広範なアテンションカーネルのファミリーを学習することを可能にした。
2. 理論的表現力: Flexformerのカーネルファミリーがソフトマックスカーネルを厳密に包含していること（定常バリアントの場合）、および非定常バリアントにおいてより高い表現力を提供することを証明し、固定または経験的に設計されたカーネルの限界に対処した。
3. 蒸留と転移可能性: Flexformerが、事前学習済みTransformerから効果的に蒸留してソフトマックスアテンションの性能を回復できること、および強力なクロスドメインのカーネル転移性を示すことを示した。

実験結果

著者らは、言語モデリング、シーケンス分類、および蒸留タスクにおいてFlexformerを評価した。

• Long-Range Arena (LRA): Flexformerは、すべてのタスク（ListOps, Text, Retrieval, Image, Pathfinder）において最高の平均精度を達成した。具体的には、Flexformer$_n$は文書検索（document retrieval）、画像分類（image classification）、Pathfinderにおいて最高の成績を収め、既存の最良のリニアアテンション・ベースラインを平均精度で相対的に4.4%上回った。
• 自己回帰言語モデリング (WikiText-103):
  • 小規模モデル（41M パラメータ）において、Flexformerの各バリアントは他のすべてのリニアアテンション・ベースラインを上回り、バニラTransformerの性能に接近した。
  • 大規模モデル（247M パラメータ）において、Flexformer$_n$はバニラTransformerを凌駕した。これは、学習可能なカーネルが増大したモデル容量から大きな恩恵を受けることを示唆している。
• 効率性: 文書検索タスクにおいて、FlexformerはバニラTransformerと比較して2.6倍のトレーニング高速化と84%のメモリ使用量削減を達成しつつ、時間およびメモリにおける線形スケーリングを維持した。
• 蒸留 (GLUEベンチマーク): 事前学習済みRoBERTaモデルを蒸留する際、Flexformer$_n$はソフトマックスアテンションの挙動をほぼ完全に復元し、非学習型のリニア手法（Performer, RFA）を上回り、CoLAデータセットにおいては標準的なソフトマックスさえも上回った。
• カーネルの転移可能性: クロスドメイン実験において、Flexformer$_n$によって学習されたカーネルは、Hedgehogによって学習されたものよりもGLUEの他のデータセットへより効果的に転移し、アウトオブドメインのタスクにおける性能低下が小さいことが示された。

意義と主張

本論文は、Flexformerがソフトマックスアテンションに代わる効果的かつスケーラブルな選択肢であることを主張している。その主な意義は、線形アテンションにおける効率性と表現力の間の緊張関係を解消することにある：

• ソフトマックスが最適であるという仮定を超え、データから直接カーネルを学習する。
• 経験的なヒューリスティックに頼るのではなく、スペクトル表現理論（via spectral representation theory）を通じて、柔軟なカーネルを構築するための原理的なアプローチを提供する。
• 線形アテンションが、特にスケールアップまたは蒸留を行う際に、計算効率を維持しながら二次的アテンションの性能に匹каけ、あるいはそれを超えることができることを実証した。

限界: 著者らは、Flexformerは学習されたアテンションカーネルが正定値であることを前提としていると述べている。正定値カーネルが理論的にアテンションにとって最適な選択であるかどうかは、今後の研究における未解決の課題である。