OptMuon: Closed-Loop Orthogonalized Momentum Methods for Stochastic Optimization with Zero-Noise Optimality

Resumen Técnico de "OptMuon: Métodos de Momento Ortogonalizados de Bucle Cerrado para la Optimización Estocástica con Optimalidad de Ruido Cero"

1. Planteamiento del Problema

El artículo aborda el problema de la optimización estocástica de primer orden para objetivos no convexos con valores matriciales:
$$\min_{X \in \mathbb{R}^{m \times n}} f(X) := \mathbb{E}_{\xi \sim \mathcal{D}}[f(X; \xi)]$$
El objetivo es encontrar una iterada aleatoria $X_{out}$ tal que se minimice la norma de Frobenius esperada del gradiente, $\mathbb{E}[\|\nabla f(X_{out})\|_F]$, dentro de $T$ llamadas de gradiente estocástico.

Si bien los optimizadores de estilo Muon han demostrado estabilidad empírica al transformar las matrices de momento en direcciones aproximadamente ortogonales (mediante la factorización polar), los métodos existentes suelen emparejar estas direcciones con reglas de bucle abierto o de magnitud constante. Estas reglas suelen depender de constantes de suavidad ($L$) o niveles de varianza ($\sigma^2$) prescritos, lo que las hace sensibles al ajuste de hiperparámetros y capaces de no calibrar explícitamente las magnitudes de actualización basadas en la trayectoria de optimización observada. Por el contrario, los métodos adaptativos como STORM y AdaGrad-Norm ofrecen una adaptación escalar de bucle cerrado y adaptatividad al ruido, pero generalmente operan sobre gradientes brutos en lugar de direcciones de momento ortogonalizadas.

El desafío central abordado es si la adaptación escalar de bucle cerrado puede combinarse con éxito con el momento ortogonalizado tipo Muon para retener la estabilidad geométrica y, al mismo tiempo, lograr la adaptatividad al ruido y la optimalidad de ruido cero sin un reajuste manual.

2. Metodología: OptMuon

Los autores proponen OptMuon, un marco que desacopla la dirección de la actualización de la magnitud de la actualización.

• Dirección (Ortogonalización): La dirección de la actualización es el factor polar de la matriz de momento $M_t$, denotado como $\text{Orth}(M_t) = U W^\top$ (donde $M_t = U \Sigma W^\top$ es la SVD). Esto asegura que la dirección esté normalizada geométricamente y separe la orientación del escalamiento de los valores singulares.
• Magnitud (Adaptación de Bucle Cerrado): El tamaño del paso se determina mediante un coeficiente escalar $\gamma_t$ derivado de la trayectoria observada (historial de gradiente y momento), en lugar de constantes fijas del problema.
  • Programación del Coeficiente: El núcleo del control de la magnitud es una programación de AdaGrad-Norm corregida por máximo móvil que acumula la energía del gradiente estocástico en el denominador para suprimir los pasos en regímenes ruidosos, pero incluye un término en el numerador ($\max g_i^2$) para evitar que el coeficiente colapse excesivamente debido a picos aislados de gradiente.
  • Libre de Lipschitz: El algoritmo no requiere el conocimiento de la constante de suavidad $L$, la varianza $\sigma^2$ o la constante de gradiente acotado $\hat{G}$ para la selección de parámetros, aunque estos aparecen en las constantes teóricas.

Dos Variantes:

1. OptMuon-A (Suavidad Promedio): Diseñado para el supuesto de suavidad promedio. Utiliza un solo gradiente estocástico por iteración y una acumulación de momento estándar ($M_t = G_t + (1-\alpha_t)M_{t-1}$). El exponente en la programación es $q=1/2$.
2. OptMuon-I (Suavidez Individual): Diseñado para el supuesto de suavidad individual. Emplea un estimador de momento recursivo tipo STORM ($M_t = (1-\alpha_t)(M_{t-1} - \nabla f(X_{t-1}; \xi_t)$), utilizando dos evaluaciones de gradiente sobre la misma muestra para reducir la varianza. El exponente es $q=2/3$.

Ambas variantes actualizan los parámetros mediante:
$$X_{t+1} = X_t - \theta \gamma_t \|M_t\|_F \text{Orth}(M_t)$$

3. Contribuciones Clave

El artículo reclama las siguientes contribuciones específicas:

• Ortogonalización de Momento de Bucle Cerrado: Un marco que separa limpiamente la geometría de la matriz (vía el factor polar) de la calibración de la magnitud (vía un esquema escalar auto-normalizado).
• Programación Corregida por Máximo Móvil: Un diseño de coeficiente novedoso que preserva la supresión auto-normalizada de AdaGrad-Norm mientras mitiga el riesgo de colapso del coeficiente causado por picos de gradiente individuales.
• Tasas de Convergencia Adaptativas al Ruido: Garantías teóricas que muestran que los términos estocásticos en los límites de convergencia mejoran automáticamente a medida que la varianza $\sigma$ disminuye.
• Optimalidad de Ruido Cero: La demostración de que cuando $\sigma = 0$, los límites se reducen automáticamente a la tasa determinista casi óptima $\tilde{O}(T^{-1/2})$ sin ningún reajuste de hiperparámetros.

4. Resultados Teóricos

Bajo supuestos de acotación inferior, gradientes estocásticos no sesgados con varianza acotada, suavidad (promedio o individual) y gradientes estocásticos casi-seguramente acotados, el artículo demuestra:

• OptMuon-A (Suavidad Promedio): Logra una tasa de convergencia de:
  $$\tilde{O}\left(T^{-1/2} + \sigma^{1/2}T^{-1/4}\right)$$
• OptMuon-I (Suavidad Individual): Logra una tasa estocástica más aguda de:
  $$\tilde{O}\left(T^{-1/2} + \sigma^{1/3}T^{-1/3}\right)$$

Significado de los Resultados:

• Recuperación de Ruido Cero: En el régimen de ruido cero ($\sigma = 0$), ambos límites se reducen a $\tilde{O}(T^{-1/2})$, coincidiendo con la tasa determinista de primer orden casi óptima. Esto contrasta con los análisis de Muon existentes que a menudo producen $\tilde{O}(T^{-1/4})$ o requieren tamaños de paso fijos.
• Adaptatividad al Ruido: Los términos estocásticos ($\sigma^{1/2}T^{-1/4}$ y $\sigma^{1/3}T^{-1/3}$) son adaptativos; se estrechan a medida que el nivel de ruido disminuye, a diferencia de los métodos con tamaños de paso fijos.
• Comparación: El artículo señala que OptMuon-A mejora la tasa $\tilde{O}(T^{-1/4})$ establecida en variantes recientes de Muon con reducción de varianza (por ejemplo, Chang et al., 2025) en el régimen de bajo ruido, manteniendo la optimalidad de ruido cero.

5. Significado y Limitaciones

Signancia:
El artículo establece que la adaptación escalar de bucle cerrado puede acoplarse con el momento ortogonalizado de estilo Muon sin sacrificar la estabilidad geométrica. Extiende el concepto de "adaptatividad de ruido cero" —previamente encontrado en reglas de tamaño de paso escalar adaptativas— a las actualizaciones de momento de factor polar. Esto permite una optimización robusta que se ajusta automáticamente a los niveles de ruido y recupera el rendimiento determinista óptimo sin intervención manual.

Limitaciones (según los autores):

• Ortogonalización Exacta: La teoría asume el cálculo exacto del factor polar. Aunque las iteraciones de Newton-Schulz pueden aproximar esto, el análisis actual no cuantifica el error introducido por la ortogonalización inexacta.
• Gradientes Acotados: El análisis depende de una condición de gradiente estocástico casi-seguramente acotado (Supuesto 4). El artículo no pretende tener una teoría libre de gradientes acotados.
• Validación Empírica: El trabajo se centra en garantías de convergencia teórica y no incluye evaluaciones empíricas a gran escala ni extensiones a objetivos compuestos o con restricciones, las cuales se señalan como direcciones para trabajos futuros.