Riepilogo Tecnico: Mean Mode Screaming e Residui MV-Split per Deep Diffusion Transformers

1. Enunciato del Problema

La scalabilità dei Diffusion Transformers (DiT) verso profondità estreme (da centinaia a migliaia di livelli) introduce una vulnerabilità strutturale distinta dalle euristiche standard di esplosione o vanishing dei gradienti. Gli autori identificano una specifica modalità di fallimento denominata Collasso Dominato dalla Media, in cui le rappresentazioni dei token si omogeneizzano e la variazione centrata dei token viene soppressa. Ciò porta a uno stato in cui la rete prevede una media banale, causando la divergenza della perdita di training fino ai livelli di inizializzazione, senza evidenti NaN o saturazione in avanti.

L'evento di ingresso improvviso in questo stato è caratterizzato come Mean Mode Screaming (MMS). L'MMS è definito da un picco acuto nella componente del gradiente coerente con la media, un rapido apertura dei rami residui e la successiva soppressione dei gradienti Query/Key (Q/K). Questo fenomeno può verificarsi anche quando l'addestramento appare stabile per migliaia di step prima di divergere entro pochi aggiornamenti.

I stabilizzatori di profondità esistenti, come ReZero e LayerScale, applicano un gating isotropo (scalare o per canale) all'intero ramo residuo. Sebbene questi metodi stabilizzino l'addestramento, sopprimono simultaneamente sia la componente media sia quella centrata. Questo smorzamento isotropo rallenta la convergenza attenuando il segnale centrato necessario per l'apprendimento di caratteristiche variabili spazialmente.

2. Metodologia: Analisi Meccanicistica e MV-Split

2.1 Analisi Meccanicistica

Gli autori scompongono le dinamiche di fallimento in due meccanismi primari:

1. Asimmetria Geometrica nell'Attention: L'attention stocastica per riga preserva rigorosamente gli stati di pura media (Proposizione 1) ma agisce in modo contrattivo sul sottospazio centrato (Proposizione 2). Man mano che i livelli si approfondiscono, la componente centrata dipende interamente dai rami residui per il rifornimento. Se gli aggiornamenti residui diventano dominati dalla componente media, la rete collassa in uno stato di pura media.
2. Decomposizione e Amplificazione del Gradiente: Gli autori dimostrano che i gradienti per le mappe lineari token-wise (scrittori residui) ammettono una decomposizione additiva esatta in una componente coerente con la media ($\Delta W_\mu$) e una componente centrata ($\Delta W_c$).
   • La componente coerente con la media scala come $O(T)$ quando i token si allineano (dove $T$ è la lunghezza della sequenza), entrando in un regime di accumulo coerente.
   • La componente centrata scala in modo diffusivo.
   • Trigger MMS: Man mano che le rappresentazioni dei token e gli aggiunti si omogeneizzano, i termini fuori diagonale con segno nell'espansione del gradiente cessano di annullarsi. Ciò causa un picco nel gradiente coerente con la media, dominando l'aggiornamento residuo.
3. Estinzione dei Gradienti Q/K: Una volta che i valori si omogeneizzano, il nucleo nullo della Jacobiana Softmax sopprime la componente costante del gradiente dei pesi di attention. Ciò causa la scomparsa dei gradienti Q/K (Lemma 1), bloccando la rete nello stato collassato perché il meccanismo di attention non può più ripristinare la variazione dei token.

2.2 Soluzione Proposta: Residui Mean–Variance Split (MV-Split)

Per affrontare ciò, gli autori propongono i Residui MV-Split, che disaccoppiano il controllo dei sottospazi della media e di quello centrato. Invece di un unico guadagno isotropo, l'MV-Split impiega una fusione instradata per sottospazio:

$$Z_l = X_l + \beta \odot (P F_l) + \alpha \odot J(F_l - X_l)$$

Dove:

• $X_l$ è l'input del tronco, $F_l$ è l'output del ramo.
• $J$ è il proiettore sul sottospazio della media della sequenza, e $P = I - J$ è il proiettore sul sottospazio centrato.
• $\alpha$ e $\beta$ sono vettori apprendibili (trasmettiti in broadcast attraverso i token).

Dinamiche Chiave:

• Percorso Centrato: La componente centrata riceve un aggiornamento residuo standard con guadagno $\beta$, preservando la capacità di apprendere caratteristiche variabili spazialmente.
• Percorso della Media: La componente media subisce una "sostituzione di media del tronco con perdita". La media del tronco viene contratta di $(1-\alpha)$ prima che venga aggiunta una nuova correzione. Ciò previene l'accumulo di aggiornamenti coerenti con la media senza restringere il percorso centrato.
• Dinamiche Inverse: Il flusso del gradiente è similmente diviso, permettendo alla componente del gradiente coerente con la media ($\Delta W_\mu$) di essere smorzata da $\alpha$ senza influenzare la componente del gradiente centrato ($\Delta W_c$) controllata da $\beta$.

3. Contributi Chiave

1. Caratterizzazione del Collasso: Il documento distingue tra lo stato di collasso dominato dalla media e l'evento di ingresso improvviso, l'MMS. Dimostra che l'inizializzazione standard porta a un collasso progressivo per profondità, mentre scrittori inizializzati a zero portano al picco ritardato dell'MMS.
2. Insight Meccanicistico: Gli autori forniscono una prova teorica che l'attention stocastica per riga preserva gli stati di pura media mentre contrae il sottospazio centrato. Derivano la Legge di Amplificazione dell'Allineamento, collegando la transizione all'accumulo coerente del gradiente all'allineamento delle rappresentazioni dei token e degli aggiunti. Dimostrano inoltre che l'omogeneizzazione dei valori sopprime strutturalmente l'apprendimento Q/K tramite il nucleo nullo della Jacobiana Softmax.
3. Metodo e Validazione: La proposta dei Residui MV-Split, che stabilizza l'addestramento smorzando il percorso della media mentre preserva il percorso centrato.
   • Validazione a 400 Livelli: In esperimenti abbinati, l'MV-Split previene il collasso divergente osservato nei baseline non stabilizzati e converge significativamente più velocemente di LayerScale.
   • Validazione di Scala a 1000 Livelli: Gli autori hanno addestrato con successo un DiT a 1000 livelli utilizzando l'MV-Split, dimostrando che l'architettura rimane addestrabile in modo stabile a profondità estreme dove i metodi standard falliscono.

4. Risultati Sperimentali

• Stabilità: Nelle esecuzioni a 400 livelli su latenti ImageNet-2012, sia il baseline non stabilizzato sia il baseline con tasso di apprendimento ridotto sono divergenti (collasso). LayerScale è rimasto stabile ma ha convergito più lentamente. L'MV-Split è rimasto stabile e ha ottenuto le prestazioni migliori tra le esecuzioni non divergenti.
• Prestazioni: A 50k step di addestramento, il modello MV-Split a 400 livelli ha raggiunto un FID-50K di 2.60 e un Inception Score di 185.5, superando LayerScale (FID 2.90, IS 165.5).
• Dinamiche del Gradiente: L'analisi dei gradienti degli scrittori ha mostrato che l'MV-Split limita la componente coerente con la media ($G_{mean}$) mantenendo una banda più alta e stabile per la componente centrata ($G_{ctr}$). Al contrario, LayerScale ha soppresso entrambe le componenti in modo uguale.
• Esecuzione a 1000 Livelli: Una singola esecuzione separata a 1000 livelli utilizzando lo stesso design MV-Split è rimasta stabile per tutta la durata dell'addestramento, raggiungendo un FID di 2.77 e un IS di 217.3 a 50k step.
• Risultati Qualitativi: Il modello a 1000 livelli ha generato campioni testo-immagine ad alta fedeltà attraverso diverse categorie semantiche (ad esempio, animali, oggetti, paesaggi) condizionati da didascalie in linguaggio naturale.

5. Significato e Affermazioni

Il documento afferma che l'MV-Split affronta un percorso di fallimento specifico e precedentemente non affrontato nei DiT Post-Norm ultra-profondi: lo squilibrio tra un aggiornamento dello scrittore coerente con la media in crescita $O(T)$ e un percorso centrato contrattivo.

• Necessità Architettonica: Gli autori sostengono che il gating isotropo (come LayerScale) è insufficiente per profondità estreme perché smorza proprio i segnali (variazione centrata) necessari per l'apprendimento delle caratteristiche. L'MV-Split fornisce il necessario controllo anisotropo.
• Scalabilità: L'addestramento riuscito di un DiT a 1000 livelli funge da validazione di scala limite, stabilendo che l'architettura può essere scalata a profondità estreme se l'interfaccia residua è gestita correttamente.
• Comprensione Meccanicistica: Il lavoro fornisce una spiegazione geometrica e algebrica rigorosa del motivo per cui i DiT profondi collassano, andando oltre le osservazioni euristico per arrivare a una scomposizione delle modalità di gradiente e delle dinamiche di attention.

Gli autori notano che la loro analisi è specifica per l'attention Softmax e le architetture Post-Norm. Non affermano che l'MMS spieghi tutte le instabilità dei Transformer profondi, ma piuttosto identificano un specifico percorso di fallimento del sottospazio residuo che l'MV-Split risolve all'interfaccia residua.