PaperlayerView exact revision on arXiv

Source-linked AI summary

Rethinking Sequence Compression with Test-time Regression

Sihan Yang, Jie Huang, Hanxuan Li, Yifan Guo, Weigao Sun, Yiran Zhong, Weiyang Liu

arXiv:2610.31574v1cs.LG

TL;DR

Sequence compression is an important information bottleneck, yet common pooling and projection approaches do not model inter-token relationships. TTR-Compressor makes tokens define a regression problem whose fitted coefficients form the compressed representation, learned end-to-end and solved in closed form. It improves retrieval and visual-token-merging results, matching the MLP merger with substantially fewer visual tokens.

  • Problem

    Common pooling and projection-based compressors do not account for inter-token correlations, despite sequence compression being a critical information bottleneck in efficient language and multimodal models.

    Source: §1 · pp. 1–2

  • Method

    TTR-Compressor constructs a ridge-regression problem from input tokens, learns its parameters end-to-end, and uses the closed-form fitted coefficients as the compressed representation.

    Source: §1 · p. 2, §3.1 · p. 3

  • Results

    TTR-Compressor improves RULER retrieval over attention-based compressors and matches MLP visual-merger performance with 55.6% fewer visual tokens.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The results establish test-time regression as an effective primitive for sequence compression in block-wise sparse attention and multimodal token merging.

    Source: §1 · p. 2, §6 · p. 9

Abstract

from arXiv · show

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.

1 INTRODUCTION

The introduction frames sequence compression as an important but underexplored information bottleneck, then proposes TTR-Compressor and reports gains in retrieval and visual-token merging.

1 INTRODUCTION

  • Existing compressors often use pooling or learned projections despite compression being a critical information bottleneck in sparse attention and multimodal models.Information discarded during compression cannot be recovered downstream.Source: §1 · pp. 1–2
  • The introduction argues that pooling and projection aggregate tokens without accounting for inter-token relationships, allowing redundant evidence to reinforce while distinctive information is diluted.These methods primarily use first-order statistics and ignore second-order inter-token correlations.Source: §1 · p. 2
  • TTR-Compressor makes tokens define a self-supervised regression problem whose closed-form fitted coefficients become the compressed representation, with the regression problem meta-learned for task utility.The introduction describes it as an effective drop-in replacement for sequence compression.Source: §1 · p. 2
  • TTR-Compressor outperforms traditional baselines on RULER retrieval tasks, surpasses attention-based compressors by a considerable margin, and matches MLP merger performance with 55.6% fewer visual tokens.It also consistently outperforms the MLP merger across diverse compression ratios.Source: §1 · p. 2
  • The paper presents TTR-Compressor as a practical operator with an efficient dual closed-form solution and no hyperparameter tuning for large-scale pretraining.The stated contributions also include empirical effectiveness in block-wise sparse attention and visual token merging.Source: §1 · p. 2

2 RELATED WORK

Related work situates TTR-Compressor among sequence-compression operators and test-time regression methods, emphasizing its distinct use of regression to directly produce a token-block representation.

Sequence compression

  • Existing sequence-compression operators include parameter-free pooling, learned projection after concatenation, and attention-based weighted averaging, but do not explicitly account for inter-token correlations.TTR-Compressor instead derives its representation from a linear system whose Gram matrix encodes those correlations.Source: §2 · pp. 2–3

Test-time regression and test-time training

  • Related test-time regression work spans robustness, sequence-modeling layers, adapted MLP projections, and online-regression interpretations of linear-attention descendants.MesaNet solves its inner optimization with conjugate gradients, while LaCT updates fast weights over large unordered chunks.Source: §2 · p. 3
  • Prior test-time regression methods adapt model weights or fast-weight memories, while TTR-Compressor uses regression to produce a token-block representation directly.The fitted coefficients serve as a compressed representation, such as a visual token or retrieval key.Source: §2 · p. 3
  • Unlike iterative inner optimization or single-step approximations in related methods, TTR-Compressor solves an exact closed-form regression for each small, disjoint block.Its blockwise setting permits a single T × T dual solve without cross-block dependencies.Source: §2 · p. 3

3 JOINT SEQUENCE COMPRESSION VIA TEST-TIME REGRESSION

TTR-Compressor formulates sequence compression as an input-specific ridge-regression problem, learns its construction through task loss, and computes the representation using an efficient dual solution.

3 JOINT SEQUENCE COMPRESSION VIA TEST-TIME REGRESSION

  • Given a token sequence, TTR-Compressor constructs a regression problem and uses its fitted coefficients as the compressed representation.The regression is solved for each input during training and inference, while the learned model parameters remain fixed at inference.Source: §3 · p. 3, §3.1 · p. 3

3.1 COMPRESSION AS TEST-TIME REGRESSION

  • Tokens provide regression observations, and the fitted coefficients summarize their key-target relationships while regularization controls coefficient magnitude.Each token is associated with a regression key and scalar target.Source: §3.1 · p. 3

3.2 CONSTRUCTING THE REGRESSION PROBLEM

  • Keys are normalized input features, while targets are learned projections of unnormalized features, preserving token magnitudes and avoiding a sequence-invariant fit.Targets are input-derived and need no additional labels; an optional local convolution can aggregate neighboring context before projection.Source: §3.2 · pp. 3–4
  • A learnable positive diagonal regularizer replaces fixed ridge regularization, and the regression solution incorporates pairwise token interactions.The diagonal regularizer is meta-learned through the task loss.Source: §3.2 · p. 4

3.3 EFFICIENT DUAL-FORM SOLUTION

  • The dual formulation replaces the feature-space system with a T × T token-space system and computes the representation from token similarities.For scalar regularization, the stated cost is O(T^2d + T^3), versus O(Td^2 + d^3) for the primal; with diagonal regularization, keys are rescaled before Cholesky factorization.Source: §3.3 · pp. 4–5

3.4 FULLY END-TO-END TRAINING

  • The target projection and diagonal parameters are trained end-to-end through the differentiable closed-form solution, while each input’s representation is recomputed independently.Each compressor’s learned parameters are shared across input sequences, and no iterative gradient-based test-time optimization is required.Source: §3.4 · p. 5

4 UNDERSTANDING TTR-COMPRESSOR

TTR-Compressor separates learning a task-useful regression objective from solving that regression for each input. Its outer loop shapes the objective through downstream task loss, while its inner solution adapts to input-specific token correlations.

4 UNDERSTANDING TTR-COMPRESSOR

  • The learning process has two mechanisms: outer-loop training learns a useful regression problem, and the inner solution adapts the representation to current input statistics.The outer objective evaluates the compressed representation for the final task rather than minimizing the inner regression error alone.Source: §4 · p. 5, §4.1 · p. 5

4.1 TASK-SUPERVISED LEARNING OF THE COMPRESSION OBJECTIVE

  • The downstream task loss trains the regression targets and coordinate-wise regularization so fitted coefficients yield representations useful to later layers.The scalar targets need no predefined semantics; they act as learned measurements, while the regression supplies a structural inductive bias for compression.Source: §4.1 · p. 5

4.2 INPUT-DEPENDENT SECOND-ORDER ADAPTATION

  • The inner solution jointly weights tokens using regression targets and pairwise correlations, giving highly similar tokens diminishing rather than additive influence.The Gram matrix captures input-dependent token similarities, and learned diagonal regularization further adapts the feature dimensions used to measure similarity.Source: §4.2 · p. 6

5 APPLICATIONS AND EXPERIMENTAL RESULTS

The section evaluates TTR-Compressor in sparse-attention retrieval and visual token merging, then analyzes its efficiency and design choices. It reports stronger retrieval and visual-merging results than conventional compressors, with benefits under higher compression and from learned regularization.

5.1 BLOCK-WISE SPARSE ATTENTION

  • In block-wise sparse attention, the compressor determines which distant information remains accessible, because an uninformative summary can prevent a relevant block from being selected regardless of its token representations.Source: §5.1 · p. 6

5.2 VISUAL TOKEN MERGER

  • TTR-Compressor outperforms sparse-attention compressor baselines on RULER retrieval and exceeds the MLP merger by 2.58 average points under 3 × 3 merging.With 3 × 3 merging, its average of 54.85 is close to the MLP merger’s 54.93 under 2 × 2 merging, using fewer than half as many visual tokens.Source: §5.1 · p. 7, Table 1 · p. 7, §5.2 · p. 8

5.3 EFFICIENCY ANALYSIS

  • When T ≪ d, TTR-Compressor avoids O(d^2) overhead and uses only 2d parameters while retaining data-dependent expressiveness.The stated regime includes T ∈ {4, 9, 64} and d ≥ 1024.Source: §5.3 · p. 8, Table 4 · p. 8

5.4 ABLATION STUDY AND ANALYSIS

  • Removing key normalization collapses retrieval performance because raw keys bias the fitted representation toward a fixed projection and let large-norm tokens dominate correlations.Normalization breaks the trivial alignment with static parameters and prevents large-norm tokens from dominating the Gram matrix.Source: §5.4 · p. 9, Table 5 · p. 9
  • A learnable diagonal regularizer achieves the best reported performance by assigning separate penalties to feature channels.A learnable scalar reaches average accuracy 0.4971; fixed settings perform worse when λ is too small or too large.Source: §5.4 · p. 9, Table 6 · p. 9

6 CONCLUDING REMARKS

The conclusion frames sequence compression as input-adaptive ridge regression, with fitted coefficients serving as compact representations. The appendix supports the practical formulation by deriving dual solutions that reduce the solve to a system sized by the number of tokens.

6 CONCLUDING REMARKS

  • TTR-Compressor uses input-dependent ridge regression to capture inter-token correlations while remaining lightweight and parameter-efficient.The paper reports strong results in block-wise sparse attention and multimodal token reduction, including matching standard visual mergers with significantly fewer tokens.Source: §6 · p. 9

A DERIVATION OF THE DUAL-FORM SOLUTIONS

  • With diagonal regularization, rescaling the keys yields the dual solution, and the stated computation remains O(T^2d + T^3).The appendix notes that forming the rescaled keys costs O(Td) and applying D^-1/2 costs O(d).Source: A DERIVATION OF THE DUAL-FORM SOLUTIONS · p. 15
  • The dual form obtains coefficients by solving a single T × T system, since the fitted representation lies in the row space of K.The appendix proves equivalence between primal and dual ridge regression forms.Source: A DERIVATION OF THE DUAL-FORM SOLUTIONS · p. 15

B REDUNDANCY-AWARE WEIGHTING OF TTR-COMPRESSOR

TTR-Compressor weights tokens according to target information unexplained by the jointly fitted representation, so redundant tokens have diminishing influence. The analysis extends this behavior to groups of correlated tokens and learnable diagonal regularization.

B REDUNDANCY-AWARE WEIGHTING OF TTR-COMPRESSOR

  • A token’s effective weight reflects the part of its target left unexplained by the joint fit, and a fully explained token has no influence on the representation.By contrast, additive pooling always adds the token’s target-key contribution.Source: B REDUNDANCY-AWARE WEIGHTING OF… · p. 16
  • The n-th repeated token adds a coefficient of O(n^-2), so repeated evidence has diminishing returns.For m repeated tokens, their relative influence is bounded by λ + 1, while additive pooling’s corresponding ratio is m.Source: B REDUNDANCY-AWARE WEIGHTING OF… · p. 17
  • As λ approaches zero, each orthogonal group contributes its mean target regardless of group size.The group-saturation result assumes shared unit keys within each group and orthogonality between groups.Source: B REDUNDANCY-AWARE WEIGHTING OF… · pp. 16–17
  • With learnable diagonal regularization, the diminishing-influence mechanism remains, while the learned metric determines which feature directions count as redundant.The extension’s group result applies when groups are orthogonal under that learned metric.Source: B REDUNDANCY-AWARE WEIGHTING OF… · p. 17

C IMPLEMENTATION DETAILS

The implementation uses local depthwise convolutions to construct targets, retains original inputs for regression keys, and solves the regression in FP32. The language and vision variants add convolutional parameters while retaining O(T 2d) complexity for fixed kernel sizes and T ≪ d.

C IMPLEMENTATION DETAILS

  • Regression keys and targets are cast to FP32 for the dual solve, which uses Cholesky factorization without explicitly forming an inverse.Convolution and target projection use BF16 input precision, and the compressed representation is cast back to the input dtype.Source: C IMPLEMENTATION DETAILS · p. 18
  • Only regression targets use convolutional context; language uses an identity-initialized, non-causal depthwise 1D convolution within each block, while vision uses a 3×3 depthwise convolution.The language convolution has kernel size 5 and two-sided zero padding; the vision convolution has no cross-channel mixing before scalar target projection.Source: C IMPLEMENTATION DETAILS · p. 18
  • The language and vision variants require 7d and 11d parameters, respectively, and retain O(T 2d) computation for fixed kernel sizes and T ≪ d.The base compressor has 2d effective trainable parameters; adding a depthwise convolution with s kernel elements adds sd parameters and O(sTd) computation.Source: C IMPLEMENTATION DETAILS · p. 18

D EXPERIMENTAL DETAILS

The experimental details specify separate configurations for block-wise sparse attention and visual token merging. The attention evaluation uses 345M-parameter models, while the visual-merger training uses BF16 and solves TTR regression systems in FP32.

D EXPERIMENTAL DETAILS

  • The block-wise sparse attention evaluation uses 345M-parameter models with 16 Transformer layers, hidden dimension 1,024, and grouped-query attention with 16 query heads and 2 key–value heads.The setup also uses rank-64 query calibration and AdamW with 0.01 weight decay, gradient clipping at 1.0, and a 0.7% learning-rate warmup.Source: D EXPERIMENTAL DETAILS · p. 18
  • Visual token merger training uses BF16, DeepSpeed ZeRO-1, gradient checkpointing, and a maximum sequence length of 16,384 tokens.The two stages use per-device batch sizes of 4 and 1, gradient accumulation of 4 and 8 steps, and 3% warmup followed by cosine decay.Source: D EXPERIMENTAL DETAILS · p. 18
  • For TTR-Compressor in the visual-merger experiments, the regression systems are solved in FP32.The two stages use minimum learning rates of 10^-5 and 0, and weight decay of 0.01 and 0.02, respectively.Source: D EXPERIMENTAL DETAILS · p. 18

E.1 SCALING UP MODEL SIZE AND TRAINING BUDGET

At 1B parameters and 100B training tokens, TTR-Compressor outperforms attention-weighted pooling across almost all tested RULER subtasks and context lengths. Its average accuracy is higher, and the performance gap widens at longer contexts.

E.1 SCALING UP MODEL SIZE AND TRAINING BUDGET

  • The scale-up evaluation compares both methods at 1B parameters and 100B training tokens, using attention-weighted pooling as the strongest baseline from the 345M-parameter experiments.It tests the same RULER subtasks over context lengths ranging from 4K to 32K.Source: §E.1 · p. 18
  • TTR-Compressor outperforms Attention-Weighted Pooling across almost all tested context lengths and subtasks, averaging 0.7996 accuracy versus 0.735; the gap widens with longer contexts despite substantially increased model size and training budget relative to the main paper’s 345M-parameter setting.Source: §E.1 · p. 18
  • The performance gap widens as context length increases, indicating more graceful scaling to longer contexts in this evaluation.The comparison uses the same RULER subtasks across context lengths from 4K to 32K.Source: §E.1 · p. 18

E.2 ABLATION ON THE VALUE PARAMETERIZATION

The latency comparison finds similar prefill and decode latency across compression methods, with all substantially outperforming full attention beyond 16K tokens. In the value-target ablation, a linear readout over local 3 × 3 features achieves the best overall score.

E.2 ABLATION ON THE VALUE PARAMETERIZATION

  • 54.85 is the best overall score, achieved by a linear readout over local 3 × 3 features.The ablation compares value-target parameterizations under the 3 × 3 compression setup.Source: §E.2 · pp. 18–19, Table 8 · p. 19
  • The constant target scores 53.56, the linear target gains +0.74, and plain MLP matches the linear baseline at 54.30.The bilinear head scores 54.55, while its convolutional counterpart scores 54.37; convolution degrades the MLP score to 53.91.Source: §E.2 · p. 19

F INFERENCE LATENCY OF BLOCK-WISE SPARSE ATTENTION

Related papers

Loading 2610.31574v1…