Riassunto Tecnico: NoiseTilt: Noise-Tilted Reverse Kernels per l'Allineamento della Ricompensa nella Diffusione

Problematica

La scalabilità durante l'inferenza è emersa come un potente meccanismo per allineare i modelli di diffusione pre-addestrati alle preferenze dell'utente senza costosi processi di ri-addestramento. Quando le preferenze sono espresse come funzioni di ricompensa (reward), il processo iterativo di denoising dei modelli di diffusione offre un modo agnostico rispetto al modello per guidare il campionamento verso regioni ad alta ricompensa. Tuttavia, i metodi esistenti affrontano un compromesso fondamentale tra guida del gradiente e compatibilità con il rumore:

1. Metodi a Spostamento della Media (Guida del Gradiente): Approcci come DPS [10] iniettano i gradienti della ricompensa direttamente nella media del kernel gaussiano inverso. Sebbene efficaci nel guidare la generazione, essi spostano gli stati campionati al di fuori del "regime compatibile con il rumore" — la stretta regione dello spazio latente in cui il modello pre-addestrato è stato istruito per operare. Questo spostamento rende gli stati intermedi fuori distribuzione (out-of-distribution), degradando la qualità della generazione e rendendo il sistema incline al "reward hacking".
2. Metodi Basati sulla Ricerca: Approcci come SVDD [38] preservano la media inversa di base estraendo $K$ candidati dal kernel stocastico e selezionando il migliore. Sebbene ciò mantenga la compatibilità con il rumore e la qualità del campione, rinuncia alla guida diretta del gradiente, affidandosi al campionamento casuale per trovare direzioni che migliorino la ricompensa. Questo è computazionalmente inefficiente, richiedendo un numero proibitivo di campioni per ottenere un forte allineamento.

Nessun metodo precedente riesce a ottenere simultaneamente la guida del gradiente e la compatibilità con il rumore.

Metodologia: Noise-Tilted Reverse Kernel (NTRK)

Gli autori propongono i Noise-Tilted Reverse Kernels (NTRK), un campionatore che risolve questo compromesso mantenendo fissa la media inversa e instradando l'informazione della ricompensa esclusivamente attraverso il termine del rumore.

Meccanismo Centrale

Il passo inverso standard è definito come $\bm{x}_{t-1} = \bm{\mu}_\theta(\bm{x}_t, t) + \sigma_t \bm{\epsilon}_t$, dove $\bm{\epsilon}_t \sim \mathcal{N}(\bm{0}, \bm{I})$.
NTRK modifica il termine del rumore in:
$$\bm{x}_{t-1} = \bm{\mu}_\theta(\bm{x}_t, t) + \sigma_t \tilde{\bm{\epsilon}}_t$$
dove il rumore guidato $\tilde{\bm{\epsilon}}_t$ è una miscela di una direzione di ricompensa sbiancata (whitened) e stocasticità non polarizzata:
$$\tilde{\bm{\epsilon}}_t = \sqrt{\rho_t} \, \mathcal{W}(\nabla_{\bm{x}_t} r(\hat{\bm{x}}_{0|t})) + \sqrt{1-\rho_t} \, \bm{\epsilon}_t$$
Qui, $\rho_t \in [0, 1]$ controlla la forza della guida, e $\hat{\bm{x}}_{0|t}$ è la stima di Tweedie dell'immagine pulita.

L'Operatore di Sbiancamento ($\mathcal{W}$)

La sfida centrale è che un gradiente di ricompensa grezzo è strutturato e deterministico, il che viola le proprietà statistiche del tipico rumore gaussiano richiesto dal modello pre-addestrato. Iniettarlo direttamente creerebbe un "rumore atipico", causando artefatti.

Per risolvere questo problema, gli autori introducono un operatore di sbiancamento $\mathcal{W}$ che trasforma il gradiente della ricompensa in un vettore di rumore "tipico" prima dell'iniezione. Questo operatore assicura che la perturbazione rimanga nel regime compatibile con il rumore preservando al contempo il segnale direzionale del gradiente. L'operatore è costruito tramite una sequenza di proiezioni su insiemi ad alta confidenza derivati dalla statistica gaussiana standard:

1. Proiezione a Due Livelli di Statistica d'Ordine (2OS): Il vettore di input viene rimodellato in tile. I valori vengono ordinati all'interno di ogni tile e poi tra le tile. Le statistiche risultanti dal doppio ordinamento sono tagliate ai livelli di confidenza del 99,99% derivati dalla distribuzione Beta. Ciò sopprime gli artefatti strutturati localizzati mantenendo l'istogramma globale.
2. Statistiche per Tile: L'operatore vincola ulteriormente la media e l'energia centrata di ogni tile per farle corrispondere alle distribuzioni attese del rumore gaussiano standard.
3. Domini di Trasformazione Ortogonale: Le stesse proiezioni di intervallo di confidenza sono applicate in domini trasformati (ad esempio, mix Fourier compatto e stile Hadamard) per sopprimere le strutture globali residue che non sono visibili nel dominio spaziale.

Fondamentalmente, per input che sono già rumore gaussiano tipico, $\mathcal{W}$ li lascia quasi invariati, mentre gli input strutturati vengono sostanzialmente sbiancati.

Contributi Chiave

1. Risoluzione del Compromesso: NTRK è il primo metodo che riesce simultaneamente a ottenere la guida basata sul gradiente e la compatibilità con il rumore. Preserva la media inversa pre-addestrata (a differenza dei metodi a spostamento della media) pur utilizzando un singolo campione per passo (a differenza dei metodi basati sulla ricerca).
2. L'Operatore di Sbiancamento: Il paper identifica che la qualità dell'allineamento della ricompensa è direttamente legata alla qualità dell'operatore di sbiancamento. La proiezione basata su 2OS proposta si dimostra superiore alla semplice proiezione della norma o ai vincoli di uguaglianza rigidi (come WGNC [24]), poiché bilancia la soppressione della struttura con la preservazione della naturale variabilità del rumore.
3. Efficienza: Instradando i gradienti attraverso il termine del rumore, NTRK raggiunge prestazioni allo stato dell'arte con un numero significativamente inferiore di valutazioni della funzione (NFE).

Risultati Sperimentali

Gli autori hanno valutato NTRK attraverso tre compiti primari utilizzando FLUX [35] per le immagini e Wan2.1 [70] per i video:

• Generazione di Immagini Estetiche: NTRK ha ottenuto la ricompensa target (Aesthetic Score) più alta tra tutti i baseline. Notevolmente, NTRK con soli 25 NFE ha superato la ricompensa del miglior baseline (SVDD) che utilizzava 500 NFE, rappresentando una riduzione del calcolo di 20$\times$. Ha inoltre mantenuto o migliorato le metriche di controllo come ImageReward e HPSv2.
• Generazione di Immagini Allineate al Testo: Su prompt composizionali complessi (T2I-CompBench++), NTRK ha superato i baseline in PickScore (ricompensa target) e nelle metriche di controllo come VQA Score e HPSv2.
• Generazione di Video Allineata alle Preferenze: Utilizzando VideoReward come target, NTRK ha generato video con migliore allineamento testuale, qualità del movimento e qualità visiva rispetto a DPS e FreeDoM.
• Ulteriori Applicazioni: Il metodo è stato applicato con successo alla generazione consapevole della quantità (conteggio di oggetti) e all'allineamento della ricompensa basato su VLM (posizionamento dell'orizzonte e della profondità), dimostrando una vasta applicabilità.
• Integrazione con il Fine-Tuning: È stato dimostrato che NTRK è ortogonale al fine-tuning (ad esempio, MixGRPO [37]), fornendo ulteriori guadagni quando applicato su modelli già sottoposti a fine-tuning.

Significato e Rivendicazioni

Il paper sostiene che NTRK offra un meccanismo semplice e ampiamente applicabile per l'allineamento della ricompensa durante l'inferenza che lascia intatte le dinamiche di campionamento pre-addestrate. Trattando il gradiente della ricompensa come una perturbazione da sbiancare piuttosto che come uno spostamento della media, NTRK evita i problemi di fuori distribuzione che affliggono la guida basata sul gradiente.

Gli autori sottolineano che l'efficacia del metodo dipende fortemente dall'operatore di sbiancamento. I loro studi di ablazione mostrano che operatori più deboli non riescono a rimuovere la struttura dannosa, mentre operatori troppo rigidi distorcono il rumore genuino. L'approccio proposto basato su 2OS naviga con successo questo compromesso, consentendo una generazione ad alta ricompensa senza sacrificare la qualità inerente al modello pre-addestrato.

Il lavoro suggerisce che la futura scalabilità all'inferenza possa andare oltre la ricerca a forza bruta o gli spostamenti distruttivi della media, sfruttando invece le proprietà statistiche del rumore per guidare la generazione in modo efficiente.