Technische Samenvatting: Groot-schaal score-gebaseerde variationale posterior-inferentie voor Bayesiaanse diepe neurale netwerken

Probleemstelling

Bayesiaanse diepe neurale netwerken (BNN's) bieden aanzienlijke voordelen ten opzichte van deep learning met punt-schattingen, waaronder robuuste kwantificering van onzekerheid, weerstand tegen ruis en mitigatie van overfitting. Echter is exacte posterior-inferentie in BNN's over het algemeen onuitvoerbaar. Variationale Inferentie (VI) is een dominante benaderingsmethode die inferentie formuleert als een optimalisatieprobleem door de doel-posterior $\pi(\theta)$ te benaderen met een uitvoerbare verdeling $q_\lambda(\theta)$.

De huidige dominante aanpak, ELBO-gebaseerde VI (minimaliseren van de reverse KL-divergentie), lijdt onder een bekende "mode collapsing"-problematiek, waarbij de benadering faalt om alle modi van de doeldichtheid te vangen. Hoewel score-gebaseerde VI-methoden (die het verschil minimaliseren tussen de doelscore $\nabla_\theta \log \pi(\theta)$ en de modelscore $\nabla_\theta \log q_\lambda(\theta)$) theoretisch mode collapsing kunnen aanpakken, zijn bestaande score-gebaseerde benaderingen niet schaalbaar naar groot-schaal BNN's. Specifiek:

1. Computational Onuitvoerbaarheid: Methoden zoals Gaussian Score Matching (GSM) en Batch-and-Match (BaM) vertrouwen op matrixinversies of -vermenigvuldigingen die kubisch of kwadratisch schalen met de statische dimensie, waardoor ze onaanvaardbaar zijn voor hoog-dimensionale neurale netwerkparameters.
2. Incompatibiliteit met Stochastische Gradienten: Bestaande score-gebaseerde methoden vereisen doorgaans exacte scores of kunnen geen omgaan met ruisachtige, onbevooroordeelde mini-batch scores afgeleid van stochastische gradienten, welke essentieel zijn voor training op grote datasets.
3. Reparametriseringsbeperkingen: Veel methoden vertrouwen op gereparametriseerde sampling, wat computatief duur of instabiel kan zijn in hoge dimensies.

Methodologie

De auteurs stellen een nieuw Proximal Stochastic-Gradient Score-based VI-algoritme voor dat is ontworpen om te schalen naar groot-schaal BNN's, waaronder Vision Transformers (ViT) en ResNets.

Kernoptimalisatiedoelstelling

De methode update iteratief de variationale verdeling $q_t(\theta)$ naar $q_{t+1}(\theta)$ door het volgende optimalisatieprobleem op te lossen bij elke iteratie $t$:

$$q_{t+1} = \arg \min_{q \in \mathcal{Q}} \mathbb{E}_{\theta \sim q_t(\theta)} \left[ \alpha_t \|\nabla_\theta \log q(\theta) - \nabla_\theta \log q_t(\theta)\|_2^2 + \|\nabla_\theta \log q(\theta) - s(\theta)\|_2^2 \right]$$

Waarbij:

• $s(\theta) = \nabla_\theta \log \pi(\theta)$ de doelscore is.
• $\alpha_t$ een adaptieve proximaal impactparameter is (lineair gepland van 0 tot een constante $\le 1$).
• De eerste term een proximaal strafterm is die de nieuwe iteratie dwingt dicht bij de huidige iteratie te blijven in termen van scores.
• De tweede term de score-matching loss is.

Belangrijkste Technische Innovaties

1. Omgaan met Ruisachtige Mini-batch Scores: In tegenstelling tot eerdere score-gebaseerde methoden die exacte scores vereisen, maakt deze aanpak gebruik van ruisachtige, onbevooroordeelde mini-batch scores $\hat{s}(\theta)$ verkregen via stochastische gradienten. De lineariteit van de scoreterm in de verliesfunctie zorgt ervoor dat de gradient van het verlies met betrekking tot variationale parameters $\lambda$ een onbevooroordeelde schatting blijft, zelfs bij ruisachtige scores.
2. Vermijden van Reparametrering: De proximaal term maakt het mogelijk dat de verwachting wordt genomen ten opzichte van de vaste huidige verdeling $q_t$ in plaats van de variabele $q$. Dit omzeilt de noodzaak van de reparametreringstruc (gebruikt in ADVI) en vermijdt het tweemaal differentiëren ten opzichte van $\theta$ (vereist door traditionele Fisher/Score-methoden), wat de computatieve overhead en het geheugengebruik aanzienlijk reduceert.
3. Schaalbaarheid: Door het gebruik van Euclidische normen in plaats van door covariantie geïnduceerde normen (die matrixinversie vereisen) en het ondersteunen van stochastische gradienten, schaalt de methode naar netwerken met honderden miljoenen parameters (bijv. ViT-L-32 met ~305M parameters).

Theoretische Convergentie

De auteurs leveren een convergentiestelling (Stelling 3.1) die stelt dat, onder de aanname dat de variationale familie rijk genoeg is om gladde vectorvelden te representeren, de score van de variationale verdeling convergeert naar de doelscore naarmate de iteraties toenemen. Hoewel de aanname van representatieve rijkdom ideaal is, merken de auteurs op dat deze exact geldt als zowel de ware posterior als de variationale familie Gaussisch zijn.

Belangrijkste Bijdragen

1. Nieuw Algoritme: Introductie van een proximaal stochastic-gradient score-matching VI-algoritme dat score-matching loss combineert met een proximaal strafterm.
2. Schaalbaarheid naar Groot-schaal BNN's: De methode is de eerste score-gebaseerde VI-aanpak die effectief werkt op groot-schaal diepe netwerken (ResNets, ViTs) en grote datasets, waarbij de computatieve knelpunten van GSM en BaM worden overwonnen.
3. Robuustheid tegen Ruis: Het algoritme verwerkt effectief ruisachtige, onbevooroordeelde mini-batch scores, waardoor het gebruik van stochastische gradienten mogelijk is zonder volledige batch-scoreberekeningen te vereisen.
4. Eliminatie van Reparametrering: De methode vermijdt de reparametreringstruc, wat het optimalisatielandschap vereenvoudigt en de efficiëntie verbetert in vergelijking met ELBO-gebaseerde ADVI.

Experimentele Resultaten

De auteurs hebben de methode geëvalueerd op synthetische data, de PosteriorDB-benchmark en real-world groot-schaal taken:

• Synthetische Data: Op Gaussische en Gaussische mengsel-doelen toonde de voorgestelde methode snellere convergentie en betere robuustheid tegen initiële parameterkeuzes (specifiek wat betreft kleine covariantie-eigenwaarden) in vergelijking met ADVI, GSM en BaM. Het slaagde erin modelmismatch en hoog-dimensionale statenruimtes te hanteren waar GSM/BaM faalden door numerieke problemen.
• MNIST (Klein-schaal BNN): In regimes met weinig data bereikte de methode vergelijkbare testfoutpercentages als ADVI, maar toonde superieure kwantificering van onzekerheid, bewezen door lagere Expected Calibration Error (ECE) en Negative Log-Likelihood (NLL).
• Groot-schaal Visuele Herkenning (ResNet-101 & ViT-L-32): Op datasets zoals Pets, Flowers, Aircraft en DTD, boekte de methode een testvoorspellingsnauwkeurigheid die gelijk was aan die van ADVI, terwijl het consistent beter presteerde op onzekerheidsmetrieken (ECE en NLL). Concurrerende methoden (GSM, BaM, traditionele score-matching) faalden door geheugenbeperkingen of numerieke explosies.
• Tijdreeksvoorspelling: Met behulp van de Koopa-netwerkarchitectuur overtrof de methode ELBO-gebaseerde VI en niet-Bayesiaanse baselines in MSE, MAE en NLL over meerdere datasets (ECL, ETTh2, Ex. Rate, ILI, Traffic, Weather).
• Normalizing Flows: De methode werd uitgebreid naar Normalizing Flow variationale dichtheden, wat consistente verbeteringen toonde ten opzichte van ADVI op tijdreeks-taken.
• Efficiëntie: Ondanks het behouden van zowel de vorige als de huidige iteraties, waren de wandkloktijd en het GPU-geheugengebruik vergelijkbaar met ADVI, wat de praktische haalbaarheid voor groot-schaal training bewijst.

Betekenis en Claims

Het artikel claimt dat dit werk een schaalbaar, flexibel alternatief biedt voor de dominante ELBO-gebaseerde VI voor Bayesiaans diep leren op grote schaal. Door het mode-collapsing-probleem inherent aan reverse KL-doelen aan te pakken en de computatieve beperkingen van eerdere score-gebaseerde methoden te overwinnen, stelt het voorgestelde algoritme rigoureuze Bayesiaanse inferentie mogelijk op moderne, hoog-capaciteit architecturen zoals Vision Transformers.

De auteurs benadrukken dat hun aanpak empirisch effectief is over diverse benchmarks, en betere kwantificering van onzekerheid biedt zonder in te leveren op voorspellende prestaties. Ze merken bescheiden op dat hoewel theoretische convergentie berust op sterke aannames over de expressiviteit van de variationale familie, de empirische resultaten de bruikbaarheid van de methode in de praktijk ondersteunen. Het werk positioneert score-gebaseerde VI als een levensvatbare, schaalbare weg voor Bayesiaans diep leren, verdergaand dan de beperkingen van alleen-Gaussische of volledige-batch vereisten van eerdere werken.