How the Optimizer Shapes Learned Solutions in Equivariant Neural Networks

Riepilogo Tecnico: Come l'Ottimizzatore Modella le Soluzioni Apprese nelle Reti Neurali Equivarianti

Enunciato del Problema
Le reti neurali equivarianti codificano le simmetrie geometriche direttamente nella loro architettura, fornendo un pregiudizio induttivo naturale per dati con simmetrie geometriche note. Tuttavia, i vincoli di equivarianza rigidi rendono spesso l'ottimizzazione difficile, portando a paesaggi di perdita contenenti punti critici e minimi spurii. Di conseguenza, questi modelli possono performare peggio rispetto ad alternative meno vincolate su larga scala. La risposta dominante nella letteratura recente è stata quella di ammorbidire i vincoli architetturali attraverso metodi di equivarianza approssimata o rilassata. Al contrario, il ruolo dell'ottimizzatore nel modellare le soluzioni dei modelli equivarianti rimane poco esplorato. Questo lavoro indaga se la scelta dell'ottimizzatore, indipendentemente dalle modifiche architetturali, possa alterare significativamente le soluzioni raggiunte durante l'addestramento e migliorare le prestazioni.

Metodologia
Gli autori conducono uno studio comparativo tra Adam (Kingma & Ba, 2015) e Muon (Jordan et al., 2024), un ottimizzatore proposto di recente che ortogonalizza il suo buffer di momento tramite iterazioni di Newton-Schulz. La valutazione abbraccia tre contesti distinti:

1. Classificazione di Nuvole di Punti 3D: Utilizzando il dataset ModelNet40 (e la sua variante corrotta, ModelNet40-C) con tre architetture dotate di diversi pregiudizi induttivi geometrici:
   • EGNN: Rete neurale a grafo esplicitamente equivariante sotto $E(3)$.
   • DGCNN: Grafo CNN equivariante per permutazione con connettività $k$-NN dinamica.
   • PointNet: Architettura invariante per permutazione tramite pooling simmetrico.
2. Apprendimento Molecolare: Utilizzando il dataset QM9 con GotenNet (equivariante sotto $E(3)$) e GINE (equivariante per permutazione) sui dataset peptide e ZINC.
3. Analisi Geometrica: Per comprendere perché le prestazioni differiscono, gli autori analizzano i checkpoint addestrati utilizzando:
   • Geometria Locale della Perdita: Visualizzazione di fette 2D della perdita e superfici 3D attorno ai checkpoint, e stima di riassunti dell'Hessiano (autovalore massimo e traccia) tramite iterazione di potenza e stimatori di Hutchinson.
   • Struttura Spettrale: Calcolo del rank stabile e del rank effettivo (basato sull'entropia di Shannon dei valori singolari) sia per le matrici dei pesi che per le rappresentazioni degli strati intermedi.

Risultati Chiave

• Miglioramenti delle Prestazioni: Su ModelNet40, Muon supera costantemente Adam su tutte le architetture (EGNN, PointNet, DGCNN) sia in contesti puliti che corrotti. Ad esempio, su ModelNet40 pulito, Muon ha migliorato l'accuratezza di EGNN del +5,17%, di PointNet del +2,67% e di DGCNN del +1,96%. Tendenze simili sono state osservate sul dataset molecolare QM9, dove Muon ha migliorato la maggior parte delle proprietà target.
• Caratteristiche del Paesaggio di Perdita: Le visualizzazioni della superficie di perdita attorno ai checkpoint addestrati rivelano che le soluzioni di Muon esibiscono una geometria locale più liscia e regolare rispetto ai contorni irregolari spesso osservati con Adam. Tuttavia, le stime dell'Hessiano contraddicono l'intuizione secondo cui una geometria più liscia implichi una curvatura inferiore; i checkpoint di Muon possiedono in realtà autovalori massimi e stime di traccia più grandi (indicando una curvatura più elevata) rispetto ai checkpoint di Adam.
• Proprietà Spettrali: Muon produce costantemente rank stabili ed effettivi più elevati sia per i pesi che per le rappresentazioni intermedie attraverso tutti gli strati e le architetture. Ciò suggerisce che Muon distribuisce la capacità lungo le direzioni singolari in modo più uniforme, evitando la concentrazione degli spettri in poche direzioni dominanti spesso associate ai metodi basati sul gradiente.

Significato e Affermazioni
Il lavoro postula che l'ottimizzatore sia una fonte significativa di pregiudizio induttivo che interagisce con i vincoli geometrici. Gli autori affermano che:

1. La Scelta dell'Ottimizzatore Conta: Cambiare l'ottimizzatore da Adam a Muon può migliorare le prestazioni nei modelli equivarianti senza modificare l'architettura o rilassare i vincoli.
2. Spazi di Soluzione Distinti: Muon e Adam convergono verso soluzioni qualitativamente diverse, caratterizzate da geometrie locali della perdita e strutture spettrali differenti.
3. Ipotesi sul Meccanismo: Gli aggiornamenti del momento ortogonalizzati in Muon sembrano fornire un bias a livello di traiettoria verso regioni più favorevoli del paesaggio di perdita, trattando le direzioni singolari degli aggiornamenti in modo più uniforme.
4. Principio di Progettazione: I risultati suggeriscono un potenziale principio di progettazione per ottimizzatori equivarianti: regole di aggiornamento che impediscono agli spettri singolari dei pesi e delle rappresentazioni di concentrarsi in poche direzioni possono migliorare l'addestramento.

Gli autori mantengono una posizione modesta, notando che, sebbene la correlazione tra le proprietà spettrali e i guadagni di accuratezza sia coerente, un nesso causale non è stato stabilito. Concludono che l'interazione tra la progettazione dell'ottimizzatore e il pregiudizio induttivo geometrico merita ulteriori indagini, in particolare riguardo a come gli ottimizzatori possano sfruttare la struttura dei gruppi di simmetria invece di trattare l'equivarianza come un vincolo a scatola nera.