A Framework for Variational Inference of Lightweight Bayesian Neural Networks with Heteroscedastic Uncertainties

1. Enunciato del Problema

Le Reti Neurali Bayesiane (BNN) sono essenziali per quantificare l'incertezza predittiva, distinguendo tra incertezza epistemica (incertezza del modello dovuta a dati limitati) e incertezza aleatoria (rumore dei dati).

• La Sfida: In molte applicazioni, in particolare quelle che richiedono incertezze eteroschedastiche (dove i livelli di rumore variano in base all'input), gli approcci standard richiedono che la BNN outputti la varianza aleatoria come un'ulteriore testa di rete accanto alla media predittiva.
• Il Limite: L'aggiunta di un nodo di output per la varianza aumenta il numero di parametri apprendibili e la complessità architetturale. Questo è dannoso per le BNN leggere destinate alla distribuzione su sistemi con vincoli hardware o dove il tempo di calcolo è limitato.
• Il Divario: I metodi esistenti spesso costringono a un compromesso tra le dimensioni del modello e la capacità di catturare strutture di rumore complesse dipendenti dall'input.

2. Metodologia

Gli autori propongono un nuovo framework di inferenza che elimina la necessità di un nodo di output separato per l'incertezza aleatoria, incorporando direttamente sia le incertezze aleatorie che quelle epistemiche nelle varianze dei parametri apprendibili della BNN.

A. Reinterpretazione delle Distribuzioni dei Parametri

Invece di trattare la varianza dei parametri dei neuroni bayesiani come rappresentante esclusivamente l'incertezza epistemica, gli autori reinterpretano queste distribuzioni per codificare l'incertezza predittiva totale (aleatoria + epistemica).

• Approccio Standard: $Var(\hat{y}) = \sigma^2_{epistemic}$. La varianza aleatoria ($\sigma^2_{aleatoric}$) viene prevista come un output separato.
• Approccio Proposto: La varianza totale $\sigma^2_k$ per un dato input $k$ è definita come la somma delle varianze aleatoria ed epistemica:
  $$\sigma^2_k = \sigma^2_{a,k} + \sigma^2_{e,k}$$
  Durante l'addestramento, la rete minimizza il Limite Inferiore dell'Evidenza (ELBO) dove la perdita di ricostruzione assume che la distribuzione di output sia Gaussiana con varianza $\sigma^2_k$. La varianza $\sigma^2_k$ è calcolata come la varianza dell'output della BNN derivata dalle distribuzioni dei parametri, "assorbendo" efficacemente il rumore aleatorio nelle varianze dei parametri.

B. Inferenza Variazionale Senza Campionamento tramite Propagazione dei Momenti

Per mantenere l'efficienza computazionale ed evitare il costo del campionamento Monte Carlo, gli autori impiegano uno schema di propagazione dei momenti.

• Meccanismo: La media e la varianza dell'input di ogni livello vengono propagate analiticamente attraverso la rete per calcolare la media e la varianza dell'output.
• Gestione dei Livelli:
  • Fully Connected & Convolutional: Propagazione analitica esatta di medie e varianze assumendo l'indipendenza tra i neuroni.
  • Leaky-ReLU: Propagazione approssimata assumendo che l'input alla funzione di attivazione sia distribuito normalmente. Gli autori derivano espressioni in forma chiusa per la media e la varianza dell'output utilizzando la legge della speranza totale e della varianza, coinvolgendo distribuzioni normali troncate.
• Vantaggio: Questo permette un'inferenza deterministica senza la necessità di campionamenti ripetuti, rendendo il framework adatto ad applicazioni leggere e in tempo reale.

3. Contributi Chiave

1. Incorporazione della Varianza dei Parametri: Un cambiamento concettuale in cui le varianze dei pesi e degli bias della BNN vengono addestrate per rappresentare la somma delle incertezze aleatorie ed epistemiche, eliminando la necessità di una testa di output dedicata alla varianza.
2. Architettura Leggera: Il framework permette alle BNN di catturare incertezze eteroschedastiche con meno parametri apprendibili rispetto agli approcci standard di "varianza appresa".
3. Inferenza Senza Campionamento: Integrazione di una tecnica semplificata di propagazione dei momenti che fornisce stime analitiche dell'incertezza, evitando il sovraccarico computazionale del dropout Monte Carlo o del campionamento.
4. Framework Unificato: Una pipeline completa che combina la reinterpretazione delle fonti di incertezza con una propagazione efficiente dei momenti per l'addestramento e l'inferenza.

4. Risultati Sperimentali

Gli autori hanno valutato il loro approccio "Varianza Incorporata" rispetto all'approccio standard "Varianza Appresa" (dove la varianza è un output separato) utilizzando un compito di regressione polinomiale con rumore Gaussiano eteroschedastico.

• Prestazioni in Regimi Leggeri:
  • Con un piccolo numero di neuroni nascosti (ad esempio, 4), il modello a Varianza Incorporata ha previsto accuratamente sia la tendenza che l'incertezza per i dati in-distribuzione ed ha espresso correttamente un'alta incertezza per i dati out-of-distribution.
  • Il modello a Varianza Appresa non è riuscito a recuperare l'incertezza reale per i dati in-distribuzione e ha sottostimato significativamente l'incertezza per i dati out-of-distribution con la stessa architettura piccola.
• Efficienza dei Parametri:
  • Per un numero fisso di neuroni, il modello a "Varianza Appresa" ha richiesto circa il 34% di parametri apprendibili in più rispetto al modello a "Varianza Incorporata" per ottenere prestazioni comparabili.
  • L'approccio a "Varianza Incorporata" ha ottenuto una perdita di ricostruzione inferiore (log-verosimiglianza negativa) nel regime a basso numero di parametri.
• Generalizzazione:
  • I modelli a "Varianza Incorporata" hanno dimostrato una migliore generalizzazione ai dati out-of-distribution, mostrando un aumento più graduale della perdita di ricostruzione man mano che i dati si allontanavano dalla distribuzione di addestramento rispetto ai modelli a "Varianza Appresa".
• Scalabilità: All'aumentare delle dimensioni della rete (ad esempio, 1024 neuroni), il divario di prestazioni si è ristretto e entrambi i metodi hanno performance simili, suggerendo che il metodo proposto è più critico per modelli leggeri con vincoli di risorse.

5. Significato

Questo lavoro affronta un collo di bottiglia critico nella distribuzione delle Reti Neurali Bayesiane per il computing edge e i sistemi safety-critical.

• Efficienza: Rimuovendo la necessità di nodi di output aggiuntivi per prevedere il rumore, il metodo riduce significativamente le dimensioni del modello e il costo computazionale, rendendo le BNN fattibili per hardware con vincoli rigorosi.
• Robustezza: La capacità di catturare incertezze eteroschedastiche senza aumentare la complessità del modello garantisce che i modelli leggeri non sacrifichino l'affidabilità o le valutazioni di sicurezza.
• Praticità: La natura senza campionamento dell'inferenza permette tempi di previsione più rapidi, essenziali per applicazioni in tempo reale dove la latenza è un problema.

In sintesi, il documento dimostra che l'incertezza totale può essere efficacemente "incorporata" nelle distribuzioni dei parametri di una BNN, offrendo un'alternativa superiore alle tradizionali teste multi-output per la quantificazione dell'incertezza eteroschedastica leggera.