Beyond Homogeneous Attention: Memory-Efficient LLMs via Fourier-Approximated KV Cache

要約

大規模な言語モデルは、コンテキストの長さが増加するにつれて、成長するキー価値（kV）キャッシュからのメモリの要求に苦しんでいます。
既存の圧縮方法は、頭の寸法を均一にするか、注意誘導トークンの剪定に依存し、多くの場合、精度を犠牲にしたり、計算オーバーヘッドを導入したりします。
ForierAttentionを提案します。これは、トランスヘッド寸法の不均一な役割を活用するトレーニングフリーのフレームワークです。低次元はローカルコンテキストに優先順位を付け、上部が長距離依存関係をキャプチャします。
長いコンテキストに無感覚な寸法を直交フーリエ塩基に投影することにより、フーリエラットは、固定長のスペクトル係数を持つ時間的進化を近似します。
Llamaモデルでの評価は、フーリエラットがロングベンチとヘイスタック（NIAH）で最高の長いコンテストの精度を達成することを示しています。
また、カスタムトリトンカーネルであるFlashFourierAttentionは、合理化された読み取りワイト操作を介してメモリを最適化するように設計されており、パフォーマンスの妥協なしに効率的な展開を可能にします。

要約(オリジナル)

Large Language Models struggle with memory demands from the growing Key-Value (KV) cache as context lengths increase. Existing compression methods homogenize head dimensions or rely on attention-guided token pruning, often sacrificing accuracy or introducing computational overhead. We propose FourierAttention, a training-free framework that exploits the heterogeneous roles of transformer head dimensions: lower dimensions prioritize local context, while upper ones capture long-range dependencies. By projecting the long-context-insensitive dimensions onto orthogonal Fourier bases, FourierAttention approximates their temporal evolution with fixed-length spectral coefficients. Evaluations on LLaMA models show that FourierAttention achieves the best long-context accuracy on LongBench and Needle-In-A-Haystack (NIAH). Besides, a custom Triton kernel, FlashFourierAttention, is designed to optimize memory via streamlined read-write operations, enabling efficient deployment without performance compromise.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google