MoA: Mixture of Sparse Attention for Automatic Large Language Model Compression

要約

スパース注意は、長い文脈における大規模言語モデル（LLM）のメモリとスループットの大きな要求を効果的に軽減することができる。既存の手法では一般的に、異なる注意ヘッドと入力長に対して同じスパースパターンを適用する、一様なスパース注意マスクを採用している。しかし、この画一的なアプローチでは、LLMに固有の多様な注意パターンを捉えることができず、精度と遅延のトレードオフを無視することになる。この課題を解決するために、我々は、異なるヘッドとレイヤーに異なるスパースアテンション構成を自動的に調整するミクスチャーオブアテンション（MoA）を提案する。MoAは様々な注意パターンと、入力シーケンスの長さに対するそのスケーリングルールからなる探索空間を構築し、ナビゲートする。MoAはモデルをプロファイリングし、潜在的な構成を評価し、最適なスパースアテンション圧縮計画を特定する。MoAは様々な入力サイズに適応し、ある注意ヘッドはより長いシーケンスに対応するためにその焦点を拡大し、他のヘッドは一貫して固定長の局所コンテキストに集中することを明らかにする。実験によると、MoAは同じ平均注意スパンで有効コンテキスト長を$3.9times$増加させ、Vicuna-{7B,13B}とLlama3-{8B,70B}モデルにおいて、一様注意ベースラインよりも検索精度を$1.5-7.1times$向上させた。さらに、MoAはスパースモデルとデンスモデル間の能力差を縮め、2つの長文脈理解ベンチマ ークにおいて、最大相対性能低下を$9%-36%から$5%以内に抑えた。MoAは1.2-1.4times$のGPUメモリ削減を達成し、FlashAttention2とvLLMと比較して、デコードスループットを$6.6-8.2times$と$1.7-1.9times$向上させ、性能への影響は最小です。私たちのコードは⦿URL{https://github.com/thu-nics/MoA}で入手可能です。

要約(オリジナル)

Sparse attention can effectively mitigate the significant memory and throughput demands of Large Language Models (LLMs) in long contexts. Existing methods typically employ a uniform sparse attention mask, applying the same sparse pattern across different attention heads and input lengths. However, this uniform approach fails to capture the diverse attention patterns inherent in LLMs, ignoring their distinct accuracy-latency trade-offs. To address this challenge, we propose the Mixture of Attention (MoA), which automatically tailors distinct sparse attention configurations to different heads and layers. MoA constructs and navigates a search space of various attention patterns and their scaling rules relative to input sequence lengths. It profiles the model, evaluates potential configurations, and pinpoints the optimal sparse attention compression plan. MoA adapts to varying input sizes, revealing that some attention heads expand their focus to accommodate longer sequences, while other heads consistently concentrate on fixed-length local contexts. Experiments show that MoA increases the effective context length by $3.9\times$ with the same average attention span, boosting retrieval accuracy by $1.5-7.1\times$ over the uniform-attention baseline across Vicuna-{7B,13B}, and Llama3-{8B,70B} models. Moreover, MoA narrows the capability gaps between sparse and dense models, reducing the maximum relative performance drop from $9\%-36\%$ to within $5\%$ across two long-context understanding benchmarks. MoA achieves a $1.2-1.4\times$ GPU memory reduction, boosting decode throughput by $6.6-8.2\times$ and $1.7-1.9\times$ compared to FlashAttention2 and vLLM, with minimal impact on performance. Our code is available at \url{https://github.com/thu-nics/MoA}.

arxiv情報

提供元, 利用サービス

arxiv.jp, DeepL