Leave No Context Behind: Efficient Infinite Context Transformers with Infini-attention — part 1/3

Cleaned-text code points [0, 15277). Bibliography omitted here; source.md includes it.

# Leave No Context Behind: Efficient Infinite Context Transformers with Infini-attention

###### Abstract

This work introduces an efficient method to scale Transformer-based Large Language Models (LLMs) to infinitely long inputs with bounded memory and computation. A key component in our proposed approach is a new attention technique dubbed Infini-attention. The Infini-attention incorporates a compressive memory into the vanilla attention mechanism and builds in both masked local attention and long-term linear attention mechanisms in a single Transformer block. We demonstrate the effectiveness of our approach on long-context language modeling benchmarks, 1M sequence length passkey context block retrieval and 500K length book summarization tasks with 1B and 8B LLMs. Our approach introduces minimal bounded memory parameters and enables fast streaming inference for LLMs.

## 1 Introduction

Memory serves as a cornerstone of intelligence, as it enables efficient computations tailored to specific contexts. However, Transformers ( Vaswani et al., 2017 ) and Transformer-based LLMs ( Brown et al., 2020 ; Touvron et al., 2023 ; Anil et al., 2023 ; Groeneveld et al., 2024 ) have a constrained context-dependent memory, due to the nature of the attention mechanism.

The attention mechanism in Transformers exhibits quadratic complexity in both memory footprint and computation time. For example, the attention Key-Value (KV) states have 3TB memory footprint for a 500B model with batch size 512 and context length 2048 ( Pope et al., 2023 ) . Indeed, scaling LLMs to longer sequences (i.e. 1M tokens) is challenging with the standard Transformer architectures and serving longer and longer context models becomes costly financially.

Compressive memory systems promise to be more scalable and efficient than the attention mechanism for extremely long sequences ( Kanerva, 1988 ; Munkhdalai et al., 2019 ) . Instead of using an array that grows with the input sequence length, a compressive memory primarily maintains a fixed number of parameters to store and recall information with a bounded storage and computation costs. In the compressive memory, new information is added to the memory by changing its parameters with an objective that this information can be recovered back later on. However, the LLMs in their current state have yet to see an effective, practical compressive memory technique that balances simplicity along with quality.

In this work, we introduce a novel approach that enables Transformer LLMs to effectively process infinitely long inputs with bounded memory footprint and computation. A key component in our proposed approach is a new attention technique dubbed Infini-attention (Figure 1 ). The Infini-attention incorporates a compressive memory into the vanilla attention mechanism ( Bahdanau et al., 2014 ; Vaswani et al., 2017 ) and builds in both masked local attention and long-term linear attention mechanisms in a single Transformer block.

Such a subtle but critical modification to the Transformer attention layer enables a natural extension of existing LLMs to infinitely long contexts via continual pre-training and fine-tuning.

Our Infini-attention reuses all the key, value and query states of the standard attention computation for long-term memory consolidation and retrieval. We store old KV states of the attention in the compressive memory, instead of discarding them like in the standard attention mechanism. We then retrieve the values from the memory by using the attention query states when processing subsequent sequences. To compute the final contextual output, the Infini-attention aggregates the long-term memory-retrieved values and the local attention contexts.

In our experiments, we show that our approach outperforms baseline models on long-context language modeling benchmarks while having 114x comprehension ratio in terms of memory size. The model achieves even better perplexity when trained with 100K sequence length. A 1B LLM naturally scales to 1M sequence length and solves the passkey retrieval task when injected with Infini-attention. Finally, we show that a 8B model with Infini-attention reaches a new SOTA result on a 500K length book summarization task after continual pre-training and task fine-tuning.

In summary, our work makes the following contributions: 1. We introduce a practical and yet powerful attention mechanism – Infini-attention with long-term compressive memory and local causal attention for efficiently modeling both long and short-range contextual dependencies.

2. Infini-attention introduces minimal change to the standard scaled dot-product attention and supports plug-and-play continual pre-training and long-context adaptation by design.

3. Our approach enables Transformer LLMs to scale to infinitely long context with a bounded memory and compute resource by processing extremely long inputs in a streaming fashion.

## 2 Background

Recurrent Neural Networks (RNNs) process a single token x t x_{t} at each step t t and computes a recurrent hidden state h t h_{t} to represent an entire input sequence ( Hochreiter & Schmidhuber, 1997 ; Maass et al., 2002 ) : h t = R ​ N ​ N ​ ( x t , h t − 1 ) . h_{t}=RNN(x_{t},h_{t-1}). (1) The RNN computation is very efficient since the model maintains only a fixed-size vector h t h_{t} for input sequence. However, for processing long sequences it becomes difficult to store entire contextual information into a single fixed-size vector and this limitation had implications on RNNs utility in certain tasks ( Kaiser & Sutskever, 2015 ) . To address the limitation, people extended the standard RNNs with an external memory component that can be read from and written to. One such an instance is Metalearned Neural Memory (MNM) ( Munkhdalai et al., 2019 ) : h t , θ t = M ​ N ​ M ​ ( x t , h t − 1 , θ t − 1 ) . h_{t},\theta_{t}=MNM(x_{t},h_{t-1},\theta_{t-1}). (2)

MNM learns an additional memory state θ \theta parameterized by a feed-forward neural network (FFN) and uses query, key and value vectors (QKV) to interact with the memory, similar to the attention mechanism. To store information, it modifies the parameters of the FFN by using the key vectors as input and the value vectors for the target, and to read memory entries, it forward-passes the query vectors through the memory FFN and retrieves its corresponding value. Like RNNs, the memory state is still bounded in MNM.

Unlike the RNNs, the attention mechanism however doesn’t maintain a recurrent state and only performs a feed-forward computation on input sequence segment X s X_{s} : O s = a ​ t ​ t ​ e ​ n ​ t ​ i ​ o ​ n ​ ( X s ) . O_{s}=attention(X_{s}). (3)

The attention output O s O_{s} is simply passed to the next layer and no state is carried over to the next input sequence X s + 1 {X_{s+1}} at the same attention layer . In the attention layer, in order to capture the dependency between the consequent segments X s X_{s} and X s + 1 X_{s+1} , one needs to process them altogether at the same time and this process becomes a bottleneck requiring large computational resources as the length of input sequence grows more and more. To improve the efficiency while still being able to benefit from the expressiveness of the attention mechanism, this work introduces a recurrent attention layer .

## 3 Method

Figure 2 compares our model, Infini-Transformer, and Transformer-XL ( Dai et al., 2019 ) . Similar to Transformer-XL, Infini-Transformer operates on a sequence of segments. We compute the standard causal dot-product attention context within each segment. So the dot-product attention computation is local in a sense that it covers a total N N number of tokens of the current segment with index S S ( N N is the segment length).

The local attention ( Dai et al., 2019 ) , however, discards the attention states of the previous segment when processing the next one. In Infini-Transformers, instead of leaving out the old KV attention states, we propose to reuse them to maintain the entire context history with a compressive memory. So each attention layer of Infini-Transformers has both global compressive and local fine-grained states. We call such an efficient attention mechanism Infini-attention, which is illustrated in Figure 1 and described formally in the following sections.

### 3.1 Infini-attention

As shown Figure 1 , our Infini-attention is a recurrent attention mechanism that computes both local and global context states and combine them for its output. Similar to multi-head attention (MHA), it maintains H H number of parallel compressive memory per attention layer ( H H is the number of attention heads) in addition to the dot-product attention and like the RNNs and MNM, it maintains a recurrent memory state to efficiently track the long sequence context: O s , M s = i ​ n ​ f ​ i ​ n ​ i ​ - ​ a ​ t ​ t ​ e ​ n ​ t ​ i ​ o ​ n ​ ( X s , M s − 1 ) O_{s},M_{s}=infini\text{-}attention(X_{s},M_{s-1}) (4)

#### 3.1.1 Scaled Dot-product Attention

The multi-head scaled dot-product attention ( Vaswani et al., 2017 ) , specially its self-attention variant ( Munkhdalai et al., 2016 ; Cheng et al., 2016 ) , has been the main building block in LLMs. The MHA’s strong capability to model context-dependent dynamic computation and its conveniences of temporal masking have been leveraged extensively in the autoregressive generative models.

A single head in the vanilla MHA computes its attention context A d ​ o ​ t ∈ I ​ R N × d v ​ a ​ l ​ u ​ e A_{dot}\in{\rm I\!R}^{N\times d_{value}} from sequence of input segments X ∈ I ​ R N × d m ​ o ​ d ​ e ​ l X\in{\rm I\!R}^{N\times d_{model}} as follows. First, it computes attention query, key, and value states: K = X W K , V = X W V and Q = X W Q . K=XW_{K},\text{ }V=XW_{V}\text{ and }Q=XW_{Q}. (5) Here, W K ∈ I ​ R d m ​ o ​ d ​ e ​ l × d k ​ e ​ y , W V ∈ I ​ R d m ​ o ​ d ​ e ​ l × d v ​ a ​ l ​ u ​ e ​ and ​ W Q ∈ I ​ R d m ​ o ​ d ​ e ​ l × d k ​ e ​ y W_{K}\in{\rm I\!R}^{d_{model}\times d_{key}},W_{V}\in{\rm I\!R}^{d_{model}\times d_{value}}\text{ and }W_{Q}\in{\rm I\!R}^{d_{model}\times d_{key}} are trainable projection matrices. Then, the attention context is calculated as a weighted average of all other values as A d ​ o ​ t = softmax ​ ( Q ​ K T d m ​ o ​ d ​ e ​ l ) ​ V . A_{dot}=\text{softmax}\left(\frac{QK^{T}}{\sqrt{d_{model}}}\right)V. (6)

For MHA, we compute H H number of attention context vectors for each sequence element in parallel, concatenate them along the second dimension and then finally project the concatenated vector to the model space to obtain the attention output.

#### 3.1.2 Compressive Memory

In Infini-attention, instead of computing new memory entries for compressive memory, we reuse the query, key and value states ( Q Q , K K and V V ) from the dot-product attention computation. The state sharing and reusing between the dot-product attention and compressive memory not only enables efficient plug-in-play long-context adaptation but also speeds up training and inference. Similar to the prior work ( Munkhdalai et al., 2019 ) , our goal is to store bindings of key and value states in the compressive memory and retrieve by using the query vectors.

While there are different forms of compressive memory proposed in the literature ( Hopfield, 1982 ; Kanerva, 1988 ; Schlag et al., 2019 ; Munkhdalai et al., 2019 ) , for simplicity and computational efficiency, in this work we parameterize the memory with an associative matrix ( Schlag et al., 2020 ) . This approach further allows us to cast the memory update and retrieval process as linear attention mechanism ( Shen et al., 2018 ) and to leverage stable training techniques from the related methods. Specially, we adopt the update rule and retrieval mechanism by Katharopoulos et al. (2020) mainly due to its simplicity and competitive performance.

Memory retrieval. In Infini-attention, we retrieve new content A m ​ e ​ m ∈ I ​ R N × d v ​ a ​ l ​ u ​ e A_{mem}\in{\rm I\!R}^{N\times d_{value}} from the memory M s − 1 ∈ I ​ R d k ​ e ​ y × d v ​ a ​ l ​ u ​ e M_{s-1}\in{\rm I\!R}^{d_{key}\times d_{value}} by using the query Q ∈ I ​ R N × d k ​ e ​ y Q\in{\rm I\!R}^{N\times d_{key}} as: A m ​ e ​ m = σ ⁡ ( Q ) ​ M s − 1 σ ⁡ ( Q ) ​ z s − 1 . A_{mem}=\frac{\sigma({Q})M_{s-1}}{{\sigma(Q)}z_{s-1}}. (7)

Here, σ \sigma and z s − 1 ∈ I ​ R d k ​ e ​ y z_{s-1}\in{\rm I\!R}^{d_{key}} are a nonlinear activation function and a normalization term, respectively. As the choice of the non-linearity and the norm method is crucial for training stability, following Katharopoulos et al. (2020) we record a sum over all keys as the normalization term z s − 1 z_{s-1} and use element-wise ELU + 1 as the activation function ( Clevert et al., 2015 ) .

Memory update. Once the retrieval is done, we update the memory and the normalization term with the new KV entries and obtain the next states as M s ← M s − 1 + σ ​ ( K ) T ​ V ​ and ​ z s ← z s − 1 + ∑ t = 1 N σ ⁡ ( K t ) . M_{s}\leftarrow M_{s-1}+\sigma(K)^{T}V\text{ and }z_{s}\leftarrow z_{s-1}+\sum_{t=1}^{N}\sigma(K_{t}). (8) The new memory states M s M_{s} and z s z_{s} are then passed to the next segment S + 1 S+1 , building in a recurrence in each attention layer. The right side term σ ​ ( K ) T ​ V \sigma(K)^{T}V in Eq. ( 8 ) is known as an associative binding operator ( Smolensky, 1990 ; Hebb, 2005 ; Schlag et al., 2020 ) .

Inspired by the success of delta rule ( Munkhdalai et al., 2019 ; Schlag et al., 2020 ; Schlag et al., 2021 ) , we have also incorporated it into our Infini-attention. The delta rule attempts a slightly improved memory update by first retrieving existing value entries and subtracting them from the new values before applying the associative bindings as new update. M s ← M s − 1 + σ ​ ( K ) T ​ ( V − σ ⁡ ( K ) ​ M s − 1 σ ⁡ ( K ) ​ z s − 1 ) . M_{s}\leftarrow M_{s-1}+\sigma(K)^{T}(V-\frac{\sigma({K})M_{s-1}}{{\sigma(K)}z_{s-1}}). (9)

This update rule ( L ​ i ​ n ​ e ​ a ​ r + D ​ e ​ l ​ t ​ a Linear+Delta ) leaves the associative matrix unmodified if the KV binding already exists in the memory while still tracking the same normalization term as the former one ( L ​ i ​ n ​ e ​ a ​ r Linear ) for numerical stability.

Long-term context injection. We aggregate the local attention state A d ​ o ​ t A_{dot} and memory retrieved content A m ​ e ​ m A_{mem} via a learned gating scalar β \beta : A = sigmoid ​ ( β ) ⊙ A m ​ e ​ m + ( 1 − sigmoid ​ ( β ) ) ⊙ A d ​ o ​ t . A=\textit{sigmoid}(\beta)\odot A_{mem}+(1-\textit{sigmoid}(\beta))\odot A_{dot}. (10) This adds only a single scalar value as training parameter per head while allowing a learnable trade-off between the long-term and local information flows in the model ( Wu et al., 2022 ) .

Similar to the standard MHA, for the multi-head Infini-attention we compute H H number of context states in parallel, and concatenate and project them for the final attention output O ∈ I ​ R N × d m ​ o ​ d ​ e ​ l O\in{\rm I\!R}^{N\times d_{model}} : O = [ A 1 ; … ​ A H ] ​ W O O=[A^{1};\dotso A^{H}]W_{O} (11) where W O ∈ I ​ R H × d v ​ a ​ l ​ u ​ e × d m ​ o ​ d ​ e ​ l W_{O}\in{\rm I\!R}^{H\times d_{value}\times d_{model}} is trainable weights.

### 3.2 Memory and Effective Context Window