Sintesi Tecnica: LiDAR (Lookahead Sample Reward Guidance) per lo Scaling al Tempo di Test dei Modelli di Diffusione

Definizione del Problema

I modelli di diffusione hanno raggiunto prestazioni generative straordinarie in ambiti di immagini, video e linguaggio; tuttavia, i campioni generati spesso non si allineano pienamente con l'intento umano. Sebbene esistano funzioni di ricompensa che catturano le preferenze umane, l'integrazione di queste nel processo di campionamento presenta sfide significative. L'obiettivo centrale dello scaling al tempo di test è campionare da una distribuzione inclinata dalla ricompensa $p_\theta^r(x_0|c) \propto p_\theta(x_0|c) \exp(\lambda \cdot r(x_0, c))$.

Raggiungere questo obiettivo richiede il calcolo della Ricompensa Futura Attesa (EFR - Expected Future Reward), $r_t^\lambda(x_t, c)$, per le particelle intermedie $x_t$ al fine di guidare il processo di diffusione inversa. I metodi esistenti affrontano limitazioni critiche:

1. Backward Rollout: Il campionamento diretto da $x_t$ a $x_0$ per stimare l'EFR è computazionalmente proibitivo a causa della necessità di un denoising iterativo ad ogni timestep.
2. Approssimazioni basate su Tweedie: Metodi come la Gradient Guidance e il Sequential Monte Carlo (SMC) si affidano ad approssimazioni di Taylor del primo ordine (formula di Tweedie) per stimare l'EFR. Questi introducono bias, specialmente quando la forza del segnale di ricompensa ($\lambda$) è elevata, e spesso richiedono la backpropagation attraverso reti neurali (Gradient Guidance) o soffrono di collasso del campione in spazi ad alta dimensionalità (SMC).

Metodologia: LiDAR Sampling

Gli autori propongono LiDAR (Lookahead Sample Reward Guidance), un framework che disaccoppia la dipendenza neurale tra la particella intermedia $x_t$ e l'EFR, consentendo una guida a forma chiusa e priva di derivate (derivative-free).

1. Riformulazione della Ricompensa Futura Attesa

Il paper deriva una nuova formulazione per l'EFR (Teorema 3.1) che esprime l'aspettativa della ricompensa utilizzando solo campioni marginali da un modello di diffusione pre-addestrato e dal kernel di perturbazione forward, piuttosto che condizionarsi sulla particella corrente $x_t$ tramite reti neurali:
$$r_t^\lambda(x_t, c) = \log \mathbb{E}_{p_\theta(x_0|c)} \left[ \frac{p(x_t|x_0)}{\mathbb{E}_{p_\theta(x_0|c)}[p(x_t|x_0)]} \exp(\lambda \cdot r(x_0, c)) \right]$$
Questo approccio di "forward rollout" permette di calcolare l'EFR per qualsiasi $x_t$ utilizzando un insieme pre-generato di campioni marginali $\{\hat{x}_0^i\}$, eliminando la necessità di backpropagation attraverso la rete di score $s_\theta$ o il modello di ricompensa durante lo step di guida.

2. Campionamento Lookahead Efficiente

Per mitigare il costo della generazione di campioni marginali per ogni prompt, gli autori introducono una strategia di lookahead sampling efficiente. Invece di utilizzare l'intero e lento processo di diffusione, utilizzano un generatore veloce $q(x_0|c)$, come un solver ODE a pochi step (es. DPM-Solver) o un modello distillato, per generare un insieme di $n$ campioni lookahead $\{\hat{x}_0^i\}_{i=1}^n$. Questi campioni sono annotati con valori di ricompensa $r(\hat{x}_0^i, c)$.

3. Guida a Forma Chiusa Priva di Derivate

Utilizzando l'insieme finito di campioni lookahead, gli autori derivano un'espressione in forma chiusa per il gradiente della ricompensa empirica lookahead (Teorema 3.3). Il termine di guida orienta la particella corrente $x_t$ verso i campioni lookahead ad alta ricompensa e la allontana da quelli a bassa ricompensa:
$$\nabla_{x_t} \hat{r}_t^\lambda(x_t, c) = \sum_{i=1}^n (w_i^r - w_i) \frac{\hat{x}_0^i}{\sigma_t^2}$$
dove $w_i^r$ e $w_i$ sono pesi softmax basati sui valori di ricompensa e sulla distanza tra $x_t$ e i campioni lookahead $\hat{x}_0^i$. Questa formulazione non richiede la differenziazione neurale.

4. Leggi di Scaling Teoriche

Il paper stabilisce le leggi di scaling per il metodo:

• Accuratezza Lookahead ($\delta$): La distanza di variazione totale tra la distribuzione target e la distribuzione generata da LiDAR diminuisce come $O(1/\sqrt{\delta})$, dove $\delta$ è il numero di step di discretizzazione del solver lookahead.
• Dimensione del Campione ($n$): La varianza della stima della ricompensa empirica diminuisce come $O(1/\sqrt{n})$, implicando che aumentare il numero di campioni lookahead migliora l'accuratezza dell'approssimazione.

Risultati Chiave

Gli esperimenti sono stati condotti su SD v1.5, SDXL, FLUX (flow matching) e UDLM (diffusione discreta) utilizzando le metriche ImageReward, CLIP, HPS e GenEval.

• Performance vs. Efficienza: Su SDXL, LiDAR raggiunge prestazioni su GenEval comparabili alle più recenti tecniche di gradient guidance (DATE), ma con un incremento di velocità di 9.5× nel tempo di inferenza.
• Efficienza della Memoria: A differenza dei metodi di gradient guidance che richiedono circa 3× più memoria a causa della backpropagation, l'uso della memoria di LiDAR è quasi identico al campionamento vanilla, poiché la fase di campionamento target (Algoritmo 2) non richiede il calcolo del gradiente.
• Robustezza al Reward Hacking: Mentre i metodi di gradient guidance spesso degradano in metriche come CLIP e HPS quando ottimizzano per ImageReward (reward hacking), LiDAR migliora simultaneamente tutte le metriche. Ciò è attribuito alla correlazione tra ImageReward e altre metriche nei campioni lookahead, che guida naturalmente il campionatore in una direzione bilanciata.
• Miglioramenti Ortogonali: LiDAR fornisce ulteriori guadagni di performance quando applicato a modelli già affinati con Direct Preference Optimization (DPO), dimostrando la sua ortogonalità all'allineamento basato su training.
• Scalabilità: Il metodo scala efficacemente con il numero di campioni lookahead ($n$) e gli step del solver ($\delta$), e può essere combinato con altre tecniche come SMC e Best-of-N.

Significato e Rivendicazioni

Il paper sostiene che LiDAR affronta sistematicamente i limiti di efficienza e di reward-hacking dei metodi esistenti di scaling al tempo di test. Riformulando il calcolo dell'EFR per fare affidamento su forward rollout e campioni marginali, gli autori consentono una guida priva di derivate che è sia computazionalmente efficiente che teoricamente fondata.

Gli autori sottolineano che LiDAR raggiunge un trade-off efficienza-performance superiore rispetto ai baseline esistenti. Nello specifico, evita l'overhead computazionale della backpropagation neurale e i problemi di collasso del campione tipici di SMC, fornendo al contempo stime dell'EFR più accurate rispetto ai metodi basati su approssimazioni di Taylor. Il lavoro suggerisce che il disaccoppiamento della dipendenza neurale del segnale di ricompensa dalla particella intermedia sia una via percorribile verso un modellazione generativa scalabile e allineata ai valori umani.