PaperlayerView exact revision on arXiv

Source-linked AI summary

Low-Rank Friction for Memory-Efficient Transformer Pretraining

Rajit Rajpal, Benedict Leimkuhler

arXiv:2609.30342v2stat.MLcs.LG

TL;DR

iKFAD matches Adam’s performance but stores a full friction tensor with the same per-layer O(mn) memory overhead as Adam’s second-moment buffer. This paper introduces R-iKFAD, which factors friction into a rank-1 form while retaining momentum. It reports benchmark parity with iKFAD at roughly half the auxiliary optimizer memory, and proves convergence results in both damped and undamped regimes.

  • Problem

    iKFAD’s full friction tensor has the same O(mn) per-layer memory overhead as Adam’s second-moment buffer, leaving a need to reduce friction state without discarding momentum.

    Source: §1 · p. 2

  • Method

    R-iKFAD replaces iKFAD’s full friction tensor with a rank-1 estimate based on row and column momentum statistics while retaining full momentum.

    Source: §2 · p. 2

  • Results

    Across four benchmarks, R-iKFAD matches iKFAD in convergence speed, final performance, and hyperparameter robustness at roughly half the auxiliary optimizer memory.

    Source: §5 · p. 9

  • Takeaways & Limitations

    The analysis proves exponential convergence under strong convexity for γ > 0 and convergence to the minimizer for γ = 0, with energy rates t^-1 when ε_stab = 0 and t^-1/2 when ε_stab > 0.

    Source: §5 · p. 9

    The formal crossover calculation relies on time-scale separation and an averaging closure, although Theorem 3 obtains the same exponents without those assumptions.

    Source: §3.1 · p. 4

Abstract

from arXiv · show

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 $\xi\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 $\xi$ 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 ($\gamma>0$) we prove exponential convergence under strong convexity. For $\gamma=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 $\epsilon_{\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.

1 Introduction

Adaptive optimizers can impose substantial auxiliary-memory costs, while continuous-time dynamics offer a way to analyze optimizer updates. iKFAD moves coordinate-wise adaptation from the position equation to momentum through dynamic friction, but its full friction tensor leaves the memory gap that motivates R-iKFAD.

1 Introduction

  • Adam’s two auxiliary moment states require 2N scalars, while Adafactor reduces per-layer state from O(mn) to O(m+n) by factoring the second moment and eliminating the first-moment buffer.The paper frames Adam’s auxiliary storage as a large-model training bottleneck.Source: §1 · p. 1
  • Continuous-time ODEs represent optimizer updates as discrete time-stepping approximations and support analysis of their dynamics.Source: §1 · p. 1
  • iKFAD replaces constant damping γ with coordinate-wise friction ξ, moving adaptation into the momentum dynamics.Source: §1 · pp. 1–2
  • Because iKFAD stores a full friction tensor alongside momentum, its auxiliary state is 2N scalars, matching Adam’s storage and motivating a factored counterpart.R-iKFAD applies a rank-1 factorisation to friction, reducing its auxiliary state while retaining full momentum.Source: §1 · p. 2

2 Methodology: Rank-1 iKFAD

R-iKFAD applies rank-1 factorisation to iKFAD’s friction tensor while retaining full momentum, reducing the adapted state from O(mn) to O(m+n). The method uses a split dynamics discretization and requires approximately N auxiliary scalars, roughly half the memory of iKFAD or Adam.

Rank-1 factorisation

  • The stabilizer ε_stab is a small regularization coefficient that ensures numerical stability near R = 0.Source: §2 · p. 2
  • R-iKFAD replaces iKFAD’s full friction tensor with a rank-1 estimate built from aggregate squared momenta, which damps momentum while preserving full momentum evolution.Adafactor-m instead factors the second moment from squared gradients and applies it in the position equation; both reduce adapted state from O(mn) to O(m+n).Source: §2 · p. 2

Splitting Discretization

  • A geometric-integration splitting decomposes the dynamics into four components, with three integrated analytically and the fourth approximated by an extra step.The cited geometric numerical integration foundation is Leimkuhler and Reich [2005].Source: §2 · pp. 2–3
  • Sub-flows A, B, and D implement position drift, a gradient kick, and linear friction; D becomes the identity when γ = 0.For C, momentum is damped for half a step, factors are updated for a full step with momentum held fixed, then momentum is damped again.Source: §2 · p. 3
  • The applied first-order Lie–Trotter sequence is B → A → C → D, and position updates use the gradient-kicked momentum.Appendix A’s Table 3 gives explicit analytical update rules over time step h.Source: §2 · p. 3

Memory Complexity

  • R-iKFAD requires approximately N auxiliary scalars, effectively halving optimizer memory relative to iKFAD and Adam; non-2D parameters use element-wise iKFAD with negligible additional state.Adafactor-m has similar auxiliary memory and is the primary experimental baseline.Source: §2 · p. 3

3 Theoretical Results

The section establishes convergence results for both damping regimes, including algebraic rates when γ=0. Numerical tests recover the predicted decay powers, while related-work comparisons distinguish the evolving-friction setting from earlier analyses.

3 Theoretical Results

  • For γ>0, a Lyapunov argument proves exponential convergence when f is strongly convex and has an Lf-Lipschitz gradient.The theorem assumes positive γ, α, µ, and εstab, with componentwise nonnegative initial factors.Source: §3 · p. 4
  • For γ=0, every solution converges to (X⋆, 0, 0, 0), where X⋆ is the unique minimizer of f.The theorem assumes strong convexity, a locally Lipschitz gradient, positive α and µ, and componentwise nonnegative initial factors.Source: §3 · p. 4
  • At γ=0, E(t) has matching upper and lower bounds of order (1+t)^-1 for εstab=0 and (1+t)^-1/2 for εstab>0.For εstab=0, the theorem additionally assumes R(0)=C(0)=0; the regimes cross when E is of order αµεstab.Source: §3 · p. 4

3.1 Formal rate calculation

  • The formal averaging calculation identifies the decay exponents using time-scale separation and an averaging closure, but Theorem 3 proves them without those assumptions.The reduced regimes meet when E is of order αµεstab.Source: §3.1 · p. 4

3.2 Numerical validation of proved results

  • Numerical tests recover both predicted powers across three matrix dimensions; the approach to slope −1/2 is slower for larger mn.The slower regime begins once E is approximately at or below αµεstab, and the crossover occurs later in larger systems.Source: §3.2 · p. 5, Table 1

3.3 Related Work: Algebraic Rates and Factored Optimisers

  • Earlier continuous-time rate analyses considered prescribed time-dependent damping, while the damping here evolves with factors and can vanish with them.The cited H-Fac continuous-time result uses positive damping; the γ=0 results also carry over to KFAD and iKFAD, which lack stabilizers and therefore have only the t^-1 rate.Source: §3.3 · pp. 5–6

4 Numerical experiments

Across four benchmarks, R-iKFAD0 matches or improves on iKFAD0 while using roughly half the auxiliary optimizer memory. The experiments also compare damping choices, hyperparameter robustness, and the accuracy of the rank-1 friction approximation.

4 Numerical experiments

  • The experiments span GPT2-Nano, TinyViT, DistilBERT, and GPT2-S, with mean and standard deviation over 10 seeds except for single-seed GPT2-S.Hyperparameters were tuned with approximately 80 trials per optimizer per model; the authors note that this fixed trial budget disproportionately favors Adafactor-m, which tunes only two hyperparameters.Source: §4 · p. 6

4.1 Main Results

  • R-iKFAD0 matches iKFAD0 across all four benchmarks while halving auxiliary optimizer state, and Figure 3 shows the loss-curve comparison.The authors characterize the methods as performing like Adam at roughly half Adam’s optimizer memory, without claiming R-iKFAD0 outperforms Adam.Source: §4.1 · p. 6, Figure 3 · p. 7, Table 2 · p. 7

4.2 Ablation: γ = 0 versus γ > 0

  • Across GPT2-Nano, TinyViT, and DistilBERT (SST-2), adding linear damping (γ > 0) gave no consistent improvement over R-iKFAD0 (γ = 0), which authors recommend as the practical default with one fewer hyperparameter.The γ > 0 variant used the same h, α, and µ settings as R-iKFAD0 and swept γ log-uniformly over [10^-6, 10].Source: §4.2 · pp. 6–7

4.3 Robustness to Hyperparameters

  • R-iKFAD retains broad high-performance regions across the tested γ–h grid; its negative feedback increases friction as momentum grows.The authors explain that this feedback helps account for why γ adds little: constant damping acts in a direction already covered by adaptive damping, without responding to state changes.Source: §4.3 · p. 7, Figure 5 · p. 8

4.4 How Close is ˜ξ to the Best Rank-1 Friction?

  • Across four benchmarks, R-iKFAD’s friction approximation is 0.07–0.85 percentage points above the best rank-1 error; most approximation error comes from the rank-1 restriction.On GPT2-S, the median gap across 74 two-dimensional layers is 0.06 percentage points, and 61 layers are within one percentage point of the minimum.Source: §4.4 · pp. 8–9, Figure 6 · p. 9

5 Conclusion

R-iKFAD factors iKFAD’s friction into a rank-1 outer product, reducing friction state from O(mn) to O(m + n) per layer while retaining a full momentum buffer. It matches iKFAD’s reported performance across four benchmarks, and the paper proves convergence results for both damped and undamped continuous-time dynamics.

5 Conclusion

  • R-iKFAD replaces iKFAD’s full friction tensor with a rank-1 outer product of momentum row and column sums, reducing per-layer friction state from O(mn) to O(m + n).The momentum buffer remains full-sized, while the optimizer’s auxiliary memory is approximately halved.Source: §5 · p. 9
  • Across GPT2-Nano, TinyViT, DistilBERT, and GPT2-S, R-iKFAD matches iKFAD in convergence speed, final performance, and hyperparameter robustness at roughly half the auxiliary optimizer memory.The analysis is continuous-time and deterministic, and assumes strong convexity.Source: §5 · p. 9
  • Under strong convexity, the continuous-time analysis proves exponential convergence for γ > 0 and convergence to the minimizer for γ = 0, with energy rates t^-1 when ϵ_stab = 0 and t^-1/2 when ϵ_stab > 0.Extending the analysis to discrete-time stochastic-gradient dynamics remains open.Source: §5 · p. 9

Reproducibility Statement

The reproducibility materials include code and training configurations for the benchmarks and numerical rate verification. The appendix also documents the BACD splitting updates and its treatment of a coupled sub-flow without a closed-form joint solution.

Reproducibility Statement

  • The repository provides benchmark implementations, training and sweep scripts, environment specifications, visualization code, and MATLAB code for numerical rate verification.Appendix B documents the hyperparameter sweeps and training configurations.Source: Reproducibility Statement · p. 9

A Update Rules for Splitting Discretisation

  • Appendix A gives explicit sub-step updates for the BACD splitting sequence and addresses the sub-flow without a closed-form joint solution.Table 3 presents the analytical sub-step update rules over time step h.Source: A Update Rules for Splitting… · p. 10, Table 3 · p. 11
  • For the coupled momentum and factor sub-flow, the method symmetrically applies friction around an exact exponential factor update; using intermediate momentum for factor updates reduces forcing magnitude.The stated accuracy improvement applies within sub-flow C alone.Source: A Update Rules for Splitting… · p. 10

B Optimal Hyperparameters

Hyperparameters were selected by validation-loss sweeps, with reported results evaluated on held-out test data. The best friction scale µ varied substantially across models and tasks.

B Optimal Hyperparameters

  • The best-performing µ varied by several orders of magnitude across models and tasks, and momentum normalization is proposed as future work to potentially reduce this sensitivity.The authors suggest that the variation is likely related to momentum scale because µ controls adaptive-friction strength.Source: B Optimal Hyperparameters · p. 11
  • Hyperparameters were selected by minimum validation loss during Optuna sweeps, and all reported results use held-out test data.Source: B Optimal Hyperparameters · p. 11
  • Sweeps covered TinyViT, DistilBERT, GPT2-Nano, and GPT2-S, with 80–500 trials and task-specific ranges for step size and friction scales.GPT2-S used a narrower µ range for R-iKFAD0 than for iKFAD0; iKFAD0 and R-iKFAD0 use γ = 0.Source: B Optimal Hyperparameters · pp. 11–12, Table 4 · p. 12

C Memory Analysis

Rank-1 friction factors reduce R-iKFAD's auxiliary-state memory relative to iKFAD and Adam. On the two classification benchmarks, its accuracy curves match those of iKFAD0 and Adam.

C Memory Analysis

  • R-iKFAD achieves approximately 2× auxiliary-state memory savings over both iKFAD and Adam, with its friction factors negligible beside the full momentum buffer.Adafactor-m with β1 > 0 has auxiliary memory identical to R-iKFAD for a matched-memory comparison.Source: Table 5 · p. 12

D Accuracy Curves

  • R-iKFAD0 matches iKFAD0 and Adam in accuracy averaged over 10 seeds.The accuracy curves correspond to the two classification benchmarks; error bars show one standard deviation.Source: Figure 7 · p. 13, D Accuracy Curves · p. 13

E Prefactor Dependence

The tested prefactor ratio changes modestly despite a wider variation in α and µ.

E Prefactor Dependence

  • With α and µ each varied by a factor of 4, their product αµ varied by 16× while the prefactor ratio moved by less than 30%.The table evaluates parameter dependence at t = 10^4 with m = 3 and n = 2.Source: Table 6 · p. 13

F Numerical Setup for Section 3.2

The numerical checks compare BACD splitting with ODE integration and smaller step sizes. The analysis establishes bounded trajectories and exponential convergence under positive damping and strong convexity.

F Numerical Setup for Section 3.2

G.1 Dynamical properties and boundedness

  • Under the stated coercivity and smoothness assumptions, solutions remain in a compact set for all t ≥ 0, and the friction tensor is bounded.The boundedness result assumes componentwise nonnegative initial R and C.Source: §G.1 · p. 14

G.2 Geometric convergence

  • For γ > 0, the theorem gives exponential convergence when f is strongly convex with Lipschitz gradient, α, µ, and ε_stab are positive, and initial factors are nonnegative.Source: §G.2 · p. 15
  • When γ = 0, the cross-term control used in the exponential argument is absent, so that regime is treated separately.Source: §G.2 · p. 16

G.3 The case γ = 0

With γ = 0, the strongly convex system converges to the unique minimizer when α, μ, and εstab are positive and the initial factors are nonnegative.

G.3 The case γ = 0

  • Experiments found γ = 0 optimal or near-optimal, motivating the convergence analysis for this damping regime.Source: §G.3 · p. 16
  • Every solution converges to the unique minimizer when f is strongly convex with locally Lipschitz gradient, α, μ, εstab > 0, and initial factors are componentwise nonnegative.The theorem concerns the continuous system (19a)–(19d) with γ = 0.Source: §G.3 · p. 16
  • LaSalle’s principle identifies the largest invariant set as X = X*, P = 0, R,C ≥ 0; the factor formulas then give R(t), C(t) → 0.Source: §G.3 · p. 17

G.3.1 Algebraic rates

For γ = 0, the energy has matching algebraic upper and lower bounds, with the rate depending on whether εstab is zero.

G.3.1 Algebraic rates

  • E(t) is bounded above and below by positive constants times (1+t)^-1/2 when εstab > 0, and by constants times (1+t)^-1 when εstab = 0.The zero-stabilizer result additionally assumes R(0) = C(0) = 0 and E(0) > 0.Source: §G.3.1 · p. 17
  • Fixed-window energy-loss estimates provide upper bounds, while bounds on instantaneous loss provide lower bounds.Source: §G.3.1 · pp. 17, 20
  • For εstab = 0, setting the friction to zero when the factor total s = 0 yields a locally Lipschitz friction tensor on the specified phase space.The zero-stabilizer phase-space argument establishes the needed Lipschitz property.Source: §G.3.1 · pp. 17–18
  • A fixed-window kinetic-energy lower bound applies once X → X*, P → 0, and the friction norm tends to zero.This estimate supplies a quantitative ingredient in the energy-rate proof.Source: §G.3.1 · pp. 17–18
  • The γ = 0 proof uses nonincreasing energy and LaSalle’s principle to obtain X → X*, P → 0, while the friction norm also tends to zero.Source: §G.3.1 · p. 19

Related papers

Loading 2609.30342v2…