Запись архива

Как делают окно в миллион токенов: RoPE, YaRN и ring-attention

Разбираем, из каких частей собирают длинный контекст в LLM: что делает RoPE, зачем нужны YaRN и LongRoPE, и почему без blockwise/ring-attention миллион токенов не помещается в память.

Схема длинного контекста в LLM: RoPE вращает позиции, YaRN и LongRoPE растягивают окно, а ring-attention распределяет блоки K/V между несколькими GPU для обработки около миллиона токенов

Окно в миллион токенов не появляется из одной «секретной настройки». В волне работ 2023–2024 годов оно сложилось из трех слоев: позиционного кодирования на базе RoPE, методов растяжения этого кодирования вроде YaRN и LongRoPE, а также системных приемов, которые позволяют вообще посчитать внимание на такой длине.

Ниже — исторический разбор именно этого стека по первоисточникам 2021–2024 годов: от RoFormer и Position Interpolation до RingAttention, открытой 1M-пайплайн работы Liu et al. и анонса Gemini 1.5 от 15 февраля 2024 года.

Коротко

  • RoPE — это не «миллион токенов», а удобная позиционная база: она вращает пары координат в запросах и ключах так, чтобы расстояние между токенами естественно влияло на attention.
  • Проблема начинается при выходе за длину, на которой модель училась: прямое RoPE-экстраполирование дает нестабильные attention-оценки. Поэтому в 2023 году появились методы интерполяции позиций, а затем YaRN и LongRoPE.
  • YaRN решает позиционную часть сравнительно дешево: в ICLR 2024 авторы сообщали о 10-кратном сокращении числа обучающих токенов и 2,5-кратном сокращении числа шагов по сравнению с более ранними схемами расширения окна.
  • RingAttention решает уже не геометрию позиций, а память и коммуникации: длинная последовательность разбивается на блоки, а блоки K/V циркулируют по кольцу устройств, пока считается точное attention без аппроксимации.
  • Даже при заявленном длинном окне модель может использовать контекст неравномерно. Работы Lost in the Middle и RULER показывают, что длинный ввод и эффективное использование этого ввода — не одно и то же.

Контекст

У длинного контекста в трансформерах есть две разные проблемы, и их часто путают. Первая — позиционная: модель должна понимать, где находится токен на длинах, которых она не видела при обучении. Вторая — системная: даже если позиционное кодирование работает, стандартное self-attention упирается в память, потому что классическая реализация материализует матрицы размера, квадратичного по длине последовательности.

RoPE появилось раньше этой «гонки за миллионом». По arXiv, работа RoFormer: Enhanced Transformer with Rotary Position Embedding была впервые выложена 20 апреля 2021 года и предложила вращать пары признаков так, чтобы абсолютная позиция задавалась через матрицу поворота, а относительное расстояние естественно проявлялось в скалярном произведении запросов и ключей. Это оказалось удобной базой для LLM, но само по себе не решило задачу выхода далеко за пределы обучающей длины.

Следующий этап начался в 2023 году. В работе Extending Context Window of Large Language Models via Positional Interpolation, впервые выложенной 27 июня 2023 года, авторы показали, что прямое экстраполирование RoPE за исходную длину ведет к «катастрофически высоким» attention-значениям, и предложили не экстраполировать позиции, а сжимать их обратно в диапазон, на котором модель обучалась. Это был первый практический рецепт расширения окна без полного переобучения модели с нуля.

Дальше стек разделился на две ветки. Первая ветка улучшала именно позиционную часть: YaRN, а затем LongRoPE. Вторая ветка улучшала исполнение attention на длинных последовательностях: сначала blockwise attention и blockwise feedforward, затем RingAttention. Только их комбинация и привела к публично описанным окнам порядка 1M токенов в 2024 году.

Метод: из чего собирается окно в миллион токенов

1. База: что делает RoPE

Идея RoPE проста: для каждой позиции p пары координат вектора поворачиваются на угол, зависящий от p. В результате позиция зашивается не добавлением отдельного эмбеддинга, а преобразованием самих запросов и ключей.

q_p = R(p) · q
k_p = R(p) · k
score(p, t) = (R(p)q) · (R(t)k)

Практический смысл в том, что attention начинает учитывать не только содержание токенов, но и их относительное расстояние. В abstract RoFormer авторы отдельно выделяют три свойства: гибкость по длине последовательности, затухание межтокенной зависимости с ростом расстояния и совместимость с линейным attention. Именно поэтому RoPE стало стандартной основой для многих открытых LLM.

2. Почему одного RoPE недостаточно

Проблема не в самом вращении, а в том, что модель обучалась только на ограниченном диапазоне позиций. Когда мы подаем сильно большие индексы, углы поворота попадают в режим, где внимание начинает вести себя нестабильно. В работе про positional interpolation авторы прямо противопоставляют extrapolation и interpolation: вместо выхода за обученный диапазон они линейно сжимают новые позиции обратно внутрь исходного окна.

Это важный сдвиг в мышлении. Длинный контекст здесь получается не потому, что модель вдруг «научилась понимать бесконечную нумерацию», а потому, что мы меняем геометрию позиций так, чтобы новая длина выглядела для RoPE как более плотная разметка старой длины.

3. YaRN: более дешевая настройка RoPE

YaRN — это уже не просто линейное сжатие индексов. В ICLR 2024 авторы комбинируют частичную интерполяцию частот RoPE с дополнительным масштабированием attention. Исторически это важно потому, что YaRN показал: расширение окна не обязательно требует очень долгого дообучения на длинных текстах.

По материалам ICLR 2024 и arXiv-версии, авторы YaRN заявляли два ключевых эффекта: примерно в 10 раз меньше обучающих токенов и в 2,5 раза меньше шагов по сравнению с более ранними методами расширения контекста. В той же работе сообщается, что модели Llama 2 7B и 13B удавалось довести до 128k контекста, а на passkey retrieval для 128k авторы показывали более 99% точности для некоторых конфигураций. Это важно, но здесь стоит помнить: это цифры самих авторов, а не независимая репликация.

Практически YaRN ценен тем, что почти не меняет архитектуру. Если RoPE уже есть, вы обычно меняете способ вычисления частот и добавляете короткое long-context дообучение, а не перестраиваете весь трансформер.

4. LongRoPE: прогрессивное растяжение до миллионов

Хотя в теме статьи вынесены RoPE, YaRN и ring, следующий логический шаг — LongRoPE, опубликованный на arXiv 21 февраля 2024 года. Это хорошая иллюстрация того, как из «сотен тысяч» переходят к «миллионам» на позиционном уровне.

LongRoPE не прыгает сразу из короткого окна в 1M+. Авторы описывают прогрессивное расширение: сначала дообучение на 256k, затем повторная интерполяция уже расширенной модели, чтобы получить 2048k. В abstract они отдельно подчеркивают, что это делается при «minor modifications» к positional embedding, то есть снова без радикальной смены архитектуры. На наш взгляд, именно этот прогрессивный curriculum — один из недооцененных выводов всей линии работ.

5. Blockwise и ring: как это вообще помещается в память

Даже идеально масштабированное RoPE бесполезно, если модель не может посчитать attention на такой длине. Здесь важен переход от обычного attention к блочному точному attention. В работе Blockwise Parallel Transformer for Large Context Models, впервые выложенной 30 мая 2023 года, Hao Liu и Pieter Abbeel предлагают считать attention по блокам и одновременно «сплавлять» с ним feedforward, чтобы не хранить лишние промежуточные активации. В paper авторы пишут, что это позволяло обучать последовательности в 32 раза длиннее vanilla Transformer и до 4 раз длиннее более ранних memory-efficient методов.

RingAttention, впервые выложенный 3 октября 2023 года и затем опубликованный на ICLR 2024, добавляет следующий слой: блоки длинной последовательности раскладываются по нескольким устройствам, а блоки ключей и значений передаются по кольцу. Пока устройство считает attention для своего query-блока, оно одновременно отправляет текущий K/V-блок соседу и принимает следующий.

GPU1: Q1 attends to KV1 -> KV2 -> KV3 ...
GPU2: Q2 attends to KV2 -> KV3 -> KV4 ...
GPU3: Q3 attends to KV3 -> KV4 -> KV1 ...

Критический момент: это не sparse-attention и не приближение. RingAttention заявлен как способ считать точное attention, но распределенно. В paper ICLR 2024 авторы формулируют идею почти линейно: если некий blockwise-трансформер может обучаться на длине s на n устройствах, то ring-схема позволяет дойти примерно до n × s.

Результаты

Ниже — компактная карта того, как накапливался стек длинного контекста. Это не единый лидерборд: цифры взяты из разных работ, на разных моделях и с разными протоколами. Их стоит читать как исторические вехи, а не как честное «A лучше B».

Работа Что меняется Сообщенный масштаб Что это значит
RoFormer / RoPE (arXiv, 20.04.2021) Позиция кодируется поворотом запросов и ключей Числа длин окна здесь не главный результат Создает удобную базу для последующего растяжения контекста
Position Interpolation (arXiv, 27.06.2023) Позиции не экстраполируются, а сжимаются внутрь исходного диапазона По отчету авторов — до 32768 при дообучении менее 1000 шагов Первый практичный рецепт расширения окна RoPE-моделей
YaRN (ICLR 2024) Улучшенная интерполяция RoPE плюс scaling attention По авторам — 64k/128k; примерно в 10 раз меньше токенов и в 2,5 раза меньше шагов, чем у ранних схем Делает long-context дообучение заметно дешевле
RingAttention (ICLR 2024) Точное blockwise attention распределяется по кольцу устройств По авторам — более 1M токенов для 7B на 32 A100 и более 30M на TPUv4-512 Снимает ограничение памяти одного устройства
LongRoPE (arXiv, 21.02.2024) Прогрессивное расширение RoPE в несколько этапов По авторам — 2048k при дообучении до 256k и до 1k шагов Показывает путь от сотен тысяч к миллионам на позиционном уровне

На уровне публичных систем 1M-контекст в 2024 году был уже не только лабораторной гипотезой. Google объявила 15 февраля 2024 года private preview Gemini 1.5 Pro с окном до 1M токенов, а в работе World Model on Million-Length Video And Language With Blockwise RingAttention, выложенной 13 февраля 2024 года, Liu et al. описали открытую пайплайн-схему с прогрессивным расширением от 4K до 1M и выпустили 7B-модели для документов и видео длиннее 1M токенов.

Интерпретация

Наш комментарий

Если убрать маркетинговые формулировки, рецепт «миллион токенов» выглядит так. RoPE дает удобную координатную систему. YaRN и LongRoPE перенастраивают эту систему так, чтобы модель не разваливалась на дальних позициях. Blockwise и ring делают вычисление и память масштабируемыми. И почти всегда между ними нужен еще четвертый компонент — curriculum, то есть постепенное дообучение на все более длинных последовательностях.

Отсюда и главный практический вывод: длинный контекст — это не одна feature flag в inference-сервере. Если у вас есть только масштабирование RoPE, но нет подходящего дообучения, модель формально примет длинный prompt, но будет использовать его плохо. Если у вас есть только ring-параллелизм, но позиционная схема не адаптирована, вы сможете посчитать длинную последовательность, но не получите надежного качества.

Поэтому корректнее говорить не «модель умеет 1M», а уточнять: каким способом растянуты позиции, на какой длине модель дообучалась, какой benchmark использован и решена ли системная часть обучения и инференса.

Ограничения

  • Длинное окно не равно равномерному использованию контекста. В работе Lost in the Middle авторы показали U-образный профиль: модели лучше используют начало и конец, чем середину, причем это наблюдается даже у long-context моделей.
  • Простые retrieval-тесты могут переоценивать реальную полезность окна. В RULER авторы пишут, что модели, которые показывают почти идеальный результат на needle-in-a-haystack и passkey retrieval, заметно деградируют на задачах трассировки, агрегации и QA при росте длины.
  • Миллион токенов не бесплатен по вычислениям. В paper RingAttention авторы отдельно оценивают, что переход маленьких моделей к 1M-контексту ведет к росту FLOPs примерно в 20–40 раз; это уже не проблема только памяти, а полноценная проблема бюджета.
  • Прогрессивное расширение остается хрупким. И LongRoPE, и 1M-pipeline Liu et al. подчеркивают важность постепенного роста длины, подбора данных и отдельной адаптации короткого окна после расширения. Иначе модель может формально поддерживать длинный ввод, но просесть на коротких задачах.
  • Цифры между работами плохо сравнимы напрямую. Где-то измеряется perplexity, где-то passkey retrieval, где-то synthetic retrieval, а где-то мультимодальная recall-задача. С инженерной точки зрения это значит, что «128k», «1M» и «2M» без протокола оценки почти ничего не говорят.

Заключение

Окно в миллион токенов — это не один трюк, а сборка из позиционной математики и системной инженерии. RoPE дает базовое представление о позиции, YaRN и LongRoPE делают эту базу пригодной для длинных окон, а blockwise/ring-attention позволяют не утонуть в памяти и коммуникациях.

Практически это означает простую вещь: если вы оцениваете long-context модель, смотрите не только на заявленную длину, но и на метод растяжения RoPE, длину long-context дообучения и benchmark, на котором это окно проверялось. Именно там скрывается разница между «модель принимает 1M токенов» и «модель надежно работает с 1M токенов».

Источники

FAQ

Можно ли получить миллион токенов только настройкой RoPE?

Обычно нет. Одного масштабирования позиций недостаточно: нужен хотя бы короткий этап long-context дообучения, а для реального обучения или инференса на такой длине — еще и память/коммуникации уровня blockwise или ring.

Зачем тогда вообще нужен ring-attention, если есть YaRN?

YaRN решает позиционную часть, а ring-attention — вычислительную. YaRN помогает модели не терять устойчивость на дальних позициях, но не делает квадратичное attention бесплатным по памяти.

Почему модели с длинным окном все равно «теряют середину»?

Потому что длинное окно и эффективное использование длинного окна — разные свойства. Это показали как Lost in the Middle, так и RULER: модель может формально принять длинный prompt, но хуже использовать факты из середины или из сложных, зашумленных частей контекста.

Заменяет ли окно в миллион токенов RAG?

Не полностью. Длинный контекст уменьшает потребность в агрессивном отборе фрагментов, но не отменяет стоимость инференса и не гарантирует лучшую точность на длинных шумных документах. На практике это скорее другой компромисс, а не универсальная замена.

Читайте также