Cartridges: Lichtgewicht en veelzijdige lange-context representaties via zelfstudie

Samenvatting

Grote taalmodellen worden vaak gebruikt om vragen te beantwoorden die gebaseerd zijn op grote tekstcorpora (bijvoorbeeld codebases, juridische documenten of chatgeschiedenissen) door het volledige corpus in het contextvenster te plaatsen en gebruik te maken van in-context learning (ICL). Hoewel huidige modellen contexten van 100K-1M tokens ondersteunen, is deze opzet kostbaar om te bedienen omdat het geheugengebruik van de KV-cache schaalt met de invoerlengte. Wij onderzoeken een alternatief: het offline trainen van een kleinere KV-cache voor elk corpus. Tijdens inferentie laden we deze getrainde KV-cache, die we een Cartridge noemen, en decoderen we een antwoord. Cruciaal is dat de kosten van het trainen van een Cartridge kunnen worden verdeeld over alle queries die verwijzen naar hetzelfde corpus. We ontdekken echter dat de naïeve aanpak van het trainen van de Cartridge met next-token prediction op het corpus niet concurrerend is met ICL. In plaats daarvan stellen we self-study voor, een trainingsmethode waarbij we synthetische gesprekken over het corpus genereren en de Cartridge trainen met een context-distillatie-objectief. We ontdekken dat Cartridges die met self-study zijn getraind, de functionaliteit van ICL repliceren, terwijl ze aanzienlijk goedkoper zijn om te bedienen. Op uitdagende long-context benchmarks presteren Cartridges die met self-study zijn getraind even goed als ICL, terwijl ze 38,6x minder geheugen gebruiken en een 26,4x hogere doorvoer mogelijk maken. Self-study breidt ook de effectieve contextlengte van het model uit (bijvoorbeeld van 128k naar 484k tokens op MTOB) en leidt, verrassend genoeg, tot Cartridges die tijdens inferentie kunnen worden samengesteld zonder hertraining.

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: Lichtgewicht en veelzijdige lange-context representaties via zelfstudie

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

Samenvatting

Support