Sintesi Tecnica: UNIGP – Addomesticare il Diffusion Transformer per la Generazione e la Percezione Unificate con Preservazione dei Prior

Definizione del Problema

Gli approcci recenti nei modelli di diffusione hanno ottenuto un successo straordinario sia nella generazione controllabile di immagini (ad esempio, Text-to-Image guidato da profondità o bordi) sia nei compiti di predizione densa (ad esempio, stima della profondità monoculare e delle normali di superficie). Tuttavia, gli approcci esistenti trattano tipicamente questi come domini separati. I metodi di generazione controllabile (come ControlNet) sono limitati ai compiti di generazione e richiedono spesso dataset massicci per un controllo preciso, mentre i metodi di predizione densa (come Marigold o i modelli basati su regressione) sono spesso affinati partendo da backbone generativi, portando al "dimenticanza catastrofica" (catastrophic forgetting) dei versatili prior generativi del modello.

I tentativi di unificazione attuali (ad esempio, JointNet, UniCon) spesso non riescono a sfruttare appieno la relazione complementare tra generazione e percezione. Essi incorrono o in un enorme overhead computazionale duplicando l'intero backbone, o compromettono le capacità generative effettuando il fine-tuning diretto del backbone per la percezione. La sfida centrale è unificare questi compiti eterogenei all'interno di un singolo framework che preservi i prior generativi pre-addestrati pur abilitando una precisa predizione densa e una generazione controllabile, senza richiedere architetture o perdite (loss) specifiche per il compito.

Metodologia: Framework UNIGP

Gli autori propongono UNIGP, un framework costruito sulla base dell'architettura Multi-Modal Diffusion Transformer (MMDiT). Il metodo consiste in due innovazioni primarie: un ramo architettonico disgiunto e una strategia di addestramento unificata.

1. Generazione e Percezione Unificate Disgiunte (DUGP)

Seguendo il principio del rasoio di Occam, UNIGP evita di duplicare l'intero backbone MMDiT. Invece, introduce il ramo DUGP, che copia solo il ramo d'immagine ($\mathcal{F}_i$) del pre-addestrato MMDiT.

• Architettura: Il ramo DUGP è inizializzato copiando i parametri del ramo d'immagine. È diviso in un Condition Patcher (correttore di condizione), Stacked Control Layers (strati di controllo impilati) e Stacked Perception Layers (strati di percezione impilati).
• Preservazione dei Prior (Isolamento del Gradiente): Una scelta di design critica è il "bypass" della connessione additiva nei livelli di percezione. Nei livelli di controllo, l'output del ramo DUGP viene aggiunto all'output del backbone tramite uno strato lineare inizializzato a zero ($\mathcal{L}_{zero}$), permettendo alla condizione di guidare la generazione. Tuttavia, nei livelli di percezione, questo percorso additivo viene interrotto. Ciò assicura che i gradienti dalla perdita di predizione densa fluiscano esclusivamente attraverso il ramo DUGP, impedendo loro di alterare i pesi pre-addestrati del backbone. Questo isola efficacemente i prior generativi, permettendo al contempo al ramo DUGP di apprendere la percezione come un "lettore" non invasivo.
• Gestione Multi-Condizione: Per scenari che richiedono più condizioni, il condition patcher di DUGP viene duplicato per ogni condizione, con le interazioni gestite tramite concatenazione di sequenze per mantenere l'efficienza dell'inferenza.

2. Dataset Unificato e Strategia di Addestramento

Per colmare il divario tra generazione e percezione, UNIGP impiega una strategia di addestramento unificata su un dataset eterogeneo che combina dati di generazione filtrati e dati di percezione sintetici.

• Ponderazione Binaria della Loss: Invece di una complessa bilanciazione delle perdite, il modello utilizza un semplice meccanismo di routing binario. Per ogni campione in un batch misto, i pesi della loss ($\lambda_g, \lambda_p$) sono impostati a 1 o 0 in base alla sorgente dei dati. Se il campione proviene dal dataset di generazione, la loss di generazione è attiva; se proviene dal dataset di percezione, la loss di percezione è attiva.
• Obiettivo: La rete impara a predire il campo di velocità per i latenti RGB, profondità e normali simultaneamente, ottimizzando sia la qualità generativa che l'accuratezza geometrica all'interno di un singolo passaggio in avanti (forward pass).

Contributi Chiave

1. Framework UNIGP: Un modello Diffusion Transformer unificato che integra la generazione controllabile e la predizione densa. Utilizza un ramo DUGP che copia solo il ramo d'immagine di MMDiT, aderendo al rasoio di Occam per minimizzare l'overhead computazionale pur preservando i versatili prior del backbone.
2. Strategia di Addestramento Unificata: Un nuovo dataset e un approccio di addestramento che apprendono congiuntamente compiti eterogenei utilizzando uno schema semplice di ponderazione binaria della loss, eliminando la necessità di design architettonici specifici per il compito o di una complessa bilanciazione delle perdite.
3. Sinergia Dimostrata: Il documento fornisce prove che l'addestramento congiunto produce benefici complementari: i prior generativi arricchiscono i dettagli percettivi (catturando texture fini), mentre l'apprendimento percettivo impone un rigoroso allineamento spaziale nella generazione.

Risultati Sperimentali

Esperimenti estesi dimostrano che UNIGP supera i precedenti approcci unificati e performa in linea con i metodi "esperti" specializzati.

• Generazione Controllabile: UNIGP supera i modelli a controllo singolo (ControlNet) e multi-controllo (UniControlNet) in termini di FID e punteggi CLIP attraverso varie condizioni (profondità, normali, schizzo, posa, Canny). Notevolmente, raggiunge l'RMSE più basso nei compiti depth-to-image e normal-to-image, indicando un allineamento strutturale superiore.
• Predizione Densa: Nella stima zero-shot di profondità e normali di superficie, UNIGP raggiunge prestazioni comparabili ai modelli di regressione allo stato dell'arte (ad esempio, DepthAnything V2, DSINE) nonostante l'addestramento su molti meno immagini (59K contro 62.6M per DepthAnything). Supera significativamente altri metodi di percezione generativa (Marigold, GeoWizard) e modelli unificati (JointNet, UniCon).
• Generazione Congiunta: UNIGP supporta la generazione simultanea di mappe RGB, profondità e normali di superficie da prompt testuali, superando i baseline come LDM3D e JointNet sia nella qualità dell'immagine (FID) che nella coerenza geometrica (metriche GD/GN).
• Studi di Ablazione:
  • Sinergia: Rimuovere l'addestramento alla percezione degrada l'allineamento strutturale delle immagini generate (RMSE più alto), mentre rimuovere l'addestramento alla generazione risulta in mappe di profondità eccessivamente levigate e prive di dettagli fini.
  • Scelte di Design: Sostituire DUGP con una duplicazione completa del backbone (stile JointNet) o con un fine-tuning completo (stile Marigold) risulta in prestazioni significativamente inferiori, confermando che il flusso di gradiente disgiunto è la chiave del successo.
  • Bilanciamento dei Livelli: Un'allocazione uguale tra strati di controllo e di percezione impilati (12:12) produce l'equilibrio ottimale tra controllo generativo e accuratezza percettiva.

Significato e Rivendicazioni

Il documento sostiene che UNIGP riesce a colmare i campi separati della sintesi basata su diffusione e della predizione geometrica densa. Il suo significato primario risiede nel dimostrare che i prior generativi e i vincoli percettivi non sono mutuamente esclusivi ma si rinforzano a vicenda.

• Preservazione dei Prior: Isolando il flusso del gradiente, UNIGP mantiene le capacità generative di alta qualità del backbone MMDiT pre-addestrato, che vengono spesso perse nel fine-tuning specifico per la percezione.
• Efficienza: Il metodo raggiunge risultati allo stato dell'arte con un incremento di parametri relativamente piccolo (da 2B a 3B) rispetto alla duplicazione dell'intero backbone, rendendolo una soluzione unificata computazionalmente efficiente.
• Versatilità: Il framework supporta una vasta gamma di compiti all'interno di un singolo modello, inclusi text-to-image, stima di profondità/normali e generazione con qualsiasi condizione, offrendo un'alternativa più versatile al mantenimento di molteplici modelli specializzati.

Gli autori riconoscono che, sebbene UNIGP introduca un ulteriore overhead computazionale rispetto al modello base (aumentando il tempo di inferenza da ~5s a ~9.5s), questo compromesso è giustificato dalla capacità di sostituire molteplici architetture specializzate con un singolo modello unificato altamente versatile.