Latent Compaction with Test-Time Training

Summary (Overview)

• Novel approach to KV cache compression: Latent Compaction learns a compact, document-specific KV cache through test-time training, treating the cache as a trainable neural network with frozen model weights.
• Two-stage training pipeline: Stage 1 fits attention outputs and log-normalizers to the full cache (with reconstruction-selected entries frozen), followed by Stage 2 end-to-end distillation from the full-cache model using generic document prompts.
• Complementary initialization strategy: Cache entries are initialized using three rules—KVZip+ reconstruction scores (preservers), teacher attention (anchors), and key-space diversity (diversifiers).
• Superior empirical results: Achieves 96.7% accuracy at 32× compression on RULER (vs. 40.0% for KVZip+), with strong gains on LongBench and Phonebook Lookup across multiple model scales and families (Qwen3, MiniCPM5).
• No access to downstream questions: Training uses only generic prompts about the document, making the optimized cache reusable across unseen queries.

---

Introduction and Theoretical Foundation

Background and Motivation

Large key–value (KV) caches are a critical bottleneck in efficient LLM inference. Long contexts increase memory use and attention costs during generation, limiting batch sizes and increasing memory-bandwidth demands. Frontier-model platforms already offer context compaction (e.g., Claude's text summaries), but text summaries can omit details whose importance only becomes clear later.

The key research question: Can gradient descent on a compact KV cache preserve document information under aggressive memory budgets?

Theoretical Foundation

The paper builds on several theoretical threads:

1. KV-cache eviction (H₂O, SnapKV, PyramidKV, KVZip+): Selecting original KV entries using attention or reconstruction scores.
2. Learned context representations (Prefix-tuning, AutoCompressors, Cartridges, Attention Matching): Optimizing compact continuous representations with frozen model weights.
3. Test-time training (TTT): Adapting model parameters using self-supervised objectives on unlabeled test inputs before prediction.

The core insight is that the KV cache can be treated as a neural network itself, optimized through test-time training to preserve document information for unseen questions.

Attention Preliminaries

For a query head with query $q$, keys $k_i$, values $v_i \in \mathbb{R}^d$, and cache entries indexed by set $\mathcal{I}$:

$$A(q) = \frac{\sum_{i \in \mathcal{I}} \exp\left(q^{\top} k_i / \sqrt{d}\right) v_i}{Z(q)}, \qquad Z(q) = \sum_{i \in \mathcal{I}} \exp\left(q^{\top} k_i / \sqrt{d}\right) \tag{1}$$

where $Z(q) > 0$ is the scalar normalizer and $s = 1/\sqrt{d}$ is the attention scale.

---

Methodology

Memory Architecture

Given an average budget of $B$ KV pairs per layer and KV head, Latent Compaction initializes $LHB$ KV slots total. Each slot carries a scalar attention bias $b_{\text{cmp}}$.

For a query $q \in \mathbb{R}^d$ with compacted keys $K_{\text{cmp}}$, values $V_{\text{cmp}} \in \mathbb{R}^{n \times d}$ and biases $b_{\text{cmp}} \in \mathbb{R}^n$:

$$\ell(q) = s K_{\text{cmp}} q + b_{\text{cmp}} \tag{2}$$ $$A_{\text{cmp}}(q) = V_{\text{cmp}}^{\top} \mathrm{softmax}(\ell(q)) \tag{3}$$

Handling subsequent tokens: Question and generated tokens append ordinary KV pairs after the document. The combined output is:

$$A_t = \frac{Z_{\text{cmp}}(q_t) A_{\text{cmp}}(q_t) + Z_{\text{new},t} A_{\text{new},t}}{Z_{\text{cmp}}(q_t) + Z_{\text{new},t}} \tag{5}$$

where $Z_{\text{cmp}}(q_t) = \mathbf{1}_n^{\top} \exp(\ell(q_t))$ is the compressed-cache normalizer.

Initialization Strategy

Three complementary entry types:

Anchor scores are computed as:

$$a_i = \max_{q \in \mathcal{Q}} \left[ \mathrm{softmax}(s K q) \right]_i \tag{6}$$

Diversifiers are selected via farthest-first traversal:

$$i^{\star} = \arg\max_{i \notin S} \min_{j \in S} \| k_i - k_j \|_2^2 \tag{7}$$

Budget allocation: $M = B/4$ slots per head for anchors + diversifiers; remaining $LH(B - M)$ slots from global KVZip+ ranking.

Training Stage 1: Attention Fitting

Fits each layer independently using captured teacher queries:

$$\mathcal{L}_{\text{attn}} = \mathrm{MSE}\left(\hat{A}_{\text{cmp}}, \hat{A}_{\text{full}}\right) + \mathrm{MSE}\left(\widehat{\log Z}_{\text{cmp}}, \widehat{\log Z}_{\text{full}}\right) \tag{8}$$

Hats denote standardization using fixed teacher statistics. Preservers are frozen during this stage; anchors and diversifiers receive updates.

Training Stage 2: End-to-End Refinement

All entries are jointly optimized using KL divergence and language modeling loss:

$$\mathcal{L}_{\mathrm{KL}} = \frac{1}{T} \sum_{t=1}^{T} \left[ \sum_{v \in \mathcal{T}_t} p_{t,v} \log \frac{p_{t,v}}{r_{t,v}} + p_{t,\text{tail}} \log \frac{p_{t,\text{tail}}}{r_{t,\text{tail}}} \right] \tag{9}$$ $$\mathcal{L}_{\mathrm{LM}} = -\frac{1}{T} \sum_{t=1}^{T} \log r_{t,c_t}, \qquad \mathcal{L}_{\mathrm{stage2}} = \mathcal{L}_{\mathrm{KL}} + \mathcal{L}_{\mathrm{LM}} \tag{10}$$

The teacher's top-20 token set $\mathcal{T}_t$ is retained individually; the rest is grouped into a tail category.

Teacher Rollouts

• 480 generic text prompts (24 document requests × 20 response instructions)
• Each rollout generates up to 24 continuation tokens from the frozen full-cache model
• No downstream questions or answers are used
• Teacher forcing: student receives teacher tokens, not its own predictions

---

Empirical Validation / Results

Experimental Setup

• Default model: Qwen3-1.7B (frozen weights, BF16 computation, FP32 KV cache)
• Training: 1,000 stage-1 updates + 1,000 stage-2 updates, AdamW with learning rates $3 \times 10^{-3}$ and $10^{-4}$
• Baselines: KVZip+, Attention Matching, SnapKV, PyramidKV, full-cache, text-summary references

Main Results

RULER (16K context): Latent Compaction achieves 96.7% accuracy at 32× compression vs. 40.0% for KVZip+.

LongBench: Outperforms text summarization throughout and KVZip+ at 16× and 32× compression across Qwen3-0.6B, 1.7B, and 4B models.

Phonebook Lookup: Large gains at 16× compression; all methods approach zero at 64×+ compression.

Transfer Results

Retrieval advantages extend to:

• Qwen3: 0.6B, 1.7B, 4B, 8B model scales
• MiniCPM5: Different model family

Ablation Results

Table 1: Initialization ablations (Phonebook Lookup, 16× compression)

Table 2: Stage-1 ablations (Phonebook Lookup, 16× compression)

Table 3: Stage-2 ablations (Phonebook Lookup, 16× compression)

Table 4: Architectural ablations (Phonebook Lookup, 16× compression)

Key Findings from Ablations

1. Diversifiers are crucial: Replacing them with random or Gaussian vectors causes catastrophic accuracy drops (~60 points)
2. Stage 1 is essential: Removing it reduces accuracy from 80.67% to 13.92%
3. Freezing preservers in stage 1 matters: Trainable preservers reduce accuracy by ~47 points
4. Log-normalizer loss is critical: Removing it drops accuracy by ~44 points
5. Normalizer weighting (Eq. 5) is vital: Equal averaging collapses accuracy to ~0%
6. Optimal training duration: 1,000 updates for both stages balances accuracy and cost

---

Theoretical and Practical Implications

Theoretical Contributions

1. KV cache as a trainable neural network: The paper demonstrates that treating the KV cache as a parameterized model and optimizing it with gradient descent can preserve retrieval accuracy under aggressive compression—a paradigm shift from traditional eviction-based methods.

2. Two-stage curriculum: The staged approach (attention fitting → end-to-end distillation) mirrors curriculum learning, first matching local attention statistics then refining global prediction quality.

3. Complementary selection theory: The combination of reconstruction importance (KVZip+), teacher attention, and key-space diversity provides theoretical justification for covering both information-theoretic importance and geometric coverage.

Practical Implications

1. Question-agnostic compression: Since training uses only generic prompts, the optimized cache can be reused across arbitrary downstream queries—critical for real-world deployment where questions are unknown in advance.

2. Model transferability: The procedure transfers across model scales (0.6B–8B) and families (Qwen3, MiniCPM5), suggesting broad applicability.

3. Logical vs. physical budget: The paper explicitly notes that results concern logical KV-pair budgets; physical memory savings depend on precision, biases, and padding choices.

Limitations

• Per-document training cost: 1,000 + 1,000 updates per document is computationally expensive, though amortizable across repeated use
• Extreme compression degradation: All methods approach zero accuracy at 64×+ compression
• Text summarization advantage at extreme ratios: On RULER at the most aggressive ratios, text summaries outperform Latent Compaction

---

Conclusion

Latent Compaction introduces a novel paradigm for KV cache compression: treating the cache as a trainable document-specific neural network optimized through test-time training. The method achieves substantial gains over budget-matched baselines across RULER, LongBench, and Phonebook Lookup benchmarks, with strong transfer across model scales and families.

Key takeaways:

• Initialization via complementary selection rules (preservers + anchors + diversifiers) is critical for success
• Two-stage training (attention fitting followed by end-to-end distillation) provides complementary benefits
• Attention normalization (biases and normalizer weighting) is essential for correct joint attention over document and subsequent tokens

Future directions:

• Reducing per-document preparation cost
• Studying broader task types beyond retrieval and long-context understanding
• Exploring physical memory optimization (precision, padding) alongside logical KV-pair budgets

The results demonstrate that treating the KV cache as a neural network and optimizing it with test-time training can preserve retrieval accuracy under aggressive compression, opening new avenues for efficient long-context LLM inference.