Riepilogo Tecnico: I Registri Sono Fondamentali per i Diffusion Transformers nello Spazio dei Pixel

Enunciato del Problema

È noto che i Vision Transformers (ViT) presentano valori anomali (outlier) nei token delle patch con norme elevate, che degradano la qualità delle mappe delle caratteristiche e generano artefatti nelle mappe di attenzione. Questo problema viene efficacemente mitigato nei ViT introducendo "token di registro", che agiscono come pozzi per tali valori anomali. Man mano che i Diffusion Transformers (DiT) adottano sempre più architetture basate su transformer e si spostano verso l'addestramento direttamente nello spazio dei pixel (allontanandosi dallo spazio latente), assomigliano strutturalmente ai ViT in misura maggiore. Ciò solleva due domande critiche:

1. I DiT ereditano i valori anomali nei token delle patch con norme elevate osservati nei ViT?
2. I token di registro sono vantaggiosi per i DiT, in particolare nei regimi di addestramento nello spazio dei pixel?

Sebbene i pozzi di attenzione e i comportamenti dei token speciali siano stati studiati nei Large Language Models (LLM) e nei modelli di diffusione linguistica, la loro presenza e il loro ruolo funzionale nei transformer di diffusione per immagini rimangono poco esplorati.

Metodologia

Gli autori conducono un'analisi sistematica dei token di registro attraverso diversi spazi di modelli di diffusione (spazio dei pixel, spazio VAE e spazio RAE) e li confrontano con ViT standard (in particolare DINOv2).

1. Analisi Empirica dei Valori Anomali: Lo studio investiga le norme delle caratteristiche a livello di token e le mappe di attenzione nei DiT (in particolare i pDiT addestrati con flow matching su ImageNet) rispetto ai ViT. Analizzano se i token delle patch presentano valori anomali con norme elevate in assenza di registri.
2. Integrazione dei Token di Registro: I token di registro sono implementati come token apprendibili aggiuntivi aggiunti alla sequenza dei token delle patch, seguendo l'approccio utilizzato nei ViT. Questi token sono esclusi dalla funzione obiettivo di addestramento (funzione di perdita).
3. Analisi delle Rappresentazioni Interne: Gli autori analizzano le mappe delle caratteristiche intermedie utilizzando la Variazione Totale (TV) per misurare la regolarità spaziale e l'analisi tramite sonda lineare (linear probing) per valutare il contenuto semantico dei token di registro. Esaminano come i registri influenzano le norme delle caratteristiche e la coerenza delle rappresentazioni a diversi istanti temporali di diffusione (livelli di rumore).
4. Esplorazione Architettonica: Motivati dai ruoli distinti dei registri e dei token delle patch, il documento propone un'architettura a doppio flusso efficiente in termini di parametri. Questo design elabora i token di registro e quelli delle patch attraverso percorsi specializzati (in particolare disaccoppiando i livelli RMSNorm, MLP e adaLN), condividendo al contempo i meccanismi di attenzione e la maggior parte dei parametri del transformer.
5. Studi di Ablazione: Lo studio esegue ablationi sul numero di registri, sulla profondità alla quale vengono introdotti e sui componenti specifici da disaccoppiare per determinare le configurazioni ottimali.

Risultati Chiave e Contributi

1. I DiT Differiscono Fondamentalmente dai ViT per quanto Riguarda i Valori Anomali

Contrariamente ai ViT, i DiT standard (sia nello spazio latente che in quello dei pixel) non presentano valori anomali nei token delle patch con norme elevate. Nei ViT, i valori anomali sono localizzati su token specifici (spesso in regioni a bassa informazione), mentre nei DiT le norme dei token delle patch rimangono relativamente uniformi e le mappe di attenzione mancano degli artefatti a bassa informazione osservati nei ViT.

2. I Registri Migliorano i DiT nello Spazio dei Pixel Nonostante l'Assenza di Valori Anomali Preesistenti

Sorprendentemente, l'aggiunta di token di registro ai DiT migliora significativamente la convergenza e la qualità della generazione, in particolare nei modelli nello spazio dei pixel.

• Meccanismo: A differenza dei ViT, dove i registri assorbono valori anomali preesistenti, nei DiT i registri creano al loro interno token con norme elevate.
• Effetto: La presenza di registri riduce costantemente le norme delle caratteristiche di tutti i token delle patch. Ciò porta a mappe delle caratteristiche intermedie più pulite e regolari, specialmente a livelli di rumore elevati (fasi iniziali della diffusione).
• Doppio Ruolo: I token di registro si specializzano in due ruoli distinti:
  • Pozzi di Norma: Token con norme estremamente elevate che assorbono la magnitudine dai token delle patch.
  • Vettori di Informazione Globale: Token con norme moderate che codificano informazioni semantiche diversificate sull'immagine (ad esempio, oggetti in primo piano, contesto di sfondo), verificate tramite sonda lineare e visualizzazione delle mappe di attenzione.

3. Efficacia Specifica dello Spazio

I benefici dei token di registro dipendono fortemente dallo spazio di addestramento:

• Spazio dei Pixel: I registri producono i miglioramenti più significativi nella qualità della generazione (FID). I modelli nello spazio dei pixel producono intrinsecamente rappresentazioni intermedie più rumorose con norme delle caratteristiche più elevate, rendendo cruciale l'effetto di regolarizzazione dei registri.
• Spazio Latente (VAE/RAE): I registri offrono guadagni moderati o, nel caso di modelli basati su RAE, degradano leggermente le prestazioni. Gli spazi latenti sono più strutturati e a dimensionalità inferiore, riducendo la necessità della forte regolarizzazione fornita dai registri.

4. Registri Impliciti nelle Architetture Esistenti

Gli autori osservano che le recenti architetture DiT nello spazio dei pixel (ad esempio JiT) che utilizzano condizionamento di classe nel contesto (aggiungendo embedding di classe alla sequenza di token) incorporano implicitamente meccanismi simili ai registri. Questi token nel contesto esibiscono un comportamento simile ai registri espliciti: alcuni agiscono come pozzi di norma, mentre altri codificano informazioni semantiche globali. Ciò suggerisce che i guadagni di prestazioni in tali modelli derivino in gran parte da questo comportamento simile ai registri piuttosto che dalle informazioni di classe stesse.

5. Architettura a Doppio Flusso Efficiente in Termini di Parametri

Per sfruttare queste intuizioni, gli autori propongono un'architettura compatta a doppio flusso che specializza l'elaborazione dei token di registro.

• Design: I registri vengono introdotti solo negli strati più profondi (a partire dal blocco 4), dove viene stabilita la struttura semantica. L'architettura utilizza un'attenzione condivisa ma disaccoppia RMSNorm, MLP e adaLN per registri e patch.
• Efficienza: Questo approccio aumenta le dimensioni del modello solo del ~14% (rispetto al ~77% di una duplicazione completa ingenua) mantenendo costi di inferenza simili (GFLOPs).
• Prestazioni: Questa architettura migliora costantemente la qualità della generazione attraverso diverse scale di modello (B, L, H) e risoluzioni (256x256, 512x512) rispetto alle baseline a flusso singolo.

Risultati

• Qualità della Generazione: Su ImageNet 256x256, i pDiT con registri superano costantemente quelli senza. Ad esempio, un modello pDiT-B/16 ha migliorato il suo FID da 7.39 a 5.30 a 200 epoche.
• Regolarità delle Caratteristiche: I modelli con registri mostrano rapporti di Variazione Totale (TV) più bassi nelle caratteristiche intermedie a livelli di rumore elevati, indicando rappresentazioni più regolari e coerenti.
• Ablazione: I registri sono più efficaci quando introdotti dopo i primi 4 strati e quando si utilizza un numero maggiore di token (ad esempio, 32 registri per pDiT-B, rispetto a 4 per i ViT).

Significato e Affermazioni

Il documento afferma che i token di registro sono un componente critico, ma finora trascurato, per la generazione di alta qualità tramite diffusione nello spazio dei pixel. Il significato risiede in:

1. Ridefinizione del Ruolo dei Registri: Dimostrare che i registri non sono semplicemente una soluzione per valori anomali specifici dei ViT, ma svolgono una funzione di regolarizzazione più ampia nei modelli di diffusione gestendo le norme delle caratteristiche e migliorando la coerenza spaziale.
2. Spiegazione dei Recenti Successi: Fornire una motivazione per le forti prestazioni empiriche dei recenti DiT nello spazio dei pixel che utilizzano il condizionamento nel contesto, attribuendo parte del loro successo al comportamento implicito dei registri.
3. Progettazione Efficiente di Architetture: Offrire un design a doppio flusso efficiente in termini di parametri che sfrutta i ruoli distinti di registri e patch, raggiungendo una qualità di generazione all'avanguardia con un sovraccarico di runtime trascurabile.

Gli autori concludono che, sebbene i loro design catturino la promessa della specializzazione dei token, il pieno potenziale dell'utilizzo dei registri nei transformer di diffusione nello spazio dei pixel rimane una domanda aperta per la ricerca futura.