Rotary Position Embedding (RoPE)

Ротационное позиционное кодирование (RoPE) передаёт трансформеру информацию о порядке токенов, поворачивая каждый вектор запроса и ключа на угол, пропорциональный позиции токена. Самовнимание само по себе не может определить порядок слов, поэтому моделям нужно как-то передать информацию о позициях. При использовании RoPE оценка внимания между запросом в позиции m и ключом в позиции n зависит только от смещения m − n. RoPE представили в 2021 году вместе с моделью RoFormer, а популяризировали GPT-J и GPT-NeoX от EleutherAI. Сейчас RoPE — стандартный способ позиционного кодирования в открытых больших языковых моделях (LLM), включая Llama, Mistral, Qwen, Gemma и DeepSeek.

Как работает RoPE#

RoPE делит каждый вектор запроса и ключа размерности d на d/2 пар признаков и представляет каждую пару как точку на плоскости. В позиции m пара i поворачивается на угол m·θᵢ, где θᵢ = base^(−2i/d), а значение base в исходной статье равно 10 000. Первые пары поворачиваются быстро и различают близкие позиции, а последние — медленно и позволяют различать токены, отстоящие друг от друга на тысячи позиций. Поворот запроса на mθ и ключа на nθ приводит к тому, что их скалярное произведение зависит только от разности (m − n)θ; длина вектора при этом не меняется. Поворот применяется внутри каждого слоя внимания после проекций запроса и ключа; векторы значений не поворачиваются, а к эмбеддингам токенов на входе ничего не добавляется:

Приведённый ниже фрагмент реализует поворот с помощью комплексных чисел, как в эталонном коде Llama от Meta, и проверяет свойство относительных позиций:

import torch

torch.manual_seed(0)
dim = 64
freqs = 10000 ** (-torch.arange(0, dim, 2) / dim)  # one rotation frequency per pair of features


def rope(x, pos):
    """Rotate each consecutive feature pair of x by the angle pos * frequency."""
    pairs = torch.view_as_complex(x.reshape(-1, 2))
    return torch.view_as_real(pairs * torch.polar(torch.ones_like(freqs), pos * freqs)).flatten()


q, k = torch.randn(2, dim)
# The same query and key, three tokens apart, at two different places in the sequence
print(f"positions 5 and 2:     {rope(q, 5) @ rope(k, 2):.4f}")
print(f"positions 105 and 102: {rope(q, 105) @ rope(k, 102):.4f}")
print(f"positions 5 and 4:     {rope(q, 5) @ rope(k, 4):.4f}")  # one token apart
print(f"Vector length kept: {rope(q, 5).norm():.4f} vs {q.norm():.4f}")

Для обеих пар, отстоящих на три токена друг от друга, выводится -8.2811 — независимо от того, находятся ли они в позициях 5 и 2 или 105 и 102. Для тех же векторов, отстоящих на один токен, выводится -11.4563. Длина повёрнутого вектора запроса остаётся равной 8.3705.

RoPE и другие методы позиционного кодирования#

• Синусоидальное абсолютное кодирование: исходный трансформер один раз добавляет к входным эмбеддингам фиксированные векторы синуса и косинуса перед первым слоем. Относительное расстояние приходится выводить из смешанного сигнала.
• Обучаемые абсолютные эмбеддинги: BERT и GPT-2 обучают по одному вектору на позицию — до 512 и 1 024 позиций соответственно — и не имеют эмбеддингов для позиций за этими пределами.
• Относительные смещения позиций: T5 и ALiBi добавляют к оценкам внимания обучаемое или фиксированное смещение, зависящее от расстояния. RoPE кодирует относительное расстояние самим поворотом и не добавляет параметров.

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

Модель, обученная с RoPE на последовательностях длиной L, никогда не видела углы поворота для позиций за пределами L, поэтому при встрече с такими позициями качество внимания ухудшается. Методы расширения контекста изменяют масштаб углов, не обучая модель с нуля. Интерполяция позиций сжимает новые позиции до обученного диапазона и позволила расширить контекстное окно моделей Llama до 32 768 токенов за 1 000 шагов дообучения. Масштабирование с учётом NTK и YaRN сильнее растягивают медленные низкочастотные пары, чем быстрые; YaRN расширил контекст Llama 2 до 128K токенов, используя в 10 раз меньше обучающих токенов, чем предыдущие методы. Во время предварительного обучения Llama 3 увеличила базовое значение с 10 000 до 500 000, поэтому самые медленные пары поворачиваются меньше на токен, и удалённые позиции остаются различимыми.

RoPE в моделях компьютерного зрения#

Изображения требуют двух координат, поэтому визуальные трансформеры используют двумерное, или осевое, RoPE: половина пар признаков каждой головы поворачивается по столбцу патча, а половина — по строке. Qwen2-VL добавляет третью, временную ось для видео в мультимодальном M-RoPE. SAM 2 от Meta, доступная в пакете Ultralytics, применяет осевое RoPE в четырёх слоях своего механизма внимания к памяти, где признаки текущего видеокадра обращаются к сохранённым признакам предыдущих кадров. Ключи каждого сохранённого кадра получают ту же сетку углов, что и запросы текущего кадра, поэтому внимание между ними зависит от пространственного смещения патчей.