【RetrievalAttention】長い文脈を処理するLLMの効率を向上！
3つの要点
✔️ 長文脈を扱う大規模言語モデル（LLM）の推論速度を向上させる「RetrievalAttention」という手法を提案
✔️ 膨大なメモリ使用量と計算コストを削減しつつ、高い精度を維持したまま高速な推論が可能になった
✔️ 長文脈タスクでの推論効率を大幅に改善
RetrievalAttention: Accelerating Long-Context LLM Inference via Vector Retrieval
written by Di Liu, Meng Chen, Baotong Lu, Huiqiang Jiang, Zhenhua Han, Qianxi Zhang, Qi Chen, Chengruidong Zhang, Bailu Ding, Kai Zhang, Chen Chen, Fan Yang, Yuqing Yang, Lili Qiu
(Submitted on 16 Sep 2024)
Comments: 16 pages
Subjects: Machine Learning (cs.LG); Computation and Language (cs.CL)
code：
本記事で使用している画像は論文中のもの、紹介スライドのもの、またはそれを参考に作成したものを使用しております。
概要
この論文の研究背景は、主に「長い文脈を持つ大規模言語モデル（LLM）」が抱える問題を解決することを目的としています。
TransformerベースのLLMは、多様な分野で広く利用されていますが、長い文脈を扱う際の計算コストが非常に高くなります。特に「アテンション」という計算は、文脈が長くなるほど処理時間とメモリの使用量が増大し、これがボトルネックとなります。この課題を解決するために多くの技術が開発されていますが、どれも完全な解決には至っていません。
本研究の手法は、既存のアテンション最適化技術と比べ、ほぼ同等の精度を維持しながら、長い文脈に対して非常に効率的な推論を可能にするという点で優れています。特に、大きなモデル（8Bパラメータモデル）を比較的低スペックなGPUで効率的に動作させることが可能です。
要するに、RetrievalAttentionは、長い文脈を扱うLLMの推論において、メモリと時間の効率を劇的に改善する技術であり、実用化に向けた大きなステップとなる技術です。
研究背景
大規模言語モデルは、非常に長いテキストデータを処理できるため、自然言語処理の様々なタスクで非常に高い性能を発揮しています。例えば、膨大な文章を読んで、その内容に基づいて回答を生成したり、要約を行ったりすることが可能です。しかし、こうしたモデルの中核を担う「アテンションメカニズム」には大きな課題があります。
アテンションメカニズムは、入力されたテキストのどの部分が重要かを判断し、それに基づいて次の単語を予測するための手法です。しかし、アテンションの計算量は「クエリ」と「キー・バリュー」ベクトルという2つのベクトルセットを比較して行われるため、テキストが長くなると、計算量が急激に増加してしまいます。これにより、推論速度が低下し、メモリ使用量も膨大になるという問題が発生します。特に、文脈が非常に長い場合には、GPUのメモリがすぐに限界に達してしまうことが大きなボトルネックとなっています。
従来の解決策として、「KVキャッシング」という技術が使われてきました。これは、計算に必要な過去の「キー」と「バリュー」の状態を保持し、再利用することで無駄な計算を省く手法です。しかし、この方法でも、長い文脈を処理する際には大量のメモリを消費してしまいます。例えば、1つのGPUで大量のトークンを処理しようとすると、500GB以上のメモリが必要になることもあります。これでは、現実的なシステムでの利用が難しいため、より効率的な手法が求められていました。
そこで、この論文では、アテンションメカニズムにおける「動的なスパース性」という特徴に注目しています。実際には、すべてのトークンが次の単語の予測に重要なわけではなく、特定の一部のトークンのみが重要な役割を果たしています。つまり、全てのトークンを計算に含める必要はなく、重要なトークンだけに焦点を当てて計算すれば、計算コストを大幅に削減できるのではないか、という発想です。
この背景を基に、新しい手法である「RetrievalAttention」が提案されており、効率的なアテンション計算の実現を目指しています。
提案手法
ransformerモデルでは、アテンション機構が入力されたテキストの中でどの部分が重要かを判断し、それに基づいて次のトークンを予測します。しかし、文脈が長くなると、すべてのトークンに対してアテンションを計算するため、計算量が膨大になります。
RetrievalAttentionの大きな特徴は、クエリとキーのベクトル間の分布の違い（OOD問題）に対応している点です。通常の近似最適検索では、クエリとキーが同じ分布に属していることが前提ですが、アテンション計算においては、クエリベクトルとキー・バリューベクトルが異なる分布を持つことが多く、これがパフォーマンスの低下につながります。この問題を解決するために、RetrievalAttentionでは、アテンション特有の分布に適応する新しい検索アルゴリズムを使用しています。この手法により、クエリに対してわずか1〜3%のデータしかスキャンしなくても、高精度なアテンション結果を得ることが可能となっています。
また、RetrievalAttentionは、GPUとCPUのメモリをうまく使い分けています。具体的には、重要な「キー・バリュー」ベクトルはGPUに保持し、残りのデータはCPUにオフロードすることで、GPUのメモリ消費を抑えつつ、計算の効率を保っています。
「RetrievalAttention」は、このアテンション計算を効率化するために、次の2つの大きなアイデアを使っています：
動的スパース性の活用
アテンション計算では、全てのトークンが等しく重要なわけではなく、実際には一部のトークンのみが次のトークンの予測に重要な役割を果たしています。これを「動的スパース性」と呼びます。「RetrievalAttention」は、この特性を利用して、重要なトークンだけに注目し、その他のトークンを省略します。
ベクトル検索による最適化
次に、すべてのトークンを対象にするのではなく、近似的に最も重要なトークンを選び出すために「近似最適検索（ANNS）」という技術を使います。この技術は、膨大なデータから重要なデータを高速に検索するもので、通常のアテンション計算と比べて計算量を大幅に減らします。
実験
この論文で提案されている「RetrievalAttention」の実験では、長い文脈を扱う大規模言語モデル（LLM）の推論効率を向上させる手法の有効性を検証しています。実験では、複数の大規模モデルとベンチマークを使用し、提案手法がどれだけ優れた性能を発揮するかを詳細に評価しています。
まず、実験環境として、NVIDIA RTX 4090 GPU（24GBメモリ）を使用し、Llama-3-8BやYi-6Bといった複数のLLMが対象になっています。これらのモデルは、それぞれ128,000トークンまでの長文脈を処理する能力を持っています。実験の目的は、RetrievalAttentionが他の手法と比較して、推論精度を維持しながらどれだけ高速化できるかを確認することです。
実験では、精度と速度の両面で提案手法が評価されています。まず、精度の観点では、RetrievalAttentionはフルアテンションとほぼ同等のパフォーマンスを示しました。これは、重要なトークンのみを効率的に抽出することで、モデルの推論結果に影響を与えることなく、計算コストを削減できていることを意味しています。この結果は、ベンチマークタスクである∞-Benchの結果からも確認できます。
一方、推論速度に関しても大幅な改善が見られました。特に、128,000トークンの長い文脈を処理する際、RetrievalAttentionは従来のフルアテンションに比べて5倍近く速く推論が行えることが確認されました。この速度向上は、不要なトークンへのアクセスを大幅に減らすことで実現されています。Needle-in-a-haystackのタスクでは、膨大なデータの中から特定の情報を抽出する際の効率性が特に際立っており、RetrievalAttentionの優れた検索性能が示されています。
また、RetrievalAttentionは、GPUメモリの使用量の点でも優れています。通常、長い文脈を処理するには膨大なメモリが必要となりますが、提案手法は16GBのGPUメモリで128,000トークンを処理できるように設計されており、これによりハードウェアコストを抑えつつ、効率的な推論が可能です。
このように、RetrievalAttentionは精度を維持しつつ、推論速度とメモリ効率を大幅に向上させる手法であり、特に長文脈を扱うタスクにおいては、実用的な性能を発揮しています。
まとめ
この論文の結論は、提案した「RetrievalAttention」手法が、長文脈を扱う大規模言語モデル（LLM）の推論を効率化するために非常に有効であるということです。
通常のアテンション計算では、膨大な量のトークンを扱うため、時間やメモリのコストが大きくなりがちです。特に、文脈が長くなればなるほど、アテンション計算の負荷は指数的に増えてしまいます。しかし、RetrievalAttentionは、アテンションにおける「動的なスパース性」に着目し、必要な部分だけを効果的に処理することにより、この問題を解決しました。
この手法は、動的に重要なトークンを選び出し、従来の手法に比べて少ない計算量とメモリ使用量で同等の精度を達成しています。また、実験の結果、RetrievalAttentionを使うことで、推論の速度が大幅に改善され、特に長文脈タスクにおいては、最大で5倍の高速化が確認されました。
さらに、RetrievalAttentionは、限られたGPUメモリでも長い文脈を効率的に処理できるため、メモリ使用量の削減にも成功しています。この特徴により、従来は非常に高価なハードウェアが必要だったタスクを、より手頃な環境でも実現できる可能性が広がりました。
結論として、RetrievalAttentionは、長文脈を扱う際の推論コストを大幅に削減しつつ、精度を維持するための強力な手法であり、今後のLLMの発展において非常に重要な技術であると言えます。
この記事に関するカテゴリー
