最大400万のトークンコンテキストと22倍の推論速度を備えたStreamingLLMは人気を博し、GitHubで2.5Kのスターを獲得しました。

会話型 AI ロボットとのコミュニケーションを経験したことがあるなら、非常にイライラした瞬間を間違いなく思い出せるはずです。例えば、前日の会話であなたが言った要点はAIによって完全に忘れられていました...

これは、現在の LLM のほとんどは、試験のために詰め込み勉強する学生と同じように、限られたコンテキストしか記憶できず、少し質問するだけでその内容が明らかになってしまうためです。

AI アシスタントがチャット中に数週間または数か月前の会話を文脈的に参照できるとしたらどうでしょう。あるいは、AI アシスタントに何千ページものレポートを要約するよう依頼できるとしたら、それはうらやましい機能ではないでしょうか。

研究者たちは、LLM がより多くのことをより良く記憶できるようにするために絶えず取り組んでいます。最近、MIT、Meta AI、CMU の研究者らは、言語モデルが無限のテキストをスムーズに処理できるようにする「StreamingLLM」と呼ばれる手法を提案しました。

• 論文アドレス: https://arxiv.org/pdf/2309.17453.pdf
• プロジェクトアドレス: https://github.com/mit-han-lab/streaming-llm

StreamingLLM は、モデルの固有の「注意シンク」を識別して保存し、その推論の初期トークンを固定することによって機能します。最近のトークンのローリング キャッシュと組み合わせることで、StreamingLLM は精度を犠牲にすることなく推論を 22 倍高速化します。わずか数日で、このプロジェクトは GitHub プラットフォームで 2.5K 個のスターを獲得しました。

具体的には、StreamingLLM により、言語モデルは、前回の試合のタッチダウン スコア、新生児の名前、長い契約、討論の内容などを正確に記憶できるようになります。これは、AI アシスタントのメモリをアップグレードして、より重いタスクを完璧に処理できるようにするのと同じです。

それでは技術的な詳細を見てみましょう。

方法の革新

通常、LLM は事前トレーニング中の注意ウィンドウによって制限されます。このウィンドウ サイズを拡張し、トレーニングと推論の効率を向上させるための多くの作業が行われてきましたが、LLM の許容シーケンス長はまだ限られており、長期的な展開には適していません。

この論文では、研究者らはまず LLM ストリーム アプリケーションの概念を紹介し、「効率とパフォーマンスを犠牲にすることなく、無限に長い入力で LLM を展開できるか」という疑問を提起しました。

LLM を無限長の入力ストリームに適用する場合、主に 2 つの課題があります。

1. デコードフェーズでは、トランスフォーマーベースの LLM は、図 1 (a) に示すように、以前のすべてのトークンのキーと値の状態 (KV) をキャッシュします。これにより、メモリの使用量が過剰になり、デコード待ち時間が長くなる可能性があります。

2. 既存のモデルでは長さの外挿機能が限られているため、シーケンスの長さが事前トレーニング中に設定された注意ウィンドウのサイズを超えると、パフォーマンスが低下します。

直感的な方法は、ウィンドウ アテンションと呼ばれます (図 1 b を参照)。この方法では、最新のトークンの KV 状態に固定サイズのスライディング ウィンドウを維持するだけです。キャッシュがいっぱいになった後も安定したメモリ使用量とデコード速度を確保できますが、シーケンスの長さがキャッシュ サイズを超えると、最初のトークンの KV を削除するだけでもモデルが崩壊します。もう 1 つの方法は、スライディング ウィンドウを再計算することです (図 1 c を参照)。この方法では、生成されたトークンごとに最新のトークンの KV 状態が再構築されます。パフォーマンスは優れていますが、ウィンドウ内でセカンダリ アテンションを計算する必要があるため、大幅に遅くなり、実際のストリーム アプリケーションには適していません。

ウィンドウ注意の失敗を理解する過程で、研究者は自己回帰 LLM の興味深い現象を発見しました。図 2 に示すように、初期トークンには、これらのトークンが言語モデリング タスクに関連しているかどうかに関係なく、多数の注意スコアが割り当てられます。

研究者たちはこれらのトークンを「注目プール」と呼んでいます。意味論的な意味はないものの、注目スコアの大部分を占めています。研究者たちは、この現象はソフトマックス（すべてのコンテキスト トークンの注目スコアの合計が 1 になることを要求する）によるものだと考えています。現在のクエリが以前の多くのトークンに強い一致がない場合でも、モデルはこれらの不要な注目値をどこかに割り当てて、合計が 1 になるようにする必要があります。最初のトークンが「プール」になる理由は直感的です。自己回帰言語モデリングの特性により、最初のトークンは後続のほぼすべてのトークンから見えるため、注意プールとしてトレーニングしやすくなります。

上記の洞察に基づいて、研究者らは、限られた注意ウィンドウでトレーニングされた注意モデルが微調整なしで無限に長いテキストを処理できるようにするシンプルで効率的なフレームワークである StreamingLLM を提案しました。

StreamingLLM は、注目プールの注目値が高いという事実を活用し、これらの注目プールを保持することで注目スコア分布を正規分布に近づけることができます。したがって、StreamingLLM では、アテンション プール トークンの KV 値 (必要な初期トークンは 4 つだけ) とスライディング ウィンドウの KV 値を保持して、アテンション計算を固定し、モデルのパフォーマンスを安定させるだけで済みます。

StreamingLLMを使用すると、Llama-2-[7,13,70] B、MPT-[7,30] B、Falcon-[7,40] B、Pythia [2.9,6.9,12] Bなどのモデルは、400万以上のトークンを確実にシミュレートできます。

唯一の実行可能なベースラインであるスライディング ウィンドウの再計算と比較すると、StreamingLLM はパフォーマンスを損なうことなく 22.2 倍高速です。

評価

実験では、図 3 に示すように、20K トークンの範囲のテキストで、StreamingLLM のパープレキシティは Oracle ベースライン (スライディング ウィンドウの再計算) のパープレキシティに匹敵します。一方、入力長が事前トレーニング ウィンドウを超えると、高密度アテンションは失敗しますが、入力長がキャッシュ サイズを超えると、ウィンドウ アテンションは停止し、初期トークンが削除されます。

図 5 は、StreamingLLM がさまざまなモデル ファミリとスケールにわたって、400 万を超えるトークンを含む異常なサイズのテキストを確実に処理できることをさらに確認しています。これらには、Llama-2-[7,13,70] B、Falcon-[7,40] B、Pythia-[2.8,6.9,12] B、およびMPT-[7,30] Bが含まれます。

その後、研究者らは「注意プール」仮説を確認し、言語モデルは事前トレーニング可能であり、ストリーミング方式で展開する場合、1 つの注意プール トークンのみが必要であることを証明しました。具体的には、指定された注意プールとして、すべてのトレーニング サンプルの始めに追加の学習可能なトークンを追加することを提案しました。研究者らは、1億6000万のパラメータを持つ言語モデルを最初から事前トレーニングすることで、この方法がモデルのパフォーマンスを維持できることを実証した。これは、同じレベルのパフォーマンスを達成するために複数の初期トークンを注意プールとして再導入する必要がある現在の言語モデルとはまったく対照的です。

最後に、研究者らは、StreamingLLM のデコード待ち時間とメモリ使用量をスライディング ウィンドウの再計算と比較し、Llama-2-7B および Llama-2-13B モデルを使用して単一の NVIDIA A6000 GPU でテストしました。図 10 に示すように、バッファ サイズが増加すると、StreamingLLM のデコード速度は直線的に増加します。後者のデコード遅延は二次曲線で増加します。実験では、StreamingLLM はトークンあたり最大 22.2 倍の驚異的な高速化を実現することが示されています。

研究の詳細については原著論文を参照してください。