El optimizador Muon del MIT estabiliza Transformers con límites de Lipschitz
Entrenar modelos Transformer a gran escala de manera estable ha sido durante mucho tiempo un desafío significativo en el aprendizaje profundo, particularmente a medida que estos modelos continúan creciendo en tamaño y complejidad. Investigadores del MIT han abordado un problema fundamental: el crecimiento descontrolado de los valores de activación y los picos resultantes en la pérdida durante el entrenamiento, a menudo causados por normas de peso y activación sin restricciones.
Su innovadora solución implica la imposición de “límites de Lipschitz demostrables” en los transformadores. Esto se logra regulando directamente las propiedades espectrales de los pesos del modelo, sin depender de técnicas de estabilización comunes como la normalización de activación, la normalización QK o el softcapping de logits.
Comprendiendo los límites de Lipschitz y su importancia
Un límite de Lipschitz en una red neuronal cuantifica la tasa máxima a la que la salida de la red puede cambiar en respuesta a perturbaciones en su entrada o pesos internos. En términos más simples, un límite de Lipschitz más bajo indica que la red es menos sensible a pequeños cambios o ruido, lo que la hace más robusta y predecible. Esta propiedad es crucial para garantizar la estabilidad durante el entrenamiento, mejorar la robustez adversarial (resistencia a manipulaciones maliciosas de entrada), mejorar la privacidad y promover una mejor generalización a nuevos datos.
El problema con los métodos de estabilización tradicionales
Históricamente, lograr la estabilidad en grandes transformadores ha implicado una variedad de soluciones “parche”, como la normalización de capas (Layer Normalization), la normalización QK (QK Normalization) y el softcapping Logit Tanh (Logit Tanh Softcapping). Si bien estos métodos ofrecen cierta estabilidad, no abordan directamente la causa subyacente de la inestabilidad: el crecimiento descontrolado de la “norma espectral” (el valor singular más grande) dentro de las matrices de peso. Este crecimiento sin restricciones es un motor principal de las activaciones explosivas y la inestabilidad del entrenamiento, especialmente en modelos muy grandes.
La hipótesis central del equipo del MIT es que al regular directamente las propiedades espectrales de los pesos mismos —yendo más allá del optimizador o las activaciones— pueden mantener un control estricto sobre la Lipschitz de la red, abordando así la inestabilidad en su origen.
Innovaciones clave: Optimizador Muon y Regulación Espectral de Pesos
El enfoque de los investigadores se basa en el optimizador “Muon”, que ya regulariza espectralmente los gradientes, asegurando que cada paso de gradiente no aumente la norma espectral más allá de un límite establecido. La innovación clave del MIT extiende esta regulación a los pesos del modelo: después de cada paso de entrenamiento, aplican operaciones para limitar los valores singulares de cada matriz de peso. Los valores singulares son componentes matemáticos que describen cuánto una matriz estira o encoge las entradas; limitarlos controla directamente el factor de amplificación de los pesos.
Un resultado notable de esta regulación de pesos es que las normas de activación —las magnitudes de los valores dentro de las capas de la red— permanecen excepcionalmente pequeñas. En sus transformadores a escala GPT-2, las entradas de activación máximas nunca excedieron aproximadamente 100. Esto contrasta fuertemente con las líneas base sin restricciones, donde las activaciones máximas podrían superar los 148,000. Crucialmente, esta estabilidad se logró sin usar ninguna de las normalizaciones de capa tradicionales, la norma QK o los trucos de logit tanh. Las pequeñas magnitudes de activación también hacen que estos modelos sean compatibles con formatos de datos de baja precisión como FP8, lo cual es muy beneficioso para una implementación eficiente en hardware.
Métodos para imponer restricciones de Lipschitz
Los investigadores exploraron y compararon varios métodos para imponer restricciones de norma de peso, evaluando su capacidad para mantener un alto rendimiento, garantizar un límite de Lipschitz y optimizar el equilibrio entre rendimiento y Lipschitz:
Decaimiento de Peso (Weight Decay): Un método de regularización estándar, pero no siempre preciso en el control de la norma espectral.
Normalización Espectral (Spectral Normalization): Limita el valor singular más grande de una matriz de peso, pero puede afectar todos los valores singulares globalmente.
Soft Cap Espectral (Spectral Soft Cap): Una técnica novedosa que limita de manera suave y eficiente todos los valores singulares en paralelo. Este método fue co-diseñado específicamente para funcionar eficazmente con las actualizaciones de rango estable de Muon, permitiendo límites más estrictos.
Martillo Espectral (Spectral Hammer): Un método que establece solo el valor singular más grande a un máximo, más adecuado para usar con el optimizador AdamW.
Resultados experimentales y perspectivas
La investigación demostró hallazgos significativos en varias escalas de modelos:
Evaluación del Modelo: Para transformadores más pequeños (como Shakespeare, con un límite de Lipschitz demostrable por debajo de 2), el método logró una precisión de validación del 60% y superó a las líneas base sin restricciones en la pérdida de validación. Para modelos más grandes como NanoGPT (145M parámetros), un límite de Lipschitz estricto de menos de 10 produjo una precisión de validación del 21.2%. Para igualar el rendimiento de una línea base fuerte sin restricciones (39.4% de precisión), se requirió un límite superior mucho mayor (por ejemplo, 10^264). Esto destaca el equilibrio actual entre las restricciones de Lipschitz muy estrictas y el logro de la expresividad máxima a escalas mayores.
Eficiencia de los métodos de restricción: La combinación del optimizador Muon con Spectral Soft Cap lideró consistentemente la frontera en el equilibrio pérdida-Lipschitz, logrando constantes de Lipschitz más bajas con una pérdida de validación comparable o mejor en comparación con AdamW con decaimiento de peso.
Estabilidad y Robustez: Los modelos entrenados con una constante de Lipschitz restringida mostraron una robustez adversarial significativamente mayor, experimentando caídas de precisión mucho más leves bajo ataques adversariales en comparación con las líneas base sin restricciones.
Magnitudes de Activación: Como se señaló, la regulación espectral de pesos mantuvo las activaciones máximas consistentemente pequeñas, incluso a escala. Esto abre nuevas vías para el “entrenamiento e inferencia de baja precisión” en hardware, donde activaciones más pequeñas pueden reducir drásticamente los costos computacionales, de memoria y de energía.
Limitaciones y Direcciones Futuras
A pesar de estos avances, la investigación identifica varias preguntas abiertas y limitaciones:
La selección del equilibrio óptimo entre las normas de peso, el escalado de logits y el escalado de atención aún depende en gran medida de barridos empíricos en lugar de métodos basados en principios.
Los límites de Lipschitz globales actuales calculados para los modelos pueden ser astronómicamente grandes (por ejemplo, 10^264), incluso cuando las normas de activación reales permanecen muy pequeñas. Esto indica que los límites teóricos son a menudo mucho más flojos que el comportamiento observado.
Sigue sin estar claro si igualar el rendimiento de las líneas base sin restricciones con límites de Lipschitz estrictamente pequeños es factible a medida que la escala del modelo continúa aumentando. Se necesita más investigación en esta área.
Conclusión
El trabajo de los investigadores del MIT demuestra que la regulación espectral de pesos, particularmente cuando se integra con el optimizador Muon, proporciona un método poderoso para entrenar de manera estable grandes transformadores con límites de Lipschitz impuestos. Este enfoque elimina la necesidad de la normalización de activación tradicional y otros trucos de estabilización ad-hoc, abordando la inestabilidad a un nivel más profundo y fundamental. Al mantener las activaciones dentro de un rango compacto y predecible, el método mejora significativamente la robustez adversarial y ofrece un potencial sustancial para mejorar la eficiencia del hardware a través de la implementación de IA de baja precisión. Esta investigación allana el camino para nuevas y eficientes primitivas computacionales para la regulación de redes neuronales, con amplias implicaciones para la seguridad, privacidad y despliegue práctico de sistemas de IA avanzados.