Position-Agnostic Pre-Projection for Transformer Attention: Nonlinear Feature Construction and Content Skip Before Q/K/V

Resumen Técnico: Pre-Proyección Posición-Agnóstica para la Atención de Transformadores

1. Planteamiento del Problema

Los bloques de atención en los transformadores modernos siguen una tubería fija: normalización de capa, proyecciones lineales para Q/K/V (Query, Key, Value), codificación posicional (como RoPE) y, finalmente, la atención de producto punto escalado. El autor identifica dos limitaciones fundamentales en este diseño:

1. Cuello de botella de características lineales: Las proyecciones Q/K/V son puramente lineales. Si los patrones de atención útiles dependen de combinaciones no lineales de características, el modelo debe depender de las capas anteriores (FFN) para construirlas, lo que requiere múltiples capas para una operación que podría resolverse en una sola.
2. Filtrado posicional obligatorio: Toda la información que llega a la salida debe pasar a través de la atención sensible a la posición. La información de contenido relevante independientemente de la posición (tipo semántico, categoría sintáctica, identidad de entidad) se ve forzada a pasar por el mismo enrutamiento posicional que la información dependiente de la posición, lo que puede diluir señales importantes.

2. Metodología Propuesta

El artículo propone dos modificaciones complementarias que se insertan en el bloque del transformador, ambas diseñadas para ser agnósticas a la posición (no contienen codificación posicional) y sin sobrecarga en la caché K/V:

A. Pre-Proyección No Lineal (Pre-Projection)

• Ubicación: Se inserta un pequeño MLP no lineal entre la normalización de la capa (RMSNorm) y las proyecciones Q/K/V.
• Funcionamiento: Transforma la entrada $x$ en características enriquecidas $\tilde{x}$ mediante una proyección hacia arriba, una activación SiLU y una proyección hacia abajo, manteniendo una conexión residual para una inicialización casi identidad.
• Objetivo: Construir características más ricas y no lineales antes de que se aplique cualquier codificación posicional, permitiendo que las proyecciones Q/K/V operen sobre un espacio de características más expresivo.

B. Conexión de Salto de Contenido (Content Skip)

• Mecanismo: Las mismas características enriquecidas $\tilde{x}$ de la pre-proyección se proyectan a través de un mapa lineal aprendido ($W_{skip}$) y se suman directamente a la salida de la capa, bypaseando el mecanismo de atención.
• Fórmula de salida: $out = x + \text{attn\_out} + W_{skip} \cdot \tilde{x}$.
• Aprendizaje: Los pesos de $W_{skip}$ se inicializan cerca de cero. El modelo aprende dinámicamente, capa por capa, cuánto del contenido debe bypassar la atención posicional. No se requiere un parámetro de mezcla ($\alpha$) separado; la magnitud de los pesos aprendidos cumple esta función.

Insight Clave: La pre-proyección realiza una "doble función": enriquece las entradas para la atención (camino sensible a la posición) y proporciona características de contenido que pueden eludir la atención (camino agnóstico a la posición). El modelo decide en qué ruta encajar la información según la capa.

3. Contribuciones Clave

1. Arquitectura Híbrida: Introduce un flujo de información que permite que el contenido semántico bypass la atención posicional cuando es beneficioso, separando la preparación de contenido de la conciencia posicional.
2. Eficiencia de Caché: Ninguna de las modificaciones añade sobrecarga a la caché K/V, ya que son transformaciones de costo fijo por token que ocurren antes y después del cálculo de atención.
3. Descubrimiento Arquitectónico: Se demuestra que las capas finales de los transformadores se benefician significativamente de un bypass de contenido, sugiriendo que las capas profundas manejan representaciones semánticas maduras donde el enrutamiento posicional es menos crítico.
4. Aplicabilidad Multimodal: El diseño es ideal para arquitecturas multimodales (texto, imagen, video), donde se pueden construir características independientes del modo antes de aplicar codificaciones posicionales específicas de cada modo.

4. Resultados Experimentales

Los experimentos se realizaron utilizando un protocolo de "sonda congelada" (frozen probe) en modelos Pythia-160M y Pythia-410M, entrenados en WikiText-103. Solo los parámetros inyectados (pre-proyección, salto, o LoRA) se entrenaron.

• Rendimiento General: La combinación de Pre-Proyección + Salto de Contenido obtuvo los mejores resultados en todas las métricas clave.
  • Escala 160M: Logró un aumento del +40.6% en la precisión de LAMBADA (comprensión del discurso) y una reducción del -39% en la perplejidad (PPL) en comparación con la línea base. Fue el único método que mejoró HellaSwag.
  • Escala 410M: Alcanzó la mejor perplejidad (17.0), superando incluso a una implementación LoRA con muchos más parámetros (17.7).
• Complementariedad:
  • La pre-proyección sola mejoró la comprensión (LAMBADA).
  • LoRA sola mejoró la perplejidad pero degradó la comprensión en algunos casos.
  • La combinación (Pre-proyección + Salto) unió lo mejor de ambos mundos, mejorando simultáneamente la comprensión y la perplejidad.
• Análisis por Capas: El análisis de los pesos del salto de contenido reveló un patrón consistente: las capas finales activan el bypass de contenido mucho más fuertemente que las capas iniciales.
  • En Pythia-160M, la norma del salto creció de 0.13 (capas 0-7) a 0.30 (capa 11).
  • Esto confirma la hipótesis de que las capas tempranas necesitan atención posicional para el enrutamiento sintáctico, mientras que las capas tardías se benefician de un flujo directo de información semántica.

5. Significado e Impacto

Este trabajo ofrece una visión profunda sobre cómo los transformadores procesan la información:

• Desacoplamiento de Contenido y Posición: Propone una arquitectura donde la preparación de características (contenido) y el enrutamiento (posición) se manejan de manera más flexible, permitiendo que el contenido semántico maduro fluya sin la "distorsión" de la atención posicional en las etapas finales.
• Eficiencia en Fine-Tuning: Dado que funciona bien en configuraciones de parámetros congelados (frozen probe), es una técnica prometedora para la adaptación eficiente de grandes modelos de lenguaje (LLMs) sin necesidad de reentrenamiento completo.
• Futuro Multimodal: La naturaleza agnóstica a la posición de la pre-proyección y el salto de contenido los hace candidatos ideales para arquitecturas futuras que integran texto, imágenes y video, donde las definiciones de "posición" varían drásticamente entre modalidades.

En conclusión, la propuesta demuestra que introducir no linealidad antes de la atención y permitir que el contenido bypass la atención en capas profundas mejora significativamente tanto la comprensión semántica como la eficiencia predictiva de los transformadores.