Rennala NSGD: ottimizzazione stocastica asincrona oltre la geometria euclidea

Abstract degli autori

Il successo empirico degli ottimizzatori che tengono conto della struttura matriciale, come Muon (Jordan et al., 2024) e Shampoo (arXiv:1802.09568), mette in luce il potenziale della geometria non euclidea per il deep learning, compreso l’addestramento di modelli linguistici di grandi dimensioni (LLM). Nonostante i progressi, la comprensione teorica dell’ottimizzazione stocastica negli spazi normati in generale rimane limitata, soprattutto nei contesti distribuiti. Analizziamo il minibatch stochastic gradient descent (SGD) per l’ottimizzazione liscia non convessa rispetto a una norma arbitraria su uno spazio finito-dimensionale. Utilizzando il quadro degli spazi normati $κ$-regolari (arXiv:0809.0813), ricaviamo limiti di complessità dell’oracolo ad alta probabilità sotto una condizione di code leggere, formulata direttamente nella norma duale indotta dalla norma degli aggiornamenti. I limiti rendono esplicita la dipendenza dalla geometria attraverso la costante di regolarità $κ$ dello spazio duale. Nel caso euclideo, il nostro limite recupera il tasso standard del minibatch SGD, a meno di fattori logaritmici. Per alcune geometrie non euclidee e modelli di rumore, è più stretto dei limiti ottenuti convertendo un’ipotesi sul rumore euclideo, riducendo una dipendenza lineare dalla dimensione a una dipendenza logaritmica. Basandoci su questo quadro, introduciamo Rennala NSGD, un algoritmo distribuito semi-asincrono, e stabiliamo garanzie sulla complessità in tempo di esecuzione con tempi di calcolo dei worker fissi ed eterogenei. Nel caso euclideo ($κ=1$), il nostro limite corrisponde alla dipendenza dai parametri del Rennala SGD classico, a meno di fattori logaritmici.