Stable Global Weighting of Flow Mixtures using Simplex Exponential Moving Average

Sintesi Tecnica: Ponderazione Globale Stabile di Miscele di Flusso tramite Simplex Exponential Moving Average

Problematica
I flussi normalizzanti hanno ampliato significativamente la flessibilità dell'inferenza variazionale (VI) trasformando distribuzioni di base semplici in densità di probabilità espressive. Tuttavia, le singole architetture di flusso (ad esempio, RealNVP, MAF, RBIG) spesso non riescono a generalizzare su geometrie posteriori eterogenee. Diversi flussi possiedono distinti bias induttivi, dinamiche di ottimizzazione e caratteristiche numeriche, portando a una situazione in cui nessuna singola architettura performa in modo costante attraverso diverse distribuzioni target. Sebbene gli approcci basati su miscele offrano una strategia naturale per combinare esperti complementari, i framework esistenti soffrono di limitazioni significative:

1. Ottimizzazione Congiunta: L'apprendimento simultaneo dei pesi della miscela e dei parametri degli esperti porta spesso a un'instabilità nell'addestramento o al "dominio dell'esperto" a causa dell'accoppiamento di dinamiche di ottimizzazione eterogenee.
2. Ponderazione Globale Fissa: La stima di coefficienti fissi durante l'addestramento e il loro mantenimento costante durante l'inferenza manca della flessibilità necessaria per sfruttare il comportamento complementare degli esperti, risultando spesso in un'allocazione inefficiente della capacità.

Di conseguenza, è necessario un framework di ottimizzazione stabile e principato che possa combinare esperti eterogenei senza richiedere la conoscenza preventiva dell'architettura più appropriata o incorrendo nell'instabilità dell'addestramento congiunto.

Metodologia: AMF-VI-sEMA
Il documento propone AMF-VI-sEMA (Adaptive Mixture of Flows - Variational Inference with Simplex Exponential Moving Average), un framework a due stadi progettato per disaccoppiare l'apprendimento dell'esperto dall'ottimizzazione della miscela.

• Stadio 1: Specializzazione Indipendente dell'Esperto
  Un insieme eterogeneo di esperti (specificamente RealNVP, MAF e RBIG) viene addestrato indipendentemente per specializzarsi in distinti regimi strutturali. Ogni esperto $q_k(z|\phi_k)$ è ottimizzato per massimizzare il proprio log-likelihood sui dati di addestramento. Una volta addestrato, i parametri dell'esperto $\phi_k$ vengono congelati. Questo stadio assicura che ogni componente converga sotto le proprie dinamiche di ottimizzazione, evitando le patologie di coordinamento dell'ottimizzazione congiunta.

• Stadio 2: Ponderazione Globale Simplex EMA (sEMA)
  Con i parametri degli esperti congelati, il framework apprende un vettore di peso della miscela globale $\pi$ sul simplesso di probabilità. Ciò avviene attraverso un meccanismo guidato dalla verosimiglianza e controllato dalla temperatura:

1. Calcolo della Responsabilità: Per un batch di dati di validazione tenuti da parte (held-out), le responsabilità $r_k(z)$ sono calcolate utilizzando una softmax delle medie dei log-likelihood degli esperti, scalate da un parametro di temperatura $\tau$.
2. Smoothing e Mediazione: Le responsabilità sono mediate sul batch di validazione e opzionalmente smussate verso una distribuzione uniforme per attenuare i transitori.
3. Aggiornamento Simplex EMA: I pesi globali $\pi$ sono aggiornati tramite una Media Mobile Esponenziale (EMA) direttamente sul simplesso: $\pi \leftarrow \alpha \pi + (1-\alpha)\bar{r}$. Viene applicato un piccolo valore minimo $\epsilon$ per evitare lo azzeramento, seguito dalla rinormalizzazione.

Questo design crea un meccanismo di gating trattabile e agnostico rispetto ai dati che rialloca la capacità in modo adattivo senza gating per singolo campione o backpropagation attraverso i pesi. Il framework monitora il numero effettivo di esperti ($N_{eff} = \exp(H(\pi))$) per diagnosticare il collasso dei componenti.

Contributi Chiave

1. Framework di Ottimizzazione Disaccoppiato: Il documento introduce una strategia a due stadi che separa l'apprendimento dell'esperto dall'adattamento della miscela. Congelando gli esperti prima di aggiornare i pesi, evita l'instabilità dell'ottimizzazione ponderata per responsabilità congiunta pur rimanendo agnostico rispetto all'architettura.
2. Meccanismo di Ponderazione Adattiva Pratico: Gli autori propongono un aggiornamento Simplex EMA basato sulla responsabilità e computazionalmente efficiente. Questo metodo stima i pesi globali della miscela senza richiedere l'ottimizzazione congiunta o il gating per singolo campione, fornendo una cooperazione stabile tra esperti con un overhead computazionale minimo.
3. Valutazione Empirica Completa: Il framework è valutato su dieci problemi di benchmark, inclusi sei gruppi sintetici 2D canonici (Banana, X-Shaped, Bimodal, Multimodal, Two-moons, Rings) e quattro target bayesiani reali/a bassa dimensionalità (BLR, BPR, Weibull, Real-GMM2). La valutazione si estende oltre il log-likelihood standard per includere la distanza Wasserstein-2, il Maximum Mean Discrepimento (MMD), l'analisi della calibrazione e la robustezza attraverso i seed casuali.
4. Analisi delle Dinamiche di Ottimizzazione: Il documento fornisce un'analisi dell'utilizzo degli esperti, dell'evoluzione del peso della miscela e della sensibilità agli iperparametri, offrendo spunti su come le miscele di flussi eterogenei abbiano successo e come il meccanismo sEMA stabilizzi l'addestramento.

Risultati
Le valutazioni empiriche dimostrano che AMF-VI-sEMA ottiene miglioramenti costanti rispetto al suo predecessore (AMF-VI) ed evita i fallimenti catastrofici del trasporto osservati nelle singole baseline (come ResFlow e EM-Mix).

• Performance: Il framework raggiunge il miglior o il quasi-miglior Negative Log-Likelihood (NLL) su tutti i dataset sintetici e reali. Mantiene un comportamento di trasporto robusto (Wasserstein-2) e metriche di discrepanza (MMD), evitando i fallimenti geometrici visti nelle singole architetture (ad esempio, NICE e ResFlow mostrano alti errori di trasporto su geometrie curve).
• Stabilità: Il metodo mantiene traiettorie di peso stabili con un numero effettivo di esperti $N_{eff} > 1.4$ su tutti i dataset, indicando che la miscela non collassa su un singolo esperto.
• Robustezza: Le valutazioni cross-seed mostrano una bassa varianza in NLL, divergenza KL e distanza Wasserstein, confermando che il framework è riproducibile e robusto rispetto all'inizializzazione casuale.
• Adattività: I pesi globali appresi si adattano alla geometria target. Ad esempio, su posteri di tipo regressione (BLR, BPR), il meccanismo sopprime RealNVP a favore di MAF e RBIG, mentre su strutture ad anello (Rings), concentra fortemente la massa su RealNVP.

Significatività e Rivendicazioni
Il documento sostiene che AMF-VI-sEMA fornisca una strategia di ottimizzazione robusta e affidabile per l'approssimazione della posteriore attraverso paesaggi eterogenei senza richiedere la conoscenza preventiva di quale architettura di flusso sia la più appropriata. La significatività del lavoro risiede in:

• Stabilità: Il meccanismo sEMA introduce un effetto di regolarizzazione temporale che stabilizza le dinamiche di gating e mitiga il volatile switching tra esperti.
• Semplicità ed Efficienza: L'approccio è agnostico rispetto all'architettura, non richiede modifiche ai modelli di flusso sottostanti e introduce un overhead computazionale trascurabile rispetto all'addestramento congiunto.
• Affidabilità: Evitando la fragilità dell'addestramento congiunto ponderato per responsabilità, il framework offre una soluzione pratica per modellare geometrie posteriori complesse, non convesse e multimodali.

Gli autori osservano che, sebbene le architetture a singolo flusso specializzate possano rimanere preferibili per target unimodali semplici, AMF-VI-sEMA supera costantemente le altre in termini di robustezza attraverso diverse geometrie. Il lavoro futuro suggerisce di esplorare strategie di gating condizionali o amortizzate e estensioni a contesti ad alta dimensionalità.