Глава 3.2 — Длинный контекст

Содержание

1. Две разные проблемы, скрывающиеся за термином «длина контекста»
2. Позиционные кодировки заново: что на самом деле означает «не обобщается»
3. Позиционная интерполяция: сжатие позиции в привычный диапазон
4. YaRN: неравномерная обработка частот вместо единообразной
5. Скользящее окно внимания в продакшене: ограничение самого вычисления
6. Как эти элементы складываются в реальной системе с длинным контекстом
7. Взгляд с точки зрения собеседования
8. Вопросы для самопроверки
9. Источники

1. Две разные проблемы, скрывающиеся за термином «длина контекста»

Глава 3.1 показала, что внимание (attention) дорого обходится при больших длинах последовательности, и рассмотрела несколько способов удешевить его: разрежение паттерна, переформулировка алгебры или оптимизация паттерна обращения к памяти без изменения самой математики. После этой главы легко решить, что «длинный контекст» — уже решённая инженерная задача: сделай внимание достаточно дешёвым, и модель сможет работать с любой длиной контекста. Такой вывод упускает вторую, независимую проблему, которой целиком посвящена эта глава: даже если бы внимание было бесплатным, модель, обученная на последовательностях длиной, скажем, 4096 токенов, никогда не видела во время обучения индекс позиции больше 4096, и нет никакой гарантии, что схема позиционного кодирования будет вести себя разумно, если на этапе инференса ей просто подать позицию 50 000.

Это стоит сформулировать максимально резко, потому что здесь часто возникает путаница, в том числе на собеседованиях: стоимость внимания (глава 3.1) — это проблема вычислений и памяти, решаемая лучшими алгоритмами и лучшим использованием оборудования. Обобщение позиционных кодировок (эта глава) — это статистическая и репрезентационная проблема: продолжает ли функция, которую выучила модель и которая была подогнана под определённый диапазон значений позиции, выдавать осмысленные результаты за пределами этого диапазона. У вас может быть неограниченный бюджет на внимание — и всё равно модель развалится на длинном контексте; и наоборот, у вас может быть модель с идеальным обобщением позиций на большие расстояния, которую при этом вычислительно невозможно запустить на такой длине. Продакшен-системам с длинным контекстом нужно решать обе проблемы одновременно, и именно поэтому эта глава напрямую продолжает предыдущую, а не заменяет её.

2. Позиционные кодировки заново: что на самом деле означает «не обобщается»

Часть II подробно разбирала механику RoPE и ALiBi, поэтому здесь мы не будем выводить их заново, но важно точно понимать, как конкретно выглядит «сбой экстраполяции», поскольку именно это мотивирует всё дальнейшее изложение. RoPE кодирует позицию, поворачивая подпары измерений каждого вектора запроса/ключа (query/key) на угол, пропорциональный позиции, причём разные пары измерений вращаются с разной частотой — одни быстро (совершая много полных циклов по мере роста позиции), другие медленно. Такая конструкция даёт RoPE очень полезное свойство в пределах диапазона позиций, встречавшихся при обучении: скалярное произведение между запросом в позиции $m$ и ключом в позиции $n$ зависит только от относительного расстояния $m - n$, а не от абсолютных позиций — именно такое индуктивное смещение и нужно для языка, где относительный порядок обычно важнее абсолютной позиции в документе. Проблемы возникают, когда и $m$, и $n$ значительно превышают всё, что встречалось при обучении. Пары измерений с высокой частотой вращаются так быстро, что при очень больших позициях их углы поворота фактически становятся неинтерпретируемыми — модель никогда не училась понимать, что означает данный угол на таком удалении, поскольку никогда не видела таких углов при обучении, и, что ещё хуже, выученные моделью представления запросов/ключей для данного относительного расстояния могли быть согласованы по смыслу только в пределах обучающего окна.

ALiBi, рассмотренный в главе 2.2, пошёл совершенно другим путём: вместо вращения векторов он добавляет фиксированный, пропорциональный расстоянию штраф прямо к оценкам внимания, штрафуя внимание к далёким токенам тем сильнее, чем больше расстояние. Поскольку этот штраф — простая монотонная функция расстояния без каких-либо обучаемых частотных компонент, которые могли бы выйти за пределы диапазона, ALiBi был специально спроектирован и эмпирически показал гораздо более качественное поведение при экстраполяции — модель, обученная на коротких последовательностях с ALiBi, может оцениваться на заметно более длинных последовательностях с плавной деградацией качества, а не с резким обвалом. Это стоит держать в уме при чтении остальной части главы, потому что это переосмысляет проблему экстраполяции RoPE как в некотором смысле самим же RoPE и порождённую: RoPE получает отличное поведение относительной позиции и сильную эмпирическую производительность в пределах обучающего диапазона в обмен на хрупкость за его пределами, и два метода, рассматриваемых далее — позиционная интерполяция и YaRN — существуют именно потому, что почти все крупные продакшен-модели на практике используют RoPE, а не ALiBi, что делает исправление экстраполяционного поведения RoPE практически насущной задачей, а не чисто академической.

3. Позиционная интерполяция: сжатие позиции в привычный диапазон

Наиболее прямое решение проблемы экстраполяции RoPE, предложенное Chen et al. из Meta AI, отталкивается от наблюдения о том, что именно ломается, а что нет. Проблема не в том, что математика RoPE перестаёт работать на больших позициях — вращение математически прекрасно определено при любой позиции; проблема в том, что выученные параметры модели когда-либо подвергались воздействию — и настраивались давать хорошие предсказания — только на углах, соответствующих позициям вплоть до длины обучения $L$. Основная идея позиционной интерполяции обезоруживающе проста: вместо того чтобы позволять индексам позиции доходить до некоторой новой, большей длины $L'$, которую модель никогда не видела, их масштабируют вниз с коэффициентом $L/L'$, так что наибольший индекс позиции, с которым модель когда-либо фактически вычисляет, всё ещё остаётся равным $L$ — значению, на котором она обучалась. Конкретно: если вы хотите расширить модель, обученную на 4096 токенах, до 16384 токенов, вы не подаёте в RoPE сырые индексы позиции от 0 до 16383; вы умножаете каждый индекс позиции на $4096/16384 = 0.25$ перед вычислением углов поворота, так что токен на истинной позиции 16 000 обрабатывается — для целей вычисления его угла поворота — так, как если бы он находился на позиции 4000, что вполне комфортно укладывается в диапазон, который модель уже понимает.

Компромисс здесь именно такой, какого и следует ожидать из названия: вы обмениваете разрешение на диапазон. Втискивание в четыре раза большего числа истинных позиций в тот же диапазон углов означает, что два токена, находящиеся, скажем, на расстоянии 40 позиций друг от друга в сжатом представлении, на самом деле могут отстоять на 160 позиций в реальном документе — тонкое ощущение модели «насколько далеко друг от друга эти два токена» грубеет по мере роста коэффициента сжатия. Эмпирически оказывается, что модель может достаточно хорошо адаптироваться к этому огрублению при небольшом объёме дополнительного дообучения (fine-tuning) на новой длине контекста, поскольку модель по-прежнему работает целиком в пределах диапазона углов, который она уже умеет интерпретировать; ей нужно лишь выучить новое (или скорректированное) соответствие между «этим сжатым расстоянием» и «этой степенью релевантности», а не разбираться с углами, которых она по-настоящему никогда не встречала. Настоящая привлекательность метода в том, что это дообучение дёшево — сравнительно небольшое число дополнительных шагов обучения на новой длине, а не полный повторный прогон обучения, — и именно поэтому расширение контекста через интерполяцию стало стандартным практическим рецептом, а не чисто теоретическим исправлением.

4. YaRN: неравномерная обработка частот вместо единообразной

Равномерное масштабирование позиционной интерполяции просто, но обрабатывает частоту вращения каждой пары измерений совершенно одинаково, и эта единообразность на самом деле не согласуется с тем, как ведут себя частоты RoPE. Напомним, что RoPE присваивает каждой паре измерений свою частоту, обычно охватывающую широкий диапазон — от высоких частот (совершающих множество полных вращений даже на короткой последовательности) до низких частот (едва совершающих долю одного оборота даже на полной длине обучения модели). Пара измерений с высокой частотой, уже совершившая множество полных вращений в пределах обучающего диапазона, фактически уже «увидела» плотную выборку возможных углов — экстраполяция для неё не столько требует от модели интерпретации незнакомого угла, сколько продолжения паттерна, который уже был тщательно опробован. Пара измерений с низкой частотой, напротив, может пройти лишь малую долю полного оборота даже за весь обучающий диапазон, а значит, модель когда-либо видела лишь узкую полоску возможных углов для этой пары измерений — дальнейшее продвижение этого измерения представляет собой гораздо более рискованную экстраполяцию, поскольку обучающего сигнала о том, что вообще должны кодировать бо́льшие углы для этой пары измерений, сравнительно немного.

Это частотно-зависимое рассуждение впервые появилось как NTK-aware scaling — схема интерполяции, предложенная в комьюнити-заметках ещё до формализации в литературе, которая рассуждает в терминах теории Neural Tangent Kernel о том, какие частоты сеть способна научиться экстраполировать, а какие нет, и интерполирует каждую частоту на свою величину соответственно, а не применяет единый равномерный коэффициент позиционной интерполяции. NTK-aware scaling — прямой предшественник более тщательно проработанной схемы ниже: само название YaRN — это отсылка к тому, что это один из нескольких методов в стиле «NTK-aware», последовавших за исходной идеей.

YaRN (Yet another RoPE extensioN method), предложенный Peng, Quesnelle, Fan и Shippole, берёт это наблюдение и превращает его в основной принцип проектирования: вместо применения единого равномерного коэффициента масштабирования позиционной интерполяции ко всем частотам, YaRN интерполирует высокочастотные измерения менее агрессивно (поскольку они и сами по себе экстраполируются достаточно хорошо) и интерполирует низкочастотные измерения более агрессивно (поскольку им требуется больше всего помощи), с плавным переходом между этими двумя режимами, откалиброванным по частоте. Такой прицельный подход устраняет конкретную слабость равномерной интерполяции: поскольку в обычной позиционной интерполяции каждое измерение получает одинаковое сжатие, высокочастотные измерения — которым на самом деле не нужна была особая помощь — тоже без нужды сжимаются, ненужно размывая как раз ту тонкую относительную позиционную информацию, которую высокие частоты RoPE умеют хорошо кодировать. YaRN также включает дополнительную корректировку температуры внимания, компенсирующую тонкий сдвиг в распределении оценок внимания, который вносит интерполяция, — авторы обнаружили, что это дополнительно улучшает качество. Практический итог, о котором сообщается в статье, состоит в том, что YaRN достигает сильных результатов расширения контекста при существенно меньшем объёме данных для дообучения и меньшем числе шагов обучения, чем требуется равномерной позиционной интерполяции для достижения сопоставимого качества — прямая выгода от учёта того факта, что частотные диапазоны RoPE не взаимозаменяемы.

5. Скользящее окно внимания в продакшене: ограничение самого вычисления

Всё, что рассматривалось в этой главе до сих пор, касалось того, как заставить позиционное понимание модели обобщаться дальше, при этом неявно предполагая, что глава 3.1 уже сделала само внимание доступным по цене на такой длине. Продакшен-системы иногда вместо этого используют более прямолинейный, архитектурно жёсткий подход: ограничить вычисление внимания так, чтобы оно вообще никогда не масштабировалось с номинальной длиной контекста, независимо от того, насколько номинально длинным является диалог или документ. Mistral 7B — конкретный, широко используемый пример этого подхода. Его слои внимания используют скользящее окно: каждый токен обращает внимание только на окно фиксированного размера из самых последних токенов (окно в несколько тысяч токенов в дизайне Mistral), независимо от того, сколько текста предшествует этому окну во всей последовательности. В пределах одного слоя это ограничивает стоимость внимания тем же паттерном $O(n \cdot w)$, что и локальный компонент Longformer из главы 3.1 — но мотивация и постановка задачи здесь несколько иные: это не позиционируется в первую очередь как стратегия аппроксимации для удешевления модели с фиксированным контекстом, это нативный дизайн модели для обработки последовательностей, номинально гораздо более длинных, чем прямое рецептивное поле любого отдельного слоя.

Ключевая идея, которая делает это работоспособным, а не калечащим модель, состоит в том, что рецептивное поле накапливается по слоям — так же, как в свёрточной сети с фиксированным размером ядра. Токен в слое 1 может видеть только $w$ токенов непосредственно позади себя, но токен в слое 2 построен из представлений слоя 1, каждое из которых само видело $w$ токенов позади себя, так что эффективное рецептивное поле токена в слое 2 составляет примерно $2w$; после $k$ слоёв скользящего окна внимания информация в принципе может распространиться примерно на $k \times w$ токенов исходной последовательности, даже если ни один отдельный слой никогда напрямую не выходил вниманием за пределы собственного окна. Дизайн Mistral опирается на достаточное число слоёв, чтобы это накопленное рецептивное поле уверенно покрывало предполагаемую рабочую длину контекста. Тем не менее компромисс здесь реальный, и об этом стоит говорить честно: это косвенный поток информации, опосредованный несколькими преобразованиями, а не прямое сопоставление токен-токен, которое даёт полное внимание, и нет никакой гарантии, что конкретный фрагмент информации, находящийся на расстоянии во много слоёв, переживёт этот косвенный путь так же чисто, как его сохранила бы прямая связь внимания. На практике скользящее окно внимания в сочетании с достаточной глубиной сети оказалось работоспособным, эффективным рецептом для целого ряда продакшен-задач, но это настоящий компромисс проектирования, а не бесплатный обед — это тот же фундаментальный компромисс, на который пошёл Longformer в главе 3.1, здесь принятый как основной архитектурный выбор, а не как доработка поверх уже существующей модели.

6. Как эти элементы складываются в реальной системе с длинным контекстом

Стоит сделать шаг назад и явно проговорить, как техники этой главы сочетаются друг с другом и с главой 3.1, потому что реальная продакшен-система с длинным контекстом обычно не выбирает лишь одну из этих идей — она комбинирует их. Модель может использовать RoPE, расширенный через YaRN-подобную частотно-зависимую интерполяцию, чтобы обрабатывать индексы позиции далеко за пределами исходного диапазона обучения, вычисляемый эффективно через FlashAttention, чтобы сырая стоимость вычислений и памяти для полного внимания на этой расширенной длине оставалась приемлемой, и в некоторых развёртываниях дополнительно ограниченный скользящим окном, чтобы даже линейная по памяти стоимость FlashAttention не росла без предела по мере того, как диалоги становятся сколь угодно длинными. Ни одна из этих техник не заменяет другую: исправление позиционного обобщения никак не уменьшает вычислительную сложность $O(n^2)$ полного внимания на длинной последовательности, а удешевление внимания никак не устраняет путаницу модели на основе RoPE в отношении индексов позиции, которых она никогда не встречала. Считать «длинный контекст» единой проблемой с единым решением — ошибка, часто встречающаяся в неформальных обсуждениях этой темы, и умение чётко отделить «это проблема позиционного обобщения» от «это проблема вычислений/памяти» — и назвать конкретную технику, решающую каждую из них, — один из более надёжных сигналов подлинной глубины понимания этой темы на техническом собеседовании.

Стоит явно назвать ещё одно ограничение, поскольку это не проблема ни позиционного обобщения, ни вычислений/памяти, и ни одна из техник выше её не затрагивает: даже модель, корректно справляющаяся с заданной длиной контекста в смысле «может ли она вообще туда обратиться внимание», о котором шла речь выше, часто использует информацию из середины длинного контекста заметно хуже, чем информацию у начала или конца. Liu с соавторами задокументировали это напрямую, обнаружив устойчивую U-образную кривую качества на задачах, требующих извлечения информации с конкретной позиции в длинном документе: точность максимальна, когда нужная информация находится в самом начале или самом конце контекста, и заметно падает для той же самой информации, помещённой в середину, — у моделей, которые справляются с сырой длиной контекста без какого-либо архитектурного сбоя. Это по-настоящему отдельная проблема от всего остального в этой главе: позиционная интерполяция и YaRN — про то, способна ли модель вообще представить позицию и обратить на неё внимание, скользящее окно внимания и FlashAttention — про то, по карману ли обратить на неё внимание, а «потерянность в середине» — про то, действительно ли модель хорошо использует информацию, которую технически может увидеть и позволить себе рассмотреть, — это вопрос скорее обучения и поведения, чем архитектуры, и с ним всё равно столкнётся система, полагающаяся исключительно на удлинение контекста без внимания к качеству извлечения внутри него.

Теперь, когда внимание одновременно доступно по цене (глава 3.1) и расширяемо за пределы обучающих позиций (эта глава), следующая глава обращается к совершенно другой оси масштабирования: не удешевлению или обобщению фиксированного объёма вычислений, а увеличению самой модели — а именно, к резкому росту числа её параметров — без пропорциональной платы в виде вычислений на токен. Это область смеси экспертов (Mixture of Experts), которая заменяет плотный полносвязный слой, через который до сих пор безусловно проходил каждый токен, гораздо большим банком специализированных полносвязных сетей, из которых любой данный токен фактически использует лишь несколько.

7. Взгляд с точки зрения собеседования

«Команда хочет расширить окно контекста своей модели с 8K до 32K токенов. Какие разные проблемы им нужно решить — и одна ли это проблема или две разные?» Сильный ответ явно разделяет их: нужно, чтобы внимание оставалось вычислительно приемлемым на 32K токенах (проблема главы 3.1 — решается через FlashAttention для точного вычисления или через разреженное/линейное внимание, если этого недостаточно), и отдельно нужно, чтобы схема позиционного кодирования вела себя разумно на индексах позиции, которые она никогда не видела при обучении (проблема этой главы — решается через позиционную интерполяцию или YaRN, если используется RoPE, и в значительной степени уже решена, если используется ALiBi). Сильный кандидат отмечает, что эти проблемы требуют разных техник и решение одной не решает другую.

«Почему RoPE нуждается в исправлениях на основе интерполяции для длинного контекста, а ALiBi, как правило, нет?» Сильный ответ противопоставляет механизмы: RoPE кодирует позицию через основанные на частоте вращения, которые обучались только на углах вплоть до некоторой максимальной позиции, поэтому новые, очень большие позиции порождают углы, которые модель никогда не училась интерпретировать. ALiBi же добавляет простой монотонный, пропорциональный расстоянию штраф к оценкам внимания без каких-либо обучаемых частотных компонент, поэтому аналогичного сбоя «незнакомого угла» не возникает, что даёт гораздо более качественную нативную экстраполяцию — цена, которую RoPE платит за своё отличное поведение относительной позиции и эмпирическую производительность внутри диапазона, — большая хрупкость за его пределами.

«Объясните, почему YaRN обрабатывает высоко- и низкочастотные измерения RoPE по-разному, а не применяет равномерную позиционную интерполяцию.» Сильный ответ объясняет, что пары измерений RoPE охватывают диапазон частот вращения, и высокочастотные пары совершают множество вращений даже в пределах исходной длины обучения, поэтому они уже фактически опробовали плотный диапазон углов и экстраполируются достаточно хорошо сами по себе, тогда как низкочастотные пары проходят лишь малую долю оборота даже на полной длине обучения, поэтому их расширение — гораздо более рискованная экстраполяция, требующая более агрессивной интерполяции. Применение одного равномерного коэффициента сжатия ко всем частотам, как в обычной позиционной интерполяции, без необходимости размывает тонкую относительную позиционную информацию, которую несут высокочастотные измерения, не нуждавшиеся в помощи.

«Каков компромисс скользящего окна внимания, используемого в Mistral, и как информация извне окна вообще достигает данного токена?» Сильный ответ объясняет, что прямое внимание каждого слоя ограничено фиксированным окном, но рецептивное поле накапливается по глубине так же, как в CNN — эффективный информационный горизонт токена после $k$ слоёв составляет примерно $k$-кратный размер окна, — так что информация издалека всё же может достичь токена, но лишь через несколько косвенных преобразований, а не одно прямое сравнение внимания, что представляет собой реальный (хотя часто приемлемый) компромисс по точности по сравнению с полным вниманием.

«Если бы вычисления и память для внимания были буквально бесплатными, были бы проблемы длинного контекста решены?» Сильный ответ говорит «нет» и объясняет почему: обобщение позиционного кодирования — отдельная, статистическая проблема о том, осмысленны ли выученные представления модели при значениях позиции за пределами её обучающего распределения, и никакое удешевление вычислений не меняет того, чему модель на самом деле научилась. Именно вокруг этого различия построена вся глава, и его явное озвучивание — сигнал подлинного понимания, а не поверхностной ассоциации «длинный контекст = дорогое внимание».

8. Вопросы для самопроверки

1. Объясните своими словами, почему «внимание дорого на длинном контексте» и «позиционные кодировки не обобщаются на длинный контекст» — две разные проблемы, используя RoPE как конкретный пример для второй.
2. Почему RoPE гарантирует, что скалярное произведение запроса и ключа зависит только от их относительного расстояния, и почему это свойство нарушается на практике при индексах позиции, значительно выходящих за пределы обучающего диапазона?
3. Пройдите по механике позиционной интерполяции: какая конкретно величина масштабируется, на какой коэффициент, и почему это удерживает модель в пределах диапазона углов, который она уже понимает?
4. В чём заключается компромисс «разрешение против диапазона», вносимый позиционной интерполяцией, и почему небольшого объёма дообучения обычно достаточно, чтобы адаптировать модель к нему?
5. Объясните, почему YaRN интерполирует разные частотные диапазоны RoPE в разной степени и почему ожидается, что для этого потребуется меньше данных дообучения, чем для равномерной интерполяции, чтобы достичь того же качества.
6. В скользящем окне внимания Mistral, как эффективное рецептивное поле токена растёт с глубиной сети, и какой конкретный компромисс по точности вносит этот косвенный поток информации по сравнению с полным вниманием?
7. Если продакшен-система использует RoPE с расширением в стиле YaRN, FlashAttention и скользящее окно все вместе, объясните, какую проблему решает каждый из трёх компонентов по отдельности.

9. Источники

• Press, O., Smith, N. A., & Lewis, M. (2021). Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation (ALiBi). ICLR 2022. arXiv:2108.12409. https://arxiv.org/abs/2108.12409
• Chen, S., Wong, S., Chen, L., & Tian, Y. (2023, Meta AI). Extending Context Window of Large Language Models via Positional Interpolation. arXiv:2306.15595. https://arxiv.org/abs/2306.15595
• Peng, B., Quesnelle, J., Fan, H., & Shippole, E. (2023). YaRN: Efficient Context Window Extension of Large Language Models. ICLR 2024. arXiv:2309.00071. https://arxiv.org/abs/2309.00071
• Jiang, A. Q., et al. (2023, Mistral AI). Mistral 7B. arXiv:2310.06825. https://arxiv.org/abs/2310.06825
• Liu, N. F., Lin, K., Hewitt, J., Paranjape, A., Bevilacqua, M., Petroni, F., & Liang, P. (2023). Lost in the Middle: How Language Models Use Long Contexts. TACL 2024. arXiv:2307.03172. https://arxiv.org/abs/2307.03172