ホモジニアスな注意機構を超えて：フーリエ近似KVキャッシュによるメモリ効率的なLLM

コアコンセプト

FourierAttentionは、トランスフォーマーのヘッド次元が異種的な役割を持つことを発見することで、Key-Valueキャッシュからのメモリボトルネックに対処します。下位の次元はローカルコンテキストを処理し、上位の次元は長距離依存関係をキャプチャします。本手法は、直交フーリエ基底関数を使用してコンテキスト非依存な次元を圧縮し、優位なスペクトル係数のみを保持しながら長コンテキスト性能を維持します。

アーキテクチャ概要

• 次元分析ステージ: ノイズ摂動実験を通じて、完全保持が必要なヘッド次元と圧縮対象の次元を識別します
• フーリエ圧縮: プリフィリング中に平行移動フーリエ変換を使用してシーケンスを固定長スペクトル表現に変換します（k ≪ L係数）
• 選択的保持: 初期トークン、最近のローカルトークン、長コンテキスト依存次元は非圧縮のまま保持し、重要度の低い次元の中間範囲トークンを圧縮します
• カスタムTritonカーネル: デコード中のメモリ効率を最適化し、パフォーマンスを損なわない実装を提供します

実装

ステップ1: 注意ヘッドの次元を分析

ノイズ摂動を使用して、コンテキスト非依存で圧縮可能な

...