Tapered Language Models

Resumen Técnico: Modelos de Lenguaje Cónicos (Tapered Language Models)

Planteamiento del Problema
Los modelos de lenguaje (LM) modernos, incluyendo los Transformers, sus variantes recurrentes y las arquitecturas basadas en memoria, comparten un defecto estructural común: una pila de capas idénticas donde los parámetros se asignan de manera uniforme a lo largo de la profundidad. A pesar de la creciente evidencia que sugiere que las capas contribuyen de forma no uniforme al resultado final —donde las capas iniciales suelen capturar patrones sintácticos superficiales y las capas posteriores refinan el flujo residual en lugar de transformarlo—, esta asignación uniforme sigue siendo el estándar. El artículo postula que este enfoque de "talla única" no logra reflejar la asimetría de la importancia de las capas, lo que potencialmente desperdicia la capacidad de los parámetros en las capas posteriores donde estos se utilizan menos.

Metodología
Los autores introducen los Modelos de Lenguaje Cónicos (TLMs), un principio arquitectónico donde un componente portador de parámetros se reduce de forma monótona a través de la profundidad bajo un presupuesto total fijo.

• Componente Objetivo: El estudio instancia este principio en el ancho intermedio del MLP ($d_{ff}$). Se eligen los MLPs porque dominan el recuento de parámetros en los LM modernos y exponen el ancho como un eje de variación limpio e independiente del módulo de mezcla de tokens.
• Esquemas de Reducción (Tapering Schedules): En lugar de asignaciones discretas por bloques, los autores proponen esquemas de decaimiento suaves y continuos para $d_{ff}(l)$ a través de las capas $l \in \{0, \dots, L-1\}$. Se evalúan tres esquemas:
  1. Lineal: Tasa de decaimiento constante.
  2. Coseno: Sigue una curva de medio coseno, decayendo lentamente en los extremos y abruptamente en el medio.
  3. Sigmoide: Concentra la transición en una banda estrecha alrededor del punto medio.
• Restricciones: La reducción es estrictamente preservadora del presupuesto. El ancho promedio por capa se restringe para ser igual al ancho uniforme base ($\frac{1}{L}\sum d_{ff}(l) = d_{ff}^{baseline}$), asegurando que el recuento total de parámetros y los FLOPs (tanto de entrenamiento como de inferencia) sean idénticos al baseline uniforme.
• Configuración Experimental: El principio se evalúa en cuatro arquitecturas distintas (Transformer Estándar, Atención con Compuerta [Gated Attention], Hope-attention y Titans) y tres escalas de modelo (440M, 760M y 1.3B de parámetros). El estudio utiliza una configuración de entrenamiento fija (tokenizador Llama 3, longitud de secuencia de 4K) y compara los modelos cónicos contra los baselines uniformes en perplejidad de validación in-distribución, perplejidad fuera de distribución (WikiText, LAMBADA) y ocho bancos de pruebas de razonamiento de sentido común.

Resultados Clave

1. La Dirección de la Asignación Importa: Un experimento motivador en un Transformer de 440M particionó las capas en grupos tempranos, medios y tardíos. Asignar más capacidad a las capas tempranas ("Wider-early") mejoró la perplejidad en 0.32 puntos respecto al baseline uniforme. Por el contrario, asignar más capacidad a las capas tardías ("Wider-late") degradó el rendimiento en más de 1.0 punto, demostrando que la dirección de la capacidad es crítica.
2. Esquema y Ratio Óptimos: Un barrido de ratios de ancho y esquemas en el modelo de 440M reveló que el esquema de Coseno superó consistentemente a los esquemas Lineal y Sigmoide. La configuración óptima se encontró en un ratio de inicio a fin de 1.50 $\to$ 0.50, logrando una perplejidad de validación de 14.44 (una mejora de 1.84 puntos sobre el baseline uniforme de 16.28).
3. Generalización: Aplicar la configuración óptima fija (Coseno, 1.5/0.5) a los modelos de 760M y 1.3B en las cuatro arquitecturas mejoró consistentemente la perplejidad y la precisión en tareas de razonamiento de sentido común. Las ganancias se observaron independientemente de si el módulo de mezcla de tokens utilizaba atención softmax, atención con compuerta o mecanismos recurrentes/basados en memoria.
4. Perspectiva Mecanicista: El análisis de los transformers preentrenados reveló que las salidas de los MLP se alinean progresivamente más con el flujo residual existente a medida que aumenta la profundidad (medido mediante similitud de coseno). Esto sugiere que los MLPs posteriores refuerzan el contenido existente en lugar de computar nuevas características. La reducción (tapering) alinea la arquitectura con este patrón al reducir la capacidad donde es menos útil (capas posteriores) y reasignarla a las capas anteriores donde es más ortogonal y útil.

Significancia y Reivindicaciones
El artículo establece la asignación de capacidad consciente de la profundidad como una palanca de diseño simple y agnóstica a la arquitectura para los modelos de lenguaje. Los autores afirman que la reducción del ancho de los MLP ofrece una mejora de rendimiento "gratuita": mejora la perplejidad y la precisión en tareas de seguimiento sin aumentar el recuento de parámetros o el costo computacional.

La significancia radica en desafiar la uniformidad heredada del chasis del Transformer. Los resultados sugieren que el principio de la reducción concierne a la asignación general de parámetros a través de la profundidad, más que a propiedades específicas de los mecanismos de atención. Aunque los autores señalan que el esquema óptimo específico (Coseno, 1.5/0.5) se derivó de un modelo de 440M y puede no ser el óptimo global para todas las escalas, las ganancias consistentes en diversas arquitecturas y tamaños indican que el principio es robusto. El artículo concluye que la asignación consciente de la profundidad es una "palanca oculta" a plena vista que puede aplicarse a modelos fundacionales más allá del modelado de lenguaje, incluyendo modelos de visión y multimodales.