Technische Samenvatting: Dead-Direction Conditioners (DDC)

1. Probleemstelling

Verliesfuncties van diepe neurale netwerken vertonen continue symmetrieën (gauge-symmetrieën) in hun parameterruimtes. Deze omvatten logit-shifts, ReLU/SwiGLU-rescalings, LayerNorm-schaalinvarianties en rotaties van attention heads. Langs de banen van deze symmetrieën is het verlieslandschap vlak (dead directions), wat betekent dat de Fisher-informatie-metriek verdwijnt.

Standaard adaptieve optimizers zoals Adam houden geen rekening met deze symmetrieën. De per-coördinaat preconditioner van Adam ($1/\sqrt{\hat{v}}$) is alleen equivariënt onder coördinaatpermutaties. Wanneer architecturale groepen coördinaten mengen (bijv. het invers schalen van een paar gewichten of het roteren van een vectorruimte), "lekt" de update van Adam in de gauge-baan. Deze lekkage veroorzaakt twee primaire problemen:

1. Drift: De optimizer-traject drift langs de vlakke baan in plaats van op een vaste sectie van de symmetrie-quotiënt $\bar{\Theta} = \Theta/G$ te blijven. Deze drift beweegt de optimizer naar een ander minimum dan het zou bereiken als de symmetrie gerespecteerd zou worden.
2. Vervaging van de Rate Readability: Theoretische resultaten (Shirodkar, 2026) stellen vast dat de singuliere leersnelheid (de afname van de kleinste Fisher-eigenwaarde) alleen leesbaar is uit de traject als de optimizer schoon naar de quotiënt afdaalt. De lekkage van Adam introduceert ruis en bias, waardoor deze rate vervaagt en de geometrie van het minimum onleesbaar wordt.

2. Methodologie: Dead-Direction Conditioners (DDC)

Het artikel introduceert DDC, een preconditioning-framework dat een basis-optimizer (bijv. Adam of Muon) transformeert naar een $G$-equivariante optimizer. De kernmechanisme omvat het decomponeren van de parameterruimte en gradiënt-updates in verticale (baan-tangente) en horizontale (quotiënt) componenten met behulp van een $G$-invariante Riemanniaanse metriek $\rho$.

De Update-regel

Voor een gradiënt $g_t$ splitst DDC deze in $g_t = \Pi_V g_t + \Pi_H g_t$. Het houdt aparte tweede-moment-schatters bij:

• Verticale Component ($\hat{v}_V$): De verticale gradiënt wordt gecomprimeerd tot een enkele scalar per baan-dimensie (bijv. via row-mean voor CE-shifts). Deze scalar wordt gebruikt om de verticale update te normaliseren, waarbij de gehele baan-dimensie effectief als één eenheid wordt behanden.
• Horizontale Component ($\hat{v}_H$): De horizontale gradiënt wordt per coördinaat gepreconditionerd, maar strikt binnen een $G$-equivariante horizontale frame. Cruciaal is dat de update na de per-coördinaat divisie wordt geherprojecteerd ($\Pi_H$) om te garanderen dat de traject horizontaal blijft.

De update-stap is:
$$\theta_{t+1} = \theta_t - \eta \left[ \Pi_H \left( \frac{\Pi_H \hat{m}_t}{\sqrt{\hat{v}_H} + \epsilon} \right) + \text{lift}_V \left( \frac{\text{vert-to-scalar}(\Pi_V \hat{m}_t)}{\sqrt{\hat{v}_V} + \epsilon} \right) \right]$$

Gehandelde Architecturale Gauges

DDC construeert specifieke projectoren en reducties voor vier architecturale gauges:

1. Cross-Entropy Shift: Additieve verschuivingen in logits (bias of unembed weights).
2. ReLU/SwiGLU Rescaling: Multiplicatieve schaling van aangrenzende gewichtsparen ($W_\ell, W_{\ell+1} \to cW_\ell, c^{-1}W_{\ell+1}$).
3. LayerNorm/RMSNorm Scale: Per-kanaal schaling van normalisatieparameters en downstream gewichten.
4. Attention Head Rotation: Niet-abeliaanse $O(d_{head})$ rotaties van QK en VO gewichten.

Voor de niet-abeliaanse rotatie gauge faalt de standaard per-coördinaat $\hat{v}$. DDC stelt vier $O(d_{head})$-bewuste constructies voor, waarvan de geïmplementeerde standaard body_frame_topk is. Dit behoudt per-coördinaat adaptiviteit in de eigenbasis van de Gram-matrix $M = W_Q^T W_Q + W_K^T W_K$, maar alleen op "levende" eigenrichtingen (die eigenwaarden boven een drempelwaarde hebben), terwijl dode richtingen enkel met momentum worden behandeld.

Samenstelling van de Basis-optimizer

• DDCAdam: Past de volledige DDC-constructie toe op Adam. Het behandelt alle vier de gauges exact (of tot op de eerste orde voor body-frame rotaties).
• DDCMuon: Past DDC toe op de Muon optimizer (die updates orthogonaliseert). De orthogonalisatie van Muon commuteert exact met rotatie gauges, waardoor DDCMuon exact is voor rotaties. Echter, de orthogonalisatie van Muon verwijdert de per-kanaal schaalvrijheid, wat betekent dat DDCMuon de rescaling gauges (ReLU/LN scale) niet kan conditioneren en terugvalt op Adam voor die componenten.

3. Belangrijkste Bijdragen

1. Equivariante Preconditioning: Een algemene constructie die adaptieve optimizers transformeert om respect te hebben voor continue architecturale symmetrieën, waardoor de traject een gepreconditioneerde gradiëntstroom op de symmetrie-quotiënt $\bar{\Theta}$ blijft.
2. Rate Readability: Theoretisch bewijs (Theorem 1) dat DDC de "trajectory-rate readability" eigenschap herstelt. Op de quotiënt vervalt de kleinste Fisher-eigenwaarde als $\bar{t}^{2(k-1)}$, waardoor de singuliere orde $k$ direct afgelezen kan worden uit het traject, een eigenschap die in standaard Adam wordt vertroebeld.
3. Gauge Containment: Empirische demonstratie dat DDC de drift van gauge-modi (bijv. logit shifts, weight rescalings) voorkomt die Adam teistert, door deze modi op hun initiële waarden of vaste punten te houden.
4. Verbeterde Optimalisatie-uitkomsten:
   • Resistentie tegen Over-training: Op taalmodellen getraind voorbij het punt van fit, weerstaat DDCAdam de "over-training collapse" waarbij de validatieverlies divergeert van het trainingsverlies.
   • Grokking Betrouwbaarheid: Op diepe transformers (diepte 24), maakt DDCMuon "grokking" (plotselinge generalisatie) mogelijk op 10/11 seeds waar vanilla Muon volledig faalt (0/11).
   • Compressie: Op Vision Transformers comprimeert DDCAdam de overbodige feed-forward capaciteit agressiever onder weight decay dan AdamW, waarbij een lagere validatieverlies wordt bereikt.

4. Experimentele Resultaten

• Taalmodellen (Diepte-12, Over-training):
  • Rate Readability: DDCAdam leest de singuliere rate in 32 van de 65 layer-observable cellen, terwijl AdamW dit slechts in 7 cellen doet. In AdamW stijgen de Fisher-eigenwaarden weg van het minimum; in DDCAdam dalen ze naar de vloer.
  • Over-training: DDCAdam behoudt een validatie-train verliesverschil van 0.67, vergeleken met 5.88 voor AdamW.
• Grokking Transformers:
  • Diepte-24: DDCMuon bereikt 10/11 grokking seeds met een nauwkeurigheid >0.9, terwijl vanilla Muon 0/11 bereikt.
  • Diepte-8 (RoPE): Het matchen van de rotatie gauge aan de specifieke torus-symmetrie van Rotary Position Embeddings ($SO(2)^{d/2}$) levert een hogere nauwkeurigheid (0.977) en lagere variantie op dan generieke $O(d_{head})$ gauges (0.955).
• Vision Transformers (ViT):
  • Getraind vanaf nul op ImageNet-100, bereikt DDCAdam een lager validatieverlies (1.71 vs 2.12 nats) en comprimeert het netwerk (activatie singuliere waarden 4x kleiner) bij dezelfde weight decay waarbij AdamW de capaciteit intact laat.
• Gauge Containment:
  • Op synthetische taken vermindert DDCAdam de cross-seed drift van gauge-modi met 14x (ReLU rescale) en 65x (LayerNorm scale) vergeleken met AdamW, terwijl de uiteindelijke taakverlies gelijk blijft.

5. Betekenis en Claims

Het artikel claimt dat het inbouwen van symmetrie in de optimizer fundamenteel de geometrie van het gevonden oplossing verandert. Door de gauge-symmetrie te respecteren:

• Het Minimum is Scherper: De optimizer settleert in een minder gedegenereerd minimum (hogere Fisher-eigenwaarden buiten de singuliere vloer).
• De Geometrie is Meetbaar: Het traject wordt een leesbaar signaal van de onderliggende singuliere leerscoëfficiënt, waardoor de singulariteit direct uit de trainingsdynamiek kan worden gediagnosticeerd.
• Betrouwbaarheid Neemt Toe: De optimizer wordt robuust tegen de "drift" die falen veroorzaakt in diepe of over-getrainde regimes, met name in taken die "grokking" vereisen.

De auteurs benadrukken dat de constructie generiek is aan de preconditioner en dat het voordeel niet enkel voortkomt uit weight decay, maar uit de geometrische beperking om op de quotiënt te blijven. Het artikel concludeert dat de gauge-symmetrie van een netwerk, wanneer gerespecteerd door de optimizer, "het minimum dat het vindt verscherpt en de geometrie van dat minimum verandert in iets dat het traject kan meten."

Genoteerde Beperkingen:

• De snelheidswinst van Muon-gebaseerde grokking wordt gedeeltelijk toegeschreven aan de gebundelde weight decay die inherent is aan de equivariante update, hoewel de projectie zelf ook zorgt voor betrouwbaarheid en een schonere structuur.
• De exacte equivariantie voor de body-frame rotatie constructie is benaderend (eerste-orde) vanwege eigenvector-tekenambiguïteit, hoewel deze drift klein is ($\sim 10^{-3}$ over 20 stappen).
• Het lezen van de volledige singuliere structuur (orde en multipliciteit) van canonical alignment vereist machinerie die niet volledig in dit werk is ontwikkeld.