Sequence compression is a fundamental operation and a critical information bottleneck in efficient language and multimodal models, and yet prevailing approaches based on pooling or learned projections remain underperforming. We revisit sequence compression through the lens of test-time regression: input tokens define a regression problem whose fitted coefficients form the compressed representation. Importantly, the regression problem is learned end-to-end via backpropagation, while its coefficients are computed efficiently in closed form during the forward pass. This leads to TTR-Compressor, a general alternative to conventional pooling and projection-based compression. We evaluate TTR-Compressor as the compressor for block-wise sparse attention and the visual token merger for multimodal large language models. On block-wise sparse attention, TTR-Compressor improves RULER performance over attention-based compressors by a large margin. For visual token merging, it consistently outperforms the standard MLP merger and matches its performance with 55.6% fewer visual tokens. These results establish test-time regression as an effective primitive for sequence compression in efficient language and multimodal models.
Figures & tables
Figure 1: Overview of TTR-Compressor. (a) Unlike pooling and projection, which process tokens independently, TTR-Compressor fits all tokens jointly and thereby captures the interplay among them. (b) TTR-Compressor outperforms existing methods on both block-wise sparse attention and visual token merging.
Figure 2: Overview of TTR-Compressor. (a) Bilevel learning: The inner solve formulates sequence compression as test-time ridge regression with an exact dual Cholesky solve to obtain the representation h∗ , while the regression parameters are learned end-to-end via task gradients. (b) Input-dependent second-order adaptation: By jointly coupling tokens through the Gram matrix, TTR naturally diminishes the relative influence of redundant tokens compared to additive pooling. For intuitive illustration, we set D=I and uniform targets v=1 ; as typically learned in practice, a smaller regularizer D further amplifies this deduplication effect, yielding an even higher relative share for distinct tokens.
Model
S-NIAH-1
S-NIAH-2
S-NIAH-3
MK-NIAH-1
Avg.
4K
8K
16K
32K
4K
8K
16K
32K
4K
8K
16K
32K
4K
8K
16K
32K
Full Attention
0.834
0.428
0
0
0.850
0.412
0
0
0.572
0.470
0
0
0.374
0.272
0
0
0.2632
Sliding Window Attention
0.116
0.050
0.044
0.016
0.148
0.048
0.034
0.030
0.112
0.094
0.030
0.034
0.126
0.080
0.026
0.028
0.0635
Mean Pooling
0.938
0.536
0.470
0.270
0.840
0.258
0.042
0.050
0.292
0.282
0.040
0.042
0.366
0.230
0.042
0.056
0.2971
Concat & Projection
0.892
0.840
0.792
0.582
0.764
0.232
0.124
0.068
0.512
0.324
0.174
0.076
0.476
0.278
0.084
0.052
0.3919
Attention-Weighted Pooling
0.944
0.914
0.876
0.740
0.926
0.274
0.136
0.060
0.340
0.082
0.026
0.026
0.474
0.258
0.124
0.072
0.3920
Table 1: Accuracy on Single Needle-In-A-Haystack (S-NIAH) and Multi-Key Needle-In-A-Haystack (MK-NIAH) tasks from RULER across varied context lengths. Best values for each context length are bolded , and second-best values are underlined .
Model
ARC-c
ARC-e
BoolQ
COPA
Hella.
LMB.
OBQA
PIQA
SciQ
Wino.
Avg.
acc ↑
acc ↑
acc ↑
acc ↑
acc ↑
acc ↑
acc ↑
acc ↑
acc ↑
acc ↑
acc ↑
Full Attention
24.74
42.09
60.15
68.00
39.85
35.71
29.60
67.74
62.70
51.07
48.16
Sliding Window Attention
24.91
41.96
55.02
69.00
39.53
35.47
29.00
67.14
63.00
50.75
47.58
Mean Pooling
24.66
42.38
55.99
66.00
39.62
35.88
28.40
69.31
65.50
53.04
48.08
Concat & Projection
24.83
42.30
58.65
68.00
39.58
36.62
30.00
67.74
62.20
51.78
48.17
Attention-Weighted Pooling
25.00
41.75
61.22
65.00
39.61
35.61
29.40
68.93
63.80
50.51
48.08
Table 2: Performance comparison on language modeling and zero-shot common-sense reasoning. The evaluation context lengths are too short to activate sparse attention, so the compressor is not involved during inference; these results therefore only reflect the common-sense reasoning ability acquired during language model pretraining. Note that Concat & Projection has significantly more parameters than the others. Best values are bolded , and second-best values are underlined .
Method
General VQA
Knowledge Reasoning
Document Understanding
Spatial Under.
Avg.
MMBench
MMStar
RWQA
MMMU
AI2D
OCRBench
InfoVQA
ChartQA
MMSI
ERQA
2×2 Token Compression
MLP
76.63
49.93
64.44
46.22
75.49
62.80
41.72
70.64
26.20
35.25
54.93
TTR-Compressor
76.89
53.94
63.40
45.44
76.94
62.90
43.74
73.04
26.90
36.00
55.92
TTR-Compressor (w/ Conv)
75.86
52.26
66.93
45.67
76.65
63.80
44.33
72.52
26.30
36.00
56.03
3×3 Token Compression
Table 3: Evaluation across multimodal benchmarks under different token compression ratios. Results are grouped by 2×2 and 3×3 token compression. Best values within each compression group are bolded , and second-best values are underlined .
Paradigm
Complexity
Params
Mean Pooling
O(Td)
0
Linear Projection
O(Td2)
Td2
Attention (Reusing)
O(Td)
d
Attention (Standalone)
O(Td2)
∼4d2
TTR-Compressor
O(T2d)
2d
Table 4: Complexity comparison of different token compression methods ( T≪d ).
Figure 3: Wall-clock training time comparisons.
Model
S-NIAH-1
S-NIAH-2
S-NIAH-3
MK-NIAH-1
Avg.
4K
8K
16K
32K
4K
8K
16K
32K
4K
8K
16K
32K
4K
8K
16K
32K
TTR-Compressor
0.996
0.996
0.992
0.650
0.982
0.566
0.288
0.126
0.770
0.376
0.188
0.066
0.528
0.302
0.212
0.110
0.5092
w/o normalize k
0.010
0.000
0.002
0.000
0.938
0.254
0.136
0.104
0.764
0.592
0.186
0.092
0.500
0.258
0.098
0.116
0.2531
Table 5: Ablation study on key normalization in TTR-Compressor. We compare the full model with a variant without normalizing k (w/o normalize k ). Best values for each context length are bolded .
Model
S-NIAH-1
S-NIAH-2
S-NIAH-3
MK-NIAH-1
Avg.
4K
8K
16K
32K
4K
8K
16K
32K
4K
8K
16K
32K
4K
8K
16K
32K
λ=0.0001
0.986
0.896
0.826
0.628
0.888
0.348
0.134
0.054
0.654
0.528
0.086
0.108
0.384
0.242
0.118
0.102
0.4364
λ=0.001
0.998
0.996
0.990
0.948
0.998
0.448
0.186
0.112
0.686
0.420
0.186
0.072
0.466
0.274
0.174
0.082
0.5023
λ=0.01
1.000
1.000
1.000
1.000
0.936
0.486
0.148
0.076
0.692
0.374
0.042
0.022
0.484
0.288
0.146
0.068
0.4851
λ=0.1
0.990
0.986
0.982
0.978
0.932
0.312
0.054
0.068
0.728
0.424
0.088
0.078
0.470
0.252
0.112
0.074
0.4705
λ=1
0.872
0.552
0.826
0.794
0.896
0.204
0.062
0.072
0.772
0.496
0.184
0.088
0.526
0.278
0.090
0.068
0.4238
Table 6: Ablation study on the regularization parameter λ in TTR-Compressor. Best values for each context length are bolded , and second-best values are underlined .
Appendix figures & tables3 assets
Supplementary material from the paper’s appendix.
Appendix
Model
S-NIAH-1
S-NIAH-2
S-NIAH-3
MK-NIAH-1
Avg.
4K
8K
16K
32K
4K
8K
16K
32K
4K
8K
16K
32K
4K
8K
16K
32K
Attention-Weighted Pooling
0.978
0.972
0.984
0.990
1.000
0.982
0.860
0.732
0.898
0.856
0.456
0.550
0.494
0.358
0.336
0.314
0.735
TTR-Compressor
1.000
1.000
1.000
1.000
1.000
0.956
0.862
0.750
0.980
0.918
0.740
0.662
0.626
0.442
0.446
0.412
0.7996
Appendix
Table 7: Accuracy on Single Needle-In-A-Haystack (S-NIAH) and Multi-Key Needle-In-A-Haystack (MK-NIAH) tasks from RULER across varied context lengths. Results are based on a 1B-parameter model trained on 100B tokens. Best values for each context length are bolded .
Method
vt
General VQA
Knowledge Reasoning
Document Understanding
Spatial Under.
Avg.
MMBench
MMStar
RWQA
MMMU
AI2D
OCRBench
InfoVQA
ChartQA
MMSI
ERQA
Constant value
Constant
1
75.26
51.63
64.31
45.67
76.07
55.90
37.84
67.20
26.80
35.00
53.56
Linear value
TTR-Compressor
a⊤xt
76.46
52.40
65.62
45.89
75.23
58.70
38.27
67.16
26.80
36.50
54.30
TTR-Compressor (w/ Conv)
a⊤Conv(X)t
76.63
51.16
65.88
47.44
75.16
60.00
39.70
67.16
28.90
36.50
54.85
Appendix
Table 8: Ablation on the parameterization of the value target vt in the TTR objective h⋆=argminh21∥Kh−v∥22+21h⊤Dh . All variants use 3×3 token compression. Gray rows denote the TTR-Compressor configurations used in the main paper. Best values are bolded , and second-best values are underlined .
Figure 4: Prefill and decode latency across context lengths.
Transformer language models process sequences token by token in an autoregressive manner, making growing contexts increasingly expensive. Yet many adjacent token spans are highly predictable or frequently occur as stable units, suggesting that their representations may be compressible. We introduce a method for distilling sequential computation by replacing spans of input tokens with collapsed representations, computed on the fly by a lightweight merge module. This module generates a single surrogate embedding from a sequence of static token embeddings that captures the functional role of the multiple tokens, allowing pretrained models to operate on compressed inputs without architectural changes or re-training. We apply this approach during inference to compress both prompts and intermediate decoding steps, using a rollback mechanism to substitute stored multi-token KV cache entries with their single-step surrogates. Experiments across diverse models show that the merge module can be used to reduce effective sequence length by up to 40% with minimal accuracy degradation across language modeling evaluations and downstream tasks, including question answering, summarization, commonsense reasoning, and long-form mathematical reasoning. Additional lightweight adaptation of the merge module further improves the accuracy-compression trade-off in selected settings. These results demonstrate that sequential token computation in Transformers can be effectively approximated through condensed surrogate representations that approximate the original behavior without model updating.
Zixuan Lan, Jessica Yang, Yanhong Li +2
The University of Chicago · Toyota Technological Institute at Chicago · Independent Researcher +1
Large Language Models (LLMs) incur significant computational and memory costs when processing long prompts, as full self-attention scales quadratically with input length. Token compression aims to address this challenge by reducing the number of tokens representing inputs. However, existing prompt-compression approaches primarily operate in token space and overlook inefficiencies in the latent embedding space. In this paper, we propose K-Token Merging, a latent-space compression framework that merges each contiguous block of K token embeddings into a single embedding via a lightweight encoder. The compressed sequence is processed by a LoRA-adapted LLM, while generation remains in the original vocabulary. Experiments on structural reasoning (Textualized Tree), sentiment classification (Amazon Reviews), and code editing (CommitPackFT) show that K-Token Merging lies on the Pareto frontier of performance vs. compression, achieving up to 75% input length reduction with minimal performance degradation. Code is available at https://github.com/shsjxzh/K-Token-Merging.
Zihao Xu, John Harvill, Ziwei Fan +3
♢Rutgers University · ♡Mistral AI · ♠AWS AI Labs +3
Post-training compression of LLM attention is often formulated as independent matrix approximation, ignoring both the shared structure among attention projections and the representation shift introduced by earlier compression. We propose FTC, a sequential structured compression framework that adapts the approximation to the current compressed model while jointly exploiting the native Q/K/V head structure under a fixed storage budget. The output projection is handled separately to account for the changed post-attention representation. FTC requires neither fine-tuning nor gradient-based recovery. Across seven decoder-only LLMs from 6B to 32B parameters, FTC achieves the lowest WikiText-2 perplexity among the compared methods at every tested keep ratio on five modern GQA models, with the largest gains under aggressive compression. The improvements transfer to downstream tasks and remain substantial at the 32B scale.
Jiangfeng Chen, Xinyu Wang, Tianshuo Yan +4
University of Manitoba · McGill University · Simpleway +2