✗ Failed on benchmark 2026

Rotation-aware low-rank Adam

Implementation & benchmark of arXiv:2607.05872 — No Subspace to Track: Non-Identifiability and Optimizer State in Low-Rank Training

Usefulness7/10

Difficulty4/10

Novelty7/10

Optimization Stability Memory Accuracy Optimizer Fine tuning Training Memory Optimization Linear algebra Geometry Stochastic processes

Source paper: No Subspace to Track: Non-Identifiability and Optimizer State in Low-Rank Training arXiv:2607.05872 ⓘ · analyzed Aug 30, 2026

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

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.

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}.$$

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.

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-beta1step); vhat = v/(1-beta2step); 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

Failed on benchmark

Stage 1 · Toy mechanism gate: Passed ✓

Stage 2 · Mechanism transferred to benchmark: Not reproduced ✗

Stage 2 · Practical benchmark result: No measurable effect

Methodology: 8 paired seeds, permutation test (p<0.05) against a learning-rate-tuned baseline. How verification works

Stage 1 — Mechanism check agent confidence 8/10

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.

Agent confidence: 8/10

Baseline: Canonical carried-coordinate second-moment relative covariance error ranged from 0.205 to 0.497 across tested ranks and beta2 values. Rotating-quadratic final losses were 2.72e-5, 0.01798, and 7.924 for refresh intervals 5, 20, and 100 at beta2=0.99; corresponding beta2=0.999 losses were 0.287, 0.03963, and 7.925.
Idea: Isotropic second-moment relative covariance error ranged from 0.104 to 0.216 and was consistently lower than canonical. Rotating-quadratic final losses were 2.95e-5, 0.01713, and 7.953 for refresh intervals 5, 20, and 100 at beta2=0.99; corresponding beta2=0.999 losses were 0.288, 0.05342, and 8.040. Mean refresh update norms were lower for the proposed method in every tested setting.

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.

How to run: python3 rotation_aware_experiment.py > results.json

Stage 2 — Benchmark vs. tuned baseline (arXiv:2607.05872) Sep 2, 2026

No significant difference vs. the tuned baseline (+0.2%, p=0.3209, wins 3 / 8 paired seeds).

Benchmark: Tabular regression (Friedman)
Model: mlp_tiny
Paired seeds: 8
Baseline mean: 199.4986
Idea mean: 199.8595
Effect (Δ): 0.3609 (+0.2%; negative = idea better)
Wins: 3 / 8 paired seeds
p-value: 0.3209 (permutation test, 20 000 shuffles)
Smallest detectable effect: ±0.4%
Mechanism: Not confirmed ✗
Practical effect: No effect
Baseline tuning: swept over 6 configs

Limitations:

Only the registered built-in tabular Friedman#1 optimizer track was tested, with mlp_tiny, 400 training samples, 400 test samples, 12 epochs, rank 8, and refresh interval 5. GPU/CPU timing, optimizer-state memory, larger models, alternative refresh intervals, and GaLore-scale language-model training were not tested.

How to run: python3 rotation_aware_bench.py

Verdict computed by deterministic test code from paired-seed statistics — not by the language model.

Stage 2 — Benchmark Aug 30, 2026 latest Failed ✗

Agent confidence: 9/10

The complete benchmark used the registered built-in tabular track with paired eight-seed training, a six-configuration baseline sweep, and an equal-sized idea sweep. Rotation-aware low-rank Adam slightly worsened mean test MSE (199.8595 versus 199.4986), with paired delta +0.3609 and permutation p=0.32085, so the required significant-win criterion was not met. The model-derived mechanism signature was also not confirmed because independently refreshed subspaces produced observed momentum reconstruction error 2.0965 rather than the zero predicted for an exact within-subspace coordinate rotation.

Baseline: Best baseline lr=0.006, beta2=0.99, rank=8, refresh=5; full paired mean MSE 199.4985733.
Idea: Best idea lr=0.006, beta2=0.99, rank=8, refresh=5; full paired mean MSE 199.8595085; paired delta +0.3609352; permutation p=0.32085; 3 of 8 paired wins.

Benchmark result