Low-Rank Friction for Memory-Efficient Transformer Pretraining

Rajit Rajpal; Benedict Leimkuhler

stat.MLSep 24, 2026

Low-Rank Friction for Memory-Efficient Transformer Pretraining

Authors: Rajit Rajpal, Benedict Leimkuhler

Organizations: School of Mathematics University of Edinburgh Edinburgh, UK EH9 3FD

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 limitation is that the full friction tensor $ξ\in\mathbb{R}^{m\times n}$ carries the same $\mathcal{O}(mn)$ memory overhead per layer as Adam's second-moment buffer. Here we replace iKFAD's friction tensor $ξ$ with a rank-1 outer-product factorisation built from row and column momentum statistics, resulting in Rank-1 iKFAD (R-iKFAD). This reduces the friction memory footprint from $\mathcal{O}(mn)$ to $\mathcal{O}(m+n)$ per layer, which approximately halves iKFAD's total optimiser state. Despite this reduction, R-iKFAD maintains parity in performance with iKFAD: experiments on GPT2-Nano, TinyViT, DistilBERT and GPT2-S confirm that it matches or exceeds iKFAD while nearly halving the memory footprint and remaining comparably robust to hyperparameters. We analyse the continuous-time dynamics in two damping regimes. For linear damping ( $γ>0$ ) we prove exponential convergence under strong convexity. For $γ=0$ , the preferred option in our experiments, the friction is generated entirely from past momentum and switches off as the momentum vanishes, so geometric convergence cannot be shown. We nonetheless prove convergence to the minimiser, together with matching upper and lower bounds on the energy: of order $t^{-1}$ when the regularisation scale $ε_{\mathrm{stab}}$ is zero, and of order $t^{-1/2}$ when it is positive. To our knowledge this is the first convergence rate for a rank-1 factored optimiser in continuous time, and the first such result that does not require positive damping.

Figures & tables

Panel of Figure 1 — Figure 1: The same rank-1 compression applied in two different equations. Adafactor-m factorizes Adam’s second moment $\zeta$ which scales the position update while R-iKFAD factorizes iKFAD’s friction $\xi$ which damps the momentum update. Both reduce the adapted state per layer from $\mathcal{O}(mn)$ to $\mathcal{O}(m+n)$ while keeping full momentum, so total auxiliary memory falls from $2N$ to roughly $N$ .

Figure 2: Energy decay from BACD splitting with

m=4

n=3

\gamma=0

\alpha=\mu=1

h=0.02

, and final time

10^{7}

. Both runs have

R(0)=C(0)=0

. The blue dashed curve has

\epsilon_{\mathrm{stab}}=0

, and the orange dash-dotted curve has

\epsilon_{\mathrm{stab}}=0.1

. The solid lines have slopes

-1

and

-1/2

, respectively.

$\epsilon_{\mathrm{stab}}=0$ (proved exponent $-1$ )
	$[10^{3},10^{4}]$	$[10^{4},10^{5}]$	$[10^{5},10^{6}]$	$[10^{6},10^{7}]$
$m=n=2$	$-0.955$	$-0.980$	$-0.986$	$-0.989$
$m=4,n=3$	$-0.991$	$-0.996$	$-0.996$	$-0.998$
$m=6,n=5$	$-0.993$	$-0.997$	$-0.998$	$-0.999$
$\epsilon_{\mathrm{stab}}=0.1$ (proved exponent $-1/2$ )
$m=n=2$	$-0.499$	$-0.497$	$-0.496$	$-0.497$

Table 1: Local slopes of

\log\mathcal{E}

against

\log t

on successive decades, from the BACD splitting at

h=0.02

with

\alpha=\mu=1

, integrated to

T=10^{7}

. Both exponents are independent of the matrix dimensions.

Figure 3: Loss curves for all four benchmarks, averaged over 10 random seeds (GPT2-S: a single seed). R-iKFAD0 matches or improves upon iKFAD0 on all tasks, while using approximately half the optimizer memory. Both methods are competitive with Adam and Adafactor-m, which require comparable or greater memory. Shaded regions denote one standard deviation.

Dataset	Adam	Adafactor-m	iKFAD0	R-iKFAD0 (ours)
GPT2-Nano	$1.643\pm 0.008$	$1.661\pm 0.007$	$1.640\pm 0.008$	$\mathbf{1.633\pm 0.011}$
TinyViT	$\mathbf{0.593\pm 0.015}$	$0.596\pm 0.010$	$0.606\pm 0.011$	$0.606\pm 0.012$
DistilBERT	$0.292\pm 0.005$	$0.280\pm 0.007$	$0.269\pm 0.005$	$\mathbf{0.268\pm 0.003}$
GPT2-S	$3.301$	$3.287$	$3.263$	$\mathbf{3.242}$

Table 2: Test loss at the checkpoint selected by lowest validation loss (mean

\pm

std over 10 seeds, GPT2-S: single seed due to computational limitations). The test set was never used for model or hyper-parameter selection. Lower is better. R-iKFAD0 matches or improves upon iKFAD0 across all benchmarks while using half the optimizer memory.

Figure 4: Training and test loss for R-iKFAD (

\gamma>0

) vs. R-iKFAD0 (

\gamma=0

) on GPT2-Nano, TinyViT, and DistilBERT (SST-2), averaged over 10 seeds (shaded = std). The two variants are indistinguishable in performance, supporting R-iKFAD0 as the recommended default.

Figure 5: Two-dimensional $\gamma$ – $h$ sensitivity for iKFAD (top) and R-iKFAD (bottom). Each cell shows best performance at a fixed $(h,\gamma)$ grid point. Brighter = better. R-iKFAD retains broad high-performance plateaus.

Figure 6: R-iKFAD’s rank-1 friction against the best rank-1 approximation of the same tensor. Left: layer-averaged relative error

\|\xi-\tilde{\xi}\|_{F}/\|\xi\|_{F}

(solid) against the Eckart–Young floor

\sqrt{1-\sigma_{1}^{2}/\sum_{i}\sigma_{i}^{2}}

(dashed). The shaded region is the excess attributable to the construction, and is never wider than

0.85

pp. The horizontal axis is the fraction of each run completed, as the four benchmarks used different step budgets. Right: distribution over layers of the stable rank of

\xi

at the final step, where

1

would be exactly rank-1. Measurement is taken along iKFAD0 trajectories carrying R-iKFAD’s factors as shadow states.

Appendix figures & tables5 assets

Supplementary material from the paper’s appendix.

Appendix

Sub-step	Update Equation
B	$P^{(B)}=P_{n}-h\,\nabla f(X_{n})$
A	$X_{n+1}=X_{n}+h\,P^{(B)}$
C	$\tilde{\xi}_{n}=\dfrac{R_{n}C_{n}^{\top}}{\mathbf{1}_{m}^{\top}R_{n}+\epsilon_{\mathrm{stab}}}$
	$P^{(1)}=P^{(B)}\odot\exp\!\left(-\tfrac{h}{2}\,\tilde{\xi}_{n}\right)$
	$R_{n+1}=\mathrm{e}^{-\alpha h}R_{n}+\dfrac{1-\mathrm{e}^{-\alpha h}}{\mu\alpha}[P^{(1)}]^{2}\mathbf{1}_{n}$
	$C_{n+1}=\mathrm{e}^{-\alpha h}C_{n}+\dfrac{1-\mathrm{e}^{-\alpha h}}{\mu\alpha}([P^{(1)}]^{2})^{\top}\mathbf{1}_{m}$

Appendix

Table 3: Analytical sub-step update rules for the BACD splitting sequence over time step

h

	$h$	$\alpha$	$\mu$	$\beta_{1}$	$\beta_{2}$
TinyViT (CIFAR-10)
iKFAD0	0.0725	0.0899	$1.2643\times 10^{-5}$	–	–
R-iKFAD0	0.0860	0.5327	$9.4765\times 10^{-7}$	–	–
Adam	$5.5311\times 10^{-4}$	–	–	0.8608	0.8852
Adafactor-m	$5.6050\times 10^{-4}$	–	–	0.8907	–
DistilBERT (SST-2)

Appendix

Table 4: Optimized hyperparameters for all experiments. All iKFAD0 and R-iKFAD0 entries have

\gamma=0

. Dashes indicate parameters not used by the method.

Model	$N$	Adam	Adafactor-m	Adafactor	iKFAD	R-iKFAD
TinyViT	4.0M	30.7	15.6	0.2	30.7	15.6 ( $2.0\times$ )
DistilBERT	67.0M	510.8	256.1	0.7	510.8	256.1 ( $2.0\times$ )
GPT2-Nano	0.8M	6.1	3.1	$<0.1$	6.1	3.1 ( $2.0\times$ )
GPT2-S	124.4M	948.9	475.4	0.9	948.9	475.4 ( $2.0\times$ )

Appendix

Table 5: Optimizer auxiliary state (MiB, float32), computed over all model parameters. R-iKFAD achieves approximately

2\times

savings over both iKFAD and Adam because the rank-1 friction factors

(R_{\ell},C_{\ell})

are negligible relative to the full momentum buffer. The Adafactor-m baseline used throughout enables its optional first moment (

\beta_{1}>0

, see Section 4 ) and therefore carries an auxiliary footprint identical to that of R-iKFAD, making the two directly comparable at matched memory. Momentum-free Adafactor is included for reference.

Figure 7: Accuracy averaged over 10 seeds. R-iKFAD0 matches iKFAD0 and Adam. Error bars denote one std.

quantity	parameters varied	mean	spread (max/min)
$t\mathcal{E}/(\alpha\mu)$ , $\epsilon_{\mathrm{stab}}=0$	$\alpha,\mu\in\{\tfrac{1}{2},1,2\}$	$4.69$	$1.29$
$\mathcal{E}\sqrt{t}\big/\bigl(\alpha\mu\sqrt{\epsilon_{\mathrm{stab}}/2}\bigr)$ , $\epsilon_{\mathrm{stab}}>0$	$\alpha,\mu\in\{\tfrac{1}{2},1\}$ , $\epsilon_{\mathrm{stab}}\in\{0.05,0.2\}$	$2.31$	$1.14$
$\mathcal{E}\big/(\alpha\mu\,\epsilon_{\mathrm{stab}})$ at the crossing	$\alpha,\mu\in\{\tfrac{1}{2},1\}$ , $\epsilon_{\mathrm{stab}}\in\{0.05,0.2\}$	$0.87$	$1.21$

Appendix

Table 6: Parameter dependence of the prefactors, evaluated at

t=10^{4}

with

m=3

n=2

. In the first row

\alpha

and

\mu

each range over a factor of

4

, so

\alpha\mu

ranges over a factor of

16

, yet the ratio moves by less than

30\%

Explore similar work

PowerStep: Memory-Efficient Adaptive Optimization via $\ell_p$ -Norm Steepest Descent
May 11, 2026Yao Lu, Dengdong Fan, Shixun Zhang +1Adam Neural Network Training
Pengcheng Laboratory · Peking University
AdaPreLoRA: Adafactor Preconditioned Low-Rank Adaptation
May 9, 2026Ziyun Liu, Fengmiao Bian, Jian-Feng CaiSpectral Preconditioning Low-Rank Structure
Department of Mathematics The Hong Kong University of Science and Technology
Gefen: Optimized Stochastic Optimizer
Jun 11, 2026Nadav Benedek, Tomer Koren, Ohad FriedScaling Laws
Reichman University · Tel Aviv University, Google Research · Google Research

May 11, 2026, cs.LGJ/K move · Enter open · S save

PowerStep: Memory-Efficient Adaptive Optimization via $\ell_p$ -Norm Steepest Descent

Yao Lu, Dengdong Fan, Shixun Zhang, Yonghong Tian

Pengcheng Laboratory · Peking University

Adam Neural Network Training

Open paper AskPDF arXiv

Adaptive optimizers such as Adam are standard for training Transformers, but storing gradient first and second moments incurs substantial memory overhead. We introduce PowerStep, a memory-efficient optimizer that achieves coordinate-wise adaptivity without storing second-moment statistics. Motivated by

\ell_p

-norm steepest descent, PowerStep applies a signed-power transform directly to one momentum buffer. We establish a finite-horizon stationarity bound for exact, unregularized updates, with an

O(1/\sqrt{T})

term and a noise-dependent residual. Experiments on Transformers from 124M to 235B parameters show competitive validation quality while halving

\texttt{fp32}

optimizer-state memory relative to AdamW. Combined with uniform

\texttt{int8}

quantization, PowerStep remains numerically stable and reduces optimizer-state memory by

\sim8\times

compared to

\texttt{fp32}

AdamW. PowerStep thus provides a simple, memory-efficient alternative for large-scale training.

Abstract

Explore similar work

PowerStep: Memory-Efficient Adaptive Optimization via ℓp\ell_pℓp​-Norm Steepest Descent

AdaPreLoRA: Adafactor Preconditioned Low-Rank Adaptation

Gefen: Optimized Stochastic Optimizer

PowerStep: Memory-Efficient Adaptive Optimization via $\ell_p$ -Norm Steepest Descent