Uniform Scaling Limits in AdamW-Trained Transformers

Resumen Técnico: Límites de Escalado Uniforme en Transformers Entrenados con AdamW

Enunciado del Problema
El artículo aborda el desafío teórico de comprender la dinámica de entrenamiento de Modelos de Lenguaje Grandes (LLM), específicamente Transformers, a medida que escalan en profundidad ($L$) y anchura (número de cabezas de atención, $H$). Aunque las leyes de escalado empírico impulsan a la industria a aumentar el tamaño de los modelos, una comprensión teórica rigurosa de cómo las decisiones de diseño afectan la generalización y la robustez en el límite de profundidad infinita y anchura infinita sigue siendo incompleta.

La literatura existente sobre la convergencia del descenso de gradiente hacia la optimalidad global a menudo se basa en modelos límite (por ejemplo, Ecuaciones Diferenciales Ordinarias Neuronales o límites de Campo Medio). Sin embargo, transferir resultados de estos límites a modelos prácticos requiere cotas de convergencia que se mantengan uniformemente sobre la entrada del modelo. Existe una brecha significativa en la literatura respecto a los Transformers entrenados con el optimizador AdamW. Trabajos anteriores a menudo asumen descenso de gradiente o requieren que los parámetros estén compartidos o dependientes entre capas, lo cual no refleja la práctica estándar de inicialización independiente e idénticamente distribuida (i.i.d.). Además, el crecimiento cúbico de los gradientes en la autoatención softmax (debido a los productos de matrices $Q, K, O, V$) complica las garantías teóricas, lo que a menudo conduce a cotas que crecen exponencialmente con los pasos de entrenamiento. Finalmente, las cotas de convergencia estándar a menudo dependen del número de tokens ($N$) mediante argumentos de recubrimiento (por ejemplo, el Teorema de la Integral de Entropía de Dudley), lo que limita su aplicabilidad a la convergencia uniforme sobre distribuciones de entrada arbitrarias.

Metodología
Los autores modelan la dinámica del estado oculto de un Transformer sin enmascarar, solo de atención, entrenado con AdamW, como un sistema de partículas interactuantes (IPS) acoplado a través del mecanismo de atención. El análisis procede mediante los siguientes pasos:

1. Formulación de Campo Medio: La dinámica discreta de los $N$ tokens y las $H$ cabezas se aproxima mediante un límite de tiempo continuo. Se demuestra que los estados ocultos y las variables retropropagadas (adyuntas) convergen a la solución de un sistema hacia adelante-hacia atrás de Ecuaciones Diferenciales Ordinarias (EDO).
2. EDO de McKean–Vlasov (MVODE): Cuando falta el enmascaramiento causal, el sistema límite se identifica como una EDO de McKean–Vlasov. El campo de velocidad de esta EDO está definido por la esperanza sobre la distribución de parámetros.
3. Manejo de AdamW y No Linealidad: Un componente crítico de la metodología es el tratamiento del optimizador AdamW. A diferencia del SGD estándar, AdamW desacopla la decadencia de pesos. Los autores aprovechan esta propiedad para demostrar que el soporte de las medidas de parámetros permanece dentro de un conjunto compacto uniformemente a través de todos los pasos de entrenamiento. Esto permite constantes de Lipschitz globales para el mecanismo de atención y sus derivadas, evitando el crecimiento exponencial de cotas típicamente asociado con el crecimiento cúbico de gradientes en la autoatención.
4. Convergencia Uniforme sin Recubrimiento: Para lograr cotas independientes del número de tokens ($N$), los autores evitan los argumentos de recubrimiento tradicionales. En su lugar, utilizan:
   • Concentración de la Medida: Aplicando técnicas para acotar la norma $L^1$ de las diferencias de martingala asociadas con la inicialización de parámetros.
   • Aproximación de Taylor por Bloques: Dividiendo las capas en bloques y aproximando los mapas de flujo dentro de cada bloque mediante expansiones de Taylor.
   • Estructura Tensorial: Explotando la estructura algebraica específica del mecanismo de atención (expresada como ratios de funciones lineales de medidas) para mostrar que el supremo sobre las condiciones iniciales puede reducirse a un espacio de dimensión finita, eliminando así la dependencia de $N$.
5. AdamW por Bloques: Los autores introducen una variante de AdamW (AdamW por Bloques) que regulariza bloques de parámetros por separado. Esta adaptación permite que las cotas derivadas sean independientes de la dimensión de incrustación de tokens ($d$) también.

Contribuciones Clave

• Teorema de Convergencia Uniforme: El resultado principal (Teorema 7) establece que el Transformer sin enmascarar, solo de atención, con parámetros inicializados i.i.d. desde una distribución con soporte compacto, converge en $L^2$ uniformemente sobre las condiciones iniciales a un sistema de partículas interactuantes de EDOs medias. La tasa de convergencia es $O(L^{-1} + L^{-1/3}H^{-1/2})$.
• Independencia del Número de Tokens: A diferencia de las cotas anteriores que escalan con $\sqrt{\log(N)}$ o dependen de $N$ mediante argumentos de recubrimiento, las cotas derivadas son independientes del número de tokens $N$. Esto implica que los resultados se mantienen para condiciones iniciales no presentes en los datos de entrenamiento.
• Adaptación a AdamW: El trabajo extiende los resultados de convergencia de gran profundidad a Transformers entrenados con AdamW. Demuestra que la decadencia de pesos desacoplada de AdamW asegura que los parámetros permanezcan en un conjunto compacto, proporcionando continuidad Lipschitz global en el entrenamiento para el mecanismo de atención.
• Independencia de Dimensión: Al utilizar AdamW por Bloques, se demuestra que las cotas son independientes de la dimensión de incrustación de tokens $d$.
• Derivadas Explícitas: El artículo calcula explícitamente las derivadas de Lions para el mecanismo de atención sobre el componente de medida, estableciendo continuidad Lipschitz local con respecto a la distancia $p$-Wasserstein para todo $p \in [1, \infty]$.

Resultados
El artículo demuestra que, bajo un escalado apropiado de las cabezas de atención y con parámetros inicializados i.i.d., la dinámica conjunta de los estados ocultos y las variables retropropagadas converge a la solución de un sistema hacia adelante-hacia atrás de EDOs.

• Tasa de Convergencia: El error entre el modelo discreto de Transformer y el límite continuo de campo medio está acotado por $C(L^{-1} + L^{-1/3}H^{-1/2})$, donde $C$ es una constante independiente de $N$ y, bajo AdamW por Bloques, independiente de $d$.
• Uniformidad: La convergencia es uniforme sobre conjuntos compactos de condiciones iniciales.
• Cotas de Parámetros: Los autores muestran que las actualizaciones de AdamW confinan los parámetros a un conjunto compacto $S(R_\theta)$ uniformemente sobre todos los pasos de entrenamiento, una propiedad crucial para establecer constantes de Lipschitz globales.

Significado y Afirmaciones
El artículo afirma proporcionar una base teórica rigurosa para comprender el comportamiento de escalado de los Transformers entrenados con el optimizador dominante, AdamW. El significado radica en:

1. Cerrar la Brecha: Conecta la dinámica de entrenamiento discreta de Transformers prácticos y a gran escala con los límites continuos de campo medio, ofreciendo un entorno manejable para el análisis.
2. Robustez ante el Tamaño de Entrada: Al establecer cotas independientes del número de tokens, los resultados sugieren que el comportamiento teórico del modelo es intrínseco a su arquitectura e inicialización, en lugar de ser artefactos del tamaño específico del conjunto de datos. Esto es particularmente relevante para la generalización a datos no vistos.
3. Estabilidad de Optimización: El análisis destaca el papel estabilizador de la decadencia de pesos desacoplada de AdamW en el control del crecimiento de parámetros, lo cual es esencial para las garantías teóricas de continuidad Lipschitz en presencia de mecanismos de atención no lineales.

Los autores reconocen una compensación: la tasa de convergencia $O(L^{-1} + L^{-1/3}H^{-1/2})$ es ligeramente más lenta que la tasa $O(L^{-1} + (LH)^{-1/2})$ observada en algunos estudios de ResNet, un sacrificio realizado para lograr cotas uniformes independientes de $N$. Señalan que extender estos resultados al enmascaramiento causal (estándar en el modelado de lenguaje) y a regímenes de escalado de EDEs (Ecuaciones Diferenciales Estocásticas) sigue siendo un área abierta para futuras investigaciones.