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

Resumo Técnico de "OptMuon: Métodos de Momento Ortogonalizados de Malha Fechada para Otimização Estocástica com Otimalidade de Ruído Zero"

1. Definição do Problema

O artigo aborda o problema da otimização estocástica de primeira ordem para objetivos não convexos de valores matriciais:
$$\min_{X \in \mathbb{R}^{m \times n}} f(X) := \mathbb{E}_{\xi \sim \mathcal{D}}[f(X; \xi)]$$
O objetivo é encontrar uma iteração aleatória $X_{out}$ tal que a norma de Frobenius esperada do gradiente, $\mathbb{E}[\|\nabla f(X_{out})\|_F]$, seja minimizada dentro de $T$ chamadas de gradiente estocástico.

Embora otimizadores do estilo Muon tenham demonstrado estabilidade empírica ao transformar matrizes de momento em direções aproximadamente ortogonais (via fatorização polar), os métodos existentes tipicamente combinam essas direções com regras de malha aberta ou de magnitude constante. Essas regras frequentemente dependem de constantes de suavidade ($L$) ou níveis de variância ($\sigma^2$) prescritos, tornando-as sensíveis ao ajuste de hiperparâmetros e incapazes de calibrar explicitamente as magnitudes de atualização com base na trajetória de otimização observada. Por outro lado, métodos adaptativos como STORM e AdaGrad-Norm ofere vezes "adaptação escalar de malha fechada" e adaptatividade de ruído, mas geralmente operam sobre gradientes brutos em vez de direções de momento ortogonalizadas.

O desafio central abordado é se a adaptação escalar de malha fechada pode ser combinada com sucesso com o momento ortogonalizado do estilo Muon para manter a estabilidade geométrica, ao mesmo tempo em que alcança adaptatividade de ruído e otimalidade de ruído zero sem necessidade de reajuste manual.

2. Metodologia: OptMuon

Os autores propõem o OptMuon, um framework que desacopla a direção da atualização da magnitude da atualização.

• Direção (Ortogonalização): A direção da atualização é o fator polar da matriz de momento $M_t$, denotado por $\text{Orth}(M_t) = U W^\top$ (onde $M_t = U \Sigma W^\top$ é a SVD). Isso garante que a direção seja geometria-normalizada e separe a orientação da escala dos valores singulares.
• Magnitude (Adaptação de Malha Fechada): O tamanho do passo é determinado por um coeficiente escalar $\gamma_t$ derivado da trajetória observada (histórico de gradiente e momento), em vez de constantes fixas do problema.
  • Cronograma de Coeficiente: O núcleo do controle de magnitude é um cronograma AdaGrad-Norm corrigido pelo máximo móvel que acumula a energia do gradiente estocástico no denominador para suprimir passos em regimes ruidosos, mas inclui um termo no numerador ($\max g_i^2$) para evitar que o coeficiente colapse excessivamente devido a picos isolados de gradiente.
  • Livre de Lipschitz: O algoritmo não requer conhecimento da constante de suavidade $L$, da variância $\sigma^2$ ou da constante de gradiente limitado $\hat{G}$ para a seleção de parâmetros, embora estes apareçam nas constantes teóricas.

Duas Variantes:

1. OptMuon-A (Suavidade Média): Projetado para a suposição de suavidade média. Utiliza um gradiente estocástico por iteração e um acúmulo de momento padrão ($M_t = G_t + (1-\alpha_t)M_{t-1}$). O expoente no cronograma é $q=1/2$.
2. OptMuon-I (Suavidade Individual): Projetado para a suposição de suavidade individual. Emprega um estimador de momento do tipo STORM recursivo ($M_t = (1-\alpha_t)(M_{t-1} - \nabla f(X_{t-1}; \xi_t)$), utilizando duas avaliações de gradiente na mesma amostra para reduzir a variância. O expoente é $q=2/3$.

Ambas as variantes atualizam os parâmetros via:
$$X_{t+1} = X_t - \theta \gamma_t \|M_t\|_F \text{Orth}(M_t)$$

3. Principais Contribuições

O artigo reivindica as seguintes contribuições específicas:

• Ortogonalização de Momento de Malha Fechada: Um framework que separa claramente a geometria da matriz (via fator polar) da calibração de magnitude (via um cronograma escalar auto-normalizado).
• Cronograma Corrigido pelo Máximo Móvel: Um design de coeficiente inovador que preserva a supressão auto-normalizada do AdaGrad-Norm enquanto mitiga o risco de colapso do coeficiente causado por picos únicos de gradiente.
• Taxas de Convergência Adaptativas ao Ruído: Garantias teóricas mostrando que os termos estocásticos nos limites de convergência melhoram automaticamente conforme a variância $\sigma$ diminui.
• Otimalidade de Ruído Zero: A demonstração de que, quando $\sigma = 0$, os limites reduzem-se automaticamente para a taxa determinística quase ótima $\tilde{O}(T^{-1/2})$ sem qualquer reajuste de hiperparâmetros.

4. Resultados Teóricos

Sob suposições de limitação inferior, gradientes estocásticos não viesados com variância limitada, suavidade (média ou individual) e gradientes estocásticos quase-certamente limitados, o artigo prova:

• OptMuon-A (Suavidade Média): Alcança uma taxa de convergência de:
  $$\tilde{O}\left(T^{-1/2} + \sigma^{1/2}T^{-1/4}\right)$$
• OptMuon-I (Suavidade Individual): Alcança uma taxa estocástica mais nítida de:
  $$\tilde{O}\left(T^{-1/2} + \sigma^{1/3}T^{-1/3}\right)$$

Significância dos Resultados:

• Recuperação de Ruído Zero: No regime de ruído zero ($\sigma = 0$), ambos os limites reduzem-se para $\tilde{O}(T^{-1/2})$, correspondendo à taxa determinística de primeira ordem quase ótima. Isso contrasta com as análises de Muon existentes que frequentemente resultam em $\tilde{O}(T^{-1/4})$ ou exigem tamanhos de passo fixos.
• Adaptatividade de Ruído: Os termos estocásticos ($\sigma^{1/2}T^{-1/4}$ e $\sigma^{1/3}T^{-1/3}$) são adaptativos; eles se estreitam conforme o nível de ruído diminui, ao contrário de métodos com tamanhos de passo fixos.
• Comparação: O artigo observa que o OptMuon-A melhora a taxa $\tilde{O}(T^{-1/4})$ estabelecida em variantes recentes de redução de variância do Muon (ex: Chang et al., 2025) no regime de baixo ruído, mantendo a otimalidade de ruído zero.

5. Significância e Limitações

Significância:
O artigo estabelece que a adaptação escalar de malha fechada pode ser acoplada ao momento ortogonalizado do estilo Muon sem sacrificar a estabilidade geométrica. Ele estende o conceito de "adaptatividade de ruído zero" — anteriormente encontrado em regras de tamanho de passo escalar adaptativo — para atualizações de momento de fator polar. Isso permite uma otimização robusta que se ajusta automaticamente aos níveis de ruído e recupera o desempenho determinístico ideal sem intervenção manual.

Limitações (conforme declarado pelos autores):

• Ortogonalização Exata: A teoria assume o cálculo exato do fator polar. Embora iterações de Newton-Schulz possam aproximar isso, a análise atual não quantifica o erro introduzido pela ortogonalização inexata.
• Gradientes Limitados: A análise depende de uma condição de gradiente estocástico quase-certamente limitado (Suposição 4). O artigo não reivindica uma teoria livre de gradiente limitado.
• Validação Empírica: O foco do artigo são as garantias teóricas de convergência e não inclui avaliações empíricas de larga escala ou extensões para objetivos compostos/restritos, que são notas como direções para trabalhos futuros.