Gradient-Free Noise Optimization for Reward Alignment in Generative Models

Riepilogo Tecnico: ZeNO (Ottimizzazione del Rumore di Ordine Zero)

Enunciato del Problema
I metodi esistenti di allineamento della ricompensa per i modelli generativi (ad esempio, modelli di diffusione e modelli di flusso) si basano tipicamente sull'ottimizzazione di traiettorie stocastiche multi-step. Sebbene efficaci per questi scenari, tali approcci sono difficili da estendere a generatori deterministici o one-step (ad esempio, modelli di coerenza, modelli di diffusione distillati), che non espongono naturalmente la stessa struttura di traiettoria. Un'alternativa naturale è ottimizzare direttamente il rumore di input. Tuttavia, l'ottimizzazione del rumore basata su gradienti standard richiede la retropropagazione sia attraverso il generatore sia attraverso la funzione di ricompensa. Ciò limita l'applicabilità a pipeline differenziabili e può risultare inaffidabile quando i gradienti riflettono sensibilità locali spurie piuttosto che miglioramenti semantici significativi. Inoltre, molte funzioni di ricompensa pratiche (ad esempio, la progettabilità delle proteine, l'allineamento testo-immagine compositivo) sono non differenziabili o black-box, rendendo di fatto inapplicabili i metodi basati su gradienti.

Metodologia: ZeNO
Gli autori propongono ZeNO (Ottimizzazione del Rumore di Ordine Zero), un framework privo di gradienti che formula l'ottimizzazione del rumore come un problema di controllo integrale di percorso (PIC). L'idea fondamentale è trattare l'evoluzione del rumore di input come un processo stocastico controllabile, ottimizzandolo utilizzando esclusivamente valutazioni di ricompensa di ordine zero.

1. Formulazione del Controllo Integrale di Percorso:
   Il metodo modella l'ottimizzazione del rumore come un'Equazione Differenziale Stocastica (SDE) controllata. L'obiettivo è guidare la traiettoria del rumore verso regioni che producono stati terminali ad alta ricompensa, regolarizzando al contempo lo sforzo di controllo. Il controllo ottimale $u^*(z, t)$ è derivato in una forma integrale di percorso:
   $$u^*(z, t) = \frac{\mathbb{E}_{z:t \to T} [\exp(r(G_\theta(z_T))/\lambda) d w_t]}{\mathbb{E}_{z:t \to T} [\exp(r(G_\theta(z_T))/\lambda)]}$$
   Crucialmente, questo stimatore dipende esclusivamente dalle valutazioni del costo terminale (ricompense) e non richiede la differenziazione attraverso il generatore $G_\theta$ o la funzione di ricompensa $r$.

2. Dinamiche di Riferimento (Ornstein–Uhlenbeck):
   Per garantire che il rumore aggiornato rimanga compatibile con i generatori pre-addestrati, ZeNO adotta il processo di Ornstein–Uhlenbeck (OU) come dinamica di riferimento. A differenza delle SDE con esplosione di varianza, il processo OU preserva la distribuzione stazionaria Gaussiana $N(0, I)$, che corrisponde al prior della maggior parte dei generatori pre-addestrati. Questa scelta collega l'aggiornamento del rumore alla dinamica di Langevin che mira implicitamente a una distribuzione inclinata dalla ricompensa $p^*(z) \propto p(z) \exp(r(G_\theta(z))/\lambda)$.

3. Stimatore Linearizzato e Riduzione della Varianza:
   Lo stimatore integrale di percorso standard soffre di un'alta varianza a causa della ponderazione esponenziale. Gli autori introducono un'approssimazione linearizzata (Proposizione 1) sottraendo la ricompensa media $\bar{r}$ come baseline e linearizzando il termine esponenziale sotto l'ipotesi di una piccola varianza della ricompensa:
   $$u^*(z_t, t) \approx \mathbb{E}[(r(G_\theta(z_T)) - \bar{r}) d w_t] / \lambda$$
   Questo stimatore coincide con lo stimatore della funzione di punteggio utilizzato nelle Strategie Evolutive Naturali e supera empiricamente le varianti esponenziali.

4. Orizzonte Riducente in un Solo Passo:
   Per evitare il costo computazionale di roll-out a lungo orizzonte, ZeNO impiega un'approssimazione a orizzonte ridotto in un solo passo ($H=1$). Ad ogni iterazione, $N$ candidati di rumore perturbati vengono campionati dalle dinamiche di riferimento, valutati dal modello di ricompensa e lo stimatore di controllo linearizzato viene utilizzato per aggiornare lo stato corrente del rumore. Ciò consente un'espansione della scalabilità al momento dell'inferenza aumentando il numero di particelle ($N$) o iterazioni ($M$) senza modificare il generatore.

Contributi Chiave

• Allineamento Privo di Gradienti: ZeNO abilita l'allineamento della ricompensa per generatori deterministici e one-step senza richiedere la retropropagazione attraverso il generatore o il modello di ricompensa.
• Compatibilità Black-Box: Il metodo è applicabile a funzioni di ricompensa non differenziabili e black-box, un contesto in cui i metodi esistenti di ottimizzazione del rumore basati su gradienti (ad esempio, ORIGEN, ReNO) falliscono.
• Connessione Teorica: Il framework stabilisce un legame teorico tra l'ottimizzazione del rumore di ordine zero e il campionamento di Langevin che mira a una distribuzione inclinata dalla ricompensa, fornendo una giustificazione principiale per la regola di aggiornamento.
• Scalabilità al Momento dell'Inferenza: ZeNO introduce un meccanismo flessibile per bilanciare il budget computazionale (tramite $N$ e $M$) con prestazioni di ricompensa migliorate.

Risultati Sperimentali
Gli autori valutano ZeNO su generatori diversificati (SDXL-Turbo, DMD2, LCM, DFGAN) e funzioni di ricompensa (punteggio estetico, PickScore, GenEval e progettabilità delle proteine).

• Generazione di Immagini: ZeNO migliora costantemente i punteggi di ricompensa su tutti i generatori e i tipi di ricompensa testati. Raggiunge prestazioni competitive o superiori rispetto alle baseline basate su gradienti (ORIGEN, $\Psi$-Sampler) e al campionamento Best-of-N. Notevolmente, sul generatore LCM, i metodi basati su gradienti sono peggiorati al di sotto della baseline non controllata, mentre ZeNO ha mantenuto il miglioramento, evidenziando la robustezza dei metodi di ordine zero contro le sensibilità locali dei gradienti.
• Ricompense Non Differenziabili: Nel compito GenEval (conteggio e posizionamento di oggetti), dove le ricompense sono non differenziabili, ZeNO ha superato Best-of-N nei compiti di conteggio, dimostrando che l'ottimizzazione iterativa del rumore estrae più valore dallo stesso budget computazionale rispetto al campionamento indipendente.
• Generazione di Strutture Proteiche: Applicato a un generatore one-step di backbone proteico Riemannian MeanFlow (RMF), ZeNO ha migliorato significativamente la progettabilità (misurata tramite RMSD di auto-consistenza). Ha raggiunto livelli di progettabilità comparabili a una baseline RMF a 5 step utilizzando una sola valutazione del flusso, dimostrando la sua capacità di compensare i passi di campionamento ridotti tramite l'ottimizzazione del rumore.
• Diversità: L'analisi tramite il Vendi Score indica che ZeNO migliora la ricompensa preservando la diversità dei campioni, evitando il collasso su un singolo modo spesso osservato negli approcci basati su gradienti.

Significato e Affermazioni
Il documento afferma che ZeNO fornisce una soluzione pratica e teoricamente fondata per allineare la crescente classe di modelli generativi efficienti, deterministici e one-step con funzioni di ricompensa. Il suo significato primario risiede nella capacità di operare in contesti in cui i gradienti non sono disponibili o sono inaffidabili, come nel caso di modelli di ricompensa black-box o compiti complessi di progettazione biologica. Formulando l'ottimizzazione del rumore come un problema di controllo integrale di percorso, ZeNO offre un framework scalabile e agnostico rispetto al generatore per l'allineamento al momento dell'inferenza che non richiede il riaddestramento o la modifica del modello generativo sottostante. Gli autori notano una limitazione: il metodo richiede multiple valutazioni della ricompensa per passo di aggiornamento, il che può essere computazionalmente oneroso se il modello di ricompensa stesso è costoso.