SQUEEZED ATTENTION: 長文脈LLM推論を加速する新手法

今回は、長文脈を扱う大規模言語モデル（LLM）の効率性を飛躍的に向上させる新しい手法「SQUEEZED ATTENTION」を紹介します。本論文は、固定文脈を効率的に処理することで、計算コストとメモリ使用量を大幅に削減し、長文脈タスクにおける推論効率を最適化することを目指しています。

論文情報

タイトル: SQUEEZED ATTENTION: Accelerating Long Context Length LLM Inference
リンク: arXiv
発表日: 2024年11月14日
著者: Coleman Hooper, Sehoon Kim, Hiva Mohammadzadeh, Monishwaran Maheswaran, June Paik, Michael W. Mahoney, Kurt Keutzer, Amir Gholami
DOI: 10.48550/arXiv.2411.09688

LLMの文脈長が増加する中で、多くの応用タスクにおいて以下の課題が顕在化しています：

メモリ使用量の増加:
- 注意計算（Attention Mechanism）は文脈長 (L) に比例して計算コストが増加します。たとえば、文脈長が32Kを超えると、モデル推論に必要なKVキャッシュが膨大なメモリを占有します。
固定文脈の非効率利用:
- 実際のアプリケーションでは、多くの入力文脈が固定されており、これを逐次計算することはリソースの無駄につながります。
既存手法の限界:
- QUESTやSparse Attentionは長文脈の一部をスパース化するアプローチを取っていますが、計算コストの大幅な削減には至っていません。

本研究の目的は、固定文脈を効率的に利用する「SQUEEZED ATTENTION」を提案し、以下の目標を達成することです：

提案手法は、以下の3つのステップで構成されています。

目的: 固定文脈のキー（Key）をセマンティックに類似したグループに分割し、各グループの代表（セントロイド）を計算します。
アルゴリズム:
- K-meansクラスタリングを適用し、クラスタ (C_j) を生成。
- 各クラスタのセントロイドを以下で計算：
  $$ C_j = \frac{1}{|K_j|} \sum_{k \in K_j} k $$
- セントロイドは注意計算時に固定文脈の代表として使用されます。

クエリ処理: 推論時にユーザー入力（クエリ）が与えられると、以下の手順でキー検索を実行します。
- クエリ (q) とセントロイド (C_j) の内積を計算。
  $$ S_i = \frac{\exp(q \cdot C_i)}{\sum_{j} N_j \cdot \exp(q \cdot C_j)} $$
- スコアが閾値を超えたクラスタのみを使用し、注意計算を実行します。

SQUEEZED ATTENTIONは、長文脈LLMにおける推論効率の課題に対し、革新的な解決策を提供しました。特に固定文脈の有効活用という観点で、他手法と一線を画す性能を示しています。

今後の研究課題として以下が挙げられます：

この記事が、読者の皆さんの研究や実務に役立つことを願っています。ご質問やフィードバックは、ぜひコメント欄にお寄せください！