Как работает RoPE (Rotary Position Embeddings) для экстраполяции на длинные контексты?

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

RoPE (Position Embeddings) — это метод кодирования позиций токенов в трансформерах, основанный на вращении векторов в комплексной плоскости. Он позволяет attention-механизму зависеть только от относительного расстояния между токенами, что даёт хорошую экстраполяцию на последовательности длиннее тех, на которых модель обучалась. Однако при значительном превышении длины контекста возникают проблемы из-за выхода значений вращения за обученный диапазон. Для решения применяются методы NTK-aware RoPE (изменение базовой частоты) и YaRN (интерполяция позиций), которые масштабируют позиционные кодировки на произвольную длину.

1. Зачем нужны позиционные эмбеддинги в трансформерах

Трансформер (Transformer) — архитектура, основанная на механизме self-attention, который по своей природе инвариантен к порядку токенов. Без дополнительной информации модель не различает последовательности «собака укусила человека» и «человек укусила собаку». Для внесения информации о порядке используются позиционные эмбеддинги (Encoding|positional embeddings).

Исторически применялись:

• Sinusoidal embeddings (из оригинального Transformer) — фиксированные синусоидальные функции.
• Learnable embeddings — обучаемые векторы для каждой позиции.
• Relative position biases (T5, ALiBi) — добавление смещения в attention scores]] на основе относительного расстояния.

RoPE относится к классу относительных позиционных кодировок (Encoding|Position Encoding|Position Encoding|Position Encoding|relative Encoding|positional encoding), но реализуется через вращение, а не через аддитивные смещения.

2. Как работает RoPE: математическая основа

RoPE кодирует позицию токена путём вращения его query и key векторов в комплексной плоскости. Рассмотрим двумерный случай (для каждого измерения пары).

Пусть у нас есть вектор x размерности d (чётное). Разобьём его на пары (x_{2k}, x_{2k+1}) для k = 0, 1, ..., d/2 - 1. Для каждой пары определим угол поворота, зависящий от позиции m:

θ_k = base^{-2k/d}

где base — базовая частота (обычно 10000, как в оригинальном Transformer).

Тогда для позиции m вращение задаётся матрицей:

R(m) = [[cos(mθ_k), -sin(mθ_k)],
        [sin(mθ_k),  cos(mθ_k)]]

Вектор x на позиции m преобразуется в:

RoPE(x, m) = [x_{2k} * cos(mθ_k) - x_{2k+1} * sin(mθ_k),
              x_{2k} * sin(mθ_k) + x_{2k+1} * cos(mθ_k)]

В комплексной форме это эквивалентно умножению комплексного числа (x_{2k} + i x_{2k+1}) на e^{i m θ_k}.

Ключевое свойство: скалярное произведение (attention score) между query на позиции m и key на позиции n зависит только от разности m - n:

⟨RoPE(q, m), RoPE(k, n)⟩ = Re[ (q_complex) * conj(k_complex) * e^{i (m-n) θ_k} ]

Таким образом, attention score содержит множитель cos((m-n)θ_k), что делает его функцией относительного расстояния.

3. Свойство относительности и его преимущества

Относительное позиционирование означает, что модель обучается учитывать расстояние между токенами, а не их абсолютные номера. Это даёт:

• Инвариантность к сдвигу: если последовательность сдвинуть, attention между парой токнов не изменится.
• Лучшая экстраполяция: модель может работать с длинами, не виденными во время обучения, если относительные расстояния остаются в том же диапазоне.
• Эффективность: не нужно хранить отдельные эмбеддинги для каждой позиции.

RoPE используется в большинстве современных LLM: LLaMA, Mistral, GPT-NeoX, Qwen и др.

4. Проблема экстраполяции на длинные контексты

Хотя RoPE теоретически поддерживает любые позиции, на практике при значительном превышении обученной длины контекста (например, модель обучалась на 2048 токенах, а мы подаём 8192) возникают проблемы:

• Выход за диапазон косинуса: значения cos(mθ_k) для больших m могут принимать значения, которые модель никогда не видела на этапе обучения (например, очень малые или очень большие). Attention scores становятся нестабильными.
• Разреженность градиентов: для высокочастотных компонент (маленькие θ_k) вращение происходит быстро, и при больших m функция cos осциллирует, что затрудняет обучение.
• Потеря точности: модель не может корректно различать далёкие позиции, так как разница в углах становится слишком малой.

Экстраполяция (extrapolation) — способность модели обрабатывать последовательности длиннее, чем те, на которых она обучалась, без дополнительного fine-tuning. RoPE даёт ограниченную экстраполяцию (обычно до 2-4×), но не бесконечную.

5. Решение 1: NTK-aware RoPE

NTK-aware RoPE (Neural Tangent Kernel) — метод, предложенный в контексте LLaMA и других моделей. Идея: изменить базовую частоту base так, чтобы частоты вращения стали более равномерными и покрывали больший диапазон позиций.

В оригинальном RoPE θ_k = base^{-2k/d}. Для больших k (высокочастотные компоненты) θ_k очень малы, что приводит к медленному вращению. Для маленьких k (низкочастотные) — быстрое вращение.

NTK-aware предлагает увеличить base (например, с 10000 до 500000 или 1e6). Это сдвигает все частоты вниз, делая вращение более медленным для всех компонент. В результате:

• Для тех же позиций m значения cos(mθ_k) остаются в более плавном диапазоне.
• Модель может обрабатывать более длинные последовательности без резких осцилляций.

Формула:

new_base = base * scale^{d/(d-2)}

где scale — коэффициент удлинения контекста (например, 4 для 8k вместо 2k).

Преимущества: не требует дополнительного обучения (zero-shot экстраполяция). Недостатки: может снизить точность на коротких последовательностях из-за изменения частот.

6. Решение 2: YaRN (Yet another RoPE extensioN)

YaRN — более продвинутый метод, комбинирующий интерполяцию позиций и NTK-масштабирование. Основные шаги:

1. Интерполяция позиций: вместо того чтобы подавать реальную позицию m, подают m / scale, где scale — коэффициент удлинения. Это уменьшает эффективную длину последовательности, возвращая значения cos в обученный диапазон.

2. NTK-масштабирование: дополнительно изменяют base аналогично NTK-aware, но с учётом интерполяции.

3. Температурный коэффициент: вводят параметр temperature для attention logits, чтобы компенсировать изменение масштаба.

Формула YaRN:

θ_k' = (base * scale^{d/(d-2)})^{-2k/d}
m' = m / scale

Затем применяют стандартное вращение с новыми параметрами.

YaRN позволяет экстраполировать до 32× обученной длины без fine-tuning. Используется в моделях LLaMA-2, Mistral и др.

Сравнение методов:

7. Практические рекомендации для RAG и длинных контекстов

В Agentic RAG часто требуется обрабатывать документы длиной 10k-100k токенов. Для таких случаев:

• Используйте модели, уже обученные с YaRN или NTK-aware (например, LLaMA-3.1 с контекстом 128k).
• Если модель не поддерживает длинные контексты, примените YaRN как пост-обработку весов (без обучения) — это работает для большинства LLaMA-подобных моделей.
• Для fine-tuning на длинных контекстах используйте Position Interpolation (PI) с небольшим количеством шагов.
• Следите за perplexity на длинных последовательностях: если она резко растёт, значит, экстраполяция не работает.

8. Реализация RoPE на Python (упрощённый пример)

import torch
import math

def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0):
    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
    t = torch.arange(end, device=freqs.device)
    freqs = torch.outer(t, freqs)
    freqs_cis = torch.polar(torch.ones_like(freqs), freqs)
    return freqs_cis

def apply_rotary_emb(xq: torch.Tensor, xk: torch.Tensor, freqs_cis: torch.Tensor):
    # xq, xk: (batch, seq_len, n_heads, head_dim)
    # freqs_cis: (seq_len, head_dim/2)
    xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))
    xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))
    freqs_cis = freqs_cis.unsqueeze(0).unsqueeze(2)  # (1, seq_len, 1, head_dim/2)
    xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3)
    xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3)
    return xq_out.type_as(xq), xk_out.type_as(xk)

Для экстраполяции с YaRN нужно изменить theta и масштабировать позиции.

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

Задача: Реализовать RoPE с нуля и сравнить экстраполяцию оригинального RoPE, NTK-aware и YaRN на синтетических данных.

Инструменты: Python, PyTorch, библиотека transformers (для загрузки маленькой модели, например, TinyLLaMA).

Шаги:

1. Реализовать precompute_freqs_cis и apply_rotary_emb.
2. Загрузить предобученную модель с RoPE (например, TinyLlama/TinyLlama-1.1B-Chat-v1.0).
3. Измерить perplexity на последовательностях разной длины (от 256 до 8192 токенов) для:
   • Оригинального RoPE (без изменений).
   • NTK-aware (увеличить base до 500000).
   • YaRN (интерполяция с scale=4 и base=10000).
4. Построить график зависимости perplexity от длины.
5. Сделать вывод: какой метод лучше сохраняет качество на длинных контекстах.

Ожидаемый результат: Вы увидите, что оригинальный RoPE резко теряет качество после обученной длины, NTK-aware даёт плавное ухудшение, а YaRN сохраняет низкую perplexity вплоть до 4× обученной длины.

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

Навигация

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