Riepilogo Tecnico: Muon può Affinare Modelli Pre-addestrati con Adam?

Enunciato del Problema

Muon (MomentUm Orthogonalizzato da Newton-Schulz) si è affermato come un ottimizzatore altamente efficiente per il pre-addestramento di grandi modelli linguistici (LLM), offrendo un'efficienza computazionale di circa $2\times$ rispetto ad Adam e riducendo l'utilizzo di memoria eliminando la stima del secondo momento. Tuttavia, la sua applicazione pratica all'affinamento (fine-tuning) rimane limitata. L'ostacolo principale è il problema di incompatibilità degli ottimizzatori: la maggior parte dei modelli open disponibili pubblicamente è pre-addestrata con Adam, e passare in modo ingenuo a Muon per l'affinamento porta a prestazioni degradate. Viceversa, l'uso di Adam per affinare modelli pre-addestrati con Muon produce anch'esso risultati subottimali. Questa incompatibilità limita severamente l'adozione di Muon nel paradigma dominante di affinamento, nonostante i suoi successi nel pre-addestramento.

Metodologia e Analisi

Gli autori indagano questa incompatibilità attraverso esperimenti controllati e analisi teorica, concentrandosi sui distinti bias impliciti di Adam e Muon.

1. Caratterizzazione del Bias Implicito:

   • Adam utilizza una precondizionamento elemento per elemento basato sui primi e secondi momenti, convergendo verso soluzioni che massimizzano il margine della norma-max ($\|W\|_{\max}$).
   • Muon applica un precondizionamento a livello di matrice tramite ortogonalizzazione di Newton-Schulz, convergendo verso soluzioni che massimizzano il margine della norma-spettrale ($\|W\|_2$).
   • L'analisi teorica su un problema di regressione lineare semplificato dimostra che questi ottimizzatori convergono verso soluzioni fondamentalmente diverse ($W^*_{\max}$ vs. $W^*_2$). Sperimentalmente, i pesi addestrati con Muon mostrano un rango stabile più elevato e distribuzioni di valori singolari più uniformi rispetto ai pesi addestrati con Adam.

2. Meccanismo del Degrado:

   • L'incompatibilità interrompe la conoscenza pre-addestrata. Quando viene applicato un ottimizzatore con un bias implicito diverso, forza i pesi del modello in una direzione strutturale incompatibile con la fase di pre-addestramento.
   • Gli esperimenti mostrano che l'affinamento con incompatibilità aumenta la sensibilità alla forza degli aggiornamenti. Aggiornamenti più forti causano una maggiore perturbazione della struttura pre-addestrata, portando a una maggiore perplessità e a un oblio catastrofico più severo.

3. Mitigazione Proposta: LoRA:

   • Gli autori ipotizzano che vincolare l'estensione degli aggiornamenti possa mitigare questa incompatibilità. Impiegano l'Adattamento a Basso Rango (LoRA), che congela i pesi pre-addestrati ($W_0$) e vincola gli aggiornamenti a un sottospazio a basso rango ($\Delta W = BA$).
   • Questo approccio limita l'entità delle variazioni dei pesi, impedendo all'ottimizzatore di alterare aggressivamente la struttura pre-addestrata. L'analisi teorica suggerisce che l'"inflazione dell'incompatibilità" (il divario tra prestazioni con e senza incompatibilità) è limitata dal rango LoRA $r$, svanendo quando $r \to 1$.

Contributi Chiave

• Riproduzione e Analisi: Il documento riproduce il fenomeno dell'incompatibilità degli ottimizzatori a scale accessibili (NanoChat, 561M parametri) e lo collega ai distinti bias impliciti (norma-max vs. norma-spettrale) di Adam e Muon.
• LoRA come Soluzione: Gli autori dimostrano che LoRA mitiga efficacemente l'incompatibilità. Vincolando gli aggiornamenti, LoRA permette a Muon di affinare modelli pre-addestrati con Adam con prestazioni paragonabili o superiori a LoRA-Adam.
• Validazione Completa: I risultati sono validati su:
  • Comprensione del Linguaggio Naturale (NLU): Benchmark GLUE con T5-Base.
  • Generazione del Linguaggio Naturale (NLG): Affinamento istruito di Llama 2-7B/13B su matematica, codice e ragionamento di senso comune.
  • Visione Artificiale: Affinamento di CLIP ViT-B/32 su sei compiti di classificazione di immagini.
• Studi su Rango e Oblio:
  • Sensibilità al Rango: Il beneficio di LoRA-Muon è più pronunciato a ranghi bassi e moderati. All'aumentare del rango (avvicinandosi all'affinamento completo), l'incompatibilità riemerge e le prestazioni peggiorano rispetto ad Adam.
  • Oblio Catastrofico: L'affinamento completo con Muon su modelli pre-addestrati con Adam causa un oblio più severo della conoscenza pre-addestrata rispetto ad Adam. LoRA riduce significativamente questo oblio, preservando la struttura pre-addestrata.
• Varianti LoRA: Lo studio valuta le varianti LoRA esistenti (es. rsLoRA, PiSSA, LoRA-One). Si rileva che le varianti progettate per rendere gli aggiornamenti LoRA più vicini all'affinamento completo (come PiSSA) non necessariamente beneficiano Muon e possono persino esacerbare l'incompatibilità, suggerendo la necessità di adattamenti specifici per Muon.

Risultati

• Riduzione del Divario di Prestazioni: Nell'affinamento completo, Muon performa costantemente peggio di Adam su modelli pre-addestrati con Adam (es. divario di perplessità da +0,009 a +0,023 su WikiText-2). Sotto LoRA, questo divario si riduce significativamente (del 39% per modelli pre-addestrati con Muon e del 78% per modelli pre-addestrati con Adam).
• Superiorità in LoRA: Attraverso compiti NLU, NLG e di visione, LoRA-Muon eguaglia o supera LoRA-Adam. Ad esempio, sul benchmark GLUE, LoRA-Muon-PE ha raggiunto la massima accuratezza media (89,20%), superando sia LoRA-Adam (88,93%) che Full-Adam (89,14%).
• Efficienza: Sebbene Muon comporti un sovraccarico modesto per passo a causa dell'iterazione di Newton-Schulz (1,1–1,2$\times$ più lento di Adam nelle impostazioni LoRA), la sua efficienza di memoria (archiviazione di un solo buffer di momento contro due per Adam) rimane un vantaggio pratico, risparmiando circa il 50% della memoria dello stato dell'ottimizzatore.

Significato e Affermazioni

Il documento afferma di fornire la prima analisi approfondita del problema di incompatibilità degli ottimizzatori tra Adam e Muon. La sua principale rilevanza risiede nel rendere possibile l'uso pratico di Muon per l'affinamento del vasto ecosistema di modelli pre-addestrati con Adam.

Gli autori concludono che:

1. L'incompatibilità degli ottimizzatori è strutturale: Deriva da bias impliciti conflittuali che interrompono la conoscenza pre-addestrata quando la forza degli aggiornamenti non è vincolata.
2. Gli aggiornamenti vincolati sono la chiave: Tecniche come LoRA, che limitano l'entità delle variazioni dei pesi, neutralizzano efficacemente questa incompatibilità.
3. Raccomandazione Pratica: I praticanti possono utilizzare Muon come sostituto diretto di Adam nell'affinamento LoRA per ottenere efficienza di memoria e prestazioni competitive senza subire le penalità di incompatibilità osservate nell'affinamento completo.

Il documento nota modestamente che, sebbene LoRA mitighi il problema, le varianti LoRA esistenti ottimizzate per Adam non si trasferiscono automaticamente a Muon, suggerendo che sono necessari lavori futuri per sviluppare adattamenti specifici per Muon. Inoltre, la caratterizzazione teorica dell'incompatibilità rimane una questione aperta.