Rotate All, Store Few: Decoupling Coverage from Storage for Memory-Efficient Training

Adam improves when the gradient is normalized in its singular basis, which we call $\text{rotation}$. Full-rank rotation exceeds the memory of Adam by storing a basis and states for every direction, whereas low-rank gradient projection rotates only the top-$r$ directions but degrades sharply as $r$ decreases. This is because a single rank $r$ sets both the $\text{rotation coverage}$, the number of rotated directions, and the $\text{state storage}$, the number of directions with optimizer states. The two decisions nonetheless need not be coupled. The rotation coverage should be full since the whitened update magnitude stays nearly flat along the spectral index, whereas the state storage can be reduced since adjacent second moments have high cosine similarity. We, therefore, propose $\text{COAST}$, which decouples rotation $\underline{\text{CO}}$verage $\underline{\text{A}}$nd state $\underline{\text{ST}}$orage. $\text{COAST}$ rotates every direction but maintains optimizer states for only $r$ stored directions and estimates the updates of the rest from their stored neighbors. On LLaMA pre-training, $\text{COAST}$ outperforms Adam with $0.42$–$0.63\times$ Adam's optimizer memory. It also improves on low-rank methods at similar memory. Halving the storage budget changes perplexity by less than 1% at every scale.