RoPE

Одна идея

Прошлый урок закончился разложением, в котором мешали перекрёстные слагаемые. Они появились потому, что позицию прибавили к содержанию. RoPE не прибавляет — он поворачивает.

$$\big(\htmlData{k=rot}{R_i} q\big)^\top \big(\htmlData{k=rot}{R_j} k\big) = q^\top R_i^\top R_j k = q^\top R_{\htmlData{k=offset}{j-i}} k$$

Всё содержание метода — в двух равенствах: $R^\top = R^{-1}$ для ортогональной матрицы (блок 1, урок 060) и $R_i^{-1}R_j = R_{j-i}$, потому что повороты складывают углы (блок 1, урок 100). Абсолютные позиции сократились тождественно, а не приблизительно.

Композиция углов — это то самое, что видно на циферблате: поверните на $\theta$, затем на $\varphi$ и убедитесь, что оказались там же, куда ведёт поворот на $\theta + \varphi$.

матрица поворота на θ

0.77 -0.64 0.64 0.77

cos

0.77

sin

0.64

радиан

0.7

θ + φ

85°

первый угол θ

40°

второй угол φ

45°

Как это устроено внутри

Поворот $R_p$ — блочно-диагональная матрица: координаты головы разбиты на пары, и каждая пара поворачивается на свой угол $p\,\omega_i$ с теми же частотами $\omega_i = \beta^{-2i/d}$, что были у синусоид.

$R_p = \begin{pmatrix} \cos p\omega_0 & -\sin p\omega_0 & & \\ \sin p\omega_0 & \cos p\omega_0 & & \\ & & \ddots & \\ \end{pmatrix}$

Реализуется это без матриц вовсе — по две координаты за раз:

def rope(x, pos, base=10000.0):
    d = x.shape[-1]
    i = torch.arange(d // 2)
    theta = pos * base ** (-2 * i / d)
    a, b = x[..., 0::2], x[..., 1::2]
    out = torch.empty_like(x)
    out[..., 0::2] = a * theta.cos() - b * theta.sin()
    out[..., 1::2] = a * theta.sin() + b * theta.cos()
    return out

Стоимость — $O(nd)$ на голову, обучаемых параметров нуль, и применяется это только к $q$ и $k$, но не к $v$. Последнее логично: позиция должна влиять на то, куда смотреть, а не на то, что забрать.

Есть и компактная запись. Если пару координат считать комплексным числом, поворот — это умножение на $e^{ip\omega}$, и вся оценка сворачивается в

$q^\top R_\Delta k = \operatorname{Re}\sum_{i} q_i \bar{k}_i\, e^{-i\Delta\omega_i}$

Проверено численно: при $d = 64$ обе формы дают одно и то же с расхождением $10^{-15}$.

Проверьте относительность

Переключите виджет между двумя схемами и смотрите на «разброс по якорям».

синусоидальнаяRoPE

смещение j − i оценка

позиция 0позиция 8позиция 24позиция 40

размер головы d

64

смещений

32

размах оценок

2.95

разброс по якорям

2.121

норма q после

1.24

Два readout’а стоят рядом не случайно. Прибавление вектора меняет длину $q$ на величину, зависящую от позиции: измерено $1.003$ на позиции $24$ и $0.988$ на позиции $8$. Значит масштаб оценок слегка дрейфует по последовательности. Поворот не может изменить длину вовсе — норма ровно $1.000$ при любой позиции, и это второе, более тихое преимущество RoPE.

Поправка: «долгосрочное затухание» — не то, чем кажется

Здесь стоит быть точным, потому что утверждение «RoPE обеспечивает затухание внимания с расстоянием» повторяется часто и в сильной форме неверно.

Проверим прямо: возьмём случайные $q$ и $k$ и посмотрим на средний модуль оценки как функцию смещения. Две тысячи пар, $d = 64$:

смещение	средний $\vert$оценка$\vert$
$0$	$0.790$
$16$	$0.795$
$256$	$0.802$
$4096$	$0.801$

Никакого затухания. И это ожидаемо: $R_\Delta$ — ортогональная матрица, она сохраняет норму, а значит и типичную величину скалярного произведения при любом $\Delta$. В работе про RoPE доказано затухание верхней оценки при дополнительном предположении об убывании частичных сумм — это утверждение про границу, а не про сами оценки, и для случайных векторов оно не проявляется вовсе.

Как есть на самом деле: RoPE даёт архитектуре возможность выразить зависимость от расстояния, а локальность модель выучивает. При инициализации никакого предпочтения близким токенам нет. Сравните с ALiBi, где штраф $-m|i-j|$ задан жёстко и затухание есть с первого шага — это разные инженерные ставки, и путать их не стоит.

Расширение контекста

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

Посмотрим на углы при $d = 64$, $\beta = 10^4$ и длине обучения $2048$:

пара	оборотов за 2048 позиций
$0$	$326$
$8$	$32.6$
$16$	$3.26$
$24$	$0.33$
$31$	$0.044$

Быстрые пары повторяются сотни раз — им безразлично, где мы. А медленные за всё обучение не прошли даже десятой доли оборота: самая медленная пара добралась до угла $0.27$ радиана. При переходе на $8192$ она попадёт в $1.09$ радиана — область, которую модель никогда не видела. Вот где ломается экстраполяция, и ломается она у медленных пар, а не у быстрых.

Отсюда приёмы, каждый из которых решает ровно эту проблему:

приём	что делает	цена
линейная интерполяция позиций	$p \to p \cdot \frac{L_{\text{обуч}}}{L_{\text{нов}}}$	сжимает и быстрые пары, теряя различение соседей
увеличение $\beta$ (NTK-aware)	растягивает все длины волн	меняет и то, что уже работало
YaRN	интерполирует медленные пары, быстрые оставляет	сложнее, требует немного дообучения

Логика YaRN теперь читается прямо из таблицы: пары, прошедшие много оборотов, экстраполируются сами, и трогать их не нужно; интерполировать надо те, что не прошли и одного. Линейная интерполяция бьёт по всем без разбора — поэтому она работает, но заметно портит различение близких токенов.

Итог

• Одна замена: прибавить → повернуть. Из ортогональности и сложения углов следует точная зависимость только от смещения.
• Ноль обучаемых параметров, $O(nd)$ работы, применяется к $q$ и $k$ и не применяется к $v$.
• Норма запросов и ключей не меняется — масштаб оценок не дрейфует по последовательности.
• Затухания внимания с расстоянием RoPE не даёт; это выученное свойство, а не встроенное.
• Расширение контекста — это вопрос про углы медленных пар, и все известные приёмы отличаются тем, как именно они с ними обходятся.