# Rotation-aware low-rank Adam

- ID: 450
- Canonical URL: https://synthcore.org/idea/450/rotation-aware-low-rank-adam
- API JSON: https://synthcore.org/api/idea/450.json
- API Markdown: https://synthcore.org/api/idea/450.md
- Verification status: failed_benchmark
- Source: [arXiv:2607.05872](https://arxiv.org/abs/2607.05872)
- Category: optimization
- Solves: stability, memory, accuracy
- ML areas: optimizer, fine-tuning, training, memory
- Math tags: optimization, linear-algebra, geometry, stochastic-processes
- Ratings: usefulness 7/10; difficulty 4/10; novelty 7/10

## Idea description

Treat every low-rank basis refresh as a change of coordinates instead of assuming that old optimizer coordinates remain aligned with the new basis. Transport the first moment with the basis-overlap matrix, but collapse the second moment to a rotation-blind isotropic estimate rather than applying the same coordinate transformation to elementwise squared moments. This should eliminate second-moment staleness while preserving the memory savings of low-rank optimization.

## Mathematical statement

Adam maintains projected first and second moments m_t and v_t according to m_t=beta_1 m_{t-1}+(1-beta_1) tilde_g_t and v_t=beta_2 v_{t-1}+(1-beta_2) tilde_g_t^2, where tilde_g_t is the gradient in the current rank-r coordinate system and the square is elementwise. The paper gives the relaxation times tau_m=1/(1-beta_1) and tau_v=1/(1-beta_2), which are approximately 10 and 1000 for beta_1=0.9 and beta_2=0.999. Let U_old and U_new be d by r matrices with orthonormal columns, representing the old and new subspaces, and define the coordinate-overlap matrix Q=U_new^T U_old. A vector represented by old coordinates m_old is represented in the new coordinates by Q m_old; when the subspaces coincide this is an exact orthogonal coordinate change. In contrast, elementwise squaring is not equivariant: generally (Q m)^2 is not equal to Q^2 m^2, so blindly carrying v_old or rotating it coordinatewise creates basis-dependent stale curvature. Under an isotropic projected-gradient model with covariance E[tilde_g tilde_g^T]=sigma^2 I_r, the rotation-blind second-moment estimator is the scalar average v_bar=(1/r) sum_i (v_old)_i and v_new=v_bar 1_r, where 1_r is the r-dimensional all-ones vector. The proposed update combines m_new=Q m_old with v_new=v_bar 1_r and then applies standard Adam bias correction and parameter updates.

## Key formulas

- $$m_t=\beta_1m_{t-1}+(1-\beta_1)\tilde g_t,\qquad v_t=\beta_2v_{t-1}+(1-\beta_2)\tilde g_t^{\,2}.$$
- $$\tau_m=\frac{1}{1-\beta_1},\qquad \tau_v=\frac{1}{1-\beta_2},\qquad (\beta_1,\beta_2)=(0.9,0.999)\Rightarrow(\tau_m,\tau_v)\approx(10,1000).$$
- $$Q=U_{new}^{\top}U_{old},\qquad m_{old\to new}=Qm_{old},\qquad (Qm_{old})^2\neq Q^2(m_{old}^{\,2})\ \text{in general}.$$
- $$\bar v=\frac{1}{r}\sum_{i=1}^{r}(v_{old})_i,\qquad v_{old\to new}=\bar v\,\mathbf 1_r,\qquad \hat m_t=\frac{m_t}{1-\beta_1^t},\quad \hat v_t=\frac{v_t}{1-\beta_2^t}.$$

## Implementation notes

1. Integration point: modify the optimizer state of a GaLore-style linear layer. At a refresh step, the layer has an old orthonormal projection basis U_old and computes a new basis U_new from the current gradient or minibatch. Store the low-rank first moment m and second moment v in r-dimensional coordinates, rather than storing full d-dimensional states. Before processing the first gradient in the new coordinates, transport the state. 2. Pseudocode: `g_tilde = U_new.T @ grad`; `Q = U_new.T @ U_old`; `m = Q @ m`; `v_scalar = mean(v)`; `v = v_scalar * ones(r)`; `m = beta1*m + (1-beta1)*g_tilde`; `v = beta2*v + (1-beta2)*(g_tilde*g_tilde)`; `mhat = m/(1-beta1**step)`; `vhat = v/(1-beta2**step)`; `delta = learning_rate*mhat/(sqrt(vhat)+epsilon)`; `W -= U_new @ delta`, or use the exact parameter-update convention of the selected GaLore implementation. After the refresh, set `U_old = U_new`. 3. The overlap Q and transported momentum are computed directly from the coordinate-change rule. The scalar second-moment replacement is the isotropic-gradient approximation; estimate its validity by logging the per-coordinate variance of g_tilde and comparing it with mean(v). Test a variant with beta2=0.99 because the relaxation-time argument predicts that shorter memory is safer when bases refresh frequently. 4. First cheap experiment: train Pythia-160M or a small Transformer on WikiText-103 using rank r=32 or 64 and refresh intervals T in {100, 500, 1000}. Compare canonical GaLore, GaLore with only momentum transport, GaLore with isotropic second moments, and the complete rotation-aware optimizer with beta2 in {0.99, 0.999}, using identical seeds and FLOPs. Measure training loss and perplexity at equal steps, update-norm spikes immediately after refreshes, and optimizer-state memory. Success is faster or smoother loss descent, lower post-refresh spikes, and improved perplexity at the same low-rank-state memory. Failure is no improvement over canonical GaLore or a large degradation from isotropic second moments.

## Verification

- Status: failed_benchmark
- Mechanism evidence: yes
- Mechanism confirmed: no
- Practical verdict: no_effect
- Verdict: Built a deterministic NumPy toy verification of rotation-aware low-rank Adam, including exact first-moment coordinate transport, non-equivariant elementwise second moments, isotropic replacement, and a rotating quadratic optimizer comparison. The mechanism manifested: transport reconstruction error was 7.3e-16 versus the predicted zero; second-moment discrepancy was exactly 0 at rank 1 and increased from 0.323 at r=2 to 0.769 at r=32; isotropic covariance error was consistently below blindly carried canonical error. Optimizer-level results were mixed: rotation-aware Adam improved final loss at refresh=20,beta2=0.99 (0.01713 vs 0.01798) but degraded it at beta2=0.999 (0.05342 vs 0.03963), so the mechanism is supported but broad optimization superiority is not established.

### Mechanism check

- Verdict: Built a deterministic NumPy toy verification of rotation-aware low-rank Adam, including exact first-moment coordinate transport, non-equivariant elementwise second moments, isotropic replacement, and a rotating quadratic optimizer comparison. The mechanism manifested: transport reconstruction error was 7.3e-16 versus the predicted zero; second-moment discrepancy was exactly 0 at rank 1 and increased from 0.323 at r=2 to 0.769 at r=32; isotropic covariance error was consistently below blindly carried canonical error. Optimizer-level results were mixed: rotation-aware Adam improved final loss at refresh=20,beta2=0.99 (0.01713 vs 0.01798) but degraded it at beta2=0.999 (0.05342 vs 0.03963), so the mechanism is supported but broad optimization superiority is not established.
- Confidence: 8/10
- Limitations: This is a small NumPy toy rather than a transformer or real GaLore implementation. It does not measure GPU speed, actual optimizer-state memory, language-model perplexity, minibatch gradient noise, multiple random seeds, or statistically reliable training improvements.

### Practical benchmark

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

## Artifacts

- [bench_report.json](https://synthcore.org/code/229/bench_report.json)
- [report.md](https://synthcore.org/code/229/report.md)
- [report_bench_2026-08-30T211214.md](https://synthcore.org/code/229/report_bench_2026-08-30T211214.md)
- [results.json](https://synthcore.org/code/229/results.json)
- [rotation_aware_bench.py](https://synthcore.org/code/229/rotation_aware_bench.py)
- [rotation_aware_experiment.py](https://synthcore.org/code/229/rotation_aware_experiment.py)
- [Download all files as ZIP](https://synthcore.org/download/229)

## Disclaimer

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