Technische Samenvatting: Het sturen van posterieure exploratie met optimizer-afgeleide geometrie

Probleemstelling

Bayesiaans Diep Leren (BDL) streeft naar een gefundeerde kwantificering van onzekerheid voor neurale netwerken, een cruciale vereiste voor toepassingen met een hoog risico. Hoewel steekproefgebaseerde methoden zoals Markov Chain Monte Carlo (MCMC) een rigoureuze aanpak bieden voor het benaderen van posterieure verdelingen, wordt hun praktische toepassing gehinderd door de enorme computationele kosten van het verkennen van hoogdimensionale, multimodale en slecht geconditioneerde posterieure landschappen.

Recente schaalbare inferentiestrategieën, zoals de twee-fasen Bayesian Deep Ensemble (BDE), proberen dit te verzachten door eerst te optimaliseren om meerdere onderscheidende modi te vinden (vergelijkbaar met Deep Ensembles) en vervolgens onafhankelijke MCMC-ketens te initialiseren vanuit deze oplossingen om lokale attractiebekkens te verkennen. Echter, een significante inefficiëntie blijft bestaan: de optimalisatie- en steekproefstadia worden doorgaans als gescheiden processen behandeld. Bijgevolg wordt de geometrische informatie over het verlieslandschap (specifiek de lokale kromming) die tijdens de optimalisatiefase is verzameld, weggegooid. De daaropvolgende steekproeffase wordt gedwongen om deze lokale kromming vanaf nul opnieuw te ontdekken, wat vaak een langdurige en kostbare "warmup" of schaaladaptatieperiode vereist om de preconditioning-matrix van de sampler aan te passen.

Methodologie

Dit artikel stelt een methode voor om de kloof tussen optimalisatie en steekproeven te overbruggen door statistieken die worden gegenereerd door adaptieve optimizers (specifiek AdamW) te hergebruiken om de initialisatie van Stochastic Gradient MCMC (SGMCMC) samplers te informeren.

Optimizer-afgeleide Geometrie

Adaptieve optimizers zoals AdamW houden exponentiële voortschrijdende gemiddelden bij van gradiënten: een eerste-moment-vector ($\mu_t$) die fungeert als momentum en een tweede-moment-vector ($\nu_t$) die dient als een proxy voor gradiëntvariantie (en daarmee voor lokale kromming). De auteurs stellen dat deze statistieken, berekend als bijproduct van de optimalisatiefase, dienen als een "gratis", hoogwaardige preconditioner voor de daaropvolgende steekproeffase.

Preconditioning Strategieën

De auteurs formaliseren verschillende strategieën om deze optimizer-statistieken over te dragen naar de massamatrix $M$ van een SGHMC-sampler:

1. Gradiënt-gebaseerd: $M = \text{diag}(g_T^2 + \epsilon I)$, gebruikmakend van de laatste gradiënt. Dit is puur lokaal maar gevoelig voor ruis.
2. Eerste-moment: $M = \text{diag}(\mu_T^2 + \epsilon I)$, gebruikmakend van de momentumterm. Dit biedt temporele uitvlakking.
3. Tweede-moment: $M = \text{diag}(\sqrt{\nu_T} + \epsilon I)$, gebruikmakend van de gekwadrateerde gradiënt-accumulator. Dit benadert de diagonale Fisher Informatie Matrix (FIM) en vangt robuuster globale kromminginformatie op.

Laag-gewijze Normalisatie

Om numerieke stabiliteit te verbeteren en extreme stapgroottes voor individuele parameters te voorkomen, evalueren de auteurs een laag-gewijze strategie. In plaats van een unieke preconditioning-waarde toe te wijzen aan elke parameter, berekenen zij een enkele scalar per laag door de optimizer-statistieken van die laag te middelen en over de lagen te normaliseren. Dit fungeert als een regularisator en stelt een bovengrens aan de maximale parameter-gewijze stapgrootte.

Theoretische Rechtvaardiging

Het artikel biedt een theoretische link (Propositie 3.1) die aantoont dat als de posterieure potentiaal lokaal kwadratisch is, het instellen van de massamatrix $M$ gelijk aan de Hessiaan $H^*$ bij de modus, een isotrope Gaussische benadering met een eenheidsconditienummer induceert. Dit bereikt optimale lokale menging. Omdat de volledige Hessiaan onhandelbaar is, dient de tweede-moment-statistiek $\nu_T$ als een computationeel efficiënte diagonale proxy voor $H^*$, wat effectief de "effectieve metriek-tensor" overdraagt die het optimalisatielandschap succesvol heeft genavigeerd naar de steekproeffase.

Belangrijkste Bijdragen

1. Theoretisch Kader: De auteurs leiden theoretisch gefundeerde methoden af voor het gebruik van optimizer-eigenschappen (specifiek van AdamW) om SGMCMC-samplers te preconditioneren, waarmee de optimalisatie- en steekproeffasen worden overbrugd.
2. Empirische Evaluatie: Zij evalueren meerdere strategieën voor het overdragen van optimizer-afgeleide geometrie over diverse modaliteiten (tabulair, beeld, taal) en architecturen (FCNs, CNNs, ResNets, Transformers).
3. Efficiëntie- en Stabiliteitsanalyse: De studie beoordeelt het vermogen van deze strategieën om de convergentie van samplers te versnellen, de behoefte aan warmup te verminderen of te elimineren, en de robuustheid ten opzichte van hyperparameters (specifiek stapgrootte) te verbeteren zonder extra computationele kosten.
4. Geometrische Analyse: Het artikel analyseert lokale geometrische patronen in realistische BNN's, waarbij wordt onthuld dat statische, optimizer-afgeleide preconditioning de onderscheidende geometrische informatie van elke gevonden modus in Fase I behoudt, terwijl adaptieve samplers (zoals AdaSGHMC) deze initiële geometrie hebben de neiging te "overschrijven", wat de ensemble-diversiteit vermindert.

Resultaten

De voorgestelde aanpak werd geëvalueerd op datasets inclusief bikesharing, protein, F-MNIST, CIFAR-10 en een gemoderniseerde Shakespeare-dataset.

• Voorspellende Prestaties: Optimizer-afgeleide preconditioning behoudt consistent of verbetert de voorspellende prestaties (RMSE en Accuraatheid) vergeleken met standaard Deep Ensembles en Vanilla SGHMC.
• Onzekerheidskwantificering (UQ): De methode levert significante verbeteringen op in UQ-metrieken (Log Pointwise Predictive Density en Perplexity), waarbij het de Deep Ensemble baseline vaak met meer dan 10% in LPPD overtreft.
• Eliminatie van Warmup: Strategieën gebaseerd op gradiënten en eerste-momenten bereiken hun piekprestaties onmiddellijk met nul warmup, waardoor de kostbare schaaladaptatie-fase effectief wordt geëlimineerd. De tweede-moment-strategie profiteert van een langere warmup, maar presteert nog steeds beter dan de baselines.
• Robuustheid: Preconditioned varianten behouden stabiele prestaties over een breed scala aan stapgroottes, terwijl Vanilla SGHMC scherp degradeert bij suboptimale stapgroottes.
• Vergelijking met Adaptieve Methoden: Verrassend genoeg presteert de statische optimizer-afgeleide preconditioning beter dan adaptieve methoden zoals AdaSGHMC. De auteurs schrijven dit toe aan het feit dat de interne adaptatie van AdaSGHMC de initiële geometrische diversiteit van de ensemble-ketens overschrijft, wat leidt tot zelf-gelijke massamatrices en verminderde generalisatie.

Betekenis

Het artikel stelt dat het overdragen van optimizer-afgeleide geometrie naar SGMCMC-samplers een effectieve, kosteloze strategie is voor het verbeteren van Bayesian Deep Ensembles. Door gebruik te maken van statistieken die al tijdens de optimalisatiefase zijn berekend, stelt de methode de gebruiker in staat om:

• Computationele Kosten te Verlagen: Het elimineert of verkort aanzienlijk de sampling warmup-fase, wat een efficiënter pad naar posterieure exploratie biedt.
• Stabiliteit te Verbeteren: Het biedt een numeriek stabiel startpunt dat robuust is tegen de keuze van hyperparameters.
• Diversiteit te Preserveren: Door een statische preconditioner te gebruiken die afgeleid is van de specifieke modus gevonden in Fase I, wordt de geometrische diversiteit van het ensemble behouden, wat volgens de auteurs cruciaal is voor superieure generalisatie vergeleken met adaptieve samplers die de geometrie van het ensemble homogeniseren.

De auteurs concluderen dat deze aanpak het mogelijk maakt voor praktijkgebruikers om state-of-the-art onzekerheidskwantificering en voorspellende prestaties te bereiken met minimale extra computationele overhead, waardoor Bayesiaanse inferentie toegankelijker wordt voor domeinen met een hoog risico.