Sintesi Tecnica: Modulazione Consapevole della Qualità per i Diffusion Transformer

Definizione del Problema

I moderni modelli di diffusione text-to-image, in particolare i Diffusion Transformer (DiT) come Stable Diffusion 3.5 (SD3.5), si affidano a segnali di condizionamento — specificamente embedding del timestep e embedding del prompt aggregati (pooled) — per modulare il processo di denoising tramite l'Adaptive Layer Normalization (AdaLN). Sebbene questi segnali comunichino efficacemente il livello di rumore attuale e l'intento semantico, essi mancano di informazioni sulla qualità (quality-aware). Di conseguenza, il meccanismo di modulazione non può reagire direttamente ai fallimenti specifici del campione, come l'accumulo di errori, l'interpretazione errata del prompt o le caratteristiche allucinate che emergono mentre l'immagine latente evolve. Il condizionamento di base è indipendente dallo stato latente corrente ($\partial M_{base} / \partial x_t = 0$), impedendo al modello di correggere i problemi di fedeltà basandosi sulla qualità effettiva dell'immagine intermedia.

Metodologia

Gli autori propongono il Quality Representation Module (QRM), un adapter leggero basato su transformer progettato per iniettare segnali consapevoli della qualità nel processo di diffusione senza modificare il backbone pre-addestrato.

Architettura

QRM è implementato come un decoder transformer che predice offset di modulazione additivi ($M_{qrm}$) ai parametri AdaLN di base ($M_{base}$).

• Input: Il modulo riceve l'immagine latente corrente ($x_t$), l'embedding del prompt aggregato ($p$), l'embedding del timestep ($t$) e i parametri di modulazione di base ($M_{base}$).
• Elaborazione:
  1. L'immagine latente $x_t$ viene proiettata in token latenti tramite una convoluzione $1\times1$ appresa e un pooling adattivo.
  2. I parametri di modulazione di base vengono normalizzati e proiettati per formare i token di query ($Z_0$).
  3. Una memoria key-value ($KV$) viene costruita dai token latenti, dal prompt e dagli embedding del timestep.
  4. Attraverso strati di cross-attention e self-attention, i token di query attendono la memoria $KV$ per generare una rappresentazione raffinata.
  5. L'output viene proiettato nuovamente alle ampiezze native di AdaLN per produrre il vettore di offset $M_{qrm}$.
• Applicazione: La modulazione finale applicata ai blocchi del transformer è $M_t = M_{base,t} + M_{qrm,t}$.

Strategia di Addestramento (ReFL)

QRM è addestrato utilizzando una strategia di Reward Feedback Learning (ReFL):

• Componenti Congelati: Il backbone della diffusione (denoiser, encoder testuali, VAE), il sampler e il modello di reward rimangono congelati. Solo i parametri di QRM vengono aggiornati.
• Obiettivo: Il modulo impara a massimizzare un punteggio di reward scalare derivato da un modello di reward congelato (ad es., HPSv2.1 o CLIPScore) applicato a un'immagine decodificata da uno step di denoising intermedio.
• Funzioni di Perdita: Il paper utilizza una perdita di massimizzazione del reward o una perdita di ricerca del margine (margin-seeking loss), che spinge l'immagine modulata da QRM a superare il baseline di un margine specifico.
• Selezione dei Timestep: Per migliorare l'efficienza e concentrarsi sulla formazione strutturale, QRM viene attivato solo durante un sottoinsieme predefinito di timestep ad alto rumore (ad es., $t \in \{50, \dots, 26\}$), piuttosto che durante l'intero schema di diffusione.

Inferenza

Durante l'inferenza, il modulo QRM addestrato agisce come un adapter plug-in. Predice gli offset di modulazione basandosi sullo stato latente corrente e li aggiunge ai parametri di base. Ciò consente al modello di guidare la traiettoria di denoising verso stati di qualità superiore senza riaddestrare il backbone o alterare lo schema di campionamento.

Contributi Chiave

1. Architettura QRM: Introduzione di un modulo transformer leggero che apprende rappresentazioni consapevoli della qualità dagli input esistenti del modello (latente, prompt, timestep) per raffinare i parametri di modulazione.
2. Integrazione Non Invasiva: Un meccanismo per aumentare la modulazione del diffusion transformer con segnali di qualità dell'immagine latente senza modificare il backbone di diffusione, il VAE o lo schema di campionamento sottostante.
3. Modulazione Guidata dal Reward: Dimostrazione che gli aggiornamenti di modulazione predetti ($M_{qrm}$) possono effettivamente raffinare il processo di denoising ad ogni timestep, migliorando l'aderenza al prompt e la fedeltà dell'immagine.
4. Valutazione Comprensiva: Studi di ablazione estesi che analizzano gli effetti della dimensione dell'architettura QRM, degli intervalli di timestep attivi e delle metriche di reward dell'addestramento sulla qualità dell'immagine.

Risultati Sperimentali

Gli autori hanno valutato QRM utilizzando Stable Diffusion 3.5 (SD3.5) sul benchmark Parti-Prompts.

• Performance Quantitative:
  • Rispetto al baseline di SD3.5, QRM addestrato con HPSv2.1 come segnale di reward ha ottenuto i punteggi più alti tra tutti i metodi valutati (inclusi prompt refinement, fine-tuning e approcci ReNeg).
  • Miglioramenti: +0.20 nello score Aesthetic e +0.82 nello score HPSv2.1.
  • Quando addestrato con CLIPScore, QRM ha migliorato l'allineamento con la preferenza umana (HPSv2.1) mantenendo una qualità percettiva comparabile.
• Analisi Qualitativa:
  • Le visualizzazioni delle traiettorie di denoising mostrano che l'applicazione di QRM durante i primi timestep ad alto rumore porta alla comparsa di strutture semantiche riconoscibili significativamente prima rispetto al baseline.
  • L'applicazione di QRM solo nelle fasi intermedie o tardive ha prodotto miglioramenti minori, confermando che la struttura globale si stabilisce precocemente nel processo di diffusione.
• Risultati delle Ablazioni:
  • Capacità: Modelli QRM più grandi ottimizzati più fortemente per il reward specifico di addestramento (ad es., HPSv2.1) hanno mostrato però lievi compromessi in altre metriche (ad es., CLIPScore).
  • Intervallo di Timestep: La modulazione durante i primi timestep ($50 \to 26$) è stata costantemente più efficace per l'allineamento rispetto alla modulazione nelle fasi successive.
  • Metrica di Reward: La scelta del modello di reward (CLIPScore vs. HPSv2.1) ha plasmato direttamente i bias percettivi appresi da QRM.

Significato e Rivendicazioni

Il paper sostiene che gli interventi a livello di modulazione forniscono una via scalabile ed efficace per allineare i grandi modelli di diffusione ai segnali di preferenza umana. A differenza dei metodi che richiedono un costoso fine-tuning dell'intero backbone (ad es., LoRA o fine-tuning completo), QRM ottiene miglioramenti costanti della qualità operando interamente attraverso aggiornamenti di modulazione additivi condizionati dalla rappresentazione latente in evoluzione.

Gli autori posizionano QRM come un adapter flessibile e consapevole della qualità che può essere guidato da diversi segnali di reward e applicato a varie architetture di diffusion transformer. Riconoscono che QRM eredita i bias dei modelli di reward utilizzati per l'addestramento e suggeriscono che il lavoro futuro potrebbe esplorare l'addestramento multi-reward e la selezione adattiva dei timestep per migliorare ulteriormente robustezza e generalizzazione. Il lavoro dimostra che guadagni significativi nella fedeltà dell'immagine e nell'allineamento al prompt possono essere ottenuti senza alterare i pesi pre-addestrati dei modelli allo stato dell'arte.