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

Technische Samenvatting: Don't Stop Me Yet: Sampling van verliesminima via Dissipatieve Riemanniaanse Mechanica

Probleemstelling
Moderne neurale netwerken zijn doorgaans overgeparameteriseerd, wat resulteert in verlieslandschappen waarbij de globale minima geen geïsoleerde punten zijn, maar samenhangende, continue subvariëteiten van herparameterisatie-invariante oplossingen vormen. Hoewel deze oplossingen even goed presteren op trainingsdata, kunnen ze diverse functionele gedragingen vertonen op onzichtbare invoer. Bestaande benaderingen voor onzekerheidskwantificering staan voor een dichotomie: methoden die lage-verliesgebieden bemonsteren, slagen er vaak niet in om exacte interpolatie van de trainingsdata af te dwingen, terwijl methoden die beperkt zijn tot de herparameterisatie-invariante set vaak inherent lokaal zijn, wat hun vermogen beperkt om verschillende mindalvallen te verkennen. Het artikel adresseert de uitdaging om de ruimte van "goede oplossingen" te navigeren om diverse modellen te bemonsteren die de trainingsverlies strikt minimaliseren terwijl ze een hoge voorspellende onzekerheid behouden op data buiten de verdeling (OOD).

Methodologie: Dissipatieve Minima Sampler (DIMS)
De auteurs stellen de Dissipatieve Minima Sampler (DIMS) voor, een door geometrie gedreven dynamisch systeem dat ontworpen is om exact te bemonsteren vanuit de minimum-niveaulijnen van de verliesfunctie. De aanpak is gebaseerd op Riemanniaanse mechanica en klassieke dissipatieve systemen.

1. Riemanniaanse Formulering: Het verliesoppervlak wordt behandeld als een Riemanniaanse variëteit $(M, g)$ ingebed in $\mathbb{R}^{K+1}$. De metriektensor $G(\theta)$ wordt gedefinieerd als een teruggetrokken metriek, $G(\theta) = I_K + \nabla_\theta L(\theta)\nabla_\theta L(\theta)^\top$, die de kromming van het verlieslandschap integreert in de geometrie van de parameter ruimte.
2. Dynamisch Systeem: Het bemonsteringsproces modelleert een deeltje dat zich op deze variëteit beweegt. De dynamiek wordt beheerst door een tweede-orde gewone differentiaalvergelijking (ODE) afgeleid uit de Euler-Lagrange-vergelijkingen. Het systeem omvat:
   • Kinetische Energie: $T(\alpha, \dot{\alpha}) = \frac{1}{2}\dot{\alpha}^\top G(\alpha)\dot{\alpha}$.
   • Potentiële Energie: $V(\alpha) = \kappa(L(\alpha) - L(\theta^*_{glob}))$, dienend als een zwaartekrachtstrek naar minima.
   • Dissipatie: Een wrijvingsterm $-\eta(t)\dot{\alpha}$ wordt geïntroduceerd om energiebehoud te doorbreken.
3. Snelheidsafhankelijke Wrijving: Een belangrijke innovatie is de definitie van de dissipatiefunctie $\eta(t)$, die schaalt met de snelheid van het deeltje (de vierkantswortel van de kinetische energie) en de hoek tussen de snelheid en de gradiënt. Dit zorgt ervoor dat energie sneller wordt gedissipeerd wanneer men tegen de gradiënt in beweegt, terwijl het deeltje de minimum-niveaulijn (waar de gradiënt orthogonaal is op de snelheid) met minimale wrijving kan doorkruisen.
4. Bemonsteringsprocedure: De sampler initialiseert een deeltje bij een getraind minimum $\theta^*$ met een willekeurige snelheid bemonsterd uit een verdeling (doorgaans een Laplace-benadering die de lokale geometrie respecteert). De ODE wordt numeriek geïntegreerd totdat het deeltje convergeert naar een nieuw punt op de minimum subvariëteit $S^*(D)$.

Belangrijkste Bijdragen

1. Deterministische Sampler voor Minimum Subvariëteiten: DIMS biedt een methode om continu verschillende punten op de minimum subvariëteiten van het trainingsverlies te bezoeken. Het is geformuleerd als een door geometrie gedreven dynamisch systeem met fysisch interpreteerbare hyperparameters (wrijvingscoëfficiënt $\eta_0$) die exploratie controleren.
2. Theoretische Waarborgen: Het artikel stelt vast dat de voorgestelde dynamiek het verliesoppervlak verkent voordat het in eindige tijd convergeert naar de minimumvariëteit. Specifiek garandeert Stelling 1 convergentie naar de set van evenwichten (lokale minima) onder de strikte-zadel-eigenschap, terwijl Stelling 2 en 3 grenzen bieden aan de convergentietijd en de afgelegde booglengte, waardoor niet-triviale exploratie wordt gewaarborgd.
3. Empirische Validatie: De auteurs demonstreren dat posterior-benaderingen geconcentreerd op de minimumvariëteit goed generaliseren. DIMS verbetert de voorspellende prestaties, kalibratie en specifiek OOD-onzekerheidskwantificering in vergelijking met bestaande Laplace-gebaseerde benaderingen (Standaard LA, Linearized LA en Riemanniaanse LA).

Experimentele Resultaten
De methode werd geëvalueerd op regressie (Snelson dataset) en classificatietaken (Banana, UCI datasets, MNIST, Fashion MNIST).

• Regressie: Op de Snelson dataset convergerde DIMS succesvol naar oplossingen die trainingsdata perfect pasten terwijl ze hoge onzekerheid OOD vertoonden, terwijl de Riemanniaanse Laplace-benadering (RLA) er niet in slaagde interpolatie te garanderen wanneer geïnitieerd vanuit suboptimale posities.
• Classificatie: Op UCI-datasets behaalde DIMS concurrerende of superieure Negative Log-Likelihood (NLL) en kalibratiemetingen in vergelijking met RLA en andere baselines. Opvallend was dat DIMS robuust was voor de wrijvingshyperparameter $\eta_0$, vaak goed presterend met een vaste standaardwaarde over diverse datasets.
• OOD-prestaties: Op MNIST en Fashion MNIST overtrof DIMS andere methoden consistent in AUROC-scores voor OOD-detectie, wat wijst op een superieur vermogen om in-distribution te onderscheiden van out-of-distribution samples.

Betekenis en Claims
Het artikel claimt dat DIMS een principiële aanpak biedt voor onzekerheidskwantificering in deep learning door expliciet te richten op de geometrie van de minimum-verliesvariëteiten. In tegenstelling tot traditionele benaderde posteriors die zich over lage-verliesgebieden kunnen verspreiden zonder exacte data-interpolatie af te dwingen, garandeert DIMS convergentie naar de minimumvariëteit. De auteurs stellen dat deze aanpak toelaat om diverse functionele ruimte-samples te genereren die overeenkomen op trainingsdata maar elders variëren, een eigenschap die cruciaal is voor Bayesiaanse inferentie. Het werk suggereert dat door de natuurkunde geïnspireerde, geometrie-bewuste dynamiek effectief de complexe verlieslandschappen van overgeparameteriseerde netwerken kan navigeren, en een robuust alternatief biedt voor bestaande bemonsterings- en benaderingstechnieken. De auteurs erkennen dat de methode rekenkundig veeleisend is vanwege de noodzaak om ODE's in hoogdimensionale ruimten op te lossen, maar betogen dat de afweging leidt tot superieure onzekerheidsschattingen.