Please explain the following research article like I am a layperson: "" License: arXiv.org perpetual non-exclusive license
arXiv:2404.07143v1 [cs.CL] 10 Apr 2024
Leave No Context Behind: Efficient Infinite Context Transformers with Infini-attention
Tsendsuren Munkhdalai, Manaal Faruqui and Siddharth Gopal
Google <PRESIDIO_ANONYMIZED_EMAIL_ADDRESS>
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.

1Introduction
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.

Refer to caption
Figure 1:Infini-attention has an additional compressive memory with linear attention for processing infinitely long contexts. 
{
�
⁢
�
}
�
−
1
 and 
{
�
⁢
�
}
�
 are attention key and values for current and previous input segments, respectively and 
�
�
 the attention queries. PE denotes position embeddings.
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.
2Method
Refer to caption
Figure 2:Infini-Transformer (top) has an entire context history whereas Transformer-XL (bottom) discards old contexts since it caches the KV states for the last segment only.
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 
�
 number of tokens of the current segment with index 
�
 (
�
 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.

2.1Infini-attention
As shown Figure 1, our Infini-attention computes both local and global context states and combine them for its output. Similar to multi-head attention (MHA), it maintains 
�
 number of parallel compressive memory per attention layer (
�
 is the number of attention heads) in addition to the dot-product attention.

2.1.1Scaled 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 
�
�
⁢
�
⁢
�
∈
I
⁢
R
�
×
�
�
⁢
�
⁢
�
⁢
�
⁢
�
 from sequence of input segments 
�
∈
I
⁢
R
�
×
�
�
⁢
�
⁢
�
⁢
�
⁢
�
 as follows. First, it computes attention query, key, and value states:

�
=
�
⁢
�
�
,
 
⁢
�
=
�
⁢
�
�
⁢
 and 
⁢
�
=
�
⁢
�
�
.
(1)
Here, 
�
�
∈
I
⁢
R
�
�
⁢
�
⁢
�
⁢
�
⁢
�
×
�
�
⁢
�
⁢
�
,
�
�
∈
I
⁢
R
�
�
⁢
�
⁢
�
⁢
�
⁢
�
×
�
�
⁢
�
⁢
�
⁢
�
⁢
�
⁢
 and 
⁢
�
�
∈
I
⁢
R
�
�
⁢
�
⁢
�
⁢
�
⁢
�
×
�
�
⁢
�
⁢
�
 are trainable projection matrices. Then, the attention context is calculated as a weighted average of all other values as

�
�
⁢
�
⁢
�
=
softmax
⁢
(
�
⁢
�
�
�
�
⁢
�
⁢
�
⁢
�
⁢
�
)
⁢
�
.
(2)
For MHA, we compute 
�
 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 attention the output.

2.1.2Compressive Memory
In Infini-attention, instead of computing new memory entries for compressive memory, we reuse the query, key and value states (
�
, 
�
 and 
�
) 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 
�
�
⁢
�
⁢
�
∈
I
⁢
R
�
×
�
�
⁢
�
⁢
�
⁢
�
⁢
�
 from the memory 
�
�
−
1
∈
I
⁢
R
�
�
⁢
�
⁢
�
×
�
�
⁢
�
⁢
�
⁢
�
⁢
�
 by using the query 
�
∈
I
⁢
R
�
×
�
�
⁢
�
⁢
�
 as:

�
�
⁢
�
⁢
�
=
�
⁢
(
�
)
⁢
�
�
−
1
�
⁢
(
�
)
⁢
�
�
−
1
.
(3)
Here, 
�
 and 
�
�
−
1
∈
I
⁢
R
�
�
⁢
�
⁢
�
 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 
�
�
−
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

�
�
←
�
�
−
1
+
�
⁢
(
�
)
�
⁢
�
⁢
 and 
⁢
�
�
←
�
�
−
1
+
∑
�
=
1
�
�
⁢
(
�
�
)
.
(4)
The new memory states 
�
�
 and 
�
�
 are then passed to the next segment 
�
+
1
, building in a recurrence in each attention layer. The right side term 
�
⁢
(
�
)
�
⁢
�
 in Eq. (4) 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; 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.

�
�
←
�
�
−
1
+
�
⁢
(
�
)
�
⁢
(
�
−
�
⁢
(
�
)
⁢
�
�
−
1
�
⁢
(
�
)
⁢
�
�
−
1
)
.
(5)
This update rule (
�
⁢
�
⁢
�
⁢
�
⁢
�
⁢
�
+
�
⁢
�
⁢
�
⁢
�
⁢
�
) 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 (
�
⁢
�
⁢
�
⁢
�
⁢
�
⁢
�
) for numerical stability.

Long-term context injection. We aggregate the local attention state 
�
�
⁢
�
⁢
�
 and memory retrieved content 
�
�
⁢
�
⁢
�
 via a learned gating scalar 
�
:

�
=
𝑠𝑖𝑔𝑚𝑜𝑖𝑑
⁢
(
�
)
⊙
�
�
⁢
�
⁢
�
+
(
1
−
𝑠𝑖𝑔𝑚𝑜𝑖𝑑
⁢
(
�
)
)
⊙
�
�
⁢
�
⁢
�
.
(6)
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 
�
 number of context states in parallel, and concatenate and project them for the final attention output 
�
∈
I
⁢
R
�
×
�
�
⁢
�
⁢
�
⁢
�
⁢
�
:

�
=
[
�
1
;
…
⁢
�
�
]
⁢
�
�
(7)
where 
�
�
∈
I
⁢
R
�
×
�
�
⁢
�
⁢
�
⁢
�
⁢
�
×
�
�
⁢
�
⁢
�
⁢
�
⁢
�
 is trainable weights.

2.2Memory and Effective Context Window
Model	Memory (cache) footprint	Context length	Memory update	Memory retrieval
Transformer-XL	
(
�
�
⁢
�
⁢
�
+
�
�
⁢
�
⁢
�
⁢
�
⁢
�
)
×
�
×
�
×
�
�
×
�
Discarded	Dot-product attention
Compressive Transformer	
�
�
⁢
�
⁢
�
⁢
�
⁢
�
×
(
�
+
�
)
×
�
(
�
×
�
+
�
)
×
�
Discarded	Dot-product attention
Memorizing Transformers	
(
�
�
⁢
�
⁢
�
+
�
�
⁢
�
⁢
�
⁢
�
⁢
�
)
×
�
×
�
×
�
�
×
�
None	kNN + dot-product attention
RMT	
�
�
⁢
�
⁢
�
⁢
�
⁢
�
×
�
×
�
×
2
�
×
�
Discarded	Soft-prompt input
AutoCompressors	
�
�
⁢
�
⁢
�
⁢
�
⁢
�
×
�
×
(
�
+
1
)
×
�
�
×
�
Discarded	Soft-prompt input
Infini-Transformers	
�
�
⁢
�
⁢
�
×
(
�
�
⁢
�
⁢
�
⁢
�
⁢
�
+
1
)
×
�
×
�
�
×
�
Incremental	Linear attention
Table 1:Transformer models with segment-level memory are compared. For each model, the memory size and effective context length are defined in terms of their model parameters (
�
: input segment length, 
�
: the number of segments, 
�
: the number of layers, 
�
: the number of attention heads, 
�
: Compressive Transformer memory size, 
�
: compression ratio, 
�
: the number of soft-prompt summary vectors and 
�
: summary vector accumulation steps).
Our Infini-Transformer enables an unbounded context window with a bounded memory footprint. To illustrate this, Table 1 lists the previous segment-level memory models with their context-memory footprint and effective context length defined in terms of model parameters and input segment length. Infini-Transformer has a constant memory complexity of 
�
�
⁢
�
⁢
�
×
�
�
⁢
�
⁢
�
⁢
�
⁢
�
+
�
�
⁢
�
⁢
�
 for storing compressed context in 
�
�
 and 
�
�
 for each head in single layer while for the other models, the complexity grows along with the sequence dimension - the memory complexity depends either on the cache size for Transformer-XL (Dai et al., 2019), Compressive Transformer (Rae et al., 2019) and Memorizing Transformers (Wu et al., 2022) or on the soft-prompt size for RTM (Bulatov et al., 2022) and AutoCompressors (Ge et al., 2023).

Refer to caption
Figure 3:There are two types of heads emerged in Infini-attention after training: specialized heads with gating score near 0 or 1 and mixer heads with score close to 0.5. The specialized heads either process contextual information via the local attention mechanism or retrieve from the compressive memory whereas the mixer heads aggregate both current contextual information and long-term memory content together into single output.
Transformer-XL computes attention over KV states cached from the last segment in addition to the current states. Since this is done for each layer, Transformer-XL extends the context window from 
�
 to 
�
×
�
 tokens with an additional memory footprint of 
(
�
�
⁢
�
⁢
�
+
�
�
⁢
�
⁢
�
⁢
�
⁢
�
)
×
�
×
�
×
�
. Compressive Transformer adds a second cache to Transformer-XL and stores compressed representations of past segment activations. So it extends the Transformer-XL’s context window by 
�
×
�
×
�
 but still has a large context-memory complexity. Taking the idea further, Memorizing Transformers opt to store the entire KV states as context for input sequences. Since the storage becomes prohibitively expensive in this case, they restrict the contextual computation to a single layer only. By utilizing a fast kNN retriever, Memorizing Transformers then build a context window covering the entire sequence history of length 
�
×
�
 at an increased cost of storage. Our experiments show that Infini-Transformer LM can achieve more than 100x compression rate on top of Memorizing Transformers while further improving the perplexity score.

RMT and AutoCompressors allow for a potentially infinite context length since they compress the input into summary vectors and then pass them as extra soft-prompt inputs for the subsequent segments. However, in practice the success of those techniques highly depends on the size of soft-prompt vectors. Namely, it is necessary to increase the number of soft-prompt (summary) vectors to achieve a better performance with AutoCompressors (Chevalier et al., 2023) and with that, the memory and compute complexity grow quickly resulting in diminished efficiency. It was also observed in AutoCompressors (Chevalier et al., 2023) that an efficient compression objective is needed for training such prompt compression techniques (Ge et al., 2023).

3Experiments
We evaluated our Infini-Transformer models on benchmarks involving extremely long input sequences: long-context language modeling, 1M length passkey context block retrieval and 500K length book summarization tasks. For the language modeling benchmark, we train our models from scratch while for the passkey and book summarization tasks, we continually pre-train existing LLMs in order to highlight a plug-and-play long-context adaptation capability of our approach.

3.1Long-context Language Modeling
We trained and evaluated small Infini-Transformer models on PG19 (Rae et al., 2019) and Arxiv-math (Wu et al., 2022) benchmarks. Our setup closely resembles that of Memorizing Transformers (Wu et al., 2022). Namely, all our models have 12 layers and 8 attention heads of dimension 128 each and FFNs with hidden layer 4096.

Model	Memory size (comp.)	XL cache	Segment length	PG19	Arxiv-math
Transformer-XL	50M (3.7x)	2048	2048	11.88	2.42
Memorizing Transformers	183M (1x)	2048	2048	11.37	2.26
RMT	2.5M (73x)	None	2048	13.27	2.55
Infini-Transformer (Linear)	1.6M (114x)	None	2048	9.65	2.24
Infini-Transformer (Linear + Delta)	1.6M (114x)	None	2048	9.67	2.23
Table 2:Long-context language modeling results are compared in terms of average token-level perplexity. Comp. denotes compression ratio. Infini-Transformer outperforms memorizing transformers with memory length of 65K and achieves 114x compression ratio.
We set the Infini-attention segment length 
�
 to 2048 for all attention layers and the input sequence length to 32768 for training. This allows the Infini-attention to unroll over 16 steps w.r.t its compressive memory states. For the RMT baseline, we performed several runs with summary prompt lengths 50, 100 and 150 and sequence lengths 4096, 8196 and 32768. RMT with 100 summary vectors gave the best result when trained on 8196 length sequences.

The main results from the language modeling experiments are summarized in Table 2. Our Infini-Transformer outperforms both Transformer-XL (Dai et al., 2019) and Memorizing Transformers (Wu et al., 2022) baselines while maintaining 114x less memory parameters than the Memorizing Transformer model with a vector retrieval-based KV memory with length of 65K at its 
9
�
⁢
ℎ
 layer.

Zero-shot
32K	128K	256K	512K	1M
Infini-Transformer (Linear)	14/13/98	11/14/100	6/3/100	6/7/99	8/6/98
Infini-Transformer (Linear + Delta)	13/11/99	6/9/99	7/5/99	6/8/97	7/6/97
FT (400 steps)
Infini-Transformer (Linear)	100/100/100	100/100/100	100/100/100	97/99/100	96/94/100
Infini-Transformer (Linear + Delta)	100/100/100	100/100/99	100/100/99	100/100/100	100/100/100
Table 3:Infini-Transformers solved the passkey task with up to 1M context length when fine-tuned on 5K length inputs. We report token-level retrieval accuracy for passkeys hidden in a different part (start/middle/end) of long inputs with lengths 32K to 1M.
100K length training. We further increased the training sequence length to 100K from 32K and trained the models on Arxiv-math dataset. 100K training further decreased the perplexity score to 2.21 and 2.20 for 
�
⁢
�
⁢
�
⁢
�
⁢
�
⁢
�
 and 
�
⁢
�
⁢
�
⁢
�
⁢
�
⁢
�
+
�
⁢
�
⁢
�
⁢
�
⁢
�
 models.

Gating score visualization. Figure 3 visualizes the gating score, 
𝑠𝑖𝑔𝑚𝑜𝑖𝑑
⁢
(
�
)
 for the compressive memory for all attention heads in each layer. There are two types of heads emerged in Infini-attention after training: specialized heads with a gating score near 0 or 1 and mixer heads with a score close to 0.5. The specialize"""