Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Mahdi Karami, Ali Behrouz, Praneeth Kacham, Vahab Mirrokni

TL;DR

Trellis tackles the inefficiency of Transformer attention by introducing a bounded-memory transformer that learns to compress its key-value memory at test time. It uses a two-pass, non-linear compression with an online gradient-descent updater and a forget gate to store and selectively forget contextual information in a fixed-size memory. Across language modeling, commonsense reasoning, recall-heavy tasks, and time-series data, Trellis outperforms strong baselines and shows stronger gains as context length increases, highlighting its suitability for long-context settings. This approach enables efficient memory management and scalable inference without sacrificing performance, with promising directions for future integration with non-linear recurrent architectures and pretrained models.

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.

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

TL;DR

Trellis tackles the inefficiency of Transformer attention by introducing a bounded-memory transformer that learns to compress its key-value memory at test time. It uses a two-pass, non-linear compression with an online gradient-descent updater and a forget gate to store and selectively forget contextual information in a fixed-size memory. Across language modeling, commonsense reasoning, recall-heavy tasks, and time-series data, Trellis outperforms strong baselines and shows stronger gains as context length increases, highlighting its suitability for long-context settings. This approach enables efficient memory management and scalable inference without sacrificing performance, with promising directions for future integration with non-linear recurrent architectures and pretrained models.

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.
Paper Structure (15 sections, 14 equations, 4 figures, 5 tables)

This paper contains 15 sections, 14 equations, 4 figures, 5 tables.

Table of Contents

  1. Introduction
  2. Our approach and main contributions
  3. Related Work and Background
  4. Method
  5. State Decay
  6. Parallel and Hardware Efficient Implementation
  7. Related Works and Discussion
  8. Experiments
  9. Overview:
  10. Baseline Models:
  11. Language Modeling and Common-sense Reasoning
  12. Needle in Haystack Tasks
  13. Ablations
  14. Conclusion
  15. Experiment Details

Figures (4)

  • Figure 1: (Left) Block diagram of the language model. (Right) The Trellis block. Each sequence mixing block is composed of a short $\operatorname{Conv1D}$ for $\{q, k \}$ and the Trellis is followed by a post normalization and the $\operatorname{GeLU}$ post-gate.
  • 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 dataset with 125M parameters. Training Transformers from scratch on very long sequence contexts (e.g., $T \in \{16k, 32k\}$) can yield poor perplexity, hence the standard practice for such contexts is typically to finetune a Transformer pre-trained on shorter sequences touvron2023llama. Here, for the Transformer baseline trained from scratch in these results, performance is only reported up to a context length of $T=8k$.
  • Figure 2: Ablations on improving from linear $\operatorname{DeltaNet}$yang2024parallelizing and also $\operatorname{TTT}$sun2024ttt. All models have 125M parameters trained on The Pile dataset.
  • Figure 3: Scaling pattern of models w.r.t. Perplexity vs. FLOPs.