Rethinking Sequence Compression with Test-time Regression

AI summary

Overview

Research area: Efficient sequence modeling and multimodal large language models (MLLMs), specifically the design of sequence compressors — the operators that collapse a block of tokens into a single summary vector.

Technical level: Intermediate to Advanced. The core idea is stated simply, but the method relies on ridge regression, its dual form, and Cholesky factorization, with gradient derivations in the appendix.

Scope: The paper proposes TTR-Compressor, a drop-in sequence compression operator that replaces pooling or learned projections with the closed-form solution of a learned ridge-regression problem, and evaluates it on block-wise sparse attention (345M-parameter language models) and visual token merging (Qwen3-VL ViT + Qwen3-4B-Instruct).

What This Paper Is About

Sequence compression is a critical information bottleneck in efficient language and multimodal models: block-wise sparse attention compresses token blocks into summary vectors used to score and select relevant blocks, and MLLMs merge neighboring visual features into compact tokens before the language backbone. Despite this, most models use simple compressors — mean pooling, projection over concatenated tokens, or attention-weighted pooling — all of which combine tokens using weights that ignore inter-token relationships, so redundant tokens reinforce each other while distinctive information is diluted. The paper reframes compression as test-time regression: the input tokens define a self-supervised ridge-regression problem whose fitted coefficients are the compressed representation, with the regression problem learned end-to-end by backpropagation while its solution is computed in closed form during the forward pass.

Key Contributions

1. A new compression framework. The authors recast token summarization as an instance-adaptive optimization problem, "test-time regression," in which a token block defines a regression problem and the fitted coefficients serve as the compressed vector (a retrieval key in sparse attention, a visual token in multimodal models). This frames compression as a bi-level optimization: an outer task loss learns which regression problem is useful, and an inner solve determines how it is fitted to each input.

2. A practical operator. They derive an efficient dual closed-form solution requiring only a T×T Cholesky factorization (with T the number of tokens compressed), which avoids iterative gradient-based inner optimization, and they replace the ridge hyperparameter with a learnable diagonal regularizer to eliminate tuning during large-scale pretraining. The compressor adds only 2d learnable parameters.

3. Second-order, input-adaptive behavior. Because the solution's Gram matrix K Kᵀ encodes pairwise token correlations, correlated tokens are fitted jointly, so redundant evidence has diminishing rather than additive influence — behavior the authors argue cannot be reproduced by a fixed projection or additive token combination.

4. Empirical validation in two distinct regimes. As a compressor for block-wise sparse attention, TTR-Compressor improves RULER performance over attention-based compressors by a large margin; as a visual token merger, it consistently outperforms the standard MLP merger and matches its performance with 55.6% fewer visual tokens.

Main Findings

• Retrieval quality is dominated by the compressor. On RULER S-NIAH-1/2/3 and MK-NIAH-1 at 4K–32K context, TTR-Compressor reaches an average accuracy of 0.5092, versus 0.3920 for attention-weighted pooling, 0.3919 for concat & projection, 0.2971 for mean pooling, 0.2632 for full attention, and 0.0635 for sliding-window attention. The convolutional variant reaches 0.5144, which the authors describe as 12.24 points over the strongest baseline.

• Advantages persist under length extrapolation. The gains hold when extrapolating beyond the 8K training length to 16K and 32K, where TTR-Compressor leads the attention-weighted pooling baseline at each reported context length in Table 1.

• Zero-shot reasoning is preserved. On the ten zero-shot benchmarks in Table 2 (ARC-c, ARC-e, BoolQ, COPA, HellaSwag, LAMBADA, OBQA, PIQA, SciQ, WinoGrande), TTR-Compressor averages 48.16, matching full attention's 48.16; the convolutional variant averages 48.12. The paper notes these evaluation contexts are too short to activate sparse attention, so the compressor is not involved at inference and the numbers reflect only reasoning ability acquired during pretraining.

• Visual token merging improves at both compression ratios. Under 2×2 merging, average score rises from 54.93 (MLP) to 55.92 (TTR) to 56.03 (TTR w/ Conv). Under 3×3 merging, from 52.27 (MLP) to 54.30 (TTR) to 54.85 (TTR w/ Conv).

• Gains concentrate on text-rich benchmarks. Under 3×3 merging, the convolutional variant beats the MLP merger by 4.10 on OCRBench, 4.54 on InfoVQA, and 3.56 on ChartQA, and by 2.58 points on average. TTR-Compressor (w/ Conv) with 3×3 merging scores 54.85, close to the 54.93 of the MLP merger with 2×2 merging while using fewer than half as many visual tokens.

• Key normalization is essential. Removing key normalization (k_t = x_t / ‖x_t‖₂) collapses average retrieval accuracy from 0.5092 to 0.2531. The authors explain that raw keys make the static vector a a trivial exact fit (k_tᵀ a ≡ v_t), biasing h* toward an input-invariant projection, and that large-norm outlier tokens then dominate the Gram matrix.

• Learnable anisotropic regularization beats fixed scalars. Fixed λ values give averages of 0.4364 (λ = 10⁻⁴), 0.5023 (λ = 10⁻³), 0.4851 (λ = 10⁻²), 0.4705 (λ = 0.1), and 0.4238 (λ = 1); a learnable scalar gives 0.4971; the learnable diagonal D gives the best overall 0.5092.

• Cost stays low in theory and in practice. Complexity is O(T²d + T³) for the dual form versus O(Td² + d³) for the primal, with 2d parameters versus Td² for a linear projection and roughly 4d² for standalone attention. Wall-clock training on 8×H100 GPUs (sparse attention) shows TTR comparable to other variants except the more limited mean pooling; on 2×8×H100 GPUs (visual merging) it is marginally slower than MLP under 2×2 merging but faster under 3×3.

Methodology in Plain English

The authors treat a block of tokens as data points for a small regression problem. Each token is turned into a "key" by normalizing it to unit length, so that dot products between keys measure semantic direction rather than magnitude. Each token also produces a scalar "target" through a learned projection — optionally after a local convolution that first mixes in neighboring tokens (1D for language, 2D for vision). The compressed vector is the coefficient vector that best predicts those targets from the keys, with a penalty on its size.

The key point is how this differs from pooling. In mean pooling or projection, each token contributes independently. Here, solving the regression requires inverting a matrix built from all pairwise key similarities, so tokens that are near-duplicates get jointly accounted for and stop compounding, while a distinctive token that other tokens fail to explain retains a larger relative share.

Solving the feature-space system would cost d×d work, so the authors use the equivalent dual form, which is a T×T system (T is typically much smaller than the channel dimension d). With learnable diagonal regularization they rescale the keys and reduce the final operator to D^{-1/2} K̃ᵀ (I_T + K̃ K̃ᵀ)^{-1} v, computed via a Cholesky factorization and two triangular solves — no explicit matrix inverse, and no iterative optimization at test time. Only two things are learned: the projection a that produces targets and the per-channel regularization strengths ρ (passed through softplus). Both are trained by backpropagating the downstream task loss through the closed-form solution, so the regression problem itself is shaped to produce task-useful representations rather than merely to minimize its own fitting error.

Why This Matters

Impact on research. The paper argues that sequence compression, despite being ubiquitous, is under-designed — most architectures default to pooling or a projection. It offers a principled replacement that adds second-order information at low parameter cost (2d parameters) and without hyperparameter tuning, and it draws a sharp distinction from prior test-time-training work: because compression operates on small disjoint blocks with no cross-block dependencies, the regression can be solved exactly in closed form, avoiding both iterative inner optimization (as in TTT, ViT³, In-Place TTT, and MesaNet) and single-step approximations like delta-rule updates. It also repurposes test-time regression: instead of producing input-dependent weights or a fast-weight memory to be queried, the fitted coefficients are themselves the representation.

Real-world applications:

• Long-context retrieval and question answering, where a sparse-attention model must decide which distant blocks are worth attending to.
• Document, chart, and infographic understanding, where high-resolution images produce many visual tokens and fine-grained text is easily lost during merging.
• Multi-image and video multimodal assistants, which need stronger compression to handle longer multimodal contexts.
• Agentic workflows with repeated visual observations, which the paper cites as motivation for aggressive merging such as 3×3.

Industry relevance. The operator is a drop-in replacement for existing compressors in two production-style pipelines — block-wise sparse attention and the 2×2 spatial merger used in modern MLLMs — with complexity that avoids O(d²) overhead when T ≪ d (e.g., T ∈ {4, 9, 64} versus d ≥ 1024), making it attractive for pretraining at scale.

Future Directions

• Beyond scalar targets. The paper states that alternative parameterizations of v_t are explored in Appendix E.2, leaving open how far target design (including richer convolutions) can be pushed.
• Compression beyond small disjoint blocks. The exact closed-form solve depends on blocks being independent and small; extending the framework to settings where compression windows interact across blocks is a natural unresolved question.
• Where else can test-time regression serve as a primitive? The conclusion frames the result broadly — compression as an input-dependent inference problem — inviting application to other bottlenecks in efficient foundation models.
• Understanding the learned geometry. The learnable diagonal D means training decides which feature dimensions matter when measuring token similarity; how that learned correlation geometry relates to task structure is not analyzed.

Target Audience

Researchers and engineers working on efficient transformers, long-context language modeling, and multimodal token reduction who are comfortable with attention, ridge regression, and training at scale. It is most directly useful to practitioners choosing a compressor for block-sparse attention or a visual merger, and to researchers interested in test-time training and fast-weight methods who want an exact-solve alternative in the block-compression regime. Readers wanting full derivations will need the paper's appendix, since the provided content includes only the beginning of Appendix A.

Authors’ abstract

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.