The Implicit Bias of Adam and Muon on Smooth Homogeneous Neural Networks

Questo studio dimostra che ottimizzatori basati su momentum come Adam e Muon, quando applicati a reti neurali omogenee lisce, seguono traiettorie approssimate di discesa ripida normalizzata con un tasso di apprendimento decrescente, portando a una bias implicito verso punti KKT di problemi di massimizzazione del margine la cui identità dipende dalla norma scelta dall'ottimizzatore.

L'essenziale

Immagina di dover insegnare a un gruppo di studenti (la tua rete neurale) a distinguere le foto di gatti da quelle di cani. Il tuo obiettivo è trovare la regola perfetta che separi i due gruppi. Ma c'è un problema: ci sono infinite regole che funzionano perfettamente sui dati di allenamento. Quale sceglierai?

La ricerca di Eitan Gronich e Gal Vardi del Weizmann Institute of Science ci dice che la risposta non dipende solo dai dati, ma da come insegni agli studenti a imparare. In termini tecnici, ogni algoritmo di ottimizzazione ha un "pregiudizio implicito" (una tendenza nascosta) che lo spinge verso una soluzione specifica, anche se non glielo dici esplicitamente.

Ecco una spiegazione semplice di cosa hanno scoperto, usando metafore quotidiane.

1. Il Concetto di Base: La "Soglia di Sicurezza"

Immagina che ogni studente (ogni punto dati) debba stare a una certa distanza dal confine tra "gatto" e "cane". Più sono lontani dal confine, più sono sicuri di non sbagliare. Questo è il margine.

• L'obiettivo: Trovare la regola che massimizza questa distanza di sicurezza per tutti gli studenti.
• Il problema: Ci sono molti modi per massimizzare questa distanza, a seconda di come misuri la "distanza".

2. I "Metodi di Studio" (Gli Ottimizzatori)

Il paper analizza tre modi principali in cui gli algoritmi "studiano" e correggono i loro errori. Ognuno ha una personalità diversa:

A. La Discesa del Gradiente (GD) - Il "Misuratore Classico"

Immagina un escursionista che scende una montagna. Se usa un metodo classico (Gradient Descent), cammina sempre nella direzione più ripida, ma misura la distanza con un metro standard (la norma L2, come la distanza in linea retta).

• Risultato: Tende a trovare la soluzione che massimizza il margine usando questo metro standard. È il metodo "classico" e prevedibile.

B. Adam - Il "Corridore Frettoloso"

Adam è l'algoritmo più usato oggi (specialmente per le Intelligenze Artificiali che scrivono testi). È come un corridore che guarda solo il passo immediato e si adatta velocemente, ignorando il passato lontano.

• La scoperta: Gli autori hanno scoperto che Adam, quando la "velocità di apprendimento" diminuisce (come quando si avvicina alla fine della corsa), smette di usare il metro standard. Inizia a comportarsi come se usasse un metro a righe (la norma L-infinito).
• In pratica: Adam cerca di massimizzare la sicurezza basandosi sul "peggior caso possibile" (il margine più piccolo), ignorando le piccole variazioni. È come se dicesse: "Non importa quanto sono sicuro degli altri, devo solo assicurarmi che questo studente non sbagli".

C. Muon - Il "Regista delle Forme"

Muon è un nuovo algoritmo molto popolare per le grandi reti neurali. Immagina che invece di guardare i singoli studenti, Muon guardi l'intera classe come un blocco unico e cerchi di mantenere una certa "forma" o struttura geometrica (usando le norme spettrali delle matrici).

• La scoperta: Muon ha un pregiudizio verso una misura di distanza molto specifica, legata alla "forza" delle sue connessioni interne (norma spettrale).
• Il mix (Muon-Adam): Se usi Muon per le connessioni pesanti (i neuroni) e Adam per le altre parti, l'algoritmo crea un ibrido. Massimizza la sicurezza secondo una regola che è il "massimo" tra le due misure: la forma geometrica di Muon e la rigidità di Adam.

3. L'Analogia della "Pista da Corsa"

Immagina che l'addestramento della rete neurale sia una corsa su una pista infinita.

• Il percorso: È la funzione di perdita (quanto sbagliamo).
• La meta: È il punto dove non sbagliamo più nulla.
• Il pregiudizio implicito: È la traiettoria che l'algoritmo sceglie di seguire.
  • Se usi un algoritmo "lento e costante" (Gradiente), seguirai un sentiero che porta a un punto specifico della meta.
  • Se usi Adam, il tuo sentiero si piegherà verso un punto diverso, anche se arrivi allo stesso risultato finale.
  • Se usi Muon, il tuo sentiero seguirà una curva ancora diversa, basata sulla geometria complessa della pista.

4. Perché è importante?

Prima di questo studio, sapevamo che gli algoritmi "trovavano" soluzioni che generalizzavano bene (funzionavano bene su dati nuovi), ma non sapevamo perché sceglievano proprio quella soluzione tra le infinite possibili.

Gli autori hanno dimostrato che:

1. Non è magia: È una conseguenza matematica precisa di come l'algoritmo si muove.
2. La scelta conta: Se vuoi che la tua AI sia robusta contro certi tipi di errori, devi scegliere l'algoritmo giusto. Se vuoi massimizzare la sicurezza del "caso peggiore", Adam è ottimo. Se vuoi una struttura geometrica specifica, Muon è la scelta.
3. La teoria conferma la pratica: Hanno fatto esperimenti su immagini reali (MNIST, i numeri scritti a mano) e hanno visto che gli algoritmi facevano esattamente quello che la teoria prediceva: Adam massimizzava il margine "L-infinito", Muon quello "spettrale", e così via.

In sintesi

Questo paper ci dice che l'algoritmo che scegli non è solo uno strumento per arrivare alla meta, ma è anche l'architetto che decide quale meta raggiungere.

• Adam è come un avvocato difensore che si concentra solo sul caso più debole del cliente per garantirne l'assoluzione totale.
• Muon è come un architetto che progetta un edificio basandosi sulla forza strutturale massima delle sue travi.
• Gradient Descent è il geometra che misura tutto con il metro classico.

Capire queste "personalità" nascoste ci aiuta a costruire Intelligenze Artificiali più sicure, robuste e prevedibili.

Sommario tecnico

1. Il Problema

Le reti neurali profonde, sebbene spesso sovraparametrizzate e addestrate senza regolarizzazione esplicita, mostrano una capacità di generalizzazione eccezionale. Un filone di ricerca consolidato attribuisce questo fenomeno al bias implicito degli algoritmi di ottimizzazione basati sul gradiente, che tendono a convergere verso soluzioni con margini ampi (max-margin).

Mentre il bias implicito della Gradient Descent (GD) è stato ampiamente studiato (massimizzazione del margine $\ell_2$), l'analisi degli ottimizzatori moderni e più complessi come Adam e Muon rimane meno esplorata, specialmente al di fuori dei modelli lineari. La domanda centrale è: qual è il bias implicito di Adam e Muon quando applicati a modelli omogenei non lineari (smooth)? In particolare, quale norma viene massimizzata e verso quali punti di ottimalità (KKT) convergono le traiettorie?

2. Metodologia

Gli autori adottano un approccio teorico basato sull'analisi asintotica delle traiettorie di ottimizzazione in tempo continuo (flussi gradienti).

• Modelli: Si concentrano su modelli omogenei lisci (smooth homogeneous), dove la funzione di uscita $f(x; \alpha\theta) = \alpha^L f(x; \theta)$ per $\alpha > 0$. Questo include reti lineari profonde e reti con attivazioni lisce omogenee (es. $ReLU_q$ con $q>1$, attivazioni quadratiche).
• Ottimizzatori Analizzati:
  • Steepest Descent Normalizzato: Estensione dei risultati esistenti a schemi di apprendimento (learning rate) variabili.
  • Momentum Steepest Descent (MSD): Include Muon (norma spettrale), MomentumGD ($\ell_2$) e Signum ($\ell_\infty$).
  • Adam: Analizzato senza il costante di stabilità ($\epsilon=0$), riflettendo meglio il comportamento pratico dove il gradiente domina.
  • Compositi: Muon-Signum e Muon-Adam.
• Framework Teorico:
  • Introducono il concetto di "Approximate Steepest Descent" (Discesa del gradiente approssimata). Dimostrano che, sotto un regime di learning rate decrescente, gli algoritmi basati su momentum (come Adam e Muon) asintoticamente si comportano come traiettorie di discesa del gradiente approssimata rispetto a una specifica norma.
  • Utilizzano le condizioni KKT (Karush-Kuhn-Tucker) per caratterizzare i punti limite delle direzioni dei parametri $\frac{\theta_t}{\|\theta_t\|}$.
  • Assumono la convergenza direzionale (ipotesi (T2)) e l'assenza di divisione per zero (per Adam).

3. Contributi Chiave

1. Estensione della Steepest Descent Normalizzata:
   Dimostrano che per qualsiasi modello omogeneo localmente Lipschitz (inclusi i ReLU network sotto condizioni deboli), ogni punto limite della traiettoria normalizzata di Steepest Descent con un learning rate $\eta(t)$ tale che $\int_0^\infty \eta(t) dt = \infty$, è un punto KKT del problema di massimizzazione del margine rispetto alla norma $\|\cdot\|$ scelta.

2. Bias di Muon e Momentum Steepest Descent:
   Dimostrano che algoritmi come Muon (che utilizza la norma spettrale sulle matrici dei pesi) e Muon-Signum convergono asintoticamente verso i punti KKT del problema di massimizzazione del margine rispetto a una norma definita dalla combinazione delle norme spettrali delle matrici dei pesi (norma $\|\cdot\|_{msp}$).

   • Risultato: Muon massimizza il margine rispetto alla norma spettrale massima tra le matrici dei pesi.

3. Bias di Adam:
   Dimostrano che Adam (senza costante di stabilità) sotto un learning rate decrescente ha un bias implicito verso la massimizzazione del margine $\ell_\infty$. Questo estende risultati precedenti limitati ai modelli lineari a una classe molto più ampia di modelli omogenei lisci.

4. Analisi di Algoritmi Ibridi (Muon-Adam):
   Analizzano il caso in cui diverse parti del modello (es. matrici vs vettori) sono ottimizzate con algoritmi diversi (Muon per le matrici, Adam per i vettori). Dimostrano che la traiettoria complessiva converge verso un punto KKT rispetto a una norma ibrida:
   $$\|\theta\| = \max \left( \frac{\eta_A}{\eta_M} \|(W_1, ..., W_K)\|_{msp}, \|u\|_\infty \right)$$
   dove $\eta_A$ e $\eta_M$ sono i learning rate base per Adam e Muon.

4. Risultati Sperimentali

Gli esperimenti sono stati condotti su reti neurali a due livelli (una hidden layer) per classificare cifre MNIST (pari/dispari) utilizzando la perdita esponenziale.

• Confronto: Sono stati confrontati NGD (con/senza momentum), Signum, Adam, Muon e Muon-Adam.
• Attivazioni: Testati sia con ReLU che con ReLU quadrato (smooth).
• Risultati:
  • NGD: Massimizza il margine $\ell_2$.
  • Signum e Adam: Massimizzano il margine $\ell_\infty$.
  • Muon: Massimizza il margine rispetto alla norma spettrale massima ($\|\cdot\|_{msp}$).
  • Muon-Adam: Massimizza la norma ibrida prevista dalla teoria.
  • Le curve di convergenza mostrano che la direzione dei parametri converge (similitudine cosinica > 0.99 nella seconda metà dell'addestramento), validando l'ipotesi di convergenza direzionale.

5. Significato e Implicazioni

Questo lavoro è significativo per diversi motivi:

• Unificazione Teorica: Fornisce un quadro unificante ("Approximate Steepest Descent") che spiega il comportamento di ottimizzatori moderni basati su momentum, collegandoli alla massimizzazione del margine in modelli non lineari.
• Comprensione di Adam e Muon: Chiarisce perché Adam e Muon, pur essendo algoritmi adattivi complessi, tendono a soluzioni con proprietà geometriche specifiche (margine $\ell_\infty$ o spettrale), offrendo una spiegazione teorica alla loro efficacia pratica.
• Progettazione di Ottimizzatori: Suggerisce che la scelta dell'ottimizzatore determina implicitamente la "geometria" della soluzione finale. Questo potrebbe guidare la progettazione di nuovi ottimizzatori per migliorare la generalizzazione o la robustezza in scenari specifici.
• Limiti e Futuro: Il lavoro lascia aperte questioni sulla convergenza direzionale per modelli non lisci (come le reti ReLU standard) senza assunzioni aggiuntive e sull'estensione di questi risultati a modelli non omogenei.

In sintesi, il paper dimostra che il "segreto" del successo di ottimizzatori come Adam e Muon risiede in un bias implicito verso la massimizzazione di margini specifici (rispettivamente $\ell_\infty$ e norme spettrali), estendendo la teoria classica della discesa del gradiente a un'era di ottimizzatori adattivi complessi.