Rotary Position Embedding (RoPE)

La incrustación rotatoria de posición (RoPE) proporciona a un transformador información sobre el orden de los tokens al rotar cada vector de consulta y clave un ángulo proporcional a la posición del token. La autoatención por sí sola no puede distinguir el orden de las palabras, así que los modelos deben incorporar la posición de algún modo; con RoPE, la puntuación de atención entre una consulta en la posición m y una clave en la posición n depende únicamente de su desplazamiento m − n. Presentada en 2021 con el modelo RoFormer y popularizada por GPT-J y GPT-NeoX de EleutherAI, RoPE es ahora la codificación posicional estándar de los modelos de lenguaje grandes (LLM) abiertos, como Llama, Mistral, Qwen, Gemma y DeepSeek.

Cómo funciona RoPE#

RoPE divide cada vector de consulta y clave de dimensión d en d/2 pares de características y trata cada par como un punto de un plano. En la posición m, el par i gira un ángulo m·θᵢ, donde θᵢ = base^(−2i/d) y la base es 10 000 en el artículo original. Los primeros pares giran rápidamente y resuelven el orden local con precisión, mientras que los últimos giran despacio y siguen diferenciando tokens separados por miles de posiciones. Al rotar una consulta mθ y una clave nθ, su producto escalar depende únicamente de la diferencia (m − n)θ, y la rotación nunca cambia la longitud del vector. La rotación se aplica dentro de cada capa de atención, después de las proyecciones de consulta y clave; los vectores de valor no se rotan y no se añade nada a las incrustaciones de los tokens de entrada:

El fragmento siguiente implementa la rotación con números complejos, como hace el código de referencia de Llama de Meta, y comprueba la propiedad de posición relativa:

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}")

Imprime -8.2811 para ambos pares separados por tres tokens, tanto si ocupan las posiciones 5 y 2 como si ocupan las posiciones 105 y 102, e imprime -11.4563 cuando los mismos vectores están separados por un token. La consulta rotada conserva su longitud de 8.3705.

RoPE frente a otras codificaciones posicionales#

• Codificación absoluta sinusoidal: El transformador original añade vectores fijos de seno y coseno a las incrustaciones de entrada una sola vez, antes de la primera capa. La distancia relativa debe inferirse a partir de la señal combinada.
• Incrustaciones absolutas aprendidas: BERT y GPT-2 aprenden un vector por posición, hasta las posiciones 512 y 1024, respectivamente, y no tienen incrustaciones para posiciones posteriores a esos límites.
• Sesgos de posición relativa: T5 y ALiBi añaden a las puntuaciones de atención un sesgo aprendido o fijo basado en la distancia. RoPE codifica la distancia relativa mediante la propia rotación y no añade parámetros.

Ampliación de la ventana de contexto#

Un modelo entrenado con RoPE en secuencias de longitud L nunca ha visto los ángulos de rotación de posiciones posteriores a L, y la atención se degrada cuando se encuentra con ellas. Los métodos de ampliación del contexto reescalan los ángulos en lugar de volver a entrenar desde cero. Position Interpolation comprime las nuevas posiciones dentro del intervalo de entrenamiento y amplió los modelos Llama a una ventana de contexto de 32 768 tokens en 1000 pasos de ajuste fino. El escalado adaptado al NTK y YaRN estiran los pares lentos y de baja frecuencia más que los rápidos; YaRN amplió Llama 2 a 128 000 tokens con 10 veces menos tokens de entrenamiento que los métodos anteriores. Llama 3 aumentó la base de 10 000 a 500 000 durante el preentrenamiento, de modo que los pares más lentos giran menos por token y las posiciones distantes siguen siendo distinguibles.

RoPE en modelos de visión#

Las imágenes necesitan dos coordenadas, por lo que los transformadores de visión usan RoPE 2D, o axial: la mitad de los pares de características de cada cabeza gira según la columna del parche y la otra mitad, según la fila. Qwen2-VL añade un tercer eje temporal para el vídeo en su M-RoPE multimodal. El SAM 2 de Meta, que funciona en el paquete Ultralytics, aplica RoPE axial en las cuatro capas de su atención de memoria, donde las características del fotograma de vídeo actual atienden a las características almacenadas de fotogramas anteriores. Las claves de cada fotograma almacenado reciben la misma cuadrícula de ángulos que las consultas del fotograma actual, por lo que la atención entre ambos depende del desplazamiento espacial entre los parches.