Что такое selective attention в контексте long context обработки?

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

Selective attention — это механизм, при котором модель обрабатывает только подмножество наиболее важных токенов из длинного контекста, а не все токены. Это снижает вычислительную сложность с O(n²) до O(n·k), где k — число выбранных токенов (k << n). Примеры реализаций: Quest (2025) и InfLLM. Attention|Selective attention позволяет LLM работать с контекстами длиной до миллионов токенов, сохраняя при этом качество ответов за счёт фокусировки на релевантной информации.

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

Стандартный механизм внимания (attention) в трансформерах вычисляет попарные веса между всеми токенами последовательности. Для последовательности из n токенов сложность составляет O(n²) по времени и памяти. Это делает обработку длинных контекстов (сотни тысяч или миллионы токенов) непрактичной: даже для n=100k требуется ~10¹⁰ операций, что превышает возможности современных GPU.

Пример ограничений

• GPT-4 (128k токенов) — уже дорого.
• Claude 3 (200k) — использует оптимизации.
• Для миллионов токенов нужны принципиально иные подходы.

Selective attention решает эту проблему, обрабатывая не все пары токенов, а только те, которые признаны важными.

2. Определение selective attention

Selective attention (Attention|избирательное внимание) — это семейство методов, при которых модель на каждом шаге вычисляет внимание только к выбранным токенам из контекста. Выбор осуществляется на основе некоторой метрики важности (score|importance score). Таким образом, вместо матрицы внимания размером n×n мы получаем разреженную матрицу размером n×k, где k — фиксированное число (обычно 1024–8192).

Ключевые компоненты

• Механизм выбора (selection mechanism) — определяет, какие токены важны.
• Метрика важности — может быть основана на attention scores]], косинусной близости эмбеддингов, эвристиках (например, TF-IDF) или обученной голове.
• attention|Разреженное внимание (sparse attention) — вычисление только для выбранных пар.

3. Как работает selective attention: общая схема

1. Вход: последовательность токенов длины n.
2. Вычисление важности: для каждого токена (или для каждого запроса) получаем score важности.
3. Отбор top-k: выбираем k токенов с наибольшими scores.
4. Вычисление внимания: стандартное внимание (softmax, взвешенная сумма) только для выбранных k токенов.
5. Выход: обновлённые представления для каждого токена.

Псевдокод

def selective_attention(query, keys, values, k):
    # query: [batch, n, d]
    # keys, values: [batch, n, d]
    scores = compute_importance(query, keys)  # [batch, n]
    topk_indices = topk(scores, k)            # [batch, k]
    selected_keys = gather(keys, topk_indices)
    selected_values = gather(values, topk_indices)
    attn_weights = softmax(query @ selected_keys.T / sqrt(d))
    output = attn_weights @ selected_values
    return output

4. Метрики важности токенов

Выбор «важных» токенов — ключевой момент. Используются разные подходы:

Важно: метрика должна быть дешёвой, иначе выигрыш от разреженности теряется.

5. Примеры реализаций selective attention

5.1 Quest (2025)

Quest (Query-aware Sparse Attention) — метод, который для каждого запроса выбирает k наиболее релевантных ключей. Использует двухуровневую схему:

• Сначала грубая фильтрация через кэш ключей (key cache) с помощью бинарного поиска по кластеризованным ключам.
• Затем точное внимание только к выбранным.

Особенности

• Сложность O(n·k) на практике.
• Поддерживает контексты до 1M токенов.
• Не требует дополнительного обучения — работает с предобученными моделями.

5.2 InfLLM

InfLLM — подход, который хранит memory bank из важных токенов, обновляемый по мере генерации. На каждом шаге выбираются токены из memory bank на основе косинусной близости с текущим скрытым состоянием.

Особенности

• Использует скользящее окно для локального контекста + selective attention к глобальному memory bank.
• Может обрабатывать контексты до 2M токенов.
• Требует небольшого fine-tuning для адаптации.

6. Сравнение с другими подходами к long context

Selective attention сочетает адаптивность (выбирает важные токены динамически) и эффективность (линейная сложность по n).

7. Преимущества и ограничения selective attention

Преимущества

• Масштабируемость: позволяет обрабатывать контексты до миллионов токенов.
• Адаптивность: выбирает именно те токены, которые важны для текущего запроса.
• Совместимость: может быть добавлен в существующие трансформеры без полного переобучения (например, Quest).

Ограничения

• Потеря информации: если важный токен не попал в top-k, модель его не увидит.
• Дополнительные накладные расходы: вычисление важности и сортировка могут быть дорогими при очень больших n.
• Зависимость от метрики: плохая метрика важности снижает качество.
• Необходимость fine-tuning: некоторые методы (InfLLM) требуют дообучения.

8. Связь с Agentic RAG

В контексте Agentic RAG (агентные RAG-системы) selective attention может использоваться как альтернатива или дополнение к retrieval. Вместо того чтобы искать документы во внешней базе, агент может применять selective attention к собственному длинному контексту (например, истории диалога или набору инструкций). Это особенно полезно, когда:

• Контекст уже загружен в память модели (например, в in-context learning).
• Нужно быстро переключаться между разными частями контекста без обращения к векторной БД.
• Retrieval даёт много шума — selective attention может отфильтровать нерелевантные части.

Пример: агент, который анализирует длинный лог событий, использует selective attention, чтобы сфокусироваться на критических моментах, игнорируя фоновый шум.

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

Задача: Реализовать простой selective attention механизм для генерации текста с длинным контекстом (например, суммаризация книги).

Инструменты: Python, PyTorch, Hugging Face Transformers.

Шаги:

1. Загрузите предобученную модель (например, GPT-2 или TinyLlama).
2. Реализуйте класс SelectiveAttention, который:
   • Принимает на вход последовательность токенов.
   • Вычисляет важность каждого токена как среднее attention scores по всем головам последнего слоя.
   • Выбирает top-k токенов (k = 512).
   • Выполняет forward pass только с выбранными токенами.
3. Сравните время инференса и качество (perplexity) с полным вниманием на тексте длиной 10k токенов.
4. Поэкспериментируйте с разными метриками важности (косинусная близость, случайный выбор).

Ожидаемый результат:

• Ускорение в 5–10 раз при k << n.
• Perplexity ухудшается не более чем на 5–10% при правильной метрике.
• Понимание, как выбор k влияет на компромисс скорость/качество.

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

11. Навигация

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

Навигация

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