研究者たちは、ChatGPTのような大規模言語モデルにおいて、そうでなければモデルの性能に深刻な影響を及ぼす難解な問題に対する簡潔で効果的な解決策を発見した。
複数ラウンドの連続対話を含むAI会話では、ChatGPTのようなチャットボットを駆動する強力で大規模な言語機械学習モデルが突然失敗し、ボットのパフォーマンスが劇的に低下することがあります。
マサチューセッツ工科大学（MIT）などの研究チームは、この問題の意外な原因を発見し、チャットボットがクラッシュしたり速度が落ちたりすることなく会話を続けられるようにする簡単な解決策を考え出した。
彼らのアプローチは、多くの大規模言語モデルの中心にあるキーキャッシュ（対話のメモリストアに相当）を適応させる。キャッシュが実際の容量よりも多くの情報を保存する必要がある場合、最も古いデータが置き換えられ、モデルの故障につながることがある。
最初のデータポイントのいくつかをメモリに保持することで、研究者たちのスキームでは、チャットボットはどれだけ会話が続いてもコミュニケーションを続けることができる。
StreamingLLMと呼ばれるこのアプローチは、対話が400万語を超えてもモデルの高い効率を維持する。以前の対話の一部を常に再計算することでシステムクラッシュを回避する別のアプローチと比較すると、StreamingLLMは22倍以上高速に動作する。
これにより、チャットボットは再起動することなく、一日中長い会話を続けることができるようになり、コードを書いたり、編集したり、生成したりする作業のためのAIアシスタントを効果的に提供することができるようになる。
「今日、このアプローチによって、これらの大規模な言語モデルを継続的に展開することができます。いつでもコミュニケーションでき、最近の会話内容に基づいて応答できるチャットボットを作成することで、このチャットボットを多くの新しい応用分野で活用することができます」と、電気工学・コンピューターサイエンス（EECS）の大学院生で、論文の主執筆者であるGuangxuan Xiaoは述べている。
シャオの共著者には、彼のメンターであるソン・ハン氏（EECS准教授、MIT-IBMワトソンAIラボのメンバー、NVIDIAの特別科学者）、ユアンドン・ティアン氏（Meta AI研究員）、ベイディ・チェン氏（カーネギーメロン大学助教授）、マイク・ルイス氏（Meta AI研究員、論文のシニア執筆者）らがいる。本研究は、International Conference on Learning Representationsで発表される。
不可解な現象
大規模な言語モデルは、ユーザーのクエリに含まれる語彙などのデータを「トークン」と呼ばれる表現形式にエンコードする。多くのモデルは、これらのトークンを使用して新しいテキストを作成する、いわゆる「注意メカニズム」を採用している。
一般的に、AIチャットボットは見たテキストに基づいて新しいコンテンツを書く。そのため、KVキャッシュと呼ばれるメモリバンクに最近のトークンを保存し、後で使用する。アテンション・メカニズムは、キャッシュの全トークンを含むグリッド・マップ、すなわち「アテンション・マップ」を作成し、そこに各単語と他の単語との相関の強さをマークする。
このような関係を理解することは、このような大規模な言語モデルが人間のような文章を作成できるようになるのに役立つ。
しかし、キャッシュ量が膨大になると、アテンション・マップも膨大になり、計算速度の低下を招く。
さらに、コンテンツをエンコードするのに必要なトークンがキャッシュの上限を超えると、モデルのパフォーマンスが低下する。例えば、広く使われているモデルは4,096個のトークンを保存することができるが、学術論文には約10,000個のトークンが含まれることがある。
これらの問題を回避するために、研究者たちは「スライディング・キャッシュ」戦略を採用し、新しいトークンを追加するために最も古いトークンを削除した。しかし、最初のトークンが削除されると、モデルの性能は急落する傾向があり、新しく生成される語彙の品質に直接影響する。
この新しい論文で研究者たちは、スライディング・キャッシュの最初のトークンを一定に保つと、キャッシュ・サイズを超えた場合でも、このモデルの性能が維持されることを発見した。
しかし、この発見は直感に反するように思えるかもしれない。小説の最初の単語が最後の単語に直結することはほとんどないのに、なぜ最初の単語がモデルの新語生成能力にとって重要なのだろうか？
研究者たちは新しい論文の中で、この現象の背後にある理由も明らかにしている。
注目
ソフトマックス演算では、すべてのアテンションスコアの合計が1にならなければならない。モデルは残りの注目スコアを最初の語彙要素に割り当てる。
研究者たちは、この最初の語彙要素を「注意の焦点」と呼んでいる。
「注目点が必要であり、モデルは最初の語彙要素をこの注目点として選ぶ。モデルのダイナミクスを維持するためには、この注目点を常にキャッシュに留めておく必要があることがわかりました」。ハンはこう言った。
StreamingLLMの開発において、研究者たちは、4つの注意を集中させるポイント語彙要素をキャッシュの一番最初に配置することで、最適なパフォーマンスが得られることを発見した。
また、新しい語彙要素が追加され、古い語彙要素が削除されても、各語彙要素の位置エンコーディングは同じままでなければならないことにも注意している。たとえば、5番目の語彙要素が削除された場合、6番目の語彙要素がキャッシュの5番目になったとしても、6番目としてのエンコーディングを維持しなければならない。
これら2つのアイデアを組み合わせることで、StreamingLLMは首尾一貫した対話を維持し、再計算技術を使用する一般的なアプローチを凌駕することができる。
例えば、キャッシュ容量が256ワードの場合、再計算技法を用いた方法では新しい単語のデコードに63ミリ秒かかるのに対し、StreamingLLMでは31ミリ秒しかかからない。しかし、キャッシュサイズを4096ワードに増やすと、再計算手法では新しい語彙のデコードに1411ミリ秒かかるのに対し、StreamingLLMでは65ミリ秒しかかからない。
"StreamingLLMは、革新的なアテンションプーリング技術を採用しており、最大400万トークンのテキストを処理する際のパフォーマンスとメモリの安定性の問題に対処している。"シンガポール国立大学コンピューターサイエンス学長若手教授のヨン・ヤン氏は、自分はこの研究に関与していないものの、次のように述べている。「StreamingLLMの性能と多用途性は、アプリケーションを生成するためにAIを使用する方法を推進する革命的な技術になることを予兆しています。"
この見解に同意したのは、カーネギーメロン大学の機械学習とコンピューターサイエンスの助教授、ティアンチー・チェンである。StreamingLLMを使えば、大規模な言語モデルの対話の長さをスムーズに拡張することができます。私たちはこれを用いて、MistralのモデルをiPhoneに展開することに成功しました。"
研究チームはまた、すべてのトレーニングサンプルの前にいくつかのプレースホルダーを含めることで、モデルトレーニング中のアテンションプーリングの適用を検討した。
彼らは、アテンションプーリングと組み合わせて訓練されたモデルは、キャッシュ内のアテンションプールを1つだけ使用することで、パフォーマンスを維持できることを発見した。
しかし、StreamingLLMによってモデルが継続的な対話を行えるようになっても、キャッシュに保存されていない単語を記憶することはできない。将来的には、削除されたトークンを回収する方法や、モデルに過去の対話を記憶させる方法を模索することで、この制限を克服する予定である。
StreamingLLMは、NVIDIAのビッグ言語モデル最適化ライブラリに統合されました。TensorRT-LLMミディアム。
本研究の一部は、MIT-IBMワトソン人工知能研究所、MITサイエンスセンター、全米科学財団の支援を受けた。
