Purrception: Variational Flow Matching for Vector-Quantized Image Generation

1. Il Problema: Il Dilemma dello Spazio Latente VQ

La generazione di immagini ad alta risoluzione tramite spazi latenti vettoriali quantizzati (VQ) presenta una tensione fondamentale tra due nature dei dati:

1. Natura Discreta: I dati sono rappresentati come indici discreti tratti da un vocabolario (codebook) finito.
2. Natura Continua: Questi indici corrispondono a vettori di embedding continui che possiedono relazioni geometriche (prossimità, direzione).

Le approcci esistenti falliscono nel gestire simultaneamente queste due proprietà:

• Metodi Continui (es. Flow Matching Continuo - CFM, Diffusion Latente): Operano nello spazio degli embedding continui. Preservano la geometria e permettono un trasporto fluido, ma ignorano la struttura categorica. Non ricevono un segnale di apprendimento categorico diretto, non possono esprimere incertezza su più codici plausibili e mancano di "logits" per controlli come il temperature scaling.
• Metodi Discreti Puri (es. Flow Matching Discreto - DFM, VQ-Diffusion): Predicono direttamente gli indici discreti. Sebbene allineati alla struttura quantizzata, collassano la geometria: trattano codici semanticamente correlati come token non correlati. Le previsioni diventano "teletrasporti" discreti tra indici, rendendo inutile la modellazione dell'incertezza e il controllo tramite temperatura.

2. Metodologia: Purrception (VQ-VFM)

Gli autori introducono Purrception, un adattamento del Variational Flow Matching (VFM) specifico per latenti VQ. L'approccio ibrido risolve il compromesso sopra descritto combinando la supervisione categorica con la dinamica di trasporto continuo.

Concetti Chiave:

• Posteriore Variazionale Categoriale: Invece di predire direttamente il vettore di velocità o l'indice finale, il modello apprende una distribuzione di probabilità posteriore categorica $q_\theta(c | z_t)$ sugli indici del codebook, data la posizione corrente nello spazio continuo $z_t$.
• Trasporto nello Spazio Continuo: La velocità del flusso ($v_\theta$) viene calcolata come l'aspettazione di questa distribuzione categorica sui vettori del codebook.
  $$v_\theta(z_t) = \sum_{k=1}^K \pi_{\theta, k}(z_t) \left( \frac{e_k - z_t}{1-t} \right)$$
  Dove $\pi_{\theta, k}$ è la probabilità di scegliere il codice $k$ e $e_k$ è il vettore di embedding corrispondente.
• Ottimizzazione: L'obiettivo di apprendimento si riduce a una semplice perdita di entropia incrociata (Cross-Entropy) tra la distribuzione predetta e l'indice ground-truth del codice quantizzato.
• Controllo tramite Temperatura ($\tau$): Poiché la distribuzione è ottenuta tramite una softmax sui logits, il modello eredita un parametro di temperatura. Questo permette di controllare l'incertezza durante l'inferenza:
  • $\tau$ basso: La distribuzione collassa sull'indice più probabile (output nitidi ma potenzialmente semplificati).
  • $\tau$ alto: La probabilità si distribuisce su codici vicini, aggiungendo dettaglio e varietà (ma rischiando rumore).

3. Contributi Principali

1. Nuovo Framework Ibrido: Purrception è il primo metodo che applica il Variational Flow Matching agli spazi latenti VQ, unendo la supervisione categorica (per imparare la struttura discreta) con la dinamica di trasporto continuo (per preservare la geometria).
2. Efficienza di Addestramento: Il metodo dimostra una convergenza significativamente più rapida rispetto ai baselines puramente continui (CFM) e puramente discreti (DFM).
3. Controllabilità Intrinseca: Introduce il temperature scaling come meccanismo di controllo per la qualità e la diversità delle immagini generate, una funzionalità assente nei metodi continui puri e inefficace in quelli discreti puri.
4. Quantificazione dell'Incertezza: Il modello è in grado di esprimere incertezza su codici plausibili e tradurla in un movimento geometrico fluido nello spazio degli embedding, evitando i salti discreti bruschi.

4. Risultati Sperimentali

Gli esperimenti sono stati condotti sulla generazione di immagini ImageNet-1k a 256x256, utilizzando tokenizer VQ (Stable Diffusion vq-f8 e LlamaGen vq-ds8-c2i) e backbone DiT (L/2 e XL/2).

• Velocità di Convergenza:
  • Purrception raggiunge le prestazioni di base (FID) molto prima rispetto ai competitor.
  • Con il backbone DiT-L/2, è 1.65 volte più veloce di CFM-endpoint e 3.0 volte più veloce di DFM.
  • Con il backbone DiT-XL/2, la velocità di convergenza è 2.3 volte superiore ai baselines CFM e 3.5 volte superiore a DFM.
• Qualità del Campionamento (FID):
  • Purrception ottiene un FID di 3.88, risultando competitivo con i migliori modelli autoregressivi e superando tutti i modelli di diffusione discreta e generativi mascherati.
  • Sebbene non superi i modelli di diffusione continua di fascia alta (come DiT-XL/2 o SiT-XL/2) che usano autoencoder VAE di qualità superiore e tempi di addestramento più lunghi, Purrception si avvicina notevolmente a tali prestazioni mantenendo l'efficienza del VQ.
• Effetto della Temperatura:
  • È stata osservata una relazione a "U" tra la temperatura di softmax e il FID. Un valore ottimale intorno a $\tau \approx 0.8-0.9$ massimizza la qualità, bilanciando nitidezza e dettaglio.
  • Questo permette di migliorare la qualità dei campioni in fase di inferenza senza riaddestrare il modello.

5. Significato e Implicazioni

Il lavoro di Purrception è significativo perché risolve il compromesso storico tra la preservazione della geometria e la supervisione categorica nella generazione di immagini VQ.

• Efficienza: Dimostra che l'uso di un obiettivo categorico (cross-entropy) all'interno di un framework di flusso continuo accelera drasticamente l'ottimizzazione.
• Flessibilità: Fornisce un meccanismo di controllo (temperatura) che era precedentemente inaccessibile ai modelli di flusso continuo, permettendo di adattare il trade-off tra fedeltà e diversità.
• Futuro: Apre la strada all'applicazione di questo approccio ibrido ad altri domini (audio, video, forme 3D) e suggerisce che la combinazione di trasporto continuo e supervisione discreta è una direzione promettente per la prossima generazione di modelli generativi efficienti.

In sintesi, Purrception dimostra che è possibile "avere il meglio di entrambi i mondi": la fluidità geometrica dei metodi continui e la precisione strutturale dei metodi discreti, portando a un addestramento più rapido e a una generazione di immagini di alta qualità.