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

TL;DR

This work tackles the high memory cost of using large corpora in in-context learning by proposing Cartridges: small, offline-trained KV caches that encode a corpus for a given LM. Cartridges are trained via Self-Study, which generates synthetic conversations about the corpus and uses a context-distillation objective to mimic in-context behavior with dramatically less memory. On challenging long-context benchmarks, Cartridges match ICL performance while delivering up to $38.6\times$ memory reductions and $26.4\times$ higher throughput, and they extend effective context lengths (e.g., from $128k$ to $484k$ tokens) and support composability without retraining. The approach integrates smoothly with existing serving infrastructure and opens avenues for context-aware tools that can reason over very large repositories or documents.

Abstract

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.

Figures (12)

• Figure 1: Producing Cartridges via self-study. For a given document corpus, we train a Cartridge by distilling the corpus into a parameterized KV cache through a process we call Self-Study. At inference time, this Cartridge can be loaded into an LLM, which can then be used to answer diverse queries about the corpus, simulating in-context analysis of the corpus while requiring substantially less memory.
• Figure 2: Comparing KV caching strategies.Cartridge improves memory efficiency, while retaining the quality of in-context learning across a broad set of prompts. ✓ indicates a strength and ✗ indicates a limitation.
• Figure 3: Cartridges trained with Self-Study balance the generality and memory consumption tradeoff. We compare four methods on the GenConvo dataset: Cartridges trained with next-token prediction over $\mathcal{C}$, Cartridges trained with Self-Study, full ICL, and truncated ICL, a prompt compression method in which we truncate the $\mathcal{C}$ to the first $k$ tokens. (Left) We evaluate on different slices from the GenConvo dataset. Cartridges trained with next-token prediction performs well on memorization queries, which resemble it's training distribution, but cannot generalize to other queries like the other methods. (Center) The $x$-axis measures the size of the KV cache in GB for the different methods. The $y$-axis shows log-perplexity on the GenConvo dataset averaged over the query types. (Right) Peak throughput (tokens/s) measured for different cache sizes for Llama-3B and Llama-8B with SGLang zheng2024sglang on an 1xH100 (See \ref{['app:results']}).
• Figure 4: Cartridgesmatches ICL quality with lower memory costs. We measure Llama-3B response quality ($y$-axis) against KV cache memory ($x$-axis) for different methods, at different KV cache sizes. The dashed line marks the quality of standard ICL.
• Figure 5: Scaling Self-Study compute. These plots show how quality improves as we scale the training compute with Self-Study. In all plots, the $x$-axis shows the total number of global training steps with batch size 64 and maximum sequence length 1024. No synthetically generated data is reused (i.e. training proceeds for one epoch). Curves are provided for Cartridges of varying sizes ($p \in \{128, 512, 2048, 8192\}$). (Left) The $y$-axis shows accuracy on LongHealthadams2024longhealth with Llama-8B. (Middle) The $y$-axis shows the chrF on MTOBtanzer2023benchmark with Llama-3B. (Right) The $y$-axis shows log-perplexity (lower is better) on QASPERdasigi2021dataset with Llama-3B.

Theorems & Definitions (22)

• Definition 1
• Definition 2
• Definition 3
• Definition 4
• Definition 5
• Definition 6
• Lemma 1
• Lemma 2
• Definition 7
• proof