Ottimizzazione consapevole dei moduli per le reti neurali a grafo

Abstract degli autori

Ottimizzatori diversi impongono geometrie diverse agli aggiornamenti dei parametri. I metodi adattivi come Adam e AdamW applicano aggiornamenti per coordinata, mentre gli ottimizzatori che tengono conto delle matrici, come Muon, sfruttano la struttura delle matrici dei pesi attraverso aggiornamenti ortogonalizzati. Nelle reti neurali a grafo (GNN), tuttavia, la geometria effettiva delle matrici dei pesi è determinata non solo dai pesi appresi che mescolano i canali, ma anche dalle rappresentazioni propagate attraverso il grafo su cui agiscono. Di conseguenza, anche se un aggiornamento è ben normalizzato come matrice, il suo effetto sull’output del layer può essere amplificato quando viene applicato a caratteristiche propagate attraverso il grafo. Per colmare questa lacuna, introduciamo GraphMuon, un ottimizzatore consapevole del grafo che adatta gli aggiornamenti matriciali di Muon alla struttura di covarianza degli input del layer propagati attraverso il grafo. GraphMuon esegue l’ortogonalizzazione in una metrica definita da queste rappresentazioni propagate e normalizza l’ampiezza dei passi in base all’effetto dell’aggiornamento sulla rappresentazione del layer. Mostriamo che, sul piano teorico, Muon nella sua forma standard può subire un’amplificazione degli aggiornamenti indotta dal grafo, mentre GraphMuon può attenuare questo effetto. GraphMuon fornisce quindi un analogo, nella metrica del grafo, del passo di aggiornamento polare di Muon. Sul piano empirico, GraphMuon ottiene accelerazioni da $1.32\times$ a $2.42\times$ nel tempo necessario a raggiungere il target di loss di validazione rispetto ad AdamW, su tre benchmark OGB su larga scala. In 15 configurazioni di dataset e architettura che comprendono classificazione dei nodi, classificazione induttiva dei grafi e regressione sui grafi, raggiunge un tempo inferiore per arrivare al target in 13 casi, mantenendo prestazioni predittive competitive o superiori. Il codice è disponibile su https://github.com/guy-hadad/GraphMuon.