Trellis: Learning to Compress Key-Value Memory in Attention Models

AI summary

Overview

Research area: efficient sequence modeling, long-context Transformers, and bounded-memory attention. Technical level: Advanced. This paper introduces Trellis, a Transformer architecture that replaces the linearly growing Key-Value (KV) cache with a fixed-size memory and learns to compress keys and values dynamically at test time.

What This Paper Is About

Transformers achieve strong sequence modeling performance but suffer from quadratic computational complexity and an ever-growing KV cache during inference. Trellis addresses this by storing new keys and values into a bounded memory with (m) slots using a two-pass recurrent compression mechanism, and by learning how to selectively forget less useful past information. The goal is to retain long-context modeling ability while reducing memory and computation costs.

Key Contributions

1. Trellis replaces the standard KV cache with a fixed-size memory of (m) slots and trains a two-pass recurrent compression mechanism to store new keys and values into that memory.
2. It introduces an online gradient descent update with a forget gate and state decay, enabling the compressed memory to be updated recursively while retaining important contextual information at test time.
3. It derives a hardware-efficient chunkwise parallel form for the nonlinear recurrence, extending the State Space Duality formulation introduced in Mamba2 and enabling efficient use of GPU matmul operations.
4. It reports extensive experiments on language modeling, common-sense reasoning, recall-intensive tasks, needle-in-haystack tasks, and time series forecasting, showing performance gains that increase as sequence length grows.

Main Findings

• Language modeling and common-sense reasoning: In Table 1 at 790M params / 30B tokens, Trellis achieves LMB. perplexity 20.28 and average accuracy 48.38, outperforming Gated-DeltaNet (21.40, 47.52), TTT (27.05, 46.71), Mamba2 (28.91, 46.23), and Transformer++ (25.89, 46.09). The evaluation includes LMB., PIQA, Hella., Wino., ARC-e, ARC-c, and SIQA.
• Context length scaling: Figure 2 shows Trellis attains the lowest perplexity compared to all baselines across tested context lengths. It was evaluated on the Books dataset with context lengths of {512, 1024, 2048, 4096, 8192, 16384, 32748}, on C4 with {2048, 4096, 8192, 16384}, and on the Pile with (T \in {2048, 8192}). Performance gains over other linear RNNs increase as sequence length grows.
• Needle-in-haystack / RULER: Table 3 reports 1B-parameter results on S-NIAH-PK, S-NIAH-N, and S-NIAH-W. Trellis achieves an average of 79.8, outperforming the second-best model, Gated-DeltaNet, at 75.8. This is about +4% on average over the second-best model and about +6% in longer sequences. Trellis is second best on S-NIAH-PK at 4K (95.2) and 8K (97.8), which the paper associates with the observation that simple NIAH tasks with repeated synthetic context require long-term retention that a forget gate can damage.
• Ablations: Table 2 at 125M parameters on the Pile dataset shows Trellis at 10.87, compared with DeltaNet at 11.58 and TTT at 11.44. Removing the forget gate gives 11.28; changing the intermediate activation gives 10.98 for L2-SiLU and 11.29 for Softmax; linear recurrence with Softmax gives 12.71, and linear recurrence with LN-SiLU gives 11.65. Memory size (m=32) gives 11.14, (m=128) gives 10.87, and chunk size (b=1) gives 10.75. The ablation baseline uses 125M parameters with (d=768) and (m=64).
• Additional language modeling results: Table 4 compares 125M/2.4B-token and 350M/7.5B-token models on Pile (2k, 8k) and Books (512 to 32k) against Transformer++, Linear-Attention (LA), DeltaNet, Mamba2, Gated-DeltaNet, and TTT. Trellis achieves the best highlighted results, including Pile (2k) 10.87 and Books (32k) 17.18 at 125M/2.4B tokens, and Pile (2k) 8.26 at 350M/7.5B tokens.
• Scaling: Figure 3 shows Trellis follows a consistent trend and achieves better perplexity than baselines with a fixed FLOPs budget, which the paper describes as Pareto frontier results in the efficiency-effectiveness trade-off.
• Time series: The abstract and experiments overview state that time series forecasting was evaluated and that Trellis outperforms strong baselines, but the provided content does not report specific time-series datasets, metrics, or numerical results.

Methodology in Plain English

Trellis treats memory compression as an online regression problem. For each token, the model uses a compression layer with memory matrix (M_t) to reconstruct a latent target vector (\alpha_t) from the token’s key embedding (k_t), minimizing an (\ell_2) reconstruction loss. The memory is updated with one gradient descent step per token, so (M_t) acts as fast weights that learn in context while the rest of the model’s parameters are trained normally in an outer loop.

To prevent the fixed-size memory from overflowing or overfitting to early tokens, Trellis adds (\ell_2) regularization on the memory states. This produces a state

Authors’ abstract

Transformers, while powerful, suffer from quadratic computational complexity and the ever-growing Key-Value (KV) cache of the attention mechanism. This paper introduces Trellis, a novel Transformer architecture with bounded memory that learns how to compress its key-value memory dynamically at test time. Trellis replaces the standard KV cache with a fixed-size memory and train a two-pass recurrent compression mechanism to store new keys and values into memory. To achieve this, it leverages an online gradient descent procedure with a forget gate, enabling the compressed memory to be updated recursively while learning to retain important contextual information from incoming tokens at test time. Extensive experiments on language modeling, common-sense reasoning, recall-intensive tasks, and time series show that the proposed architecture outperforms strong baselines. Notably, its performance gains increase as the sequence length grows, highlighting its potential for long-context applications.