Technical Summary: Dead-Direction Conditioners (DDC)

1. Problem Statement

Deep neural network losses exhibit continuous symmetries (gauge symmetries) in their parameter spaces. These include logit shifts, ReLU/SwiGLU rescalings, LayerNorm scale invariances, and attention head rotations. Along the orbits of these symmetries, the loss landscape is flat (dead directions), meaning the Fisher information metric vanishes.

Standard adaptive optimizers like Adam fail to respect these symmetries. Adam's per-coordinate preconditioner ($1/\sqrt{\hat{v}}$) is only equivariant under coordinate permutations. When architectural groups mix coordinates (e.g., scaling a pair of weights inversely or rotating a vector space), Adam's update "leaks" into the gauge orbit. This leakage causes two primary issues:

1. Drift: The optimizer trajectory drifts along the flat orbit rather than staying on a fixed section of the symmetry quotient $\bar{\Theta} = \Theta/G$. This drift moves the optimizer to a different minimum than it would reach if the symmetry were respected.
2. Blurred Rate Readability: Theoretical results (Shirodkar, 2026) establish that the singular-learning rate (the decay of the smallest Fisher eigenvalue) is readable from the trajectory only if the optimizer descends cleanly on the quotient. Adam's leakage introduces noise and bias, blurring this rate and making the geometry of the minimum unreadable.

2. Methodology: Dead-Direction Conditioners (DDC)

The paper introduces DDC, a preconditioning framework that lifts a base optimizer (e.g., Adam or Muon) into a $G$-equivariant one. The core mechanism involves decomposing the parameter space and gradient updates into vertical (orbit-tangent) and horizontal (quotient) components using a $G$-invariant Riemannian metric $\rho$.

The Update Rule

For a gradient $g_t$, DDC splits it into $g_t = \Pi_V g_t + \Pi_H g_t$. It maintains separate second-moment estimators:

• Vertical Component ($\hat{v}_V$): The vertical gradient is collapsed to a single scalar per orbit dimension (e.g., via row-mean for CE shifts). This scalar is used to normalize the vertical update, effectively treating the entire orbit dimension as a single unit.
• Horizontal Component ($\hat{v}_H$): The horizontal gradient is preconditioned per-coordinate but strictly within a $G$-equivariant horizontal frame. Crucially, after the per-coordinate division, the update is re-projected ($\Pi_H$) to ensure the trajectory remains horizontal.

The update step 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]$$

Architectural Gauges Handled

DDC constructs specific projectors and reductions for four architectural gauges:

1. Cross-Entropy Shift: Additive shifts in logits (bias or unembed weights).
2. ReLU/SwiGLU Rescaling: Multiplicative scaling of adjacent weight pairs ($W_\ell, W_{\ell+1} \to cW_\ell, c^{-1}W_{\ell+1}$).
3. LayerNorm/RMSNorm Scale: Per-channel scaling of normalization parameters and downstream weights.
4. Attention Head Rotation: Non-abelian $O(d_{head})$ rotations of QK and VO weights.

For the non-abelian rotation gauge, standard per-coordinate $\hat{v}$ fails. DDC proposes four $O(d_{head})$-aware constructions, with the deployed default being body_frame_topk. This maintains per-coordinate adaptivity in the eigenbasis of the Gram matrix $M = W_Q^T W_Q + W_K^T W_K$, but only on "live" eigendirections (those with eigenvalues above a threshold), while treating dead directions with momentum only.

Base Optimizer Composition

• DDCAdam: Applies the full DDC construction to Adam. It handles all four gauges exactly (or to leading order for body-frame rotations).
• DDCMuon: Applies DDC to the Muon optimizer (which orthogonalizes updates). Muon's orthogonalization commutes exactly with rotation gauges, making DDCMuon exact for rotations. However, Muon's orthogonalization removes the per-channel scale freedom, meaning DDCMuon cannot condition the rescaling gauges (ReLU/LN scale) and defaults to Adam for those components.

3. Key Contributions

1. Equivariant Preconditioning: A general construction that lifts adaptive optimizers to respect continuous architectural symmetries, ensuring the trajectory remains a preconditioned gradient flow on the symmetry quotient $\bar{\Theta}$.
2. Rate Readability: Theoretical proof (Theorem 1) that DDC restores the "trajectory-rate readability" property. On the quotient, the smallest Fisher eigenvalue decays as $\bar{t}^{2(k-1)}$, allowing the singular order $k$ to be read directly from the trajectory, a property obscured in standard Adam.
3. Gauge Containment: Empirical demonstration that DDC prevents the drift of gauge modes (e.g., logit shifts, weight rescalings) that plagues Adam, holding these modes at their initial values or fixed points.
4. Improved Optimization Outcomes:
   • Over-training Resistance: On language models trained past the point of fit, DDCAdam resists the "over-training collapse" where validation loss diverges from training loss.
   • Grokking Reliability: On deep transformers (depth 24), DDCMuon enables "grokking" (sudden generalization) on 10/11 seeds where vanilla Muon fails completely (0/11).
   • Compression: On Vision Transformers, DDCAdam compresses spare feed-forward capacity more aggressively under weight decay than AdamW, reaching lower validation loss.

4. Experimental Results

• Language Models (Depth-12, Over-training):
  • Rate Readability: DDCAdam reads the singular rate in 32 of 65 layer-observable cells, whereas AdamW reads it in only 7. In AdamW, the Fisher eigenvalues rise away from the minimum; in DDCAdam, they collapse toward the floor.
  • Over-training: DDCAdam maintains a validation-train loss gap of 0.67, compared to 5.88 for AdamW.
• Grokking Transformers:
  • Depth-24: DDCMuon achieves 10/11 grokking seeds at accuracy >0.9, while vanilla Muon achieves 0/11.
  • Depth-8 (RoPE): Matching the rotation gauge to the specific torus symmetry of Rotary Position Embeddings ($SO(2)^{d/2}$) yields higher accuracy (0.977) and lower variance than generic $O(d_{head})$ gauges (0.955).
• Vision Transformers (ViT):
  • Trained from scratch on ImageNet-100, DDCAdam achieves lower validation loss (1.71 vs 2.12 nats) and compresses the network (activation singular values 4x smaller) at the same weight decay where AdamW leaves capacity intact.
• Gauge Containment:
  • On synthetic tasks, DDCAdam reduces the cross-seed drift of gauge modes by 14x (ReLU rescale) and 65x (LayerNorm scale) compared to AdamW, while matching final task loss.

5. Significance and Claims

The paper claims that building symmetry into the optimizer fundamentally changes the geometry of the solution found. By respecting the gauge symmetry:

• The Minimum is Sharper: The optimizer settles into a less degenerate minimum (higher Fisher eigenvalues off the singular floor).
• The Geometry is Measurable: The trajectory becomes a readable signal of the underlying singular learning coefficient, allowing for the diagnosis of model complexity directly from training dynamics.
• Reliability Increases: The optimizer becomes robust to the "drift" that causes failure in deep or over-trained regimes, particularly in tasks requiring "grokking."

The authors emphasize that the construction is generic to the preconditioner, lifting Adam, Muon, and Shampoo, and that the benefits are not merely due to weight decay but to the geometric constraint of staying on the quotient. The paper concludes that a network's gauge symmetry, when respected by the optimizer, "sharpens the minimum it finds and turns that minimum's geometry into something the trajectory can measure."

Limitations Noted:

• The Muon-base speed-up in grokking is partially attributed to the bundled weight decay inherent in the equivariant update, though the projection itself provides reliability and cleaner structure.
• Exact equivariance for the body-frame rotation construction is approximate (leading-order) due to eigenvector sign ambiguity, though this drift is small ($\sim 10^{-3}$ over 20 steps).
• Reading the full singular structure (order and multiplicity) off canonical alignment requires machinery not fully developed in this work.