Звёздное внимание: Эффективный вывод LLM для длинных последовательностей

Современные крупные языковые модели (LLM) способны обрабатывать контексты до миллионов токенов, открывая возможности для анализа кода на уровне репозитория, суммаризации нескольких документов и поиска по большим корпусам данных. Однако обработка таких длинных последовательностей требует значительных вычислительных и памяти ресурсов из-за квадратичной сложности механизма само-внимания. В этой статье мы рассмотрим новый подход к решению этой проблемы, названный Star Attention (Звёздное внимание), который предлагает значительные улучшения в эффективности вывода LLM при обработке длинных последовательностей.

LLM, такие как Gemini, Anthropic и Meta AI, могут поддерживать контексты длиной до миллионов токенов, что позволяет решать задачи, требующие анализа больших объемов текста. Однако, основное препятствие в этом процессе — это квадратичная сложность само-внимания, которая делает обработку таких контекстов дорогой и медленной. Различные методы были предложены для снижения потребления памяти и ускорения вывода, такие как FlashAttention и RingAttention, но они имеют свои ограничения.

Star Attention — это новый алгоритм, который использует двухфазный подход для эффективного вывода LLM на длинных последовательностях. Этот метод основан на наблюдении, что при выводе LLM обычно имеются две стадии: кодирование подсказки и генерация токенов. В контексте длинных задач, входные данные часто включают длинный контекст, за которым следует короткий запрос и короткий ответ. Информация, необходимая для ответа на запрос, часто локализована в небольших частях контекста, что позволяет использовать локальное внимание для контекста и глобальное внимание для запроса.

Алгоритм Star Attention

Фаза 1: Кодирование контекста

В первой фазе контекст делится на континуальные блоки и распределяется по "контекстным" хостам. Каждый хост также получает копию первого блока, называемого "якорным блоком". Хосты вычисляют само-внимание только для своих назначенных блоков, без обмена данными друг с другом, что снижает сложность внимания с квадратичной до линейной относительно длины контекста.

**Контекст:** [c1, c2, ..., cn]
**Якорный блок:** c1
**Распределение:** [c1, (c1, c2), (c1, c3), ..., (c1, cn)]

Фаза 2: Кодирование запроса и генерация токенов

Во второй фазе запрос реплицируется на все хосты, где он сначала обрабатывается с использованием локального кэша ключ-значение (KV) каждого хоста. Затем глобальное внимание вычисляется путем агрегации результатов на обозначенном "запросном" хосте, который эффективно обменивается одним вектором и скаляром на каждый токен от каждого контекстного хоста. Только запросный хост обновляет свой KV-кэш на этом этапе.

**Запрос:** q
**Агрегация:** Query host aggregates results from all hosts.

Эксперименты и результаты

Star Attention был оценен на нескольких моделях Llama, с длинами последовательностей от 16K до 1M токенов, на бенчмарках RULER и BABILong. Сравнение с Ring Attention показало, что Star Attention достигает значительного ускорения вывода (до 11 раз) при сохранении 95-100% точности.

Модель: Llama-3-8B-Instruct
Длина последовательности: 1048K
Размер блока: 32K
Ускорение: 2.7x
Точность: 77.39% (Star Attention) vs 76.16% (Ring Attention)

Для более крупных моделей, таких как Llama-3.1-70B-Instruct, ускорение было еще более выраженным:

Модель: Llama-3.1-70B-Instruct
Длина последовательности: 128K
Размер блока: 16K
Ускорение: 4.7x
Точность: 88.54% (Star Attention) vs 90.17% (Ring Attention)

Заключение

Star Attention представляет собой значительный шаг вперед в области эффективного вывода LLM на длинных последовательностях. Этот метод позволяет масштабировать длину контекста линейно с увеличением числа хостов, сохраняя при этом высокую точность и значительно снижая затраты на память и время вывода. Он совместим с большинством трансформеров, обученных с глобальным вниманием, и может быть дополнительно оптимизирован с другими методами, такими как FlashAttention или сжатие KV-кэша.

Будущие исследования могут быть направлены на дальнейшее улучшение механизма якорных блоков и повышение производительности на более сложных задачах с длинным контекстом, чтобы сделать Star Attention еще более масштабируемым и устойчивым.