Resumen Técnico: OrScale – Optimización Ortogonalizada con Escalado de Relación de Confianza por Capa

1. Declaración del Problema

El optimizador Muon ha surgido como una herramienta práctica para el entrenamiento a gran escala, utilizando actualizaciones ortogonalizadas para fijar la dirección de los pasos matriciales mediante iteraciones de Newton–Schulz (aproximando el factor polar). Sin embargo, Muon sufre de una limitación crítica: mientras optimiza la dirección de las actualizaciones, deja la magnitud controlada por una única tasa de aprendizaje esencialmente global.

Este escalado global es subóptimo porque diferentes capas (por ejemplo, proyecciones de atención, proyecciones de subida/bajada de MLP, incrustaciones) exhiben distribuciones de gradiente y curvaturas efectivas muy diferentes. Aunque la variante Moonlight introdujo un factor estático por forma para igualar el RMS por elemento, no puede adaptarse a la deriva dinámica de las escalas de las capas durante el entrenamiento.

Los intentos previos de portar el Escalado Adaptativo de Tasa por Capa (LARS) y el Escalado Adaptativo de Momento por Capa (LAMB)—que desacoplaron con éxito los tamaños de paso de las tasas de aprendizaje globales en SGD y Adam—a Muon han fallado. El artículo identifica que los puertos ingenuos colapsan debido a tres razones estructurales específicas:

1. Denominadores Degenerados: Usar la norma de Frobenius de la dirección ortogonalizada ($\|Q_\ell\|_F$) como denominador de la relación de confianza falla porque es casi una constante determinada por la forma de la matriz (rango), no por la magnitud de la actualización.
2. Incompatibilidad de Unidades y Saturación de Recorte: Usar la norma de momento cruda ($\|\tilde{M}_\ell\|_F$) introduce una incompatibilidad de unidades (unidades de gradiente vs. unidades de peso), haciendo que la relación de confianza alcance el límite superior de recorte ($r_{max}$) casi universalmente, volviendo inefectiva la adaptación.
3. Descontrol de Decaimiento de Pesos Desacoplado: Combinar denominadores calibrados con decaimiento de pesos desacoplado crea un bucle de retroalimentación positiva donde las normas de los pesos crecen descontroladamente (10–50$\times$), desestabilizando el entrenamiento.

2. Metodología

Los autores proponen OrScale, un marco basado en un único principio de diseño: El denominador de una relación de confianza por capa debe medir la norma de Frobenius de la dirección de actualización real que el optimizador restará de los parámetros.

Principio Central

Sea $D_{\ell,t}$ la dirección real en el espacio de parámetros que se resta de los pesos $W_{\ell,t}$ antes de aplicar la tasa de aprendizaje global. La relación de confianza se define como:
$$r_{\ell,t} = \frac{\|W_{\ell,t}\|_F}{\|D_{\ell,t}\|_F + \epsilon}$$
Críticamente, $D_{\ell,t}$ debe incluir tanto el término de decaimiento de pesos como la dirección ortogonalizada.

Dos Variantes

El marco produce dos implementaciones específicas:

1. OrScale (General/Defecto para Visión):

   • Dirección de Actualización: $D_{\ell,t} = \lambda W_{\ell,t} + Q_{\ell,t}$ (donde $\lambda$ es el decaimiento de pesos y $Q$ es la dirección ortogonalizada).
   • Factor de Forma: $s_\ell = 1$ (sin escalado estático).
   • Calibración: Ninguna.
   • Decaimiento de Pesos: Acoplado directamente en el paso de actualización escalado por la relación de confianza.

2. OrScale-LM (Defecto para LLM):

   • Dirección de Actualización: $D_{\ell,t} = \lambda W_{\ell,t} + s_\ell Q_{\ell,t}$.
   • Factor de Forma: Adopta el factor estático de Moonlight $s_\ell = 0.2\sqrt{\max(m_\ell, n_\ell)}$ para igualar el RMS por elemento de AdamW.
   • Calibración: Introduce una constante de calibración por capa perezosa ($c_{\text{denom}, \ell}$) calculada una vez en la inicialización. Esto ancla la relación de confianza inicial exactamente a 1 ($r_{\ell,0}=1$).
   • Significado de la Calibración: Esto permite que OrScale-LM herede los hiperparámetros de tasa de aprendizaje y decaimiento de pesos ajustados para Moonlight (y, por extensión, AdamW) sin requerir un nuevo barrido de hiperparámetros. Garantiza la invarianza de ancho en la inicialización.

Detalles Algorítmicos

• Decaimiento de Pesos Acoplado: El decaimiento de pesos se integra en el cálculo de la relación de confianza ($\lambda W$ es parte del denominador) en lugar de aplicarse como un retorno separado. Esto previene el bucle de retroalimentación descontrolado observado en variantes desacopladas.
• Eficiencia: El método añade una sobrecarga insignificante (<1% de tiempo de reloj), requiriendo solo un escalar adicional por capa matricial y dos cálculos de norma de Frobenius por paso, los cuales están dominados por las iteraciones de Newton–Schulz existentes.

3. Contribuciones Clave

• Algoritmo: Una receta unificada y ligera en parámetros (Algoritmo 1) que integra con éxito el escalado de relación de confianza por capa en Muon. Los autores identifican la única "esquina viable" en el espacio de diseño (denominador de dirección de actualización real + WD acoplado).
• Análisis de Fallos: Un mapeo riguroso del espacio de diseño de la relación de confianza de Muon (Tabla 1) que caracteriza analítica y empíricamente tres modos de fallo (denominador degenerado, saturación de recorte y crecimiento descontrolado de pesos), demostrando por qué fallan y por qué la solución propuesta tiene éxito.
• Teoría:
  • Convergencia: Prueba una garantía de convergencia no convexa de norma nuclear $O(1/\sqrt{T})$.
  • Ganancia Adaptativa: Establece una constante de descenso adaptativa por capa $\kappa_{\text{eff}} > 1$ bajo heterogeneidad de capas medible, mejorando estrictamente sobre la constante homogénea de Muon.
  • Separación: Prueba una separación estricta de MuTrust/MuScale bajo condiciones de saturación de recorte.
  • Calibración: Demuestra que OrScale-LM preserva la transferencia de tasa de aprendizaje estilo $\mu$P en la inicialización.
• Empírica: Evaluación exhaustiva en tareas de modelos de visión y lenguaje.

4. Resultados Experimentales

Visión (CIFAR-10 / DavidNet)

• Rendimiento: OrScale logró una precisión top-1 de validación de 94.05% $\pm$ 0.08, ocupando el primer lugar entre cinco líneas base.
• Mejora: Esto representa una mejora de +0.35 puntos sobre Muon estándar (93.70%) y de +0.30 sobre Muon+Moonlight (93.75%).
• Estabilidad: OrScale demostró una ventana de tasa de aprendizaje estable más amplia en comparación con Muon, consistente con las predicciones teóricas de ganancias adaptativas por capa.
• Ablación: Un estudio de ablación de modos de fallo confirmó que las variantes correspondientes a los modos de fallo identificados (por ejemplo, usando $\|Q_\ell\|_F$ o momento crudo) exhibieron los síntomas predichos (por ejemplo, saturación de recorte del 99.5–100% o descontrol de la norma de los pesos) y no lograron igualar el rendimiento de OrScale.

Modelos de Lenguaje Grande (Pre-entrenamiento FineWeb-Edu)

• Escalado: Evaluado en cuatro escalas de modelo (de 125M a 1.1B parámetros).
• Comparación: OrScale-LM superó a Muon+Moonlight en tres de cuatro escalas (125M, 545M, 1.1B) y superó a AdamW en todas las escalas.
• Reducción de Pérdida: En la escala de 1.1B, OrScale-LM logró una entropía cruzada final de validación de 2.6251, comparado con 2.6360 para Muon+Moonlight y 2.7304 para AdamW.
• Eficiencia: La ventaja se mantuvo en todo el rango de cómputo, mostrando OrScale-LM una victoria clara en el régimen donde el cómputo es más importante para los LLM.

5. Significado y Afirmaciones

El artículo afirma que OrScale cierra con éxito la brecha entre la geometría espectral de Muon y la adaptabilidad por capa de LARS/LAMB. Su principal significado radica en:

1. Resolver el Problema de Magnitud: Proporciona el primer control de magnitud variable en el tiempo y por capa para Muon que es robusto frente a modos de fallo estructurales.
2. Transferibilidad Práctica: A través de la variante OrScale-LM y la calibración perezosa, permite a los practicantes transferir tasas de aprendizaje ajustadas para AdamW o Moonlight directamente a OrScale-LM sin re-ajuste, reduciendo la barrera de adopción.
3. Solidez Teórica: Ofrece una justificación teórica rigurosa de por qué los intentos híbridos previos fallaron y prueba que el método propuesto produce constantes de convergencia estrictamente mejores bajo heterogeneidad.
4. Validación Empírica: Demuestra que estas mejoras teóricas se traducen en ganancias tangibles tanto en tareas de visión a pequeña escala como en pre-entrenamiento de LLM a gran escala, estableciendo a OrScale como un optimizador de vanguardia para capas matriciales.

Los autores mantienen modestia respecto a los resultados de LLM, señalando que la ventaja en la escala de 1.1B es comparable al ruido de una sola semilla en escalas más pequeñas, y por lo tanto, los resultados deben verse como evidencia direccional de la eficacia del método en lugar de una superioridad estadística definitiva en todos los regímenes. Sin embargo, la mejora consistente a través de las escalas y la ablación exitosa de los modos de fallo apoyan fuertemente el principio de diseño propuesto.