FOCUS: Optimal Control for Multi-Entity World Modeling in Text-to-Image Generation

1. Il Problema

I modelli di generazione di immagini da testo (Text-to-Image o T2I) hanno raggiunto livelli eccezionali di fedeltà visiva per prompt contenenti un singolo soggetto. Tuttavia, faticano notevolmente con scene a multi-soggetto. Le modalità di fallimento comuni includono:

• Perdita di attributi (Attribute Leakage): Gli attributi di un oggetto (es. "cappello rosso") vengono erroneamente assegnati a un altro soggetto.
• Intreccio di identità (Identity Entanglement): Le identità dei soggetti si mescolano, rendendo difficile distinguere chi possiede quale caratteristica.
• Omissione di soggetti: Alcuni elementi richiesti nel prompt non appaiono nell'immagine finale.

Questi errori derivano da una rappresentazione del mondo debole a livello di oggetti, dove i vincoli diventano non locali (modificare un oggetto influisce involontariamente sugli altri). Le soluzioni esistenti sono spesso euristiche, prive di un obiettivo di ottimizzazione unificato e difficili da trasferire tra diverse architetture di modelli moderni.

2. Metodologia

Gli autori propongono un quadro teorico unificante basato sulla Teoria del Controllo Stocastico Ottimo (Stochastic Optimal Control - SOC) applicato ai modelli di Flow Matching (FM).

A. Formulazione come Controllo Ottimo

Il processo di generazione viene reinterpretato come un problema di controllo stocastico. L'obiettivo è minimizzare una funzione di costo che bilancia la fedeltà al modello base e la disentanglement (separazione) dei soggetti.

• Equazione di stato: Il flusso latente $X_t$ evolve secondo un'equazione differenziale stocastica (SDE) controllata da un termine $u_t$.
• Funzione di costo: Si introduce un costo di esecuzione $f(X_t, t)$ che misura l'entanglement dei soggetti.
• Soluzione: Utilizzando le condizioni di ottimalità di Pontryagin, si deriva un controllo ottimo $u^*_t$ che dipende dal gradiente della funzione di costo rispetto allo stato latente.

B. Due Implementazioni Pratiche

Sulla base di questa formulazione, il paper propone due algoritmi architetturalmente agnostici:

1. Controllo a Tempo di Esecuzione (Test-Time Controller):

   • È un metodo senza riaddestramento (training-free).
   • Durante il campionamento, il vettore di velocità del modello base viene perturbato on-the-fly.
   • Poiché calcolare l'esatto stato co-adiunto (adjoint) richiederebbe una risoluzione all'indietro (impraticabile), l'algoritmo approssima il controllo in un singolo passaggio (single-pass) linearizzando il sistema attorno allo stato corrente.
   • La velocità controllata diventa: $v^*_t \approx v_{base} - \eta(t) \nabla_X f(X_t, t)$, dove $\eta(t)$ è un passo temporale derivato teoricamente.

2. Adjoint Matching (Fine-tuning):

   • È un metodo di fine-tuning leggero per apprendere una rete di controllo $u_\theta$.
   • Utilizza la tecnica Adjoint Matching per addestrare la rete a regressare il segnale di controllo verso un segnale co-adiunto "snello" (lean adjoint) calcolato su traiettorie congelate del modello base.
   • Questo permette di ottenere la disentanglement con un singolo passaggio in avanti durante l'inferenza, eliminando l'overhead computazionale del controllo a tempo di esecuzione.

C. La Funzione di Costo FOCUS

Il cuore del metodo è la definizione della funzione di costo $f$, chiamata FOCUS (Flow Optimal Control for Unentangled Subjects).

• Sfrutta le mappe di attenzione incrociata (cross-attention) del modello, che correlano i token di testo con le posizioni spaziali nell'immagine.
• Tratta le mappe di attenzione come distribuzioni di probabilità spaziali.
• Utilizza la Divergenza di Jensen-Shannon (JSD) per misurare l'entanglement:
  • Termine 1: Penalizza la dispersione entro lo stesso soggetto (promuove coerenza e binding).
  • Termine 2: Massimizza la divergenza tra i soggetti (promuove la separazione spaziale).
• Questo approccio probabilistico è scalabile e non richiede annotazioni spaziali manuali.

3. Contributi Chiave

1. FOCUS: Un obiettivo probabilistico basato sulla divergenza per la disentanglement multi-soggetto, che supera le euristiche precedenti.
2. Due istanze pratiche: Un controller a tempo di esecuzione (training-free) e un metodo di fine-tuning basato su Adjoint Matching, entrambi derivati dallo stesso principio teorico.
3. Prospettiva SOC: Fornisce un'interpretazione unificata delle tecniche di guida (guidance) esistenti, offrendo un unico "knob" (parametro di controllo) per bilanciare fedeltà e separazione, rendendo i metodi portabili su diverse architetture (es. SD 3.5, FLUX.1).

4. Risultati Sperimentali

Gli esperimenti sono stati condotti su Stable Diffusion 3.5 e FLUX.1, confrontando FOCUS con metodi basati sull'attenzione come Attend&Excite, CONFORM, e Divide&Bind.

• Metriche Automatiche: FOCUS ottiene costantemente i punteggi più alti o competitivi nelle metriche di allineamento immagine-testo (CLIP, SigLIP) e testo-testo (BLIP, Qwen2), nonché nelle metriche di preferenza umana (PickScore, ImageReward).
• Studio di Preferenza Umana: In un test con 50 partecipanti (2.000 confronti), FOCUS ha ottenuto il tasso di vittoria più alto e il punteggio Elo più elevato sia per SD 3.5 che per FLUX.1, superando tutte le altre euristiche sia nella modalità test-time che nel fine-tuning.
• Efficienza:
  • Il controllo a tempo di esecuzione aumenta la latenza di circa il 90-138% ma funziona su GPU consumer.
  • Il modello fine-tuned non aggiunge alcun overhead all'inferenza e generalizza bene a prompt mai visti durante l'addestramento.
• Generalizzazione: Il metodo è stato validato anche su Stable Diffusion 1.5 (basato su U-Net), dimostrando la sua portabilità oltre i modelli Flow Matching moderni.

5. Significato e Impatto

Il lavoro di FOCUS rappresenta un passo avanti significativo nella generazione di immagini complesse:

• Unificazione Teorica: Trasforma approcci euristici frammentati in un quadro matematico coerente (Controllo Ottimo), spiegando perché e quando funzionano.
• Qualità Compositiva: Risolve in modo affidabile problemi critici come l'attribuzione errata degli attributi e l'omissione di oggetti, abilitando applicazioni pratiche come illustrazioni di storie, composizioni multi-pannello e comunicazione scientifica.
• Flessibilità: Offre agli utenti la scelta tra una soluzione immediata (test-time) e una soluzione ottimizzata per la produzione (fine-tuning), mantenendo lo stile e la diversità del modello base.

In sintesi, FOCUS dimostra che trattare la generazione di immagini multi-soggetto come un problema di controllo ottimo su spazi di attenzione probabilistici porta a risultati superiori, più stabili e teoricamente fondati rispetto alle tecniche attuali.