Riepilogo Tecnico: Modello di Valore Assemblato per l'Allineamento della Diffusione (StitchVM)

1. Enunciato del Problema

I modelli generativi basati su diffusione e flussi eccellono nella cattura delle distribuzioni dei dati, ma spesso richiedono un allineamento con ricompense specifiche per il compito (ad esempio, fedeltà al prompt, preferenze estetiche, preferenze umane) per essere praticamente utili. La sfida centrale in questo allineamento è che le ricompense sono tipicamente definite per immagini di output pulite ($x_0$), mentre le procedure di allineamento (sia al momento dell'inferenza che durante l'addestramento) richiedono la stima della ricompensa attesa di campioni puliti indotti da latenti intermedi rumorosi ($z_t$). Questa stima è nota come funzione di valore.

I metodi esistenti per approssimare questa funzione di valore soffrono di compromessi significativi:

• Approssimazione di Tweedie: Stima la media a posteriori del campione pulito e calcola la ricompensa su tale proxy. Sebbene efficiente, introduce un bias sostanziale in regimi ad alto rumore e richiede valutazioni aggiuntive costose del denoiser e del decoder VAE.
• Approssimazione Monte Carlo (MC): Esegue multiple traiettorie di denoising da $z_t$ e media le ricompense. Sebbene più accurata, comporta costi computazionali proibitivi a causa della necessità di multiple esecuzioni complete.
• Modelli di Valore Appresi: Addestrare direttamente una funzione di valore per latenti rumorosi è un problema aperto. I tentativi esistenti si basano spesso su addestramenti su piccola scala o dataset di preferenze ristretti, risultando in accuratezza e generalizzazione inferiori rispetto ai modelli di ricompensa nello spazio dei pixel su scala fondazionale. Inoltre, addestrare un nuovo modello di valore per ogni nuovo backbone di diffusione o modello di ricompensa è computazionalmente insostenibile.

2. Metodologia: StitchVM

Gli autori propongono StitchVM (Stitched Value Model), un framework che trasferisce le capacità di modelli di ricompensa nello spazio dei pixel preaddestrati e di alta qualità nel regime dei latenti rumorosi con costi computazionali minimi. L'approccio sfrutta il model stitching (assemblaggio di modelli), una tecnica che ricombina livelli di diverse reti neurali.

Architettura Principale

StitchVM costruisce un modello ibrido collegando un backbone di diffusione troncato e congelato (la "testa") a un modello di ricompensa preaddestrato e sezionato (la "coda") tramite un livello di assemblaggio leggero.

• Testa: Un backbone di diffusione (ad esempio, SD 3.5, FLUX) troncato al livello $i$. Questo componente elabora nativamente latenti rumorosi ed estrae caratteristiche consapevoli del rumore.
• Coda: Un modello di ricompensa preaddestrato (ad esempio, CLIP, HPSv2, Aesthetic Predictor) a partire dal livello $j$. Questo componente mantiene la capacità di ricompensa robusta e specifica per il compito appresa durante il preaddestramento su larga scala.
• Livello di Assemblaggio ($s_\psi$): Un modulo leggero che allinea lo spazio delle caratteristiche della testa di diffusione con la coda di ricompensa.

Procedura di Addestramento

L'addestramento di StitchVM avviene in due fasi:

1. Selezione dell'Interfaccia (Forma Chiusa): Gli autori cercano gli indici di assemblaggio ottimali $(i, j)$ in cui le caratteristiche di diffusione (da latenti rumorosi) e le caratteristiche di ricompensa (da immagini pulite) sono più compatibili. Ciò è ottenuto adattando una mappatura lineare $W$ tra caratteristiche accoppiate estratte da un piccolo set di sonde di immagini pulite e i loro corrispondenti latenti rumorosi. Viene selezionata la coppia che minimizza la perdita di matching delle caratteristiche.
2. Fine-tuning Leggero: Una volta fissata l'interfaccia, il livello di assemblaggio e la coda di ricompensa troncata vengono affinati utilizzando immagini pulite non etichettate. L'obiettivo è regressare la funzione di valore $V(z_t)$ verso il punteggio di ricompensa originale $r(z_0)$ dell'immagine pulita. Questo obiettivo off-policy permette al modello di apprendere la ricompensa attesa di campioni puliti condizionati a latenti rumorosi senza richiedere esecuzioni complete.

Applicazione all'Allineamento

StitchVM funge da stimatore diretto della funzione di valore, sostituendo le approssimazioni sia nell'allineamento al momento dell'inferenza che durante l'addestramento:

• Al Momento dell'Inferenza:
  • DPS (Diffusion Posterior Sampling): Sostituisce il gradiente basato su Tweedie con il gradiente diretto di StitchVM ($\nabla_{z_t} V(z_t)$), eliminando la necessità di backpropagation attraverso il denoiser e il VAE.
  • FK Steering: Abilita la "proporzione di proposta", dove ogni particella genera multiple proposte locali. Queste vengono valutate in modo economico da StitchVM per selezionare il candidato migliore, un processo che sarebbe troppo costoso con le approssimazioni di Tweedie.
• Durante l'Addestramento:
  • Fine-tuning Diretto della Ricompensa (ad esempio, AlignProp, DRaFT): Permette di fermare l'addestramento a latenti rumorosi intermedi $z_\tau$ e utilizzare la previsione di StitchVM come obiettivo di supervisione, evitando esecuzioni complete di traiettorie e lunghe catene di backpropagation.
  • DiffusionNFT: Sostituisce le ricompense finali pulite con stime di valore intermedie, abilitando un RL online efficiente sul processo forward.

3. Contributi Chiave

• Un Nuovo Paradigma per la Modellazione del Valore: StitchVM dimostra che potenti modelli di ricompensa nello spazio dei pixel possono essere trasferiti efficacemente nel regime dei latenti rumorosi senza un riaddestramento su larga scala, semplicemente assemblandoli a un backbone di diffusione.
• Efficienza e Accuratezza: Il metodo raggiunge prestazioni all'avanguardia nella stima del valore essendo eccezionalmente leggero. Ad esempio, assemblare CLIP ViT-L in un modello di valore SD 3.5 Medium richiede solo $\approx 10$ ore GPU.
• Miglioramenti delle Prestazioni nell'Allineamento:
  • Inferenza: DPS diventa 3,2$\times$ più veloce con una riduzione del 50% della memoria GPU di picco. FK steering con StitchVM supera significativamente lo standard FK steering e il campionamento Best-of-$N$, in particolare abilitando una selezione efficiente delle proposte locali.
  • Addestramento: DiffusionNFT diventa 2,3$\times$ più veloce, e il fine-tuning diretto della ricompensa (DRaFT) diventa 1,3$\times$ più veloce raggiungendo al contempo una qualità di generazione superiore (ad esempio, +30% su GenEval) grazie a una supervisione efficace in passaggi ad alto rumore.
• Generalizzazione: L'approccio funziona su diversi backbone di diffusione (SD 3.5 Medium/Large, FLUX) e modelli di ricompensa (CLIP, DFN-CLIP, HPSv2, Aesthetic Predictor).

4. Risultati Sperimentali

Il documento valuta StitchVM su diversi benchmark:

• Qualità della Funzione di Valore: Su latenti rumorosi con livelli di rumore variabili ($\sigma \in [0.1, 0.9]$), StitchVM corrisponde da vicino alle prestazioni dei modelli di ricompensa puliti originali nel recupero zero-shot (MSCOCO, Flickr30K), nella previsione delle preferenze (HPDv2, ImageReward) e nella valutazione estetica (AVA). Supera significativamente le linee di base come l'assemblaggio VAE (che collassa ad alto rumore) e i modelli addestrati da zero su dati rumorosi (ad esempio, NoisyCLIP, DiNa-LRM).
• Allineamento all'Inferenza: Negli esperimenti DPS e FK steering su DrawBench, i metodi potenziati da StitchVM migliorano costantemente metriche come ImageReward, HPSv2 e GenEval riducendo drasticamente il costo computazionale. Notabilmente, aumentare il numero di proposte locali ($M$) in FK steering utilizzando StitchVM si rivela più efficiente che aumentare il numero di particelle ($N$).
• Allineamento durante l'Addestramento: Le curve di addestramento mostrano che le varianti di StitchVM raggiungono punteggi finali più alti con meno ore GPU rispetto a DRaFT e DiffusionNFT standard. I risultati qualitativi indicano che i modelli aumentati con StitchVM convergono più velocemente e producono immagini più nitide e di qualità superiore nelle fasi iniziali dell'addestramento.

5. Significato e Affermazioni

Il documento afferma che StitchVM fornisce un modello generico per combinare backbone di diffusione latente con modelli feedforward nello spazio dei pixel senza sacrificare l'esteso preaddestramento di nessuno dei due. Elevando potenti modelli di ricompensa allo spazio latente, abilita un nuovo stile di allineamento della diffusione in cui la corretta funzione di valore per latenti rumorosi è costruita una volta sola e ammortizzata su molti campioni e iterazioni.

Gli autori posizionano StitchVM non come un sostituto del preaddestramento dei modelli fondazionali, ma come un meccanismo pratico e a basso costo per adattare questi modelli a compiti specifici. Sottolineano che questo approccio rende l'allineamento basato su modelli di valore accessibile ed efficiente, potenzialmente accelerando la ricerca verso una modellazione generativa più sicura e affidabile. Il lavoro suggerisce che l'"assemblaggio" di componenti preaddestrati è un'alternativa praticabile ai costi proibitivi dell'addestramento di funzioni di valore da zero o del ricorso a approssimazioni distorte.