Test-Time Alignment of Text-to-Image Diffusion Models via Null-Text Embedding Optimisation

Sintesi Tecnica: Allineamento del Modello di Diffusione Testo-Immagine tramite Null-Text al Tempo di Test (Null-Text Test-Time Alignment)

Problematica
I modelli di diffusione testo-immagine pre-addestrati sono spesso addestrati su dataset web su larga scala e non curati, che contengono contenuti disallineati o indesiderabili. Allineare questi modelli a specifici valori umani o obiettivi target (ad esempio, qualità estetica, preferenze specifiche) durante l'inferenza è fondamentale per un dispiegamento affidabile. Gli approcci di allineamento esistenti rientrano generalmente in due categorie, entrambe affette da limitazioni significative:

1. Metodi di Fine-tuning: Questi aggiornano i parametri del modello per ottimizzare una funzione di ricompensa (reward). Sebbene efficaci, sono computazionalmente costosi e soggetti all'over-optimization della ricompensa (reward hacking), dove il modello si adatta eccessivamente a ricompense proxy, perdendo la generalizzazione su altre ricompense e sulla diversità dell'output.
2. Metodi di Allineamento al Tempo di Test (TTA): Questi ottimizzano variabili latenti o di rumore, campionano da posteriori condizionati alla ricompensa o utilizzano algoritmi di ricerca discreta senza aggiornare i pesi del modello. Tuttavia, questi metodi soffrono spesso di sotto-ottimizzazione o over-optimization della ricompensa a causa della natura altamente non strutturata degli spazi latenti e di rumore. L'ottimizzazione in questi spazi può portare al "reward hacking", in cui il modello sfrutta pattern di rumore non semantici per massimizzare la ricompensa target a scapito della qualità visiva, della coerenza semantica e delle prestazioni su altri parametri.

Metodologia: Null-Text Test-Time Alignment (Null-TTA)
Gli autori propongono Null-TTA, un framework training-free che allinea i modelli di diffusione ottimizzando l'embedding incondizionale (null-text) all'interno della Classifier-Free Guidance (CFG), anziché manipolare variabili latenti o di rumore.

• Intuizione Centrale: Nella CFG, l'embedding incondizionale ($\phi$) funge da ancora geometrica per la distribuzione generativa del modello. A differenza degli spazi di rumore non strutturati, lo spazio degli embedding testuali possiede una natura semantica strutturata. L'ottimizzazione all'interno di questo spazio garantisce che l'allineamento avvenga su un manifold semanticamente coerente, impedendo al modello di sfruttare artefatti di rumore non semantici.
• Funzione Obiettivo: Il metodo formula un obiettivo regolarizzato che guida la distribuzione generativa del modello verso una ricompensa target, penalizzando al contempo la deriva distributiva dal modello pre-addestrato. L'obiettivo massimizza la ricompensa attesa $R(x_0)$ minimizzando la divergenza KL tra la distribuzione indotta dall'embedding ottimizzato $\phi'$ e la distribuzione originale indotta da $\phi$:
  $$\max_{\phi'} \left( \lambda_1 \mathbb{E}_{p(x_0|\phi')}[R(x_0)] - \lambda_2 \text{KL}(p(x_{0:T}, \phi') || p(x_{0:T}, \phi)) \right)$$
  Il termine KL è decomposto in divergenze locali tra passi di denoising consecutivi e una divergenza tra le distribuzioni degli embedding. Questo viene semplificato in un'espressione in forma chiusa che coinvolge la differenza nel rumore predetto ($\tilde{\epsilon}$) e la differenza negli embedding ($\phi' - \phi$).
• Strategia di Ottimizzazione:
  • Regolarizzazione Adattiva: Il peso di regolarizzazione $\lambda_2$ viene attenuato in modo adattivo man mano che il processo di denoising procede ($t \to 0$). Una forte regolarizzazione nei primi passi stabilizza l'ottimizzazione, mentre una regolarizzazione più debole nei passi successivi consente un allineamento più fine.
  • Ricerca Greedy: Ad ogni passo di denoising, il metodo impiega una ricerca greedy leggera. Campiona $K$ candidati latenti dal kernel di transizione, stima il campione pulito tramite la formula di Tweedie, li valuta con il modello di ricompensa e seleziona deterministicamente il candidato con il punteggio più alto per il passo successivo.
  • Ricompense Non-Differenziabili: Per scenari in cui la funzione di ricompensa è una black box (non differenziabile), il metodo impiega la stima del gradiente di ordine zero tramite perturbazioni casuali dell'embedding null-text.

Contributi Chiave

1. Framework Null-TTA: Un nuovo framework TTA training-free che esegue l'allineamento nello spazio semantico strutturato del condizionamento testuale, evitando le insidie dell'ottimizzazione di variabili latenti/rumore non strutturate.
2. Obiettivo Principato: Un obiettivo regolarizzato tramite KL che riorienta la distribuzione generativa del modello stesso verso la ricompensa target, garantendo un allineamento stabile ed efficiente senza aggiornamenti dei parametri. Ciò previene il reward hacking e mantiene la coerenza semantica.
3. Prestazioni allo Stato dell'Arte: Esperimenti estesi dimostrano che Null-TTA raggiunge una superiore performance della ricompensa target mantenendo una forte generalizzazione tra le diverse ricompense e la qualità percettiva.

Risultati Sperimentali
Gli autori hanno valutato Null-TTA su Stable Diffusion v1.5 e SDXL utilizzando varie ricompense target (HPSv2, PickScore, Aesthetic Score, ImageReward) e obiettivi non differenziabili (compressibilità JPEG).

• Ottimizzazione della Ricompensa vs. Generalizzazione: Null-TTA sposta costantemente la frontiera di Pareto verso l'esterno rispetto ai baseline (inclusi DNO, DAS, DSearch e TDPO). Raggiunge ricompense target più elevate degradando le ricompense non considerate significativamente meno rispetto ai competitor, indicando un allineamento riuscito senza over-optimization.
• Performance Qualitative: In categorie impegnative (conteggio, composizionalità, ragionamento spaziale, colori insoliti, trasferimento di stile, scene impossibili), Null-TTA produce immagini che soddisfano più fedelmente i vincoli del prompt preservando la coerenza globale, evitando modalità di fallimento come conteggi errati di oggetti o relazioni intrecciate.
• Studio Utente: Uno studio cieco condotto da utenti con 16 partecipanti ha mostrato una preferenza statisticamente significativa per Null-TTA rispetto ai baseline, confermando che i guadagni in termini di ricompensa si traducono in qualità percettiva.
• Efficienza Computazionale: Null-TTA raggiunge una performance della ricompensa target più forte con un minor utilizzo di memoria GPU e tempi di esecuzione (wall-clock time) comparabili o più rapidi rispetto ai baseline. Questa efficienza deriva dal backpropagate solo attraverso i livelli di cross-attention (aggiornando l'embedding) piuttosto che l'intero U-Net.
• Ricompense Non-Differenziabili: Il metodo ha ottimizzato con successo la dimensione del file JPEG (un obiettivo black-box) utilizzando gradienti di ordine zero, ottenendo tassi di compressione competitivi pur preservando meglio la qualità visiva rispetto ai baseline basati su gradiente.

Significatività e Rivendicazioni
Il paper stabilisce l'ottimizzazione dello spazio semantico come un paradigma nuovo, efficace e principato per il Test-Time Alignment. Sfruttando la natura strutturata dello spazio degli embedding testuali, Null-TTA fornisce una regolarizzazione intrinseca che previene il reward hacking e mantiene la conoscenza pre-addestrata del modello. Gli autori affermano che questo approccio offre un'alternativa semplice ma potente agli esistenti metodi TTA, consentendo un allineamento all'inferenza affidabile, efficiente e generalizzabile dei modelli di diffusione senza il costo computazionale del fine-tuning.