Длинный контекст

Контекстное окно — это сколько токенов модель может учитывать одновременно: промпт, документы, весь диалог и её собственный ответ. Оно выросло с 2048 токенов у GPT-3 до 128 000 у многих открытых моделей сегодня, а некоторые коммерческие модели принимают миллион и больше. Длинный контекст позволяет модели прочитать целую кодовую базу или книгу за раз. Но он приносит три проблемы: стоимость, точность и само кодирование позиций.

Проблема 1: стоимость

С длиной контекста $n$ растут две вещи:

Память. KV-кэш хранит ключи и значения каждого токена: он растёт линейно, $\propto n$ .
Вычисления. Запрос каждого токена сравнивается с каждым предыдущим ключом. Число пар растёт квадратично, $\propto n^2$ .

Все остальные умножения матриц (проекции и MLP) стоят одинаково на каждый токен, поэтому растут линейно. Вычисления префилла примерно равны

\text{FLOP} \approx \underbrace{2 P n}_{\text{веса}} + \underbrace{2 L d\, n^2}_{\text{внимание}}

для $P$ параметров, $L$ слоёв и ширины модели $d$ (длины вектора каждого токена). Первое слагаемое — одно умножение и одно сложение на каждый параметр для каждого токена. Второе считает само внимание: примерно $n^2/2$ пар токенов (каждый смотрит только на предыдущие), и в каждом слое каждая пара стоит скалярного произведения длины $d$ ради оценки и ещё $d$ умножений-сложений, чтобы подмешать значение, — около $4d$ операций.

На коротких промптах побеждает первое слагаемое. После длины, где они равны, $n = P / (L d)$ , доминирует внимание, и кривая стоимости загибается вверх. Для Llama 3 8B это $8 \times 10^9 / (32 \times 4096) \approx 61\,000$ токенов.

Рисунок 12.1 · Сколько стоит длинный контекст

Контекст128k токенов

17 ГБKV-кэш одной последовательности1,1× от весоввеса модели (16 бит)

Формы как у Llama 3 8B и 70B; «без GQA» — у каждой головы запросов свои ключи и значения.

Обе оси логарифмические, поэтому прямая означает степенной закон: у памяти наклон 1, у вычислений внимания — 2. Для формы 8B KV-кэш одной последовательности длиной 128k уже равен весам модели, а внимание обгоняет всё остальное примерно на 60k токенов. Сравните с моделью без группового внимания.

Декодирование тоже чувствует длину — через память, а не через арифметику. Каждый шаг декодирования читает весь кэш (глава 10 объясняет, почему скорость декодирования задаёт чтение, а не вычисления). На 128k токенах кэш модели на 8B размером с её веса, так что каждый новый токен выходит примерно вдвое дольше, чем с коротким промптом.

Групповое внимание, квантизация KV и другие приёмы из главы Кэширование бьют по линии памяти. FlashAttention (Быстрый инференс) делает квадратичные вычисления быстрыми и экономными по памяти, но они остаются квадратичными. Чтобы изменить сам показатель степени, нужен другой шаблон внимания — об этом ниже.

Проблема 2: а пользуется ли им модель?

Модель, которая принимает миллион токенов, не обязательно хорошо их использует. Это показали два эксперимента.

Потеря в середине (Lost in the Middle, Liu и др., 2023): моделям давали 20 найденных документов, ответ был только в одном, и этот документ передвигали. Точность была выше всего, когда он стоял первым или последним, и резко падала в середине — в худшем случае почти до уровня ответа вообще без документов. Иголка в стоге сена (needle in a haystack) прячет одно чужеродное предложение на заданной глубине длинного текста и просит его найти, перебирая и глубину, и длину.

Рисунок 12.2 · Где ломается поиск в длинном контексте

Условные числа, повторяющие форму опубликованных результатов, а не замеры конкретной модели.

Условные кривые, повторяющие форму опубликованных результатов. Ранние модели дают U-образную кривую и проваливаются в середине длинных входов; новые гораздо ровнее. Переключитесь на тест с иголкой: каждая клетка — одна длина документа и глубина иголки.

Новые модели проходят простой тест с иголкой почти идеально, поэтому теперь используют более сложные (RULER, тесты с несколькими иголками и многошаговые, где ответ требует связать несколько фактов). Они по-прежнему показывают, что эффективный контекст — длина, на которой модель сохраняет качество короткого контекста, — часто заметно меньше заявленного и что качество медленно падает по мере добавления нерелевантного текста. Это называют «гниением контекста» (context rot).

Почему U-образная кривая? Причины всё ещё обсуждают; обычно называют три. В обучающих данных редко нужна информация из середины очень длинных текстов; шаблоны внимания тяготеют к началу («сток внимания» из главы Кэширование) и к самым последним токенам; а при причинном внимании ранние токены обрабатываются, когда вопрос ещё неизвестен.

Проблема 3: позиции, которых модель не видела

RoPE кодирует позицию, поворачивая пары чисел в каждом запросе и ключе, каждую пару со своей скоростью. Быстрые пары делают много полных оборотов в пределах обучающей длины, медленные — меньше одного. Модель, обученная на 4096 токенах, никогда не видела медленные пары повёрнутыми дальше их угла на 4k. Дайте ей позицию 30 000 — эти пары окажутся на незнакомых углах, и внимание сломается.

Все исправления сжимают позиции обратно в обученный диапазон:

Интерполяция позиций (PI) делит каждую позицию на $s$ = целевая длина / обучающая длина. Все углы помещаются, но соседи теперь в $s$ раз ближе, и быстрые пары размываются.
NTK-aware масштабирование (название — от теории о том, как учатся сети, «нейронного касательного ядра», neural tangent kernel) вместо этого увеличивает основание RoPE (то самое 10 000, которое задаёт скорости вращения, глава 4): быстрые пары почти не меняются, медленные растягиваются сильнее всего.
YaRN обрабатывает каждую пару отдельно: пары, уже сделавшие много оборотов, не трогает, медленные интерполирует, между ними — плавный переход, плюс небольшая поправка температуры внимания — множителя, который делает softmax внимания чуть резче, примерно как температура при сэмплировании. Многие открытые модели с длинным контекстом используют YaRN или близкие методы.

Рисунок 12.3 · Растягиваем RoPE на длинный контекст

углы, виденные при обучении (4k) углы на целевой длине не встречались

пара 0длина волны 6виден весь кругшаг ×1

пара 16длина волны 63виден весь кругшаг ×1

пара 32длина волны 628виден весь кругшаг ×1

пара 40длина волны 2kвиден весь кругшаг ×1

пара 44длина волны 3kвиден весь кругшаг ×1

пара 48длина волны 6kне встречалисьшаг ×1

пара 52длина волны 11kне встречалисьшаг ×1

пара 56длина волны 19kне встречалисьшаг ×1

пара 60длина волны 35kне встречалисьшаг ×1

пара 63длина волны 53kне встречалисьшаг ×1

5 из 10 показанных пар выходят за обученный диапазон

Позиции дальше 4k поворачивают медленные пары на углы, которых при обучении не было. Оценки внимания там непредсказуемы, и качество разваливается.

Десять из 64 вращающихся пар одной головы, обучение на 4k токенах. Внутреннее кольцо — углы, встречавшиеся при обучении. Внешнее — углы, нужные на целевой длине, оранжевым — никогда не встречавшиеся. Быстрые пары (вверху) видели весь круг; проблема в медленных. Сравните методы и следите за шагом — насколько меняется угол между соседними токенами.

Одного масштабирования для лучшего качества мало. На практике модели предобучают на умеренной длине (4k–8k), а затем продолжают обучение на длинных документах, наращивая контекст по этапам до 128k и дальше. Эта последняя фаза занимает малую долю вычислений, но учит модель действительно пользоваться далёкими токенами.

Более дешёвые шаблоны внимания

Чтобы изменить квадратичный показатель, токен должен перестать смотреть на каждый предыдущий токен. Разреженное внимание оставляет структурированное подмножество. В главе 5 было видно, какие оценки эти шаблоны пропускают; здесь вопрос в дальности. Если токен не видит далёкое прошлое напрямую, может ли информация оттуда всё же до него добраться?

Рисунок 12.4 · Шаблоны разреженного внимания

O(n·w)58/136пар оценивается · 43% от полного внимания

внимание напрямую достижимо через предыдущие слои недостижимо

После этих слоёв токен 15 может опираться на 4 из 16 токенов

Нажмите на строку, чтобы выбрать токен-запрос.

Каждый токен видит только последние w = 4. Стоимость линейная; информация всё равно проходит w − 1 токенов за слой, поэтому глубокая сеть дотягивается далеко (Mistral 7B, многие гибридные модели).

Каждая строка — запрос, каждый столбец — ключ (16 токенов). Выделенная строка: на что выбранный токен смотрит напрямую, а с ростом числа слоёв — до чего он может дотянуться косвенно, через другие токены. Даже узкое скользящее окно дотягивается далеко, если сложить слои.

Скользящее окно (Mistral 7B, а в моделях вроде Gemma — чередуется со слоями полного внимания) стоит $O(n \cdot w)$ — работа растёт как длина, умноженная на окно, а не как квадрат длины, — и держит KV-кэш размером $w$ токенов на слой.
Глобальные токены и стоки добавляют несколько позиций, видимых всем (Longformer, BigBird, StreamingLLM).
Обучаемое разреженное внимание позволяет модели для каждого запроса выбирать, какие блоки предыдущих токенов читать, как в разреженном внимании DeepSeek и Native Sparse Attention.
Некоторые архитектуры вообще отказываются от растущего кэша внимания и хранят состояние фиксированного размера: см. За пределами классического трансформера.

Очень длинным контекстам нужно больше одного GPU на последовательность. Кольцевое внимание (ring attention) и параллелизм по контексту делят последовательность между устройствами: каждое хранит часть запросов, а блоки ключей и значений передаются по кольцу, и передача идёт одновременно с вычислениями.