Sintesi Tecnica: Sparse Context – Generazione Condizionata da Riferimento Efficiente tramite Token Dropping

Definizione del Problema

I modelli di diffusione basati su riferimenti sono diventati fondamentali per la generazione di immagini controllabile, abilitando compiti come l'editing di immagini e la sintesi guidata da soggetti, condizionando la generazione su immagini di input. Tuttavia, questi modelli affrontano significativi colli di bottiglia computazionali. Nei moderni modelli di diffusione basati su transformer (DiT), le immagini di riferimento vengono tipicamente codificate in griglie spaziali dense di token tramite un VAE e concatenate con i token rumorosi di input. Poiché le operazioni di self-attention scalano quadraticamente con il numero di token, l'inclusione di una o più immagini di riferimento ad alta risoluzione aumenta drasticamente i tempi di inferenza e l'uso della memoria. Questo costo diventa proibitivo nelle applicazioni reali che coinvolgono molteplici riferimenti. Sebbene esistano tecniche di efficienza per la generazione guidata da prompt, esse rimangono poco esplorate per i modelli condizionati da riferimenti, che presentano sfide uniche a causa della rappresentazione densa degli input visivi.

Metodologia

Il documento introduce Sparse Context, un framework progettato per costruire rappresentazioni di riferimento sparse mantenendo solo un sottoinsieme ridotto di token di riferimento. La metodologia affronta il problema attraverso tre componenti principali:

1. Osservazione della Ridondanza

Gli autori dimostrano innanzitutto che i token di riferimento contengono una significativa ridondanza. Attraverso un esperimento di "Naïve Context Pruning" su un modello pre-addestrato FLUX.2-Klien, mostrano che l'eliminazione casuale fino all'80% dei token di riferimento durante l'inferenza preserva ampiamente il layout grossolano della scena, nonostante una perdita nei dettagli fini. Ciò suggerisce che i modelli condizionati da riferimenti sono intrinsecamente robusti agli input sparsi, ma il divario di prestazioni tra riferimenti completi e parziali è esacerbato da un disallineamento tra la distribuzione di addestramento e quella di test.

2. Strategia di Addestramento: Random Token Dropping

Per colmare il divario di distribuzione e abilitare una generazione sparsa robusta, gli autori affinano un generatore condizionato da riferimenti allo stato dell'arte utilizzando il random token dropping durante l'addestramento.

• Processo: Durante l'addestramento, una frazione casuale di token di riferimento viene eliminata (mantenendo una frazione $f \in [0.05, 0.25]$) prima della concatenazione con i token target rumorosi.
• Obiettivo: Questo incoraggia il modello a imparare rappresentazioni robuste che funzionino efficacemente con informazioni di riferimento parziali, scindendo il modello da qualsiasi regola specifica di selezione dei token.
• Implementazione: Un leggero LoRA (Low-Rank Adaptation) viene applicato al modello base FLUX-2-Klein-9B su un dataset curato che copre l'editing guidato da istruzioni e la personalizzazione.

3. Strategia di Inferenza: Selezione dei Token Consapevole del Task

Mentre l'addestramento utilizza l'eliminazione casuale per garantire la robustezza, l'inferenza impiega strategie di selezione dei token consapevoli del task per massimizzare la qualità entro un budget computazionale fisso:

• Editing di Immagini Guidato da Istruzioni: Per preservare la struttura spaziale, il metodo utilizza una strategia di campionamento consapevole dei bordi (edge-aware). Viene estratta una mappa dei bordi Canny dall'immagine di riferimento, e i token vengono campionati in base a questa distribuzione di probabilità, dando priorità ai confini strutturali rispetto alle regioni uniformi.
• Personalizzazione: Per preservare l'identità del soggetto, il metodo utilizza un campionamento basato sulla salienza (saliency-based). Una mappa di salienza approssima la maschera del soggetto, permettendo al modello di concentrare i token sulle regioni più informative dell'immagine di riferimento.

Contributi Chiave

1. Framework Sparse Context: Un approccio innovativo che costruisce rappresentazioni di riferimento sparse eliminando i token, riducendo significativamente il costo quadratico della self-attention senza modificare l'architettura sottostante del modello.
2. Strategia di Fine-Tuning Robusta: Un regime di addestramento che utilizza il random token dropping che consente al modello di generalizzare attraverso diversi livelli di sparsità, permettendo un controllo flessibile sul budget computazionale all'inferenza.
3. Inferenza Consapevole del Task: La proposta di specifici euristici di selezione dei token (basati sui bordi per l'editing, basati sulla salienza per la personalizzazione) che superano il random dropping dando priorità alle informazioni critiche per il task.
4. Ortogonalità rispetto ai Metodi Esistenti: Dimostrazione che Sparse Context è compatibile e complementare ad altre tecniche di efficienza, come il Token Merging (ToMe) e il KV-caching, consentendo guadagni di velocità composti.

Risultati Sperimentali

Il metodo è stato valutato sull'editing di immagini guidato da istruzioni (PIE-Bench) e sulla personalizzazione basata su riferimenti (Subject-200K, CustomDiffusion-105) utilizzando il modello FLUX-2-Klein-9B.

• Efficienza:
  • Singolo Riferimento: Ottiene un'accelerazione dell'inferenza fino a 2× (ad esempio, 1.83× con il 10% di ritenzione dei token).
  • Multi-Riferimento: Ottiene un'accelerazione fino a 4× per scenari con molteplici riferimenti (ad esempio, 6 riferimenti), poiché i token di riferimento dominano il calcolo dell'attenzione.
  • Efficienza Combinata: Integrando Sparse Context con il KV-caching per la generazione multi-riferimento, il metodo raggiunge accelerazioni fino a 6.72×.
• Qualità:
  • Editing di Immagini: Il metodo preserva la struttura della scena e l'identità significativamente meglio della baseline di Naïve Context Pruning. Le metriche quantitative (LPIPS, DINO, CLIP-I) mostrano che Sparse Context mantiene un'alta fedeltà anche a basse frazioni di token (ad esempio, $f=0.1$), mentre la baseline soffre di corruzione dell'identità.
  • Personalizzazione: Il campionamento basato sulla salienza preserva con successo i dettagli fini del soggetto (ad esempio, colori specifici, pattern) che vengono persi con il random dropping.
• Integrazione: Combinare Sparse Context con il Token Merging migliora ulteriormente le prestazioni, dimostrando che il modello affinato sfrutta efficacemente altre tecniche di riduzione dei token.

Significato e Rivendicazioni

Il documento afferma che Sparse Context sblocca il potenziale di efficienza inespresso intrinseco nei modelli di diffusione basati su riferimenti. Riconoscendo che le griglie dense di token sono spesso non necessarie, gli autori dimostrano che la generazione basata su riferimenti può essere resa significativamente più efficiente senza compromettere la controllabilità o la fedeltà visiva.

Gli autori sottolineano che il loro approccio è semplice e modulare, richiedendo solo un fine-tuning leggero e offrendo un controllo flessibile del budget computazionale all'inferenza. Essi asseriscono che questo metodo affronta un collo di bottiglia critico nello spiegamento di modelli generativi ad alta risoluzione e multi-riferimento, rendendoli più pratici per le applicazioni del mondo reale. Il lavoro evidenzia come la ridondanza nelle rappresentazioni di riferimento sia una proprietà fondamentale che può essere sfruttata per ridurre i costi computazionali, una scoperta che si estende oltre i modelli specifici testati alla classe più ampia di sistemi di generazione di immagini e video basati su DiT.