Rethinking Sequence Compression with Test-time Regression

Paper

Introduction

In models that handle long contexts, the compression operation that reduces multiple tokens to a single vector becomes a performance bottleneck. This paper reframes compression not as a choice of weighted averaging, but as the problem of finding the coefficients of a regression problem defined by the tokens. TTR-Compressor constructs a label-free ridge regression problem from the input tokens and uses its closed-form solution as the compressed representation.

Why existing compressors discard information

Sequence compression is an operation that transforms T d-dimensional token vectors into a single d-dimensional summary vector. In block-sparse attention, this vector becomes the retrieval key that selects the blocks a query reads, while in visual token merging, this vector itself becomes the visual token seen by the language model. In both cases, the compressor stores the tokens as a single vector, and subsequent layers read only that vector. Information discarded here cannot be recovered later.

Existing approaches fall into three categories. (1) Mean pooling assigns the same weight to every token. (2) Concatenating tokens and passing them through a learned linear projection uses fixed weights per position. (3) Attention-weighted pooling scores each token individually but evaluates each token independently of the others. What all three share is that their weights do not reflect relationships among tokens. They use only first-order statistics—how important each token is individually—and ignore second-order statistics, the correlations between tokens. Consequently, duplicate tokens reinforce one another, and rare features are diluted.

A small example illustrates this difference. The example below is for explanation and is not an actual measured result.

Illustrative example. Suppose four tokens are given, three of which carry essentially the same content A, and only one carries distinct information B. Set all target values to v=(1,1,1,1). Additive pooling (mean) assigns a weight of 1/4 to each token, yielding an influence of 3/4 for the A group and 1/4 for B, a ratio of 3:1. TTR groups token similarities into a Gram matrix and fits them jointly, so the three correlated A tokens are treated as duplicates; the group influence is still 3/4, but B becomes 1/2. The ratio becomes 3:2, weakening the dominance of the duplicate group and increasing the relative weight of the distinct token B. In general, according to Appendix B (Proposition 4, Corollary 1), the relative influence ratio between m fully overlapping tokens and a single distinct token is not m but m(λ+1)/(λ+m) < m, so duplicated evidence has diminishing returns. This number comes from the example in Figure 2(b) of the paper (normalized D=I, λ=1) and its general formula; it is not a measured value.

What changes when tokens become a regression problem

Given tokens X=[x₁,…,x_T], TTR-Compressor defines a ridge regression problem on the tokens themselves and uses the fitted coefficients h* as the compressed representation (§3.1, Eq. 1). Ridge regression fits target values while penalizing the magnitude of coefficients. The key is to solve, for each input, “which coefficients best explain these tokens” and use those coefficients as the summary.

Specifically, it determines two things. One is each token’s key k_t = x_t / ‖x_t‖, retaining only the vector direction so that inner products between keys measure only directional similarity. The other is a scalar target v_t = aᵀx_t, which extracts a single number from each token using one shared learnable vector a (§3.2). This target requires no labels and is derived from the input.

The reason key normalization matters is clear. Since targets are generated linearly as v_t = aᵀx_t, without key normalization (k_t = x_t), the fixed vector a itself becomes the correct answer, and h* converges to a fixed vector independent of the input. In addition, outlier tokens with large norms dominate the Gram matrix, so the model sees magnitude instead of similarity. Normalization removes both problems at once. Indeed, removing key normalization drops the RULER average from 0.509 to 0.253 (§5.4, Table 5).

Solving the regression problem yields the following compressed representation.

Here K is the key matrix, v is the target vector, and D is the normalization matrix. The result is determined jointly by the statistics of keys and targets, Kᵀv, and by KᵀK, which captures the similarity structure among tokens. KᵀK is the Gram matrix in feature space, and because it groups tokens pointing in the same direction, it produces the property from the earlier example: “duplicates have reduced influence.”

Conventional ridge regression requires tuning the regularization strength λ as a hyperparameter, but tuning it during large-scale pretraining is expensive and cumbersome when model size changes. Instead of a scalar λ, TTR uses a learnable diagonal matrix D = diag(softplus(ρ)) (§3.2, Eq. 3). It learns a different regularization strength for each channel, penalizing less important feature directions more strongly and preserving directions important to the task. The learned diagonal outperforms the best fixed λ (λ=0.001, 0.502) with 0.509 (Table 6), and a learned scalar is also competitive at 0.497, showing that the task loss can find the regularization strength on its own.

Solving small blocks in dual form

Working directly in d dimensions requires solving a d×d linear system, but the number of tokens to compress, T, is typically much smaller than the channel dimension d (e.g., T∈{4, 9, 64} versus d≥1024). In this case, we can switch to the dual form and solve only a T×T system (§3.3, Eq. 7). The dual form is an equivalent transformation that turns a d×d system into a T×T system.

The computation proceeds as follows. Rescale the keys by D to form K̃, factor the T×T Gram matrix I_T + K̃K̃ᵀ via Cholesky decomposition, solve for α, and then apply K̃ᵀα to obtain the compressed representation. Cholesky decomposition factors a symmetric positive-definite matrix into a product involving a lower-triangular matrix, solving the linear system with two substitutions. The complexity is O(T²d + T³), much lighter than the O(Td² + d³) of solving the d×d system directly when T≪d. The actual computation never forms an explicit inverse and uses only triangular solves (§5.3, Table 4; Appendix C).

This structure makes training bi-level (§3.4, Figure 2(a); §4.1). The outer loop learns parameters a and ρ that determine “which regression problem is useful” using the task loss, while the inner loop solves “how to fit” that regression problem in closed form for each input. The only learnable parameters are the target projection a (d values) and the regularization ρ (d values) per compressor, and because the inner solution is a differentiable closed form, gradients flow through it directly. At inference time, the learned parameters are fixed, and only h* is recomputed for each input. This differs from existing test-time training approaches that run gradient descent for every input or approximate it with a single delta-rule step (§3): because computation is independent per block, a single T×T dual solve yields the exact solution.

  flowchart LR
  X["Input tokens: x1,x2,x3≈A, x4=B"] --> K["Normalized keys k_i=x_i/||x_i||"]
  X --> V["Targets v_i=a^T x_i, e.g., (1,1,1,1)"]
  K --> G["Gram matrix G=K K^T: 1 among A, 0 with B"]
  V --> S["α=(G+I)^-1 v = (1/4,1/4,1/4,1/2)"]
  G --> S
  S --> H["h*=K^T α: A influence 0.75, B influence 0.50"]

The matrix coordinate (i,j) is the directional similarity between tokens i and j, where a value of 1 means the same direction and 0 means orthogonal. The fourth value of α grows to 1/2 because B captures the residual not explained by the A group.

Results: block-sparse attention

The first application is block-sparse attention based on HiLS-Attention (§5.1). The context is divided into 64-token blocks, and for each block’s compressed representation, up to 32 relevant blocks are selected and attended to together with a 512-token local window. The compressor stores each block as a retrieval key vector, and the query uses those keys to select blocks to read. Retrieval accuracy therefore directly reflects compression quality. We trained a 345M model on FineWeb with an 8K context for about 30B tokens and evaluated it on RULER’s needle-retrieval tasks (S-NIAH, MK-NIAH).

Before the evaluation, consider one question. With the same attention architecture, how much more hidden information can be found in contexts extending to 32K if only the compressor changes? The table below extracts only the average across compressors and representative long-context numbers (Table 1).

The interpretation has two parts. First, the choice of compressor itself strongly affects retrieval quality. The gap between mean pooling and attention-weighted pooling reaches about 9.5 points on average. Second, TTR substantially outperforms all competing compressors on average. At 16K extrapolation beyond the 8K training length, the advantage extends to individual tasks, but under 32K extrapolation, standard TTR loses S-NIAH-1 (0.650) to attention-weighted pooling (0.740), and on the 32K four-task average, TTR’s 0.238 versus attention-weighted pooling’s 0.224 narrows the lead. The convolution variant largely mitigates this, recording 0.988 on 32K S-NIAH-1. The paper attributes this to convolution making the target v_t more informative and thereby improving compression quality; the convolution variant leads the strongest baseline by 12.24 points (§5.1).

One caveat is that the gains are confirmed on “finding” tasks, and whether they extend to general language quality is a separate question. In the commonsense reasoning benchmark in Table 2, evaluation contexts are short, so sparse attention is not activated; all compressors are similar in the 48-point range. This is discussed in the limitations section below.

Results: visual token merging and efficiency

The second application is visual token merging in multimodal models (§5.2). High-resolution images produce many visual tokens, making it expensive to feed them directly into the language backbone. The standard design concatenates adjacent 2×2 features and reduces them to one token through a two-layer MLP. Here, TTR fits one regression per spatial group (2×2 or 3×3) to compress 4 or 9 features into a single vector, and a two-layer MLP projects it to the language model dimension. We trained with Qwen3-VL’s ViT encoder and Qwen3-4B-Instruct using a two-stage procedure that trains the merger and backbone sequentially.

The comparison question before the table is this. With the same encoder and language model, how much less information is lost on tasks where fine-grained details such as documents and charts matter when the merger changes?

TTR outperforms the MLP at both ratios, and the gain grows with the stronger 3×3 compression (for 3×3, the convolution variant is +2.58 on average versus the MLP). The difference is most pronounced on text-heavy benchmarks where fine details are easily lost during merging: at 3×3, the convolution variant records OCRBench +4.10 (55.90→60.00), InfoVQA +4.54 (35.16→39.70), and ChartQA +3.56 (63.60→67.16) versus the MLP (Table 3). The key point is that TTR (w/ Conv) with 3×3 merging averages 54.85, almost the same as the 2×2 MLP’s 54.93. In other words, it reduces the token count by 55.6% (N/4 → N/9, less than half in terms of token count) while maintaining performance at the level of the standard 2×2 MLP. The paper explains this with a second correlation-structure interpretation: when a merger mixes features with fixed weights, redundant features such as background patches dominate the token and erase details, whereas TTR fits correlated features jointly and reduces the influence of redundancy.

Efficiency must be examined at two levels. The operator itself costs O(T²d) for TTR, heavier than mean pooling O(Td, 0 parameters) but lighter than concat+projection (O(Td²), Td² parameters) or independent attention compressors (O(Td²), ~4d² parameters), and its 2d learnable parameters are relatively few (Table 4). The convolution variant adds kernel parameters, bringing the count to 7d for language and 11d for vision (Appendix C). At the system level, wall-clock time in Figure 3 shows that in block-sparse attention training (8×H100), TTR takes 22.50 hours, slower than mean pooling’s 19.33 hours but similar to other learned compressors (23.12 and 21.57 hours), while in visual merging (2×8×H100) it is slightly slower than the MLP at 2×2 but actually faster at 3×3. Inference latency is similar across compressors and substantially lower than full attention after 16K (Appendix F, Figure 4). In other words, most of TTR’s cost falls on training time, and at inference it is effectively the same as other compressors.

Limitations and when to apply

We separate limitations directly exposed by the authors’ data from those the reviewer must infer.

Constraints directly confirmed in the authors’ experiments. First, training cost is not free. Training is slower than the lightest mean pooling, and at 2×2 visual merging it is even slightly slower than the MLP (Figure 3). Second, under 32K extrapolation, more than four times the training length, standard TTR drops sharply on some subtasks (S-NIAH-1 32K = 0.650). The convolution variant largely mitigates this (0.988), but it adds parameters and convolution operations. Third, the design is sensitive to choices, to the point that removing key normalization collapses performance (0.509→0.253; Table 5).

Unverified points from the reviewer’s perspective. First, it has not been confirmed that improved retrieval accuracy translates into better general language quality. In the commonsense reasoning benchmark in Table 2, contexts are short, so sparse attention is not activated and the compressor does not affect performance; all compressors are similar in the 48-point range (though Concat & Projection has far more parameters and may have benefited in pretraining itself). In other words, the gains are currently limited to tasks of a “finding/retrieval” nature. Second, the validation scale is small. The advantage was shown at 345M, and scaling to a 1B model with 100B tokens maintains the lead at an average of 0.7996 versus 0.735, with the gap widening as context grows (Appendix E.1, Table 7). However, as HiLS-Attention notes, large models can compensate for retrieval errors with greater information bandwidth, leaving room for compressor differences to shrink. Third, the applications are limited to block-sparse attention and visual merging, so generalization to other compression points is unverified, and TTR is efficient only when T≪d; as blocks grow, the O(T²d + T³) cost increases rapidly.

Given these characteristics, TTR-Compressor is well suited to points where small blocks (tens of tokens or fewer) are reduced to vectors with large channel dimensions. Examples include retrieval-key summaries for block-sparse attention and visual token merging that preserves fine details such as those in documents and charts. Conversely, the gains are not yet proven for retrieval in large models where attention-weighted pooling is already sufficient, compression with large block sizes, or situations where inference latency itself must be reduced.