2024-04-15: Infini-Attention

🔷 Subscribe to get breakdowns of the most important developments in AI in your inbox every morning.

Here’s today at a glance:

• Infini-Attention

• AI artwork of the day

⚠️ Infini-Attention

Google designed a memory system for transformers that allows storing of state. This compressive memory system scales better for very long sequences. A compressive memory 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 the objective that this information can be recovered back later on.

What’s happening is that a summary of the information to date (a compressed version of it) is kept in the memory store, allowing the model to not compute anew for each successive token. This potentially breaks transformers out of the quadratic increase in computation with increasing length problem.

This system is similar in some ways to the other non-Transformer architectures, such as Mamba, RWKW, etc.

Anyway, we should all be thankful, Google is publishing again!

Paper Title:

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

Who:

• Researchers at Google: Tsendsuren Munkhdalai, Manaal Faruqui, and Siddharth Gopal.

Why:

• Motivation: Traditional Transformer-based Large Language Models (LLMs) have limited context windows due to the quadratic complexity of the attention mechanism, making it challenging to handle long sequences efficiently.

• Gaps addressed: The research aimed to address the limitations of existing LLMs in processing long sequences and the high memory and computational costs associated with scaling them.

• Impact: This study has the potential to enable LLMs to efficiently handle much longer contexts, opening up possibilities for various applications, including long-form document understanding, summarization, and question answering.

How:

• Infini-attention: The researchers introduce a novel attention mechanism called Infini-attention that combines a compressive memory with the standard attention mechanism.

• Compressive Memory: This memory efficiently stores key-value pairs from previous segments of the input sequence, allowing the model to access a much larger context window without significantly increasing memory usage.

• Local and Global Attention: Infini-attention utilizes both local attention within a segment and global attention across segments retrieved from the compressive memory, enabling the model to capture both short and long-range dependencies effectively.

• Datasets: The researchers evaluated their approach on long-context language modeling benchmarks like PG19 and Arxiv-math, as well as tasks like passkey retrieval and book summarization.

• Models: They experimented with different sizes of LLMs (1B and 8B parameters) and compared their approach with baseline models like Transformer-XL and Memorizing Transformers.

What did they find:

• Improved Performance: Infini-Transformer models achieved better perplexity scores O(n) long-context language modeling benchmarks compared to baseline models, demonstrating their effectiveness in capturing long-range dependencies.

• Memory Efficiency: The compressive memory in Infini-attention enabled significant memory savings, achieving up to 114x compression ratio compared to models with explicit memory storage.

• Length Generalization: Infini-Transformers exhibited impressive length generalization capabilities, successfully solving tasks with input lengths of up to 1 million tokens even when trained on shorter sequences.

• State-of-the-art Summarization: A 8B LLM with Infini-attention achieved new state-of-the-art results on the BookSum dataset for long-form book summarization, highlighting the potential for real-world applications.

What are the limitations and what's next:

• Limitations: The paper does not provide a detailed analysis of the computational complexity of Infini-attention during training and inference, which could be important for practical deployment.

• Future Research: Exploring alternative compressive memory techniques, investigating the trade-off between memory compression and model performance, and applying Infini-attention to other tasks beyond language modeling.

Why it matters:

• Significant Contribution: This research proposes a novel and efficient approach to enable LLMs to handle infinitely long contexts with bounded memory and computation, addressing a critical limitation of current models.

• Potential Impact: Infini-attention has the potential to revolutionize the field of natural language processing by allowing LLMs to process and understand much longer sequences of text, leading to advancements in various applications such as document summarization, question answering, and machine translation.

Additional notes:

• This research paper has been submitted to a conference or journal for publication.

• The code and models are not yet publicly available.

Share this story

🖼️ AI Artwork Of The Day

That’s it for today! Become a subscriber for daily breakdowns of what’s happening in the AI world: