PaperlayerView exact revision on arXiv

Source-linked AI summary

FlexComp: One Model for Every Ratio in Context Compression

Kaiyan Zhao, Zhongtao Miao, Akiko Aizawa, Yoshimasa Tsuruoka

arXiv:2609.11192v1cs.CL

TL;DR

Soft context compressors fix one ratio for all inputs and require separate models for different ratios, despite inputs needing different budgets. FlexComp trains one any-ratio compressor with Matryoshka-style budget sampling and selects K per input through cascade routing or a learned predictor. Across three compressors on MRQA, it matches fixed-ratio specialists while achieving 158–266× average compression and improving serving efficiency.

  • Problem

    Existing soft context compressors require a separately trained model for each fixed ratio and apply that ratio uniformly despite inputs needing different budgets.

  • Method

    FlexComp samples the memory budget K during Matryoshka training and selects it per input with confidence-based cascade routing or a lightweight K predictor.

  • Results

    Across ICAE, 500xCompressor, and SAC on MRQA, one model matches fixed-ratio specialists while adaptive selection reaches 158–266× average compression and improves serving efficiency.

  • Takeaways & Limitations

    Compression ratio can be treated as a per-input decision supported by one model rather than a fixed deployment constant.

  • Takeaways & Limitations

    Sampling budgets during pretraining loses 2.3 ID and 2.5 OOD F1 versus sampling budgets only during supervised fine-tuning.

Abstract

from arXiv · show

Soft context compression condenses a context into a few memory tokens that a frozen LLM consumes in place of the raw text, but existing compressors fix the compression ratio at training and inference: each deployed ratio requires a separately trained model, and the chosen ratio is applied uniformly to all inputs, whose actual needs vary drastically. We propose FlexComp, a method-agnostic framework that decouples the ratio from both training and deployment: Matryoshka-style training samples the memory budget $K$ per instance, turning one model into an any-ratio compressor, and the budget is then chosen per input by: (1) confidence-based cascade routing or (2) a lightweight learned $K$ predictor. Across ICAE, 500xCompressor, and SAC on MRQA, a single FlexComp model matches separately trained fixed-ratio specialists with minimal degradation. Cascade routing preserves over 98% of the mildest ratio's accuracy at up to 266x average compression; the $K$ predictor, in a single compression-decoding pass, reaches 158-236x within 0.7 F1 of the mildest ratio. At serving-scale batch sizes, the $K$ predictor cuts context KV cache by 50% and improves decoding throughput by 47%.

1 Introduction

FlexComp addresses the fixed-ratio, uniform-budget limitations of soft context compression with one model that supports multiple budgets and selects them per input. Across three compressors on MRQA, it preserves specialist accuracy while substantially improving compression and serving efficiency.

  • Motivation: Different inputs require different budgets: many are solved with K=1, while a minority needs larger budgets, making any fixed K wasteful or insufficient.The minimal sufficient budget varies across inputs and domains.
  • FlexComp: FlexComp replaces separately trained fixed-ratio compressors with one Matryoshka-trained model supporting any budget without architectural changes.The method samples the soft prompt length K per instance during training and supports ICAE, 500xCompressor, and SAC.
  • Results: A single Matryoshka-trained model matches fixed-ratio specialists within 1.3 F1 in the worst case across in-domain and out-of-domain MRQA tasks.It often performs better on out-of-domain data while reducing the number of trained models from n to one.
  • Adaptive budget selection: Cascade routing retains over 98% of the mildest ratio’s accuracy while increasing average compression from 15× to as high as 266×.It escalates to larger budgets only when decoder confidence is insufficient.
  • Adaptive budget selection: The K predictor reaches 158–236× compression within 0.7 F1 of the mildest ratio in a single compression-decoding pass.At serving-scale batch sizes, it cuts context KV cache by 50% and raises aggregate decoding throughput by 47%.
  • Analysis: Larger budgets can sometimes hurt because extreme compression removes misleading surface details that would otherwise confuse the decoder.The accuracy-compression trade-off is therefore non-monotonic for some individual inputs.

2 Related Work

Prior context-compression research includes soft memory-token methods and hard prompt-shortening methods, while FlexComp extends nested, budget-conditioned representations to adaptive soft compression.

  • Context compression: Soft compression encodes contexts into continuous memory tokens consumed by a frozen LLM, whereas hard compression filters or summarizes discrete text.FlexComp builds on the soft-compression family.
  • Soft context compression: Existing soft compressors differ in how they acquire memory tokens and supervise them, but share the core idea of representing contexts through soft tokens.Examples include ICAE, 500xCompressor, and SAC.
  • Nested and budget-conditioned representations: FlexComp transfers the one-model-many-budgets principle from Matryoshka representation learning to context compression by conditioning the encoder on memory slot count.Unlike embedding truncation, its budgets use distinct encoder forward passes for different K values.
  • Adaptive computation and cascades: FlexComp applies adaptive-computation ideas within one model, escalating over compression budgets rather than switching between models.Its K predictor commits to a compute level before execution, while the cascade uses confidence-based escalation.
  • KV-cache efficiency: Soft-token compression avoids materializing the full-context KV cache, unlike eviction and quantization methods that reduce memory after prefill.Hard compression and KV-cache methods address different parts of the serving pipeline.

3 Method

FlexComp makes the compression budget an input-dependent choice rather than a fixed architectural constant, using Matryoshka training and two inference-time selection strategies. The framework supports confidence-based escalation after decoding and learned budget prediction before compression.

  • 3.1 Preliminaries: Soft-Token Context Compression: Soft-token compression encodes an L-token context into K memory tokens that a frozen decoder uses instead of the raw context, shrinking its KV-cache footprint from L to K entries.The framework applies this interface to ICAE, 500xCompressor, and SAC, whose fixed-budget designs otherwise support only one K.
  • 3.2 Matryoshka Training: Sampling K per training instance turns one fixed-ratio compressor into a single model supporting multiple budgets without architectural changes or same-budget inference overhead.The encoder separately learns representations for each supported K, replacing n training runs and deployed models with one.
  • 3.3 Inference I: Cascade Routing: Cascade routing starts at the most aggressive budget and re-encodes only when length-normalized decoder confidence falls below a calibrated threshold.Thresholds are tuned on held-out data to trace an accuracy–average-budget operating curve, and confidence correlates with token-level F1 up to Spearman’s ρ = 0.58.
  • 3.4 Inference II: Learned K Prediction: The learned K predictor commits to one budget before compression using a lightweight MLP over budget-independent, mean-pooled context features.Its labels are the smallest budgets within ε = 0.1 of the best Rouge-1 F1, with low-quality full-budget instances excluded and classes balanced during training.

4 Experiments

Experiments show that Matryoshka training preserves fixed-ratio compressor quality while enabling per-input budget selection across three compressors on MRQA. Cascade routing and a learned K predictor improve the compression–accuracy trade-off without retraining, with the predictor supporting single-pass serving.

  • 4.1 Experimental Setup: Experiments use ICAE, 500xCompressor, and SAC on six ID and six OOD MRQA tasks with budgets K={1,10,34}, corresponding to 510×, 51×, and 15× compression.Fixed-ratio baselines use the same architectures, data, and compute budgets for controlled comparison.
  • 4.2 One Model Matches Ratio-Specific Models: At most about one F1 point separates the single Matryoshka-trained model from fixed-ratio specialists, while four of nine settings improve and seven of nine OOD deltas are positive.The clearest gain is 500xCompressor at 510×, with +3.72 ID and +3.12 OOD F1.
  • 4.3 Cascade Routing: Cascade routing reaches 235× average compression on ICAE while remaining within 0.6 ID F1 of always using K=34.OOD accuracy is slightly above the fixed baseline at every operating point, and the most aggressive point averages 19.8 memory tokens.
  • 4.4 K Predictor: 158–236× compression costs at most 0.7 F1 versus the mildest budget when the K predictor selects each input’s budget before compression.The predictor uses a single compression-decoding pass and roughly halves the average budget.
  • 4.4 K Predictor: The predictor outperforms random routing by about 5 F1 on both ID and OOD at the same average budget, showing that input-aware assignment matters.Random routing matches the predictor’s per-budget proportions but ignores the input.
  • 4.3 Cascade Routing: Cascade thresholds provide a tunable compression–accuracy trade-off, while the predictor supports single-query compression and budget-grouped batching.The two strategies target repeated-query compression and on-the-fly single-pass serving, respectively.

5 Analysis

FlexComp analysis shows that training-budget placement, confidence calibration, and adaptive allocation determine whether variable compression preserves quality and improves serving efficiency. Larger budgets can sometimes hurt, while balanced prediction and adaptive routing support strong compression across scales.

  • 5.1 Matryoshka in Pretraining Hurts: Sampling budgets during pretraining loses 2.3 ID and 2.5 OOD F1 at K=34 versus sampling only during SFT.FlexComp therefore keeps pretraining at a fixed budget and applies Matryoshka training during fine-tuning.
  • 5.2 K Predictor: Balanced resampling prevents the K predictor from collapsing onto K=1, after raw-label training gets 90% of the majority class right but only 13–16% of K=10.A second training epoch adds little, indicating that label imbalance rather than undertraining explains the failure.
  • 5.3 Cascade Routing: Decoder confidence correlates with answer quality at every budget, with Spearman ρ = 0.56 to ρ = 0.58, supporting confidence-based cascade thresholds.The relationship also holds at K=1, where confidence must distinguish solved inputs from confidently wrong answers.
  • 5.4 From KV Memory to Serving Throughput: More than 96% KV-cache reduction from 15× to 510× compression enables the K predictor to cut memory by 50% and raise throughput by 47% at longer contexts.The efficiency evaluation uses the Matryoshka-trained SAC model at batch size 96.
  • 5.5 More Memory Tokens Are Not Always Better: 11.8–12.8% of solvable in-domain samples score higher at K=1 than K=34, while 11.0–26.0% strictly require the full budget.Budget sensitivity increases with compressor strength, making adaptive allocation more important for stronger compressors.
  • 5.6 Scaling to 8B: A single Matryoshka model matches or exceeds fixed-ratio specialists across all six Llama-3.1-8B ID/OOD cells, while adaptive selection reaches 156×–241× compression.The K predictor reaches 183× compression with −0.42 ID and −0.72 OOD F1 against the Matryoshka 15× baseline.

6 Conclusions

FlexComp makes soft context compression adaptable to both model training and per-input deployment. One Matryoshka-trained model supports variable budgets, while routing or prediction preserves accuracy at high compression and improves serving efficiency.

  • 6 Conclusions: FlexComp decouples compression ratio from training and deployment, using one any-ratio model plus per-input budget selection to achieve 158–266× compression with 50% less KV memory and 47% higher throughput.The framework also finds that larger budgets can sometimes hurt by preserving misleading surface details.
Loading 2609.11192v1…