🤖 machine learning

MESH: Memory-Efficient Sinkhorn Optimization for Mixture-of-Experts Training

Questo articolo introduce MESH, un metodo di ottimizzazione Sinkhorn efficiente dal punto di vista della memoria che incorpora aggiornamenti del momento nascosto e precondizionamento a blocchi opzionale per abilitare con successo l'addestramento a bassa memoria di modelli Mixture-of-Experts, affrontando l'instabilità temporale dei gradienti degli esperti instradati che causa il fallimento dei metodi Sinkhorn stateless standard.

Autori originali: Masato Fujitake

Pubblicato 2026-08-06

📖 4 min di lettura☕ Lettura da pausa caffè

CC BY 4.0

Autori originali: Masato Fujitake

Articolo originale sotto licenza CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Questa è una spiegazione generata dall'IA dell'articolo qui sotto. Non è stata scritta né approvata dagli autori. Per precisione tecnica, consulta l'articolo originale. Leggi il disclaimer completo

Immagina di cercare di insegnare a un robot gigante e super intelligente come scrivere storie. Per rendere questo robot davvero bravo, hai bisogno di un cervello enorme fatto di miliardi di minuscoli interruttori. Ma c'è il problema: il cervello del robot è così grande che entra a malapena nella memoria del computer. È come cercare di riprodurre un film blockbuster su un telefono con una batteria minuscola; il telefono si scalda, rallenta e potrebbe bloccarsi.

Per risolvere questo problema, gli scienziati usano dei "allenatori" speciali chiamati ottimizzatori. L'allenatore più famoso, chiamato AdamW, è come un coach molto attento che tiene un diario dettagliato per ogni singolo interruttore del cervello del robot, ricordando esattamente come dare un colpetto la prossima volta. Questo funziona benissimo, ma quei diari occupano un sacco di spazio. Così, i ricercatori hanno inventato un allenatore più leggero chiamato Sinkhorn. È come un coach che non tiene alcun diario; invece, guarda solo l'errore attuale e fa un rapido aggiustamento senza usare la memoria. Questo risparmia una quantità enorme di spazio, ma si scopre che per un tipo specifico di cervello robotico chiamato "Mixture-of-Experts" (MoE), questo coach senza memoria si confonde e il robot smette di imparare bene. La domanda che gli scienziati si pongono è: possiamo mantenere i vantaggi del risparmio di memoria del coach veloce senza perdere la capacità del robot di imparare?

Questo articolo presenta un nuovo metodo chiamato MESH (Sinkhorn per Esperti a Memoria Efficiente) per risolvere esattamente questo problema. I ricercatori hanno scoperto che il coach veloce e senza memoria fallisce perché gli "esperti" nel cervello del robot non sempre riescono a lavorare. In un modello Mixture-of-Experts, il robot sceglie solo alcuni esperti specifici per gestire ogni frase, un po' come un ristorante dove vengono chiamati in cucina solo alcuni chef per un ordine specifico. Poiché gli chef cambiano da un ordine all'altro, il coach veloce vede un segnale confuso e rumoroso e si perde.

L'articolo suggerisce che la soluzione non è dare al coach un diario completo, ma dargli una "memoria nascosta". Invece di guardare solo l'ordine attuale, il coach ora ricorda la media dei precedenti ordini prima di prendere una decisione. L'autore chiama questo "momento nascosto" (hidden momentum). Hanno testato questo metodo su un piccolo modello da 110 milioni di parametri (un "nanowhale") e hanno scoperto che questo semplice trucco — levigare il rumore prima di regolare i pesi — ripristina la capacità di apprendimento del robot.

Ecco cosa hanno scoperto e cosa hanno escluso:

La Soluzione Principale: L'articolo mostra che il fallimento avviene perché gli "esperti instradati" (gli chef che lavorano solo a volte) hanno bisogno che i loro segnali vengano levigati nel tempo prima che avvenga l'aggiustamento senza memoria. Mediando la storia recente di questi esperti (usando un buffer nascosto che non conta come memoria permanente), il nuovo metodo MESH funziona molto meglio.
Cosa Non Ha Funzionato: I ricercatori hanno provato molte altre idee per risolvere il problema, ma le hanno scartate. Hanno scoperto che semplicemente rendere gli aggiustamenti più piccoli, cambiare il modo in cui gli esperti sono raggruppati o usare solo informazioni di "segno" (sapere solo se l'errore è salito o sceso, non di quanto) non ha risolto il problema. La chiave era specificamente la levigatura temporale, non solo un tipo qualsiasi di dati extra.
Quanto Sono Sicuri? I risultati si basano su esperimenti con un modello piccolo specifico. In questi test, il nuovo metodo ha ridotto lo spazio occupato dall'ottimizzatore dello stato del 62,5% (scendendo da 0,883 GB a 0,331 GB) e ha ridotto l'uso della memoria di picco del computer di circa il 12,6%. Tuttavia, il punteggio finale del test del robot (perdita di valutazione) era ancora leggermente peggiore rispetto al pesante coach che tiene il diario (AdamW), con punteggi intorno a 3,64 per il nuovo metodo rispetto a 3,59 per il vecchio. L'autore suggerisce che questo piccolo divario potrebbe essere dovuto al fatto che altre parti del cervello del robot (come il vocabolario) hanno ancora bisogno del coach pesante, ma sono fiduciosi che il "momento nascosto" sia il pezzo mancante per gli esperti.

In breve, l'articolo sostiene che non serve un diario completo per addestrare questi cervelli robotici speciali; serve solo un modo per ricordare il passato recente senza scriverlo in modo permanente. Questo approccio del "momento nascosto", chiamato MESH, permette di ottenere quasi tutto il necessario, risparmiando molta memoria pur mantenendo il robot abbastanza intelligente da imparare.

Riepilogo Tecnico: MESH – Ottimizzazione Sinkhorn a Efficienza di Memoria per l'Addestamento di Mixture-of-Experts

1. Definizione del Problema

Il paper affronta la sfida di ottimizzare i modelli linguistici Mixture-of-Experts (MoE) con efficienza di memoria. Sebbene AdamW sia l'ottimizzatore standard per il pretraining grazie alla sua combinazione di smoothing del primo momento temporale e scalatura adattiva per coordinata, esso comporta un elevato costo di memoria memorizzando due interi tensori di stato per ogni parametro.

Le alternative efficienti dal punto di vista della memoria, come gli ottimizzatori matriciali di tipo Sinkhorn e gli ibridi SAGE/Sinkhorn, riducono questo costo eliminando gli stati del momento per le matrici dense. Tuttavia, l'autore dimostra che l'applicazione diretta di questi metodi stateless agli esperti MoE fallisce. In un setting controllato "nanowhale" in stile DeepSeek da 110 milioni di parametri, un tipico ibrido SAGE/Sinkhorn ha ridotto significativamente lo stato dell'ottimizzatore (da 0,883 GB a 0,331 GB) ma ha causato una grave degradazione della perdita di valutazione (3,8265 rispetto al baseline di AdamW di 3,58–3,64).

Il fallimento principale è identificato come un problema di ordinamento specifico per gli esperti routati:

I Gradienti Routati sono Condizionali: A differenza dei gradienti delle MLP dense, i gradienti degli esperti routati sono campioni condizionali basati sulle decisioni del router e sui pesi di gating, che variano temporalmente.
Non linearità di Sinkhorn: La normalizzazione Sinkhorn è un'operazione non lineare. Applicare Sinkhorn a un gradiente routato istantaneo e ad alta varianza ( $Sinkhorn(G_t)$ ) si comporta diversamente rispetto all'applicazione su un segnale mediato temporalmente ( $Sinkhorn(E[G_t])$ ).
Amplificazione del Rumore: Il Sinkhorn stateless normalizza il rumore indotto dal routing ad ogni step, mentre il fallimento dei metodi stateless deriva dalla mancanza di smoothing temporale prima che avvenga la normalizzazione della matrice.

2. Metodologia: MESH e MESH-B

L'autore propone MESH (Memory-Efficient Sinkhorn for Experts), un ottimizzatore ibrido progettato specificamente per le architetture MoE.

Meccanismo Core: Hidden Momentum

MESH sostituisce l'aggiornamento Sinkhorn stateless per gli esperti MoE con un approccio a hidden-momentum (momento nascosto). Invece di memorizzare un esplicito tensore del primo momento nello stato dell'ottimizzatore (il che annullerebbe l'obiettivo di risparmio di memoria), MESH mantiene il primo momento all'interno del ciclo di vita del buffer del gradiente:

Accumulazione: Viene mantenuto un accumulatore di momentum dell'esperto non normalizzato: $H_{e,t} = \beta_1 H_{e,t-1} + G_{e,t}$ .
Ciclo di Vita del Buffer: L'accumulatore $H_{e,t}$ è memorizzato nel buffer del gradiente anziché nel dizionario dello stato dell'ottimizzatore.
Ciclo di Aggiornamento: Prima del passaggio backward, il buffer viene riattaccato come gradiente ( $p.grad \leftarrow \beta_1 H_{e,t}$ ). Il passaggio backward aggiunge il gradiente corrente, aggiornando efficacementamente l'accumulatore per il passo successivo.
Normalizzazione: L'aggiornamento Sinkhorn è calcolato sul segnale smussato: $U_{e,t} = Sinkhorn(H_{e,t})$ .

Questo approccio ripristina il segnale del primo momento temporale richiesto per gli esperti routati senza memorizzare un esplicito tensore di stato dell'ottimizzatore per essi.

Estensione Opzionale: MESH-B

Il paper valuta anche MESH-B, che aggiunge un moltiplicatore inverse-RMS a livello di blocco/neurone alla direzione Sinkhorn. Questo funge da precondizionatore a livello di blocco. Le statistiche del blocco sono calcolate dal segnale sorgente smussato ( $H_{e,t}$ ) piuttosto che da un flusso separato di gradienti grezzi per mantenere la semplicità di implementazione.

Strategia di Assegnazione dei Ruoli

MESH mantiene una strategia ibrida per diverse classi di parametri:

Embedding / LM Head: Gestiti da SAGE (aggiornamenti a stato leggero).
Matrici Dense Non-Expert: Gestite da Sinkhorn stateless.
Esperti Routati/Shared: Gestiti da MESH (Sinkhorn con hidden-momentum).
Norme / Bias: Gestiti da SAGE.

3. Contributi Chiave

Identificazione del Punto di Fallimento: Il paper identifica le matrici degli esperti MoE routati come il punto di fallimento specifico dell'ottimizzazione diretta SAGE/Sinkhorn nel pretraining, distinguendole dalle matrici dense e dagli embedding.
Algoritmo MESH: Introduzione di un aggiornamento Sinkhorn con hidden-momentum che ripristina un segnale del primo momento per gli esperti senza aggiungere uno stato dell'ottimizzatore esplicito, utilizzando il ciclo di vita del buffer del gradiente.
Analisi di Ablazione: Analisi negative sistematiche che mostrano come le interviste basate solo sul secondo momento, solo segno/momento, solo confidenza o solo scala non spieghino il divario di prestazioni. L'ingrediente causale primario è identificato come lo smoothing temporale prima della normalizzazione della matrice.
Valutazione MESH-B: Valutazione di una variante precondizionata a blocchi/neuroni, che mostra come, sebbene possa migliorare la frontiera memoria-qualità in determinati seed, non sia universalmente necessaria rispetto allo smoothing temporale fornito da MESH.

4. Risultati Sperimentali

Gli esperimenti sono stati condotti su un modello MoE in stile DeepSeek da 110 milioni di parametri addestrato su FineWeb-Edu per 5.000 step.

Efficienza di Memoria:
- AdamW: 0,883 GB di stato dell'ottimizzatore.
- MESH / MESH-B: 0,331 GB di stato dell'ottimizzatore (una riduzione del 62,5% rispetto ad AdamW).
- Allocazione CUDA di Picco: Ridotta di circa il 12,6% (da ~4,37 GB a ~3,82 GB).
Prestazioni (Perdita di Valutazione):
- Baseline AdamW: ~3,5877 (media tra i seed 100/512).
- SAGE/Sinkhorn Stateless: 3,8265 (Degradazione significativa).
- MESH: 3,6430 (Vicino ad AdamW, ma leggermente superiore).
- MESH-B: 3,6365 (Leggermente migliore di MESH, ma comunque inferiore ad AdamW).
Risultati Diagnostici (Seed 42):
- Le varianti a stato completo (MESH-FULL, MESH-B-FULL) che memorizzano esplicitamente il primo momento recuperano una qualità simile ad AdamW (ad esempio, MESH-B-FULL con V(H) ha raggiunto 3,6102) con circa la metà dello stato dell'ottimizzatore di AdamW.
- Ciò conferma che lo smoothing temporale è l'ingrediente mancante critico, mentre l'implementazione dell'hidden-momentum approssima con successo questo processo senza il costo del pieno stato.

5. Significato e Rivendicazioni

Il paper sostiene modestamente che MESH fornisce un punto con stato significativamente inferiore sulla frontiera memoria-qualità per l'addestramento MoE.

Conclusione Primaria: Gli esperti MoE richiedono lo smoothing temporale prima della normalizzazione della matrice per funzionare efficacemente con gli ottimizzatori di tipo Sinkhorn. Questo smoothing è il fattore causale dominante per colmare il divario di prestazioni con AdamW.
Conclusione Secondaria: Il precondizionamento a blocchi/neurone (MESH-B) può migliorare i risultati in configurazioni specifiche, ma non è stabilito come componente universalmente necessario.
Limitazioni: L'autore riconosce un divario di prestazioni residuo tra MESH/MESH-B e AdamW. Ipotizza che ciò sia dovuto ai ruoli non-expert (specificamente le matrici del vocabolario e le matrici dense non-expert), che rimangono su SAGE o Sinkhorn stateless, mentre AdamW le adatta con pieni momenti.
Ambito: Lo studio è intenzionalmente ristretto, concentrandosi su un'architettura da 110M e run di 5.000 step. L'autore non afferma che MESH superi AdamW in tutti i parametri, ma che recupera la maggior parte della qualità persa dai metodi stateless, riducendo drasticamente l'overhead di memoria.

In sintesi, MESH dimostra che il fallimento degli ottimizzatori efficienti in termini di memoria nell'addestramento MoE non è inerente alla struttura MoE stessa, ma piuttosto alla mancanza di smoothing temporale per i gradienti routati. Implementando l' "hidden momentum", MESH risolve questo problema di ordinamento, offrendo una via percorribile verso un addestramento MoE efficiente dal punto di vista della memoria.