Cartridges: Lightweight and general-purpose long context representations via self-study

要約

大規模な言語モデルは、コーパス全体をコンテキストウィンドウに配置し、コンテキスト内学習（ICL）を活用することにより、大きなテキストコーパス（コードベース、法的文書、チャット履歴など）に基づいたクエリに答えるためによく使用されます。
現在のモデルは100K-1Mトークンのコンテキストをサポートしていますが、KVキャッシュのメモリ消費量が入力長のメモリ消費量が拡大するため、このセットアップはサービスを提供するのに費用がかかります。
別の方法を探ります。各コーパスでより小さなKVキャッシュをオフラインでトレーニングします。
推論時に、この訓練されたKVキャッシュをロードします。これはカートリッジと呼ばれ、応答をデコードします。
重大なことに、カートリッジをトレーニングするコストは、同じコーパスを参照するすべてのクエリで償却できます。
ただし、コーパスで次のトークン予測でカートリッジをトレーニングする素朴なアプローチは、ICLと競合していないことがわかります。
代わりに、コーパスに関する合成会話を生成し、コンテキスト指向の目的でカートリッジを訓練するトレーニングレシピである自己学習を提案します。
自習で訓練されたカートリッジは、ICLの機能を複製する一方で、サービスを大幅に安くしていることがわかります。
挑戦的なロングコンテキストベンチマークでは、38.6倍のメモリを使用し、26.4倍のスループットを有効にしながら、自習マッチのICLパフォーマンスで訓練されたカートリッジ。
自己学習は、モデルの有効なコンテキスト長（例：MTOBの128Kから484Kトークンなど）を拡張し、驚くべきことに、再訓練なしで推論時に構成できるカートリッジにつながります。

要約(オリジナル)

Large language models are often used to answer queries grounded in large text corpora (e.g. codebases, legal documents, or chat histories) by placing the entire corpus in the context window and leveraging in-context learning (ICL). Although current models support contexts of 100K-1M tokens, this setup is costly to serve because the memory consumption of the KV cache scales with input length. We explore an alternative: training a smaller KV cache offline on each corpus. At inference time, we load this trained KV cache, which we call a Cartridge, and decode a response. Critically, the cost of training a Cartridge can be amortized across all the queries referencing the same corpus. However, we find that the naive approach of training the Cartridge with next-token prediction on the corpus is not competitive with ICL. Instead, we propose self-study, a training recipe in which we generate synthetic conversations about the corpus and train the Cartridge with a context-distillation objective. We find that Cartridges trained with self-study replicate the functionality of ICL, while being significantly cheaper to serve. On challenging long-context benchmarks, Cartridges trained with self-study match ICL performance while using 38.6x less memory and enabling 26.4x higher throughput. Self-study also extends the model’s effective context length (e.g. from 128k to 484k tokens on MTOB) and surprisingly, leads to Cartridges that can be composed at inference time without retraining.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google