Sintesi Tecnica: Imparare a Campionare dai Modelli di Diffusione tramite l'Apprendimento per Rinforzo Inverso

Definizione del Problema

I modelli di diffusione sono diventati uno standard per la generazione di dati ad alta dimensionalità, eppure la procedura di campionamento rimane un collo di bottiglia critico. Sebbene la rete neurale di denoising sia tipicamente fissa dopo l'addestramento, l'algoritmo di campionamento richiede una cura attenta dei superparametri (ad esempio, programmi di rumore, scale di guida, profili di stochasticità). Attualmente, questa regolazione viene eseguita tramite costose ricerche empiriche a griglia (grid search). Per modelli complessi come quelli addestrati su ImageNet-64, il costo computazionale della ricerca esaustiva dei superparametri può avvicinarsi o addirittura superare il costo dell'addestramento del denoiser stesso. Inoltre, gli approoli di apprendimento per rinforzo (RL) o di controllo esistenti per i modelli di diffusione spesso richiedono funzioni di ricompensa esplicite o il fine-tuning del denoiser, il che contraddice l'obiettivo di ottimizzare il campionatore senza modificare il generatore pre-addestrato.

Metodologia

Gli autori propongono di inquadrare il campionamento della diffusione come un problema di decision-making sequenziale risolvibile tramite l'Apprendimento per Rinforzo Inverso (Inverse Reinforcement Learning - IRL). L'innovazione principale è l'apprendimento di una politica di campionamento che corrisponda alle distribuzioni di stato target senza accesso a traiettorie di esperti, ricompense esplicite o ri-addestramento del denoiser.

1. Formulazione MDP

Il processo di campionamento della diffusione è formalizzato come un Processo Decisionale di Markov (MDP) a orizzonte finito:

• Stato ($s_t$): Definito come $(x_t, \sigma_t)$, dove $x_t$ è il campione corrente e $\sigma_t$ è il livello di rumore.
• Azione ($a_t$): Decisioni di controllo discrete applicate alla dinamica di campionamento, come la selezione delle scale di guida, i livelli di iniezione del rumore o la decisione se ri-iniettare rumore (restart) a un precedente livello di rumore.
• Dinamica: Il kernel di transizione è indotto dalle dinamiche di diffusione discretizzate e dall'azione di controllo selezionata. Il processo termina quando il livello di rumore raggiunge $\Sigma_0$ (uno stato assorbente).

2. Apprendimento per Rinforzo Inverso Basato Solo sullo Stato

A differenza dell'IRL tradizionale che spesso assume azioni o traiettorie di esperti, questo metodo opera in un regime basato solo sullo stato. L'"esperto" è definito esclusivamente dalle distribuzioni marginali dei campioni di dati ai livelli di rumore selezionati.

• Obiettivo: L'obiettivo è trovare una politica $\pi_\theta$ la cui misura di occupazione dello stato $\mu_\theta$ (la distribuzione marginale degli stati visitati) corrisponda alla misura di occupazione dell'esperto $\mu_E$.
• Ottimizzazione: Invece di apprendere una funzione di ricompensa, gli autori minimizzano direttamente la divergenza $f$ tra le misure di occupazione dell'esperto e della politica:
  $$\min_\theta D_f(\mu_E \parallel \mu_\theta)$$
  Il paper esplora specificamente le divergenze di Kullback-Leibler (KL) e Reverse-KL (rKL), notando come la scelta della divergenza $f$ determini il compromesso tra copertura dei modi (diversità) e ricerca dei modi (fedeltà).

3. Gradiente della Politica e Ottimizzazione

Gli autori derivano un teorema del gradiente della politica per minimizzare direttamente le divergenze $f$. Il gradiente è espresso come:
$$\nabla_\theta L_f(\theta) = \mathbb{E}_{s_{1:T} \sim p_\theta} \left[ \frac{1}{T} \sum_{t=1}^T \nabla_\theta \log \pi_\theta(a_t|s_t) \sum_{t' \ge t} h_f\left(\frac{\mu_E(s_{t'})}{\mu_\theta(s_{t'})}\right) \right]$$
dove $h_f$ è una funzione derivata dal generatore della divergenza $f$.

• Stima del Rapporto di Densità: Poiché $\mu_E$ e $\mu_\theta$ sono sconosciuti, il rapporto di densità è stimato tramite un discriminatore addestrato per distinguere tra stati dell'esperto (dati con rumore) e stati generati dalla politica.
• Stabilizzazione: Per gestire l'alta varianza nelle stime del gradiente, il metodo impiega un obiettivo troncato (clipped) in stile PPO, correzioni di campionamento per importanza (importance sampling) e un baseline della media globale del batch.

4. Strategie di Campionamento Valutate

Il framework viene applicato a tre specifiche interazioni al tempo di inferenza:

1. Guida Assente di Classe Adattiva (Adaptive Classifier-Free Guidance): Apprendimento di una scala di guida dipendente dallo stato $\omega(x, \sigma)$.
2. Iniezione di Stochasticità: Apprendimento di un livello di iniezione del rumore dipendente dallo stato $\gamma_{EDM}(x, \sigma)$.
3. Renoising (Restart Sampling): Apprendimento di quando riportare il processo a un precedente livello di rumore per uscire da cattive generazioni locali.

Contributi Chiave

1. Formulazione MDP per il Campionamento: Il paper inquadra il campionamento della diffusione come un MDP per adattare i campionatori al tempo di inferenza senza ri-addestrare il denoiser.
2. Corrispondenza Diretta delle Marginali di Stato: Deriva un obiettivo di gradiente della politica che minimizza le divergenze $f$ tra le misure di occupazione degli stati, evitando la necessità di modellare ricompense esplicite o dati di azioni esperte.
3. Compromessi Interpretabili: Il framework espone un compromesso controllabile tra diversità del campione (Recall) e fedeltà (Precision) attraverso la scelta della divergenza $f$ (KL vs rKL).
4. Controllo Dipendente dallo Stato: Le politiche apprese agiscono come controlli al tempo di inferenza che intervengono selettivamente (ad esempio, applicando un'alta guida solo a determinati livelli di rumore o per specifici tipi di campioni) piuttosto che applicare programmi globali fissi.

Risultati Sperimentali

Il metodo è stato valutato su CIFAR-10, FFHQ e ImageNet-64 utilizzando modelli EDM pre-addestrati.

• Prestazioni: Le politiche apprese eguagliano o superano le prestazioni dei migliori baseline determinati manualmente (tramite grid search) in termini di metriche FID, Precision e Recall.
  • Guida Adattiva: Ha migliorato il compromesso qualità-diversità rispetto alla guida fissa. La divergenza KL ha migliorato la Recall (diversità), mentre la rKL ha migliorato la Precision (fedeltà).
  • Iniezione di Stochasticità: I profili dipendenti dallo stato appresi hanno recuperato i pattern qualitativi degli schemi ottimali di EDM (bassa stochasticità alla fine della traiettoria) e hanno superato i baseline fissi su dataset complessi come ImageNet.
  • Renoising: Ha ridotto significativamente il Numero di Valutazioni della Funzione (NFE) su ImageNet mantenendo o migliorando l'FID rispetto alle strategie di restart standard.
• Efficienza Computazionale:
  • Addestramento: Un singolo ciclo di addestramento sostituisce l'esaustiva ricerca a griglia. Su ImageNet-64, ciò riduce il budget di ottimizzazione fino a 9 volte (da ~1437 EFLOPs per una ricerca a griglia di 3 superparametri a ~161 EFLOPs per l'addestramento IRL).
  • Inferenza: La politica aggiunge un carico modesto del 16% al 40% di FLOP rispetto al solo denoiser, un valore considerato trascurabile data l'eliminazione della sintonizzazione manuale.
• Risultati delle Ablation:
  • Condizionamento: Le politiche condizionate sull'intero stato $(x, \sigma)$ superano significativamente quelle condizionate solo sul livello di rumore $\sigma$, e performano quasi quanto quelle con etichette di classe esplicite $(x, \sigma, c)$, suggerendo che la politica apprenda implicitamente caratteristiche discriminanti di classe dai contenuti dell'immagine.
  • Scelta della Divergenza: KL e rKL si sono dimostrate le più efficaci per controllare il compromesso precision-recall, mentre le divergenze simmetriche (TV, Hellinger) hanno prodotto risultati meno interpretabili.

Significato e Dichiarazioni

Il paper sostiene che questo framework fornisce un'alternativa sistematica e automatizzata alla sintonizzazione manuale dei superparametri per il campionamento della diffusione. Trattando il campionamento come un problema di IRL, permette di ottimizzare strategie di campionamento complesse e dipendenti dallo stato che sono difficili da progettare manualmente.

Gli autori sottolineano che il metodo è particolarmente prezioso per modelli su larga scala dove il costo della valutazione per ogni prova in una ricerca a griglia è proibitivo. L'approccio non modifica il modello generativo stesso, ma apprende una "politica di controllo" ottimale per la traiettoria di inferenza, offrendo una via pratica per migliorare la qualità e la diversità dei campioni senza l'onere del ri-addestramento o della ricerca esaustiva. Il paper rimane modesto, osservando che i benefici dipendono dal dataset e che l'attuale framework si basa su spazi di azione discreti, suggerendo lavori futuri in spazi di azione continui e inizializzazioni più robuste.