ROOT: Robust Orthogonalized Optimizer for Neural Network Training

Di /
stato della ricerca deep learning

Guida a ROOT optimizer: che cos’è, perché è interessante, panoramica risultati

ROOT è un nuovo optimizer progettato dal Huawei Noah’s Ark Lab per rendere l’addestramento di Large Language Models più stabile ed efficiente, partendo dalle idee di Muon ma correggendone alcune fragilità. In particolare punta a risolvere due problemi: la sensibilità della procedura di orthogonalization delle matrici dei pesi e l’instabilità dovuta a “gradient outlier” molto grandi.

L’idea centrale è combinare due meccanismi: una orthogonalization “dimension-robust”, che si adatta automaticamente alle dimensioni di ogni matrice di pesi, e un blocco di proximal optimization che attenua in modo selettivo gli outlier nel gradient (una versione più morbida e strutturata del gradient clipping). Su un Transformer da 1B di parametri pre-addestrato su FineWeb-Edu, ROOT converge più velocemente e ottiene un leggero ma sistematico miglioramento sulle classiche benchmark zero-shot rispetto sia ad AdamW che a Muon. Gli autori mostrano inoltre che lo stesso meccanismo di soft-thresholding migliora anche la performance di un piccolo Vision Transformer su CIFAR-10, indicando che il metodo è utile oltre le sole LLM.

  • GitHub: il paper annuncia che il codice sarà rilasciato nel repository “https://github.com/huawei-noah/noah-research”; al momento della pubblicazione su arXiv è indicato come “code will be available”, quindi è possibile che non sia ancora pubblico o completo.
  • Paper: versione arXiv disponibile su https://arxiv.org/abs/2511.20626.
  • Dataset: non esiste un “ROOT dataset” dedicato; gli esperimenti principali usano sottoinsiemi da 10B e 100B token del dataset FineWeb-Edu, distribuito separatamente su Hugging Face, e non è registrato come dataset che “cita” il paper nella pagina Hugging Face del lavoro.

Indice

Come funziona ROOT: approcci, tecniche e training recipe

Da Muon a ROOT: orthogonalization dei gradienti

Muon è un optimizer “matrix-aware” che non tratta i pesi solo come vettori, ma lavora su intere matrici (per esempio le proiezioni delle attention e i layer MLP) e applica una procedura di orthogonalization al momentum dei gradienti tramite una iterazione di tipo Newton-Schulz. In pratica, invece di fare un semplice step nella direzione del gradient medio, trasforma questo momentum in una direzione quasi ortogonale con norma controllata, cercando di seguire meglio la geometria spettrale della loss.

Questa strategia porta vantaggi di stabilità e velocità rispetto ad AdamW, ma Muon usa una singola tripletta di coefficienti numerici per l’iterazione Newton-Schulz, scelta in modo “one-size-fits-all”, che genera errori di approximazione molto diversi a seconda della forma della matrice (soprattutto per matrici quadrate grandi). Inoltre l’orthogonalization può amplificare gradient outlier, perché la procedura parte da un momentum normalizzato che viene molto influenzato dagli elementi di grande magnitudine.

Adaptive Newton iteration con coefficienti dipendenti dalla dimensione

La prima novità di ROOT è rendere l’orthogonalization robusta rispetto alla dimensione, introducendo una adaptive Newton iteration con coefficienti specifici per ogni shape di matrice. Invece di usare sempre gli stessi parametri dell’iterazione (come in Muon), ROOT impara o calibra offline una tripletta di coefficienti per ciascuna coppia (m, n) di dimensioni presente nel modello (per esempio layer con 2048×4096, 4096×4096, ecc.).

Gli autori mostrano che con coefficienti fissi l’errore di orthogonalization (rispetto a una soluzione “ideale” basata su SVD) può variare di ordini di grandezza a seconda della shape, mentre con coefficienti specializzati per forma l’errore diventa molto più uniforme e ridotto su tutta la rete. Questa dimension-robust orthogonalization permette a ROOT di mantenere aggiornamenti più coerenti fra layer diversi durante il training reale, riducendo l’errore rispetto all’orthogonalizzazione “ground truth” lungo le prime migliaia di step di pre-training.

Proximal optimization e soft-thresholding per gli outlier

La seconda componente chiave è un modulo di robust optimization basato su proximal operator che separa il momentum in una parte “base” e una parte di outlier, applicando una operazione di soft-thresholding elemento per elemento. In termini qualitativi significa che tutti gli elementi del gradient al di sotto di una certa soglia rimangono quasi invariati, mentre quelli troppo grandi vengono “rimpiccioliti” in modo continuo, invece di essere brutalmente saturati come con un gradient clipping hard.

Il paper modella questa decomposizione come il risultato di un problema di ottimizzazione che penalizza gli outlier con una norma sparsa e vincola la dimensione della componente base, ma in implementazione si riduce a una funzione di shrinkage con soglia basata su un quantile dinamico del modulo del momentum (per esempio il 90-esimo percentile). Orthogonalization viene poi applicata solo alla componente “pulita” B_t, scartando la parte O_t ricca di outlier, così da non propagare direzioni rumorose all’interno della costosa iterazione Newton-Schulz.

Training recipe e setup sperimentale

Per valutare ROOT gli autori utilizzano un Transformer da 1B di parametri pre-addestrato su FineWeb-Edu, con esperimenti di ablation su 10 miliardi di token e run principali su 100 miliardi di token. Il modello viene addestrato per una singola epoca, con sequence length 4096, un learning rate programmato con warmup iniziale e successivo decadimento cosinusoidale, e batch globali dell’ordine di centinaia di migliaia fino a un milione di token.

L’addestramento è eseguito su cluster distribuiti di Ascend NPU, usando strategie di parallelismo dati/modello e una tecnica di attention-mask-reset per evitare che il modello faccia self-attention tra documenti distinti all’interno della stessa sequenza. Per il confronto vengono usati AdamW e Muon (con gli iperparametri raccomandati in letteratura) e una versione “ROOT-SoftThresh” che integra solo il blocco di outlier suppression mantenendo il resto dell’optimizer come in Muon.

Risultati di ROOT: breakdown dei principali esperimenti

Precisione dell’orthogonalization e dinamica reale

In studi controllati su matrici sintetiche e distribuzioni di singular value prese dal training reale, ROOT mostra che i coefficienti dipendenti dalla shape riducono drasticamente l’errore di orthogonalization rispetto a Newton-Schulz con coefficienti fissi. Nel caso peggiorativo per Muon (matrici quadrate grandi) l’errore medio viene ridotto di molto, mentre per matrici rettangolari si passa da errori ancora significativi a valori quasi nulli, rendendo uniforme la qualità dell’aggiornamento su tutta l’architettura.

Quando si osservano i gradient reali raccolti durante le prime 10.000 iterazioni di pre-training, ROOT mantiene sistematicamente un errore relativo rispetto all’orthogonalizzazione via SVD più basso sia di Muon sia di una variante “classica” Newton-Schulz. Questo indica che la calibrazione dimension-aware non è solo un trucco teorico, ma si traduce in una migliore approssimazione delle direzioni di discesa effettive durante il training.

Pre-training LLM: loss e benchmark zero-shot

Sul pre-training da 10B token, sia ROOT-SoftThresh sia la versione completa ROOT ottengono una training loss più bassa rispetto a Muon lungo tutta la traiettoria; la variante completa, che unisce outlier suppression e dimension-robust orthogonalization, è quella che converge meglio. Alla fine del run, ROOT raggiunge una loss leggermente inferiore alla baseline Muon, confermando che il beneficio non è solo in termini di stabilità ma anche di qualità del minimo trovato.

Nel run da 100B token, gli autori valutano il modello con lm-evaluation-harness su un insieme standard di benchmark zero-shot come HellaSwag, ARC-easy/ARC-challenge, BoolQ, PIQA, SciQ, OBQA, Wino e WSC. ROOT supera AdamW e Muon nella media delle accuratezze, con un vantaggio medio di circa mezzo punto percentuale rispetto a Muon e di circa un punto rispetto ad AdamW, e in alcuni task specifici (per esempio WSC e alcune QA scientifiche) mostra margini leggermente più ampi, pur rimanendo nella fascia “incremento piccolo ma consistente”.

Ablation su soglia di outlier e calibrazione spettrale

Per il modulo di soft-thresholding, il paper confronta diverse soglie espresse come quantile del modulo del momentum (per esempio 85%, 90%, 95%, 99%). Soglie troppo conservative (come 99%) non filtrano abbastanza il “tail” del rumore e conducono a training più instabile, mentre soglie troppo aggressive (come 85%) tagliano eccessivamente il segnale utile; un valore intorno al 90% offre il miglior compromesso tra stabilità e convergenza della loss.

La calibrazione dei coefficienti di Newton-Schulz avviene offline su distribuzioni di singular value raccolte dal training con Muon, combinando matrici reali e matrici gaussiane random in diversi rapporti; la strategia “Mix 1:3” (più matrici random) risulta la più robusta quando si scala verso modelli e setup più complessi. Una calibrazione più aggressiva su dati reali (“Mix 1:1”) può dare loss leggermente minore in un singolo scenario, ma tende a produrre instabilità (per esempio spike di loss) quando combinata con soft-thresholding o con modelli più grandi, quindi viene scartata come scelta predefinita.

Generalizzazione a task di visione

Per testare se il contributo di ROOT sia legato solo alle LLM, gli autori addestrano da zero un piccolo Vision Transformer (circa 6.3M parametri) su CIFAR-10, usando Muon o ROOT per ottimizzare le matrici 2D e AdamW per bias, layer norm ed embedding. In questo setting, la variante ROOT con soft-thresholding raggiunge una top-1 accuracy fino a circa 88.4%, contro circa 84.7% della baseline Muon, dimostrando che attenuare in modo strutturato gli outlier nel gradient aiuta anche in compiti di classificazione immagini standard.

Concetti chiave da capire per leggere ROOT (guida completa)

Matrix-aware optimizers e orthogonalization

Optimizer come Muon, Shampoo e altri “matrix-whitening” non ragionano più a livello di singolo parametro, ma guardano i pesi come matrici e cercano di “normalizzare” la geometria di queste trasformazioni. Invece di scalare ogni coordinata con un learning rate diverso, applicano trasformazioni che rendono gli aggiornamenti più isotropi nello spazio degli autovalori, spesso tramite operazioni ispirate alla whitening o alla spectral normalization.

In Muon e ROOT, l’orthogonalization del momentum serve proprio a ottenere update che non si allineino troppo con poche direzioni dominanti, ma esplorino meglio lo spazio dei parametri. Capire questo passaggio concettuale (da “per-parameter scaling” a “matrix-wise orthogonalization”) è fondamentale per seguire il resto del paper.

Newton-Schulz iteration in parole semplici

La Newton-Schulz iteration è una procedura iterativa che prende una matrice e la trasforma passo dopo passo in qualcosa di “più ortogonale”, usando una combinazione di moltiplicazioni della matrice per sé stessa pesate da coefficienti scelti ad hoc. Ogni step usa gli stessi coefficienti, e dopo un numero fisso di iterazioni si ottiene una buona approssimazione dell’operazione di whitening/orthogonalization desiderata senza dover calcolare una SVD costosa.

Il punto critico che ROOT mette in luce è che la stessa scelta di coefficienti non è ugualmente buona per tutte le dimensioni e distribuzioni spettrali: matrici con shape diverse hanno pattern di autovalori diversi, e quindi “rispondono” in maniera diversa alla stessa iterazione. Da qui nasce l’idea di usare coefficienti specializzati per shape, che sono il cuore della dimension-robust orthogonalization di ROOT.

Gradient outlier, soft-thresholding e differenza dal clipping

Nei grandi dataset testuali, alcuni batch possono contenere esempi “estremi” che producono gradient con valori molto più grandi della media, soprattutto in alcune righe/colonne delle matrici legate all’attention. Se questi outlier vengono passati tali e quali all’orthogonalization, rischiano di dominare la normalizzazione e di generare update instabili, con fenomeni come logit dell’attenzione che esplodono o spike improvvisi della loss.

Il gradient clipping tradizionale “trancia” semplicemente i valori sopra una certa soglia, ma lo fa in modo rigido e non strutturato, rischiando di distorcere troppo anche il segnale buono. Il soft-thresholding di ROOT, visto come proximal operator, riduce gradualmente i valori grandi mantenendo l’ordinamento relativo e operando su una soglia che segue in modo adattivo la scala corrente del momentum (tramite quantile), il che si traduce in una protezione più dolce e spesso più efficace contro gli outlier.

Pre-training LLM e valutazione zero-shot

ROOT viene valutato in un setting molto simile a quello usato in molti lavori recenti sulle LLM: pre-training autoregressivo su un grande corpus (FineWeb-Edu) con sequence length lungo, seguito da valutazione zero-shot su benchmark di comprensione linguistica e ragionamento. In questo contesto, anche miglioramenti modesti nell’average accuracy o nella velocità di convergenza sono significativi, perché si traducono in ore o giorni di calcolo risparmiati su cluster molto costosi.

Capire cosa misurano benchmark come HellaSwag, ARC, BoolQ, PIQA e WSC aiuta anche a leggere le tabelle dei risultati: sono test di completamento di frasi, QA scientifica, ragionamento fisico e pronome resolution, cioè aspetti diversi della competenza di una LLM. ROOT non cambia l’architettura del modello, ma agisce “solo” sull’optimizer, quindi i suoi guadagni riflettono una migliore esplorazione del paesaggio di loss a parità di backbone.

Quiz: verifica se hai capito ROOT

Quali sono i due problemi principali che ROOT vuole risolvere rispetto a Muon?

ROOT identifica come principali limiti degli optimizer basati su orthogonalization la fragilità rispetto alla dimensione delle matrici (precisione di orthogonalization che cambia molto da layer a layer) e la vulnerabilità agli outlier del gradient che possono destabilizzare il training. Il primo problema nasce dall’uso di coefficienti fissi nella Newton-Schulz iteration, mentre il secondo deriva dal fatto che la normalizzazione e la trasformazione polinomiale possono amplificare valori anomali del momentum.

In che modo l’adaptive Newton iteration rende ROOT più robusto?

L’adaptive Newton iteration di ROOT sostituisce la tripletta di coefficienti fissi usata da Muon con coefficienti appresi o calibrati specificamente per ogni shape di matrice presente nel modello. In questo modo, per matrici quadrate grandi e per matrici rettangolari molto allungate si usano parametri diversi, scelti per minimizzare l’errore di orthogonalization sulla distribuzione effettiva degli autovalori di quella shape, riducendo la variazione di qualità degli update fra layer.

Perché ROOT preferisce soft-thresholding al gradient clipping classico?

ROOT usa soft-thresholding perché offre un compromesso migliore tra rimozione degli outlier e preservazione del segnale informativo nel gradient. A differenza del clipping hard, che trancia i valori oltre una soglia fissa, il soft-thresholding riduce gradualmente i valori troppo grandi secondo una funzione continua, con una soglia che segue il quantile del modulo del momentum, e applica l’orthogonalization solo alla parte “ripulita”, il che migliora la stabilità senza rallentare in modo eccessivo la convergenza.

Su quali modelli e dataset viene valutato ROOT?

Il core del paper valuta ROOT su un Transformer da 1B di parametri pre-addestrato sul dataset testuale FineWeb-Edu, con esperimenti di ablation su 10B token e esperimenti principali su 100B token. Inoltre viene testata una variante su un piccolo Vision Transformer addestrato da zero su CIFAR-10 per mostrare che il meccanismo di soft-thresholding porta benefici anche in ambito visione.

Che tipo di miglioramenti empirici ottiene ROOT rispetto ad AdamW e Muon?

Nei run LLM, ROOT raggiunge una training loss leggermente più bassa rispetto a Muon e mostra curve di convergenza più stabili lungo tutta la traiettoria di pre-training. Sulle benchmark zero-shot come HellaSwag, ARC, BoolQ e altre, ROOT migliora la media di accuratezza di circa un punto percentuale rispetto ad AdamW e di circa mezzo punto rispetto a Muon, con alcuni task individuali che mostrano vantaggi leggermente maggiori, pur restando in un regime di guadagni incrementali.

Perché ROOT è rilevante per chi progetta nuovi optimizer per LLM?

ROOT è interessante perché dimostra che, anche in una famiglia di optimizer già avanzata come quella dei matrix-whitening/orthogonalized, ci sono ancora margini significativi intervenendo su robustezza numerica e trattamento esplicito del rumore nei gradient. Il lavoro suggerisce che futuri optimizer per LLM dovranno combinare in modo modulare orthogonalization/PCA-like, variance adaptation e meccanismi di robust optimization, piuttosto che puntare solo a una singola proprietà teorica come la “perfetta” spectral normalization.

Studi correlati su optimizer orthogonalized e matrix-whitening

Muon e NorMuon

Muon è il diretto predecessore di ROOT: un optimizer matrix-aware che usa Newton-Schulz per orthogonalizzare il momentum, ottenendo update con geometria più ben condizionata e complessità comparabile ad Adam. ROOT mantiene questa filosofia ma ne raffina la robustezza numerica e la gestione del rumore.

NorMuon (“Neuron-wise Normalized Muon”) è un lavoro più recente che combina l’orthogonalized update di Muon con adaptive learning rate a livello di neurone, normalizzando le norme degli update per evitare che pochi neuroni dominino il training. Sperimentazioni su modelli fino a circa 1.1B di parametri mostrano che NorMuon migliora l’efficienza di training di oltre il 20% rispetto ad Adam e di oltre il 10% rispetto a Muon, mantenendo un footprint di memoria simile a quello di Muon.

REG: Regularization Optimizer for Robust Training Dynamics

REG è un altro optimizer pensato per sostituire Muon in contesti LLM, ma adotta una strategia diversa: rimpiazza l’operatore di “matrix sign” aggressivo con un operatore di Row-and-Column-Scaling (RACS) che bilancia le righe e le colonne della matrice di update in modo più graduale. Questo approccio regolarizza gli step senza cambiare troppo la dinamica tipica di AdamW, risultando più semplice da implementare e più compatibile con pipeline di training esistenti, soprattutto in fine-tuning.

Gli autori di REG mostrano che, su diversi benchmark di ottimizzazione e NLP, il loro optimizer ottiene performance e stabilità superiori ad AdamW e Muon, in particolare evitando i cali di performance osservati quando si fa fine-tuning di modelli Qwen con Muon. REG e ROOT, pur nati da motivazioni diverse, condividono l’obiettivo di rendere più regolare e robusta la dinamica degli update in scenari LLM su larga scala.

Cosa conta davvero nei matrix-whitening optimizers?

Il paper “What Really Matters in Matrix-Whitening Optimizers?” analizza in modo sistematico la famiglia di optimizer come SOAP, Muon e Shampoo per capire quali componenti sono davvero responsabili dei loro vantaggi rispetto a metodi elementwise tipo Adam. In un framework sperimentale controllato (per esempio su GPT-2) scompone questi optimizer in due ingredienti principali: spectral normalization (orthogonalization/whitening) e variance adaptation (scaling degli update basato sulla varianza storica, simile al secondo momento di Adam).

Il risultato chiave è che la sola accuratezza della spectral normalization non spiega le differenze di performance: in alcuni casi l’optimizer con orthogonalization meno accurata ma con variance adaptation migliore batte quello teoricamente “più corretto” dal punto di vista spettrale. Questo rafforza il messaggio implicito di ROOT: lavorare solo sulla precisione dell’orthogonalization non basta, e i progressi maggiori arriveranno dalla combinazione attenta di orthogonalization, variance adaptation e meccanismi robusti contro rumore e outlier.

Torna in alto