Riepilogo Tecnico: DiNa-LRM – Modellazione del Reward Latente Nativa alla Diffusione

1. Enunciato del Problema

L'allineamento dei modelli di diffusione e di adattamento del flusso con le preferenze umane dipende fortemente dalle funzioni di reward. Gli approcci all'avanguardia attuali utilizzano prevalentemente Modelli Vision-Language (VLM) come backbones per il reward. Sebbene i VLM offrano capacità discriminative robuste, presentano due limitazioni critiche per l'allineamento dei modelli di diffusione:

1. Sovraccarico Computazionale e di Memoria: I reward basati su VLM sono costosi da interrogare ripetutamente durante l'ottimizzazione, creando un collo di bottiglia per l'efficienza dell'addestramento.
2. Disallineamento del Dominio: I VLM operano tipicamente nello spazio dei pixel, mentre i generatori di diffusione latente sono addestrati e ottimizzati nello spazio latente. Questo disallineamento da pixel a latente complica l'allineamento, in particolare per i metodi basati sul gradiente del reward, e introduce un sovraccarico di sistema non necessario.

Al contrario, sebbene i modelli di diffusione abbiano mostrato promesse come backbones discriminativi, i tentativi precedenti di utilizzarli come reward sono stati spesso legati ad algoritmi di ottimizzazione specifici a livello di step o si sono concentrati su stati intermedi rumorosi senza stabilire un framework generico paragonabile ai VLM per la valutazione di campioni puliti.

2. Metodologia: DiNa-LRM

Gli autori propongono DiNa-LRM (Diffusion-Native Latent Reward Model), un framework che formula l'apprendimento delle preferenze direttamente sugli stati rumorosi della diffusione all'interno dello spazio latente. Il metodo è composto da tre componenti fondamentali:

A. Formulazione delle Preferenze Nativa alla Diffusione

Invece di valutare direttamente i campioni puliti ($x_0$), DiNa-LRM estende il modello di Thurstone agli stati rumorosi della diffusione ($x_t$).

• Incertezza Calibrata sul Rumore: Gli autori riconoscono che i giudizi di preferenza su stati rumorosi sono intrinsecamente meno certi all'aumentare del livello di rumore. Introducono un'incertezza di confronto calibrata sul rumore $\sigma_u^2(t)$ che scala con il livello di rumore della diffusione $\sigma^2(t)$:
  $$\sigma_u^2(t) = k\sigma^2(t) + \sigma_u^2$$
  Questa formulazione garantisce che l'obiettivo del reward rimanga coerente con la varietà rumorosa incontrata durante l'addestramento preliminare della diffusione. Agisce come un regolarizzatore consapevole dell'incertezza, prevenendo che gradienti ad alta varianza destabilizzino l'addestramento in regimi ad alto rumore.
• Obiettivo di Addestramento: Il modello è ottimizzato utilizzando una perdita di fedeltà su coppie di preferenze $(x_t^+, x_t^-)$, dove la probabilità di confronto è modulata dall'incertezza dipendente dallo step temporale.

B. Architettura del Reward Latente

DiNa-LRM è istanziato su un backbone di diffusione latente pre-addestrato (ad es. SD3.5-Medium) con una testa di valutazione leggera basata su query.

• Estrazione di Caratteristiche Multistrato: Le caratteristiche visive e testuali vengono estratte da più livelli del backbone.
• Adattamento Condizionato allo Step Temporale: Le caratteristiche sono modulate dagli embedding dello step temporale utilizzando FiLM (Feature-wise Linear Modulation) per rendere la testa del reward esplicitamente consapevole del livello di rumore.
• Valutazione Q-Former con Gate: Un Q-Former con gate aggrega queste caratteristiche multistrato condizionate allo step temporale per produrre un punteggio di reward scalare. Questa architettura opera interamente nello spazio latente VAE, evitando la decodifica nello spazio dei pixel.

C. Scaling al Momento dell'Inferenza tramite Ensemble di Rumore

Per migliorare la robustezza e ridurre la sensibilità a uno specifico step temporale di valutazione, DiNa-LRM supporta l'ensemble di rumore.

• Invece di valutare un campione pulito a un singolo livello di rumore fisso, il modello valuta il campione a più step temporali $\{t_k\}$.
• Le caratteristiche estratte a questi diversi step temporali vengono concatenate ed elaborate dalla stessa testa Q-Former per produrre un punteggio aggregato. Questo funge da meccanismo di scaling "nativo alla diffusione" al momento del test.

3. Contributi Chiave

1. Formulazione delle Preferenze Nativa alla Diffusione: Il documento estende il modello di preferenze di Thurstone agli stati rumorosi della diffusione introducendo un'incertezza calibrata sul rumore che scala con il livello di rumore della diffusione, colmando il divario distribuzionale tra pre-addestramento e apprendimento del reward.
2. Scaling al Momento dell'Inferenza: Gli autori propongono una strategia di ensemble di rumore che aggrega caratteristiche multistep temporali, fornendo un meccanismo per lo scaling al momento del test che sfrutta la traiettoria globale della diffusione piuttosto che una singola stima puntuale.
3. Validazione Empirica: Il lavoro dimostra che i reward nativi alla diffusione possono raggiungere prestazioni competitive con giudici VLM all'avanguardia operando nativamente nello spazio latente.

4. Risultati Sperimentali

Gli autori hanno valutato DiNa-LRM su diversi benchmark e scenari di allineamento:

• Accuratezza della Modellazione del Reward:

  • DiNa-LRM supera sostanzialmente le baseline esistenti basate su reward di diffusione (ad es. LRM-SD1.5, LRM-SDXL) su ImageReward, HPDv2, HPDv3 e GenAI Bench.
  • Raggiunge prestazioni competitive con i reward basati su VLM all'avanguardia (ad es. HPSv3, UnifiedReward) in media, con un divario notevolmente ridotto su dataset specifici come HPDv2 e HPDv3.
  • Ensemble di Rumore: L'applicazione dell'ensemble di rumore migliora ulteriormente l'accuratezza media (ad es. dal 71,49% al 72,48% in media), validando il beneficio dell'aggregazione delle evidenze attraverso i livelli di rumore.

• Allineamento delle Preferenze (ReFL):

  • Negli esperimenti di allineamento post-addestramento utilizzando ReFL (Reward Feedback Learning) su SD3.5-M, DiNa-LRM ha dimostrato una convergenza più rapida nei punteggi proxy e metriche "golden" migliorate (PickScore su dati tenuti fuori) rispetto a HPSv3.
  • Efficienza: DiNa-LRM riduce significativamente i costi delle risorse rispetto ai reward basati su VLM. Riduce l'utilizzo di VRAM di picco del 51,4%, i FLOPs di calcolo del reward del 71,1% e i FLOPs della fase di ottimizzazione del 46,4%.

• Generalizzazione:

  • La formulazione generalizza su diversi backbone di diffusione (SDXL, FLUX.1-Dev, Z-Image-Turbo), superando costantemente le precedenti baseline basate su diffusione.
  • Gli esperimenti di trasferibilità mostrano che un modello di reward addestrato su SD3.5-M può allineare efficacemente SD3.5-L all'interno dello stesso spazio latente.

5. Significato e Affermazioni

Il documento afferma che DiNa-LRM offre un'alternativa pratica ed efficiente ai reward VLM nello spazio dei pixel per l'allineamento delle preferenze. Il suo significato principale risiede in:

• Vantaggio Strutturale: Operando nativamente nello spazio latente, elimina il disallineamento da latente a pixel e il sovraccarico computazionale della decodifica, rendendolo altamente compatibile con i metodi di allineamento basati sul gradiente del reward.
• Ottimizzazione-Friendly: Fornisce un segnale differenziabile ed efficiente che migliora la dinamica dell'addestramento e riduce i requisiti di risorse senza sacrificare la qualità dell'allineamento.
• Scalabilità: Il meccanismo di ensemble di rumore fornisce un modo nuovo e nativo alla diffusione per scalare le prestazioni al momento dell'inferenza.

Gli autori riconoscono le limitazioni, notando che il reward è legato a spazi latenti specifici (limitando il trasferimento tra backbone diversi senza re-addestramento) e che la modellazione nello spazio latente potrebbe sottovalutare determinati artefatti a livello di pixel, portando potenzialmente a comportamenti di hacking del reward (ad es. iniezione umana spuria o deriva stilistica) durante l'ottimizzazione a lungo termine. Suggeriscono che il lavoro futuro dovrebbe concentrarsi su backbone unificati più robusti e sull'incorporazione di regolarizzazione leggera nello spazio dei pixel.