OmniAlpha: Aligning Transparency-Aware Generation via Multi-Task Unified Reinforcement Learning

1. Enunciato del Problema

I modelli di intelligenza artificiale generativa attuali, in particolare i sistemi basati sulla diffusione, eccellono nella creazione di immagini RGB ad alta fedeltà. Tuttavia, molti flussi di lavoro visivi reali (ad esempio, effetti visivi, design grafico, modifica di immagini) richiedono una generazione consapevole della trasparenza (RGBA), che implica la modellazione non solo del colore, ma anche dell'opacità per pixel (canali alfa) e della composizione tra livelli.

Le soluzioni esistenti soffrono di due limitazioni principali:

• Frammentazione: I metodi attuali sono silos specifici per task. Esistono pipeline separate per il ritaglio (image matting), la rimozione di oggetti, la decomposizione in livelli e la generazione da testo a immagine. Non esiste un modello unificato in grado di gestire flussi di lavoro RGBA diversificati.
• Limitazioni del Fine-Tuning Supervisionato (SFT): Sebbene scalare l'SFT attraverso i task sia un passo logico verso l'unificazione, si basa su obiettivi di regressione localizzati (ad esempio, perdita L1/L2 a livello di pixel). Questi obiettivi non riescono a ottimizzare direttamente proprietà critiche della generazione consapevole della trasparenza, come:
  • Coerenza e fisica tra i livelli.
  • Precisione ai confini dell'alfa (ad esempio, capelli, pelliccia).
  • Fedeltà compositiva quando i livelli vengono mescolati.
  • Di conseguenza, i modelli basati solo sull'SFT spesso raggiungono un plateau, faticando a catturare i dettagli strutturali fini necessari per una trasparenza senza soluzione di continuità.

2. Metodologia: OMNIALPHA

Gli autori propongono OMNIALPHA, un framework unificato multi-task che sposta l'attenzione dall'apprendimento puramente supervisionato all'allineamento tramite Apprendimento per Rinforzo (RL). L'architettura è composta da tre componenti fondamentali:

A. Architettura Unificata

• VAE Consapevole dell'Alfa: Un Variational Autoencoder end-to-end inizializzato da un VAE RGB pre-addestrato utilizzando una strategia di "inizializzazione opaca". Il codificatore azzera inizialmente i pesi dell'alfa per concentrarsi sull'RGB, mentre il decodificatore è inizializzato per prevedere un canale alfa completamente opaco ($\alpha=1.0$). Ciò consente un'adattamento efficiente allo spazio RGBA a 4 canali.
• Trasformatore di Diffusione Sequenza-a-Sequenza (DiT): Il backbone di denoising elabora simultaneamente più input e output.
  • Asse dei Livelli Bidirezionale: Per gestire più immagini (input e target) in un singolo passaggio in avanti, gli autori estendono i standard 2D Rotary Positional Embeddings (RoPE) con un asse z bidirezionale.
  • I latenti di input ricevono indici non negativi ($0$a$n-1$), i latenti di target ricevono indici negativi ($-1$a$-m$) e gli embedding di testo ricevono indici positivi distinti. Ciò permette al modello di differenziare i livelli spazialmente all'interno della sequenza.

B. Paradigma di Addestramento

L'addestramento segue un processo in due fasi:

1. Avvio a Freddo (SFT Multi-Task): Il modello viene prima addestrato su un dataset unificato che copre cinque categorie di task utilizzando obiettivi di Flow Matching. Ciò stabilisce una capacità di base attraverso task diversificati.
2. Post-Training (Ottimizzazione della Politica Relativa di Gruppo - GRPO): Per superare le limitazioni dell'SFT, gli autori formulano il processo di generazione come un Processo Decisionale di Markov (MDP).
   • Stato: Condizioni multimodali, passo temporale e sequenza latente intermedia.
   • Azione: Il passo di denoising guidato dal campo di velocità previsto.
   • Ricompensa: Una Ricompensa Consapevole dei Livelli viene calcolata solo alla fine della traiettoria (basata sul risultato).

C. Design della Ricompensa Consapevole dei Livelli

Invece di ricompense specifiche per task, OMNIALPHA utilizza quattro obiettivi universali definiti sulla sequenza RGBA decodificata per guidare la politica RL:

1. Fedeltà del Livello ($r_{layer}$): Misura la similarità RGB (SSIM/PSNR) su tutti i livelli validi.
2. Fedeltà Compositiva ($r_{comp}$): Valuta l'immagine composita globale rispetto alla verità fondamentale (ground truth).
3. Struttura di Sfondo ($r_{bg}$): Penalizza le deviazioni strutturali nei livelli di sfondo utilizzando LPIPS.
4. Confine in Primo Piano ($r_{fg}$): Impone precisione e richiamo sui confini dell'alfa (cruciali per capelli/pelliccia) utilizzando maschere di confine morbide.

Queste ricompense vengono aggregate dinamicamente in base al task specifico (ad esempio, la generazione da testo a immagine omette le ricompense allineate ai pixel poiché la verità fondamentale non è disponibile, mentre i task di ritaglio attribuiscono un peso elevato alle ricompense di confine).

3. Contributi Chiave

1. Primo Framework RL Unificato per RGBA: OMNIALPHA è il primo framework ad applicare l'apprendimento per rinforzo multi-task per unificare la generazione e la manipolazione consapevoli della trasparenza, superando le pipeline frammentate e specifiche per task.
2. Paradigma di Ricompensa Strutturale: Gli autori introducono ricompense basate sulla struttura che ottimizzano esplicitamente la coerenza tra i livelli e la precisione dei confini dell'alfa, superando con successo il limite di prestazioni dell'SFT standard.
3. Codifica dei Livelli Bidirezionale: Un'estensione innovativa degli embedding posizionali che permette a un singolo DiT di modellare congiuntamente più livelli RGBA di input e output in un singolo passaggio in avanti.

4. Risultati Sperimentali

Il modello è stato valutato su cinque categorie di task: da Testo a Immagine, Rimozione di Oggetti, Ritaglio Automatico, Ritaglio con Riferimento e Decomposizione in Livelli.

• Prestazioni Complessive: OMNIALPHA supera costantemente il proprio baseline SFT e i modelli esperti specializzati su benchmark diversificati.
• Metriche Chiave:
  • Decomposizione in Livelli: Ha raggiunto una riduzione relativa del 9,07% nell'L1 RGB rispetto al baseline SFT.
  • Ritaglio Automatico: Ha superato gli strumenti convenzionali con un miglioramento del 74% nel SAD (Somma delle Differenze Assolute) e un miglioramento del 68% nel Grad (Differenza del Gradiente).
  • Ritaglio con Riferimento: Ha raggiunto risultati all'avanguardia sul benchmark RefMatte-RW100 (SAD: 14,768), battendo significativamente modelli specializzati come CLIPMat.
  • Rimozione di Oggetti: Ha dimostrato una forte generalizzazione su dataset fuori dominio (OBER-Wild), ottenendo punteggi ReMOVE comparabili al modello altamente specializzato ObjectClear.
• Studi di Ablazione: La rimozione di qualsiasi singolo componente della ricompensa ($r_{layer}, r_{comp}, r_{bg}, r_{fg}$) ha portato a un degrado delle prestazioni, confermando che l'intera suite di ricompense è necessaria per un allineamento multi-task robusto.

5. Significato

OMNIALPHA rappresenta un significativo cambiamento di paradigma nell'IA generativa per la modifica visiva. Dimostrando che l'Apprendimento per Rinforzo può allineare efficacemente i modelli di diffusione per task RGBA complessi e multilivello, il documento colma il divario tra la generazione RGB opaca e le esigenze sfumate dei flussi di lavoro visivi professionali.

Suggerisce che i futuri modelli fondazionali per la modifica di immagini dovrebbero allontanarsi dal fine-tuning isolato e specifico per task verso politiche unificate ottimizzate tramite ricompense basate sul risultato. Questo approccio non solo migliora la fedeltà dei canali alfa e della composizione dei livelli, ma permette anche a un singolo modello di gestire un ampio spettro di task creativi, dalla generazione di oggetti trasparenti alla decomposizione di scene complesse.