Same Function, Different Trajectory: No Standard Optimizer Is Equivariant to Transformer Reparameterization
Organizations: The University of Tokyo
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 at 101M parameters on WikiText-103, and Soft SAGGS reduces PPL by 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 ; hard and soft rebalancing recover to within PPL of the control.
Figures & tables
| Sub-layer | Mechanism | Exact | ||
| Self-attention | Bilinear in logits | ✓ Yes | ||
| Linear composition | ✓ Yes | |||
| MLP (ReLU) | Positive homogeneity | ✓ Yes | ||
| MLP (GELU) | Non-homogeneous | ✗ No | ||
| SwiGLU | Linear in value path | ✓ Yes |
| Vanilla | SAGGS | ||||
|---|---|---|---|---|---|
| Resume condition | Final PPL |
|---|---|
| Control resume | |
| Surgery resume, no SAGGS | |
| Surgery + SAGGS | |
| Surgery + Soft SAGGS |
Appendix figures & tables31 assets
Supplementary material from the paper’s appendix.
Appendix
| Name | Heads | Params | LR | |||
| GPT-2-style | ||||||
| Small | 256 | 1024 | 12 | 4 | 22M | |
| Medium | 512 | 2048 | 24 | 8 | 101M | |
| Deep | 256 | 1024 | 48 | 4 | 51M | |
| Xlarge | 1024 | 4096 | 24 | 16 | 354M | |
| 1B (SwiGLU) | 2048 | 5504 | 22 | 16 | 1.22B | |
| Setting | Contrast | Purpose | |||
| Primary | |||||
| 101M WT-103 | V SAGGS | Core language-model protection | |||
| 31M modern | V Soft | Architecture and tokenizer | |||
| 101M Lion | V Soft | Optimizer transfer | |||
| Workflow | |||||
| 123.8M surgery | No-SAGGS control | State-reset resume | |||
| Vanilla | SAGGS | ||
|---|---|---|---|
| Mean PPL | Paired difference | ||||||
|---|---|---|---|---|---|---|---|
| Vanilla | Soft, | ||||||
| Seed | |||
|---|---|---|---|
| 1 | |||
| 2 | |||
| 3 |
| Mean PPL | Paired difference | ||||||
|---|---|---|---|---|---|---|---|
| Vanilla | Soft SAGGS | ||||||
| Vanilla | SAGGS | Soft | State | |
|---|---|---|---|---|
| Contrast | |||||
|---|---|---|---|---|---|
| Surgery no SAGGS control | |||||
| Surgery no SAGGS SAGGS | |||||
| Surgery no SAGGS Soft | |||||
| Control SAGGS | |||||
| Control Soft |
| Resume condition | Final PPL |
|---|---|
| Control, optimizer reset | |
| Surgery, optimizer reset | |
| Surgery reset + Soft SAGGS | |
| Control, optimizer preserved | |
| Surgery, state-covariant optimizer |
| Resume condition | Final PPL |
|---|---|
| Control, preserved state | |
| Surgery + stale AdamW state | |
| Surgery + covariant AdamW state | |
| Surgery + stale state + Soft SAGGS |
| Resume condition | Final PPL |
|---|---|
| Control state preserved | |
| Exact edit, covariant state | |
| Exact edit, all state stale | |
| Exact edit, deep-half stale | |
| Deep-half stale + Soft SAGGS |
| Pair | ||||
|---|---|---|---|---|
| Q–K (ReLU) | 12 | |||
| 24 | ||||
| 48 | ||||
| V–O (ReLU) | 12 | |||
| 24 | ||||
| 48 |
| Model | Params | Q–K | V–O | MLP | |
|---|---|---|---|---|---|
| Small | 22M | 12 | |||
| Medium | 101M | 24 | |||
| Deep | 51M | 48 |
| Vanilla | SAGGS+P | |||
|---|---|---|---|---|
| PPL | Gradient slope | |||||
|---|---|---|---|---|---|---|
| Vanilla | SAGGS+P | Vanilla | SAGGS+P | |||
| Method | PPL | MLP imbalance | Gradient slope | |
|---|---|---|---|---|
| Deep, | ||||
| Vanilla (baseline) | ||||
| SAGGS+P | ||||
| Vanilla | ||||
| SAGGS+P | ||||
| Small, | ||||
| Pair type | Maximum relative error (%) | |||
|---|---|---|---|---|
| Q–K | ||||
| V–O | ||||
| MLP (SwiGLU) | ||||
| Vanilla PPL | SAGGS+P PPL | ||
|---|---|---|---|
| Scale | Vanilla PPL | SAGGS+P PPL | ||
|---|---|---|---|---|
| Small, 22M | ||||
| Medium, 101M | ||||
| 354M | ||||
| Metric | Vanilla | Vanilla | SAGGS |
|---|---|---|---|
| Final validation PPL | |||
| Best validation PPL | |||
| Mean gradient norm | |||
| Q–K max-abs log-ratio | |||
| V–O max-abs log-ratio | |||
| Mean PPL | Paired difference | ||||||
|---|---|---|---|---|---|---|---|
| Vanilla | Soft, | ||||||
| Mean PPL | Paired difference | ||||||
|---|---|---|---|---|---|---|---|
| Vanilla | Soft, | ||||||
| Mean PPL | Paired difference | ||||||
|---|---|---|---|---|---|---|---|
| Vanilla | Soft, | ||||||
| Symbol | Meaning |
|---|---|
| Architecture and loss | |
| model parameters | |
| data loss | |
| , | depth and layer index |
| , , | hidden, head, and MLP widths |
| , , , | attention projections |