Resumen Técnico: Leyes de Escalamiento Espectral Inducidas por Optimizadores

Enunciado del Problema

Las leyes de escalamiento clásicas para modelos de lenguaje grandes (LLM) han establecido relaciones predecibles de ley de potencia entre la pérdida de validación y recursos como el tamaño del modelo, el volumen de datos y la capacidad de cómputo. Sin embargo, estos marcos suelen tratar al optimizador como un detalle fijo de entrenamiento, asumiendo que solo influye en la velocidad de convergencia o actúa como un factor de eficiencia multiplicativo. Este trabajo desafía esa suposición, postulando que el optimizador es un eje fundamental del escalamiento de representaciones. Específicamente, los autores investigan cómo convierten eficazmente diferentes optimizadores el aumento del ancho de la Red de Alimentación Directa (FFN) en capacidad espectral utilizada —la dimensionalidad efectiva del espacio latente—. Se preguntan si una pérdida de validación coincidente implica una geometría de representación interna coincidente, o si diferentes optimizadores realizan la misma capacidad arquitectónica a través de estructuras espectrales fundamentalmente distintas.

Metodología

El estudio emplea un marco de análisis espectral centrado en los espectros de autovalores de las representaciones de la FFN dentro de arquitecturas Transformer.

• Métricas Espectrales: Los autores utilizan la familia de rangos efectivos de Rényi para cuantificar la dispersión espectral. Anclan su análisis en dos métricas principales:
  • Rango Blando ($R_1$): Corresponde a $\alpha=1$ (entropía de Shannon), midiendo la dispersión espectral ponderada por entropía y capturando la capacidad difusa.
  • Rango Duro ($R_2$): Corresponde a $\alpha=2$ (relación de participación), midiendo la dimensionalidad efectiva sensible a la concentración y capturando la capacidad de modo dominante.
  • Asimetría Duro-Blando ($\Delta_{1,2}$): Definida como $\log R_1 - \log R_2$, esta métrica cuantifica la brecha entre la capacidad difusa y la dominante.
• Configuración Experimental: Los experimentos se realizaron en Transformers de solo decodificador estilo GPT (160M y 350M parámetros) entrenados en el conjunto de datos FineWeb-Edu. La dimensión oculta de la FFN ($D$) se varió para establecer leyes de escalamiento ($R(D) \propto D^\beta$).
• Comparación de Optimizadores: El estudio compara el AdamW estándar contra optimizadores conscientes de matrices (Muon, NorMuon) y variantes con restricción de rango (Dion con fracciones de rango $r$ variables).
• Estratificación: Para abordar la distribución Zipfiana de los datos lingüísticos, las representaciones de la FFN se estratificaron por frecuencia de token en regímenes HEAD (frecuente), MID y TAIL (raro).
• Controles Arquitectónicos: Los autores compararon los desplazamientos inducidos por el optimizador contra intervenciones arquitectónicas, específicamente aumentar el rango de atención por cabeza (reduciendo el número de cabezas) y eliminar las Incrustaciones Posicionales Rotatorias (RoPE).

Resultados Clave

1. Leyes de Escalamiento Espectral Dependientes del Optimizador

La misma arquitectura Transformer exhibe leyes de escalamiento espectral marcadamente diferentes dependiendo del optimizador.

• Escalamiento de Rango Duro: AdamW exhibe un escalamiento de rango duro débil ($\beta_{hard} \approx 0.29$ a $0.44$), particularmente para tokens raros (TAIL). En contraste, Muon logra un escalamiento casi lineal ($\beta_{hard} \approx 1.02$) en los mismos regímenes. Esto representa un aumento de 2.3$\times$ en el exponente de escalamiento.
• Asimetría Duro-Blando: AdamW muestra una gran asimetría positiva ($\Delta_{1,2} \approx 0.37$), indicando que el ancho adicional de la FFN expande la masa espectral difusa en lugar de las direcciones de representación dominantes. Muon y NorMuon reducen esta asimetría a cerca de cero ($\Delta_{1,2} \approx 0.14$), lo que significa que el ancho añadido se convierte eficientemente en capacidad de modo dominante.

2. Pérdida Coincidente $\neq$ Geometría Coincidente

El trabajo demuestra que la pérdida de validación es un proxy insuficiente para la estructura de representación.

• Entrenamiento Extendido: Entrenar AdamW durante más pasos (12K vs. 6K) le permitió igualar la perplejidad de una variante de Dion (1/16). Sin embargo, la geometría espectral permaneció distinta: el escalamiento de rango duro de AdamW casi desapareció ($\beta_{hard} \to 0.03$), mientras que Dion mantuvo un escalamiento de ley de potencia robusto.
• Tasa de Aprendizaje: El barrido de tasas de aprendizaje confirmó que la brecha espectral AdamW-Muon no es un artefacto de hiperparámetros; el exponente de rango duro más bajo de Muon superó consistentemente al más alto de AdamW en todo el barrido.

3. Dependencias de Frecuencia y Rango

• Frecuencia de Token: El efecto del optimizador es más pronunciado en los regímenes de tokens MID y TAIL. AdamW lucha por escalar la capacidad dura para tokens raros, mientras que Muon escala linealmente.
• Rango de Actualización: Usando Dion, los autores aislaron el efecto del rango de actualización. Reducir la fracción de rango ($r$) de $1/2$ a $1/16$ hizo que el escalamiento de rango duro cayera de $\beta \approx 0.88$ a $\beta \approx 0.40$ (acercándose a los niveles de AdamW), mientras que el escalamiento de rango blando se degradó más gradualmente. Esto indica que el rango de actualización actúa como un control para convertir el ancho en capacidad de modo dominante.

4. Dominio sobre Intervenciones Arquitectónicas

Los desplazamientos espectrales inducidos por el optimizador a menudo exceden o remodelan los efectos de los cambios arquitectónicos:

• Rango de Atención: Aumentar el rango de atención por cabeza (reduciendo el número de cabezas) produjo desplazamientos de escalamiento espectral menores que cambiar de optimizador en 28 de 30 regímenes de frecuencia/métrica.
• Codificación Posicional: Eliminar RoPE indujo una redistribución dependiente del optimizador. Para Muon, NoPE aumentó el escalamiento de rango duro HEAD pero disminuyó el escalamiento MID/TAIL; para NorMuon, disminuyó el escalamiento en todos los regímenes. Esto muestra que las señales arquitectónicas se expresan a través de la geometría del optimizador.

Significado y Afirmaciones

El trabajo argumenta que la optimización debe tratarse como un eje de primer orden del escalamiento de representaciones, junto con el tamaño del modelo, los datos y el cómputo.

• Co-diseño Arquitectura-Optimizador: La capacidad espectral realizada es una propiedad del par arquitectura-optimizador, no de la arquitectura sola. Las intervenciones arquitectónicas (como cambiar el rango de atención o las codificaciones posicionales) no inducen desplazamientos fijos e independientes del optimizador; sus efectos están mediados por la geometría del optimizador.
• Estructura de Representación: La elección del optimizador determina cómo se estructura la capacidad del modelo a través de los modos propios y cómo se asigna a través de los regímenes de frecuencia de token. No es meramente un control de eficiencia para la velocidad de convergencia.
• Limitaciones: Los autores notan que estos hallazgos se establecen dentro de un régimen empírico controlado (modelos estilo GPT de hasta 350M parámetros). Aunque los resultados de 350M replican los patrones de 160M, se necesita una validación adicional en escalas de vanguardia (1B+). Además, aunque los rangos espectrales cuantifican la utilización de la capacidad, no caracterizan completamente el comportamiento de tareas aguas abajo.

En conclusión, el estudio motiva un cambio hacia el co-diseño arquitectura-optimizador, sugiriendo que las futuras leyes de escalamiento deben tener en cuenta los sesgos inductivos implícitos del optimizador como un determinante primario de cómo se realiza la capacidad del modelo.