# Gauge-equivariant Adam for factored layers

- ID: 1827
- Canonical URL: https://synthcore.org/idea/1827/gauge-equivariant-adam-for-factored-layers
- API JSON: https://synthcore.org/api/idea/1827.json
- API Markdown: https://synthcore.org/api/idea/1827.md
- Verification status: mech_ok_no_baseline
- Source: [arXiv:2608.05136](https://arxiv.org/abs/2608.05136)
- Category: optimization
- Solves: generalization, accuracy, stability
- ML areas: optimizer, fine-tuning, training
- Math tags: symmetry, linear-algebra, optimization
- Ratings: usefulness 8/10; difficulty 4/10; novelty 7/10

## Idea description

Replace Adam's elementwise denominator on factor matrices with a gauge-equivariant denominator shared across latent coordinates, or with a left preconditioner constructed from Gram matrices. The resulting update treats all equivalent bases of the latent space identically, allowing the low-rank implicit bias of factorized gradient flow to survive adaptive optimization.

## Mathematical statement

The paper defines the gauge symmetry L(U,V)=f(UV^{    op}) with (U,V)\mapsto(UQ,VQ) for Q\in\mathrm{O}(k), where U\in\mathbb{R}^{d    imes k}, V\in\mathbb{R}^{p    imes k}, and Q is an orthogonal latent-basis rotation. If G_U=
abla_U L and G_V=
abla_V L, then after rotation G_U'=G_UQ and G_V'=G_VQ. An optimizer is gauge-equivariant when its updates satisfy \Delta_U(UQ,VQ,G_UQ,G_VQ)=\Delta_U(U,V,G_U,G_V)Q and the analogous equation for V. Coordinate-wise Adam violates this because its elementwise second-moment denominator does not commute with right multiplication by Q. A shared-scalar adaptive rule is equivariant because Frobenius norms satisfy \|G_UQ\|_F=\|G_U\|_F. More generally, the paper's structure theorem characterizes memoryless equivariant rules as Gram-determined left preconditioners, with P_U and P_V depending on gauge-invariant quantities such as UU^{    op} and VV^{    op}.

## Key formulas

- $$L(U,V)=f(UV^{\top}),\qquad (U,V)\longmapsto(UQ,VQ),\quad Q\in\mathrm{O}(k).$$
- $$\Delta_U(UQ,VQ,G_UQ,G_VQ)=\Delta_U(U,V,G_U,G_V)Q,\qquad \Delta_V(UQ,VQ,G_UQ,G_VQ)=\Delta_V(U,V,G_U,G_V)Q.$$
- $$m^U_t=\beta_1m^U_{t-1}+(1-\beta_1)G_{U,t},\quad m^V_t=\beta_1m^V_{t-1}+(1-\beta_1)G_{V,t},$$
$$s_t=\beta_2s_{t-1}+(1-\beta_2)\frac{\|G_{U,t}\|_F^2+\|G_{V,t}\|_F^2}{dk+pk},$$
$$U_{t+1}=U_t-\eta\frac{\widehat m^U_t}{\sqrt{\widehat s_t}+\varepsilon},\qquad V_{t+1}=V_t-\eta\frac{\widehat m^V_t}{\sqrt{\widehat s_t}+\varepsilon}.$$
- $$P_U=P_U(UU^{\top},VV^{\top}),\qquad P_V=P_V(UU^{\top},VV^{\top}),\qquad U_{t+1}=U_t-\eta P_UG_{U,t},\quad V_{t+1}=V_t-\eta P_VG_{V,t}.$$

## Implementation notes

(1) Integration point: implement this as an optimizer for any factored parameterization, beginning with matrix sensing or a LoRA linear layer W=W_0+UV^{    op}. Do not apply ordinary Adam independently to U and V. Maintain first moments mU and mV, but maintain one scalar second moment s for the complete pair of factors; decoupled weight decay can be applied separately and identically to U and V. (2) Pseudocode: initialize mU=mV=0 and s=0. At each step compute L=f(U@V.T), gradients gU,gV; update mU=beta1*mU+(1-beta1)*gU and mV=beta1*mV+(1-beta1)*gV; update s=beta2*s+(1-beta2)*(sum(gU*gU)+sum(gV*gV))/(numel(U)+numel(V)); bias-correct the moments; set U=U-eta*mUhat/(sqrt(shat)+epsilon) and V=V-eta*mVhat/(sqrt(shat)+epsilon). (3) The mathematical property to verify is exact gauge equivariance: sample an orthogonal Q, rotate U,V and their gradients, and compare the rotated original update with the update recomputed after rotation. The scalar s is invariant because Frobenius norms are right-orthogonally invariant. A stronger variant can use P_U=(UU^{    op}+VV^{    op}+lambda I)^{-1/2} and an analogous P_V, approximated with diagonal or low-rank methods. (4) First cheap experiment: use underdetermined matrix sensing with planted rank-r W_true, factor width k>r, and Gaussian measurements. Compare Adam, SGD, shared-scalar Adam, and Muon at matched FLOPs. Track recovery error, effective rank of UV^{    op}, training loss, and rotated-initialization consistency. The expected signal is lower final recovery error and smaller effective rank for shared-scalar Adam than Adam at matched training loss, while equivalent rotated initializations produce nearly identical trajectories. A failure criterion is no improvement over Adam or an equivariance error above 1e-5 in float32.

## Verification

- Status: mech_ok_no_baseline
- Mechanism evidence: yes
- Mechanism confirmed: yes
- Practical verdict: no_effect
- Verdict: Implemented shared-scalar Adam for factored matrices and verified gauge equivariance numerically. Across rotations, scalar invariance error was at most 1.52e-16, first-step shared-update error at most 1.43e-17, and eight-step error at most 1.02e-17; coordinate Adam reached 7.39e-2 error under mixed rotations. In matrix sensing, shared-scalar Adam improved over coordinate Adam on measurement MSE, recovery error, and effective rank, although SGD had lower recovery error in this single run.

### Mechanism check

- Verdict: Implemented shared-scalar Adam for factored matrices and verified gauge equivariance numerically. Across rotations, scalar invariance error was at most 1.52e-16, first-step shared-update error at most 1.43e-17, and eight-step error at most 1.02e-17; coordinate Adam reached 7.39e-2 error under mixed rotations. In matrix sensing, shared-scalar Adam improved over coordinate Adam on measurement MSE, recovery error, and effective rank, although SGD had lower recovery error in this single run.
- Confidence: 9/10
- Limitations: Only one small fixed-seed matrix-sensing problem and one learning-rate/step configuration were tested; no multi-seed statistical study, Muon comparison, Gram-matrix preconditioner, GPU benchmark, or matched-FLOP wall-clock analysis was performed.

### Practical benchmark

- Paired seeds: 8
- Baseline mean: 0
- Idea mean: 0
- p-value: 0.89185
- Paired wins: 5/8
- Benchmark verdict: no measurable effect

## Artifacts

- [bench_report.json](https://synthcore.org/code/654/bench_report.json)
- [experiment.py](https://synthcore.org/code/654/experiment.py)
- [report.md](https://synthcore.org/code/654/report.md)
- [report_bench_2026-08-31T210212.md](https://synthcore.org/code/654/report_bench_2026-08-31T210212.md)
- [results.json](https://synthcore.org/code/654/results.json)
- [stage2_bench.py](https://synthcore.org/code/654/stage2_bench.py)
- [Download all files as ZIP](https://synthcore.org/download/654)

## Disclaimer

AI-generated research hypothesis, automatically tested. Not peer-reviewed.