Taming Sampling Perturbations with Variance Expansion Loss for Latent Diffusion Models

1. Il Problema: Fragilità degli Spazi Latenti nei Modelli Diffusion

I modelli di diffusione latente (LDM) sono diventati lo standard per la generazione di immagini ad alta fedeltà, operando in uno spazio latente compresso appreso da un autoencoder (solitamente un $\beta$-VAE). Tuttavia, gli autori identificano un problema critico spesso trascurato: la mancanza di robustezza dello spazio latente rispetto alle perturbazioni stocastiche durante il campionamento.

• Osservazione Paradosso: Modelli con ricostruzioni quasi perfette e bassi errori di diffusione (diffusion loss) possono produrre generazioni visivamente inferiori rispetto a modelli con ricostruzioni leggermente peggiori ma spazi latenti più "sparsi".
• Causa Radicale: I tokenizzatori basati su $\beta$-VAE tendono a creare varianze latenti collassate (troppo piccole). Per massimizzare la fedeltà della ricostruzione, l'encoder minimizza l'incertezza ($\sigma^2$), rendendo lo spazio latente eccessivamente compatto.
• Conseguenza: Durante il processo di diffusione, le piccole perturbazioni stocastiche (rumore aggiunto) spingono facilmente i campioni fuori dal manifold dei dati valido. Questo porta a errori di decodifica e degradazione della qualità dell'immagine generata.

2. Metodologia: Variance Expansion (VE) Loss

Per risolvere il problema del collasso della varianza senza sacrificare la qualità della ricostruzione, gli autori propongono una nuova funzione di perdita chiamata Variance Expansion (VE) Loss.

Analisi Teorica

Gli autori dimostrano che la funzione di perdita di ricostruzione ($L_{rec}$) induce un gradiente che spinge la varianza latente $\sigma$ verso zero. In un contesto di diffusione, dove il rumore stocastico è intrinseco, una varianza nulla rende il sistema instabile.

• L'obiettivo è mantenere una varianza adattiva sufficiente a "assorbire" le perturbazioni del campionamento, mantenendo però il centro della distribuzione ($\mu$) preciso per la ricostruzione.

Formulazione della Loss

Invece di utilizzare la regolarizzazione KL standard (che forza una distribuzione a priori fissa e spesso danneggia la ricostruzione), il metodo propone:

1. Rimozione della KL: Si abbandona il termine KL tradizionale, poiché nei modelli di diffusione la distribuzione a priori viene appresa dinamicamente dal modello di denoising.
2. VE Loss ($L_{var}$): Si introduce un termine che penalizza le varianze troppo piccole, incoraggiando un'espansione controllata:
   $$L_{var}(\sigma) = \frac{1}{\sigma^2 + \delta}$$
   dove $\delta$ è un piccolo valore per la stabilità numerica.
3. Regolarizzazione Empirica ($L_{reg}$): Per evitare che la varianza diventi eccessivamente grande, si aggiunge un termine di regolarizzazione esponenziale basato sulla magnitudine del latente:
   $$L_{reg} = e^{|z| - \tau}$$
4. Obiettivo Totale:
   $$L = L_{rec} + \lambda_1 L_{var} + \lambda_2 L_{reg}$$

Questo crea un gioco avversario adattivo: la loss di ricostruzione spinge verso varianze piccole (alta precisione), mentre la VE loss spinge verso varianze maggiori (alta robustezza). Il modello impara automaticamente una varianza adattiva che bilancia fedeltà e stabilità.

3. Contributi Chiave

1. Identificazione del Problema: Dimostrazione empirica e teorica che la robustezza alle perturbazioni di campionamento è un fattore cruciale, spesso più importante della sola accuratezza di ricostruzione, per la qualità generativa.
2. Nuova Funzione di Perdita: Proposta della Variance Expansion Loss, un metodo semplice ed efficace che previene il collasso della varianza latente senza richiedere l'aggiunta euristica di rumore o la riaddestramento di modelli foundation.
3. Bilanciamento Adattivo: Il metodo permette allo spazio latente di adattarsi localmente alla sensibilità del decoder, mantenendo la stabilità nelle regioni ad alta curvatura e la flessibilità in quelle più lisce.

4. Risultati Sperimentali

Gli esperimenti sono stati condotti su ImageNet (256x256) utilizzando diverse architetture (DiT, LightningDiT, VA-VAE).

• Qualità Generativa: L'approccio con VE Loss ha costantemente migliorato i punteggi FID (Fréchet Inception Distance) rispetto alle baseline.
  • Con il modello LightningDiT-B, l'uso della VE Loss ha ridotto il FID da 22.25 a 19.70 (e fino a 12.89 con training esteso).
  • Nel confronto con lo stato dell'arte (SOTA), il metodo proposto ha raggiunto un FID di 1.18 in soli 530 epoch, superando metodi come VA-VAE, RAE e MDTv2 che richiedono più training o parametri.
• Qualità di Ricostruzione: Nonostante l'aumento della varianza, la qualità di ricostruzione rimane competitiva (PSNR, SSIM, LPIPS), dimostrando che non è necessario sacrificare la fedeltà per ottenere robustezza.
• Stabilità: Il metodo ha mostrato una maggiore stabilità durante l'addestramento a lungo termine, risolvendo problemi di instabilità osservati con ottimizzatori standard su spazi latenti allineati a DINOv2.

5. Significato e Impatto

Questo lavoro ribalta una convinzione comune nel campo: una ricostruzione perfetta non garantisce una generazione perfetta.

• Cambiamento di Paradigma: Suggerisce che per i modelli di diffusione latente, la regolarizzazione KL standard è non solo inutile ma dannosa, poiché limita la flessibilità necessaria per gestire il rumore stocastico.
• Soluzione Pratica: La VE Loss offre una soluzione "plug-and-play" che può essere integrata in qualsiasi architettura LDM esistente per migliorare la stabilità e la qualità visiva senza modifiche architetturali complesse.
• Fondamentale per il Futuro: Stabilisce che la progettazione di spazi latenti "generative-friendly" deve prioritizzare la robustezza alle perturbazioni, aprendo nuove direzioni per la ricerca sull'apprendimento di rappresentazioni generative.

In sintesi, il paper dimostra che espandere strategicamente la varianza latente attraverso una loss specifica è la chiave per sbloccare il pieno potenziale dei modelli di diffusione latente, rendendoli più stabili e capaci di produrre immagini di alta qualità.