PaperlayerView exact revision on arXiv

Source-linked AI summary

Latent Compaction with Test-Time Training

Yingfa Chen, Xingyu Shen, Chaojun Xiao, Xu Han

arXiv:2610.20582v1cs.CLcs.LG

TL;DR

Long contexts make KV caches costly, while compression must preserve information for unseen questions. Latent Compaction trains document-specific caches with frozen model weights, and experiments report gains across RULER, LongBench, and Phonebook Lookup; per-document fitting remains costly.

  • Problem

    Long contexts enlarge KV caches, increasing memory and attention costs, while compression must preserve information for unseen questions.

    Source: §1 · p. 1

  • Method

    Latent Compaction learns a document-specific KV cache with frozen model weights through attention fitting followed by end-to-end distillation on generic document prompts.

    Source: §1 · p. 1

  • Results

    Experiments report retrieval gains over budget-matched compression baselines across RULER, LongBench, and Phonebook Lookup, with transfer to multiple Qwen3 sizes and MiniCPM5 models.

    Source: §1 · pp. 1–2

  • Takeaways & Limitations

    The learned document KV cache is reusable, while preparation costs can be amortized across repeated use.

    Source: §1 · p. 1

    Per-document fitting remains costly, and extreme compression hurts accuracy.

    Source: §5 · p. 9

Abstract

from arXiv · show

Long contexts enlarge key--value (KV) caches, increasing memory use and attention costs during generation. Compressing these caches must preserve information for unseen questions. We introduce **Latent Compaction**, which learns a compact document-specific cache through test-time training. We initialize cache entries using reconstruction importance, teacher attention, and key-space diversity. Training first fits attention outputs and log-normalizers to those of the full cache, with reconstruction-selected entries held fixed. It then jointly optimizes the entire cache through end-to-end distillation from the full-cache model. Model weights remain frozen throughout, and training uses only generic prompts about the document, without access to downstream questions or answers. Experiments on RULER, LongBench, and Phonebook Lookup demonstrate substantial gains over budget-matched compression baselines. On a RULER, for example, our method achieves 96.7% accuracy at 32$\times$ compression, compared with 40.0% for KVZip+. The method also transfers across model scales and families, with retrieval gains on both Qwen3 and MiniCPM5. These results show that treating the KV cache as a neural network and optimizing it with test-time training can preserve retrieval accuracy under aggressive compression. The code is released at https://www.github.com/THUNLP/latent-compaction

1 INTRODUCTION

Long-context KV caches raise inference costs, while compression must preserve information for questions that are not yet known. Latent Compaction trains a document-specific cache with frozen model weights and reports gains across retrieval and long-document benchmarks.

1 INTRODUCTION

  • Long contexts make KV caches a major inference bottleneck by limiting batch sizes and increasing memory-bandwidth demands at each decoding step..Source: §1 · p. 1
  • Text summaries may omit details whose importance emerges later, motivating direct compression of the KV representation before downstream questions arrive.Preparation costs can be amortized across repeated use.Source: §1 · p. 1
  • Prior eviction and cache-learning methods motivate the question of whether gradient descent on a compact KV cache can preserve document information under an aggressive memory budget.The introduction identifies selection-based eviction, Cartridges, and Attention Matching as related approaches.Source: §1 · p. 1
  • The method initializes a document-specific cache, fits attention outputs and log-normalizers while holding preservers fixed, then jointly optimizes all entries using generic document prompts.Model weights remain frozen, and downstream questions and answers are excluded.Source: §1 · p. 1
  • Latent Compaction improves retrieval over budget-matched compression baselines across context lengths and compression ratios, with gains also reported on LongBench and across model scales and families.The reported transfer includes multiple Qwen3 sizes and MiniCPM5 models.Source: §1 · pp. 1–2

2 RELATED WORK

Related approaches select original KV entries, learn compact representations, or adapt model parameters at test time. Latent Compaction instead optimizes document-specific KV slots in a frozen Transformer before downstream questions are known.

2 RELATED WORK

  • Eviction methods retain selected original KV pairs, with methods differing in their attention, observation-window, or layer/head budget criteria.Latent Compaction uses KVZip+ scores to initialize preservers and evaluates query-agnostic adaptations of SnapKV and PyramidKV.Source: §2 · p. 2
  • Learned context representations range from reusable compression mechanisms and corpus-specific cache self-study to attention matching; Latent Compaction follows stage-1 attention fitting with end-to-end refinement of all entries.RestoreKV instead generates restoration entries through a shared LoRA mechanism, whereas Latent Compaction directly optimizes document-specific slots.Source: §2 · pp. 2–3, Figure 2 · p. 3
  • Unlike prior test-time training approaches that update model weights or use online-learned hidden states, Latent Compaction updates document-specific KV slots in a frozen Transformer.It prepares the cache before downstream questions are known and reuses it across questions without further optimization.Source: §2 · p. 3

3 METHOD: LATENT COMPACTION

Latent Compaction trains a document-specific KV cache while keeping model weights frozen, then reuses that cache across downstream questions. It initializes compact entries from reconstruction scores, teacher attention, and key-space diversity, and trains them in two stages using generic document prompts.

3 METHOD: LATENT COMPACTION

  • Latent Compaction optimizes a document-specific cache with frozen model weights before questions arrive, then reuses the optimized cache across questions.The method overview assigns initialization, attention fitting, end-to-end refinement, and generic teacher rollouts to Sections 3.2–3.5.Source: §3 · p. 3

3.1 MEMORY ARCHITECTURE

  • The cache combines compacted document entries with later question and generated-token entries, while learned per-entry biases adjust document attention mass.Subsequent-token KV pairs are unbiased and outside the document-cache budget; the combined attention is equivalent to a single softmax.Source: §3.1 · pp. 3–4

3.2 INITIALIZATION

  • Initialization combines globally ranked reconstruction preservers, teacher-attention anchors, and farthest-first diversifiers for key-space coverage.Preservers stay frozen in stage 1, while anchors and diversifiers train; all entries train in stage 2.Source: §3.2 · pp. 4–5

3.3 TRAINING STAGE 1: ATTENTION FITTING

  • Stage 1 fits compact-cache attention outputs and log-normalizers to the full document cache using captured teacher queries, with preservers fixed.The remaining anchors and diversifiers receive updates; matching log-normalizers calibrates document attention relative to later tokens.Source: §3.3 · p. 5

3.4 TRAINING STAGE 2: END-TO-END REFINEMENT

  • Stage 2 jointly refines all cache entries by matching the teacher’s top-20 token probabilities and aggregate tail mass, alongside a greedy-token language-model loss.The two loss terms have equal weight, and individual tail tokens are not distinguished.Source: §3.4 · p. 5

3.5 TEACHER ROLLOUTS FOR DOCUMENT-SPECIFIC SUPERVISION

  • Both stages use fixed teacher rollouts from generic document prompts, excluding downstream evaluation questions and answers.The 480 prompts are independent; teacher continuations stop at EOS or 24 tokens, and teacher tokens rather than student predictions supply the prefixes.Source: §3.5 · pp. 5–6

4 EXPERIMENTS

Experiments compare Latent Compaction with cache-compression and text-summary baselines across retrieval and natural-document tasks. Results favor the method across many settings, with qualifications at extreme compression and ablations identifying important initialization, training, and inference choices.

4.1 EXPERIMENTAL SETUP

  • Experiments use Qwen3-1.7B by default and compare methods at matched logical KV-pair budgets, excluding question and generated-token caches.Physical memory also depends on precision, biases, and padding; the default model uses frozen weights, BF16 computation, and an FP32 KV cache.Source: §4.1 · pp. 6–7

4.2 EVALUATION BENCHMARKS AND METRICS

  • Evaluation covers synthetic long-context tasks in RULER, natural-document understanding in LongBench, and name-to-number retrieval in Phonebook Lookup.Most ablations use Phonebook Lookup with 20 questions per document and reuse each optimized cache across those questions.Source: §4.2 · p. 7

4.3 MAIN RESULTS

  • Latent Compaction generally exceeds compression baselines on RULER and Phonebook Lookup and beats text summaries across LongBench ratios, though its advantage is not uniform.Text summarization outperforms it on RULER at the most aggressive ratios, and all compressed methods approach zero Phonebook Lookup accuracy at 64× compression and beyond; LongBench scores decline as compression increases.Source: §4.3 · p. 8

4.4 TRANSFER TO OTHER MODEL SCALES AND MODEL FAMILIES

  • On Phonebook Lookup, retrieval advantages extend to additional Qwen3 model sizes and MiniCPM5 models, using each model’s own teacher outputs rather than transferring an optimized cache between models.This tests applicability across model scales and families, not cache transfer between models.Source: §4.4 · p. 8

4.5 ABLATIONS: INITIALIZATION

  • Preserver allocation affects accuracy more than anchor fraction, while a 25% anchor and 75% diversifier mixture beats every tested replacement initializer at each context length.The highest average among tested anchor fractions is 25%.Source: Table 1 · p. 7, §4.5 · p. 8

4.6 ABLATIONS: TRAINING STAGE 1

  • Stage 1 is important: omitting it, unfreezing preservers during it, or removing log-normalizer loss each substantially reduces lookup accuracy.The ablation results are reported for Phonebook Lookup with Qwen3-1.7B at 16× compression.Source: §4.6 · p. 8, Table 2 · p. 8

4.7 ABLATIONS: TRAINING STAGE 2

  • Stage 2 improves lookup accuracy, and adding language-modeling loss to grouped KL helps at all tested lengths; full-vocabulary targets provide no consistent gain.Stage 2 gains 6.75, 15.50, and 7.75 percentage points at 4K, 8K, and 16K, respectively. Accuracy peaks at 1,000 stage-2 updates and then declines slightly; doubling stage 1 from 1,000 to 2,000 updates raises accuracy only from 61.00% to 61.25%.Source: §4.7 · pp. 8–9, Figure 6 · p. 9, Table 3 · p. 9

4.8 ABLATIONS: MODEL ARCHITECTURE

  • Removing attention biases lowers accuracy, while replacing normalizer weighting with equal averaging reduces accuracy to at most 0.25%, despite retaining stage-1 normalizer loss.Normalizer weighting recovers joint attention over document and subsequent entries.Source: §4.8 · p. 9, Table 4 · p. 9

5 CONCLUSION

Latent Compaction learns reusable document-specific KV caches with frozen model weights, and experiments report accuracy gains across three benchmarks. Per-document fitting remains costly, and extreme compression hurts accuracy.

5 CONCLUSION

  • Latent Compaction trains reusable document KV caches with frozen model weights through attention fitting and distillation, with reported gains across RULER, LongBench, and Phonebook Lookup.The conclusion also reports ablations of initialization and attention normalization.Source: §5 · p. 9
  • Per-document fitting remains costly and extreme compression hurts accuracy; future work will reduce preparation cost and study broader tasks.Source: §5 · p. 9
  • The authors used generative AI for manuscript editing and code implementation, manually reviewed the implementation, and did not use AI to develop research ideas or generate datasets.Source: AI USE STATEMENT · p. 10

A.1 MODEL CONFIGURATION

The main experiments use a frozen Qwen3-1.7B checkpoint and tokenizer, with BF16 inference and the model’s original rotary configuration.

A.1 MODEL CONFIGURATION

  • All main experiments use Qwen3-1.7B with 28 layers, 2,048 hidden dimensions, 16 query heads, eight KV heads, and 128-dimensional heads.Source: §A.1 · p. 11
  • Model weights remain frozen; inference uses BF16, raw text, and the original rotary configuration, without context extension or architecture changes.Checkpoint and tokenizer hashes are verified across hosts.Source: §A.1 · p. 11
  • Table 5 lists the Hugging Face checkpoints used in the experiments and notes that each model uses its corresponding tokenizer.Source: Table 5 · p. 12

A.2 DOCUMENT-SPECIFIC TRAINING

Document-specific cache training reuses full-cache teacher rollouts across two stages while keeping model weights frozen. The appendix specifies rollout construction, optimization, attention computation, and cache allocation.

A.2 DOCUMENT-SPECIFIC TRAINING

  • Stage 1 trains anchor and diversifier keys, values, and biases while preservers remain frozen, fitting sampled query positions to full-cache attention targets.Query positions come from prompt–continuation rollouts; pools exceeding 4,096 positions are evenly subsampled, and targets are computed in chunks of 16.Source: §A.2 · p. 12
  • Stage 2 optimizes the entire compacted cache using four independently processed teacher rollouts per update, with top-20 probabilities, tail mass, and next-token cross-entropy in distillation.No downstream-question-based checkpoint selection is used.Source: §A.2 · p. 12
  • The compacted-cache attention path computes biased, masked document attention and log-normalizers in FP32, then combines document and causal continuation branches using a shared softmax normalizer.The implementation uses explicit PyTorch operations rather than fused attention; the prose says a custom fused kernel could implement the same mathematics.Source: §A.2 · p. 13
  • Stage 1 fits the compressed-cache document log-normalizer to a per-query-head standardized teacher target; normalizers are recomputed per query rather than stored in the cache artifact.The teacher standard deviation is floored at 10^-4; saved cache artifacts contain keys, values, biases, and masks.Source: §A.2 · p. 13
  • The reference allocation assigns one quarter of the logical KV-pair budget to anchors and diversifiers, with one quarter of those slots as anchors and the remainder initialized by farthest-point selection.Training uses the deployed number of slots; evaluation accounts for BF16 deployment precision.Source: §A.2 · p. 13
  • Each document uses up to 24 generated tokens for each of 480 generic prompts; teacher rollouts are reused in both stages, with 1,000 updates per stage by default.The document-cache parameters are optimized in FP32 while model computation uses BF16; Figure 7 plots the stage losses for Qwen3-1.7B at 16K context and 16× compression.Source: §A.2 · p. 12, Table 6 · p. 12, Figure 7 · p. 14

A.3 BASELINES AND BUDGETS

The evaluation compares compression baselines and an uncompressed reference, with methods receiving the document before the evaluation question. Budgets count retained and latent KV pairs per layer and KV head.

A.3 BASELINES AND BUDGETS

  • The comparison includes KVZip+, Attention Matching, SnapKV, PyramidKV, and an uncompressed reference, with each compression method receiving the document before the question is revealed.Source: §A.3 · p. 14
  • Attention Matching uses repeat-prefill reference queries, RMS attention-based key selection, two bounded least-squares iterations, and least-squares value fitting; tested SnapKV and PyramidKV settings use a 32-token observation window.The tested Attention Matching variant omits on-policy query recollection and OMP; SnapKV and PyramidKV use mean pooling of width five and shared GQA entries.Source: §A.3 · p. 14
  • For target context length N0 and compression ratio r, the average source budget is B = N0/r KV pairs per layer and KV head, counting both retained and latent entries.Logical pair counts are distinguished from physical memory, which also depends on metadata, padding, and precision; query and continuation caches are outside the source budget.Source: §A.3 · p. 14

B TRAINING LOSS CURVES

The training curves show declining objectives for both stages, alongside the Phonebook Lookup prompt format and the generic prompts used to obtain teacher signals.

B TRAINING LOSS CURVES

  • Both stage objectives decline over 1,000 updates, but their loss values are not directly comparable because they measure different objectives.Stage 1 drops rapidly at first and continues declining; stage 2 declines more gradually later, with greater per-step variation.Source: B TRAINING LOSS CURVES · p. 14

C PROMPT TEMPLATE FOR PHONEBOOK LOOKUP

  • Figure 8 marks placeholders with angle brackets, uses an ellipsis for additional entries, and continues generation after “Answer:”.Bold section labels explain the template and are not part of the plain-text prompt.Source: Figure 8 · p. 15

D GENERIC PROMPTS FOR TEACHER ROLLOUTS

  • The fixed prompt bank combines 24 document-oriented requests with 20 response instructions, covering content, organization, and writing style without benchmark questions or answers.The same bank is used for every document.Source: D GENERIC PROMPTS FOR TEACHER ROLLOUTS · p. 15
  • Each prompt forks the full-cache teacher and generates up to 24 tokens, providing attention-query and next-token targets rather than additional training documents.Evaluation questions are held out until the compacted cache is constructed and optimized.Source: D GENERIC PROMPTS FOR TEACHER ROLLOUTS · p. 15

E RULER RESULTS BY SUBTASK

The RULER task-level figures expose accuracy patterns that aggregate scores can conceal, across compression ratios and context lengths.

E RULER RESULTS BY SUBTASK

Related papers

Loading 2610.20582v1…