Decoupling the "What" and "Where" With Polar Coordinate Positional Embeddings

Resumen Técnico: Desacoplamiento del "Qué" y el "Dónde" con la Incrustación de Posición en Coordenadas Polares

Declaración del Problema
En las arquitecturas Transformer, el mecanismo de atención empareja claves con consultas basándose en dos factores: el contenido ("qué") y la posición de la secuencia ("dónde"). Los autores argumentan que la Incrustación de Posición Rotatoria (RoPE), ampliamente adoptada en modelos de vanguardia como Llama 3 y DeepSeek, entrelaza estos dos factores. En RoPE, la rotación aplicada a los vectores de clave y consulta depende tanto de las magnitudes de su contenido como de sus posiciones relativas. Este entrelazamiento crea un conflicto "qué-dónde", lo que perjudica el rendimiento del modelo cuando las tareas requieren un emparejamiento independiente de contenido o posición por sí solos. Específicamente, el término de interacción en la puntuación de atención de RoPE provoca que la fase de la rotación se vea influenciada por las características del contenido, lo que dificulta que el modelo aprenda reglas que dependan estrictamente del indexado posicional o del emparejamiento de contenido.

Metodología: Incrustación de Posición en Coordenadas Polares (PoPE)
Los autores proponen PoPE, una modificación de RoPE que desacopla el contenido y la posición reformulando las representaciones de clave y consulta en coordenadas polares.

• Reformulación: En RoPE, los componentes se rotan basándose en la posición, pero la fase inicial está determinada por el contenido. En PoPE, los autores transforman los vectores de clave y consulta de valores reales en vectores complejos donde:
  • La magnitud ($\mu$) se deriva únicamente de las características del contenido utilizando una función de activación softplus ($\sigma(x) = \ln(1 + e^x)$) para asegurar la no negatividad.
  • La fase ($\phi$) está determinada estrictamente por la posición de la secuencia ($s$ o $t$) y una frecuencia específica de cada componente ($\theta_c$).
• Puntuación de Atención: La puntuación de atención se calcula como la parte real del producto punto entre la transpuesta conjugada de la consulta y la clave. Matemáticamente, esto resulta en:
  $$a^{PoPE}_{ts} = \sum_{c=1}^{d} \mu_{\tilde{q}_{tc}} \mu_{\tilde{k}_{sc}} \cos((s - t)\theta_c + \delta_c)$$
  donde $\delta_c$ es un sesgo aprendible.
• Diferencias Clave con RoPE:
  1. Desacoplamiento: Se elimina el término de interacción ($\phi_{k_{sc}} - \phi_{q_{tc}}$) que mezcla contenido y posición en RoPE. La magnitud (contenido) y la fase (posición) están estrictamente separadas.
  2. Resolución de Frecuencia: PoPE opera sobre $d$ componentes (elementos individuales) en lugar de $d/2$ pares, duplicando efectivamente el número de canales de frecuencia.
  3. Sesgo Aprendible: Se puede introducir un sesgo $\delta_c$ fijo pero aprendible para ajustar el desfase relativo óptimo para cada frecuencia sin reintroducir el conflicto contenido-posición.

Contribuciones Clave

1. Tarea de Diagnóstico (Indexación Indirecta): Los autores introducen una tarea que requiere que los modelos identifiquen un carácter objetivo basándose en un desfase relativo de un carácter fuente (por ejemplo, "encuentra el carácter 3 posiciones a la izquierda de 'z'"). Esta tarea requiere la manipulación independiente de contenido y posición.
2. Perspectiva Teórica: El artículo proporciona una derivación algebraica que muestra cómo RoPE confunde el contenido y la posición, mientras que PoPE permite que el mecanismo de atención implemente reglas como una conjunción de emparejamientos independientes de "qué" y "dónde".
3. Implementación Eficiente: Se proporciona una implementación eficiente utilizando Triton y Flash Attention 2, que calcula la parte real del producto punto complejo sin materializar la matriz compleja completa, añadiendo solo una sobrecarga de una sola multiplicación en comparación con la atención estándar.

Resultados

• Indexación Indirecta: En la tarea de diagnóstico, RoPE logró una precisión promedio de solo 11.16%, luchando por aprender el emparejamiento independiente requerido. En contraste, PoPE logró una precisión del 94.82%, demostrando su capacidad para desacoplar los factores.
• Modelado de Secuencias (Música y Genómica):
  • Música (JSB y MAESTRO): PoPE superó a RoPE en log-verosimilitud negativa (NLL) en ambos conjuntos de datos, logrando una menor perplejidad.
  • Genómica (Genoma de Referencia Humano): PoPE mostró una caída significativa en NLL en comparación con la línea base de RoPE.
• Modelado de Lenguaje (OpenWebText): A través de diversas escalas de modelo (124M, 253M y 774M de parámetros), PoPE logró consistentemente una menor perplejidad que RoPE. La brecha de rendimiento se mantuvo estable o aumentó con el tamaño del modelo.
• Tareas de Aval (Downstream Tasks): Las evaluaciones zero-shot en seis tareas de aval (LAMBADA, BLiMP, CBT, HellaSwag, PIQA, ARC-E) mostraron que los modelos basados en PoPE alcanzaron una mayor precisión media en todos los tamaños de modelo en comparación con los modelos basados en RoPE.
• Extrapolación de Longitud:
  • Zero-Shot: PoPE demostró fuertes capacidades de extrapolación de longitud zero-shot en secuencias hasta 10 veces más largas que el contexto de entrenamiento (1024 frente a 10240 tokens) sin necesidad de ajuste fino o interpolación de frecuencia.
  • Comparación con YaRN: Mientras que YaRN (un método diseñado para la extrapolación) mitiga la degradación de RoPE hasta su límite de ajuste fino (4096 tokens), se degrada significativamente más allá de eso. PoPE mantuvo el rendimiento en secuencias más largas sin requerir las técnicas de interpolación o el ajuste fino extensivo que YaRN necesita.
  • Análisis de Uso de Frecuencia: El análisis de las normas de las características reveló que los modelos RoPE tienden a suprimir los canales de alta frecuencia para evitar el ruido, mientras que los modelos PoPE utilizan un rango de frecuencias más distribuido, incluyendo altas frecuencias, a través de las capas.

Significancia y Reivindicaciones
El artículo afirma que PoPE ofrece una solución con principios para el entrelazamiento "qué-dónde" inherente a RoPE. Al desacoplar el contenido y la posición, PoPE mejora la eficiencia de los datos, la precisión asintótica y, notablemente, la generalización de la longitud del contexto. Los autores enfatizan que estas ganancias se logran con una modificación menor a la arquitectura RoPE existente, sin requerir cambios arquitectónicos complejos.

El trabajo se presenta como investigación fundamental sobre el modelado de secuencias. Los autores señalan que, si bien el método mejora las capacidades en dominios como la música, la genómica y el lenguaje natural (potencialmente ayudando en aplicaciones como el plegamiento de proteínas o el descubrimiento de fármacos), conlleva riesgos de doble uso comunes a los modelos fundacionales, tales como la generación de contenido dañino o deepfakes sofisticados. La contribución principal es la demostración de que separar la información de contenido y de posición conduce a modelos Transformer más robustos y generalizables.