arXiv preprint ·
Same Function, Different Trajectory: No Standard Optimizer Is Equivariant to Transformer Reparameterization
- Published
- Source
- arXiv preprint, posted before peer review
- Original
- Read it at arXiv
- Topics
Abstract
We show that two Transformers representing the same function can follow different training trajectories. For example, doubling a query matrix and halving its key matrix preserves the model's function but changes the gradient scales seen by the optimizer. Our theoretical analysis establishes that standard block-local first-order optimizers fail to scale updates with parameters under a shared learning rate. For AdamW, transforming the first- and second-moment estimates alone does not correct the mismatch between parameter and update scales. To make runs restarted from equivalent checkpoints follow the same trajectory, we balance matrix pairs and transform their moment estimates together, bringing distinct but equivalent training states to the same starting point. To correct scale imbalance between paired Transformer matrices without changing predictions, we propose Symmetry-Aware Gradient Geometry Stabilization (SAGGS), which periodically equalizes the Frobenius norms of exact symmetry pairs, and its damped variant, Soft SAGGS. SAGGS reduces perplexity (PPL) by $0.374\pm0.026$ at 101M parameters on WikiText-103, and Soft SAGGS reduces PPL by $0.780\pm0.040$ in a 31M modern-architecture replication, both under scale imbalance. In a 123.8M checkpoint-resume experiment, reciprocally rescaling matrix pairs barely changes validation loss at restart but increases final PPL by $3.78\pm0.42$; hard and soft rebalancing recover to within $0.02$ PPL of the control.
The authors’ abstract, as arXiv lists it. arXiv’s listings, abstracts included, are free for anyone to reuse.
Read it on ReadSearch
Each opens after you create an account, and brings you back to this paper.
More in Machine Learning
Collective Benefits and Efficiency in Coupled Echo State Networks
How Long Is the Coast of a Neuron? Resolution-Dependent Interaction Order and Computational Capacity
Neural Dimensional Analysis: Foundations of Dimensionless and Scale-Consistent Computation
Design and Analysis of Temporal Representations in Neural Networks
Saliency-Guided Replay for Continual Learning in Deep Spiking Neural Networks
Get the few papers worth your time in Machine Learning every morning.
ReadSearch reads the new papers in your field each day and brings you the few worth reading, each with its key finding in a sentence.