eXact-Prior Variational Autoencoder (X-VAE): Learning Data-Adaptive Gaussian Mixture Priors for Latent Distributions

Problema

I classici Autoencoder Variazionali (VAE) si affidano a un prior gaussiano isotropo fisso, $p(z) = \mathcal{N}(0, I)$, nello spazio latente. Sebbene matematicamente conveniente per il calcolo della divergenza di Kullback–Leibler (KL) nella Evidence Lower Bound (ELBO), questa assunzione spesso non riesce a catturare la vera e complessa struttura delle rappresentazioni latenti nei dataset del mondo reale. Questa discrepanza tra il prior assunto e l'aggregato della posterior porta a diversi problemi:

• Discrepanza Prior-Posterior (Prior-Posterior Mismatch): La distribuzione latente appresa devia significativamente dalla normale standard, limitando l'accuratezza della ricostruzione e la qualità del campionamento.
• Collasso della Posterior (Posterior Collapse): Il modello può ignorare le variabili latenti, portando a scarse capacità generative.
• Espressività Limitata: Il prior rigido vincola lo spazio latente, impedendo al modello di rappresentare efficacemente dati multimodali o strutturalmente vincolati.

Le soluzioni esistenti, come VampPrior, i prior basati su flow o i modelli gerarchici, migliorano la flessibilità ma spesso introducono parametri addizionali addestrabili, una maggiore complessità computazionale e procedure di ottimizzazione più difficili.

Metodologia: Il Framework X-VAE

Gli autori propongono l'eXact-Prior Variational Autoencoder (X-VAE), che scinde la stima del prior dall'apprendimento variazionale. Invece di apprendere una complessa rete per il prior, X-VAE deriva un prior adattivo ai dati dalle rappresentazioni latenti di un Autoencoder (AE) deterministico pre-addestrato.

La metodologia procede in due fasi:

1. Stima del Prior (Fase di Pre-training)
Un AE deterministico viene addestrato per minimizzare l'errore di ricostruzione. I codici latenti $\{z_i^{AE}\}_{i=1}^N$ prodotti da questo AE vengono utilizzati per stimare le statistiche del prior. Gli autori propongono due configurazioni per il prior $p_{AE}(z)$:

• Prior Gaussiano Singolo: La media empirica ($\mu_{AE}$) e la covarianza ($\Sigma_{AE}$) dei codici latenti definiscono un prior gaussiano singolo: $p_{AE}(z) = \mathcal{N}(\mu_{AE}, \Sigma_{AE})$.
• Prior Gaussian Mixture Model (GMM): Un mix di $K$ componenti gaussiane diagonali viene adattato ai codici latenti dell'AE (ad esempio, tramite Expectation-Maximization). Questo mix è congelato e funge da prior:
  $$p_{AE}(z) = \sum_{k=1}^K \pi_k^p \mathcal{N}(z; \mu_k^p, \text{diag}((\sigma_k^p)^2))$$
  Qui, i pesi del mix $\pi^p$, le medie $\mu^p$ e le deviazioni standard $\sigma^p$ sono parametri fissi derivati dall'AE, non addestrabili durante l'addestramento del VAE.

2. Apprendimento Variazionale con Routed-Transport
Il VAE viene addestrato con questo prior fisso. Per allineare la posterior approssimata $q_\phi(z|x)$ al mix del prior, gli autori introducono un meccanismo di routed-transport:

• Routing Categoriale: Per ogni coordinata latente $k$, viene selezionata una componente del mix $c_k$ dai pesi del prior $\pi^p$ (o dai pesi della posterior $\pi^q$).
• Operatore di Trasporto: Un campione di rumore $\epsilon_k$ viene estratto dalla distribuzione della componente del prior selezionata. Questo rumore viene poi "trasportato" alla distribuzione della posterior dell'encoder utilizzando una mappa lineare:
  $$T_k(\epsilon_k) = \mu_k^q + \sigma_k^q \frac{1}{\sigma_k^p} (\epsilon_k - \mu_k^p)$$
  Ciò assicura che se $\epsilon_k \sim \mathcal{N}(\mu_k^p, (\sigma_k^p)^2)$, allora $T_k(\epsilon_k) \sim \mathcal{N}(\mu_k^q, (\sigma_k^q)^2)$, preservando il trucco della riparametrizzazione per la backpropagation.
• Funzione di Perdita: L'ELBO è ottimizzata utilizzando una divergenza KL in forma chiusa. Per il caso GMM, il termine KL è limitato dalla somma di un termine categoriale (che allinea le assegnazioni delle componenti) e un termine gaussiano per componente:
  $$D_{KL}(q_\phi(z|x) \parallel p(z)) \leq D_{KL}(\pi^q \parallel \pi^p) + \sum_{k=1}^K \pi_k^q D_{KL}(\mathcal{N}_k^q \parallel \mathcal{N}_k^p)$$

3. Scaling della Varianza per la Generazione
Durante la fase di generazione, gli autori introducono un fattore di scala globale $\alpha$ applicato alla deviazione standard del prior:
$$z_{gen} \sim \mathcal{N}(\mu_{AE}, \text{diag}((\alpha \sigma_{AE})^2))$$
Questo permette un controllo esplicito sul compromesso tra fedeltà del campione ( $\alpha$ piccolo, campioni vicini al manifold empirico) e diversità ($\alpha$ grande, esplorazione di regioni meno frequenti dello spazio latente) senza necessità di ri-addestramento.

Contributi Chiave

1. Prior Adattivo ai Dati: Sostituzione del prior gaussiano isotropo con un prior stimato dalle statistiche latenti di un AE pre-addestrato, ancorando il VAE alla distribuzione empirica dei dati.
2. Trattabilità Analitica: Derivazione di una divergenza KL in forma chiusa per il prior proposto, evitando la necessità di complesse reti per il prior apprese o trasformazioni basate su flow.
3. Meccanismo di Routed-Transport: Una nuova strategia di campionamento che mappa le posterior dell'encoder alle componenti del mix, preservando il gradiente della riparametrizzazione e ancorando lo spazio latente alla struttura di clustering dell'AE.
4. Diversità Controllabile: Un semplice meccanismo di scaling della varianza ($\alpha$) che consente di regolare il compromesso diversità-fedeltà durante l'inferenza.
5. Efficienza dei Parametri: Il prior aggiunge zero parametri addestrabili al modello generativo, poiché viene memorizzato in buffer fissi.

Risultati Sperimentali

Gli autori hanno valutato X-VAE su MNIST, CelebA e un dataset sintetico raggruppato (clustered).

• Performance: Su CelebA, la migliore configurazione X-VAE ha ottenuto l'obiettivo di test più basso (164.65) e il miglior Fréchet Inception Distance (FID) di 80.40, superando i VAE standard, VampPrior e i VAE gerarchici. Su MNIST, le configurazioni X-VAE hanno migliorato l'FID rispetto al VAE standard (47.21 vs 49.67) e hanno eguagliato o superato leggermente VampPrior.
• Dati Sintetici: Su un dataset sintetico a 3 cluster, X-VAE è riuscito a preservare la struttura multimodale dei dati, mentre i VAE standard hanno fatto collassare i modi in una singola gaussiana.
• Ricostruzione: X-VAE ha mantenuto una qualità di ricostruzione comparabile ai baseline producendo al contempo campioni più realistici e diversificati.
• Efficienza: Il metodo non richiede parametri addestrabili aggiuntivi per il prior, risultando computazionalmente efficiente rispetto agli approcci con prior appreso come GM-VAE o i VAE gerarchici.

Significato e Rivendicazioni

Il paper sostiene che X-VAE offra un'alternativa semplice ed efficace ai complessi framework con prior appreso. Ancorando il prior alla struttura latente empirica di un AE deterministico, X-VAE:

• Riduce la discrepanza prior-posterior inerente ai VAE standard.
• Preserva la trattabilità analitica e l'efficienza computazionale dei VAE standard.
• Fornisce un framework pratico per la modellazione generativa vincolata, particolarmente nell'ingegneria e nel design industriale dove le soluzioni generate devono soddisfare rigorosi vincoli strutturali mantenendo la diversità.
• Dimostra che la modellazione generativa di alta qualità non richiede necessariamente reti per il prior complesse e ricche di parametri; piuttosto, sfruttare la struttura statistica delle rappresentazioni pre-addestrate può portare a un migliore allineamento e controllabilità.

Gli autori riconoscono i limiti, come la dipendenza dalla qualità dell'AE pre-addestrato e la potenziale discrepanza residua tra gli spazi latenti di AE e VAE. Notano inoltre che l'assunzione di un GMM diagonale può non catturare le correlazioni entro il modo e che il numero di componenti del mix è legato alla dimensione latente, il che può sovra-specificare il conteggio naturale delle categorie in dimensioni elevate. Il lavoro futuro suggerisce di esplorare prior non gaussiani, estensioni gerarchiche e strategie di addestramento congiunto.