Low-Rank Friction for Memory-Efficient Transformer Pretraining
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 ξ∈ℝm× n carries the same 𝒪(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 𝒪(mn) to 𝒪(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 ε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.
Related
- Jacobian-Free Unrolling: memory efficient unrolled network for 3D SPECT Reconstruction
- Efficient computation of the asymptotics of extensive-rank HCIZ integrals
- Efficient Low-Rank Approximation of Quaternion Matrices via Adaptive Cross Approximation
- Rank-Efficient LoRA via Joint Tangent-Space Optimization under Isotropic Curvature
- PowerStep: Memory-Efficient Adaptive Optimization via ℓp-Norm Steepest Descent
BibTeX
@misc{rajpal2026lowrank,
title = {{Low-Rank Friction for Memory-Efficient Transformer Pretraining}},
author = {Rajit Rajpal and Benedict Leimkuhler},
howpublished = {arXiv},
year = {2026},
month = sep,
doi = {10.48550/arxiv.2609.30342},
eprint = {2609.30342},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2609.30342},
note = {CC BY 4.0},
}
RIS
TY - UNPB
TI - Low-Rank Friction for Memory-Efficient Transformer Pretraining
AU - Rajit Rajpal
AU - Benedict Leimkuhler
PY - 2026
DA - 2026/09/28
PB - arXiv
DO - 10.48550/arxiv.2609.30342
UR - https://arxiv.org/abs/2609.30342
AB - 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 ξ∈ℝ^m× n carries the same 𝒪(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 𝒪(mn) to 𝒪(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 ε_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.
ER -