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

Sintesi Tecnica di "OptMuon: Closed-Loop Orthogonalized Momentum Methods for Stochastic Optimization with Zero-Noise Optimality"

1. Definizione del Problema

Il documento affronta il problema dell'ottimizzazione stocastica del primo ordine per obiettivi non convessi a valori matriciali:
$$\min_{X \in \mathbb{R}^{m \times n}} f(X) := \mathbb{E}_{\xi \sim \mathcal{D}}[f(X; \xi)]$$
L'obiettivo è trovare un iterato casuale $X_{out}$ tale che la norma di Frobenius attesa del gradiente, $\mathbb{E}[\|\nabla f(X_{out})\|_F]$, sia minimizzata entro $T$ chiamate di gradiente stocastico.

Sebbene gli ottimizzatori in stile Muon abbiano dimostrato stabilità empirica trasformando le matrici di momentum in direzioni approssimativamente ortogonali (tramite fattorizzazione polare), i metodi esistenti tipicamente accoppiano queste direzioni con regole ad anello aperto (open-loop) o a magnitudo costante. Tali regole dipendono spesso da costanti di regolarità ($L$) o livelli di varianza ($\sigma^2$) prescritti, rendendoli sensibili alla sintonizzazione degli iperparametri e incapaci di calibrare esplicitamente le magnitudo degli aggiornamenti basandosi sulla traiettoria di ottimizzazione osservata. Al contrario, i metodi adattivi come STORM e AdaGrad-Norm offrono un'adattabilità scalare a anello chiuso (closed-loop) e l'adattività al rumore, ma operano generalmente su gradienti grezzi piuttosto che su direzioni di momentum ortogonalizzate.

La sfida centrale affrontata è se l'adattabilità scalare ad anello chiuso possa essere combinata con successo con il momentum ortogonalizzato in stile Muon per mantenere la stabilità geometrica pur ottenendo l'adattività al rumore e l'ottimalità a rumore zero senza una ricalibrazione manuale.

2. Metodologia: OptMuon

Gli autori propongono OptMuon, un framework che scinde la direzione dell'aggiornamento dalla sua magnitudo.

• Direzione (Ortogonalizzazione): La direzione dell'aggiornamento è il fattore polare della matrice di momentum $M_t$, denotato come $\text{Orth}(M_t) = U W^\top$ (dove $M_t = U \Sigma W^\top$ è la SVD). Ciò assicura che la direzione sia normalizzata geometricamente e separi l'orientamento dalla scala dei valori singolari.
• Magnitudo (Adattabilità ad Anello Chiuso): La dimensione del passo è determinata da un coefficiente scalare $\gamma_t$ derivato dalla traiettoria osservata (storia di gradienti e momentum), anziché da costanti del problema prefissate.
  • Programma del Coefficiente: Il nucleo del controllo della magnitudo è uno schema di tipo AdaGrad-Norm corretto dal massimo corrente (running-maximum corrected), che accumula l'energia del gradiente stocastico nel denominatore per sopprimere i passi in regimi rumorosi, ma include un termine al numeratore ($\max g_i^2$) per evitare che il coefficiente collassi eccessivamente a causa di picchi isolati del gradiente.
  • Lipschitz-Free: L'algoritmo non richiede la conoscenza della costante di regolarità $L$, della varianza $\sigma^2$ o della costante del gradiente limitato $\hat{G}$ per la selezione dei parametri, sebbene questi compaiano nelle costanti teoriche.

Due Varianti:

1. OptMuon-A (Smoothness Media): Progettato per l'assunzione di regolarità media. Utilizza un singolo gradiente stocastico per iterazione e un accumulo di momentum standard ($M_t = G_t + (1-\alpha_t)M_{t-1}$). L'esponente nello schema è $q=1/2$.
2. OptMuon-I (Smoothness Individuale): Progettato per l'assunzione di regolarità individuale. Impiega un estimatore di momentum ricorsivo di tipo STORM ($M_t = (1-\alpha_t)(M_{t-1} - \nabla f(X_{t-1}; \xi_t)$), utilizzando due valutazioni del gradiente sullo stesso campione per ridurre la varianza. L'esponente è $q=2/3$.

Entrambe le varianti aggiornano i parametri tramite:
$$X_{t+1} = X_t - \theta \gamma_t \|M_t\|_F \text{Orth}(M_t)$$

3. Contributi Chiave

Il documento rivendica i seguenti contributi specifici:

• Ortogonalizzazione del Momentum ad Anello Chiuso: Un framework che separa nettamente la geometria della matrice (tramite il fattore polare) dalla calibrazione della magnitudo (tramite uno schema scalare auto-normalizzato).
• Programma Corretto dal Massimo Corrente: Un design innovativo del coefficiente che preserva la soppressione auto-normalizzata di AdaGrad-Norm mitigando al contempo il rischio di collasso del coefficiente causato da singoli picchi di gradiente.
• Tassi di Convergenza Adattivi al Rumore: Garanzie teoriche che mostrano come i termini stocastici nei limiti di convergenza migliorino automaticamente man mano che la varianza $\sigma$ diminuisce.
• Ottimalità a Rumore Zero: La dimostrazione che quando $\sigma = 0$, i limiti si riducono automaticamente al tasso deterministico quasi ottimale $\tilde{O}(T^{-1/2})$ senza alcuna ricalibrazione degli iperparametri.

4. Risultati Teorici

Sotto le assunzioni di limitatezza inferiore, gradienti stocastici non distorti con varianza limitata, regolarità (media o individuale) e gradienti stocastici quasi-sicuramente limitati, il documento dimostra:

• OptMuon-A (Smoothness Media): Raggiunge un tasso di convergenza di:
  $$\tilde{O}\left(T^{-1/2} + \sigma^{1/2}T^{-1/4}\right)$$
• OptMuon-I (Smoothness Individuale): Raggiunge un tasso stocastico più acuto di:
  $$\tilde{O}\left(T^{-1/2} + \sigma^{1/3}T^{-1/3}\right)$$

Significato dei Risultati:

• Recupero a Rumore Zero: Nel regime a rumore zero ($\sigma = 0$), entrambi i limiti si riducono a $\tilde{O}(T^{-1/2})$, eguagliando il tasso deterministico del primo ordine quasi ottimale. Ciò contrasta con le analisi di Muon esistenti che spesso producono $\tilde{O}(T^{-1/4})$ o richiedono step size fissi.
• Adattività al Rumore: I termini stocastici ($\sigma^{1/2}T^{-1/4}$ e $\sigma^{1/3}T^{-1/3}$) sono adattivi; si stringono man mano che il livello di rumore diminuisce, a differenza dei metodi con step size fissi.
• Confronto: Il documento nota che OptMuon-A migliora il tasso $\tilde{O}(T^{-1/4})$ stabilito da recenti varianti di Muon con riduzione della varianza (es. Chang et al., 2025) nel regime a basso rumore, mantenendo al contempo l'ottimalità a rumore zero.

5. Significato e Limitazioni

Significato:
Il documento stabilisce che l'adattabilità scalare ad anello chiuso può essere accoppiata con il momentum ortogonalizzato in stile Muon senza sacrificare la stabilità geometrica. Estende il concetto di "adattività a rumore zero" — precedentemente trovato in regole di step size scalari adattive — agli aggiornamenti del momentum del fattore polare. Ciò consente un'ottimizzazione robusta che si adatta automaticamente ai livelli di rumore e recupera le prestazioni deterministiche ottimali senza intervento manuale.

Limitazioni (come dichiarato dagli autori):

• Ortogonalizzazione Esatta: La teoria assume il calcolo esatto del fattore polare. Sebbene le iterazioni di Newton-Schulz possano approssimarlo, l'analisi attuale non quantifica l'errore introdotto dall'ortogonalizzazione imprecisa.
• Gradienti Limitati: L'analisi si basa su una condizione di gradiente stocastico quasi-sicuramente limitato (Assunzione 4). Il documento non rivendica una teoria priva di vincoli sul gradiente.
• Validazione Empirica: Il documento si concentra sulle garanzie teoriche di convergenza e non include valutazioni empiriche su larga scala o estensioni a obiettivi compositi/vincolati, che sono indicati come direzioni per lavori futuri.