Resumen Técnico: Revisitando la Convención de Forma de los Modelos de Lenguaje Transformer

Declaración del Problema

Los modelos de lenguaje Transformer densos han adherido históricamente a una convención arquitectónica rígida: la Red de Alimentación hacia Adelante (FFN, por sus siglas en inglés) utiliza una forma de Perceptrón Multicapa (MLP) "estrecho–ancho–estrecho". En este diseño, la dimensión del modelo ($d_{model}$) se expande a una dimensión intermedia más ancha ($d_h$) con una relación de expansión que típicamente oscila entre 2 y 4, antes de proyectar de vuelta. Esta estructura domina la asignación de parámetros, representando a menudo la mayor parte de los parámetros y las operaciones de punto flotante del modelo.

Los autores cuestionan si esta forma específica es esencial o simplemente una instanciación conveniente. El trabajo teórico reciente sugiere que los MLPs residuales tipo "ancho–estrecho–ancho" (reloj de arena o hourglass) pueden servir como aproximadores de funciones óptimos, ofreciendo potencialmente una expresividad superior con menos parámetros. Además, la evidencia empírica en modelos de visión generativa y Transformers de Mezcla de Expertos (MoE) ha demostrado que las estructuras de reloj de arena pueden superar a los diseños estándar. El problema central abordado es si la convencional FFN estrecha–ancha–estrecha es una restricción necesaria o si una FFN de reloj de arena más profunda y estrecha puede ofrecer una asignación de capacidad más eficiente entre los módulos de atención y FFN.

Metodología

El artículo propone una variante de Transformer que reemplaza la FFN convencional con una FFN de Reloj de Arena (Hourglass FFN). La metodología consiste en:

1. Rediseño Arquitectónico:

   • La FFN de Reloj de Arena reemplaza la capa de expansión única y ancha por una pila de $K$ sub-MLPs con forma de reloj de arena conectados por vías residuales.
   • Cada sub-MLP ($MLP_i$) sigue una estructura de compresión-expansión: proyecta la entrada desde $d_{model}$ hacia una dimensión de cuello de botella $d_h$ (donde $d_h < d_{model}$), aplica una activación no lineal (SwiGLU) y proyecta de vuelta a $d_{model}$.
   • La salida de la capa es el resultado de pasar la entrada a través de $K$ de estos sub-bloques residuales.

2. Estrategia de Redistribución de Parámetros:

   • Al reducir la dimensión intermedia $d_h$ (haciendo la FFN más "estrecha"), los autores ahorran parámetros que típicamente se asignarían al ancho de la FFN.
   • Estos parámetros ahorrados se reasignan ya sea a:
     • Aumentar la dimensión del modelo ($d_{model}$) y la capacidad de atención.
     • Aumentar la profundidad de la pila FFN (el valor de $K$).
     • Aumentar el número de capas del Transformer ($L$).

3. Configuración Experimental:

   • Líneas Base (Baselines): Se realizan comparaciones contra Transformers de estilo LLaMA estándar y la arquitectura de vanguardia OLMo-2.
   • Escalas: Los experimentos se llevan a cabo en tamaños de modelo que van desde 113M hasta 1B de parámetros.
   • Restricciones: Los experimentos se realizan bajo presupuestos de parámetros totales igualados (con una diferencia de máximo 0.001%) y tamaños de atención fijos para aislar los efectos de la forma de la FFN.
   • Espacio de Búsqueda: Los autores realizan una búsqueda arquitectónica variando $K$ (profundidad de la pila de reloj de arena), $L$ (número de capas de Transformer), $d_{model}$ y la relación $d_h/d_{model}$.

Contribuciones Clave

El artículo realiza cuatro contribuciones principales:

1. Desafío al Estándar: Desafía la forma de facto de la FFF "estrecho–ancho–estrecho", proponiendo una FFN de Reloj de Arena "ancho–estrecho–ancho" eficiente que mantiene la expresividad mediante sub-MLPs apilados y conectados residualmente.
2. Reasignación de Parámetros: Demuestra una estrategia donde redirigir parámetros del ancho de la FFN hacia el ancho de la atención o la profundidad interna de la FFF ($K > 1$) produce un rendimiento superior, particularmente en escalas pequeñas (113M).
3. Perspectivas Arquitectónicas: Identifica un compromiso (trade-off) de ancho-profundidad en forma de U, encontrando proporciones óptimas para $d_{model}/L$ y $d_h/d_{model}$, y establece que las estructuras de FFN de Reloj de Arena más profundas son robustas.
4. Validación de Escalabilidad: Valida la FFN de Reloj de Arena desde 113M hasta 1B de parámetros, mostrando que es una alternativa competitiva a las líneas base estándar, con configuraciones específicas que superan a los modelos convencionales en escalas menores y se igualan a ellos en escalas mayores.

Resultados

La validación empírica arroja los siguientes hallazgos:

• Rendimiento con Atención Fija: A la escala de 113M con parámetros de atención fijos, las FFN de Reloj de Arena logran una pérdida de validación y una perplejidad (PPL) más bajas que las FFN convencionales. Por ejemplo, con $K=8$, la PPL cayó de 36.441 (convencional) a 36.179.
• Configuraciones Óptimas:
  • Profundidad vs. Ancho: Una relación de $d_h/d_{model} \approx 0.4$ con $K=4$ capas produjo la menor pérdida de validación.
  • Desplazamiento de Parámetros: A 113M, desplazar los parámetros para que el módulo de atención sea dominante (por ejemplo, 66M de atención frente a 47M de FFN) resultó en el mejor rendimiento (PPL 35.101), lo que sugiere que el procesamiento contextual es más valioso que la recuperación de hechos estáticos para modelos pequeños.
  • Compromiso Ancho-Profundidad: Se observó una curva distintiva en forma de U para la relación $d_{model}/L$, con un punto óptimo entre 100 y 250.
• Escalabilidad:
  • Escalas Pequeña a Mediana (113M – 906M): Las variantes de Reloj de Arena superaron consistentemente a las líneas base convencionales en pérdida de validación, perplejidad y tareas de razonamiento descendentes (ej. Arc Easy, HellaSwag, PIQA).
  • Escala Grande (1B): A la escala de 1B de parámetros, la variante de Reloj de Arena logró un rendimiento comparable a la fuerte línea base convencional (Val PPL 20.082 frente a 20.002) y superó a la configuración de OLMo-2 (20.466).
• Estudios de Ablación: Aumentar la profundidad del reloj de arena $K$ mejoró consistentemente el rendimiento, aunque con rendimientos decrecientes más allá de $K=6$. Reducir $d_h/d_{model}$ por debajo de 0.4 provocó una degradación significativa del rendimiento, lo que indica un límite inferior para la compresión de la FFN.

Significado y Reivindicaciones

El artículo sostiene que la forma de MLP "estrecho–ancho–estrecho" de larga data en los Transformers no es esencial y puede ser redundante. Al adoptar una FFN de Reloj de Arena, los autores demuestran que:

1. Eficiencia: Los modelos pueden lograr un rendimiento competitivo o superior con una distribución de parámetros más eficiente.
2. Reasignación de Recursos: La estructura de Reloj de Arena permite un cambio fundamental donde el mecanismo de atención puede convertirse en el consumidor dominante de parámetros, desafiando el presupuesto tradicional dominado por la FFN.
3. Reevaluación del Diseño: Los hallazgos incitan a repensar el equilibrio entre la atención y la FFN, sugiriendo que para modelos de lenguaje modernos, eficientes y expresivos, una FFN más profunda y estrecha combinada con mecanismos de atención más anchos puede ser una estrategia de diseño superior.

Los autores mantienen la modestia respecto a la escalabilidad hacia modelos de decenas o cientos de miles de millones de parámetros, señalando que sus hallazgos fueron extrapolados de la escala de 113M y que la "maldición de la profundidad" (problemas de propagación de señal) requiere una investigación más profunda en arquitecturas más profundas. También señalan que la interacción entre las FFN de Reloj de Arena y los mecanismos de atención avanzados (ej. Atención de Consulta por Grupo o GQA) aún no ha sido explorada.