Muon-OGD: Muon-based Spectral Orthogonal Gradient Projection for LLM Continual Learning

Riepilogo Tecnico: Muon-OGD

Enunciato del Problema
La sfida centrale affrontata in questo lavoro è l'oblio catastrofico nei Modelli Linguistici di Grande Dimensione (LLM) durante l'apprendimento continuo. Quando gli LLM vengono sintonizzati sequenzialmente su nuovi compiti, le loro prestazioni sui compiti appresi in precedenza si degradano spesso rapidamente. I metodi esistenti di apprendimento continuo basati su proiezione (ad esempio OGD, GEM, O-LoRA) mitigano questo fenomeno limitando gli aggiornamenti dei parametri a sottospazi ortogonali alle direzioni associate ai compiti passati. Tuttavia, questi metodi sono tipicamente formulati nell'ambito della geometria euclidea dei parametri, dove le grandezze degli aggiornamenti e l'ortogonalità sono governate dalla norma di Frobenius.

Gli autori identificano una discrepanza fondamentale: gli ottimizzatori moderni per parametri a valore matriciale, come il recentemente proposto ottimizzatore Muon, operano sotto una geometria della norma spettrale (norma dell'operatore). Muon applica aggiornamenti matriciali ortogonalizzati (approssimati tramite iterazioni di Newton–Schulz) e può essere interpretato come l'esecuzione della discesa del gradiente più ripida sotto un vincolo di norma spettrale. Il lavoro sostiene che affidarsi alla geometria della norma di Frobenius per i vincoli di apprendimento continuo è subottimale per i parametri degli LLM a valore matriciale, poiché non si allinea con la geometria della dinamica di ottimizzazione sottostante.

Metodologia: Muon-OGD
Il lavoro propone Muon-OGD (Proiezione del Gradiente Ortogonale Spettrale basata su Muon), un framework che integra l'ottimizzazione consapevole della norma spettrale con i vincoli di apprendimento continuo basati su proiezione.

1. Formulazione: Invece del problema vincolato standard con norma di Frobenius, Muon-OGD formula l'aggiornamento $\Delta$ come un problema di ottimizzazione vincolato dalla norma spettrale:
   $$\min_{\Delta \in \mathbb{R}^{m \times n}} \langle G, \Delta \rangle \quad \text{s.t.} \quad \|\Delta\|_2 \leq \eta, \quad \langle C_i, \Delta \rangle = 0, \forall i=1,\dots,k$$
   Qui, $G$ è la matrice del gradiente (o del momento), $\|\cdot\|_2$ è la norma spettrale, $\eta$ è la dimensione del passo e $\{C_i\}$ sono le direzioni di vincolo (ad esempio, gradienti da compiti passati) che devono rimanere ortogonali all'aggiornamento.

2. Ottimizzazione Duale: Poiché il problema manca di una soluzione in forma chiusa, gli autori derivano una formulazione duale. L'obiettivo duale comporta la minimizzazione della norma nucleare di una matrice spostata $H = G + \sum \lambda_i C_i$. L'aggiornamento ottimale viene recuperato tramite la funzione segno della matrice: $\Delta^* = -\eta \cdot \text{msgn}(H)$.

3. Algoritmo: Il metodo risolve il problema duale iterativamente utilizzando la discesa del subgradiente sulle variabili duali $\lambda$.

   • Ciclo Interno: Aggiorna iterativamente $\lambda$ per imporre i vincoli di ortogonalità.
   • Approssimazione del Segno della Matrice: La funzione segno della matrice $\text{msgn}(H)$ è approssimata efficientemente utilizzando un piccolo numero di iterazioni di Newton–Schulz (specificamente 5 iterazioni, denotate come $NS_5$), evitando il costo computazionale della SVD esplicita.
   • Istante a Rango Basso: Per efficienza, gli autori propongono un'istanza specifica in cui i vincoli sono derivati dai vettori singolari top-$k$ dei compiti passati, permettendo alle variabili duali di essere organizzate in una matrice a rango basso $\Lambda$, riducendo la dimensionalità del problema duale.

Contributi Chiave

1. Geometria della Norma Spettrale per l'Apprendimento Continuo: Il lavoro formula l'apprendimento continuo basato su proiezione sotto la geometria della norma spettrale, allineando i vincoli di stabilità dell'apprendimento continuo con la geometria della norma dell'operatore degli ottimizzatori moderni come Muon.
2. Regola di Aggiornamento Pratica: Deriva un algoritmo pratico (Muon-OGD) che combina i vincoli di ortogonalità dei compiti passati con aggiornamenti del segno della matrice in stile Muon, risolvendo il problema vincolato in modo efficiente tramite una formulazione duale e approssimazioni di Newton–Schulz.
3. Validazione Empirica: Il framework è istanziato per la sintonizzazione continua degli LLM e valutato su diversi benchmark, dimostrando miglioramenti coerenti rispetto alle controparti basate sulla norma di Frobenius.

Risultati Sperimentali
Gli autori valutano Muon-OGD su:

• Benchmark Standard: Benchmark di apprendimento continuo a 5 e 15 compiti utilizzando T5-Large.
• Benchmark TRACE: Un benchmark per LLM sintonizzati su istruzioni (utilizzando LLaMA-2-7B-Chat) che copre competenze diverse come programmazione, matematica e comprensione multilingue.
• Curricula Specifici per Dominio: Apprendimento sequenziale su Programmazione (BigCodeBench) $\to$ Matematica (GSM8K) $\to$ Medicina (HuatuoGPT-o1) utilizzando i modelli Qwen2.5 e LLaMA3.2.

Risultati:

• Prestazioni: Muon-OGD supera costantemente la sintonizzazione sequenziale (SeqFT) e le basi di confronto con gradienti ortogonali competitive (ad esempio O-LoRA, Sculpting Subspaces, OSFT) in termini di Accuratezza Media (AA) e Trasferimento Inverso (BT).
  • Sul benchmark a 15 compiti di T5-Large, Muon-OGD ha raggiunto un'accuratezza media del 72,0%, superando O-LoRA (69,6%) e OSFT (71,3%).
  • Sul benchmark TRACE, Muon-OGD ha raggiunto la più alta AA (49,4%) e il tasso di oblio più basso (-4,7%) tra i metodi confrontati.
• Scalabilità: Il metodo rimane computazionalmente scalabile, sfruttando l'efficienza delle iterazioni di Newton–Schulz e le rappresentazioni dei vincoli a rango basso.
• Agnosticismo Architetturale: L'approccio è efficace sia su architetture encoder-decoder (T5) che decoder-only (LLaMA, Qwen).

Significato e Affermazioni
Il lavoro afferma che la geometria degli aggiornamenti consapevole della norma spettrale fornisce un'alternativa pratica ed efficace alla proiezione basata sulla norma di Frobenius per l'apprendimento continuo negli LLM. Colmando il divario tra la geometria degli ottimizzatori moderni (come Muon) e i requisiti di stabilità dell'apprendimento continuo, Muon-OGD migliora il compromesso stabilità-plasticità. Gli autori suggeriscono che i metodi esistenti basati su proiezione possono essere sistematicamente aggiornati per operare sotto la geometria della norma spettrale per preservare meglio i benefici di condizionamento e stabilità osservati in Muon, prevenendo al contempo l'oblio catastrofico. Il lavoro non afferma di risolvere tutte le sfide dell'apprendimento continuo, ma offre un'alternativa geometrica principiale che produce guadagni empirici nella ritenzione e nell'adattamento attraverso vari domini e scale di modelli.