TRELLIS: Learning to Compress Key-Value Memory in Attention Models
Authors: Mahdi Karami
Ali Behrouz
Praneeth Kacham
Vahab Mirrokni
Venue: COLM 2025 (2025)
Keywords: Sequence Models
Language models
Recurrent Neural Nets
Test Time Training
TL;DR: This paper introduces a novel approach to efficiently compress the K-V cache into a fixed number of slots

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 increases, highlighting its potential for long-context applications.

Published as a conference paper at COLM 2025
TRELLIS:
Learning to Compress Key-Value Memory in Attention Models
Mahdi Karami, Ali Behrouz, Praneeth Kacham, Vahab Mirrokni Google Research
(mahdika, al alibeerou,, pkaceeaa pkacham, mi mirrokni Jēgoogle. com
Abstract
Transformers, while powerful, suffer from quadratic computational com- plexity and the ever-growing Key-Value (KV) cache of the attention mech- anism. 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 two-pass recurrent compression mechanism to store new keys and values into memory. To achieve this, it leverages an on- line gradient descent procedure with forget gate, enabling the compressed memory to be updated recursively while learning to retain important con- textual information from incoming tokens test time. Extensive experi- ments on language modeling, common-senun reasoning, recall-intessive 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.
1 Introduction
Transformers (Vaswani et al., 2017) has established itself as the de facto architecture for sequence modeling in modern deep learning, achieving significant advances across di- verse areas, including language modeling (Devlin et al., 2018; Radford et al., 2018), com- puter vision (Arnab et al., 2021; Dosovitskiy et al., 2020), and graph learning and genera- tion (Dwivedi & Bresson, 2020; Karami, 2024; Yun et al., 2019). Their success stems from the attention mechanism, which allows models to dynamically attend to relevant parts of an input sequence while enabling parallel computation. This enables the capture of long-range dependencies and in-context learning. However, the quadratic time and space complexity of attention with respect to sequence length restricts its scalability long sequence mod- eling. Furthermore, the requirement for an unbounded cache leads to inefficient memory management, particularly in ansource-constrainet environments.
These limitations have driven the exploration of alternative architectures that aim to retain the representational power of Transformers while addressing their computational and memory complexity. One key strategy involves sparsifying the dense attention matrix through various techniques, including: blockwise attention (Parmar et al., 2018; Qiu et al., 2019); using strided or sliding window attention patterns (Beltagy et al., 2020; Child al., 2019; Zaheer et al., 2020); or cuustering/sorting tokens (Kitaev et al., 2020; Roy et al., 2020; Tay et al., 2020). Another approach involves lowo-rank approximations of the self-attention matrix, leveraging the insight that it often exhibits low-rank properties (Wang et al., 2020). different paradigm employs the kernel trick, replacing the Softmax operation with a dot product of feature maps, resulting in family of linear attentions (Choromanski et al., 2020; Katharopoulos et al., 2020; Peng et al., 2021b). While these methods substantially reduce computational costs, they may sacrifice expressiveness and performance, often requiring hybrid approaches that combine them with dense attention layers (Fu et al., 2023; Mehta et al., 2022). Additionally, global convolutions (Li et al., 2022; Poli et al., 2023; Romero et al., 2021) and their alternative sequence modeling techniques.
variants (Karami & Ghodsi, 2024) have been explored as

---

Published as a conference paper at COLM 2025
Recently, Recurrent Neural Network (RNN) architectures have re-emerged as promising attention-free solutions for sequence modeling. These models leverage the parallelization capabilities of their linear recurrence, building upon earlier linear time-invariant models (Fu et al., 2023; Mehta et al., 2022; Wang et al., 2022) and extending them to more expressive input-dependent gated RNN designs with linear memory (De al., 2024; Gu & Dao, 2023; Orvieto et al., 2023; Yang et al., 2024b) that demonstrate improved in-context learning while retaining computational advantages. However, their effective memory is limited, affecting their ability to efficiently compress and summarize information over very long contexts within their fixed-size hidden states.
Trans- Our approach and main contributions In this paper, we present Trellis, former architecture with bounded memory that learns how to compress its key-value memory at test time. To address the linearly-rrowing memory of global attention, Trellis replaces the KV-cache (i.e., the pair of K and V matrices) with a bounded memory with slots and train a meta in-context model to learn how to store new keys and values into memory. For better management of memory, Trellis uses forgetting mechanism that learns how to selectively forget unnecessary past information in its compressed KV-cache.
We perform an extensive set of experimental evaluations on various tasks including lan- guage modeling, commonsense reasoning, recall-intensive, series data. We observe that Trellis outperforms state-of-the-art baselines, including Trans- former++, and modern linear recurrent neural networks in downstream tasks. Furthermore, Trellis scales better than baselines when increasing the context length, showing promising results for long-context tasks.
-nayssik- and time
2 Related Work and Background
For an input sequence = (,, XT), where x; E I2 the causal Softmax attention mecha- nism generates output tokens y, E I2, , by attending to past tokens (historic context):
y1 V: Softmax(K:, 9:)
Here, the query, key, and value vectors are computed by linear projections of the input: qi = W Xt, k; = W₂ Xt, Wo xt, where Wq. W₁, Wo E Rdxd are learnable weight matrices. The key-value memory, represented by the caches Kr E Rdxt and Vi Rdxt, stacks the key and value vectors of each new token, leading to unbounded caches with linearly-growing size. The retrieval of relevant information from this key-value cache can be rewritten as weighted sum:
y/ = Vi at, where a; = Softmax(K,) q:) E R'.
Here, the vector a; E IR' is the collection of the attention scores between 1-th token and its historic context. Hence, the attention in equation (1) can be seen non-linear query from an unbounded memory.
The key-value cache size growth poses a significant memory bottleneck during inference, especially for long sequences. Additionally, each retrieval operation scales linearly with sequence length, resulting in an overall quadratic computational complexity ( for generating full sequence of length
To address the computational and memory bottleneck of Softmax attention, various alterna- tives have been proposed (Tay al., 2022). well-established approach involves employing the kernel trick to replace the Softmax with a dot product of feature maps, ((q;), ((k;) commonly known as Linear Attention LA (Katharopoulos et al., 2020).
To maintain bounded computational and memory requirements in attention mechanisms, an alternative strategy is to explicitly use a fixed-size key-value cache. In this approach, the memory matrices K and V are constrained to length m where m < T. A straightforward implementation of this strategy involves limiting the attention window to the most recent m tokens by maintaining frst-in-first-oud (FIFO) queue, often referred to Sliding Window
2

---

Published as a conference paper at COLM 2025
Attention (SWA). While SWA achieves linear computational complexity, it suffers from limited receptive field, restricting the model's ability to capture long-range dependencies and leading to poor recall-memory trade-off (Arora et al., 2024). Moreover, the observation, supported by many research works that key-value matrices in attention mechanisms often exhibit low-rank and sparse structures (Chen et al., 2021; Singhania et al., 2024; Wang et al., 2020) motivates the design of efficient sequence mixing layers that induce these properties. Therefore, these layers aim to compress the context, storing the important information while discarding or forgetting redundancies, rather than naively truncating memory.
In light of this insight, the nism (Peng et al., 2021a) introduces a method to compress and dynamically update fixed-size memory. This is achieved using pair of linear recurrences to update the key and value matrices:
Control (ABC) mecha-
K K + a; ky E Rmxd
V, = Vi- + a; of E Rmxd,
yt vi Softmax ( R
Here, & Softmax(Wax)) E (0, 1)"h controls the update distribution across the memory slots. Each component, 4t,j, can be interpreted as the writing intensity determining the 1-th token contribution to the j-th memory slot.
The ABC update rule can be decomposed into cascade of two LAs and presented in the following two-pass process:
(0))_1= = kA 9: E R""
LA()(())())))) D1) yt R².
where the intermediate activation function: f(-) Softmax(-). Building upon this founda-
tion, Gated Slot Attention (GSA) (Zhang et al., 2024) introduced a gated version of ABC, enhancing the two-pass process with forget-gate mechanism introduced in GLA (Yang et al., 2024b). This mechanism allows the model to forget irrelevant information, resulting in better memory management and improved performance. Importantly, unlike many linear attention variants, both ABC and GSA do not rely on kernel approximations and retain the Softmax non-linearity. Moreover, they can take advantage of the chunkwise matrix form (Hua et al., 2022; Kacham et al., 2024; Yang et al., 2024b), enabling parallel hardware-efficient implementations on tensor cores.
The simple linear additive (Hebbian-like) nature of the recurrence in ABC and GSA, however, limits the memory capacity and effective memory management in long-context retrieval tasks (Behrouzed et al., 2025b; Schlag et al., 2021). Specifically, the additive Hebbian update rule in ABC and GSA endlessly accumulates new tokens into the fixed memory space without an explicit mechanism to replace past value with new information in the memory, leading to an overcapacity regime (Schlag et al., 2021). Moreover, their tate-inndependent linear outer product modification of memory lacks dynamic interaction between the memory content and incoming keys, disabling it to selectively discard irrelevant or redundant information.
To address these limitations, this paper introduces new two-pass non-linear recurrence. Leveraging techniques from meta-learnng and test-time memorization (Behrouz et al., 2024; Sun et al., 2024a), approach is designed to dynamically compress new keys and values into the memory while minimizing information loss. For clarity, the subsequent explanation and its notations focus on key (ki) cache compression in the first pass, noting that the same principles apply to the value cache in the second pass.
3

---

Published as a conference paper at COLM 2025
3 Method
We define the compression model as regression layer that projects a key token into a latent space. Specifically, for each token embedding k; and its corresponding latent representation Rt, the compression layer aims to reconstruct target vector such that &; R à; = (M,,k))
To minimize the reconstruction error and the compression loss, we formulate the learning objective as an l2 optimization problem:
f || (M,k;) a111,, M, E R K, R², a, E IR"
(4)
We model the latent representation a using an encoder network, which is implemented as a linear projection of the input: a; Wax; with projection weight matrix Wa E R Our approach follows the framework of Fast Weight Programmers (FWPs) (Schlag et al., 2021; Schmidhuber, 1992), where the internal memory of the compression layer (a.k.a. state in the context of RNNs), M,, serves "fast weights", which are dynamically updated based streaming input data. Hence, each sequence serves as a training dataset for the learning in this inner loop. To efficiently adapt to new tokens, we design an internal learning procedure that continuously updates M,, allowing it to store in-context information.
In this framework, the outer network, also referred to the "slow" network, consists of the projection layer weights and the rest of the model parameters, jointly denoted as W. These parameters are trained in the outer loop, which follows standard deep neural network optimization process, minimizing the end-to-end loss averaged over the training dataset to learn generalizable patterns from the training set. This learning process subsequently enables fast adaptation within the inner loop (Schlag et al., 2021). Importantly, the weights of the "slow" network remain frozen during the internal state updates of the inner loop (involving the "fast" weights). This overall procedure constitutes bi-level optimization strategy (Chen et al., 2022; Liu et al., 2022) commonly used in meta-learning (also referred to learning to learn) (Andrychowicz et al., 2016; Bengio et al., 1990; Schmidhuber, 1992; Thrun & Pratt, 1998).
Given the sequential nature of the data, we approach this problem online optimization problem and update the internal memory using one gradient descent step per token:
M,+1 = M2 y M.VMC(M; D,, 2;)
(5)
This update rule generates the sequence of states (M.). where each state M, is nonlinear recurrent function of the previous state and the input token. Consequently, the memory update follows causal nonlinear recurrence, ensuring that information is continually integrated into the memory.
By the chain rule, the gradient of the loss with respect to M can be computed:
(6) where Jø is the Jacobian of (()). In practice, terms involving Jacobian products (such as " == Je (M: k,) a:) in the expression above) are computed efficiently using the m
VMC: = G, v,) = 2 (Jo (M k,) -a:)) k.T
vector-Jacobian product (vjp) method available in modern machine learning frameworks.
This avoids explicitly forming the full Jacobian matrix and leverages efficient automatic differentiation.
3.1 State Decay
While the proposed compression layer can address the quadratic time and space complex- ity of Transformers by learning how to effectively compress key-value pairs (k,, Dt) into fixed-size memory states, these memories can still overfit to early tokens of the sequence or overflow as information accumulates. To mitigate these issues, we introduce & regu- larization on the memory states analogous weight decay in standard neural network aamaing-applied within the inner loop:
L: = =1 v.) dell2 + A M
2
4

---

Published as a conference paper at COLM 2025
This regularized objective results in gradient descent recurrent update with state decay:
M, (1- M, - y (M,
-,M Y: k.
[0, 1]
(8)
where, u, := Jo ( k:) - a1) This has been recently explored in the context of test- time memorization (Behrouz et al., 2025b; Karami & Mirrokni, 2025; Wang et al., 2025). In this setting, the scalar B; = 1 A, E [0, 1 acts as forget gate, controlling the retention of prior memory. When P. + it selectively updates the memory based the interaction of state and input token without fading its magnitude, while Br + 0 erases the memory (possibly due to the change of context). Such scalar gating mechanisms have gained renewed attention in recent RNN architectures as they provide a lightweight yet effective memory update (Beck et al., 2024; Behrouz et al., 2024; Peng et al., 2021b; Sun et al., 2024b; Yang et al., 2024a).
3.2 Parallel and Hardware Efficient Implementation
The non-linear nature of the update rule (8) typically hinders straightforward parallelization. Several techniques have been proposed address this limitation (Gonzalez et al., 2024; Lim et al., 2023). Sun et al. (2024a) introduces mini-batch gradient descent, where the sequence divided into chunks, and the state at the beginning of each chunk is used compute the gradients for all time steps within that chunk. In this approach, the gradients within each chunk are approximated as: VMC, (My, &t, Dt), where My represents the state at the beginning of the chunk (i.e., the final state from the preceding chunk and I' = =- mod(t, C) with C denoting the chunk size.
This strategy enables the parallel computation of a mini-batch of stale gradients at the start of each chunk, thereby significantly enhancing scalability. Using this approximation effectively linearizes the general non-linear recurrence (equation 8) within each chunk, leading to the following recurrent update rule for the internal state and the memory readout of the compression layer:
(y:)T= = compress((q;:k) (t) M M, 1 1-2t (J+ (p k.) - -a)) k.
(9)
y/ = Miq:
This locally linear recurrence satisfies the associative property and can therefore be paral- lelized by parallel scan (a.k.a. prefix sum) (Blelloch, 1990), formulated into parallel chunkwise form which has been shown to efficiently utilize the matmul units of the modern GPUs and can be more I/O efficient (Hua et al., 2022; Kacham et al., 2024; Yang et al., 2024b). We adopt this approach and derive a hardware-optimizzed chuck-wise form for the recurrence update with state decay in equation 7. For b-th block, covering time steps where b + < (b + 1)C, for r-th step in the block, let's define the local cumulative product of decay factors as n IT; IC B. We also denote the segmented cumulative product over
the sub-block steps to (1 < < sisj S C) as f = , and the lower triangular matrix with entries [N]]; = / v i S j (and 0 otherwise). By unrolling the recurrence relation locally within the b-th block we obtain:
M
8 M
7 - k
Ent
(10)
Where 6 is the state at the end of the (b- 1) -h-thlkooc.. Then Mà can be expressed in matrix form as:
M u M
U Diag(y Qg. kb
(11)
where, O the wlement-wise product. Also, unrolling the readout step and formulating in matrix form for the entire block yields:
Yd Diag(H") M²Q -U" o C n)
(12)
5

---

Published as a conference paper at COLM 2025
Output
Lineer
Unner
Sw/OLU
Norm
Trellis
Nx
Trellis
Convil
ConvID
Linear
Linesr
Linear
Linezr
Input
Trellis Block
(Left) Block diagram of the language model. (Right) The Trellis block. Each sequence mixing block composed of a short Conv1D for kq and the Trellis is followed by post normalization and the GeLU post-gate.
Figure 1:
where n² = Ip"................ ........ p and K2 Q E RdxC and Uo Y E RT*C are matrices collecting the corresponding vectors within the block. This matrix form, which computes the states only at the end of each block, extends the State Space Duality formulation introduced in Mamba2 (Dao & Gu, 2024) and enables efficient use of the matmul operations on GPUs.
Consequently, applying the proposed recurrent compression model to key and value caches in two-pass process yields the following operations:
( = (f, (9, E R" Ty:)T1
(13)
y/ E
(14)
We explored alternative choices for the intermediate activation function in our architecture
and found that normalized SiLUI defined as f(x) [S]LUX)]]] outperforms the commonly used
SiL.U(x)
Softmax, also known the normalized exponential function. This improvement can be explained by the fact that, in our mrchitecture-whene where the cache is densely compressed into limited number of memory slots (m < T normalization function with less spikiness than Softmax is more effective for retrieval. It also worth noting the difference in input ordering between the second pass of our architecture and that of ABC and GSA. In Trellis, a; acts
shared target vector for both compression layers. Consequently, we use yt = ((M, 9:) as the readout operation of the value compression layer to ensure correct output dimensions. As a result, in the first pass (equation 13), Trellis compresses the new key embedding, k,, into its memory, generating an intermediate representation and in the second pass (equation 14), compress the new value embedding and finally outputs yi-
The overall Trellis architecture used for language models is illustrated in Figure 1.
3.3 Related Works and Discussion
Linear State Space Models (SSMs) and Recurrent Neural Networks (RNNs) have recently received renewed interest as an efficient paradigm for sequence modeling. Offering sub- quadratic scaling during training and constant-time recurrence at inference, they have proven particularly effective for modeling long-range dependencies (Gu et al., 2020). The recurrence in the linear time invariant SSMs can be reformulated as global convolution, enabling efficient computational implementations (Gu et al., 2021; Mehta et al., 2022). Furthermore, the linearity in input-dependent SSMs (Dao & Gu, 2024; Gu & Dao, 2023) and

---

Published as a conference paper at COLM 2025
modern gated RNNs (Beck et al., 2024; De et al., 2024; Orvieto et al., 2023) take advantage of parallelization through techniques like associative scan (Blelloch, 1990; De et al., 2024; Smith et al., 2023), chunkwise parallel forms (Behrouz et al., 2024; 2025a; Hua et al., 2022; Yang et al., 2024b;c). However, a potential limitation arises from state-independen The additive (gated) linear modification of the memory, doesn't consider the interaction between the memory content and incoming keys, potentially limiting the model's ability to efficiently compress and summarize information within its finite memory state.
updates.
Fast Weight Programmers (FWPs) represent a class of architectures where the parameters of one network, often termed the "fast" network, are dynamically generated or modified by another "slow" network (Schlag et al., 2021; Schmidhuber, 1992). This concept, also referred to as input-dependent parameterization (Gu Dao, 2023; Karami & Ghodsi, 2024; Karami et al., 2019), enables the model to adapt to the specific characteristics of its input, potentially capturing more complex contextual patterns. In our framework, the compression layer's memory state acts as "fast weights" that quickly learn in-context information, while the "slow" network parameters learns generalizable patterns across the training set. Therefore, we effectively adopt meta-learning (Andrychowicz et al., 2016; Bengio et al., 1990; Schmid- huber, 1992, Thrun & Pratt, 1998) style approach by deriving an internal online learning procedure that continuously minimizes a reconstruction loss.
Online Gradient Descent is widely adopted for non-stationary sequential data in classical adaptive filtering such as the Least Mean Squares (LMS) algorithm (Haykin, 2002). LMS-also known as the Delta Rule (Schlag et al., 2021; Widrow & Hoff, 1988)-typically minimizes the instantaneous squared error through a simple a linear update rule. Convex optimization- based methods have also been explored for identifying SSMs ka. linear dynamical systems); for example, Karami al. (2017) formulates the problem as a multi-view matrix factorization and proposes a global optimizer.
In contrast, our work offers a more general non-linear update rule incorporating a forgetting gate. This update rule is applied to both key and value compression within a cascading, two-pass structure. This design is particularly suitable for long-context modeling, where efficient and expressive memory compression is critical.
4 Experiments
Overview: In this section, following recent studies (Lim et al., 2023; Zhang et al., 2024), we evaluate the performance of Trellis in various downstream tasks, including language modeling, common-sense reasoning, needle in haystack, and time series forecasting. We use four scales of the proposed Trellis with (1) 125M, (2) 350M, (3) 780M, and (4) 1B parameters. For each scale, we train three different versions, each of which is solely trained on either Pile (Gao et al., 2020), C4 (Raffel et al., 2020), or Books datasets (a subset of the Pile). The main reason for our choice of datasets and training different version of Trellis are two folds: First, the dataset used for training and its characteristics can significantly affect the performance of the model. Therefore, while a dataset can be more useful for a model, might not be the best choice for others. Accordingly, to avoid cherry-picking of the dataset, and to emphasize the generalizebility and the power of Trellis, we use three different versions, each of which is trained on commonly-used dataset in the literature. Second, we use each dataset to highlight the strength of Trellis in one aspects. That is, we use Pile (resp. Books) dataset to highlight the strength of Trellis in short context (resp. long context).
Baseline Models: We compare our method against the Transformer++ architec- ture proposed in Touvron et al. (2023) and the following sub-quadratic models: Linear-Attention (LA) (Katharopoulos et al., 2020), TTT (Sun et al., 2024a), GSA (Zhang et al., 2024), DeltaNet (Yang et al., 2024c), Gated-DeltaNet (Yang et al., 2024a), Mamba2 Dao & Gu (2024). All RNN models follow the Mamba architecture where the sequence models follows by a normalization and gating before output linear projection.
To investigate the effect of context length model performance, in this part, evaluated models trained with various context lengths, comparing Trellis (both with and without its forget gate) against the baseline models. For the Books dataset, models were trained
7

---

Published as a conference paper at COLM 2025
21
Transfrmer
20
Membe2
Tells we Forget
24-
19
GSA
Tellis
23
Parolodiy
22
21
16
Traraformer
TTT
15
19
Mamba2
Trells w/o Forget Trellss
14
18
2K
4K
8K
16K
2K
4K
BC
16K
Context Length
Connext Length
The Pile Detioct
14.5
Iapstenneee + UrannAttrtinn
- Unap Mertie
DelaNie:
Narrhd
Bzllaa
13.5
Tolla
Pepleaity
Pepieske:
12.5
1.5
10
507
76 11
Cantoct Length
Contect Corged
Figure 2: The effect of context length on model's perplexity. Subplots show: (Top Left) C4 dataset with 350M parameters; (Top Right) C4 dataset with 780M parameters; (Bottom Left) Books dataset with 125M parameters; (Bottom Right) The Pile cataset with 125M parameters. Training Transformers from scratch on very long sequence contexts (e.g., TE {16k, 32k)) can yield poor perplexity, hence the standard practice for such contexts is typically to finetune a Transformer pre-trained on shorter sequences (Touvron et al., 2023). Here, for the Transformer baseline trained from scratch in these results, performance is only reported up to context length of 7 = 8k.
Table 1: Performance of Trellis and baselines on language modeling and common-sense reasoning tasks.
Model
LMB. ppl 4
LMB. acc 1
PIQA
Hella. acc 1
Wino.
ARC-e acc 1
ARC-c acc f
SIQA acc f
Avg. t
acc 1
acc f
790M params / 30B tokens
Transformer++
25.89
33.41
64.75
41.98
51.33
59.06
31.85
40.26
46.09
Mamba2
28.91
32.72
64.98
42.60
50.01
61.99
30.24
41.07
46.23
ITT
27.05
33.18
65.03
43.17
49.93
62.16
32.13
41.35
46.71
Gated-DeltaNet
21.40
34.83
65.79
43.66
50.45
64.02
32.24
41.68
47.52
Trellis
20.28
35.44
67.51
44.29
51.08
65.12
33.17
42.04
48.38
using sequence lengths of (512, 1024, 2048, 4096, 8192, 16384, 32748). We also trained the models on the C4 dataset with context lengths in 22048, 4096, 8192, 16384), and on the Pile dataset using subset of context lengths T (2048, 8192). The results presented in Figure show that Trellis achieves the lowest perplexity compared to all baselines across all tested context lengths. Interestingly, Trellis shows greater performance gains compared to other linear RNNs as the sequence length increases, demonstrating the potential of our approach for tasks requiring long-context reasoning. Furthermore, comparing against Trellis w w/o forgetting highlights the importance of this component in the overall performance of Trellis.
Language Modeling and Common-sense Reasoning Following recent studies on se- quence modeling (Behrouz et al., 2024; Yang et al., 2024a), in this section, we compare the performance of Trellis with modern linear recurrent neural networks and Transformer language modeling and common-sense reasoning tasks. The results are reported in Table 1. Trellis achieves outstanding performance across all scales and outperform all linear recurrent models and Transformer++.
8

---

Published as a conference paper at COLM 2025
Table 2:
Ablations on improving from linear DeltaNet (Yang et al., 2024c) and also TTT (Sun et al., 2024a). All models have 125M parameters trained on The Pile dataset.
Table 3: Performance of Trellis and baselines with 1B param- eters on S-NIAH task from RULER benchmark. The best results with highest accuracy are highlighted.
S-NIAH-PK
S-NIAH-N
S-NIAH-W
Average
4K
8K
4K
Mamba2
98.6
61.4
31.0
98.4
58.8
14.2
62.2
42.2
4.2
$2.0
DeltaNet
96.8
98.8
98.6
47.2
15.4
12.8
83.2
46.2
20.0
$7.9
Gated DeltaNet
24.4
75.8
Configuration
ppl
Trellis
99.2
95.2
97.8
99.4
94.2
34.6
86.4
62.8
28.4
79,8
DeltaNet
11.58
TTT
11.44
25
Trellis
10.87
23
w/o forget gate
11.28
22
with f L2-SiLU 10.98
21
with f Softmax
11.29
20
Linear = tmax
12.71
Transformer
Linear = LN-SiLU
11.65
18
Mamba2
Trellis
GSA
32
11.14
17
10
1940
m
10.87
FLOPs
=
10.75
Figure 3: Scaling pattern of models w.r.t Perplexity vs. FLOPs.
To see the scaling law in Trellis and compare it with baselines, we report the perplexity for different model sizes in Figure 3. Trellis shows a consistent trend and achieves better perplexity compared to baselines with a fixed budget of FLOPs. This pattern shows that in the trade-off of efficiency and effectiveness, Trellis achieves Pareto frontier results.
Needle in Haystack Tasks In this section, we follow recent studies on sequence modeling and evaluate the performance of Trellis in the RULER benchmark and needle in haystack tasks (Hsieh al., 2024). The results are reported in Figure 3. Trellis outperforms all the baselines with +4% performance gain on average over the second best model, showing higher performance gain of about +6% longer sequences. We attribute the superior performance of Trellis to its: (1) more powerful memory management using the hybrid of linear and non-linear recurrence with forget gate, (2) architectural design, in which, we use two-pass recurrence with the same memory to store keys and values. In the results find two exceptions that Trellis achieves the second best result (i.e., S-NIAH-PK task with 4K and 8K sequence length). These results matches the observation of Yang et al. (2024a) that simple NIAH task with repeated synthetic context require long-term retention, which forget gate can damage.
Ablations In this section, we perform ablation studies on key components and design choices within the proposed architecture to evaluate the contribution of each to the overall performance. The baseline Trellis is a 2-pass K-V cache compression layer (presented in equations (13) and (14)) with 125M parameter (d 768 and 64). Its intermediate acti- vation function: f(x) LN-SiLU(x) == LayerNorm (SiLU(x)). Starting from this baseline configuration, we evaluate the impact of the following modifications (one component at time): (1) Removing forget gate, (2, 3) Changing the intermediate activation function:
SiLU(x)
[ and standard softmax, (4, 5) Removing the compression normalized SiLU: f(x)
layer non-linearity, i.e. (xx) x in equation 4, which reduces the recurrence in both passes to the Delta Rule. This linear recurrence is then tested with f Softmax and f= LN-SiLU, (6, 7) Varying the memory size, and (8) Using fully non-linear recurrence with mini-batch (chunk) size: B 1. The results, summarized in Table 2, highlight the significance of all

---

Published as a conference paper at COLM 2025
these design choices overall performance, notably, the 2-pass compression, forget gate and the choice of intermediate gate.
Table 4: Performance comparison of 125M-parameter language models on The Pile (2k context length) and Books (various context lengths) datasets. Baselines include Transformer++ (Touvron et al., 2023), Linear-Attention (LA) (Katharopoulos et al., 2020), DeltaNet (Yang et al., 2024c), Gated-DeltaNe (Yang et al., 2024a), Mamba2 (Dao & Gu, 2024), and TTT (Sun et al., 2024a). The best results are highlighted.
Model
Pile (k)
Pile tik)
Books
Books (lk)
Books (3k)
Books Hlk) ppl
Books (Sk)
Books (16k)
Books 152 Peal
Pile PP"
ppl
ppl 1
Ppl
ppl
ppl
ppl
Fpl
123M params
2:411 tokeess
30M parants 77.0t tokens
Transformene-
18.90
18.58
18:85
Lineue-Atsmtion
12:43
13.26
23.54
19.96
19.60
19.64
20.34
21.57
25.07
9.16
DetaNe:
19.21
Mamba2
11.47
11.35
19.6
18.89
18.19
18.00
18.23
18.40
19.33
8.56
Galed DelaNee
17.40
11.44
11.29
1991
18.28
18.19
18/0
18.34
18.29
19:16
8222
Conclusion
In this paper, we introduced Trellis, meta in-context learning framework that learns how to compress the KV cache of attention into a fixed-size memory. Trellis employs two- pass recurrent memory update process, in which keys and values are stored a compact memory module. To enhance learning from long context, Trellis further incorporates a gating mechanism that learns how to filter and forget irrelevant past information. Our experimental results across diverse domains domains-includiin including language modeling, commonsense reasoning, whedle-in-a-haystack tasks, and time series data-indicate the superior performance of Trellis compared to modern linear RNNs and Transformers. The results further support the new architectural design of Trellis, showing that all its components meaningfully contribute to its overall performance. By introducing learned, dynamic memory compression, Trellis offers an effective and efficient solution for long-sequence modeling.
Looking ahead, recent advances in non-linear recurrent models such as Titans (Behrouz et al., 2024; 2025a), Lattice (Karami & Mirrokni, 2025) offer promising directions for further improving KV cache compression. Additionally, finetuning pretrained Transformers into RNNs (T2R) (Kasai et al., 2021) represents another promising direction to explore the capa- bilities of large pre-trained models with the inference efficiency of recurrent architectures.
10

---

Published as a conference paper at COLM 2025
References
Marcin Andrychowicz, Misha Denil, Sergio Gomez, Matthew W Hoffman, David Pfau, Tom Schaul, Brendan Shillingford, and Nando De Freitas. Learning to learn by gradient descent by gradient descent. Advances in neural information processing systems, 29, 2016.
Anurag Arnab, Mostafa Dehghani, Georg Heigold, Chen Sun, Mario Lučić, and Cordelia Schmid. Vivit: A video vision transformer. In Proceedings of the IEEE/CVF international conference on computer vision, Pp. 6836-6846, 2021.
Simran Arora, Sabri Byuboglu, Michael Zhang, Aman Timalsina, Silas Alberti, Dylan Zinsley, James Zou, Atri Rudra, and Christopher Ré. Simple linear attention language models balance the recall-throughput tradeoff. ar Xiv preprint arXiv:2402.11666, 2024.
Maximilian Beck, Korbinian Pöppel, Markus Spanring, Andreas Auer, Oleksandra Prud- nikova, Michael Kopp, Günter Klambauer, Johannes Brandstetter, and Sepp Hochreiter. XLSTM: Extended long short-term memory. arv preprint ar Xio:2405.00517. 2024.
Ali Behrouz, Peilin Zhong, and Vahab Mirrokni. Titans: Learning to memorize at test time.
arviv preprint ar Xrio:2501.00663, 2024.
Ali Behrouz, Zeman Li, Praneeth Kacham, Majid Daliri, Yuan Deng, Peilin Zhong, Meisam Razaviyayn, and Vahab Mirrokni. Atlas: Learning to optimally memorize the context at test time. Xiviv preprint arXiv:2505.23735, 2025a.
Ali Behrouz, Meisam Razaviyayn, Peilin Zhong, and Vahab Mirrokni. It's all connected: A journey through test-time memorization, attentional bias, retention, and online optimiza- tion, 2025b. URL
13173.
Iz Beltagy, Matthew E Peters, and Arman Cohan. Longformer: The long-document trans- former. arXiu preprint ar Xrio:2004.051500 2020.
Yoshua Bengio, Samy Bengio, and Jocelyn Cloutier. Learning synaptic learning rule. Citeseer, 1990.
Guy E Blelloch. Prefix sums and their applications. 1990.
Beidi Chen, Tri Dao, Eric Winsor, Zhao Song, Atri Rudra, and Christopher Ré. Scatterbrain: Unifying sparse and low-rank attention. In Advances in Neural Information Processing Systems (NeurIPS), 2021.
Can Chen, Xi Chen, Chen Ma, Zixuan Liu, and Xue Liu. Gradient-based bi-level optimization for deep learning: A survey. arXiv preprint ar XrXio:2207.11719, 2022.
Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. ar preprint ar 105509994..0509, 2019.
Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, et al. Rethinking attention with performers. arxiv preprint Xiv::20009 14499,, 2020.
Tri Dao and Albert Gu. Transformers are SSMs: Generalized models and efficient algorithms through structured state space duality. In International Conference on Machine Learning (ICML), 2024.
Soham De, Samuel L Smith, Anushan Fernando, Aleksandar Botev, George Eristian-Murard Albert Gu, Ruba Haroun, Leonard Berrada, Yutian Chen, Srivatsan Srinivasan, et al. Griffin: Mixing gated linear recurrences with local attention for efficient language models. arv preprint ar (rXiv:2402.19477 2024.
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre- training of deep bidirectional transformers for language understanding. arXio preprint Xrv: 0011181004800 2018.
11

---

Published as a conference paper at COLM 2025
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. Xiv preprint ar Xit:20000.11929, 2020.
Vijay Prakash Dwivedi and Xavier Bresson. A generalization of transformer networks to graphs. Xiviv preprint afXtt:2012.09699, 2020.
Daniel Y Fu, Tri Dao, Khaled K Saab, Armin W Thomas, Atri Rudra, and Christopher Ré. Hungry hungry hippos: Towards language modeling with state space models. International Conference on Learning Representations, 2023.
Leo Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He, Anish Thite, Noa Nabeshima, et al. The pile: An 800gb dataset of diverse text for language modeling. arviv preprint arXiv:2101.0000
2020.
Xavier Gonzalez, Andrew Warrington, Jimmy TH Smith, and Scott W Linderman. Towards scalable and stable parallelization of nonlinear rnns. arXio preprint ar Xxtx:2407.19115, 2024.
Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. arv preprint ar arXit:2312.00752, 2023.
Albert Gu, Tri Dao, Stefano Ermon, Atri Rudra, and Christopher Ré. HiPPO: Recurrent memory with optimal polynomial projections. Advances in neural information processing systemss, 33:1474414487, 2020.
Albert Gu, Isys Johnson, Karan Goel, Khaled Saab, Tri Dao, Atri Rudra, and Christopher Ré. Combining recurrent, convolutional, and continuous-time models with linear state space layers. Advances in neural information processing systems, 34:572-585, 2021.
Simon S Haykin. Adaptive filter theory. Pearson Education India, 2002.
Cheng-Ping
Hsieh, Simeng Sun, Samuel Kriman, Shantanu Acharya, Dima Rekesh, Fei Jia, and Boris Ginsburg. RULER: What's the real context size of your long-context language models? In First Conference on Language Modeling, 2024. URL https:// https://openreviiew net/ forum?id=kIoBbc765))
Weizhe Hua, Zihang Dai, Hanxiao Liu, and Quoc Le. Transformer quality in linear time. In International conference on machine learning, pp. 9099-9117. PMLR, 2022.
Praneeth Kacham, Vahab Mirrokni, and Peilin Zhong. Polysketchformer Fast transformers via sketching polynomial kernels. In Forty-first International Conference on Machine Learning, 2024. URL https: / http///pppeeeeeeeeee/eeeeeeeeeee hwtp:viwwwneefffooom wwt.:/ertfoott fdd
Mahdi Karami. HiGen: Hierarchical graph generative networks. In The Twelfth International Conference on Learning Representations, 2024.
Mahdi Karami and Ali Ghodsi. Orchid: Flexible and data-dependent convolution for
sequence modeling. In Thirty-eighth Conference on Advances in Neural Information Processing Systems, 2024. URL https:
18508.
Mahdi Karami and Vahab Mirrokni. Lattice: Learning to efficiently compress the memory. arXiv preprint arXio:2504.055 006644,,0.666.... 2025.
Mahdi Karami, Martha White, Dale Schuurmans, and Csaba Szepesvári. Multi-view matrix factorization for linear dynamical system estimation. Advances in Neural Information Processing Systems, 30, 2017.
Mahdi Karami, Dale Schuurmans, Jascha Sohl-Dickstein, Laurent Dinh, and Daniel Duck- worth. Invertible convolutional flow. Advances in Neural Information Processing Systems, 32, 2019.
Jungo Kasai, Hao Peng, Yizhe Zhang, Dani Yogatama, Gabriel Ilharco, Nikolaos Pappas, Yi Mao, Weizhu Chen, and Noah A Smith. Finetuning pretrained transformers into mnns. ar Xiv preprint Xii::2003.13076 2021.
12

---

Published as a conference paper at COLM 2025
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are mans: Fast autoregressive transformers with linear attention. In International Conference on Machine Learning, pp. 5156-5165. PMLR, 2020.
Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. In International Conference on Learning Representations, 2020.
Yuhong Li, Tianle Cai, Yi Zhang, Deming Chen, and Debadeepta Dey. What makes con-
volutional models great on long sequence modeling? ar arviv preprint arXio:2210.09298, 2022.
Yi Heng Lim, Qi Zhu, Joshua Selfridge, and Muhammad Firmansyah Kasim. Parallelizing non-linear sequential models over the sequence length. Xivv preprint ar XrX::2209.11222, 2023.
Bo Liu, Mao Ye, Stephen Wright, Peter Stone, and Qiang Liu. Bome! bilevel optimization made easy: A simple first-order approach. Advances in neural information processing systems, 15:17248-1726,, 2022.
Harsh Mehta, Ankit Gupta, Ashok Cutkosky, and Behnam Neyshabur. Long range language modeling via gated state spaces. Xrv preprint wwXio:2206.1394/ 2022.
Antonio Orvieto, Samuel L Smith, Albert Gu, Anushan Fernando, Caglar Gulcehre, Razvan Pascanu, and Soham De. Resurrecting recurrent neural networks for long sequences. In International Conference on Machine Learning, pp. 26670-26698. PMLR, 2023.
Niki Parmar, Ashish Vaswani, Jakob Uszkoreit, Lukasz Kaiser, Noam Shazeer, Alexander Ku, and Dustin Tran. Image transformer. Proceedings of ICML 2018, 2018.
Hao Peng, Jungo Kasai, Nikolaos Pappas, Dani Yogatama, Zhaofeng Wu, Lingpeng Kong, Roy Schwartz, and Noah A Smith. ABC: Attention with bounded-memory control. arXiv preprint ar Xix::2100002488, 2021a.
Hao Peng, Nikolaos Pappas, Dani Yogatama, Roy Schwartz, Noah A Smith, and Lingpeng Kong. Random feature attention. Proceedings of ICLR, 2021b.
Michael Poli, Stefano Massaroli, Eric Nguyen, Daniel Y Fu, Tri Dao, Stephen Baccus, Yoshua Bengio, Stefano Ermon, and Christopher Ré. Hyena hierarchy: Towards larger convolu- tional language models. International Conference on Machine Learning, 2023.
Jiezhong Qiu, Hao Ma, Omer Levy, Scott Wen-tau Yih, Sinong Wang, and lie Tang. Blockwise self-attention for long document understanding, arviv preprint arXiv:1911.00972, 2019.
Alec Radford, Karthik Narasimhan, Tim Salimans, and llya Sutskever. Improving language understanding by generative pre-training. 2018
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with unified text-to-text transformer. Journal of machtine learning research, (1(140):1667,202)
David W Romero, Anna Kuzina, Erik J Bekkers, Jakub Mikolaj Tomczak, and Mark Hoogen- doorn. Ckconv: Continuous kernel convolution for sequential data. In International Conference Learning Representations, 2021.
Aurko Roy, Mohammad Saffar, Ashish Vaswani, and David Grangier. Efficient content-based sparse attention with routing transformers. Proceedings of TACI, 2020.
Imanol Schlag, Kazuki Trie, and Jürgen Schmidhuber. Linear transformers are secretly fast weight programmers. In International Conference on Machine Learning, pp. 9355-9366. PMLR, 2021.
JH Schmidhuber.
Learning to control fast-weight memories: An alternative to recurrent nets. Neural Compulation, 1992.
13

---

Published as a conference paper at COLM 2025
Prajwal Singhania, Siddharth Singh, Shwai He, Soheil Feizi, and Abhinav Bhatele. Loki: Low-rank keys for efficient sparse attention. arviv preprint arXio:2406.02542) 2024.
Jimmy T.H. Smith, Andrew Warrington, and Scott Linderman. Simplified state space layers for sequence modeling. In The Eleventh International Conference on Learning Representations, 2023. URL https: / httpss///pppeeeeeeeeee net wwtt//eerr/eeeeee Ai AHw3AAKK.
Yu Sun, Xinhao Li, Karan Dalal, Jiarui Xu, Arjun Vikram, Genghan Zhang, Yann Dubois, Xinlei Chen, Xiaolong Wang, Sanmi Koyejo, et al. Learning to (learn at test time): Rnns with expressive hidden states. arXiv preprint ar Xii:24007.46200 2024a.
Yutao Sun, Li Dong, Yi Zhu, Shaohan Huang, Wenhui Wang, Shuming Ma, Quanlu Zhang, Jianyong Wang, and Furu Wei. You only cache once: Decoder-decoder architectures for language models. Advances Neural Information Processing Systems, 37:7339-7361, 2024b.
Yi Tay, Dara Bahri, Liu Yang, Donald Metzler, and Da-Cheng Juan. Sparse a Proceedings of ICML, 2020.
Yi Tay, Mostafa Dehghani, Dara Bahri, and Donald Metzler. Efficient transformers: A survey. ACM Computing Surveys, 55(6):1-28, 2022.
Sebastian Thrun and Lorien Pratt. Learning to learn: Introduction and overview. In Learning 10 learn, pp. 3-17. Springer, 1998.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXio:2302.1397 2023.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. volume 30, 2017.
Junxiong Wang, Jing Nathan Yan, Albert Gu, and Alexander M Rush. Pretraining without attention. ar arv preprint ar Xxxx:2212.105444 2022.
Ke Alexander Wang, Jiaxin Shi, and Emily B Fox. Test-time regression: a unifying framework for designing sequence models with associative memory. arxio preprint arXiv:2501.11235, 2025.
Sinong Wang, Belinda Z Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self- attention with linear complexity. arXiv preprint arXiv:2006.0477 2020.
Bernard Widrow and Marcian E Hoff. Adaptive switching circuits. In Neurocomputing: foundatioon of research, pp. 123-134. 1988.
Songlin Yang, Jan Kautz, and Ali Hatamizadeh. Gated delta networks: Improving mamba2 with delta rule. Xiv preprint ar Xrxi::2412.06466 2024a.
Songlin Yang, Bailin Wang, Yikang Shen, Rameswar Panda, and Yoon Kim. Gated linear at- tention transformers with hardware-efficient training. In Forty-first International Conference on Machine Learning, 2024b. URL htttps://openreiieee net wtt.foooddiidiwwfeeeeee
Songlin Yang, Bailin Wang, Yu Zhang, Yikang Shen, and Yoon Kim. Parallelizing linear transformers with the delta rule over sequence length. In The Thirty-eighth Annual Con- ference on Neural Information Processing Systems, 2024c. URL https://openreeiew.. net/ Forum??d:=y8RmAWWWPP
Seongjun Yun, Minbyul Jeong, Raehyun Kim, Jaewoo Kang, and Hyunwoo J Kim. Graph transformer networks. Advincces in neural information processing systems, 32, 2019.
Manzil Zaheer, Guru Guruganesh, Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, el. Big bird: Transformers for longer sequences. Proceedings of NeurlPS, 2020.
Yu Zhang, Songlin Yang, Ruijie Zhu, Yue Zhang, Leyang Cui, Yiqiao Wang, Bolun Wang, Freda Shi, Bailin Wang, Wei Bi, et al. Gated slot attention for efficient linear- sequence modeling. arXiv preprint ar 249.00240 2024.
14

---

Published as a conference paper at COLM 2025
A
Experiment Details
Architectural Details. We provide the architectural details of our model in Table
Table 5: Architectural Details.
Model
Block
Dim
Head
Peak I.
Token
12 125M
12
768
12
3e-3
2.4H
170M
12
768
16
3e-3
15B
350M
24
1024
16
1.5e-3
15B
780M
24
1536
16
1.25e-3
30B
The Choice of Datasets and Training. We train our models on three different datasets. Our goal is to highlight the generalizability and power of Trellis across diverse datasets. For smaller and medium-sized models (125M and 350M parameters), we included broader set of baselines and conducted detailed ablation studies. Due to computational constraints, we limited the training of larger models (780M and 1B parameters) to comparisons with the most relevant baselines. Note that Tables and 3 present results for models with 780M and 1B parameters, respectively, to show Trellis's 's performance at larger scales. The results in the bottom row of Figure are based on models trained on The Pile and Books datasets. All ablation studies and Table 4 are also conducted using models trained on The Pile. The models used in the remaining experiments (including Table 3 and Figure 3) are trained on the C4 dataset.
15


Reviews:
[{"review_id": "Uc2pAMnh71", "reviewer": "Reviewer_RaMN", "summary": "This paper introduces Trellis, a novel Transformer variant designed to address the scalability and memory bottlenecks of attention mechanisms in long-context sequence modeling. Trellis replaces the traditional key-value (KV) cache with a learned, bounded memory using a two-pass compression mechanism. The memory is updated using an online gradient descent procedure with a forget gate, allowing it to dynamically compress incoming key-value pairs while maintaining relevant context. The method is evaluated on a variety of tasks—language modeling, reasoning, and time-series forecasting—where it consistently outperforms strong baselines, particularly at longer sequence lengths.", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "* The paper addressed the fundamental limitation in Transformer attention. It proposed several techniques on top of previous SSM to improve the performance.\n\n* It has shown strong performance against previous Transformer variants.", "reasons_to_reject": "* The paper has limited novelty. The main structure is following previous SSM papers like Mamba. It has proposed some variants on top of it. However, the paper has not made clear the difference between different architectures.\n\n* The paper has only conducted Single Needle in A Haystack evaluation. For SSM models with limited memory size, memorization capability is important and should be measured in detail. The forget and state decay mechanism may be able to retrain a short info but cannot hold the entire previous information history. It is important to study the limitation in SSM-type of models to make a fair comparison against Transformer.\n\n* The current setting follows the previous pattern (1) propose a new architecture (2) train with limited number of tokens and controlled number of parameters. (3) evaluate on a set of general language modeling benchmarks.\nIt would be really interesting to jump out of this box and see experiments around more fundamental capabilities around Transformers and SSMs.", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "FzWAxMxB3O", "reviewer": "Reviewer_1Hgq", "summary": "The Trellis paper presents a novel and ambitious Transformer design to tackle the critical issue of the growing Key-Value cache. It enables the model to learn dynamic compression of this memory into a fixed-size store. This core concept of a learnable, adaptive KV cache, managed through an online gradient descent procedure and a forget gate, is a highly original and principled approach for efficient long-context modeling. The methodology is clearly explained, reflecting a solid grasp of Transformer mechanics, optimization, and hardware efficiency considerations. The work is commendably supported by extensive experiments across diverse tasks.  To further strengthen its claims and offer an even more comprehensive view of its standing, the experimental validation would benefit from including direct comparisons with other recent and highly relevant architectures designed for efficient long-context processing, such as Titans and RWKV. Benchmarking against these specific models would provide a more complete understanding of Trellis's advantages within this rapidly evolving research landscape. Despite this point for further investigation, Trellis presents a significant and innovative contribution. It offers a compelling and well-engineered solution to a crucial problem in scaling Transformers, supported by robust empirical evidence of its effectiveness and efficiency. This makes it a high-quality paper with substantial potential to influence future work in efficient sequence modeling.", "questions": "", "limitations": "", "rating": 10, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "mentioned in summary", "reasons_to_reject": "N/A", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "1x4vqwngOh", "reviewer": "Reviewer_M9Vs", "summary": "In this paper, the authors propose Trellis: a new method that allows transformers to have a bounded memory (compressed KV cache) which is updated at each step and contains a forget gate that allows the model to selectively forget past information.\n\nThis work is inspired by ABC (Attention-with-Bounded-Memory Control) and GSA (Gated Slot Attention). But differently from ABC and GSA, the authors approach the memory’s update as an online regression problem, updating the memory using a gradient descent step per token. \nSimilarly to GSA, the memory update includes a forget gate which allows the model to forget past information. This is accomplished by applying an l_2 regularization in the loss.\n\nTo ensure its applicability the authors provide a parallel and efficient implementation of Trellis.\n\nThe authors compare the approach with baselines comprising recurrent and efficient transformer architectures on language modelling, common-sense reasoning tasks, needle in the haystack, and time series.\nThese comparisons show that TRELLIS outperforms the baselines in most tested tasks across datasets and model parameters. Moreover, the authors show that for the same number of FLOPs Trellis leads to a lower perplexity than Mamba2, TTT, and GSA.\n\nOverall, I believe the proposed approach is interesting and applicable to a wide range of tasks. However, the experimental section is quite disappointing.", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "-\tAccording to the experiments, the approach proposed seems to outperform other efficient architectures.\n-\tThe authors show that Trellis’ implementation can be efficient and parallelizable.", "reasons_to_reject": "-\tThe experimental section is quite disappointing and missing details. \nFirst, I did not find an explanation as to why different baselines were used for different experiments (for example in figure 2 the two top plots have a list of baselines completely different than the that of the two bottom plots, and figures 4 or 5 even have a smaller number of baselines). \nThen, the authors mention that they train models of 4 model sizes up to 1B parameters on the 3 datasets. However, most of the comparisons show results for models with only 125M parameters or 780M parameters. Why is this?\n-\tNowadays, we see LLMs with context lengths bigger than 100k tokens. However, the paper experiments are limited to context lengths of 16k or 32k. I understand that experiments with bigger models trained with more tokens and with bigger context lengths are costly. But this limitation makes it difficult to understand the model’s applicability to real-world scenarios.", "questions_to_authors": "-\tIn lines 96 to 102, you mention “… the additive update rule in ABC and GSA endlessly accumulates new tokens into the fixed memory space without an explicit mechanism to deallocate past information …”. \nHowever, the GSA contains a gating mechanism used to forget past information. Can you explain why you make this statement?\n-\tSee “reasons to reject”.", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper introduces a method (Trellis) which compresses the KV cache of transformers, addressing the scalability and memory bottlenecks of attention mechanisms in long-context sequence modeling. The method is inspired by ABC (Attention-with-Bounded-Memory Control) and GSA (Gated Slot Attention) but adds several innovations. The compressed memory is bounded and updated using an online gradient descent procedure with a forget gate. The method is evaluated on language modeling, reasoning, and time-series forecasting with convincing performance for long sequences, outperforming baselines comprising recurrent and efficient transformer architectures, Mamba2, TTT, and GSA.

All reviewers feel positively about this paper, and so do I. While the paper has some drawbacks (somewhat limited novelty and lack of details in the experiments) the proposed approach is interesting, efficiently implemented, and appears to be effective in several tasks. I recommend acceptance, encouraging the authors to address the feedback provided by the reviewers.

[Automatically added comment] At least one review was discounted during the decision process due to quality]
Decision: Accept