Tilde Research presenta Aurora, optimizador que corrige la «muerte silenciosa» de neuronas en Muon

Tilde Research lanzó Aurora, un optimizador open source diseñado para sustituir a Muon en ciertos pasos de entrenamiento corrigiendo un fallo estructural que deja inactivas muchas neuronas en capas MLP anchas;

Tilde Research ha presentado Aurora, un optimizador pensado para reemplazar a Muon en pasos concretos del entrenamiento de redes. El equipo afirma que Aurora corrige un defecto estructural que en matrices anchas — como las de capas MLP basadas en SwiGLU — provoca la inactividad permanente de un número significativo de neuronas. El proyecto se entrega con código abierto y con resultados experimentales a escala, incluyendo un preentrenamiento de 1.1 mil millones de parámetros y una ejecución en la carrera modded‑nanoGPT que mejoró el estado del arte previo.

En términos técnicos, Aurora formula la selección de la actualización como un problema de descenso más pronunciado sujeto a dos restricciones conjuntas: semi‑ortogonalidad por la izquierda, expresada como UᵀU = I, y normas de fila uniformes apropiadas para matrices anchas. Para resolverlo, los autores publican dos implementaciones: Riemannian Aurora, que aplica proyección del gradiente sobre la variedad de Stiefel combinada con la restricción de igualdad de leverage por fila, y una versión 'vanilla' más práctica pensada para uso a escala que reduce la complejidad computacional sin renunciar al objetivo de mantener isotropía de filas.

El contexto del diseño se entiende mejor al repasar cómo funciona Muon: su paso central calcula el factor polar polar(G) = U Vᵀ para ortogonalizar el gradiente y aplica la actualización W ← W − η U Vᵀ. NorMuon introdujo una normalización por fila inspirada en escalados por parámetro y llegó a liderar temporalmente la carrera modded‑nanoGPT antes de la llegada de Aurora. Estas variantes buscaban acelerar la convergencia frente a optimizadores clásicos como AdamW, pero la operación que combina ortogonalización y normalización por fila resultó ser vulnerable en matrices anchas.

La falla detectada aparece en matrices de gran ancho — características de capas MLP como las basadas en SwiGLU — donde Muon no puede mantener al mismo tiempo la ortogonalidad y la uniformidad de normas de fila. Ese conflicto produce anisotropía de la norma por fila: algunas neuronas reciben actualizaciones amplificadas mientras que otras quedan casi sin señal. El efecto acumulativo crea una 'espiral de muerte' neuronal, en la que neuronas mediocres reciben cada vez menos gradiente hasta quedar inactivas de forma permanente. hacia el paso 500 de entrenamiento más de una de cada cuatro neuronas ya está efectivamente muerta en esas capas, y la inactividad se propaga a capas posteriores al privarlas de señales útiles.

Como respuesta inicial, los autores propusieron U‑NorMuon, una corrección intermedia que ajusta la normalización de fila de NorMuon fijando la norma objetivo de cada fila en √(n/m) para matrices anchas — la cifra correcta para una matriz column‑ortogonal—. En experimentos a 340 millones de parámetros, U‑NorMuon elimina la muerte neuronal y supera a Muon y NorMuon en esa métrica, pero a costa de sacrificar precisión del factor polar. Esa pérdida de precisión en la ortogonalización reduce el rendimiento teórico y práctico dentro del marco de Muon, lo que motivó la búsqueda de una solución que impusiera ambas restricciones de forma conjunta.

Aurora difiere porque no ajusta el factor polar después de la ortogonalización, sino que resuelve directamente la optimización con las dos restricciones simultáneas. Los autores demuestran que imponer semi‑ortogonalidad izquierda y normas de fila iguales fuerza que todos los valores singulares de U valgan exactamente 1, de modo que la solución conserva la ortogonalidad deseada sin renunciar a la isotropía de las filas. Esa propiedad evita la acumulación de anisotropía de norma por fila y, por tanto, previene la espiral de muerte que observaban en Muon.

En pruebas empíricas, Aurora se usó para preentrenar un modelo de 1.1 mil millones de parámetros que, según el informe, alcanza una eficiencia de datos 100× en datos abiertos de Internet y supera a modelos más grandes en evaluaciones generales como HellaSwag. En la carrera modded‑nanoGPT, una ejecución con Aurora mejoró el resultado previo alcanzado con NorMuon. A escala de aproximadamente 1B, los autores registran ganancias sustanciales frente a Muon y NorMuon, mientras que en la escala de 340M U‑NorMuon logró mitigar la muerte neuronal aunque con la limitación indicada en la precisión del factor polar.

Para la adopción práctica, Aurora se ofrece como reemplazo 'drop‑in' para Muon con un sobrecoste de cómputo sin afinar cercano al 6% y con código abierto disponible para inspección e integración. La ventaja operativa es que, al evitar la muerte neuronal y preservar la precisión de la ortogonalización, Aurora estabiliza el flujo de gradiente hacia capas posteriores y puede favorecer entrenamientos a gran escala donde Muon había sido adoptado. Las limitaciones reportadas por los autores son el pequeño aumento de cómputo y la necesidad de validar el optimizador en más arquitecturas y en otras escalas de modelos antes de declarar su idoneidad universal.