RoPE: Rotary Position Embeddings и расширение контекста¶

~7 минут чтения

Предварительно: Позиционное кодирование | Attention с нуля

RoPE -- метод позиционного кодирования, который стал де-факто стандартом для открытых LLM: LLaMA, Mistral, Qwen, Gemma. Вместо того чтобы прибавлять позицию к эмбеддингу (как в BERT), RoPE вращает пары измерений на угол, пропорциональный позиции. Dot product двух повернутых векторов зависит только от разности позиций -- это дает относительное кодирование через абсолютные вращения. С техниками расширения (YaRN, NTK scaling) модели растягивают контекст с 4K до 128K--1M токенов при минимальном fine-tuning (400 шагов).

Зачем нужны позиционные кодирования¶

Трансформер без позиционной информации -- мешок слов. Он не знает, что "собака укусила человека" и "человек укусил собаку" -- это разные вещи. Self-attention считает попарные скоры между токенами, но сама операция перестановочно-инвариантна: поменяй местами два токена -- скоры между ними не изменятся.

Позиционное кодирование решает эту проблему. Но как именно его реализовать -- вопрос, который определяет способность модели работать с длинными контекстами.

Эволюция подходов¶

Метод	Тип	Экстраполяция	Используется в
Sinusoidal	Абсолютное	Плохая	Оригинальный Transformer
Learned	Абсолютное	Нет	GPT-2, BERT
ALiBi	Относительное (bias)	Хорошая	BLOOM, MPT
RoPE	Относительное (rotation)	Средняя (с расширением -- отличная)	LLaMA, Mistral, Qwen, Gemma
YaRN	RoPE + NTK scaling	Отличная	LLaMA-Long, многие finetunes

Ключевой инсайт: абсолютные позиции (1, 2, 3...) плохо экстраполируют за пределы обучения. Относительные позиции (расстояние между токенами) -- гораздо лучше. RoPE элегантно кодирует относительные расстояния через вращение в комплексной плоскости.

RoPE: интуиция¶

Аналогия: часовая стрелка¶

Представьте каждую пару измерений эмбеддинга как часовую стрелку. Для позиции $m$ стрелка повернута на угол $m \cdot \theta$:

Позиция 0: стрелка на 12 часов
Позиция 1: повернута на $\theta$ градусов
Позиция 2: повернута на $2\theta$ градусов
...

Когда два токена на позициях $m$ и $n$ взаимодействуют через dot product, значение зависит только от разности $(m - n)$, а не от абсолютных позиций. Это и есть относительное позиционное кодирование.

Разные пары измерений вращаются с разной скоростью (разное $\theta$) -- как секундная, минутная и часовая стрелки. Быстрые пары кодируют локальные зависимости, медленные -- глобальные.

Математика RoPE¶

Для позиции $m$ и пары измерений $(2k, 2k+1)$ применяется вращение:

RoPE -- матрица вращения

\[ R_{\theta, m} = \begin{pmatrix} \cos(m\theta_k) & -\sin(m\theta_k) \\ \sin(m\theta_k) & \cos(m\theta_k) \end{pmatrix} \]

Частота вращения для $k$-й пары: $$\theta_k = b^{-2k/d}, \quad b = 10000$$

Для $d = 128$ (как в LLaMA):

- $k=0$: $\theta_0 = 1.0$ -- поворот на 1 радиан за позицию (быстрое вращение, локальные паттерны)
- $k=32$: $\theta_{32} = 10000^{-0.5} = 0.01$ -- медленное вращение
- $k=63$: $\theta_{63} = 10000^{-63/64} \approx 0.00015$ -- очень медленное (глобальные паттерны)

Полный эмбеддинг $\mathbf{x}$ на позиции $m$ преобразуется поэлементно:

\[ \text{RoPE}(\mathbf{x}, m) = \begin{pmatrix} x_0 \cos(m\theta_0) - x_1 \sin(m\theta_0) \\ x_0 \sin(m\theta_0) + x_1 \cos(m\theta_0) \\ x_2 \cos(m\theta_1) - x_3 \sin(m\theta_1) \\ x_2 \sin(m\theta_1) + x_3 \cos(m\theta_1) \\ \vdots \end{pmatrix} \]

Почему это работает: относительность¶

Dot product двух повернутых векторов:

\[ \langle \text{RoPE}(\mathbf{q}, m), \text{RoPE}(\mathbf{k}, n) \rangle = \langle \mathbf{q}, R_{\theta, n-m} \mathbf{k} \rangle \]

Скалярное произведение зависит только от $(n-m)$ -- разности позиций. Вот почему RoPE -- относительное кодирование, хотя формально применяется к абсолютным позициям. Это aha-момент: абсолютные вращения дают относительные расстояния.

Реализация RoPE на Python¶

Реализация RoPE (PyTorch)

import torch

def precompute_freqs(dim: int, max_seq_len: int, base: float = 10000.0):
    """Предвычисление частот вращения."""
    # theta_k = base^(-2k/dim) для k = 0, 1, ..., dim/2 - 1
    freqs = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
    # m * theta_k для каждой позиции m
    t = torch.arange(max_seq_len)
    freqs = torch.outer(t, freqs)  # [seq_len, dim/2]
    # Комплексная форма: e^{i*m*theta} = cos(m*theta) + i*sin(m*theta)
    freqs_cis = torch.polar(torch.ones_like(freqs), freqs)
    return freqs_cis  # [seq_len, dim/2] complex

def apply_rope(x: torch.Tensor, freqs_cis: torch.Tensor):
    """Применение RoPE к тензору x.

    Args:
        x: [batch, seq_len, n_heads, dim]
        freqs_cis: [seq_len, dim/2] complex
    """
    # Разбиваем dim на пары и представляем как комплексные числа
    x_complex = torch.view_as_complex(
        x.float().reshape(*x.shape[:-1], -1, 2)
    )  # [batch, seq, heads, dim/2] complex

    # Умножение на e^{i*m*theta} = вращение в комплексной плоскости
    freqs = freqs_cis[None, :x.shape[1], None, :]  # broadcast
    x_rotated = x_complex * freqs

    # Обратно в real
    return torch.view_as_real(x_rotated).flatten(-2).type_as(x)

# Пример использования
dim = 128
seq_len = 4096
freqs = precompute_freqs(dim, seq_len)
q = torch.randn(1, seq_len, 32, dim)  # LLaMA-like
q_with_pos = apply_rope(q, freqs)
print(f"Input: {q.shape}, Output: {q_with_pos.shape}")
# Input: torch.Size([1, 4096, 32, 128])
# Output: torch.Size([1, 4096, 32, 128])

Проблема экстраполяции¶

Модель обучена на позициях $[0, L)$, где $L$ -- длина контекста при обучении (например, 4096 для LLaMA 2). При инференсе на позиции $m > L$:

Высокочастотные компоненты ($\theta_k \approx 1$) видели полные циклы вращения -- экстраполируют нормально
Низкочастотные ($\theta_k \approx 0.0001$) за время обучения повернулись на малый угол -- при $m \gg L$ попадают в невиданные углы

Экстраполяция != интерполяция

Главная ошибка: думать что RoPE автоматически работает на длинных контекстах. Без специальных техник модель с $L=4096$ деградирует уже на $L=8192$. Perplexity резко растет, модель "галлюцинирует" позиции. Нужны техники расширения (см. ниже).

Техники расширения контекста¶

Position Interpolation (PI)¶

Самый простой подход (Meta, 2023): вместо экстраполяции -- интерполяция. Масштабируем позиции:

\[m' = m \cdot \frac{L}{L'}, \quad L' > L\]

Если модель обучена на $L=4096$ и хотим $L'=32768$: позиция 32768 маппится на $32768 \cdot (4096/32768) = 4096$. Все позиции попадают в знакомый диапазон $[0, L)$.

Проблема: высокочастотные пары сжимаются слишком сильно. Если $\theta_0 = 1$ и scale = 8x, то разрешение между соседними позициями падает в 8 раз -- модель теряет локальную точность.

NTK-Aware Scaling¶

Не масштабировать позиции, а изменить базу $b$:

\[b' = b \cdot s^{d/(d-2)}, \quad s = L'/L\]

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

YaRN (Yet another RoPE extensioN)¶

State-of-the-art (2023). Комбинирует три стратегии для разных частотных диапазонов:

Высокие частоты ($\theta_k$ большое): без изменений -- они и так хорошо экстраполируют
Средние частоты: NTK-интерполяция
Низкие частоты ($\theta_k$ маленькое): линейная интерполяция

Плюс attention temperature scaling -- корректировка softmax для длинных последовательностей:

\[\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d} \cdot t}\right)V\]

где $t = 0.1 \cdot \ln(s) + 1$ -- температурный коэффициент зависящий от степени расширения $s$.

Метод	Fine-tuning	Расширение	Качество
PI	1000 шагов	8-16x	Хорошее
NTK-Aware	0 (training-free)	4-8x	Среднее
YaRN	400 шагов	16-64x	Лучшее
LongRoPE	1000 шагов	128x+ (до 2M)	Отличное

ALiBi: альтернативный подход¶

Attention with Linear Biases вообще не кодирует позиции в эмбеддинги. Вместо этого добавляет линейный штраф к attention scores:

\[\text{score}(i, j) = q_i^T k_j - m \cdot |i - j|\]

где $m$ -- slope для каждого head (разные heads смотрят на разные дистанции). Экстраполяция встроена по конструкции, но выразительность ниже чем у RoPE.

Числовой пример¶

Возьмем $d = 4$ (2 пары), $b = 10000$, позиции $m = 0, 5, 100$:

theta_0 = 10000^(0/4) = 1.0        (быстрое вращение)
theta_1 = 10000^(-2/4) = 0.01      (медленное вращение)

Позиция 0:   [cos(0), -sin(0), cos(0), -sin(0)]     = [1, 0, 1, 0]
Позиция 5:   [cos(5), -sin(5), cos(0.05), -sin(0.05)] = [0.28, -0.96, 0.999, -0.05]
Позиция 100: [cos(100), -sin(100), cos(1), -sin(1)]  = [0.86, 0.51, 0.54, -0.84]

Dot product (позиция 0 и 5)  -> зависит от разности 5
Dot product (позиция 95 и 100) -> зависит от разности 5 -- ТОЖЕ САМОЕ!

Относительность в действии: dot product зависит только от расстояния.

Что используют современные модели¶

Модель	Позиционное кодирование	Контекст (base)	Максимум (с расширением)
LLaMA 3.1	RoPE (base=500000)	128K	128K
Mistral/Mixtral	RoPE	32K	128K (YaRN)
Qwen 2.5	RoPE + YaRN	128K	1M
Gemma 2	RoPE	8K	8K
GPT-4o	Неизвестно (предположительно RoPE)	128K	128K
Claude 3.5/4	Неизвестно	200K	200K
Gemini 2.0	Неизвестно	1-2M	1-2M
BLOOM	ALiBi	2K	Экстраполяция до ~8K
MPT	ALiBi	2K-65K	Хорошая экстраполяция

base != качество на длинных контекстах

Модель с 128K контекстом может плохо работать на позициях 100K+. Needle-in-a-haystack тесты показывают: качество падает на дальних позициях. 128K -- это техническая возможность, а не гарантия quality retrieval на всех позициях.

Interview Questions¶

1. Чем RoPE отличается от learned positional embeddings?¶

Red flag: "RoPE -- это просто другой тип embedding"

Strong answer: "Learned embeddings -- lookup таблица размера $L \times d$. Позиции за пределами $L$ невозможны. RoPE -- аналитическая формула вращения: работает для любых позиций, zero обучаемых параметров. Главное: RoPE кодирует относительные расстояния -- dot product $\langle R_m q, R_n k \rangle$ зависит только от $m-n$. Learned кодирует абсолютные позиции. RoPE масштабируется без увеличения параметров."

2. Как расширить контекст модели с 4K до 32K?¶

Red flag: "Просто подать более длинный вход -- RoPE же формула, работает для любого $m$"

Strong answer: "Без техник расширения модель сломается: низкочастотные компоненты попадают в невиданные углы, perplexity резко растет. Варианты: (1) Position Interpolation: масштабируем позиции $m' = m \cdot L/L'$, нужен fine-tuning ~1000 шагов, теряется локальная точность. (2) NTK-Aware: меняем base $b' = b \cdot s^{d/(d-2)}$, training-free, но среднее качество. (3) YaRN (SOTA): три зоны частот + attention temperature, 400 шагов fine-tuning, 16-64x расширение. После -- валидация needle-in-a-haystack на разных глубинах."

3. Напишите RoPE для attention в PyTorch¶

Red flag: Добавляет позиционные embeddings через сложение ($x + PE$)

Strong answer: "RoPE -- это вращение, не сложение. Пары измерений $(x_{2k}, x_{2k+1})$ вращаются на угол $m \cdot \theta_k$. Реализация через комплексные числа: (1) precompute $e^{im\theta}$ для всех позиций, (2) reshape tensor в пары -> complex, (3) умножить на $e^{im\theta}$ (= вращение), (4) обратно в real. Применяется к Q и K, не к V. См. реализацию."