Low Rank Adaptive Friction for Memory-Efficient LLM Pretraining
Rajit Rajpal ⋅ Benedict Leimkuhler
Abstract
iKFAD is a recently proposed optimiser that replaces adaptive learning rates with adaptive friction in the momentum dynamics, yet performs as well as Adam. Its full friction tensor $\xi\in\mathbb{R}^{m\times n}$ has the same $\mathcal{O}(mn)$ memory cost per layer as Adam's second-moment buffer. Here we introduce Rank-1 iKFAD (R-iKFAD), which represents this tensor using row and column statistics. This reduces the friction state dimension to $\mathcal{O}(m+n)$ values per layer and approximately halves the size of iKFAD's optimiser state. Across GPT2-Nano, TinyViT, DistilBERT, and GPT2-S, R-iKFAD matches or improves on iKFAD while using nearly half the optimiser memory. We also analyse its continuous-time dynamics. Under strong convexity, we prove exponential convergence when $\gamma>0$ and convergence when $\gamma=0$, as used in our experiments. A formal averaging calculation predicts $t^{-1}$ decay when $\epsilon_{\mathrm{stab}}=0$ and $t^{-1/2}$ decay when $\epsilon_{\mathrm{stab}}>0$. Numerical results are consistent with both predictions.
Chat is not available.
Successful Page Load