Technische Samenvatting: Mean Mode Screaming en MV-Split Residuals voor Deep Diffusion Transformers

1. Probleemstelling

Het schalen van Diffusion Transformers (DiTs) naar extreme diepten (honderden tot duizenden lagen) introduceert een structurele kwetsbaarheid die verschilt van de standaard heuristieken voor exploderende of verdwijnende gradiënten. De auteurs identificeren een specifiek faalmechanisme, genaamd Mean-Dominated Collapse (instorting door overheersing van het gemiddelde), waarbij tokenrepresentaties homogeniseren en de gecentreerde tokenvariatie wordt onderdrukt. Dit leidt tot een toestand waarin het netwerk een triviaal gemiddelde voorspelt, waardoor de trainingsfout divergeert naar niveaus van de initialisatie, zonder duidelijke NaN's of saturatie in de voorwaartse richting.

De plotselinge overgang naar deze toestand wordt gekenmerkt als Mean Mode Screaming (MMS). MMS wordt gedefinieerd door een scherpe piek in de gemiddelde-coherente gradiëntcomponent, een snelle opening van residual branches en de daaropvolgende onderdrukking van Query/Key (Q/K) gradiënten. Dit fenomeen kan optreden zelfs wanneer training duizenden stappen stabiel lijkt voordat het binnen enkele updates divergeert.

Bestaande dieptestabilisatoren, zoals ReZero en LayerScale, passen isotrope gating (scalair of per kanaal) toe op de gehele residual branch. Hoewel deze methoden de training stabiliseren, onderdrukken ze zowel de gemiddelde als de gecentreerde componenten gelijktijdig. Deze isotrope demping vertraagt de convergentie door het gecentreerde signaal te verzwakken dat nodig is voor het leren van ruimtelijk variërende kenmerken.

2. Methodologie: Mechanistische Analyse en MV-Split

2.1 Mechanistische Analyse

De auteurs ontleden de faaldynamiek in twee primaire mechanismen:

1. Geometrische Asymmetrie in Attention: Rij-stochastische attention behoudt strikt pure-gemiddelde toestanden (Propositie 1), maar werkt contractief op de gecentreerde deelruimte (Propositie 2). Naarmate de lagen dieper worden, is het gecentreerde component volledig afhankelijk van residual branches voor aanvulling. Als residual updates worden gedomineerd door het gemiddelde component, stort het netwerk in tot een pure-gemiddelde toestand.
2. Gradiënt Decompositie en Amplificatie: De auteurs tonen aan dat gradiënten voor token-weise lineaire kaarten (residual writers) een exacte additieve decompositie toelaten in een gemiddelde-coherente component ($\Delta W_\mu$) en een gecentreerde component ($\Delta W_c$).
   • De gemiddelde-coherente component schaalt als $O(T)$ wanneer tokens aligneren (waarbij $T$ de sequentielengte is), waardoor een coherent accumulatieregime wordt betreden.
   • De gecentreerde component schaalt diffuus.
   • MMS Trigger: Naarmate tokenrepresentaties en adjointen homogeniseren, stoppen de getekende niet-diagonale termen in de gradiëntexpansie met het opheffen. Dit veroorzaakt een piek in de gemiddelde-coherente gradiënt, die de residual update domineert.
3. Q/K Gradiënt Extinctie: Zodra waarden homogeniseren, onderdrukt de nulkant van de Softmax-Jacobaan de constante component van de attention-weight gradiënt. Dit zorgt ervoor dat Q/K gradiënten verdwijnen (Lemma 1), waardoor het netwerk in de ingestorte toestand wordt vergrendeld omdat het attention-mechanisme tokenvariatie niet langer kan herstellen.

2.2 Voorgestelde Oplossing: Mean–Variance Split (MV-Split) Residuals

Om dit aan te pakken, stellen de auteurs MV-Split Residuals voor, die de controle over de gemiddelde en gecentreerde deelruimten ontkoppelen. In plaats van één isotrope versterking, maakt MV-Split gebruik van een subspace-gestuurd samenvoegen:

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

Waarbij:

• $X_l$ de trunk-input is, $F_l$ is de branch-output.
• $J$ de projector is op de sequentie-gemiddelde deelruimte, en $P = I - J$ is de projector op de gecentreerde deelruimte.
• $\alpha$ en $\beta$ leerbare vectoren zijn (uitgezonden over tokens).

Kern Dynamiek:

• Gecentreerd Pad: Het gecentreerde component ontvangt een standaard residual update met versterking $\beta$, waardoor het vermogen behouden blijft om ruimtelijk variërende kenmerken te leren.
• Gemiddeld Pad: Het gemiddelde component ondergaat een "leaky trunk-mean replacement". De trunk-gemiddelde wordt gecontracteerd met $(1-\alpha)$ voordat een nieuwe correctie wordt toegevoegd. Dit voorkomt de accumulatie van gemiddelde-coherente updates zonder het gecentreerde pad te verkleinen.
• Terugwaartse Dynamiek: De gradiëntstroom wordt eveneens gesplitst, waardoor de gemiddelde-coherente gradiëntcomponent ($\Delta W_\mu$) kan worden gedempt door $\alpha$ zonder invloed op de gecentreerde gradiëntcomponent ($\Delta W_c$) die wordt gecontroleerd door $\beta$.

3. Belangrijkste Bijdragen

1. Karakterisering van Instorting: Het artikel onderscheidt tussen de instortingstoestand gedomineerd door het gemiddelde en het plotselinge instapgebeuren, MMS. Het toont aan dat standaard initialisatie leidt tot een progressieve diepte-afhankelijke instorting, terwijl nul-geinitialiseerde writers leiden tot de vertraagde MMS-piek.
2. Mechanistisch Inzicht: De auteurs leveren een theoretisch bewijs dat rij-stochastische attention pure-gemiddelde toestanden behoudt terwijl het de gecentreerde deelruimte contracteert. Zij leiden de Alignment-Amplificatiewet af, die de overgang naar coherente gradiëntaccumulatie koppelt aan de alignering van tokenrepresentaties en adjointen. Zij bewijzen verder dat waardehomogenisatie structureel Q/K-lering onderdrukt via de nulkant van de Softmax-Jacobaan.
3. Methode en Validatie: Het voorstellen van MV-Split Residuals, die training stabiliseren door het gemiddelde pad te dempen terwijl het gecentreerde pad behouden blijft.
   • 400-Laag Validatie: In overeenkomstige experimenten voorkomt MV-Split de divergente instorting die wordt gezien in niet-gestabiliseerde baselines en convergeert aanzienlijk sneller dan LayerScale.
   • 1000-Laag Schaal Validatie: De auteurs trainden succesvol een 1000-laagse DiT met MV-Split, wat aantoont dat de architectuur stabiel trainbaar blijft bij extreme diepten waar standaard methoden falen.

4. Experimentele Resultaten

• Stabiliteit: In 400-laagse runs op ImageNet-2012 latenten divergeerden zowel de niet-gestabiliseerde baseline als een baseline met verlaagde leersnelheid (instorting). LayerScale bleef stabiel maar convergeerde langzamer. MV-Split bleef stabiel en behaalde de beste prestaties onder de niet-divergente runs.
• Prestaties: Bij 50k trainingsstappen behaalde het 400-laagse MV-Split-model een FID-50K van 2.60 en een Inception Score van 185.5, waarmee het LayerScale (FID 2.90, IS 165.5) overtrof.
• Gradiënt Dynamiek: Analyse van writer-gradiënten toonde aan dat MV-Split de gemiddelde-coherente component ($G_{mean}$) begrenst terwijl het een hogere, stabiele band handhaaft voor het gecentreerde component ($G_{ctr}$). In tegenstelling hieronder onderdrukte LayerScale beide componenten gelijk.
• 1000-Laagse Run: Een aparte 1000-laagse run met hetzelfde MV-Split-ontwerp bleef stabiel gedurende de volledige trainingshorizon, met een FID van 2.77 en een IS van 217.3 bij 50k stappen.
• Kwalitatieve Resultaten: Het 1000-laagse model genereerde hoge-fideliteit tekst-naar-afbeelding samples over diverse semantische categorieën (bijvoorbeeld dieren, objecten, landschappen), geconditioneerd op natuurlijke taalbijschriften.

5. Betekenis en Claims

Het artikel claimt dat MV-Split een specifiek, eerder onbehandeld faalpad adresseert in ultra-diepe Post-Norm DiTs: het onbalans tussen een $O(T)$ groeiende gemiddelde-coherente writer-update en een contractief gecentreerd pad.

• Architectonische Noodzaak: De auteurs betogen dat isotrope gating (zoals LayerScale) onvoldoende is voor extreme diepten omdat het juist de signalen (gecentreerde variatie) dempt die nodig zijn voor kenmerklering. MV-Split biedt de benodigde anisotrope controle.
• Schalbaarheid: Het succesvol trainen van een 1000-laagse DiT dient als validatie op de grensschaal, waarmee wordt vastgesteld dat de architectuur kan worden geschaald naar extreme diepten als de residual interface correct wordt beheerd.
• Mechanistisch Begrip: Het werk biedt een rigoureuze geometrische en algebraïsche verklaring voor waarom diepe DiTs instorten, en gaat voorbij aan heuristische observaties naar een decompositie van gradiëntmodi en attention-dynamiek.

De auteurs merken op dat hun analyse specifiek is voor Softmax attention en Post-Norm architecturen. Zij claimen niet dat MMS alle instabiliteiten van diepe Transformers verklaart, maar identificeren eerder een specifiek residual-subspace faalpad dat MV-Split oplost op de residual interface.