Riepilogo Tecnico: Allineamento di Modelli Generativi a Pochi Passi mediante Ammortizzazione dell'Inferenza Variazionale Basata su Campioni

Enunciato del Problema

Allineare modelli generativi a pochi passi (ad esempio, GAN, modelli di consistenza, mappe di flusso, modelli in deriva) con obiettivi specifici, come funzioni di ricompensa o preferenze umane, presenta una sfida significativa. I framework di allineamento esistenti si basano tipicamente su assunzioni strutturali restrittive che i generatori a pochi passi violano spesso:

1. Verosimiglianze Tracciabili: Molti metodi richiedono una valutazione esplicita della verosimiglianza sulle traiettorie di denoising, non disponibile per modelli impliciti come le GAN o per mappature a un singolo passo.
2. Solver Specifici: Gli approcci basati sul Controllo Ottimo Stocastico (SOC) o sull'Apprendimento per Rinforzo (RL) assumono spesso solver specifici per Equazioni Differenziali Ordinarie (ODE) o Equazioni Differenziali Stocastiche (SDE) e processi di denoising sequenziali.
3. Specificità del Modello: Le formulazioni sono frequentemente adattate a famiglie di modelli specifiche (ad esempio, modelli di diffusione), fallendo nel generalizzare ad architetture a pochi passi diverse come VAE o mappe di flusso.

Di conseguenza, manca un framework di allineamento generale in grado di affinare generatori a pochi passi diversi preservando il loro vantaggio fondamentale: l'inferenza rapida e a singolo passo.

Metodologia: FAV

Gli autori propongono FAV (Allineamento di Modelli Generativi a Pochi Passi mediante Inferenza Variazionale Basata su Campioni), un framework generale che disaccoppia l'allineamento dalla famiglia di generatori e dalle dinamiche di campionamento. FAV formula l'allineamento come campionamento da una distribuzione inclinata dalla ricompensa ancorata a una distribuzione di riferimento:
$$q^*(x) \propto p_{\text{ref}}(x) \exp(\beta \cdot r(x))$$
dove $p_{\text{ref}}$ è una distribuzione di riferimento (ad esempio, un generatore pre-addestrato o dati empirici) e $r(x)$ è la funzione obiettivo.

Il metodo procede in tre fasi chiave:

1. Approssimazione Basata su Campioni mediante SVGD

FAV approssima la distribuzione target $q^*$ utilizzando la Discesa del Gradiente Variazionale di Stein (SVGD). Invece di richiedere una valutazione della densità, SVGD rappresenta la distribuzione con particelle e le trasporta iterativamente verso il target. Il campo di velocità di Stein ottimale $\phi^*$ è derivato per massimizzare la diminuzione della divergenza KL:
$$\phi^*_{q_\theta, q^*}(x) = \mathbb{E}_{x' \sim q_\theta} \left[ k(x', x) \nabla_{x'} \log p_{\text{ref}}(x') + \beta \cdot k(x', x) \nabla_{x'} r(x') + \nabla_{x'} k(x', x) \right]$$
Questo campo comprende tre componenti:

• Allineamento al Prior: Guida i campioni verso la distribuzione di riferimento $p_{\text{ref}}$.
• Guida della Ricompensa: Spinge i campioni verso regioni ad alta ricompensa.
• Imposizione della Diversità: Previene il collasso modale tramite il termine del gradiente del kernel.

2. Stima Non Parametrica del Punteggio (KDE)

Una sfida critica è che il punteggio $\nabla \log p_{\text{ref}}$ è spesso intrattabile per dati empirici o generatori impliciti. FAV affronta questo problema stimando il punteggio in modo non parametrico utilizzando la Stima della Densità di Kernel (KDE). Convogliando la distribuzione di riferimento con un kernel gaussiano, il punteggio è approssimato come una media ponderata delle differenze tra i campioni di riferimento e le particelle correnti, aggirando efficacemente la necessità di verosimiglianze esplicite.

3. Ammortizzazione mediante Regressione a Punto Fisso

Per preservare la velocità di inferenza rapida dei modelli a pochi passi, FAV non esegue aggiornamenti iterativi SVGD al momento dell'inferenza. Invece, ammortizza il trasporto nei parametri del generatore $\theta$ mediante regressione a punto fisso. Il generatore viene addestrato per minimizzare la distanza tra la sua uscita e il target "trasportato":
$$\mathcal{L}(\theta) = \mathbb{E}_{\epsilon} \left[ \| \text{stopgrad}(f_\theta(\epsilon) + \hat{\phi}^*_{q_\theta, q^*_\sigma}(f_\theta(\epsilon))) - f_\theta(\epsilon) \|_2^2 \right]$$
Qui, il campione trasportato agisce come un target di regressione stabile. Ciò permette al generatore di imparare a produrre campioni allineati in un singolo passaggio in avanti.

Estensione a Ricompense Black-Box: Per ricompense non differenziabili, FAV impiega la stima del gradiente di ordine zero (levigatura gaussiana) per approssimare $\nabla r(x)$, consentendo l'allineamento anche quando la funzione di ricompensa è una scatola nera.

Contributi Chiave

1. Framework di Allineamento Generale: FAV è il primo framework in grado di allineare generatori a pochi passi diversi (inclusi VAE, GAN, modelli di consistenza e mappe di flusso) senza richiedere verosimiglianze tracciabili, solver ODE/SDE specifici o traiettorie multi-passo.
2. Formulazione VI Basata su Campioni: Formulando l'allineamento come campionamento da una distribuzione inclinata dalla ricompensa utilizzando SVGD e KDE, FAV disaccoppia la procedura di allineamento dalle dinamiche interne del generatore.
3. Preservazione dell'Efficienza: Tramite l'ammortizzazione mediante regressione a punto fisso, FAV mantiene l'efficienza di inferenza a singolo passo del generatore di base, evitando il sovraccarico computazionale di processi di campionamento iterativo o distillazione.
4. Compatibilità con Black-Box: L'integrazione della stima del gradiente di ordine zero permette a FAV di ottimizzare ricompense arbitrarie di tipo black-box.

Risultati Sperimentali

Gli autori hanno valutato FAV su due domini principali: manipolazione robotica e allineamento di generatori di immagini.

1. Manipolazione Robotica (RL Offline e Offline-to-Online)

• Benchmark: 56 compiti RL offline (OGBench e D4RL) e 30 compiti RL offline-to-online.
• Prestazioni: FAV ha raggiunto prestazioni allo stato dell'arte in tutti i compiti, superando 10 baseline inclusi gradienti di politica riparametrizzati, Adjoint Matching e politiche generative distillate (ad esempio, FQL, IFQL).
• Risultato Chiave: FAV ha dimostrato capacità superiori di estrazione della politica e prestazioni di transizione dall'apprendimento offline a quello online, anche senza modelli generativi ausiliari o denoising multi-passo.

2. Allineamento di Generatori di Immagini

• Scale: Gli esperimenti hanno variato da ImageNet-256 (256x256) alla sintesi testo-immagine ad alta risoluzione (1024x1024).
• Modelli: Testati su StyleGAN-XL, modelli in deriva, iMeanFlow e SANA-Sprint.
• Metriche: Valutati rispetto a punteggi estetici, punteggi di preferenza umana (HPS) e rilevatori di sicurezza (NSFW).
• Prestazioni:
  • FAV ha costantemente raggiunto il miglior fronte di Pareto tra ottimizzazione della ricompensa e diversità/qualità del campione, mitigando l'iper-ottimizzazione della ricompensa meglio di baseline come DRaFT e Adjoint Matching.
  • Efficienza di Inferenza: FAV ha generato campioni allineati circa 180-280 volte più velocemente rispetto ai metodi di allineamento al momento dell'inferenza (ad esempio, Best-of-N, ReNO) mantenendo una qualità dell'immagine competitiva.
  • Scalabilità: Il metodo è stato scalato con successo da risoluzioni 256x256 a 1024x1024 e ha generalizzato attraverso diverse architetture di modelli.

Significato e Affermazioni

Il documento afferma che FAV fornisce una soluzione unificata e agnostica rispetto al modello per l'allineamento di modelli generativi a pochi passi. Il suo significato risiede nel:

• Rimuovere Barriere Strutturali: Elimina la necessità di verosimiglianze tracciabili o traiettorie di campionamento specifiche, rendendo l'allineamento accessibile a una classe più ampia di modelli generativi impliciti.
• Bilanciare Efficienza e Qualità: Riesce a conciliare il compromesso tra generazione ad alta ricompensa ed efficienza computazionale dell'inferenza a pochi passi, una sfida in cui i metodi precedenti spesso degradavano le prestazioni o richiedevano distillazioni complesse.
• Robustezza: Il framework dimostra robustezza su compiti diversi, dal controllo a bassa dimensionalità alla sintesi di immagini ad alta risoluzione, gestendo sia funzioni di ricompensa differenziabili che black-box.

Gli autori concludono che FAV rappresenta un approccio forte e generico per l'estrazione di politiche in RL e l'allineamento delle ricompense nella generazione di immagini, offrendo un'alternativa pratica a metodi più complessi e ricchi di assunzioni.