Entrenar transformers grandes de forma estable ha sido un reto constante en el aprendizaje profundo, sobre todo a medida que los modelos aumentan en tamaño y complejidad. Investigadores del MIT han abordado un problema clave: el crecimiento inestable de las activaciones y los picos en la pérdida, que se deben a normas no controladas de pesos y activaciones. Su propuesta consiste en imponer límites de Lipschitz comprobables directamente sobre el transformer, regulando espectralmente los pesos sin recurrir a normalizaciones de activación, ni trucos como la normalización QK o el recorte suave en los logits.
Pero, ¿qué es un límite de Lipschitz y por qué es importante? Un límite de Lipschitz establece un máximo para cuánto puede cambiar la salida de una red ante pequeñas variaciones en la entrada o en los pesos. En términos simples, si la función que representa un modelo es K-Lipschitz, entonces la diferencia entre dos salidas no puede ser mayor que K veces la diferencia entre las entradas correspondientes. Esto significa que un límite bajo implica mayor robustez y estabilidad, clave para que el modelo sea predecible, resistente a ataques adversarios, respete la privacidad y generalice mejor.
Históricamente, entrenar transformers estables a gran escala ha requerido “parches” como normalización por capas, normalización QK o recortes tanh en los logits. Sin embargo, estos trucos no solucionan el problema raíz, que es el crecimiento de la norma espectral de los pesos —es decir, el valor singular más grande— que provoca que las activaciones exploten y la formación sea inestable, especialmente en modelos grandes. La hipótesis principal del MIT es que, si regulamos directamente la norma espectral de los pesos en cada paso de entrenamiento, se puede controlar estrictamente la constante de Lipschitz y resolver la inestabilidad desde su origen.
Para ello, proponen dos innovaciones clave: Primero, el optimizador Muon regula espectralmente los gradientes para evitar que cada actualización aumente la norma espectral más allá de un límite establecido. Segundo, aplican una restricción directa sobre los pesos después de cada paso, limitando sus valores singulares y manteniendo así las activaciones sorprendentemente pequeñas —en sus transformers estilo GPT-2, las activaciones casi nunca superan valores compatibles con precisión fp8.
Lo impresionante es que, en todas sus pruebas, no emplearon normalización por capas, normalización QK ni recortes en los logits, y aun así lograron una estabilidad notable. Por ejemplo, las activaciones máximas en su transformer escala GPT-2 rondaron los 100, mientras que un modelo sin estas restricciones llegó a superar 148,000.
Probando distintas técnicas para imponer estas restricciones, como la decaimiento de pesos tradicional, normalización espectral o sus métodos novedosos de "limitación espectral suave" y "martillo espectral", lograron optimizar el equilibrio entre performance y control estrictos de Lipschitz. La limitación espectral suave es especialmente destacable por aplicar de forma eficiente un tope a todos los valores singulares usando aproximaciones polinómicas, funcionando muy bien junto con las actualizaciones de rango estable alto del optimizador Muon.
En sus experimentos, con modelos pequeños como el de Shakespeare (menos de 2-Lipschitz), alcanzaron un 60% de acierto en validación con límites de Lipschitz comprobables, superando modelos sin restricción. En modelos más grandes, como NanoGPT con 145 millones de parámetros, limitar Lipschitz estrictamente redujo la precisión, lo que señala que mantener expressividad y estabilidad al mismo tiempo sigue siendo un desafío en modelos grandes.
Con la combinación Muon + limitación espectral, lograron la mejor relación entre pérdida y Lipschitz, superando métodos estándar como AdamW con decaimiento de peso. Además, observaron que al reducir este límite, la robustez ante ataques adversarios mejoraba significativamente, con caídas de precisión mucho menores.
Otro beneficio clave de su enfoque es que las activaciones permanecen muy pequeñas, abriendo la puerta para entrenar e inferir con baja precisión numérica, lo que reduce costos de computación, memoria y energía en hardware especializado.
Sin embargo, aún quedan preguntas abiertas: elegir la mejor combinación de restricciones para pesos, escala de logits y de atención aún requiere pruebas exhaustivas. Los límites calculados pueden ser demasiado conservadores y no reflejar la verdadera magnitud de las activaciones. Tampoco está claro si a gran escala se podrá lograr el mismo rendimiento que los modelos sin límite de Lipschitz, aunque este trabajo es un paso importante hacia allí.
En resumen, regular espectralmente los pesos con el optimizador Muon permite entrenar transformers grandes de forma estable, sin depender de trucos tradicionales de normalización. Esto mejora la robustez, facilita un rango predecible de activaciones y abre posibilidades para IA eficiente en hardware de baja precisión. Esta línea de investigación apunta a nuevas herramientas para controlar redes neuronales, con aplicaciones en privacidad, seguridad y despliegue de IA ligera.
