🤖 apprendimento automatico

MKOR: Momentum-Enabled Kronecker-Factor-Based Optimizer Using Rank-1 Updates

Questo articolo introduce MKOR, un ottimizzatore del secondo ordine dotato di momentum che utilizza aggiornamenti di rango 1 e l'inversione basata su Sherman-Morrison per ridurre la complessità computazionale e di comunicazione da cubica a quadratica e lineare rispettivamente, accelerando così significativamente l'addestramento e la convergenza dei grandi modelli linguistici rispetto ai metodi del primo e del secondo ordine allo stato dell'arte.

Autori originali: Mohammad Mozaffari, Sikan Li, Zhao Zhang, Maryam Mehri Dehnavi

Pubblicato 2024-01-30✓ Author reviewed ⓘ

📖 7 min di lettura🧠 Approfondimento

Autori originali: Mohammad Mozaffari, Sikan Li, Zhao Zhang, Maryam Mehri Dehnavi

✨ Questa è una spiegazione generata dall'IA dell'articolo qui sotto. Non è stata scritta dagli autori. Per precisione tecnica, consulta l'articolo originale. Leggi il disclaimer completo

Nel vasto panorama dell'intelligenza artificiale moderna, i sistemi più potenti sono costruiti su reti neurali profonde, intricate strutture matematiche che imparano a riconoscere schemi, tradurre lingue e generare testo. Per insegnare a questi sistemi, i ricercatori devono risolvere un massiccio problema di ottimizzazione: trovare l'insieme perfetto di numeri interni, o pesi, che permettano alla rete di effettuare previsioni accurate. Per decenni, l'approccio standard è stato quello di spingere questi pesi nella direzione che riduce più rapidamente l'errore, un metodo simile a un escursionista che si orienta scendendo da una montagna nebbiosa controllando la pendenza sotto i propri piedi. Questo metodo del primo ordine è affidabile e ampiamente utilizzato, ma può essere lento, richiedendo spesso migliaia di passi per raggiungere il fondo della valle.

Una strategia più ambiziosa, nota come ottimizzazione del secondo ordine, tenta di fornire all'escursionista una mappa dell'intero terreno, inclusa la curvatura del suolo, permettendo discese molto più rapide. Comprendendo come cambia la pendenza, questi metodi possono teoricamente raggiungere la soluzione ottimale in molti meno passi. Tuttavia, questo approccio è stato storicamente troppo costoso dal punto lato computazionale per essere pratico per i modelli massicci utilizzati oggi. Calcolare le informazioni sulla curvatura necessarie per una rete di grandi dimensioni richiede una potenza di elaborazione e una memoria immense, rendendo spesso il metodo più lento del semplice approccio che mira a sostituire. La sfida è stata catturare la velocità dell'escursionista che legge la mappa senza l'onere pesante di portarla con sé.

Un team di ricercatori dell'Università di Toronto, della Rutgers University e del Texas Advanced Computing Center ha sviluppato un nuovo strumento chiamato MKOR che risolve questo dilemma. Hanno creato un ottimizzatore dotato di momentum che utilizza una specifica scorciatoia matematica per approssimare la forma del terreno senza calcolare ogni dettaglio. Invece di cercare di calcolare la curvatura completa e complessa del panorama della perdita (loss landscape), MKOR utilizza una tecnica chiamata aggiornamenti di rango 1 (rank-1 updates). Questo approccio semplifica il calcolo concentrandosi sulla direzione di cambiamento più significativa, riducendo efficacementamente il costo computazionale da una relazione cubica a una quadratica. In termini pratici, ciò significa che il metodo scala molto meglio man mano che i modelli crescono, evitando i colli di bottiglia che precedentemente impedivano l'uso dei metodi del secondo ordine nei grandi modelli linguistici.

I ricercatori hanno testato questo nuovo ottimizzatore su un grande modello linguistico noto come BERT-Large-Uncased, eseguendo esperimenti su 64 processori grafici ad alte prestazioni. Hanno confrontato MKOR con i migliori metodi del primo ordine esistenti e con le implementazioni del secondo ordine più avanzate disponibili. I risultati sono stati sorprendenti. In un compito di question answering, il metodo standard del primo ordine ha richiesto 1.536 passi per raggiungere un determinato livello di precisione. Il nuovo metodo MKOR ha raggiunto la stessa precisione o una superiore in soli 600 passi. Misurando il tempo totale necessario per completare l'addestramento, MKOR ha terminato il compito in poco più di tre ore, mentre il metodo standard ha impiegato quasi otto ore. Ciò rappresenta un'accelerazione di 2,57 volte, tagliando efficacemente il tempo di addestramento di più della metà.

Oltre alla velocità pura, lo studio ha affrontato un problema critico di stabilità che spesso affligge i metodi del secondo ordine. Poiché questi metodi regolano il processo di apprendimento in base a complessi fattori matematici, sono soggetti a "gradienti esplosivi" (exploding gradients), un fenomeno in cui gli aggiustamenti diventano così grandi che il processo di addestramento va in crash o diverge. Le soluzioni precedenti spesso richiedevano di rallentare significativamente il tasso di apprendimento per mantenere la sicurezza, il che annullava i vantaggi di velocità. MKOR include uno stabilizzatore integrato che monitora la dimensione di questi aggiustamenti. Se i valori diventano troppo grandi, il sistema li riscala automaticamente, impedendo all'addestramento di esplodere pur mantenendo il ritmo aggressivo dell'approccio del secondo ordine. Ciò consente all'ottimizzatore di muoversi rapidamente senza il rischio di instabilità.

I ricercatori hanno anche osservato che il vantaggio dei metodi del secondo ordine è più pronunciato nelle fasi iniziali dell'addestramento, quando il modello sta imparando i contorni generali dei dati. Man mano che il modello si avvicina alla soluzione ottimale, i benefici dei calcoli della curvatura complessa diminuiscono. Per affrontare questo, hanno introdotto una versione ibrida dell'ottimizzatore, chiamata MKOR-H. Questa versione inizia con il potente metodo del secondo ordine per accelerare la fase iniziale di apprendimento, ma passa automaticamente a un metodo più semplice del primo ordine una volta che la perdita di addestramento smette di diminuire rapidamente. Questo approccio ibrido ha permesso al sistema di convergere in ancora meno passi, raggiungendo la stessa precisione finale del metodo standard in meno della metà del tempo.

Lo studio ha anche esaminato perché questa specifica scorciatoia matematica funzioni così bene. I ricercatori hanno analizzato le strutture dati all'interno della rete neurale e hanno scoperto che le matrici di covarianza, che descrivono come le diverse parti della rete interagiscono, sono spesso a basso rango (low-rank). Ciò significa che, nonostante l'alta dimensionalità dei dati, l'informazione è concentrata in alcune direzioni chiave. Approssimando queste matrici tramite aggiornamenti di rango 1, MKOR cattura l'informazione essenziale necessaria per l'ottimizzazione senza l'overhead dei calcoli matriciali completi. Il team ha verificato attraverso test estesi che questa approssimazione introduce pochissimo errore, confermando che la semplificazione non avviene a scapito della precisione.

Oltre alla velocità e alla stabilità, il nuovo metodo riduce significativamente la quantità di dati che devono essere trasferiti tra diversi processori durante l'addestramento. Nel calcolo distribuito, dove migliaia di processori lavorano insieme, la comunicazione può diventare un importante collo di bottiglia. MKOR riduce questa complessità di comunicazione da una relazione quadratica a una lineare, il che significa che man mano che vengono aggiunti più processori, il sistema scala efficientemente senza essere rallentato dai ritardi nel trasferimento dei dati. Questa efficienza consente al metodo di gestire le dimensioni massicce dei batch richieste dai moderni grandi modelli linguistici, che era un limite per i precedenti approcci del secondo ordine.

Le implicazioni di questo lavoro vanno oltre il semplice addestramento di un modello specifico. Rendendo l'ottimizzazione del secondo ordine pratica per sistemi su larga scala, MKOR offre una strada verso un addestramento più efficiente dell'intelligenza artificiale. I ricercatori hanno osservato che il loro metodo potrebbe aiutare a risparmiare quantità significative di energia e tempo per i professionisti del machine learning. Sebbene la tecnologia in sé sia uno strumento neutrale, la sua applicazione in campi che vanno dalla diagnostica medica alla modellazione climatica potrebbe accelerare la scoperta scientifica. Tuttavia, gli autori hanno anche riconosciuto che gli algoritmi di ottimizzazione efficienti possono essere applicati a qualsiasi modello, inclusi quelli con impatti sociali negativi, sottolineando che l'uso responsabile di tale tecnologia dipende dalle scelte delle persone che la utilizzano.

In definitiva, il documento presenta una soluzione concreta a un problema di lunga data nel deep learning. Combinando il momentum, le approssimazioni di rango 1 e un meccanismo di commutazione ibrido, i ricercatori hanno creato un ottimizzatore che è più veloce, più stabile e più scalabile dei suoi predecessori. Gli esperimenti dimostrano che è possibile sfruttare la potenza dei metodi del secondo ordine senza i costi computazionali proibitivi che li hanno tenuti in disparte. Mentre il campo dell'intelligenza artificiale continua a spingersi verso modelli più grandi e complessi, strumenti come MKOR forniscono un modo per addestrare questi sistemi in modo più efficiente, trasformando ciò che un tempo era un vantaggio teorico in una realtà pratica.

Sintesi Tecnica: MKOR – Momentum-Enabled Kronecker-Factor-Based Optimizer Using Rank-1 Updates

1. Definizione del Problema

I metodi di ottimizzazione del secondo ordine, come il Natural Gradient Descent (NGD) e il Kronecker-Factored Approximate Curvature (KFAC), offrono tassi di convergenza superiori rispetto ai metodi del primo ordine (es. SGD, Adam, LAMB) utilizzando informazioni sulla curvatura (approssimazioni della matrice Hessiana o della matrice di informazione di Fisher). Tuttavia, la loro adozione nell'addestramento di Deep Neural Networks (DNN) su larga scala, in particolare Transformer e Large Language Models (LLM), è ostacolata da tre colli di bottiglia primari:

Complessità Computazionale: I metodi KFAC tradizionali richiedono l'inversione di matrici di covarianza di dimensione $d \times d$ (dove $d$ è la dimensione dello strato), con una complessità $O(d^3)$ . Allo stesso modo, i metodi basati su Sherman-Morrison-Woodbury (SNGD) scalano spesso con la dimensione del batch $b$ ( $O(b^3)$ ). Nei Transformer, dove le dimensioni dei batch scalano con la lunghezza della sequenza, queste complessità diventano proibitive.
Overhead di Comunicazione: L'addestramento distribuito richiede la sincronizzazione delle informazioni del secondo ordine (matrici di covarianza e le loro inverse) tra i worker. KFAC richiede la sincronizzazione di $O(d^2)$ dati, mentre SNGD richiede $O(bd + b^2)$ , creando significativi colli di bottiglia nella comunicazione che limitano la scalabilità.
Frequenza di Aggiornamento e Stabilità: A causa degli elevati costi di inversione, le implementazioni allo stato dell'arte (come KAISA) aggiornano le informazioni del secondo ordine infrequentemente (es. ogni 100–1000 iterazioni) utilizzando fattori "obsoleti" (stale). Ciò riduce l'efficacia delle informazioni sulla curvatura. Inoltre, i metodi del secondo ordine sono soggetti a instabilità numerica (gradienti esplosivi, matrici singolari) e spesso richiedono fattori di smorzamento o aritmetica ad alta precisione, il che aumenta ulteriormente i costi computazionali.

2. Metodologia: L'Algoritmo MKOR

Gli autori propongono MKOR (Momentum-Enabled Kronecker-Factor-Based Optimizer Using Rank-1 Updates), che affronta i problemi sopra citati attraverso quattro meccanismi centrali:

Approssimazione di Rango-1 e Inversione basata su SM:
Invece di calcolare e invertire intere matrici di covarianza, MKOR approssima la covarianza di attivazioni e gradienti utilizzando matrici di rango-1 derivate dalla media dei campioni ( $E[aa^T] \approx \bar{a}\bar{a}^T$ ). Applicando l'identità di Sherman-Morrison a questi aggiornamenti di rango-1, MKOR aggiorna direttamente l'inversa dei fattori di Kronecker. Ciò riduce la complessità computazionale dell'inversione dei fattori da $O(d^3)$ a $O(d^2)$ .
- Regola di Aggiornamento: I fattori inversi vengono aggiornati utilizzando il momentum ( $\gamma$ ) e la formula di Sherman-Morrison, permettendo aggiornamenti frequenti senza il costo dell'inversione completa della matrice.
Complessità di Comunicazione Lineare:
In contesti distribuiti, MKOR sincronizza solo i due vettori di approssimazione di rango-1 (dimensione $2d$ ) tra i worker, invece delle intere matrici di covarianza ( $4d^2$ ). Questo riduce la complessità di comunicazione da quadratica $O(d^2)$ a lineare $O(d)$ . Inoltre, MKOR utilizza la mezza precisione (FP16) per questi vettori, dimezzando ulteriormente i costi di comunicazione, un traguardo difficile per altri metodi del secondo ordine a causa della sensibilità numerica.
Meccanismi di Stabilizzazione:
Per prevenire l'instabilità numerica e i gradienti esplosivi comuni nei metodi del secondo ordine:
- Stabilizzatore basato sulla Norma: Monitora la norma infinito dei fattori inversi. Se superano una soglia, i fattori vengono miscelati con la matrice identità ( $\hat{L} = \zeta L + (1-\zeta)I$ ) per prevenire la divergenza.
- Riscalamento del Gradiente: I gradienti precondizionati vengono riscalati per corrispondere alla norma dei gradienti originali, garantendo che gli scheduler del tasso di apprendimento rimangano efficaci.
- Stabilità Teorica: Gli autori dimostrano (Lemma 3.1) che i fattori rimangono definiti positivi senza la necessità di fattori di smorzamento artificiali ( $\mu I$ ), e che la singola divisione scalare nella regola di aggiornamento è garantita come non nulla.
MKOR Ibrido (MKOR-H):
Riconoscendo che i vantaggi del secondo ordine spesso diminuiscono con la convergenza dell'addestramento (dove il paesaggio della perdita si avvicina all'identità), MKOR-H introduce un meccanismo di commutazione. Monitora il tasso di riduzione della perdita e passa automaticamente a un ottimizzatore del primo ordine se gli aggiornamenti del secondo ordine non forniscono più accelerazione, ottimizzando il compromesso tra costo e velocità di convergenza.

3. Contributi Chiave

Riduzione della Complessità: MKOR raggiunge una complessità computazionale $O(d^2)$ e una complessità di comunicazione $O(d)$ per gli aggiornamenti del secondo ordine, alleviando i colli di bottiglia di $O(d^3)$ e $O(d^2)$ presenti in KFAC e $O(b^3)$ in SNGD.
Aggiornamenti ad Alta Frequenza: Riducendo i costi di inversione, MKOR consente aggiornamenti del secondo ordine ad ogni iterazione (o molto frequentemente), eliminando la necessità di fattori "obsoleti" e sfruttando appieno le informazioni sulla curvatura.
Robustezza: L'introduzione della stabilizzazione basata sulla norma e del riscalamento del gradiente permette a MKOR di operare stabilmente senza fattori di smorzamento o aritmetica a doppia precisione, consentendo l'uso efficiente di hardware a mezza precisione.
Strategia Ibrida: La variante MKOR-H passa dinamicamente tra metodi del secondo ordine e del primo ordine per mantenere l'efficienza durante l'intero ciclo di vita dell'addestramento.

4. Risultati Sperimentali

Gli autori hanno valutato MKOR su BERT-Large-Uncased utilizzando 64 GPU A100, confrontandolo con l'ottimizzatore del primo ordine allo stato dell'arte LAMB e con le implementazioni del secondo ordine KAISA (KFAC) e HyLo (SNGD), oltre a Eva.

Accelerazione: MKOR ha ottenuto un'accelerazione di 2.57× rispetto a LAMB e un'accelerazione di 1.85× rispetto a KAISA nel compito di pre-training. Nello specifico, ha ridotto il tempo di addestramento di BERT-Large-Uncased da 8 ore a 3 ore.
Convergenza: Nel compito di question-answering SQuAD v1.1, MKOR-H è convergente in 600 step (rispetto ai 1.536 di LAMB) ottenendo punteggi F1 comparabili o migliori (90.64 vs 90.44).
Benchmark GLUE: MKOR ha raggiunto nuovi standard di stato dell'arte sul dataset GLUE in 1.500 step. Significativamente, mentre KAISA non è riuscito a convergere verso la precisione media di base nemmeno dopo 1.563 step, MKOR è convergente con successo.
Scalabilità: A differenza di HyLo, che è divergente o non è riuscito a entrare nella memoria della GPU a causa dell'alto overhead di memoria ( $O(bd + b^2)$ ), MKOR ha addestrato con successo con grandi dimensioni di batch (fino a 64k dimensione batch effettiva tramite accumulo) senza problemi di memoria.
Validità dell'Approssimazione: Gli esperimenti su ResNet-50 e BERT hanno mostrato che le approssimazioni di rango-1 delle matrici di covarianza comportano un errore trascurabile, validando l'assunto di approssimazione principale.

5. Significato e Rivendicazioni

Il paper sostiene che MKOR colma con successo il divario tra i benefici teorici dell'ottimizzazione del secondo ordine e i vincoli pratici dell'addestramento di moderne DNN su larga scala. Riducendo la complessità degli aggiornamenti del secondo ordine a un livello comparabile ai metodi del primo ordine (in termini di scaling con la dimensione del modello) pur mantenendo i vantaggi di convergenza, MKOR abilita:

Addestramento più Rapido: Significativa riduzione del tempo di addestramento (wall-clock time) per gli LLM.
Migliore Convergenza: La capacità di convergere verso metriche di precisione più elevate dove altri metodi del secondo ordine (come KAISA) falliscono o ristagnano.
Scalabilità: La capacità di scalare verso grandi dimensioni di batch e dimensioni del modello senza i colli di bottiglia di comunicazione e memoria che precedentemente limitavano i metodi del secondo ordine.

Gli autori sottolineano che il loro lavoro rende l'ottimizzazione del secondo ordine sostenibile per l'era dei "grandi modelli", affrontando specificamente i problemi di scalabilità introdotti dalle dimensioni dei batch dipendenti dalla lunghezza della sequenza nelle architetture Transformer. Concludono che MKOR rappresenta un nuovo stato dell'arte per l'addestramento efficiente ed efficace di reti neurali profonde.