GaLore's stale-projector descent retention is an energy-weighted sum of squared cosines between each current gradient direction and the stale subspace. It is decoupled from the principal-angle frame distance used to schedule refreshes, and much larger, so the standard drift metric overstates the descent damage by orders of magnitude and refreshes too early.
GaLore (Zhao et al. 2024) saves optimizer memory by projecting the gradient onto the top-r
left-singular subspace of a periodically recomputed gradient. Between recomputations the
projector P0 is stale while the gradient rotates. The effective update is P0 P0^T G_t,
and its descent alignment <G_t, P0 P0^T G_t> = ||P0^T G_t||_F^2 >= 0 never ascends.
Write the current gradient's SVD with energy weights w_j = sigma_j^2 / ||G_t||_F^2 and
cos^2(phi_j) = ||P0^T u_j||^2, the alignment of the j-th current singular direction with
the stale subspace. Then
rho_stale = ||P0^T G_t||_F^2 / ||G_t||_F^2 = sum_j w_j cos^2(phi_j) ,
an energy-weighted convex combination of per-direction squared cosines (matched to the
direct value at 6e-17). This is not the principal-angle distance between the two top-r
frames, and the difference is the whole point: the high-energy directions that dominate
rho_stale rotate slowly, while the tail directions that spike the frame distance carry
almost no energy.
A fresh recompute retains the Eckart-Young top-r energy rho_fresh = sum_{j<=r} w_j. The
gap decomposes exactly:
rho_fresh - rho_stale = sum_{j<=r} w_j sin^2(phi_j) - sum_{j>r} w_j cos^2(phi_j) ,
a top-r loss minus a tail-recapture. The tail term is strictly positive whenever the stale
basis overlaps any tail direction, so the effective staleness penalty is always strictly
smaller than the naive lost-top-r energy (measured 0.484 versus a 0.535 top-r loss, a
0.051 recapture).
GaLore refreshes are timed off a max or mean principal angle between the stale and current
top-r subspaces. But rho_stale is energy-weighted. In a run where the high-energy
directions are stable and one tail-boundary direction rotates 85 degrees, the max
principal angle is 86.8 degrees, so a cos^2(theta) metric predicts retention 0.003,
"descent almost entirely lost". The actual retained descent is 0.884: the frame-distance
metric understates the retention by 291x. Scheduling refreshes off the drift metric
triggers them far too early.
staleness.py: the retained-descent fraction, its energy-weighted-cosine form, and the nonnegative descent alignment.decomposition.py: the Eckart-Young fresh retention, the top-loss-minus-tail-recapture gap, and the frame-distance mis-prediction.test_stalecos.py: the energy-weighted identity, nonnegativity, the Eckart-Young fresh case, the gap decomposition, and the metric understatement.
python staleness.py
python decomposition.py
python test_stalecos.py
The pieces are classical: Parseval and subspace projection give the retained fraction,
Eckart-Young gives the fresh retention, and principal angles are Jordan's 1936 construct.
This repository claims none of them. What is here is the GaLore-specific reading: the
retained descent under a stale projector is the energy-weighted per-direction cosine
sum_j w_j cos^2(phi_j), which is decoupled from and much larger than the frame-distance
drift metric used to schedule refreshes, and the staleness penalty is strictly below the
lost top-r energy by an always-positive tail-recapture term. The GaLore paper proves
convergence under a low-rank-gradient assumption and picks the refresh interval
heuristically; it states neither. The measurements here are single full-batch trajectories
in a mildly concentrated spectrum; how the mis-prediction scales with spectral concentration
and whether refreshing on the energy-weighted metric cuts refresh count end to end are the
named extensions, and the falsifier is any regime where retained descent tracks the frame
distance rather than the energy-weighted cosine.