A failed experiment: Infini-Attention, and why we should keep trying?

TLDR: Infini-attention's performance gets worse as we increase the number of times we compress the memory, and to the best of our knowledge, ring attention, YaRN and rope scaling are still the best ways for extending a pretrained model to longer context length.

Section 0: Introduction

The context length of language models is one of the central attributes besides the model’s performance. Since the emergence of in-context learning, adding relevant information to the model’s input has become increasingly important. Thus, the context length rapidly increased from paragraphs (512 tokens with BERT/GPT-1) to pages (1024/2048 with GPT-2 and GPT-3 respectively) to books (128k of Claude) all the way to collections of books (1-10M tokens of Gemini). However, extending standard attention to such length remains challenging.

A small intro to Ring Attention: Ring Attention was first introduced by researchers from UC Berkeley in 2024 [link] (to the best of our knowledge). This engineering technique helps overcome memory limitations by performing self-attention and feedforward network computations in a blockwise fashion and distributing sequence dimensions across multiple devices, allowing concurrent computation and communication.

Even with Ring Attention, the number of GPUs required to train a Llama 3 8B on a 1-million-token context length with a batch size of 1 is 512 GPUs. As scaling laws have shown, there is a strong correlation between model size and its downstream performance, which means the bigger the model, the better (of course, both models should be well-trained). So we not only want a 1m context length, but we want a 1m context length on the biggest model (e.g., Llama 3 8B 405B). And there are only a few companies in existence that have the resources to do so.

Section 1: Reproduction Principles

We found the following rules helpful when implementing a new method and use it as guiding principles for a lot of our work:

Section 2: How does Infini-attention works

First results

Recaps