Summary (Overview)

• FlexComp is a method-agnostic framework that decouples the compression ratio from both training and deployment in soft context compression, enabling a single model to serve as an any-ratio compressor.
• Matryoshka-style training samples the memory budget $K$ per instance, turning one model into a family of fixed-ratio compressors with minimal degradation (within 1.3 F1 in the worst case).
• Two per-input budget selection strategies are introduced: confidence-based cascade routing (preserves over 98% of the mildest ratio's accuracy at up to 266× average compression) and a learned K predictor (reaches 158–236× compression in a single pass, within 0.7 F1).
• At serving-scale batch sizes, the K predictor cuts context KV cache by 50% and improves decoding throughput by 47%.
• A key finding: larger budgets can actively hurt on a non-trivial fraction of inputs, as extreme compression acts as an information bottleneck that sheds misleading surface details.

---

Introduction and Theoretical Foundation

Large language models (LLMs) rarely answer from a question alone; contexts such as retrieved passages, tool outputs, documents, or conversation history are needed for generation. Serving these context tokens is costly, as they must be prefilled on every query, and their Key-Value (KV) cache occupies memory for the whole generation, often dominating the memory footprint of inference.

Soft context compression condenses a context into a few continuous memory tokens that a frozen LLM consumes in place of raw text. However, existing compressors (e.g., ICAE, 500xCompressor, SAC) fix the compression ratio at training and inference time. This structural limitation means:

• Each deployed ratio requires a separately trained model.
• The chosen ratio is applied uniformly to all inputs, regardless of how much each input actually needs.

The paper demonstrates that different inputs require drastically different budgets. For instance, in MRQA in-domain subsets, most inputs are already solved with a single memory token ($K=1$, 510× compression), yet a non-trivial minority requires far larger budgets. A fixed $K$ therefore wastes tokens on easy inputs and starves hard ones.

The theoretical foundation draws on Matryoshka representation learning (Kusupati et al., 2022), which trains embeddings whose prefixes remain valid at multiple dimensionalities. FlexComp transfers this one-model-many-budgets principle to context compression, with a structural difference: budgets are realized by conditioning the encoder on the memory slot count rather than truncating a fixed representation.

---

Methodology

3.1 Preliminaries: Soft-Token Context Compression

An encoder LLM $E_\phi$ compresses a context $c$ of $L$ tokens into $K \ll L$ continuous memory tokens:

$$\mathbf{m}_{1:K} = E_\phi(c; K)$$

A frozen decoder LLM $D$ then consumes these tokens in place of $c$ to answer a query $q$:

$$\hat{a} = D(\mathbf{m}_{1:K}, q)$$

The training objectives (pretraining and fine-tuning) are:

$$\mathcal{L}_{\mathrm{PT}}(\phi) = \mathbb{E}_c [\ell(D(E_\phi(c; K)), t(c))],\tag{1}$$ $$\mathcal{L}_{\mathrm{SFT}}(\phi) = \mathbb{E}_{(c, q, a)} [\ell(D(E_\phi(c; K), q), a)].\tag{2}$$

3.2 Matryoshka Training

FlexComp makes the budget a conditioning variable. Let $\mathcal{K} = \{K_1 < K_2 < \cdots < K_n\}$ be the set of supported budgets. During training, for each instance a budget $K \sim p(\mathcal{K})$ is sampled:

$$\mathcal{L}(\phi) = \mathbb{E}_{(c, q, a)} \mathbb{E}_{K \sim p(\mathcal{K})} \big[ \ell(D(E_\phi(c; K), q), a) \big],\tag{3}$$

The single set of weights $\phi$ is trained to produce, for every $K \in \mathcal{K}$, a representation specialized to that budget. A key design choice: Matryoshka sampling is applied only during fine-tuning, not pretraining — applying it during pretraining degrades quality (see §5.1).

3.3 Inference I: Cascade Routing

The cascade starts from the most aggressive budget $K_1$ and iteratively escalates. At stage $i$, a confidence score is computed as the length-normalized log-probability of the decoded answer:

$$s_i = \frac{1}{|\hat{a}_i|} \sum_t \log p_\theta(\hat{a}_{i,t}).\tag{4}$$

If $s_i \geq \tau_i$, the candidate is accepted; otherwise, the model re-encodes at the next budget $K_{i+1}$. The final budget $K_n$ always accepts. Thresholds $\tau_1, \ldots, \tau_{n-1}$ are calibrated on held-out data.

3.4 Inference II: Learned K Prediction

The K predictor commits to a budget before compression. Labels are derived from the Matryoshka model itself: for each training instance, the model is decoded at every $K \in \mathcal{K}$, scored against the gold answer with Rouge-1 F1 ($F_K$), and the label is:

$$K^* = \min \left\{K \in \mathcal{K}: F_K \geq F_{K_n} - \varepsilon \right\},\tag{5}$$

with $\varepsilon = 0.1$. The predictor is a two-layer MLP $(d \to d/2$, ReLU, $d/2 \to n)$ that maps a budget-free encoder representation (mean-pooled hidden states) to a distribution over $K$. Balanced resampling (1:1:1) of the training set is necessary to prevent collapse to the majority class.

---

Empirical Validation / Results

4.2 One Model Matches Ratio-Specific Models

Table 1 compares fixed-ratio specialists with the single Matryoshka-trained model:

Table 1: F1 of fixed-ratio compressors vs. a single Matryoshka-trained model. $K=34/10/1$ correspond to $15×/51×/510×$ compression.

Key finding: Replacing three trained models with one costs at most ~1 F1 point, and in four of nine settings the single model is better. The clearest win is 500xCompressor at 510×, where the shared model exceeds its specialist by +3.72 ID and +3.12 OOD F1 — joint training regularizes the extreme K=1 case.

4.3 Cascade Routing

The cascade achieves aggressive compression with minimal accuracy loss. On ICAE, the most aggressive operating point uses an average of 19.8 tokens per input (235× average compression) while staying within 0.6 ID F1 of always running at K=34.

4.4 K Predictor

Table 2 shows the K predictor's performance:

Table 2: K predictor vs. uniform mildest budget. Random routing matches the predictor's average budget but ignores the input, isolating the benefit of routing the right inputs to smaller budgets.

5.4 From KV Memory to Serving Throughput

Table 5 shows inference efficiency at batch size 96:

Table 5: KV memory and throughput gains. The K predictor cuts memory by 50% and raises throughput by 47% at the longer context.

5.5 More Memory Tokens Are Not Always Better

On solvable samples, K=1 achieves strictly higher F1 than K=34 on 11.8–12.8% of in-domain samples; within these, on 6.8–7.7% the model fails completely at K=34 yet answers correctly from a single memory token. The fraction of samples strictly requiring the full budget (11.0–26.0%) grows with compressor strength (ICAE < 500xCompressor < SAC).

Three failure modes identified:

• Relation confusion: larger memory preserves competing relations, confusing the decoder
• Parametric-prior fallback: the decoder abandons memory and falls back on its prior
• Salient-entity anchoring: surface anchors in the memory are promoted to incorrect answers

5.6 Scaling to Llama-3.1-8B

All findings transfer to 8B scale. The single Matryoshka model now matches or exceeds its fixed-ratio specialists at every budget (+0.47 to +0.76 F1 across all six cells). The cascade's conservative point reaches 156× average compression at essentially no cost, and the K predictor reaches 183× compression at −0.42 ID and −0.72 OOD F1.

---

Theoretical and Practical Implications

Theoretical implications:

• The compression ratio need not be a design constant fixed once for all inputs; it can be a per-input decision that a single model can both support and make.
• Budget non-monotonicity: larger budgets can actively hurt on a non-trivial fraction of inputs — extreme compression acts as an information bottleneck that sheds misleading surface details. This challenges the assumption that more capacity is always better.
• The finding that Matryoshka sampling during pretraining hurts (§5.1) suggests that pretraining is where the compressor learns the basic skill of packing text, and randomizing slot count during this phase forces the model to split capacity across objectives before that skill is stable.

Practical implications:

• Model efficiency: One model replaces an entire family of fixed-ratio compressors, cutting the number of trained and deployed models from n to one.
• Deployment flexibility: Cascade routing offers a tunable, retraining-free tradeoff curve (ideal for corpora compressed once and queried repeatedly); the K predictor provides a single-pass operating point cheap enough for latency-sensitive serving.
• Serving-scale gains: At batch size 96, the K predictor cuts context KV cache by 50% and improves decoding throughput by 47%, with the effect growing with context length.
• Adaptive compression becomes more important as compressors improve: the fraction of samples strictly requiring the full budget grows with compressor strength, so adaptive budget allocation is increasingly valuable.

---

Conclusion

FlexComp decouples the fixed compression ratio of soft context compressors from both training and deployment:

• Matryoshka training turns one model into an any-ratio compressor that matches the accuracy of separately trained specialists (within 1.3 F1 in the worst case, often better).
• Per-input budget selection (cascade routing or a single-pass K predictor) preserves the mildest ratio's accuracy at 158–266× average compression, converting at serving scale into 50% less context KV memory and 47% higher decoding throughput.
• The analysis reveals that larger budgets can sometimes hurt, as extreme compression sheds misleading surface details.

The results suggest that the compression ratio need not be a design constant fixed once for all inputs, but a per-input decision that a single model can both support and make. Future directions include exploring more granular budget sets, extending to other compression architectures, and investigating the interaction between adaptive compression and other inference optimizations (e.g., KV cache eviction, quantization).