Likelihood Matching for Diffusion Models

Sintesi Tecnica: Likelihood Matching per i Modelli di Diffusione

Problema
Gli attuali modelli di diffusione mainstream, in particolare i Modelli Generativi Basati sullo Score (SGM), si affidano allo Score Matching (SM) come loro principale obiettivo di addestramento. Sebbene empiricamente efficace, lo SM minimizza la discrepanza tra una funzione di score parametrizzata e la vera funzione di score a vari livelli di rumore. Teoricamente, questo approccio minimizza solo un limite superiore della log-verosimiglianza negativa della distribuzione dei dati originali, piuttosto che la verosimiglianza stessa. Questo collegamento indiretto può portare a una perdita di efficienza statistica rispetto alla Stima di Massima Verosimiglianza (MLE), nota per produrre stimatori coerenti ed efficienti. Inoltre, lo SM standard si concentra esclusivamente sulle informazioni del primo momento (lo score/gradiente), potenzialmente trascurando le informazioni strutturali di ordine superiore contenute nelle densità di transizione.

Metodologia
Gli autori propongono il Likelihood Matching (LM), un framework che mira direttamente alla verosimiglianza della distribuzione dei dati sfruttando l'equivalenza tra la verosimiglianza dei dati target e la verosimiglianza del percorso del processo di diffusione inversa.

1. Fondamento Teorico:
   Il metodo si basa sulla proprietà per cui la verosimiglianza congiunta del percorso di diffusione inversa è equivalente alla verosimiglianza della distribuzione dei dati originali (fino a costanti legate al processo forward). Gli autori stabiliscono che massimizzare l'esatta verosimiglianza del percorso del processo inverso è equivalente a massimizzare la log-verosimiglianza dei dati.

2. Quasi-Maximum Likelihood Estimation (QMLE):
   Poiché le esatte densità di transizione inversa sono intrattabili, gli autori le approssimano utilizzando un approccio di Quasi-Maximum Likelihood. Modellano ogni densità di transizione inversa come una distribuzione Gaussiana in cui la media condizionale e la covarianza sono accoppiate al vero processo inverso.

   • Insight Chiave: La media condizionale e la covarianza del processo inverso dipendono non solo dalla funzione di score ($\nabla \log q_t$), ma anche dalla funzione Hessiana ($\nabla^2 \log q_t$).
   • Parametrizzazione: Gli autori parametrizzano sia lo score che l'Hessiana utilizzando reti neurali (specificamente U-Net). Per gestire la complessità computazionale, l'Hessiana è spesso approssimata utilizzando una struttura diagonale più low-rank.

3. L'Obiettivo di Likelihood Matching:
   La funzione obiettivo proposta, Likelihood Matching (LM), è derivata minimizzando la quasi-log-verosimiglianza negativa delle transizioni inverse.

   • L'obiettivo unifica lo score matching e il covariance matching.
   • Incorpora un meccanismo di pesatura della verosimiglianza implicita tramite la matrice di covarianza accoppiata $(I + \sigma^2 H_t)^{-1/2}$, che regola automaticamente il peso del termine di score matching in base alla curvatura locale.
   • L'addestramento utilizza un campionatore stocastico che seleziona casualmente i passi temporali lungo la traiettoria di diffusione per calcolare la perdita, garantendo un uso efficiente delle informazioni evolutive dello SDE forward.

4. Campionatore Stocastico:
   Viene introdotta una nuova procedura di campionamento che utilizza sia l'informazione dello score che quella dell'Hessiana appresi. Il passo inverso è definito come:
   $$\tilde{Y}_{t-1} = \hat{\mu}_{t-1|t}(\tilde{Y}_t) + \hat{\Sigma}_{t-1|t}^{1/2}(\tilde{Y}_t) Z_t$$
   dove la media $\hat{\mu}$ e la covarianza $\hat{\Sigma}$ sono funzioni, rispettivamente, dello score e dell'Hessiana stimati.

Contributi Chiave

• Nuovo Obiettivo di Addestramento: Il paper introduce un obiettivo di addestramento basato su quasi-verosimiglianza (LM) che approssima la log-verosimiglianza del percorso inverso, combinando efficacemente lo score matching e il covariance matching con una pesatura della verosimiglianza implicita.
• Campionamento Potenziato dall'Hessiana: Viene derivato un campionatore stocastico che sfrutta sia l'informazione dello score che dell'Hessiana, offrendo un'approssimazione della densità di transizione inversa più accurata rispetto ai campionatori standard di tipo DDPM.
• Garanzie Teoriche:
  • Convergenza: Sono fornite garanzie di convergenza non asintotica nella distanza di variazione totale. Il limite di errore scala come $O(d^3 \log^{4.5} T / T)$ per l'errore del passo inverso, con termini aggiuntivi per gli errori di stima dello score e dell'Hessiana.
  • Coerenza: Il paper dimostra la coerenza della proposta QMLE in un setting parametrico oracle, mostrando che lo stimatore converge ai veri parametri dei dati all'aumentare della dimensione del campione e dei passi di diffusione.
• Implementazione Efficiente: L'uso di approssimazioni low-rank dell'Hessiana e della formula di Sherman-Morrison-Woodbury consente un calcolo efficiente della verosimiglianza e dei passi di campionamento.

Risultati Sperimentali
Gli autori validano il metodo sia su dataset sintetici che reali:

• Dati Sintetici: Su miscele di distribuzioni Gaussiane e t, LM ha superato costantemente lo standard Score Matching (SM), in particolare nel catturare i modi ben separati. Su una distribuzione "Swiss-roll" altamente avvolgente (che rappresenta varietà curve), LM ha dimostrato prestazioni superiori, validando il beneficio della covarianza informata dall'Hessiana.
• Stima dei Parametri: In un contesto parametrico di miscela Gaussiana, LM ha ottenuto un Errore Assoluto Medio (MAE) e un errore standard inferiori rispetto a SM nella stima dei parametri, supportando le affermazioni di coerenza teorica.
• Dataset di Immagini (MNIST, CIFIA-10, CelebA):
  • LM ha ottenuto punteggi di Fréchet Inception Distance (FID) più bassi e migliori punteggi di Negative Log-Likelihood (NLL) rispetto ai baseline SM.
  • Le prestazioni sono migliorate con il rango dell'approssimazione dell'Hessiana ($r$), raggiungendo il picco intorno a $r=20-30$.
  • Trade-off Velocità-Qualità: LM ha ottenuto un FID inferiore per tempi di esecuzione comparabili e ha richiesto meno passi di campionamento inverso per raggiungere un'alta qualità visiva. Ad esempio, su MNIST, LM ha prodotto cifre riconoscibili dopo soli 20 iterazioni inverse, mentre i campioni SM rimanevano sfocati.

Significatività e Affermazioni
Il paper sostiene che il Likelihood Matching offra un'alternativa statisticamente più rigorosa allo Score Matching, mirando direttamente alla verosimiglianza dei dati attraverso un framework di quasi-verosimiglianza. Incorporando l'informazione del secondo ordine (l'Hessiana), il metodo cattura densità di transizione più ricche e migliora l'efficienza del campionamento, specialmente nei regimi a pochi passi e su varietà di dati complesse e curve. Gli autori sottolineano che, sebbene l'LM comporti costi computazionali per passo più elevati a causa della modellazione dell'Hessiana, la riduzione del numero di passi di campionamento richiesti e il miglioramento della qualità della generazione (FID e NLL) ne validano l'efficacia per l'addestramento di modelli generativi ad alta fedeltà. Il lavoro colma il divario tra l'efficienza statistica della MLE e il successo pratico dei modelli di diffusione.