Investigadores de Tilde Research han lanzado Aurora, un nuevo optimizador para entrenar redes neuronales que soluciona un defecto estructural en el optimizador Muon, ampliamente utilizado. La falla mata silenciosamente una fracción significativa de las neuronas MLP durante el entrenamiento y las mantiene permanentemente muertas. Aurora viene con un experimento de preentrenamiento de parámetros de 1.1B, un nuevo resultado de última generación en el punto de referencia de speedrun nanoGPT modificado y códigos abiertos.
¿Qué es el muón?
Para comprender a Aurora, es útil comprender primero a Muon. El optimizador Muon atrajo la atención de la comunidad de ML después de superar a AdamW en tiempo de reloj de pared para converger en la competencia speedrun nanoGPT, un punto de referencia de la comunidad que mide qué tan rápido se puede entrenar un modelo de estilo GPT hasta una pérdida de validación objetivo. Desde entonces, varios grupos de investigación han adoptado Muon en el entrenamiento de modelos a escala de frontera.
El paso algorítmico clave de Muon es calcular el factor polar de la matriz de gradiente. Para una matriz de gradiente G con descomposición de valores singulares (SVD) fina G = UΣVᵀ, Muon calcula polar(G) = UVᵀ, que es la matriz semiortogonal más cercana a G en la norma de Frobenius. Este gradiente ortogonalizado se utiliza luego para actualizar los pesos: W ← W − η UVᵀ para una tasa de aprendizaje η. El uso de algoritmos iterativos exclusivos de matmul para calcular el factor polar es lo que hace que Muon sea práctico a escala.
El rompecabezas de NorMuon: la normalización de filas ayuda, pero ¿por qué?
Antes de Aurora, NorMuon lideró el speedrun nanoGPT modificado. Introdujo un paso de normalización de filas, similar a la escala por parámetro de Adam, que ajustaba el factor polar según su norma RMS inversa. Si bien esto a menudo aleja la actualización de un gradiente estrictamente ortogonal, NorMuon aún produce resultados impresionantes. El equipo de Tilde se propuso comprender exactamente qué brecha en la formulación de Muon estaba abordando NorMuon.
El problema central: anisotropía de norma de fila y muerte neuronal en matrices altas
El equipo de investigación descubrió que el optimizador Muon “mata” involuntariamente una gran parte de las neuronas en matrices de alto peso, como las que se encuentran en las capas MLP basadas en SwiGLU. Debido a que es matemáticamente imposible que estas formas de matriz específicas permanezcan perfectamente ortogonales y al mismo tiempo mantengan uniformes las actualizaciones de las filas, el optimizador termina brindando actualizaciones masivas a algunas neuronas mientras prácticamente ignora otras. Esto da como resultado una “espiral de muerte” en la que las neuronas de bajo rendimiento reciben menos señales con el tiempo y eventualmente quedan permanentemente inactivas.
El estudio de investigación reveló que en el paso de entrenamiento número 500, más de una de cada cuatro neuronas está efectivamente muerta. Este no es sólo un problema local; la falta de actividad en estas neuronas priva de capas posteriores de datos necesarios, extendiendo la ineficiencia por todo el modelo. Aurora resuelve esto mediante el uso de un nuevo enfoque matemático que impone actualizaciones uniformes en todas las neuronas sin sacrificar los beneficios de la ortogonalización.
Antes de llegar a Aurora, la investigación introduce una solución intermedia llamada U-NorMuon. La observación clave es que NorMuon normaliza cada fila a la norma unitaria (norma = 1), pero en realidad este es el objetivo equivocado para una matriz alta. Para una matriz alta ortogonal de columnas, la norma de fila promedio matemáticamente correcta es √(n/m), no 1. U-NorMuon corrige esto normalizando las filas de la matriz alta para que tengan una norma √(n/m) en lugar de 1.
En experimentos a escala de 340M, U-NorMuon supera tanto a Muon como a NorMuon estándar y elimina por completo el fenómeno de muerte neuronal: las puntuaciones de apalancamiento se vuelven aproximadamente isotrópicas durante el entrenamiento. Fundamentalmente, U-NorMuon propaga este beneficio a capas que no toca directamente: mantener vivas las filas de entrada/arriba garantiza un flujo de gradiente isotrópico hacia la proyección descendente, estabilizando el apalancamiento de la columna sin ninguna intervención directa.
Sin embargo, U-NorMuon todavía tiene un problema: anula con fuerza el factor polar con normas de fila uniformes, sacrificando la precisión del factor polar, lo cual es teóricamente indeseable y empíricamente costoso en el marco de Muon (el artículo muestra que Muon logra una pérdida monótonamente menor con una ortogonalización más precisa). Ésta es la motivación de Aurora.
Aurora: el descenso más pronunciado bajo dos limitaciones conjuntas
Aurora reformula el problema de selección de actualizaciones desde cero. En lugar de ejecutar la ortogonalización y luego parchearla con la normalización de filas, Aurora pregunta: ¿cuál es la actualización óptima bajo la restricción conjunta de semiortogonalidad izquierda y normas de fila uniformes?
Formalmente, para matrices altas, Aurora resuelve:
U∗=argUmaxTr(G⊤U)stU⊤U=In,∥Ui:∥2=mn∀iU ∗ =arg U max Tr(G ⊤ U)stU ⊤ U=I n ,∥U i: ∥ 2 = mn ∀i
La investigación muestra que estas dos restricciones juntas obligan a todos los valores singulares de U a ser exactamente iguales a 1. Esto significa que la restricción conjunta aún produce una actualización semiortogonal izquierda válida, no una comprometida. Esta es la idea clave que separa a Aurora de NorMuon y U-NorMuon: logra uniformidad y ortogonalidad de norma de fila simultáneamente en lugar de intercambiar una por la otra.
La investigación también proporciona dos implementaciones algorítmicas de la solución de Aurora. La aurora de Riemann utiliza un enfoque de proyección de gradiente restringido a la variedad conjunta Stiefel/apalancamiento de filas iguales. Vanilla Aurora es una implementación más simple y práctica. Ambos son de código abierto. Para matrices no altas (anchas y cuadradas), la ortogonalidad ya implica la uniformidad de la norma de fila, por lo que Aurora deja esos parámetros sin cambios.
Resultados
Aurora se utilizó para entrenar un modelo de 1.100 millones que logra una eficiencia de datos 100 veces mayor en datos de Internet de código abierto y supera a modelos más grandes en evaluaciones generales como HellaSwag. A escala 1B, Aurora logra grandes ganancias sobre Muon y NorMuon. En la ejecución rápida de optimización nanoGPT modificada, la ejecución enviada por Aurora supera a la última tecnología anterior (que era NorMuon). Aurora no sintonizada conlleva solo un 6 % de sobrecarga de cómputo en comparación con Muon tradicional y está diseñada como un reemplazo directo.
El equipo de investigación también encontró que las ganancias de rendimiento de Aurora aumentan con el ancho de MLP, lo que sugiere que es particularmente efectivo para redes con grandes factores de expansión de MLP, lo cual es consistente con la hipótesis de la muerte neuronal, ya que los MLP más anchos tienen matrices más altas y más oportunidades para aprovechar la anisotropía para componer.
Conclusiones clave
La actualización del factor polar de Muon hereda la anisotropía de norma de fila en matrices altas, lo que provoca que más del 25% de las neuronas MLP mueran permanentemente ya en el paso 500 del entrenamiento. Aurora resuelve esto encontrando la actualización óptima bajo una restricción conjunta de semiortogonalidad izquierda y normas de fila uniformes, logrando ambas simultáneamente en lugar de intercambiar una por la otra. A una escala de 1.100 millones, Aurora logra una eficiencia de datos 100 veces mayor en datos de Internet de código abierto, supera a los modelos más grandes en HellaSwag y establece un nuevo SoTA en el speedrun nanoGPT modificado. Aurora es un reemplazo casi directo de Muon con solo un 6 % de sobrecarga de cómputo y sus ganancias aumentan con el ancho de MLP.
Consulte el repositorio de Paper y GitHub. Además, no dude en seguirnos en Twitter y no olvide unirse a nuestro SubReddit de más de 150 000 ML y suscribirse a nuestro boletín. ¡Esperar! estas en telegrama? Ahora también puedes unirte a nosotros en Telegram.
¿Necesita asociarse con nosotros para promocionar su repositorio de GitHub O su página principal de Hugging O su lanzamiento de producto O seminario web, etc.? Conéctate con nosotros