NeurIPS 2026 (accepted, as reported by the authors) LLM Efficiency KV cache compression region-wise quotas attention-derived quality mass adaptive segmentation long-form reasoning

Adaptive Mass-Segmented KV Compression for Long-Form Reasoning¶

Conference: NeurIPS 2026 (accepted, as reported by the authors)
arXiv: 2605.23200
Code: https://github.com/EIT-NLP/AdaptiveMassSegment
Area: LLM Efficiency
Keywords: KV cache compression, region-wise quotas, attention-derived quality mass, adaptive segmentation, long-form reasoning

Version note: This note follows arXiv v2, dated 2026-09-29. The list previously used Adaptive Mass-Segmented KV Compression for Long-Context Reasoning; v2 changes Long-Context to Long-Form without changing the arXiv ID. The filename retains the original slug.

TL;DR¶

AMS leaves token-importance scorers unchanged and instead uses attention-derived quality mass to form adaptive segments and allocate retention quotas before local selection, mitigating contiguous context loss in long-form reasoning; on Math500 with a 256-token cache budget, AMS-Expected improves over AdaKV-ExpE2 by 16.0 percentage points.

Background & Motivation¶

When a large language model (LLM) generates a long chain-of-thought (CoT), historical keys and values accumulate. Even with unchanged model weights, decoding must read an increasingly long cache, increasing memory use and memory traffic. Training-free compressors such as TOVA, Expected Attention, KeyDiff, and R-KV periodically discard KV entries using attention, geometric variation, or redundancy estimates. Their scorers answer which tokens are important, but global Top-k places all positions in one competition pool: sufficiently strong local peaks can leave an entire region of more distributed intermediate reasoning without any retained position.

This failure is more consequential than removing a few low-scoring words. An intermediate region can carry problem constraints, temporary conclusions, and their supporting derivation. Losing the region can cause subsequent generation to change the problem, overturn a correct conclusion, or repeat an earlier derivation. The authors call severe contiguous context eviction Region Wipe-out. Fixed-length chunks reduce direct token competition, but importance density is neither spatially uniform nor stationary during decoding: two equally long intervals can contain dense useful reasoning and mostly redundant narration, respectively. Partitioning by fixed positions can therefore still allocate scarce capacity to the wrong regions.

AMS separates where memory capacity should be assigned from which positions should survive inside each region. Recent attention defines the spatial distribution for the former, while the original scorer handles the latter; historical information smooths regional allocation across compression events without rewriting the base scoring rule. Core Idea: allocate feasible retention quotas to adaptive regions defined by attention-derived quality mass, then let tokens compete locally, preventing global score spikes from monopolizing the cache.

Method¶

Overall Architecture¶

AMS is a decoding-time cache-selection wrapper, not a new model, and requires no training. The model first processes the complete prompt normally and triggers compression at fixed generation intervals. At each event, selection runs independently for each layer and KV head, producing a compact cache at the target length. Two inputs have distinct roles: recent attention usage constructs quality mass, while the base scorer ranks candidates within segments; these signals need not coincide.

Each compression event constructs quality mass, optionally mixes in EMA credit, partitions the current cache order by cumulative mass, and repairs overly short or long segments. It then assigns regional budgets and must-keep priority, followed by local Top-k selection, budget correction, and gathering. A region is an interval of current cache positions, not a complete semantic reasoning step identified by the model; after compression, cache intervals no longer correspond to equally long intervals of original generation positions.

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["Current KV + recent attention"] --> B["Quality Mass and<br/>Historical Smoothing"]
    B --> C["Cumulative-Mass<br/>Adaptive Segmentation"]
    C --> D["Regional Quotas and<br/>Retention Priority"]
    D --> E["Local Selection and<br/>Physical Compaction"]
    F["Base scorer"] --> E
    E --> G["Compact KV → continue decoding"]
    E -->|Align historical state with keep indices| B

The diagram contains inference data flow and state reuse across compression events, not training supervision. The base scorer can be TOVA, Expected Attention, TriAttention, KeyDiff, or R-KV. AMS changes how scores constrain selection; it does not require all these scorers to become attention-based.

Key Designs¶

1. Quality Mass and Historical Smoothing: assign capacity by regional usage density without replacing token scoring

The method collects attention usage for each cache position from the most recent \(W\) decoding queries. Its implementation applies short-kernel one-dimensional average pooling along the sequence to reduce local noise. Recent suffix positions have fewer valid observations because of causal masking, so missing observations are padded with the maximum score in the window, avoiding an automatic penalty for newly generated positions. Usage is made nonnegative, receives a small additive constant, and is normalized:

\[ m_i=\frac{\max(u_i,0)+\epsilon}{\sum_{v=1}^{T}(\max(u_v,0)+\epsilon)},\qquad \epsilon>0. \]

Here \(u_i\) denotes usage and \(m_i\) a mass share, with all shares summing to one. Mass is not a ground-truth measure of how much correct reasoning a token contains, nor another name for the base score \(g_i\): mass determines boundaries and budgets, whereas \(g_i\) determines local winners under those budgets. With a geometric scorer, for example, regional capacity still comes from attention usage while local competition can depend on key geometry. This division of responsibilities enables the plug-in design.

A current query can temporarily concentrate on one region. Reallocating entirely from instantaneous usage would make boundaries and keep sets fluctuate. AMS maintains historical credit for every layer and KV head, updates its exponential moving average (EMA), and mixes it with current mass:

\[ c\leftarrow\lambda c+(1-\lambda)m^{(e)},\qquad m_{\mathrm{used}}=\operatorname{normalize}\!\left(\beta m^{(e)}+(1-\beta)\operatorname{normalize}(c)\right). \]

Only the mass used for subsequent segmentation and allocation changes; the base scorer remains untouched. Defaults \(\lambda=0.9\) and \(\beta=0.9\) prioritize current usage while allowing consistently useful historical positions to influence allocation. They do not mean that 90% of the final mass comes from history.

Credit must remain aligned with the actual KV positions: after compaction, historical credit needs to follow the same retained indices, and newly appended tokens require corresponding new state. Otherwise an unchanged vector index would refer to a different historical token. The paper specifies the EMA update and cross-event reuse but does not fully spell out post-compression credit reindexing and initialization for new positions. These are state-management details to verify during reproduction, not details to invent as the authors' exact implementation. Historical smoothing cannot recover already evicted KV entries.

2. Cumulative-Mass Adaptive Segmentation: use finer partitions in dense regions and longer coverage in sparse regions

AMS computes a mass prefix sum along the current cache. For cumulative thresholds \(\Delta,2\Delta,3\Delta,\ldots\), the first position reaching each threshold becomes a cut point. It therefore cuts after approximately equal amounts of attention mass, rather than equal token counts: dense regions reach a threshold quickly and form shorter segments, whereas sparse regions require more positions and form longer segments. The central boundary rule is:

\[ C(t)=\sum_{u=1}^{t}m_u,\qquad b_k=\min\{t\mid C(t)\geq k\Delta\}. \]

Initial segments have roughly equal mass, so a dense region produces more segments. Because every segment subsequently receives a minimum quota, segmentation itself supplies denser structural protection there; even with similar initial segment masses, the floor is not uniform per token. For isolated spikes and long sparse intervals, split/merge heuristics repair segment lengths: overly long segments are split into roughly equal subsegments, while short neighboring segments are merged. Repaired segments need not retain equal mass, so their masses are recomputed for allocation.

The default target segment mass is 0.1, with minimum/maximum lengths of 16/256. Length constraints avoid repeatedly assigning floors to tiny segments and allowing a long sparse interval to have only a handful of representatives. This is heuristic partitioning, not optimal semantic segmentation; its boundaries should not be described as the start and end of each reasoning step.

3. Regional Quotas and Retention Priority: floors require a feasible budget, and must-keep positions precede ordinary competition

After segmentation, each segment first receives a floor no larger than its length, and remaining capacity is distributed by segment mass. Before the additional must-keep correction, the main paper expresses allocation as:

\[ q_i^{\min}=\min(q_{\min},L_i),\qquad T_{\mathrm{rem}}=T_{\mathrm{keep}}-\sum_i q_i^{\min},\qquad q_i=q_i^{\min}+\operatorname{round}\!\left(T_{\mathrm{rem}}\frac{M_i}{\sum_j M_j}\right). \]

\(L_i\) is segment length and \(M_i\) its total mass. The default floor is one retained position per segment. Rounding can produce the wrong total, and a segment can receive more capacity than its length, requiring clipping and redistribution. This stage sets how many slots a region receives; it does not prescribe retaining its beginning, end, or complete sentences.

Keeping at least one position per segment first requires enough total capacity to pay for every floor. With must-keep positions, feasibility should additionally account for their occupied slots and the floors of regions they do not cover; checking only the segment count is insufficient. AMS normally protects the first four sink tokens and a recent suffix. The main text states that the implementation inserts must-keep positions first and allocates the remaining segment budget afterward. If the remaining budget is negative, it reduces the recent suffix while prioritizing sinks. An arbitrarily small budget cannot simultaneously retain every sink, the entire recent suffix, and every segment floor.

Taking the union of local selections and must-keep positions counts duplicates only once. Excess positions are removed by dropping low-scoring non-must-keep entries; deficits are filled with high-scoring unselected positions. The paper also allows segment quotas to be reduced when necessary. Consequently, avoiding regional erasure is a design objective under feasible budgets, not a strong theorem covering every region after every over-budget correction. It certainly does not imply preservation of all reasoning information: retaining a representative token is different from retaining complete constraints, derivations, and proofs.

4. Local Selection and Physical Compaction: change the competition pool and actually remove evicted KV entries

The base scorer supplies one score per cached position in a tensor of shape \([B,H_{kv},T]\). AMS selects high-scoring positions locally under each quota, applies union, deduplication, and budget correction, sorts retained indices, and gathers keys and values with the same indices. Different KV heads can retain different historical positions, but each receives a fixed-length compact cache. Subsequent decoding reads the retained KV entries rather than reading the full cache and masking unwanted positions.

Gather-and-compact is fundamentally different from mask-only execution. The AdaKV-ExpE2 implementation used in this paper enforces an effective budget through attention masks without physically releasing all evicted KV entries. AMS reduces the actual cache length to the budget. Equal effective lengths in accuracy comparisons therefore do not automatically imply equal physical memory in system comparisons. This distinction concerns the implementation evaluated here; it does not establish that every AdaKV implementation lacks compaction.

Appendix E additionally separates retention policy from paged storage layout. AMS produces head-wise keep indices; the runtime allocates replacement blocks, copies KV from original positions into new physical slots per head, replaces the request block table, and frees old blocks. The same compact slot can contain KV from different original positions across heads. Selection is already materialized in the contents, so steady-state attention needs no extra per-head indirection table. RoPE is already encoded in cached keys, but the next token's logical position must follow original generation progress rather than the compact cache length.

Here, vLLM compatibility means validation of the systems interface and runtime path. Main experiments use HuggingFace/KVPress, while supplementary code includes a reference adapter and vLLM-style runtime hooks. The authors explicitly distinguish this from an upstream vLLM patch and a fully optimized production-serving benchmark. Avoiding additional steady-state indirection does not eliminate compression events, GPU copies, temporary coexistence of old and replacement blocks, or batching overhead.

A Worked Example¶

The following constructed example illustrates allocation, not an additional paper experiment. Suppose one KV head currently has 16 positions and a final budget of eight. Segmentation and length repair produce four segments of length four, each with mass 0.25. A floor of one consumes four slots; distributing the remaining four equally gives each segment two slots.

Suppose mandatory beginning and ending positions are already included in the two slots for the first and last segments, leaving two retained positions per segment after correction. The base scorer now ranks candidates within each segment rather than letting all 16 positions compete for eight global slots. Even if the second segment's highest score is lower than several peaks elsewhere, it still has representatives. Global Top-8 could instead remove that segment entirely.

After new tokens arrive, AMS estimates mass and partitions the compact cache plus appended positions again. It does not freeze the original four segments, and retaining two positions from the second segment does not guarantee preservation of its full semantics.

Loss & Training¶

AMS has no training loss and updates no model weights; it is a training-free decoding-time policy. Main experiments compress every 512 generated tokens, with target effective lengths of 256, 512, or 1024 and a hidden-state buffer of 256. The recent query window for mass estimation is 128. Periodic compression allows the cache to grow between events, so the target is not an assertion that physical cache length equals that value at every instant.

AMS adds mass normalization, prefix sums, length repair, and quota allocation. Appendix H describes mass processing and prefix sums as \(O(H_{kv}T)\), but this is not the total compressor cost: base scoring, local selection, KV movement, and runtime scheduling must be accounted for separately.

Key Experimental Results¶

Main Results¶

The table below selects Math500 pass@1 results from Tables 1 and 5, using DeepSeek-R1-Distill-Qwen-7B. Columns denote target effective cache lengths after compression; values are percentages. Full KV is an uncompressed reference, not a fixed-budget policy.

Method	256	512	1024
Full KV	52.80	52.80	52.80
TOVA	29.20	44.60	48.80
AMS-TOVA	36.40	48.40	53.40
AdaKV-ExpE2	32.60	46.60	53.40
AMS-Expected	48.60	54.00	54.20
TriAttention	55.00	56.80	57.20
AMS-TriAttention	56.20	60.60	63.60

AMS-Expected improves over AdaKV-ExpE2 by 16.0, 7.4, and 0.8 percentage points; AMS-TOVA improves over TOVA by 7.2, 3.8, and 4.6 points. AMS-Expected slightly exceeds Full KV at budgets 512/1024, but this is an observation under this evaluation, not a law that compression outperforms complete context. TriAttention already exceeds Full KV and gains another 1.2, 3.8, and 6.4 points with AMS, supporting the interpretation that allocation can complement a strong scorer.

Results are not uniformly superior across tasks. In Table 6, AMS-Expected scores 27.44 on RepoBench-P versus 24.04 for Full KV, but its NIAH score is 0.2209 versus 0.3751 for Full KV. For cross-backbone transfer, AMS-Expected reaches 45.4% at budget 512 on OpenThinker3-7B, against an uncompressed reference of 45.8%. On 32B Math500 at the same budget, it reaches 43.00% versus 47.80% without compression (Tables 3 and 4).

Ablation Study¶

This table follows Table 7: Math500, the 7B backbone, and a TOVA scorer, with pass@1 in percent.

Config	512	1024	Note
Full AMS	48.4	53.4	Adaptive segmentation, mass quotas, EMA
Without mass-weighted quotas	46.0	52.8	Drops of 2.4 and 0.6 points
Without EMA credit	48.0	50.8	Drops of 0.4 and 2.6 points
Fixed-length segments	48.0	50.4	Drops of 0.4 and 3.0 points
Global head-only Top-k	42.8	49.2	Drops of 5.6 and 4.2 points

The global head-only Top-k ablation is not identical to the bare TOVA row in Table 1: their budget-512 scores are 42.8 and 44.6, respectively, and should not be merged into one configuration. Within the authors' ablation controls, removing regional quotas causes the largest loss; EMA and adaptive boundaries are more beneficial at budget 1024.

System measurements from Table 8 are listed separately below. Peak memory includes overall runtime allocations, not just KV bytes. Memory and time columns use different budgets and must not be cross-paired to calculate throughput.

Method	Peak GB: 512	Peak GB: 1024	Seconds/sample: 128	Seconds/sample: 512
StreamingLLM	15.0	14.9	50.1	48.4
TOVA	15.0	14.9	67.1	52.3
PyramidKV	15.3	15.6	67.5	51.6
AdaKV-ExpE2	39.6	39.7	91.8	62.3
AMS-Expected	15.0	14.9	41.4	44.0

Key Findings¶

Tight budgets expose the value of structural protection: at budget 256, adding AMS to KeyDiff improves accuracy from 22.80% to 42.80% (Table 5), without replacing the scorer.
Appendix C reports Region Wipe-out Rate decreasing from 11.3% to 8.7%, not to zero. Higher retained-set IoU supports more stable historical selection, not lossless semantic preservation.
Under free generation, TOVA/AMS-TOVA have repetition rates of 21.58%/16.18%, average generated lengths of 2963.6/2799.7 tokens, and times of 67.34/63.85 seconds (Table 9). Runtime gains include generating less repetitive content and are not per-step kernel speedups.
Small-sample sensitivity results require caution: Table 16 reports 50.0% for windows 16/32/64, and some appendix experiments use only 10% of Math500. This does not establish insensitivity across all windows and tasks.

Highlights & Insights¶

AMS primarily changes competition rules rather than inventing another importance score. Regional floors impose coverage constraints while strong scorers continue to handle local ranking.
Mass-balanced segmentation turns importance density into resolution: dense regions receive more, shorter segments and therefore denser floor protection. Segmentation and allocation must be understood together, not reduced to ordinary chunked Top-k.
Applying EMA to allocation rather than scoring separates historical stability from the scorer. Similar ideas could support other capacity-constrained online memories, but state must remain aligned with object identity.

Limitations & Future Work¶

Attention usage is only a proxy for long-range utility. Low-attention problem constraints or evidence needed later can still be evicted, and a few representatives do not preserve complete reasoning semantics.
Floors, must-keep positions, and total capacity have feasibility constraints; global correction can change regional counts. Reproductions should document what happens when protections cannot all be satisfied rather than claim no regional erasure at any budget.
Main systems results use HuggingFace/KVPress on A800 80 GB GPUs, with two GPUs per job by default. The vLLM path validates compatibility; dynamic batching, block-pool pressure, and transient compaction peaks require dedicated production measurements.
The paper has drifting table cross-references: GSM8K is actually Table 2, general tasks Table 6, and memory/latency Table 8. This note follows captions and values. Appendix Table 17 contains an insufficiently explained raw metric_main; it is not treated as accuracy or used to reconstruct missing results.

vs TOVA / Expected Attention: These methods supply token importance; AMS supplies temporal regional quotas within each head, allowing composition. AdaKV-ExpE2 also adapts budgets across heads, so comparison with AMS-Expected is not a single isolated local switch.
vs ChunkKV / fixed chunks: Chunking protects local structure with relatively fixed partition granularity. AMS adjusts boundaries using recent mass and continues local score-based selection rather than requiring entire chunks to survive together.
vs ReST-KV / G-KV: Related methods use history to stabilize eviction signals; AMS smooths the allocation layer. A useful next step is to separate the gains attributable to better scores from those attributable to better budget organization.

Rating¶

Novelty: 4/5 — An independent adaptive temporal-quota wrapper, distinct from changing scores or head budgets alone.
Experimental Thoroughness: 4/5 — Multiple scorers, backbones, and tasks, but no complete production-serving benchmark or extensive statistical uncertainty analysis.
Writing Quality: 3/5 — Clear central mechanism, with table references, some appendix metrics, and state-management details needing clarification.
Value: 4/5 — A composable structural-protection strategy for training-free KV compression, particularly under tight budgets in long-form reasoning.