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

We thank the reviewer for their constructive feedback and valuable observations. In the following, we address their concerns and respond to questions.

I did not find an explanation as to why different baselines were used for different experiments

Our goal in training models on three different datasets, reported in Figure 2 (and also Table 4) was to highlight the generalizability and power of Trellis across diverse datasets. For smaller and medium-sized models (125M and 350M parameters), we included a 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. We note that Table 1 and Figure 4 present results for models with 780M and 1B parameters, respectively, to show Trellis’s performance at larger scales. We acknowledge that the variation in baseline coverage across plots requires a clarification, and we will make these choices explicit in the experimental details section of the final version.

Nowadays, we see LLMs with context lengths bigger than 100k tokens.

We acknowledge the reviewer’s concern regarding context length greater than 100k tokens. In this work, we limited our evaluation to a maximum of 32K tokens, in line with the settings used in the main baseline papers. Moreover, widely used long-context evaluation suites for language modeling, such as LongBench [1], include tasks with sequences that rarely exceed 32K tokens with the average context lengths in these benchmarks typically range from ~1K to maximum of ~22K.

In 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 …”. However, the GSA contains a gating mechanism used to forget past information.

While it is true that GSA incorporates a gating mechanism, it operates within an additive (Hebbian-like) update form, where forgetting is implemented as a decay (fading) of the entire memory content. In contrast, the update rule in Trellis, derived from online optimization (Eqs. 5–6), is a form of Fast Weight Programmers (FWPs) and generalizes the classical Delta Rule (Widrow & Hoff, 1988; Schlag et al., 2021) (we can verify it by simply removing the nonlinearity (i.e., setting $\phi(\cdot)$ to the identity function), our update reduces to the Delta Rule). The Delta Rule recurrence can be written as (Schlag et al., 2021, Yang et al., 2024b):

Therefore, unlike gated additive updates in GSA, which modulates how fast past memory fades, Delta Rule provides a mechanism to deallocate or overwrite irrelevant past information based on the interaction between incoming keys and the current state, offering higher memory capacity and greater flexibility. The non-linear generalization used in the two-pass process in Trellis is demonstrated in the experiments that yields stronger empirical performance.

References:

[1] Bai, Yushi, et al. "Longbench: A bilingual, multitask benchmark for long context understanding." arXiv preprint arXiv:2308.14508 (2023).

We sincerely thank the reviewer for their insightful and encouraging feedback. We are especially grateful for their recognition of the novelty, and potential impact of our work and appreciate the acknowledgement of both the technical depth of our approach and the breadth of our experimental validation.

We thank the reviewer for their feedback and for highlighting the strong performance of our work. In the following, we address their concerns and, in particular, further clarify the key novelties of our proposed method.

The paper has limited novelty. The main structure is following previous SSM papers like Mamba.

We highlight several fundamental differences and novelties in our approach, which are key to its strong performance and distinguish it from the available SSM architectures like Mamba.

• Trellis is derived from the principles of online optimization for representation learning, leveraging techniques from meta-learning and test-time training. Its update rule is a form of online gradient descent on reconstruction losses to update the internal memory of the pairs of Keys/Values in the transformers.
• It learns how to compress Key-Value pairs into a bounded memory via a non-linear, two-pass recurrent process which also takes advantage of parallel and hardware efficient implementation.
• We tailored a different form of two-pass process in equations (13) (14) compared to prior work like ABC and incorporates a normalized SiLU activation, which our ablation studies demonstrate is highly effective for retrieval in bounded memory.
• We also derived a state decay mechanism based on $l_2$ regularization in the inner loop. The overall update rule is state and input--dependent (the update for a memory slot depends on the current value of that slot and also the layer’s input).

In the above we highlighted the key difference of the proposed method vs the selective SSM architectures like Mamba. While Trellis architecture fundamentally takes advantage of non-linear update rules which is state-dependent and input dependent in a novel two-pass process, Mamba’s gating mechanism is input dependent and updates a single associative memory in a linear form. Furthermore, the derivation of our model from an online optimization perspective is different theoretically from Mamba's derivation from continuous-time systems.

The forget and state decay mechanism may be able to retrain a short info but cannot hold the entire previous information history.

As discussed in section 2, while Transformer-based models benefit from an unbounded KV cache, the SSM and RNN alternatives have bounded state size (memory), hence their ability to recall the previous information history is inherently limited and they can’t memorize the entire context in long sequence. Given this limitation, the objective of this work is not to achieve lossless recall, but to design a non-linear recurrence that efficiently compress new keys and values into the memory while minimizing information loss. We highlight two key aspects of the proposed design for effective long-term memory management:

1. Adaptive State Decay: Regarding the concern on state decay (forget), the forget gate in Trellis is not a simple, constant decay of past information (as seen in models like S4 or the weight decay in deep neural network training). Instead, it is an adaptive input-dependent gating mechanism that learns to adjust its decay based on the current input, enabling it to selectively preserve important past information while making room for new context. Notably, we highlight that the forget gate weights are biased towards one by design, ensuring that the model's default behavior is to retain information unless the input dictates otherwise.
2. Nonlinear Memory Update Rule: Furthermore, the core of Trellis's memory management lies in its update rule, derived from online optimization (Eqs. 5-6). Its non-linear update is a generalization of the Delta Rule (Widrow & Hoff, 1988; Schlag et al., 2021). By simply removing the nonlinearity (setting $\phi(\cdot)$ to identity in Eqs. 5-6), our update reduces to the Delta Rule, which is known to offer higher memory capacity and greater flexibility compared to the simple additive (Hebbian-like) updates used in most other SSMs, such as Mamba. Importantly, it provides a mechanism to deallocate or overwrite irrelevant past information, which simple additive updates do not (https://openreview.net/forum?id=r61s1FNYlj&noteId=8dqH0mSWc2).

Together with our modified two-pass process, these mechanisms render Trellis an effective long-context sequence model. This is empirically validated in Figure 2, where Trellis's performance advantage over baselines like GSA and DeltaNet becomes more pronounced as the context length increases to 32K. Additionally, we evaluated the memorization capability of Trellis in the RULER benchmark, a widely adopted suite specifically designed to evaluate long-context language models. We assess performance on three increasingly complex tasks: passkey retrieval (S-NIAH-PK), numerical needle in haystack (S-NIAH-N), and word-based needle in haystack (S-NIAH-W). Finally, our ablation studies further characterize the effect of higher memory size ($m$) on performance improvement.