Как работает Infini-attention (Google, 2024) для бесконечного контекста?

Краткий тезис

Infini-attention — это модификация механизма внимания в Transformer, которая позволяет обрабатывать теоретически бесконечные последовательности, сохраняя линейную вычислительную сложность. Основная идея — комбинировать стандартное dot-product attention для последнего сегмента токенов с компрессированной памятью (memory|compressed memory), которая агрегирует информацию из всех предыдущих сегментов через ассоциативное обновление. Это даёт возможность работать с контекстами длиной более 1 миллиона токенов без квадратичного роста затрат.

1. Проблема: квадратичная сложность стандартного внимания

Стандартный Transformer использует scaled dot-product attention, где для последовательности длины (L) вычисляется матрица внимания размером (L \times L). Это даёт сложность (O(L^2)) по времени и памяти. При (L = 1M) это становится непрактичным (терабайты памяти). Кроме того, модели имеют фиксированный максимальный контекст (например, 4K, 8K токенов), что ограничивает обработку длинных документов, диалогов или истории агента.

Термин «контекстное окно» — максимальное количество токенов, которое модель может «увидеть» за один проход.

2. Основная идея Infini-attention

Infini-attention (предложена Google DeepMind в 2024 году) решает проблему, разделяя входную последовательность на сегменты (segments) фиксированной длины (S). Для каждого сегмента модель вычисляет:

• Стандартное внимание — только внутри текущего сегмента (локальное внимание).
• Сжатое представление — информация из всех предыдущих сегментов хранится в компрессированной памяти (memory|compressed memory) в виде ключей и значений, обновляемых через ассоциативную память (memory|associative memory).

Таким образом, модель может «вспомнить» любой предыдущий токен, обращаясь к памяти, а не пересчитывая всё заново. Сложность становится (O(L \cdot S)) (линейной по длине), так как каждый сегмент обрабатывается независимо, а память имеет фиксированный размер.

3. Архитектура: сегменты и память

Пусть входная последовательность разбита на сегменты (X_1, X_2, \dots, X_T), каждый длины (S). Для каждого слоя Transformer (с Infini-attention) алгоритм следующий:

1. Последний сегмент (X_t) обрабатывается обычным dot-product attention: [ [text](/wiki/text){Attn}(Q_t, K_t, V_t) = [text](/wiki/text){softmax}\left(\frac{Q_t K_t^T}{\sqrt{d}}\right) V_t ] где (Q_t, K_t, V_t) — запросы, ключи, значения для этого сегмента.

2. Предыдущие сегменты представлены в виде компрессированной памяти — пары матриц ((K_{[text](/wiki/text){mem}}, V_{[text](/wiki/text){mem}})) фиксированного размера (M \times d) (где (M) — число слотов памяти, гиперпараметр). Память инициализируется нулями и обновляется после каждого сегмента.

3. Выход слоя для сегмента (t) — сумма двух вкладов: [ [text](/wiki/text){Output}t = [text](/wiki/text){Attn}(Q_t, K_t, V_t) + [text](/wiki/text){MemAttn}(Q_t, K{[text](/wiki/text){mem}}, V_{[text](/wiki/text){mem}}) ] где ([text](/wiki/text){MemAttn}) — внимание к памяти (также dot-product, но с памятью фиксированного размера).

Термин «слот памяти» — один элемент в компрессированной памяти, который хранит агрегированную информацию о множестве предыдущих токенов.

4. Обновление памяти: ассоциативное обучение

Память обновляется после обработки каждого сегмента (X_t) с использованием ассоциативного правила (associative memory update), напоминающего Hebbian learning:

[ K_{[text](/wiki/text){mem}} \leftarrow K_{[text](/wiki/text){mem}} + \sigma(K_t)^T \cdot [text](/wiki/text){gate} ] [ V_{[text](/wiki/text){mem}} \leftarrow V_{[text](/wiki/text){mem}} + \sigma(K_t)^T \cdot V_t ]

Здесь (\sigma) — нелинейная функция (например, softmax по строкам), а gate — механизм забывания (learnable gating). Фактически, память накапливает взвешенные суммы ключей и значений, что позволяет хранить информацию о большом количестве токенов в ограниченном числе слотов.

Почему это линейно Размер памяти (M) фиксирован (например, 512 слотов), поэтому сложность обновления и чтения памяти — (O(M \cdot S)), а не (O(L^2)). Общая сложность на всю последовательность: (O(L \cdot S + L \cdot M)), что линейно по (L).

5. Математическая формулировка Infini-attention

Для сегмента (t):

1. Вычисляем локальное внимание: [ A_{[text](/wiki/text){local}} = [text](/wiki/text){softmax}\left(\frac{Q_t K_t^T}{\sqrt{d}}\right) V_t ]

2. Вычисляем внимание к памяти: [ A_{[text](/wiki/text){mem}} = [text](/wiki/text){softmax}\left(\frac{Q_t K_{[text](/wiki/text){mem}}^T}{\sqrt{d}}\right) V_{[text](/wiki/text){mem}} ]

3. Выход слоя: [ [text](/wiki/text){Output}t = [text](/wiki/text){LayerNorm}(A{[text](/wiki/text){local}} + A_{[text](/wiki/text){mem}}) ]

4. Обновление памяти (после сегмента): [ K_{[text](/wiki/text){mem}} \leftarrow \gamma \cdot K_{[text](/wiki/text){mem}} + (1-\gamma) \cdot [text](/wiki/text){proj}K([text](/wiki/text){mean}(K_t)) ] [ V{[text](/wiki/text){mem}} \leftarrow \gamma \cdot V_{[text](/wiki/text){mem}} + (1-\gamma) \cdot [text](/wiki/text){proj}_V([text](/wiki/text){mean}(V_t)) ] где (\gamma) — learnable gate (близкий к 1 для долгой памяти), а ([text](/wiki/text){proj}) — линейные проекции для согласования размерностей.

На практике авторы используют softmax над ключами памяти и gated averaging для обновления.

6. Преимущества Infini-attention

Ключевые плюсы:

• Бесконечный контекст — модель может обрабатывать последовательности любой длины, не теряя информацию из начала.
• Линейная сложность — позволяет работать с миллионами токенов на одном GPU.
• Совместимость — Infini-attention может быть встроена в существующие LLM (например, Mistral, Llama) путём замены слоёв внимания и дообучения на длинных последовательностях.

7. Ограничения и практические аспекты

• Сжатие информации — память фиксированного размера не может хранить все детали; возможна потеря редких, но важных фактов.
• Забывание — gating может привести к тому, что старая информация перезаписывается (как в LSTM).
• Практический предел — в статье показана работа до 1M токенов; для больших длин может потребоваться увеличение числа слотов памяти или иерархическая память.
• Обучение — требуется специальный протокол обучения с сегментами и обратным распространением через память (трюк с recurrent backpropagation или truncated BPTT).

8. Сравнение с другими подходами к длинному контексту

Infini-attention выигрывает за счёт явного сжатия всей истории в фиксированную память, в то время как Transformer-XL хранит все предыдущие сегменты (что всё ещё квадратично по числу сегментов), а sparse attention не покрывает все токены.

9. Псевдокод реализации (упрощённо)

class InfiniAttention(nn.Module):
    def __init__(self, d_model, num_heads, segment_size, memory_slots):
        super().__init__()
        self.segment_size = segment_size
        self.mem_k = nn.Parameter(torch.zeros(memory_slots, d_model))
        self.mem_v = nn.Parameter(torch.zeros(memory_slots, d_model))
        self.gate = nn.Parameter(torch.tensor(0.9))
        # проекции для Q, K, V
        self.wq, self.wk, self.wv = nn.Linear(d_model, d_model), ...

    def forward(self, x):
        # x: [batch, seq_len, d_model]
        segments = x.split(self.segment_size, dim=1)
        outputs = []
        for seg in segments:
            Q = self.wq(seg)
            K = self.wk(seg)
            V = self.wv(seg)
            # локальное внимание
            local_attn = scaled_dot_product(Q, K, V)
            # внимание к памяти
            mem_attn = scaled_dot_product(Q, self.mem_k, self.mem_v)
            out = local_attn + mem_attn
            outputs.append(out)
            # обновление памяти (упрощённо)
            with torch.no_grad():
                self.mem_k.data = self.gate * self.mem_k + (1-self.gate) * K.mean(dim=1)
                self.mem_v.data = self.gate * self.mem_v + (1-self.gate) * V.mean(dim=1)
        return torch.cat(outputs, dim=1)

Примечание: в реальной реализации используются multi-head, нормализация и learnable gating с сигмоидой.

10. Применение в Agentic RAG

В сценариях Agentic RAG агент может взаимодействовать с пользователем в длинном диалоге, читать множество документов и сохранять историю действий. Infini-attention позволяет:

• Хранить всю историю диалога (миллионы токенов) без потери контекста.
• Обрабатывать длинные цепочки рассуждений (chain-of-thought) без квадратичного роста затрат.
• Интегрировать память о предыдущих запросах к базе знаний, не пересчитывая эмбеддинги заново.

Это делает Infini-attention перспективным кандидатом для ядра долговременной памяти в AI-агентах.

Пет-проект для закрепления

Задача Реализовать упрощённую версию Infini-attention на PyTorch и сравнить производительность (время и память) со стандартным вниманием на последовательностях разной длины.

Инструменты PyTorch, Hugging Face Transformers (для baseline), библиотека для профилирования (torch.profiler).

Шаги:

1. Написать класс InfiniAttention как в псевдокоде, но с multi-head и правильным обновлением памяти (gated averaging).
2. Создать синтетический датасет: случайные последовательности длиной от 1K до 100K токенов.
3. Замерить время forward и пиковое потребление памяти для стандартного внимания (из nn.MultiheadAttention) и Infini-attention.
4. Построить график зависимости времени от длины (должен быть линейный vs квадратичный).
5. Проверить качество: обучить маленькую модель (2 слоя, d=128) на задаче language modeling с длинным контекстом и измерить perplexity.

Ожидаемый результат Infini-attention покажет линейный рост времени и памяти, в то время как стандартное внимание — квадратичный. Perplexity на длинных последовательностях будет сопоставим или лучше, чем у модели с фиксированным контекстом.

Связь с другими вопросами

Навигация

• Предыдущий: 631
• Следующий: 633
• Индекс: 00. Индекс разборов