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

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

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 d k ​ e ​ y × d v ​ a ​ l ​ u ​ e + d k ​ e ​ y d_{key}\times d_{value}+d_{key} for storing compressed context in M s M_{s} and z s z_{s} 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 RMT ( Bulatov et al., 2022 ) and AutoCompressors ( Ge et al., 2023 ) .

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 N N to N × l N\times l tokens with an additional memory footprint of ( d k ​ e ​ y + d v ​ a ​ l ​ u ​ e ) × H × N × l (d_{key}+d_{value})\times H\times N\times l . 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 c × r × l c\times r\times l 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 N × S N\times S 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 ) .

## 4 Experiments

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.

### 4.1 Implementation details

Segment chunking. we forward-pass the entire input text a Transformer model and then perform segment chunking at each Infini-attention layer - in this way, perform a minimal modification to the existing Transformer implementation. The Infini-attention layer segments the input and process it segment by segment and concatenates back the segments to pass the original-length segment as output to the next layer.

Back-propagation through time (BPTT). Each Infini-attention layer is trained with back-propagation through time ( Werbos, 1988 ) by computing the gradient w.r.t the compressive memory states, similar to how RNNs are trained. To save memory, we perform gradient checkpoint when processing the sequence segment by segment.

Position Embeddings (PE). As shown Figure 1 , we don’t use position embeddings for the key and query vectors of the compressive memory to store only global contextual information in the long-term memory. The PEs were applied to the QK vectors only after the compressive memory reading and update.

### 4.2 Long-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.

We set the Infini-attention segment length N N 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 t ​ h 9^{th} layer.

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 L ​ i ​ n ​ e ​ a ​ r Linear and L ​ i ​ n ​ e ​ a ​ r + D ​ e ​ l ​ t ​ a Linear+Delta models.

Gating score visualization. Figure 3 visualizes the gating score, sigmoid ​ ( β ) \textit{sigmoid}(\beta) 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 specialized heads either process contextual information via the local attention computation or retrieve from the compressive memory whereas the mixer heads aggregate both current contextual information and long-term memory content together into a single output. Interestingly, each layer has at least a single short-range head, allowing a forward-propagation of input signal up until the output layer. We also observed an interleaving of long and short-term content retrievals throughout the forward computation.

### 4.3 LLM Continual Pre-training

We performed a lightweight continual pre-training for long-context adaptation of existing LLMs. The pre-training data includes the PG19 and Arxiv-math corpus as well as C4 text ( Raffel et al., 2020 ) with length more than 4K tokens. The segment length N N was set to 2K throughout our experiments.

1M passkey retrieval benchmark. We replaced the vanilla MHA in a 1B LLM with Infini-attention and continued to pre-train on inputs with length of 4K. The model was trained for 30K steps with batch size of 64 before fine-tuning on the passkey retrieval task ( Mohtashami & Jaggi, 2024 ) .

The passkey task hides a random number into a long text and asks it back at the model output. The length of the distraction text is varied by repeating a text chunk multiple times. The previous work ( Chen et al., 2023a ) showed that a 8B LLaMA model can solve the task up to 32K length when fine-tuned with the same 32K length inputs with Position Interpolation. We take this challenge further and fine-tune on only 5K length inputs to test on 1M length regime.

Table 3 reports the token-level accuracy for test subsets with input lengths ranging from 32K to 1M. For each test subset, we controlled the position of the passkey so that it is either located around the beginning, middle or the end of the input sequence. We reported both zero-shot accuracy and fine-tuning accuracy. Infini-Transformers solved the task with up to 1M context length after fine-tuning on 5K length inputs for 400 steps.

500K length book summarization (BookSum). We further scaled our approach by continuously pre-training a 8B LLM model with 8K input length for 30K steps. We then fine-tuned on a book summarization task, BookSum ( Kryściński et al., 2021 ) where the goal is to generate a summary of an entire book text.

We set the input length to 32K for fine-tuning and increase to 500K for evaluating. We use a generation temperature of 0.5 and t ​ o ​ p p = 0.95 top_{p}=0.95 and set the number of decoding steps to 1024 to generate a summary of each book.

Table 4 compares our model against the encoder-decoder models that were built particularly for the summarization task ( Lewis et al., 2019 ; Xiao et al., 2021 ) and their retrieval-based long-context extension ( Bertsch et al., 2024 ) . Our model outperforms the previous best results and achieves a new SOTA on BookSum by processing the entire text from book. We have also plotted the overall Rouge score on validation split of BookSum data in Figure 4 . There is a clear trend showing that with more text provided as input from books, Our Infini-Transformers improves its summarization performance metric.

## 5 Related Work

Compressive memory. Inspired by the plasticity in biological neurons ( Munkhdalai & Yu, 2017a ; Miconi et al., 2018 ) , compressive memory approaches cast parameterized functions as memory to store and retrieve information ( Hinton & Plaut, 1987 ; Schmidhuber, 1992 ; Ba et al., 2016 ; Munkhdalai et al., 2019 ) . Unlike the Transformer KV memory array ( Vaswani et al., 2017 ; Wu et al., 2022 ) , which grows with input sequence length, compressive memory systems maintain a constant number of memory parameters for computational efficiency. The parameters are modified with an update rule to store information, which is then retrieved via a memory reading mechanism ( Graves et al., 2014 ; Sukhbaatar et al., 2015 ; Munkhdalai & Yu, 2017b ) .

Compressed input representations can be viewed as a summary of past sequence segments ( Rae et al., 2019 ; Chevalier et al., 2023 ) . Along this direction, more recent works have been utilizing a Transformer LLM itself to compress input sequence for efficient long-context modeling ( Bulatov et al., 2022 ; Chevalier et al., 2023 ; Ge et al., 2023 ; Mu et al., 2024 ; Hwang et al., 2024 ) . However, the previous segment-level compression methods, including Compressive Transformers ( Rae et al., 2019 ) still discard the memory entries of old segments in order to free up space for the new ones, limiting their context window to the most recent segments. This is in contrast to our Infini-attention that computes incremental memory updates to a fixed amount of memory parameters in a recurrent fashion.

Long-context continual pre-training. There is a line of work that extends the dot-product attention layers and continues to train LLMs for long-context ( Xiong et al., 2023 ; Fu et al., 2024 ) . The attention extensions include incorporating sparsity into the attention layer ( Chen et al., 2023b ; Ratner et al., 2022 ; Mohtashami & Jaggi, 2024 ) as well as manipulating the position encodings ( Chen et al., 2023a ; Peng et al., 2023 ) . Although the position encoding-based methods such as position interpolation techniques ( Chen et al., 2023a ) can be data efficient as they only adjust the positional bias in the attention layer, they are still costly for inference.

The attention mechanism is also prone to the issues of attention sink ( Xiao et al., 2023 ) and lost-in-the-middle ( Liu et al., 2024 ) . Consequently, they struggle in a regime where context length is longer than what was observed during training ( Press et al., 2021 ; Kazemnejad et al., 2024 ) . The proposed Infini-attention addresses those issues by enabling a segment-level streaming computation over long sequences with a fixed local attention window. Our Infini-Transformers successfully extrapolate to 1M input length regimes when trained on 32K and even 5K length sequences.

Efficient attention. The efficient attention techniques attempt to improve the efficiency of the dot-product attention with an approximation or a system-level optimization. Multiple directions have been explored for different forms of efficient attention approximation, including sparsity-based ( Child et al., 2019 ; Beltagy et al., 2020 ; Sukhbaatar et al., 2021 ; Ding et al., 2023 ; Xiao et al., 2024 ) and linear attention approximation ( Shen et al., 2018 ; Katharopoulos et al., 2020 ; Schlag et al., 2021 ) . Among those, the linear attention variants are closely related to the associative memory matrix ( Schlag et al., 2020 ; Schlag et al., 2021 ) and the metalearned neural memory ( Munkhdalai et al., 2019 ) , where KV bindings ( Smolensky, 1990 ) are stored in Fast-Weights ( Hinton & Plaut, 1987 ; Schmidhuber, 1992 ; Ba et al., 2016 ) that are modified in with respect to new contextual information. More recently, system-level optimization techniques have been proposed by leveraging specific hardware architecture to make the exact attention computation more efficient ( Dao et al., 2022 ; Liu et al., 2023 ) .

## 6 Conclusion

An effective memory system is crucial not just for comprehending long contexts with LLMs, but also for reasoning, planning, continual adaptation for fresh knowledge, and even for learning how to learn. This work introduces a close integration of compressive memory module into the vanilla dot-product attention layer. This subtle but critical modification to the attention layer enables LLMs to process infinitely long contexts with bounded memory and computation resources. We show that our approach can naturally scale to a million length regime of input sequences, while outperforming the baselines on long-context language modeling benchmark and book summarization tasks. We also demonstrate a promising length generalization capability of our approach. 1B model that was fine-tuned on up to 5K sequence length passkey instances solved the 1M length problem.

#### Acknowledgments

We would like to thank Dongseong Hwang for their help implementing efficient sequence unrolling mechanism with the jax scan function. We would also like to thank Aditya Gupta, Kalpesh Krishna, Tu Vu and Alexandra Chronopoulou for their feedback.

## Appendix A Additional Training Details

For the long-context language modeling task, we set the learning rate to 0.01 by performing small search over values of 0.003, 0.005, 0.01 and 0.03. We used the Adafactor optimizer ( Shazeer & Stern, 2018 ) with linear warmup with 1000 steps, followed by cosine decay. We applied gradient checkpointing after each segment to save to save memory. The batch size was set to 64. For the LLM experiments, we set the learning rate to 0.0001 during continual pre-training and task fine-tuning.

## Appendix B Passkey Retrieval Task

Below we showed the input format of the passkey task.

There is an important info hidden inside a lot of irrelevant text. Find it and memorize them. I will quiz you about the important information there. The grass is green. The sky is blue. The sun is yellow. Here we go. There and back again. (repeat x times) The pass key is 9054 . Remember it. 9054 is the pass key. The grass is green. The sky is blue. The sun is yellow. Here we go. There and ack again. (repeat y times) What is the pass key? The pass key is

## Instructions for reporting errors

We are continuing to improve HTML versions of papers, and your feedback helps enhance accessibility and mobile support. To report errors in the HTML that will help us improve conversion and rendering, choose any of the methods listed below: