カートリッジ：自己学習による軽量で汎用的な長文脈表現

要旨

大規模言語モデルは、大規模なテキストコーパス（コードベース、法律文書、チャット履歴など）に基づくクエリに答えるためによく使用されます。これは、コーパス全体をコンテキストウィンドウに配置し、インコンテキスト学習（ICL）を活用することで行われます。現在のモデルは100K-1Mトークンのコンテキストをサポートしていますが、この設定は、KVキャッシュのメモリ消費量が入力長に比例して増加するため、運用コストが高くなります。私たちは、各コーパスに対してオフラインで小さなKVキャッシュをトレーニングするという代替案を探ります。推論時には、このトレーニング済みのKVキャッシュ（Cartridgeと呼びます）をロードし、応答をデコードします。重要な点として、Cartridgeのトレーニングコストは、同じコーパスを参照するすべてのクエリにわたって償却することができます。しかし、コーパスに対して次のトークン予測を用いてCartridgeをトレーニングする素朴なアプローチは、ICLと比べて競争力がないことがわかりました。代わりに、コーパスに関する合成会話を生成し、コンテキスト蒸留の目的でCartridgeをトレーニングする「自己学習（self-study）」というトレーニングレシピを提案します。自己学習でトレーニングされたCartridgeは、ICLの機能を再現しつつ、運用コストが大幅に低くなることがわかりました。挑戦的な長文コンテキストのベンチマークでは、自己学習でトレーニングされたCartridgeは、ICLと同等の性能を発揮しつつ、38.6倍少ないメモリを使用し、26.4倍高いスループットを実現します。自己学習はまた、モデルの有効なコンテキスト長を拡張し（例えば、MTOBでは128kから484kトークンに）、驚くべきことに、再トレーニングなしで推論時にCartridgeを組み合わせることができるようになります。

English

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.

カートリッジ：自己学習による軽量で汎用的な長文脈表現

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

要旨

Support