Patronen: Leichtgewichtige und vielseitige Langkontext-Darstellungen durch Selbststudium

Zusammenfassung

Große Sprachmodelle werden häufig verwendet, um Anfragen zu beantworten, die auf umfangreichen Textkorpora basieren (z. B. Codebasen, juristische Dokumente oder Chatverläufe), indem das gesamte Korpus in das Kontextfenster eingebettet und In-Context-Learning (ICL) genutzt wird. Obwohl aktuelle Modelle Kontexte von 100.000 bis 1 Million Tokens unterstützen, ist diese Konfiguration kostspielig, da der Speicherverbrauch des KV-Caches mit der Eingabelänge skaliert. Wir untersuchen eine Alternative: das Offline-Training eines kleineren KV-Caches für jedes Korpus. Zum Inferenzzeitpunkt laden wir diesen trainierten KV-Cache, den wir als Cartridge bezeichnen, und dekodieren eine Antwort. Entscheidend ist, dass die Kosten für das Training einer Cartridge auf alle Anfragen, die sich auf dasselbe Korpus beziehen, verteilt werden können. Allerdings stellen wir fest, dass der naive Ansatz, die Cartridge mit Next-Token-Prediction auf dem Korpus zu trainieren, nicht mit ICL konkurrieren kann. Stattdessen schlagen wir Self-Study vor, ein Trainingsverfahren, bei dem wir synthetische Konversationen über das Korpus generieren und die Cartridge mit einem Context-Distillation-Ziel trainieren. Wir stellen fest, dass Cartridges, die mit Self-Study trainiert wurden, die Funktionalität von ICL replizieren, während sie deutlich kostengünstiger zu bedienen sind. Bei anspruchsvollen Langkontext-Benchmarks erreichen mit Self-Study trainierte Cartridges die Leistung von ICL, während sie 38,6-mal weniger Speicher verbrauchen und einen 26,4-mal höheren Durchsatz ermöglichen. Self-Study erweitert auch die effektive Kontextlänge des Modells (z. B. von 128.000 auf 484.000 Tokens bei MTOB) und führt überraschenderweise zu Cartridges, die zum Inferenzzeitpunkt ohne erneutes Training kombiniert werden können.

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.

Patronen: Leichtgewichtige und vielseitige Langkontext-Darstellungen durch Selbststudium

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

Zusammenfassung

Support