Riepilogo Tecnico: Domare i Token Outlier nei Diffusion Transformers

Enunciato del Problema

Il documento investiga il fenomeno dei token outlier all'interno dei Diffusion Transformers (DiT) utilizzati per la generazione di immagini. Mentre lavori precedenti (ad esempio, Darcet et al.) hanno stabilito che i Vision Transformers (ViT) nelle attività di riconoscimento producono spesso un piccolo sottoinsieme di token con norme di caratteristiche anormalmente elevate, che assorbono un'attenzione sproporzionata pur trasportando informazioni locali limitate, il ruolo di tali token nei modelli generativi rimaneva inesplorato.

Gli autori identificano due lacune critiche:

1. Ubiquità nelle Pipeline Generative: I token outlier appaiono non solo negli encoder ViT preaddestrati (tokenizzatori) delle moderne pipeline Representation Autoencoder (RAE), ma anche internamente all'interno degli stessi Diffusion Transformers (denoiser).
2. Natura del Difetto: Il degrado causato da questi token non è semplicemente il risultato di pochi valori di perdita estremi. Al contrario, gli autori ipotizzano che i token outlier siano sintomatici di semantica delle patch locali corrotta. Mascherare semplicemente i token ad alta norma durante l'addestramento non migliora le prestazioni, suggerendo che il problema risiede nella perdita della struttura spaziale e delle informazioni locali piuttosto che nella sola instabilità numerica.

Metodologia: Registri a Doppia Fase (DSR)

Per affrontare questi problemi, gli autori propongono i Registri a Doppia Fase (DSR), una strategia di intervento unificata che applica token di registro a entrambi i componenti della pipeline RAE-DiT: l'encoder di rappresentazione e il generatore di diffusione.

1. Registri negli Encoder Visivi

Gli autori affrontano i token outlier nella fase di encoder (il tokenizzatore) utilizzando due approcci a seconda della disponibilità di registri preaddestrati:

• Registri Addestrati: Per encoder come DINOv2 che supportano i registri, gli autori utilizzano token di registro preaddestrati per sopprimere gli artefatti ad alta norma nelle mappe di caratteristiche dell'encoder.
• Registri Ricorsivi al Tempo di Test (TTR): Per encoder privi di registri integrati (ad esempio, SigLIP2), gli autori aggiungono token non addestrati alla sequenza di input al momento dell'inferenza. Nei casi in cui viene rilevata una seconda fase di outlier nella rappresentazione risultante, applicano TTR in modo ricorsivo per stabilizzare l'output.

2. Registri nei Diffusion Transformers

Gli autori osservano che, anche con un encoder stabilizzato, il Transformer di diffusione stesso sviluppa token outlier interni, prevalentemente negli strati intermedi (a differenza dei ViT standard dove gli outlier appaiono spesso negli strati finali).

• Registri di Diffusione: Introducono un piccolo insieme di token di registro apprendibili direttamente nel Transformer di diffusione. Questi token vengono addestrati congiuntamente al generatore per assorbire i sink di attenzione e stabilizzare lo spazio di rappresentazione.
• Protocollo di Inferenza: Durante l'inferenza, gli output di tutti i token di registro (sia dall'encoder che dal modello di diffusione) vengono scartati, mantenendo solo gli output dei token immagine.

Contributi Chiave

• Scoperta di Outlier Interni: Il documento dimostra che i token outlier sono pervasivi nei DiT, apparendo sia nell'encoder che nel denoiser, con un pattern distinto di concentrazione negli strati intermedi del modello di diffusione.
• Rifiuto dell'Ipotesi "Valore Estremo": Attraverso esperimenti di mascheramento della perdita, gli autori mostrano che escludere semplicemente i token ad alta norma dalla funzione di perdita durante l'addestramento non migliora la qualità della generazione. Questo porta alla conclusione che il problema fondamentale è la corruzione della semantica delle patch locali, non la semplice presenza di valori estremi.
• Intervento a Doppia Fase: La proposta di DSR fornisce un meccanismo semplice e leggero per stabilizzare l'intera pipeline di generazione affrontando gli outlier alla fonte (encoder) e durante il processo di generazione (denoiser).
• Generalizzabilità: Il metodo si è dimostrato efficace in diversi spazi di input (spazio dei pixel, latenti VAE e spazi di rappresentazione) e su varie architetture (SiT, JiT, design basati su RAE).

Risultati Sperimentali

Gli autori valutano DSR sulla generazione condizionata per classe ImageNet-1K (256×256) e sulla generazione su larga scala da testo a immagine.

• Prestazioni su ImageNet-1K:
  • Utilizzando SigLIP2-B come encoder, DSR riduce il FID da 5.89 a 4.58 (con 800 epoche) e migliora i punteggi GenEval nelle attività da testo a immagine da 0.426 a 0.466.
  • Quando combinato con encoder DINOv2-B, il FID migliora da 4.16 a 3.95.
  • DSR accelera la convergenza, raggiungendo qualità comparabili con 4 volte meno epoche di addestramento rispetto alla baseline.
• Scalabilità: I miglioramenti sono coerenti tra le diverse dimensioni del modello (DiT-S, B, XL), con un aumento marginale del costo computazionale (aumento di circa il 10% dei GFLOPs).
• Studi di Ablazione:
  • Profondità di Inserimento: I registri sono più efficaci quando introdotti negli strati iniziali-mediali (ad esempio, a partire dal blocco 8). Introdurli troppo presto o troppo tardi produce rendimenti decrescenti.
  • Numero di Registri: Un numero moderato di registri (ad esempio, 36) produce prestazioni ottimali; un numero troppo basso (1) o troppo alto (100) degrada le prestazioni.
• Confronto con Alternative: DSR supera il "condizionamento in contesto" (che utilizza token dipendenti dall'input) man mano che la rappresentazione di input diventa più semantica, suggerendo che i registri forniscono un meccanismo più robusto per la capacità globale.

Significato e Affermazioni

Il documento afferma che il controllo dei token outlier è un componente critico, ma precedentemente trascurato, nella costruzione di pipeline di diffusione robuste basate su Transformer. Dimostrando che questi artefatti degradano la semantica delle patch locali e che gli interventi basati sui registri possono mitigarli efficacemente, gli autori posizionano DSR come un miglioramento fondamentale per la stabilità sia dell'encoder che del generatore.

Il lavoro evidenzia che la flessibilità dei Transformer, sebbene benefica per le interazioni globali, introduce artefatti specifici (token outlier) che ostacolano la fedeltà spaziale nella generazione. La soluzione proposta è presentata come un intervento semplice, plug-and-play, che migliora costantemente la qualità della generazione in contesti diversi senza richiedere cambiamenti architetturali complessi o il riaddestramento dell'intera pipeline da zero.