MOSAIC: Multi-Objective Search for Allocating Inference Cache Across LLM Decoder Layers
MOSAICは、長文脈LLM推論におけるレイヤーごとのKVキャッシュ予算を最適化するマルチオブジェクティブなニューラルアーキテクチャ探索フレームワークであり、様々なモデルやエビクション手法において、メモリ使用量を大幅に削減しながらより高いタスクスコアを達成することで、一律の割り当てを常に上回る性能を示します。
原論文は CC0 1.0 (http://creativecommons.org/publicdomain/zero/1.0/) のもとパブリックドメインに提供されています。 これは以下の論文のAI生成解説です。著者が執筆したものではありません。技術的な正確性については原論文を参照してください。 免責事項の全文を読む
大規模言語モデルは、膨大な文書を読み取ったり、長い会話を行ったりすることができる、今日の最も高度な人工知能ツールの背後にあるエンジンです。しかし、これらのモデルは、長いテキストのシーケンスを処理しようとする際に、物理的な限界に直面します。これまでに読んだ内容を把握し続けるために、コンピュータのメモリ内に増大する量のテンポラリ情報を保存しなければならないからです。このストレージ領域は「キー・バリュー・キャッシュ」として知られ、テキストの長さに比例して線形に拡大します。本一冊分の長さの文書を理解しようとしているモデルにとって、このメモリ要件は標準的なコンピュータ・ハードウェアの容量を超えるほど膨大になることがあり、その結果、システムが停止するか、情報を破棄することを余儀なくされます。
これを解決するために、研究者たちは、どの情報を保持するかについて選択的に判断する手法を開発してきました。これらの技術はフィルターのように機能し、どの部分のテキストがメモリに保持するのに十分重要であり、どの部分を安全に忘れてもよいかを決定します。これまで、これらのフィルターは、モデルの内部処理のあらゆる部分に対して、同じ厳格なルールを適用してきました。すべての層（レイヤー）がその役割を果たすために同じ量のメモリを必要とするという仮定に基づいていたのです。この新しい研究は、その仮定に異議を唱え、異なる層は実際には異なる量のスペースを必要としており、それらに等しいシェアを与えることはリソースの無駄であると提案しています。
研究チームは、「MOSIC」と呼ばれる新しい手法を導入しました。これは、大規模言語モデルの各層に割り当てる最適なメモリ量を自動的に算出するものです。モデルを、すべてのセクションがメモリというパイのケーキを等しく受け取る均一なブロックとして扱うのではなく、MOSICは各層を独自のニーズを持つ個体として扱います。このシステムは、高度な探索プロセスを用いて、メモリを分配する何千通りもの方法をテストします。それは、総メモリ制限の範囲内に収まりつつ、ある層はより多くのスペースを得て、別の層はより少ないスペースを得るという、特定のパターンを探し出します。目標は、メモリ不足に陥ることなく、モデルが最も正確に質問に答えたりテキストを要約したりできるような配置を見つけ出すことです。
研究者たちは、このアプローチを2つの人気のある大規模言語モデルでテストし、すべての層に同じ量のメモリを与えるという標準的な手法と比較しました。彼らは、複雑な読解、コード生成、長い文書の要約を含む、さまざまな困難なタスクでモデルを評価しました。その結果、この新しい手法は、圧倒的多数のケースにおいて、一律の割り当てを上回る一貫した性能を示しました。154のテスト設定のうち146の設定で、従来の均一な割り当てに勝利しました。いくつかの事例では、その改善は劇的であり、モデルは長いコンテキスト理解に関する標準化されたテストにおいて、スコアを50ポイント近く向上させました。おそらくより重要なことに、この新しい手法により、モデルは以前よりも最大4倍少ないメモリを使用して、同じ高いパフォーマンスを達成することができました。これは、モデルがより長いテキストを扱えるようになるか、あるいは、より小さく安価なコンピュータで動作できることを意味します。
最も驚くべき発見の一つは、最適なメモリ分布がランダムではなかったことです。探索の結果、一貫したパターンが明らかになりました。モデルの中間に位置する少数の層が、最初や最後にある層よりも一貫して多くのメモリを必要としていたのです。最初と最後の層は、より少ないスペースでも十分に機能することができ、一方で中間の層は、情報を保存する必要があるクリティカルなボトルネックとして機能していました。メモリを必要としない層から、必要とする層へとシフトさせることで、システムは非常に効率的になりました。この発見は、これらのモデルの内部構造が均一ではないことを示唆しています。つまり、ネットワークの異なる部分が異なる役割を果たしており、その機能を最大限に発揮するために異なるリソースを必要としているのです。
研究者たちはまた、この最適なパターンは一度発見すれば、多くの異なるメモリ制限に適用できることも発見しました。ユーザーが望むあらゆるメモリサイズに対して、毎回新しく高価な探索を実行する必要はありませんでした。彼らはメモリ分布の「勝利する形状（winning shape）」を見つけ出し、それを単に拡大または縮小させたのです。これは、たった一度の探索が、将来のあらゆるメモリ制約に対する解決策を提供できることを意味し、この手法を実用的なものにしています。このシステムは、モデルが質問に答えるために使用される前に、オフラインで理想的なバランスを学習するため、実際の使用時には追加の遅延を生じさせません。
最近の他の解決策と比較した際、この新しいアプローチは、はるかに広大な可能性の空間を探索しました。他の手法が一度にわずかなバリエーションしかテストしていなかったのに対し、このシステムは、全構成の極めて小さな断片を評価しながら、以前の手法よりも数十億倍も大きな探索空間をスキャンしました。それでも、最良のソリューションを見つけるために費やしたコンピュータ時間は、より少なくなりました。この研究は、すべての層を同じように扱うという古い考え方が非効率であることを裏付けています。一部のパーツが他のパーツよりも重要であることを認識し、それに応じてリソースを割り当てることで、これらの強力なツールをより速く、より安く、そして以前考えられていたよりもはるかに長いコンテキストを理解できるようにすることができるのです。
自分の分野の論文に埋もれていませんか？
研究キーワードに一致する最新の論文のダイジェストを毎日受け取りましょう——技術要約付き、あなたの言語で。