Amortized Factor Inference Networks for Posterior Inference

Riepilogo Tecnico: Reti di Inferenza a Fattori Ammortizzati (AFIN)

1. Enunciato del Problema

L'inferenza ammortizzata mira a imparare una singola rete neurale che mappa le osservazioni sulle distribuzioni a posteriori, offrendo un'inferenza rapida al momento del test dopo un costo di addestramento una tantum. Tuttavia, i metodi ammortizzati esistenti sono tipicamente vincolati a modelli probabilistici fissi. Se il prior, la verosimiglianza, la dimensionalità latente ($d$) o il numero di osservazioni ($N$) cambiano, la rete di inferenza richiede solitamente un riaddestramento o un costoso affinamento al momento del test. Sebbene alcuni lavori recenti generalizzino a $N$ variabili o a cambiamenti specifici dei parametri, spesso mancano della capacità di gestire combinazioni arbitrarie di famiglie di prior/verosimiglianza variabili, dimensionalità e conteggi di osservazioni simultaneamente senza sacrificare l'accuratezza o richiedere adattamenti specifici del modello.

La domanda centrale posta da questo lavoro è: Una singola rete di inferenza può generalizzare su un'intera famiglia di problemi distinti di inferenza bayesiana, gestendo prior, verosimiglianze, dimensioni e conteggi di osservazioni variabili, mantenendo al contempo l'accuratezza dei metodi non ammortizzati come NUTS?

2. Metodologia: Reti di Inferenza a Fattori Ammortizzati (AFIN)

Gli autori introducono le Reti di Inferenza a Fattori Ammortizzati (AFIN), un'architettura di codifica-fusione-decodifica progettata per mappare una rappresentazione tipizzata di una specifica del modello e delle sue osservazioni direttamente sui parametri di una variazione a posteriori.

A. Rappresentazione a Fattori Tipizzati
Invece di trattare il modello come un blocco monolitico, AFIN rappresenta la densità congiunta $p(z, y_{1:N})$ come una collezione di $N+1$ fattori: un fattore prior $p_0(z)$ e $N$ fattori di verosimiglianza $p_n(y_n|z)$. Ogni fattore è rappresentato da:

• Tipo ($t_n$): Un'etichetta che indica la famiglia di distribuzione (es. "verosimiglianza gaussiana", "verosimiglianza student-t").
• Parametri ($\theta_n$): Parametri specifici del fattore (es. scala, gradi di libertà) necessari per istanziare il fattore.
• Osservazioni ($y_n$): I dati associati ai fattori di verosimiglianza.

Questa rappresentazione permette alla rete di gestire modelli eterogenei in cui diverse osservazioni possono seguire diverse famiglie di verosimiglianza.

B. Architettura Indipendente dalla Dimensione
Per supportare dimensioni latenti variabili ($d$) e numeri di fattori ($N$), AFIN impiega moduli indipendenti da $d$ ed equivarianti rispetto alle permutazioni delle coordinate latenti:

1. BoxMLP: Un Perceptron Multistrato con pesi condivisi in cui i pesi appresi non dipendono da $d$. Elabora caratteristiche coordinate per coordinate e riepiloghi aggregati (es. media sulle dimensioni) per garantire che l'output sia equivariante rispetto alle permutazioni del vettore latente $z \in \mathbb{R}^d$.
2. BoxTransformer: Un analogo indipendente dalla dimensione del Transformer standard. Opera su una sequenza di embedding dei fattori. Crucialmente, le mappe di proiezione (Query, Key, Value, Output, FFN) sono implementate come BoxMLP applicati a token indicizzati per coordinata. Ciò permette alla rete di modellare le interazioni tra fattori senza appiattire la dimensione latente, preservando la struttura $O(d)$ o $O(d^2)$ dei parametri.

C. La Pipeline di Codifica-Fusione-Decodifica

1. Codifica: Ogni fattore è elaborato indipendentemente. Un adattatore specifico per il tipo mappa i parametri del fattore e le osservazioni in uno spazio di embedding comune, seguito da un codificatore BoxMLP condiviso. L'embedding include sia caratteristiche dei nodi (coordinate per coordinate) sia caratteristiche di coppia (interazioni tra coppie di coordinate, utili per matrici di covarianza/precisione).
2. Fusione: La sequenza di embedding dei fattori è elaborata da $M$ blocchi BoxTransformer. Questi blocchi utilizzano l'attenzione lungo l'asse dei fattori (trattando i fattori come dimensione della sequenza) per aggregare le informazioni. Il meccanismo di attenzione combina punteggi dei nodi e delle coppie per aggiornare le rappresentazioni.
3. Decodifica: Le rappresentazioni fuse sono aggregate tramite somma lungo l'asse dei fattori per produrre un singolo embedding a livello di compito. Un decodificatore BoxMLP finale mappa questo embedding sui parametri variazionali $\phi$ (es. media e precisione per una Gaussiana, o parametri di flusso per un flusso normalizzante).

D. Obiettivo di Addestramento
AFIN è addestrata tramite inferenza basata sulla simulazione. La rete è addestrata per massimizzare la densità logaritmica assegnata a campioni latenti $z$ campionati dal modello congiunto $p(z, y)$, condizionati alle osservazioni generate $y$. Questo equivale a minimizzare la divergenza KL forward attesa ($KL(p(z|y) || q_\phi(z))$) tra la vera posteriori e l'approssimazione variazionale. Si nota che l'obiettivo della KL forward è benefico per l'uso dell'approssimazione come proposta per il campionamento per importanza.

3. Contributi Chiave

• Generalizzazione: Una singola AFIN addestrata può eseguire l'inferenza a posteriori su famiglie di prior variabili, famiglie di verosimiglianza, dimensioni latenti ($d$) e numeri di osservazioni ($N$) senza riaddestramento.
• Indipendenza dalla Dimensione: L'introduzione dei moduli BoxMLP e BoxTransformer permette alla rete di operare su spazi latenti di dimensione arbitraria $d$ mantenendo un conteggio fisso di parametri, una significativa deviazione dalle reti ammortizzate standard che richiedono dimensioni di input/output fisse.
• Astrazione a Livello di Fattore: Operando su fattori tipizzati anziché su grafi di modelli fissi, AFIN supporta verosimiglianze eterogenee (es. mescolando osservazioni Gaussiane e Bernoulliane) all'interno di una singola esecuzione di inferenza.
• Efficienza: Il metodo raggiunge un'alta accuratezza con un calcolo al momento del test significativamente inferiore rispetto ai metodi iterativi.

4. Risultati Sperimentali

Gli autori hanno valutato un singolo checkpoint AFIN da 5 milioni di parametri su compiti sintetici e del mondo reale senza affinamento.

• Prestazioni Sintetiche: Su 48 configurazioni (16 combinazioni prior-verosimiglianza $\times$ 3 livelli di difficoltà), AFIN ha raggiunto un'accuratezza a posteriori comparabile a NUTS e a varie baseline di Inferenza Variazionale (VI) (Gaussiana a rango completo, IAF, MAF, NSF).
  • Velocità: L'inferenza singola di AFIN è da 2 a 4 ordini di grandezza più veloce di NUTS e dei metodi VI basati su flusso.
  • Raffinamento: Quando utilizzata come proposta per il Campionamento per Importanza Normalizzato (SNIS), AFIN+SNIS ha eguagliato o superato l'accuratezza di NUTS con un tempo di esecuzione sostanzialmente inferiore.
• Estrapolazione: AFIN ha dimostrato una generalizzazione zero-shot a contesti fuori distribuzione (OOD), incluse dimensioni latenti ($d=32$) e conteggi di osservazioni ($N=512$) maggiori di quelli visti durante l'addestramento ($d \le 16, N \le 256$).
• Trasferimento nel Mondo Reale: Su 12 dataset UCI (inclusi contesti di verosimiglianza eterogenea), AFIN+SNIS ha raggiunto un'accuratezza competitiva con i riferimenti NUTS a lungo termine quasi immediatamente, mentre le baseline iterative richiedevano un tempo di ottimizzazione significativo per convergere.
• Varianti del Decodificatore: Sono stati testati sia decodificatori Gaussiani che a Flusso. Il decodificatore a Flusso ha offerto lievi miglioramenti di accuratezza rispetto al decodificatore Gaussiano mantenendo il vantaggio di velocità dell'ammortizzazione.

5. Significato e Affermazioni

Il documento afferma che AFIN colma con successo il divario tra la velocità dell'inferenza ammortizzata e la flessibilità dei metodi non ammortizzati. Disaccoppiando la rete di inferenza dalle specifiche del modello fisse attraverso un'interfaccia a fattori tipizzati e moduli indipendenti dalla dimensione, AFIN abilita un'inferenza ammortizzata riutilizzabile.

Gli autori posizionano AFIN come un passo verso un "modello fondazionale" per l'inferenza bayesiana, capace di gestire programmi probabilistici diversificati. Riconoscono le limitazioni, come la necessità di un insieme predefinito di tipi di fattori e la scalabilità $O(d^2)$ degli embedding delle componenti di coppia, ma sottolineano che l'approccio permette a una singola rete di risolvere un'ampia famiglia di problemi di inferenza che in precedenza richiedevano modelli distinti e riaddestrati. Il lavoro suggerisce che l'ammortizzazione può eguagliare la qualità dei metodi di campionamento all'avanguardia (NUTS) offrendo al contempo enormi risparmi computazionali al momento del test.