PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Sabri Eyuboglu, Ryan Ehrlich, Simran Arora, Neel Guha, Dylan Zinsley, Emily Liu, Will Tennien, Atri Rudra, James Zou, Azalia Mirhoseini, Christopher Re

arXiv:2506.06266v3cs.CLcs.AIcs.LG

TL;DR

Serving large corpora through ICL is costly because KV-cache memory scales with context length, motivating a reusable smaller representation. The paper trains corpus-specific Cartridges offline with synthetic conversations and context distillation, and finds that they match ICL quality while substantially reducing serving costs.

  • Problem

    ICL over large corpora is costly to serve because KV-cache memory grows with input length, while efficient alternatives may lose generality and structural awareness.

  • Method

    Self-study generates synthetic conversations about each corpus and trains a smaller corpus-specific KV cache with a context-distillation objective.

  • Results

    38.6× less memory and 26.4× higher peak throughput are achieved while Cartridges match ICL response quality across diverse evaluations.

  • Takeaways & Limitations

    Cartridges provide a reusable, composable alternative to ICL for repeated queries over the same large corpora and can extend effective context length.

  • Takeaways & Limitations

    Self-study requires substantially more construction compute than standard ICL prefill, taking about 30 minutes on one 8×H100 node in the unoptimized LLAMA-8B implementation.

Abstract

from arXiv · show

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.

1 Introduction

Large language models use in-context learning over large corpora, but long contexts make serving memory-intensive and reduce throughput. The paper introduces Cartridges trained by self-study to preserve ICL-like functionality with substantially lower memory use.

  • Motivation: Long-context ICL is costly because KV-cache memory grows linearly with input length.LLaMA 70B requires 84 GB for one 128k-token question, while LLaMA 8B throughput drops 77× from 1k to 120k tokens.
  • Motivation: Existing prompt and KV-cache compression methods incur memory-quality tradeoffs, with performance degrading rapidly beyond 2× compression.
  • Cartridges: Cartridges are smaller KV caches trained offline for a specific corpus, then loaded into an LLM and reused across its queries.The approach freezes the LLM and trains key and value vectors, allowing preparation costs to be amortized across repeated references to the same corpus.
  • Self-study: Self-study generates synthetic corpus conversations and trains the Cartridge with context distillation to reproduce the model’s behavior with the corpus in context.Synthetic conversations support generality and can use chunking and seed prompts to improve structural awareness; context distillation outperforms next-token prediction.
  • Results: 38.6× less memory and 26.4× higher peak throughput are achieved while matching ICL quality across challenging long-context benchmarks.The evaluations cover corpora of 100k–484k tokens and diverse queries.
  • Results: Self-study extends effective context length, supports composable Cartridges, and KV-cache parameterization outperforms LoRA parameterization on in-domain and out-of-domain tasks.On MTOB, a Cartridge built from a 484k-token textbook outperforms ICL by 11.0 chrF points over the first 130,000 tokens.

2 Preliminaries

The paper frames long-context serving as a quality–memory problem: ICL supports diverse corpus-grounded queries but KV-cache costs grow with context length. Prior compression methods trade memory savings for degraded quality, motivating offline corpus representations.

  • The target setting is a stream of diverse queries about one shared corpus, optimizing response quality while minimizing serving memory.
  • ICL supports diverse queries over large corpora, including factual recall, reasoning, summarization, and creative generation.

3 The CARTRIDGE paradigm

The CARTRIDGE paradigm trains a small corpus-specific KV cache offline and loads it at inference time as a compact substitute for the corpus prefix. Its design targets ICL-like behavior with lower memory use while retaining generality, long-range dependencies, and potential composition.

  • 3. The CARTRIDGE paradigm: The paradigm replaces on-the-fly corpus prefill with offline training of corpus representations, which can be reused across repeated queries.
  • 3.1 Formalizing CARTRIDGES: A CARTRIDGE is a small trainable parameter set that makes an LLM behave as if the corpus were in its context window.
  • 3.1 Formalizing CARTRIDGES: CARTRIDGES should generalize across unknown query types, capture long-range dependencies, and support composition across corpora.
  • 3.2 Parameterizing CARTRIDGES: The trainable representation uses key and value vectors with p virtual-token positions, so its memory matches a KV cache for a p-token prompt.
  • 3.2 Parameterizing CARTRIDGES: Training substitutes the corpus KV pairs with Z and backpropagates loss into Z while freezing the language model parameters.
  • 3.3 Serving CARTRIDGES: CARTRIDGES can be loaded directly into existing KV-cache slots, allowing serving through inference servers that manage cached prefixes.

4 SELF-STUDY: A self-supervised method for training CARTRIDGES

SELF-STUDY addresses the poor generalization of raw-corpus next-token training by generating synthetic corpus conversations and distilling the context-aware model into a trainable cache. The resulting CARTRIDGES match ICL across diverse queries, while context distillation improves accuracy and performance scales with compute.

  • Motivating observations: 107×: Naive next-token-trained CARTRIDGES use less memory than ICL but memorize the corpus and generalize poorly to diverse queries.
  • 4. SELF-STUDY: SELF-STUDY CARTRIDGES match ICL query quality across many query types, and longer training consistently improves task performance.
  • 4.1 Synthetic data generation: SELF-STUDY generates synthetic query-response conversations by prompting the model to quiz itself about corpus content.
  • 4.1 Synthetic data generation: Chunking lets SELF-STUDY focus generation on different corpus regions and train on corpora longer than the model’s context window.
  • 4.1 Synthetic data generation: The synthetic pipeline uses generic seed prompts spanning structuring, summarization, questions, use cases, and creative tasks.
  • 4.2 Context-distillation objective: Context distillation minimizes divergence between teacher next-token distributions conditioned on the subcorpus and student distributions produced by the trainable cache.
  • 4.2 Context-distillation objective: 3.7 accuracy points: Context distillation improves LONGHEALTH accuracy when synthetic-data quantity is held constant.

5 Results

SELF-STUDY-trained CARTRIDGES match or outperform ICL across long-context settings while reducing memory and increasing throughput. Ablations show that context distillation, prefix-tuning, corpus-token initialization, and diverse seed prompts are important design choices, while independently trained CARTRIDGES can be composed without joint training.

  • Up to 10× lower memory on LONGHEALTH and up to 100× lower memory on QASPER are achieved at comparable or better performance than ICL.Compression baselines degrade at factors as low as 2×, whereas CARTRIDGES retain quality at much larger savings.
  • 26× higher throughput is available at CARTRIDGE cache sizes that match ICL performance.The smaller cache footprint enables substantially higher peak tokens-per-second throughput.
  • 11.0 chrF points separate CARTRIDGES from competitive baselines on the 484k-token MTOB textbook setting.CARTRIDGES match ICL on the hand-curated 60k-token version while using only the longer corpus, which exceeds LLAMA-8B’s 128k-token context window.
  • CARTRIDGE Parameterization: Prefix-tuning outperforms memory-matched LoRA on corpus queries and preserves unrelated MMLU accuracy as CARTRIDGE size increases.On roughly 0.6 GB MTOB CARTRIDGES, prefix-tuning exceeds LoRA by 4.5 chrF points; MMLU falls from 54.7 to 54.3 with prefix-tuning versus 54.7 to 45.3 with LoRA.
  • CARTRIDGE Initialization: 55.3% LONGHEALTH accuracy is obtained by initializing with the first p corpus tokens, versus 51.3% with random-token vectors and 29.9% with random vectors.Initialization from actual corpus tokens is reported as critical for achieving ICL-level performance.
  • SELF-STUDY Objective: An 8.6-point MTOB improvement results from context distillation over next-token prediction on the same synthetic conversations.The reported scores increase from 24.9 to 33.5, with improvements also observed on LONGHEALTH and QASPER.
  • CARTRIDGE Composition: Composing independently trained CARTRIDGES substantially outperforms a single-CARTRIDGE baseline and truncated ICL on multi-document questions.The composed caches are concatenated without additional training and produce coherent generations off the shelf.

6 Discussion and conclusion

CARTRIDGES provide an alternative to ICL for repeated queries over large corpora, matching response quality with lower memory consumption and higher throughput. The tradeoff is that SELF-STUDY construction is substantially more expensive than standard ICL prefill, while future applications include context-aware assistants and IDEs.

  • 38.6× lower memory consumption and 26.4× higher peak throughput are reported while matching ICL response quality.CARTRIDGES are also described as simple to train, composable, and compatible with existing serving infrastructure.
  • SELF-STUDY training takes approximately 30 minutes on a single 8×H100 node for LLAMA-8B, making it more costly than ICL prefill.The paper frames CARTRIDGES as a compute-for-memory tradeoff rather than a drop-in ICL replacement.
  • Medical assistants using full patient histories and IDEs understanding entire codebases are envisioned as applications of CARTRIDGES.These applications are presented as future directions enabled by context-aware AI systems.

A Extended Results

This section introduces ablations of the main CARTRIDGE and SELF-STUDY design choices.

  • The paper ablates the main design choices of CARTRIDGES and SELF-STUDY.

A.1 CARTRIDGE design choices: parameterization and initialization

The study compares prefix-tuning and LoRA parameterizations for CARTRIDGES and evaluates initialization from random vectors, random-token KV caches, or corpus prefixes. Prefix-tuning is more robust across query types, while corpus-token initialization supports stronger results than random initialization.

  • Parameterization: Prefix-tuning outperforms memory-matched LoRA on both in-domain and out-of-domain queries.Prefix-tuning occupies the high-accuracy region for both MMLU and the target dataset.
  • Parameterization: At 1.06 GB, LoRA MMLU accuracy drops from 60.0% to 45.3%, whereas prefix-tuning drops to 54.3%.The LoRA degradation is correlated with CARTRIDGE size.
  • Parameterization: At 0.96 GB on LONGHEALTH, prefix-tuning reaches 55.6% accuracy versus 41.3% for LoRA.The best LoRA accuracy is 47.25% at 0.26 GB.
  • Initialization: Allowing the first token’s key and value vectors to train causes instability because they act as an attention sink.Figure 10 shows MMLU accuracy dipping below 30% before recovery in some runs.
  • Initialization: Initializing with actual-token KV vectors is critical for ICL-level performance, while first-document-token initialization is slightly stronger than unrelated-text initialization.Random vectors achieve 29.9% accuracy on LONGHEALTH.

A.2 SELF-STUDY design choices: data-generation and objective

SELF-STUDY generates synthetic corpus conversations using seeded prompts and trains CARTRIDGES with context distillation. Diverse seed prompts improve some reasoning-heavy tasks, while context distillation consistently outperforms next-token prediction at fixed training compute.

  • Data Generation: Five seed-prompt types improve CARTRIDGE quality over a single broad prompt on MTOB and LONGHEALTH.The prompts target structuring, summarization, questions, use cases, and creative responses.
  • Data Generation: A 7.9-point ChRF improvement occurs on MTOB with 1024-token CARTRIDGES, from 24.1 to 32.0.LONGHEALTH improves by 5.5 accuracy points, from 45.8 to 51.3.
  • Data Generation: Five prompt types provide no benefit on QASPER, possibly because its queries are mostly factual and require less complex reasoning.The proposed explanation is explicitly framed as possible.
  • Training Objective: Context distillation aligns CARTRIDGE-augmented output distributions with those of the model using the corpus in context.The objective requires collecting top output probabilities from the in-context model during data generation.
  • Training Objective: On MTOB with 2048-token CARTRIDGES, context distillation improves ChRF by 8.3 points over next-token prediction, from 24.9 to 33.2.LONGHEALTH shows a 3.7-point accuracy gap, from 47.6 to 51.3.
  • Training Objective: For fixed SELF-STUDY compute, context distillation is considerably more effective than next-token prediction.Performance also improves with more SELF-STUDY compute, so additional compute might narrow the gap.

B.2.2 Synthetic Data Generation

Synthetic data generation uses language models to create training examples, while KV-cache reduction methods compress, evict, merge, or project cache states. Existing compression methods often degrade beyond moderate compression and may not match full-cache performance on long-context tasks.

  • Synthetic Data Generation: Synthetic-data methods generate training examples with language models, including domain-specific pipelines and expert-model orchestration.Examples include Bonito and MetaSynth.
  • Prompt Compression: Prompt compression shortens long inputs through filtering, summarization, or learned soft-token representations.Hard-token methods output vocabulary tokens, whereas soft-token methods output learned embeddings.
  • KV-Cache Compression: KV-cache compression can evict, merge, or low-rank-project cached states instead of retaining the full cache.These approaches operate directly on key-value representations or exploit their structure.
  • KV-Cache Compression: Eviction methods cannot recover a key once it has been permanently discarded.Selective loading can improve bandwidth use but does not reduce KV-cache memory consumption.
  • KV-Cache Compression: Existing merging methods can perform similarly to full caches at 50% or greater reduction but still fail to match full-cache performance on several tasks.The cited methods also do not evaluate effectiveness in long-context settings.
  • Architectural Changes: Architectural KV-cache reductions are generally not immediately compatible with pretrained standard-MHA models.Such approaches typically require retraining from scratch or complex architecture conversion.

B.3.5 Synthetic data generation

Synthetic training data is widely used to support fine-tuning and related adaptation methods. Prior examples include specialized synthetic-data generators and language-model training processes that incorporate substantial synthetic data.

  • Synthetic Data Generation: Synthetic training data is used broadly for fine-tuning and language-model training.The passage cites Bonito, MetaSynth, and Phi-4 as examples.

C Extended method description

SELF-STUDY generates diverse synthetic conversations about corpus chunks to train CARTRIDGES, with chunking that also supports corpora longer than the model’s context window.

  • Seed prompt generation: SELF-STUDY seeds synthetic conversation generation with prompts targeting different document aspects.The method includes creative, structuring, summarization, question, and use-case seed prompt generators.
  • Chunking: SELF-STUDY samples uniformly random token-level chunks from the corpus and generally prepends textual descriptions for contextualization.Chunk descriptions help the model focus on different corpus regions and generate diverse synthetic examples.
  • Chunking: Chunking allows CARTRIDGES to be trained on corpora longer than the model’s context window.Short subcorpora let the LLM focus on different parts of the corpus during synthetic-data generation.
  • Dataset-specific chunking: 512–4096 tokens is the chunk-size range used for LONGHEALTH, MTOB, and QASPER.AMD/FinanceBench instead uses fixed 8192-token chunks without a prepended description.

D Datasets

The evaluation covers diverse long-document tasks, including finance, clinical analysis, and low-resource translation, with synthetic questions designed to test more than fact retrieval.

  • GENCONVO: GENCONVO uses AMD’s 2022 10-K filing to test comprehension, reasoning, and varied information extraction over a long document.Its generation process uses Claude Sonnet 3.7 with the full document.
  • GENCONVO: GENCONVO generates 16 unique questions for each prompt template and then produces document-grounded answers.Templates target factual recall, synthesis, multi-hop reasoning, structure, creativity, counting, mathematical reasoning, and knowledge combining document and external facts.
  • Question diversity: The benchmark includes question types for factual details, synthesis, document structure, creative generalization, counting, mathematical reasoning, and outside knowledge.These templates are intended to move beyond simple fact retrieval.
  • LONGHEALTH: LONGHEALTH contains 20 fictional clinical case reports and 400 multiple-choice questions about them.Each report contains between 5,090 and 6,754 words, and the questions cover information extraction, negation, and sorting.
  • LONGHEALTH: LONGHEALTH experiments use ten patient reports in a roughly 100k-token panel that fits within LLaMA 3’s context length.The questions are categorized into information extraction, negation, and sorting.
  • MTOB: MTOB evaluates English–Kalamang and Kalamang–English translation using a grammar book and accompanying linguistic resources.The paper focuses on Kalamang-to-English translation, where Kalamang has virtually no web presence.

E Theoretical analysis: Relationship between attention, linear attention, and CARTRIDGES

The theoretical analysis compares transformer attention, linear attention, and gradient-descent-updated states through MQAR, clarifying their memory requirements and ability to handle repeated keys.

  • State representations: CARTRIDGES use constant memory independently of input length, but update their state with SELF-STUDY gradient descent rather than a linear-attention rule.This distinguishes their update mechanism from the constant-sized KV-state matrix used in linear attention.
  • Implication: The theoretical results provide intuition for why constant-sized CARTRIDGES can match full KV-cache performance where linear-attention architectures have struggled.This connects the MQAR analysis to the paper’s long-context motivation.
  • MQAR framework: In MQAR, models are characterized by a state, an update rule, and a query rule.The analysis uses this common framework to compare the three architectures.
  • Transformer: Ω(Nd) parameters are required by the transformer state on every MQAR input, including 1-repetitive MQAR.The state grows with the stream length because key-value pairs are appended over time.
  • Linear attention: O(d^2) parameters suffice for linear attention to solve repetitive MQAR with orthonormal keys, up to scaling.The result follows from storing key-value information in a d × d matrix.
  • Gradient descent: O(d^2) parameters allow gradient descent to exactly solve m-repetitive MQAR with orthonormal keys.The update rule includes a peel term that removes the current value for a key before inserting its replacement.
  • ε-JL keys: Linear attention cannot solve even 2-repetitive MQAR under ε-JL keys, because accumulated interference from repeated keys grows with repetition.The analysis states that this makes linear attention unfit for longer contexts or contexts with many repeats.
  • ε-JL keys: Gradient descent exactly solves m-repetitive MQAR with ε-JL keys under the stated bounds on ε and α.The theorem gives the conditions ε ≤ 1/[m^2(m−1)] and α < m−1.
Loading 2506.06266v3…