Dead-Direction Conditioners: Gauge-Equivariant Preconditioning for Deep Networks
Tejas Pradeep Shirodkar
cs.LG, math.DG, math.OC, stat.ML
Submitted: 2026-06-28
Comments: 69 pages, 28 figures, 9 tables. Builds the gauge-equivariant preconditioner left open in arXiv:2606.05957
License: http://creativecommons.org/licenses/by/4.0/
The gist: A deep network's loss is invariant to continuous symmetries of its parameters: the logit shift, the ReLU rescaling, the LayerNorm scale, the per-head attention rotation.
Terminology
Abstract
A deep network's loss is invariant to continuous symmetries of its parameters: the logit shift, the ReLU rescaling, the LayerNorm scale, the per-head attention rotation. Adam's per-coordinate preconditioner drifts along each symmetry orbit, which pulls the trajectory off the symmetry quotient where the optimization lives and blurs the singular-learning rate the quotient makes readable. We build DDC, a Dead-Direction Conditioner that lifts a base optimizer into a G-equivariant one: it conditions the optimizer's state in the orbit decomposition of a G-invariant metric, so the trajectory stays a preconditioned gradient flow on the quotient = /G. The construction carries four architectural gauges (cross-entropy shift, ReLU and SwiGLU rescaling, LayerNorm and RMSNorm scale, and a per-head O(d head) attention rotation matched to RoPE), proves exactly equivariant on an Adam base, and composes with a Muon base through a gauge-equivariant orthogonaliser. Respecting the symmetry changes both the minimum the optimizer reaches and what it leaves measurable there. On a language model trained past the point of fit, DDCAdam resists the over-training collapse AdamW falls into, holding a validation-train loss gap of 0.67 against 5.88, and reads the dead-direction rate in 32 of 65 layer-by-observable cells where AdamW reads it in 7. A vision transformer trained from scratch reaches lower validation loss (1.71 against 2.12) while compressing spare feed-forward capacity a matched AdamW leaves intact. On a Muon base, where the rotation gauge composes exactly, DDCMuon groks ten of eleven seeds at depth 24 that a plain Muon never reaches. Built into the optimizer, a network's gauge symmetry sharpens the minimum it finds and turns that minimum's geometry into something the trajectory can measure.
Sources
- On the Optimization of Deep Networks: Implicit Acceleration by Overparameterization
- Riemannian Adaptive Optimization Methods
- Modular Duality in Deep Learning
- The Z-loss: a shift and scale invariant classification loss belonging to the Spherical Family
- How do simple rotations affect the implicit bias of Adam?
- Optimal Scaling Needs Optimal Norm
- Rotational Equilibrium: How Weight Decay Balances Learning Across Neural Networks
- Symmetry-Compatible Principle for Optimizer Design: Embeddings, LM Heads, SwiGLU MLPs, and MoE Routers
- NorMuon: Making Muon more efficient and scalable
- VectorAdam for Rotation Equivariant Geometry Optimization
- Training Deep Learning Models with Norm-Constrained LMOs
- Grokking at the Edge of Numerical Stability
- Dead-Direction Signatures: A Cheap Spectral Reading of Singular Complexity
- Symmetry Breaking in Transformers for Efficient and Interpretable Training
- L2 Regularization versus Batch and Weight Normalization
- LoRA Done RITE: Robust Invariant Transformation Equilibration for LoRA Optimization
- Symmetry in Neural Network Parameter Spaces
Related papers
- Polynomial-Augmented Neural Networks (PANNs) with Weak Orthogonality Constraints for Enhanced Function and PDE Approximation
- AIRL-S: Unifying Reinforcement Learning and Search-Based Test-Time Scaling via Adversarial Inverse Reinforcement Learning
- Transformers as Bayesian In-Context Experimenters: Smoothness-Adaptive Efficient ATE Estimation
- Convergence issues in Relational Concept Analysis based on AOC-posets
- Beliefs Beyond Posteriors: The Predictive Cost of ELBO Binding
- Understanding Diffusion Models via Ratio-Based Function Approximation with SignReLU Networks