Published on

大規模言語モデル推論効率化 新しい注意機構MFAによるKVキャッシュ削減

著者
  • avatar
    名前
    Ajax
    Twitter

大規模言語モデルにおける推論効率の課題

大規模言語モデル(LLM)の利用拡大に伴い、効率的な大規模推論が重要な課題となっています。特に、従来の注意機構におけるKey-Value (KV)キャッシュは、バッチサイズやシーケンス長に比例して増大し、LLMのスケールアップを妨げる「メモリの浪費」となっています。

MQA、GQA、MLAなどの代替案が登場していますが、厳格なメモリ制約下での性能維持や、エンジニアリング上の複雑さ、互換性の問題といった課題が残されています。

新しい注意機構：Multi-matrix Factorization Attention (MFA)

Stepes、清華大学などの研究機関の研究者たちは、新しい注意機構アーキテクチャであるMulti-matrix Factorization Attention (MFA)とその変種であるMFA-Key-Reuse (MFA-KR)を発表しました。この機構は、言語モデルの推論コストを大幅に削減しつつ、性能を向上させます。

MFAとMFA-KRの主な特徴:

  • MLAを上回る性能
  • 従来のMHAと同等の性能を維持しつつ、KVキャッシュ使用量を最大93.7%削減
  • シンプルな設計、再現性の高さ、ハイパーパラメータへの低い感度
  • 様々なPos-embedding手法との互換性

MFAのアプローチと分析

研究チームは、注意機構の一般的な設計と容量を分析し、容量に関連する2つの重要な側面を特定しました。この分析に基づいて、新しい分析手法と設計原理が開発されました。

  • Generalized Multi-Head Attention (GMHA)：様々なMHA変種を理解するための統一フレームワークとして導入
  • 推論の観点からのKey-Valueの計算とストレージの検討
  • 分解の観点からのモデル容量の検証
  • Fully Parameterized Bilinear Attention (FPBA)：理論上の上限性能として確立
  • MHAとその変種はFPBAの低ランク分解であるという発見

MQAおよびMLAとの比較

分析は、Multi-Query Attention (MQA)とMulti-Head Latent Attention (MLA)という2つの代表的な改善スキームに焦点を当てました。

  • MQA：すべての注意ヘッドが同じKey-Valueパラメータセットを共有する、より積極的なパラメータ共有戦略を採用。メモリ使用量を削減しますが、モデルの表現力が影響を受ける可能性。
  • MLA：パラメータ圧縮のために共有潜在空間を導入しますが、実際の表現力は最小次元によって制限され、中間次元を増やしても性能は大幅に向上しない。

MFAの革新的な点

MFAの開発は、理論上の性能限界に近づきながらリソース消費を最小限に抑える注意機構を作成することを目標としていました。MFAの設計には、以下の3つの重要な革新が含まれています。

  1. モデル容量を最大化するために、注意ヘッドの数と次元を大幅に増加。
  2. 注意ヘッドの数と次元を拡大しながら、パラメータ効率を維持するために、積極的な低ランク分解戦略を採用。
  3. モデルの複雑さが増してもメモリ消費を最小限に抑えるために、単一のKey-Valueヘッド設計を利用。

容量の測定と比較

MFAと他の注意機構をさらに分析するために、研究チームは2つの重要なメトリックを導入しました。

  • Total Effective Rank (TER)：注意ヘッドの数とヘッドごとの分解ランク(FRH)の積。
  • Shared Latent Subspace Dimension (SLSD)：すべての注意ヘッドで共有される隠れ空間の次元。

比較結果:

  • MFAはMQAと比較して、より高いSLSDとTERを達成。
  • MLAと比較して、MFAは同等のパラメータ予算で、より小さいKVキャッシュサイズとより高いTERを達成しつつ、同等のSLSDを維持。
  • 従来のMHAと比較して、MFAはSLSDは小さいものの、より高いTERを達成。

実験結果

10Bから1Tのトレーニングデータと1Bから7Bのパラメータを持つモデルをテストし、新しいアーキテクチャのより大きなスケールでのパフォーマンスを評価するための広範な実験が行われました。

  • MFAは、従来のMHAと同等のスケーリング能力を示し、より大きなスケールでも優れたパフォーマンスを維持。
  • MFA-KRはわずかに低いパフォーマンスを示しましたが、そのスケーリング傾向はMHAと一致。
  • MFAとMFA-KRのメモリ節約の利点はモデルサイズとともに拡大し続け、MFAは最大スケールで87.5%のメモリ節約を達成し、MFA-KRはメモリ使用量を6.25%に削減。

アブレーション研究

アブレーション研究により、MFAとMFA-KRの有効性が検証されました。それらのパフォーマンス上の利点は、様々な主流のポジショナルエンコーディング手法でも確認されました。

今後の展望

MFAは、シンプルな設計で大幅な改善を提供し、エンジニアリングの複雑さを増すことなくLLM推論におけるメモリボトルネックを効果的に解決します。既存のTransformerエコシステムにシームレスに統合され、様々なシナリオでのLLMの応用を加速します。