English · Español

03 — Rotary Position Embedding (RoPE)¶

🇪🇸 RoPE rota las parejas de dimensiones \((2k, 2k+1)\) de Q y K por un ángulo que depende de la posición. La propiedad mágica: el producto escalar de dos vectores rotados \(\langle R_m q, R_n k \rangle\) depende solo de la diferencia \(n - m\). Así, el modelo ve posición relativa sin necesidad de codificarla explícitamente. Esa propiedad es la razón por la que RoPE ganó.

Este archivo deriva RoPE, demuestra su propiedad clave, y muestra la implementación vectorizada que usa el código de producción.

Planteamiento¶

Queremos un esquema de codificación posicional con dos propiedades:

El score de attention \(Q_i \cdot K_j\) depende solo de \(i - j\). La información de posición relativa emerge de forma natural.
Sin parámetros extra. La PE debe ser una función determinista de la posición.

Sinusoidal cumple (2) pero solo débilmente (1). Aprendida no cumple ninguna. RoPE cumple ambas, multiplicando en lugar de sumando.

El caso 2D (intuición)¶

Toma un único par 2D \((q_0, q_1)\). Rota por el ángulo \(\theta\) usando la matriz de rotación 2D estándar:

\[ R_\theta \begin{pmatrix} q_0 \\ q_1 \end{pmatrix} = \begin{pmatrix} \cos\theta & -\sin\theta \\ \sin\theta & \cos\theta \end{pmatrix} \begin{pmatrix} q_0 \\ q_1 \end{pmatrix} = \begin{pmatrix} q_0 \cos\theta - q_1 \sin\theta \\ q_0 \sin\theta + q_1 \cos\theta \end{pmatrix} \]

Esto es una rotación por \(\theta\) en el plano \((q_0, q_1)\). Preserva \(\|q\|\) (transformación ortogonal).

Ahora aplica rotaciones parametrizadas por la posición: en la posición \(m\), rota \(q\) por \(m \omega\) (alguna frecuencia \(\omega\)). En la posición \(n\), rota \(k\) por \(n \omega\).

El producto interno de los vectores rotados:

\[ \langle R_{m\omega} q, R_{n\omega} k \rangle \]

Usa una propiedad clave de las rotaciones: \(\langle R_\alpha u, R_\beta v \rangle = \langle u, R_{\beta - \alpha} v \rangle\) (las rotaciones preservan el producto interno, y puedes mover el ángulo de diferencia a un lado).

Por tanto:

\[ \boxed{\; \langle R_{m\omega} q, R_{n\omega} k \rangle = \langle q, R_{(n - m)\omega} k \rangle \;} \]

El score depende solo de la diferencia \(n - m\), no de \(m\) o \(n\) individualmente. Esa es la magia.

Apilando a través de las dimensiones¶

Los vectores Q y K reales no son 2D — son \(d_\text{head}\)-dimensionales. Aplica la rotación independientemente a cada par 2D consecutivo \((2k, 2k+1)\), cada uno a su propia frecuencia:

\[ \omega_k = \frac{1}{10000^{2k/d_\text{head}}} \quad \text{para } k = 0, 1, \ldots, d_\text{head}/2 - 1 \]

(Mismo programa de frecuencias que sinusoidal — geométrico.)

Para la posición \(p\):

\[ R_p = \text{block-diag}(R_{p \omega_0}, R_{p \omega_1}, \ldots, R_{p \omega_{d_\text{head}/2 - 1}}) \]

— una matriz diagonal por bloques de rotaciones 2D, cada una a su propia frecuencia.

El Q rotado por RoPE completo (en la posición \(p\)):

\[ Q^{\text{rope}}_p = R_p Q_p \]

Análogamente para K. Aplica attention a \((Q^{\text{rope}}, K^{\text{rope}}, V)\) — V no se rota.

La propiedad clave se generaliza: cada par contribuye con un término dependiente de la posición relativa, y la suma a través de los pares da una función de \(n - m\) únicamente.

Por qué V no se rota¶

V lleva los valores a devolver. Rotar V significaría "el valor en la posición \(p\) depende de \(p\) de forma estructurada" — pero se supone que V es contenido, no posición. Manteniendo V sin rotar, preservamos la propiedad de que el enrutamiento depende de la posición (interacción Q-K) pero el contenido entregado es independiente de la posición (V).

Esta separación es lo que hace a RoPE limpio: la posición es parte del mecanismo de attention, no parte de los datos que se enrutan.

🇪🇸 Por qué no rotar V: la rotación es para la atención (decidir a quién mirar), no para el contenido (qué llevarse). V se queda igual; Q y K se rotan.

El truco de implementación¶

Calcular \(R_p Q\) como un matmul diagonal por bloques es derrochador — la mayoría de las entradas son cero. El truco: escribirlo elemento por elemento.

Para el par \((2k, 2k+1)\) en la posición \(p\):

\[ q'_{2k} = q_{2k} \cos(p \omega_k) - q_{2k+1} \sin(p \omega_k) \]

\[ q'_{2k+1} = q_{2k} \sin(p \omega_k) + q_{2k+1} \cos(p \omega_k) \]

Vectoriza sobre todo el tensor. Sea cos_pe[p, k] = cos(p * omega_k) y sin_pe[p, k] = sin(p * omega_k), ambos de forma (T, d_head // 2). Entonces:

# q shape: (T, d_head)
q_even = q[:, 0::2]  # (T, d_head // 2)
q_odd  = q[:, 1::2]  # (T, d_head // 2)
q_rot_even = q_even * cos_pe - q_odd * sin_pe   # (T, d_head // 2)
q_rot_odd  = q_even * sin_pe + q_odd * cos_pe   # (T, d_head // 2)
# Interleave back to (T, d_head)
q_rope = np.empty_like(q)
q_rope[:, 0::2] = q_rot_even
q_rope[:, 1::2] = q_rot_odd

O equivalentemente usando el ayudante "rota por 90°":

def rotate_half(x):
    """Split x = [x_a, x_b] along last dim, return [-x_b, x_a]."""
    d = x.shape[-1]
    x_a = x[..., : d // 2]
    x_b = x[..., d // 2 :]
    return np.concatenate([-x_b, x_a], axis=-1)

# Apply: q * cos + rotate_half(q) * sin

La segunda forma es la que usan la mayoría de las implementaciones de producción. El emparejamiento de pares es implícito en la convención de splitting. El lab 02 implementa ambas y verifica que coinciden.

🇪🇸 Detalle de implementación: hay dos convenciones para emparejar dimensiones — adyacentes \((2k, 2k+1)\) o por mitades \((k, k + d/2)\). Lab 02 deja claro cuál usamos. Ambas son matemáticamente equivalentes; el código y las pre-computaciones de cos/sin deben ser consistentes.

Verificando la propiedad de posición relativa¶

La propiedad establece: \(\langle R_m q, R_n k \rangle = \langle q, R_{n-m} k \rangle\).

Comprobación numérica (el lab 02 lo implementa):

Elige \(q, k \in \mathbb{R}^{d_\text{head}}\) aleatorios.
Elige las posiciones \(m = 5, n = 7\) (así que \(n - m = 2\)).
Calcula \(\text{LHS} = \langle R_5 q, R_7 k \rangle\).
Calcula \(\text{RHS} = \langle q, R_2 k \rangle\).
Asegura \(|\text{LHS} - \text{RHS}| < 10^{-5}\).

Si tu implementación es correcta, esto se cumple para cualquier elección de \(m, n, q, k\). Si no, las matrices de rotación están mal.

¿Y la dimensión multi-head?¶

Cada cabezal tiene su propio \(d_\text{head}\). RoPE se aplica por cabezal, con su propio programa de frecuencias en \(d_\text{head}\).

Si \(d_\text{model} = 64\) y \(n_\text{heads} = 4\), entonces \(d_\text{head} = 16\), así que cada cabezal recibe 8 pares de frecuencia. Las frecuencias \(\omega_k = 1 / 10000^{2k/16}\) para \(k = 0, \ldots, 7\).

Los distintos cabezales tienen el mismo programa de frecuencias. La especialización entre cabezales viene de los \(W_Q, W_K\) aprendidos — RoPE solo añade estructura posicional encima.

Argumento de extrapolación¶

¿Por qué extrapola RoPE? Dos hilos:

La propiedad de posición relativa se cumple en cualquier \(m, n\). Sea \(m = 5, n = 7\) (entrenamiento) o \(m = 5000, n = 5002\) (extrapolación), la diferencia \(n - m = 2\) controla el score. El \(W_Q W_K^\top\) aprendido del modelo se entrenó para manejar patrones de posición relativa 2; ese conocimiento se transfiere.
Los valores de PE permanecen acotados. Los valores de PE sinusoidal están en \([-1, 1]\) para todas las posiciones. Las rotaciones de RoPE preservan la norma. No hay "deriva al infinito" a medida que \(p\) aumenta. Los modelos pueden evaluarse en posiciones no vistas con activaciones estables.

Esto es mejor que sinusoidal porque sinusoidal añade la posición al embedding, y luego el modelo tiene que aprender a separar posición de contenido. RoPE mantiene la posición en el mecanismo de attention, el contenido en V. La separación se impone por construcción.

🇪🇸 Resumen extrapolation: RoPE separa "qué decir" (V) de "a quién atender" (Q, K). La posición sólo afecta la atención (multiplicativamente, vía rotación). Esto se generaliza a posiciones no vistas mejor que cualquier scheme que mezcle posición y contenido aditivamente.

Salvedades prácticas¶

Ajuste fino (fine-tuning) de contexto largo. Los modelos entrenados en \(T = 1024\) con RoPE funcionan en \(T = 4096\), pero degradan. El "fine-tuning" de contexto largo (tema de Fase 19+) ayuda.
Interpolación de posición. Un truco común: en inferencia, escala las posiciones por \(T_\text{train} / T_\text{infer}\). Mantiene las frecuencias en el rango entrenado. El lab 03 lo menciona pero no lo implementa.
Escalado NTK-aware. Una versión más sofisticada de la interpolación de posición que respeta los componentes de alta frecuencia. Fuera de alcance.

API (fijada para `src/minimodel/positional/rope.py`)¶

def rope_frequencies(d_head: int, base: float = 10000.0) -> np.ndarray:
    """Return ω_k for k = 0, ..., d_head/2 - 1. Shape (d_head/2,)."""
    ...

def precompute_rope(T: int, d_head: int) -> tuple[np.ndarray, np.ndarray]:
    """Return cos_pe, sin_pe of shape (T, d_head/2) each."""
    ...

def apply_rope(
    q: np.ndarray,        # (T, d_head)
    k: np.ndarray,        # (T, d_head)
    cos_pe: np.ndarray,   # (T, d_head/2)
    sin_pe: np.ndarray,   # (T, d_head/2)
) -> tuple[np.ndarray, np.ndarray]:
    """Apply RoPE to q and k. Return (q_rope, k_rope) of same shape."""
    ...

La Fase 17 integra esto en MultiHeadAttention.forward como un modo opcional de codificación posicional.

Resumen¶

RoPE rota Q y K por ángulos dependientes de la posición (frecuencias espaciadas geométricamente, mismo programa que sinusoidal).
La propiedad clave: \(\langle R_m q, R_n k \rangle\) depende solo de \(n - m\).
V no se rota — la posición afecta solo al enrutamiento, no al contenido.
La implementación es elemento por elemento: \(q \cdot \cos + \text{rot}(q) \cdot \sin\). Vectorizada.
Extrapola porque la dependencia de posición relativa se cumple en cualquier par de posiciones.
El lab 02 verifica la propiedad de posición relativa numéricamente.

Ya has leído los cuatro archivos de teoría de la Fase 16. Antes de abrir el lab:

Escribe la fórmula de rotación de RoPE de memoria (caso 2D).
Enuncia la propiedad clave en una frase.
Explica por qué V no se rota.

Si alguno de estos te resulta tambaleante, relee la sección correspondiente.

Siguiente: fin de la teoría. Procede a ../lab/00-permutation-equivariance.md.