Don't Stop Me Yet: Sampling Loss Minima via Dissipative Riemannian Mechanics

Riepilogo Tecnico: Don't Stop Me Yet: Minimizzazione dei Minimi di Perdita tramite Campionamento mediante Meccanica Riemanniana Dissipativa

Enunciato del Problema
Le reti neurali moderne sono tipicamente sovraparametrizzate, risultando in paesaggi di perdita in cui i minimi globali non sono punti isolati ma formano sottovarietà connesse e continue di soluzioni invarianti per riparametrizzazione. Sebbene queste soluzioni performino ugualmente bene sui dati di addestramento, possono indurre comportamenti funzionali diversi su input non visti. Gli approcci esistenti alla quantificazione dell'incertezza affrontano una dicotomia: i metodi che campionano regioni a bassa perdita spesso non riescono a imporre l'interpolazione esatta dei dati di addestramento, mentre i metodi limitati all'insieme invariante per riparametrizzazione sono spesso intrinsecamente locali, limitando la loro capacità di esplorare valli di minimi distinti. Il lavoro affronta la sfida di navigare nello spazio delle "buone soluzioni" per campionare modelli diversi che minimizzano rigorosamente la perdita di addestramento mantenendo un'alta incertezza predittiva su dati fuori distribuzione (OOD).

Metodologia: Campionatore di Minimi Dissipativi (DIMS)
Gli autori propongono il Campionatore di Minimi Dissipativi (DIMS), un sistema dinamico guidato dalla geometria progettato per campionare esattamente dai livelli minimi della funzione di perdita. L'approccio si fonda sulla meccanica Riemanniana e sui sistemi dissipativi classici.

1. Formulazione Riemanniana: La superficie di perdita è trattata come una varietà Riemanniana $(M, g)$ immersa in $\mathbb{R}^{K+1}$. Il tensore metrico $G(\theta)$ è definito come una metrica di pullback, $G(\theta) = I_K + \nabla_\theta L(\theta)\nabla_\theta L(\theta)^\top$, che incorpora la curvatura del paesaggio di perdita nella geometria dello spazio dei parametri.
2. Sistema Dinamico: Il processo di campionamento modella una particella che si muove su questa varietà. La dinamica è governata da un'equazione differenziale ordinaria (ODE) del secondo ordine derivata dalle equazioni di Euler-Lagrange. Il sistema include:
   • Energia Cinetica: $T(\alpha, \dot{\alpha}) = \frac{1}{2}\dot{\alpha}^\top G(\alpha)\dot{\alpha}$.
   • Energia Potenziale: $V(\alpha) = \kappa(L(\alpha) - L(\theta^*_{glob}))$, che agisce come un richiamo gravitazionale verso i minimi.
   • Dissipazione: Viene introdotto un termine di attrito $-\eta(t)\dot{\alpha}$ per rompere la conservazione dell'energia.
3. Attrito Dipendente dalla Velocità: Una innovazione chiave è la definizione della funzione di dissipazione $\eta(t)$, che scala con la velocità della particella (la radice quadrata dell'energia cinetica) e con l'angolo tra la velocità e il gradiente. Ciò garantisce che l'energia venga dissipata più rapidamente quando ci si muove contro il gradiente, permettendo alla particella di attraversare l'insieme di livello minimo (dove il gradiente è ortogonale alla velocità) con attrito minimo.
4. Procedura di Campionamento: Il campionatore inizializza una particella in un minimo addestrato $\theta^*$ con una velocità casuale campionata da una distribuzione (tipicamente un'approssimazione di Laplace che rispetta la geometria locale). L'ODE viene integrata numericamente finché la particella non converge a un nuovo punto sulla sottovarietà minima $S^*(D)$.

Contributi Chiave

1. Campionatore Deterministico per Sottovarietà Minime: DIMS fornisce un metodo per visitare continuamente punti distinti sulle sottovarietà minime della perdita di addestramento. È formulato come un sistema dinamico guidato dalla geometria con iperparametri interpretabili fisicamente (coefficiente di attrito $\eta_0$) che controllano l'esplorazione.
2. Garanzie Teoriche: Il lavoro stabilisce che le dinamiche proposte esplorano la superficie di perdita prima di convergere alla varietà minima in tempo finito. Nello specifico, il Teorema 1 garantisce la convergenza all'insieme degli equilibri (minimi locali) sotto la proprietà di sella stretta, mentre i Teoremi 2 e 3 forniscono limiti sul tempo di convergenza e sulla lunghezza dell'arco percorso, garantendo un'esplorazione non banale.
3. Validazione Empirica: Gli autori dimostrano che le approssimazioni posteriori concentrate sulla varietà minima generalizzano bene. DIMS migliora le prestazioni predittive, la calibrazione e specificamente la quantificazione dell'incertezza OOD rispetto alle approssimazioni basate su Laplace esistenti (LA Standard, LA Linearizzata e LA Riemanniana).

Risultati Sperimentali
Il metodo è stato valutato su compiti di regressione (dataset Snelson) e classificazione (Banana, dataset UCI, MNIST, Fashion MNIST).

• Regressione: Sul dataset Snelson, DIMS ha converguto con successo a soluzioni che adattano perfettamente i dati di addestramento mostrando alta incertezza OOD, mentre l'Approssimazione di Laplace Riemanniana (RLA) non ha garantito l'interpolazione quando inizializzata da posizioni subottimali.
• Classificazione: Sui dataset UCI, DIMS ha ottenuto metriche Negative Log-Likelihood (NLL) e di calibrazione competitive o superiori rispetto alla RLA e ad altre baseline. Notabilmente, DIMS è stato robusto all'iperparametro di attrito $\eta_0$, performando spesso bene con un valore predefinito fisso su dataset diversi.
• Prestazioni OOD: Su MNIST e Fashion MNIST, DIMS ha costantemente superato altri metodi nei punteggi AUROC per il rilevamento OOD, indicando una capacità superiore di distinguere campioni in distribuzione da campioni fuori distribuzione.

Significato e Affermazioni
Il lavoro afferma che DIMS offre un approccio fondato alla quantificazione dell'incertezza nel deep learning mirando esplicitamente alla geometria delle sottovarietà di perdita minima. A differenza delle approssimazioni posteriori tradizionali che possono distribuirsi su regioni a bassa perdita senza imporre l'interpolazione esatta dei dati, DIMS garantisce la convergenza alla varietà minima. Gli autori ipotizzano che questo approccio permetta la generazione di campioni diversi nello spazio delle funzioni che concordano sui dati di addestramento ma variano altrove, una proprietà cruciale per l'inferenza bayesiana. Il lavoro suggerisce che dinamiche ispirate alla fisica e consapevoli della geometria possono navigare efficacemente i complessi paesaggi di perdita delle reti sovraparametrizzate, offrendo un'alternativa robusta alle tecniche esistenti di campionamento e approssimazione. Gli autori riconoscono che il metodo è computazionalmente oneroso a causa della necessità di risolvere ODE in spazi ad alta dimensionalità, ma sostengono che il compromesso produca stime di incertezza superiori.