技術的サマリー：高次線形注意機構（HLA）

問題定義

トランスフォーマーアーキテクチャにおける標準的なスケーリングドドット積注意機構は、シーケンス長 $n$ に対して $O(n^2)$ の計算量およびメモリ複雑性を生じさせる。この二次的なコストは、自己回帰言語モデルを長いコンテキストに拡張する際の主要なボトルネックである。既存の代替手段である線形注意機構や状態空間モデル（SSM）は、トークンあたりの $O(n)$ または $O(1)$ の複雑性を提供するが、通常は一次近似またはカーネルベースの手法に制限されている。これらの制限は、モデルの表現力を制約し、標準的な注意機構がモデル化できるトークン間の複雑な高次相互作用を捉えることを妨げる。

手法

本論文は、コンパクトなプレフィックス充分統計量を通じて高次相互作用を実現する、因果的かつストリーミング可能な機構である**高次線形注意機構（HLA）**を提案する。核心的な洞察は、高次注意様式演算子を低次モーメント（例えば、キーの外部積の和）に因数分解できることであり、これにより $n \times n$ の注意行列を具体化することなく、厳密な因果ストリーミングが可能となる点にある。

1. 二次 HLA

二次変種は、一次モーメントを捉える定数サイズの状態を維持することで、線形注意機構を一般化する。

• 状態維持: 単一の走行和の代わりに、HLA はキーモーメント $S^K_t = \sum_{i \le t} k_i k_i^\top$、クエリ - バリューアキュムレータ $C^{QV}_t = \sum_{i \le t} q_i v_i^\top$、およびクエリ質量 $m^Q_t = \sum_{i \le t} q_i$ を含むプレフィックス要約を維持する。
• 厳密な因果性: 完全な注意行列を構築することなく厳密な自己回帰的因果性（将来のトークンが現在の出力に影響を与えることを防ぐ）を強制するため、著者は拡張要約を導出する。具体的には、単純な積 $S^K_t C^{QV}_t$ における将来項の包含を補正するクロス要約 $G_t$ および $h_t$ を導入する。
• 出力定式化: 時刻 $t$ における正規化されていない出力は $o_t = q_t^\top (S^K_t C^{QV}_t - G_t)$ として計算される。この定式化により、シーケンス長に依存せず、トークンあたり $O(d^2 + d d_v)$ の計算が可能となる。
• 正規化: 同様の拡張要約から導出されたマスクされた分母で分子を割る、比正規化された変種も定義される。

2. 非対称高次線形注意機構（AHLA）

本論文は、演算子 $Q(K^\top Q)(K^\top V)$ を実装する相補的な非対称変種AHLAを提案する。この構造は、中間キーインデックスによってルーティングされる単一パスを通じて値に重み付けを行うものであり、標準的な HLA の対称的な $AA^\top V$ 構造とは異なる。これは、一連の異なるストリーミング要約（$P^{KV}_t, m^K_t, E_t, n_t$）を利用し、同じ線形なトークンあたり複雑性を達成する。

3. 三次 HLA

この枠組みは、$A = QK^\top$ である行列積 $AA^\top A$ として定義される三次相互作用に拡張される。

• ストリーミング代数: 三次のマスク付き出力には、厳密な因果性を強制するために、3 つの主要なプレフィックスモーメント（$S^K_t, S^Q_t, P^{KV}_t$）および 6 つの追加クロス要約（$G^{(1)}_t, \dots, G^{(3)}_t$ および $h^{(1)}_t, \dots, h^{(3)}_t$）の維持が必要となる。
• 複雑性: オンライン更新コストは特徴次元に対して線形のままだが、高次モーメントを収容するために状態サイズは増加する。

4. 連関スキャンによる並列トレーニング

GPU 上での純粋な再帰モデルのトレーニングの非効率性に対処するため、著者はチャンク並列トレーニング方式を開発する。

• 連関演算子: トークンのセグメントを連結するための連関演算子（具体的には半直積）を定義する。これらの演算子は、チャンクを結合する際に生じるセグメント間のクロス項（例：$S^K_B C^{QV}_A$）を考慮する。
• スキャン同等性: これらのセグメント要約に対して排他的な Blelloch スキャンを適用することで、本手法は直列再帰のフォワード活性化と厳密に同一のものを生成する。これにより、近似された時間逆伝播に頼ることなく、効率的な並列トレーニングが可能となる。
• 三次スキャン: 三次の場合、スキャン演算子は、三次修正状態の合成を厳密に処理するためのセグメントレベルの線形写像（$M^{KQP}, M^{KQm}$）によって拡張される。

主要な貢献

1. 厳密なマスク付き二次ストリーミング: 本論文は、厳密に因果的な二次 HLA を定数あたりのトークンコストで生み出す、拡張要約の完全な代数を提供する。これには、正規化されていない変種と比正規化された変種を区別し、マスク付きストリーミング恒等式およびオンライン更新を確立する形式的な証明が含まれる。
2. 直列活性化と一致する連関スキャン: 著者は、マスクなしおよびマスクあり（指数減衰の有無を含む）の両方の設定に対する連関演算子を定義する。これらの演算子に対する標準的な排他的スキャンが、直列再帰と同一のフォワード活性化を生み出すことを証明し、効率的な並列トレーニングを可能にする。
3. 三次拡張: 本作業は、完全なマスク付き三次状態、オンライン更新、厳密に因果的なストリーミングカーネル、および正確なチャンク並列アルゴリズムを提示する。これには、三次状態を厳密に合成するために必要なセグメント写像の導出が含まれる。
4. 非対称変種（AHLA）: 新しい非対称高次注意機構が導入され、ストリーミングおよび複雑性の特性を維持しつつ、異なる帰納的バイアスを提供する。

結果と主張

本論文は、特定のデータセットにおける経験的ベンチマークよりも、主にアルゴリズム構造と理論的性質に焦点を当てている。主張は数学的導出に基づいている。

• 複雑性: HLA は、推論においてトークンあたり $O(d^2 + d d_v)$ の時間およびメモリを達成する。ここで、$d$ はクエリ/キー次元、$d_v$ はバリュー次元である。これはシーケンス長 $n$ に依存しない。
• 厳密性: チャンク並列トレーニング方式は、他の線形注意トレーニング手法で一般的である近似誤差を回避し、直列再帰の活性化を厳密に再現することが証明されている。
• 因果性: 提案されたマスク付き定式化は厳密に自己回帰的因果性を強制し、時刻 $t$ における出力が $j \le t$ のトークンのみに依存することを保証する。
• スケーラビリティ: $n \times n$ 行列の具体化を回避し、連関スキャンを利用することで、HLA は長いコンテキストモデルのためのスケーラブルな構築ブロックとして位置づけられる。

重要性

本論文は、HLA を、注意様式でデータ依存混合と現代の再帰アーキテクチャの効率性の間のギャップを埋める「原理的かつスケーラブルな構築ブロック」として位置づけている。

• 表現力対効率性: 固定された特徴マップに依存する一次線形注意機構とは異なり、HLA はデータ依存のメトリクス（$S^K_t$ を通じて）を学習し、$O(1)$ のトークンあたり状態更新パラダイムを維持しながら、高次相互作用（2 次または 3 次多項式）を捉える。
• 実装: この作業は、HLA がトランスフォーマーブロック内の標準的な注意サブレイヤーのドロップイン置換として機能し、位置エンコーディングとマスクロジックを維持しながらカーネルの交換のみを必要とすることを強調している。
• 理論的基盤: 高次項に対する厳密なストリーミング恒等式と連関スキャン演算子を提供することで、本論文は、因果制約やトレーニング効率を犠牲にすることなく、線形注意機構を一次を超えて拡張するための厳密な代数的枠組みを確立している。