Variance-Tilted Diffusion Models for Diverse Sampling

Sintesi Tecnica: Modelli di Diffusione con Tilt della Varianza per il Campionamento Diversificato

Problema
I modelli di diffusione sono diventati lo standard per la modellazione generativa in vari domini, dalla sintesi di immagini ad alta risoluzione alla scoperta scientifica (ad esempio, ripiegamento proteico e docking molecolare). Tuttavia, il paradigma di campionamento standard si basa sul disegno di campioni indipendenti e identicamente distribuiti (i.i.d.). Questo approccio è spesso subottimale per compiti a valle che richiedono un insieme diversificato di candidati che coprano regioni distinte del manifold dei dati, come la generazione di ligandi strutturalmente diversi nel drug design o la proposta di traiettorie robuste nella pianificazione del movimento robotico.

I metodi esistenti per imporre la diversità, come la guida a particelle (particle guidance), introducono potenziali repulsivi nelle dinamiche di campionamento. Sebbene efficaci, questi metodi modificano il drift in modo euristico, risultando in una distribuzione target implicita che è computazionalmente intrattabile da valutare o controllare con precisione.

Metodologia
Gli autori propongono una formulazione "target-first" in cui la distribuzione target del batch è definita esplicitamente prima di derivare l'esatto campionatore. L'innovazione principale è una distribuzione del batch pesata sulla varianza che favorisce collezioni di campioni con una grande dispersione empirica dopo una mappa lineare di feature prestabilita.

Per un batch di $n$ campioni $\mathbf{x} = (x^{(1)}, \dots, x^{(n)})$, la distribuzione target $\pi^A_0(\mathbf{x})$ è definita come:
$$\pi^A_0(\mathbf{x}) \propto \text{Var}^A_n(\mathbf{x}) \prod_{i=1}^n p_0(x^{(i)})$$
dove $p_0$ è la distribuzione originale dei dati, $A: \mathbb{R}^d \to \mathbb{R}^k$ è un operatore lineare che mappa i campioni in uno spazio di feature specifico per il task, e $\text{Var}^A_n(\mathbf{x})$ è la varianza empirica delle feature $A(x^{(i)})$.

Per campionare da questa distribuzione, gli autori derivano l'esatto trasformazione di Doob $h$-transform delle dinamiche di diffusione indipendenti. Ciò comporta il calcolo della speranza condizionata della varianza del batch dato lo stato rumoroso al tempo $t$:
$$h^A_t(\mathbf{x}_t) := \mathbb{E}[\text{Var}^A_n(\mathbf{X}_0) \mid \mathbf{X}_t = \mathbf{x}_t]$$

Utilizzando le identità di Tweedie e una decomposizione bias-varianza, gli autori derivano un'espressione in forma chiusa per la correzione dello score $\nabla_{\mathbf{x}_t} \log h^A_t(\mathbf{x}_t)$. Questa correzione consiste in due termini distinti:

1. Termine di Interazione (Repulsione): Un termine che respinge le medie denoisate posteriori ($\mu_t$) delle particelle l'una dall'altra nello spazio delle feature. Questo viene calcolato efficientemente utilizzando prodotti vettoriali-Jacobiano (VJP) attraverso il denoiser.
2. Termine di Curvatura: Un termine che sposta le particelle verso regioni di maggiore varianza delle feature, coinvolgendo la divergenza dello score (specificamente $\text{div}(A^\top A \nabla s_t)$). Questo viene stimato tramite differenze finite o stimatori di traccia stocastici (ad esempio, stimatori di Hutchinson).

Il campionatore risultante è un processo a particelle interagenti in cui lo score per ogni particella è la somma dello score standard e della correzione Doob derivata.

Contributi Chiave

1. Formulazione Target-First: Il documento inquadra la generazione diversificata come un problema di campionamento esatto da una distribuzione target di batch esplicita, evitando le forze repulsive incontrollate dei metodi euristici.
2. Correzione Doob in Forma Chiusa: Gli autori forniscono una derivazione analitica della correzione dello score per un generico operatore lineare $A$, scomponendola in un termine di repulsione per le medie posteriori e un termine di curvatura per la varianza delle feature.
3. Algoritmo di Campionamento Trattabile: Viene introdotto un campionatore a particelle interagenti discretizzato. Il principale termine di interazione sfrutta i calcoli efficienti VJP, mentre il termine di curvatura è gestito tramite stima numerica, rendendo il metodo pratico per i modelli di diffusione esistenti.

Risultati
Gli autori hanno valutato l'approccio sulla generazione di immagini utilizzando Stable Diffusion v1.5 (512x512) con vari prompt. Hanno confrontato tre configurazioni:

• Standard Classifier-Free Guidance (CFG).
• Variance-Tilted (VT) senza il termine di divergenza.
• Variance-Tilted + Divergence (VT + Divergence).

I risultati qualitativi (Figure 1–5) dimostrano che i metodi proposti producono un insieme di campioni più diversificato rispetto al CFG standard. Le immagini generate coprono un intervallo più ampio di pose, sfondi e stili visivi, piuttosto che produrre solo variazioni minori di modalità ad alta probabilità. L'inclusione del termine di divergenza affina ulteriormente la diversità, sebbene gli autori notino che il metodo completo può influire leggermente negativamente sulla qualità dell'immagine e aumentare il costo computazionale dovuto alla stima della divergenza.

Significatività e Rivendicazioni
Il documento sostiene di offrire un'alternativa probabilistica alla repulsione euristica delle particelle. Derivando l'esatto campionatore per una target pesata sulla varianza, il metodo fornisce un obiettivo probabilistico trasparente invece di uno implicito. Gli autori pongono questo lavoro come un approccio rigoroso alla generazione diversificata che mantiene percorsi di implementazione pratica.

Gli autori rimangono modesti riguardo ai limiti, riconoscendo che, sebbene il metodo migliori qualitativamente la diversità, può degradare leggermente la qualità dell'immagine e comportare costi computazionali più elevati a causa della stima delle derivate del secondo ordine (divergenza). Suggeriscono che il lavoro futuro potrebbe concentrarsi sull'annealing della correzione Doob per mitigare i problemi di qualità e sullo sviluppo di stimatori più efficienti per il termine di curvatura.