Resumen Técnico: ¿Necesitan los Transformers tres proyecciones?

Planteamiento del Problema

Desde la introducción de la arquitectura Transformer, el mecanismo de atención Query, Key y Value (QKV) ha sido el estándar para agregar información a través de secuencias. Sin embargo, la necesidad de mantener tres proyecciones lineales distintas e independientes para estos componentes sigue siendo una pregunta abierta. Mientras que otras arquitecturas como las CNN o los Modelos de Espacio de Estados (SSM) suelen utilizar representaciones internas más unificadas, los Transformers persisten con una estructura tripartita. A medida que las ventanas de contexto se expanden y la demanda de inferencia en tiempo real crece, la comunidad de investigación se ha centrado en la eficiencia arquitectónica para aliviar el cuello de botella cuadrático de la autoatención. Este artículo investiga si las tres proyecciones son verdaderamente necesarias o si representan una redundancia que puede reducirse para disminuir el recuento de parámetros y, lo que es más crítico, la sobrecarga de memoria de inferencia sin una degradación significativa del rendimiento.

Metodología

Los autores proponen y evalúan sistemáticamente tres restricciones de "Compartición de Proyecciones" (Projection-Sharing) que reducen progresivamente el número de transformaciones aprendidas en el mecanismo de atención:

1. Q=K-V (Consulta-Clave-Valor Unificada): Las proyecciones de consulta y clave se unifican ($Q=K$), mientras que la proyección de valor permanece separada. Esto resulta en una matriz de atención simétrica ($KK^T$). Para abordar las limitaciones de la simetría en tareas secuenciales, los autores introducen una variante $(Q=K-V)^+$ que inyecta asimetría mediante codificaciones posicionales 2D.
2. Q-K=V (Clave-Valor Unificados): Las proyecciones de clave y valor se unifican ($K=V$), mientras que la consulta permanece separada. Esto preserva mapas de atención asimétricos ($QK^T$) e impone una forma de vinculación de pesos entre claves y valores.
3. Q=K=V (Proyección Única): Los tres roles comparten una única matriz de proyección. Esto combina la simetría de la primera variante con el cuello de botella representacional de la segunda.

El estudio también explora la ortogonalidad de estas estrategias de compartición de proyecciones con las técnicas existentes de "Compartición de Cabezales" (Head-Sharing) como la Atención de Consulta Agrupada (GQA) y la Atención de Consulta Múltiple (MQA). Los autores hipotetizan que combinar la compartición de proyecciones (restringiendo $K=V$) con la compartición de cabezales (reduciendo el número de cabezales KV) puede generar ganancias compuestas en la eficiencia de memoria.

Los experimentos se realizaron en 12 tareas diversas:

• Razonamiento Sintético: Cinco tareas de manipulación de listas (Reverse, Sort, Sub, Swap, Copy).
• Visión por Computadora: Clasificación de imágenes en MNIST, FashionMNIST, CIFAR-10, CIFAR-100 y TinyImageNet, así como detección de anomalías en conjuntos y segmentación semántica.
• Modelado de Lenguaje: Pre-entrenamiento de modelos de estilo GPT de 300M y 1.2B de parámetros con 10 mil millones de tokens del conjunto de datos SlimPajama.

Resultados Clave

1. Intercambio entre Rendimiento y Eficiencia

• Q-K=V (Equilibrio Óptimo): Esta variante emergió como la clara ganadora. En el modelado de lenguaje, logró una reducción del 50% en el tamaño del caché KV con solo una degradación del 3.1% en la perplejidad (para modelos de 300M) y del 2.48% (para modelos de 1.2B). Mantuvo un rendimiento competitivo en tareas de visión y sintéticas.
• Q=K-V (Limitación de Simetría): Aunque esta variante logró una calidad de entrenamiento comparable a la línea base, ofreció cero beneficios de memoria de inferencia porque aún requiere el almacenamiento de tensores K y V por separado. Además, los mapas de atención simétricos limitaron su efectividad en el modelado de lenguaje causal, resultando en una degradación de la perplejidad del 4.9%.
• Q=K=V (Sobre-restringido): Esta simplificación agresiva causó una pérdida catastrófica de rendimiento (25.4% de degradación de la perplejidad en el modelado de lenguaje), indicando que forzar una sola proyección para los tres roles es demasiado restrictivo para tareas complejas.

2. Sinergia con la Compartición de Cabezales

El estudio demuestra que la compartición de proyecciones y la compartición de cabezales son estrictamente complementarias.

• Combinar Q-K=V con GQA-4 (4 grupos) resultó en una reducción del caché del 87.5% con una degradación de la perplejidad del 3.9%.
• Combinar Q-K=V con MQA (1 grupo) logró una reducción del caché del 96.9% con una degradación de la perplejidad del 4.8%.
• A la escala de 1.2B, la configuración Q-MQA redujo el caché KV en 67× (de 5.9 GB a 88 MB por usuario para 32k de contexto) manteniendo una paridad casi total con la línea base en la precisión de tareas descendentes.

3. Comportamiento de Escalamiento

El ranking relativo de las variantes se mantuvo estable a medida que el tamaño del modelo aumentó de 300M a 1.2B de parámetros. Notablemente, los modelos más grandes parecieron ser ligeramente más robustos a las restricciones de proyección, con la degradación de la perplejidad para Q-K=V bajando del 3.1% a 300M al 2.48% a 1.2B.

Significación y Reivindicaciones

El artículo afirma que la compartición de proyecciones es una instancia poco explorada de la vinculación de pesos (weight tying) que ofrece beneficios cuantificables directos en la memoria de inferencia, los cuales son particularmente valiosos para el despliegue en dispositivos periféricos (edge) y aplicaciones de contexto largo.

• Por qué funciona Q-K=V: Los autores argumentan que las claves y los valores pueden ocupar espacios representacionales similares, y la atención opera efectivamente en un régimen de bajo rango donde la distinción entre K y V es menos crítica que la asimetría entre Q y la representación compartida K-V. El análisis empírico mostró una alta similitud de coseno (0.73) entre las matrices de proyección estándar de K y V, sugiriendo una redundancia inherente.
• Por qué falla Q=K-V: Forzar $Q=K$ crea mapas de atención simétricos, lo que rompe la direccionalidad requerida para las dependencias causales en el modelado de lenguaje.
• Impacto Práctico: La contribución principal no es solo la reducción de parámetros (que es modesta, ~7-14% de reducción del tamaño total del modelo) sino la reducción del 50% en la memoria del caché KV. Esto permite:
  • 2× ventanas de contexto más largas para el mismo presupuesto de memoria.
  • 2× mayor rendimiento (usuarios concurrentes por GPU).
  • Reducciones significativas en los costos (por ejemplo, una reducción proyectada del 43% en los costos de servicio para un modelo de completado de código).

Los autores concluyen que, si bien ninguna arquitectura gana en todos los dominios, la configuración Q-K=V proporciona el mejor equilibrio entre calidad y eficiencia para tareas secuenciales, y su combinación con técnicas de compartición de cabezales permite la inferencia práctica en dispositivos para modelos de miles de millones de parámetros. El código está disponible públicamente para facilitar la investigación posterior.