Source-linked AI summary
Leave No Context Behind: Efficient Infinite Context Transformers with Infini-attention
Tsendsuren Munkhdalai, Manaal Faruqui, Siddharth Gopal
TL;DR
Transformers face growing memory and computation costs on long inputs, and practical compressive memory for LLMs remains limited. The paper introduces Infini-attention, which integrates compressive memory with local and long-term attention in a single block. It reports strong long-context results, including million-token processing and state-of-the-art 500K-length book summarization, with bounded memory.
Problem
Transformer context-dependent memory and attention costs limit efficient processing of extremely long sequences, while effective practical compressive memory for LLMs remains lacking.
Method
Infini-attention integrates long-term compressive memory and local causal attention into vanilla attention, reusing prior key-value states for recurrent context processing.
Results
The approach outperforms baselines on long-context language modeling, scales a 1B LLM to 1M-token passkey retrieval, and achieves a new SOTA result on 500K-length book summarization with an 8B model.
Takeaways & Limitations
Infini-attention enables Transformer LLMs to process infinitely long contexts with bounded memory and computation through streaming processing.
Takeaways & Limitations
Related soft-prompt compression methods can require more summary vectors for better performance, causing memory and computation complexity to grow quickly.
Abstract
from arXiv · showhide
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
Transformer attention becomes costly for very long inputs because its memory and computation grow with sequence length, while existing LLMs lack a simple effective compressive memory. Infini-attention addresses this by combining local attention with long-term compressive memory, enabling bounded-resource processing of infinitely long contexts.
- 3TB of attention KV states can be required for a 500B model with batch size 512 and context length 2048.
- Infini-attention adds compressive memory to vanilla attention while combining masked local attention and long-term linear attention in one Transformer block.
- Old key-value states are stored in compressive memory and later retrieved with queries, then combined with local attention contexts.
- A 1B LLM reaches 1M sequence length for passkey retrieval, while an 8B model reaches a new SOTA result on 500K-length book summarization after adaptation.
- The approach supports plug-and-play continual pre-training and long-context adaptation with minimal changes to standard scaled dot-product attention.
2 Background
RNNs retain fixed-size recurrent states but struggle to preserve entire long-sequence context, whereas attention processes segments without carrying state between them. This makes long-range dependencies expensive and motivates recurrent attention with persistent memory.
- RNNs process one token at a time while maintaining a fixed-size hidden state representing the input sequence.
- External-memory RNN extensions read and write bounded memory using query, key, and value vectors.
- Infini-Transformer retains prior key-value states through compressive memory rather than discarding them after each segment.
- Standard attention performs feed-forward computation on one segment and carries no state to the next segment.
- Capturing dependencies between consecutive segments requires processing them together, increasing computational demands as input length grows.
3 Method
Infini-attention extends segment-based Transformer attention by retaining long-term information in compressive memory while preserving local fine-grained attention. It combines recurrent memory updates, memory retrieval, and learned gating within each attention layer, yielding bounded memory complexity and an unbounded context window.
- Infini-attention design: Infini-attention computes local causal attention within each segment while reusing previous segments’ KV states through compressive memory.Each layer retains both global compressive states and local fine-grained states instead of discarding earlier KV states.
- Memory operations: The method reuses standard attention query, key, and value states for memory consolidation and retrieves long-term values using subsequent query states.The final contextual output aggregates memory-retrieved values with local attention contexts.
- Memory operations: Infini-attention parameterizes memory with an associative matrix, casting memory update and retrieval as a linear attention mechanism.The implementation also supports a delta-rule update that leaves existing key-value bindings unchanged while maintaining normalization for numerical stability.
- Long-term context injection: A learned gating scalar β combines local attention state Adot with memory-retrieved content Amem, allowing a trade-off between local and long-term information flows.The design adds a single scalar training parameter per attention head.
- Memory and effective context window: Infini-Transformer maintains an unbounded context window with constant memory complexity of dkey × dvalue + dkey per head in a single layer.This compressed-memory footprint contrasts with segment-memory approaches whose complexity grows with cache size, compression settings, or soft-prompt size.
- Memory and effective context window: Infini-attention produces specialized heads that favor local attention or compressive-memory retrieval, alongside mixer heads that combine both sources.Specialized heads have gating scores near 0 or 1, while mixer heads have scores close to 0.5.
4 Experiments
The experiments evaluate Infini-Transformers on long-context language modeling, 1M-length passkey retrieval, and 500K-length book summarization, using both models trained from scratch and continual adaptation of existing LLMs. Across these tasks, the method combines strong performance with bounded memory and long-context generalization.
- Long-context language modeling: Infini-Transformer outperforms Transformer-XL and Memorizing Transformer baselines while using 114x fewer memory parameters than the latter.The comparison uses long-context language modeling results in Table 2.
- Long-context language modeling: Perplexity decreases to 2.21 and 2.20 for the Linear and Linear + Delta models when training length increases to 100K.These scores are reported on the Arxiv-math dataset.
- 1M passkey retrieval: Infini-Transformers solve passkey retrieval at up to 1M context length after fine-tuning on only 5K-length inputs for 400 steps.Accuracy is evaluated for passkeys placed near the start, middle, and end of inputs ranging from 32K to 1M tokens.
- 500K book summarization: BookSum Rouge performance improves as more book text is provided as input.The trend is shown on the validation split.
5 Related Work
Related work addresses long-context modeling through compressive memory, continual pre-training, sparsity, position encodings, and efficient attention approximations. Infini-attention differs by retaining long-term information in fixed-size memory while preserving local attention over segments.
- Compressive memory: Compressive memory maintains a constant number of parameters, unlike Transformer KV memory arrays whose size grows with input length.It stores and recalls information by changing parameterized memory states.
- Compressive memory: Earlier segment-level compression methods discard old memory entries to free space for new ones, limiting retained context.This limitation motivates retaining information through the proposed compressive memory.
- Long-context continual pre-training: Long-context extensions include sparse attention and position-encoding manipulation, but position interpolation remains costly during inference.These methods continue training or modify dot-product attention layers for longer contexts.
- Long-context continual pre-training: Infini-attention addresses attention-sink and lost-in-the-middle issues through segment-level streaming computation with a fixed local attention window.The cited passage connects these issues to difficulty beyond training-time context lengths.
- Efficient attention: Efficient attention research explores approximation and system-level optimization, including sparsity-based and linear-attention approaches.These directions target improved efficiency over dot-product attention.
6 Conclusion
The paper concludes that integrating compressive memory into vanilla dot-product attention enables bounded-resource processing of extremely long contexts. Experiments report million-token scaling, benchmark improvements, and length generalization, including 1M passkey retrieval after 5K-length fine-tuning.
- Conclusion: Infini-attention integrates compressive memory into vanilla dot-product attention to process infinitely long contexts with bounded memory and computation.The conclusion presents this as the paper’s central modification and capability.
- Conclusion: The approach scales to million-token inputs while outperforming baselines on long-context language modeling and book summarization tasks.The conclusion reports this pattern across the evaluated tasks.
- Conclusion: A 1B model fine-tuned on passkey instances up to 5K tokens solves the 1M-length passkey problem.This is presented as evidence of length generalization.
A Additional Training Details
The appendix specifies optimization settings for long-context language modeling and LLM experiments, including learning rates, optimizer scheduling, checkpointing, and batch size.
- Optimization: Long-context language modeling uses a 0.01 learning rate selected from four searched values and the Adafactor optimizer.Training applies 1000-step linear warmup followed by cosine decay.
- Memory and LLM training: Gradient checkpointing is applied after each segment to reduce memory use, with batch size 64.The same passage gives a 0.0001 learning rate for LLM continual pre-training and task fine-tuning.
B Passkey Retrieval Task
The passkey task hides important information within repeated irrelevant text and asks the model to retrieve it. The example presents the passkey as 9054 and later quizzes the model for it.
- The section introduces the input format used for the passkey task.
- The task embeds a passkey within large amounts of irrelevant repeated text.The passage contrasts important information with repeated statements about grass, sky, and the sun.
- The model must remember the hidden passkey and answer a later question asking for it.The example states the passkey is 9054, repeats it, and ends with a question requesting the passkey.