# Adapprox: Adaptive Approximation in Adam Optimization via Randomized Low-
Rank Matrices

Pengxiang Zhao  Ping Li  Yingjie Gu  Yi Zheng  Stephan Ludger Kölker  Zhefeng
Wang  Xiaoming Yuan

###### Abstract

As deep learning models exponentially increase in size, optimizers such as
Adam encounter significant memory consumption challenges due to the storage of
first and second moment data. Current memory-efficient methods like Adafactor
and CAME often compromise accuracy with their matrix factorization techniques.
Addressing this, we introduce Adapprox, a novel approach that employs
randomized low-rank matrix approximation for a more effective and accurate
approximation of Adam’s second moment. Adapprox features an adaptive rank
selection mechanism, finely balancing accuracy and memory efficiency, and
includes an optional cosine similarity guidance strategy to enhance stability
and expedite convergence. In GPT-2 training and downstream tasks, Adapprox
surpasses AdamW by achieving 34.5% to 49.9% and 33.8% to 49.9% memory savings
for the 117M and 345M models, respectively, with the first moment enabled, and
further increases these savings without the first moment. Besides, it enhances
convergence speed and improves downstream task performance relative to its
counterparts.

Machine Learning, Optimizers, Memory Consumption, Matrix Factorization
Techniques

##  1 Introduction

In the field of deep learning, optimization algorithms play a pivotal role in
training models both efficiently and effectively. Among the most popular
optimization algorithms is the Adam (Kingma & Ba, 2014), and its variant,
AdamW (Loshchilov & Hutter, 2018), known for their robust performance across
diverse applications. However, the shift from smaller architectures like
AlexNet (Krizhevsky et al., 2017) with fewer than 100 million parameters, to
colossal models such as GPT-3 (Brown et al., 2020), encompassing over 100
billion parameters, poses substantial memory consumption challenges. This
issue becomes particularly acute in resource-constrained environments (Steiner
et al., 2023). Despite its effectiveness, the Adam optimizer exacerbates this
challenge by necessitating significant memory to store both first and second
moments for each parameter to maintain adaptive learning rates.

Figure 1: Singular value distributions. This figure shows the top 60 singular
values from six second moment matrices, out of a full rank of 1,024, obtained
from AdamW training a GPT-2 345M model at the 45,000th iteration.

There have been notable efforts in the development of memory-efficient
optimizers (Shazeer & Stern, 2018; Anil et al., 2019; Li et al., 2023; Luo et
al., 2023). Adafactor (Shazeer & Stern, 2018) offers the option to omit Adam’s
first moment, while utilizing a factored representation for the second moment.
However, Adafactor tends to show a decrease in training effectiveness, which
is primarily linked to the unavoidable approximation errors that occur in its
matrix factorization approach (Anil et al., 2019; Luo et al., 2023). To
mitigate this performance degradation and improve training stability, the CAME
optimizer (Luo et al., 2023) extends Adafactor with a confidence-based scaling
factor for the parameter updates. Nevertheless, CAME still relies on the same
fundamental factorization technique to store the second moment and confidence
statistics, thus inheriting the core challenges encountered by Adafactor.

In this study, we present Adapprox, a novel approach designed to overcome the
limitations associated with memory-efficient optimizers that rely on matrix
factorization. Our work is motivated by a crucial observation: in many large-
scale model training scenarios, the second moment matrices often have a
limited number of dominant singular values, with the rest exhibiting
substantially lower magnitudes. This insight is corroborated by empirical
evidence, illustrated in Figure 1. Consequently, Adapprox employs randomized
low-rank matrix approximation (Liberty et al., 2007; Li et al., 2014;
Batselier et al., 2018) to effectively approximate the second moment in Adam.
Our method primarily reduces memory usage by distilling key features from
large matrices, while also ensuring a more precise representation.

Figure 1 also elucidates the limitations of fixed 1-rank approximations used
in Adafactor and CAME. This approach’s tendency to compromise accuracy due to
multiple dominant singular values is evident in the top row of plots in Figure
1. However, in these instances, a modest increase in the target rank for
approximation can substantially improve accuracy. Considering the singular
value distribution patterns of the second moment, as illustrated in Figure 1,
overestimating the target rank in low-rank approximations can lead to
computationally intensive processes with minimal precision improvements.
Conversely, underestimating the rank significantly compromises accuracy. In
response to this, Adapprox introduces a dynamic rank-selection mechanism,
designed to adaptively select an optimal rank for each target matrix within
the low-rank approximation process.

Furthermore, to mitigate systematic errors inherent in low-rank
approximations, our approach optionally integrates a cosine similarity
guidance mechanism. This method computes the cosine similarity between each
update and the first moment, subsequently adjusting the learning rate for that
step proportionally. A low cosine similarity results in a reduced update,
while a high similarity prompts an increase in the update magnitude.

We evaluated our method through the pretraining of GPT-2 across various
configurations and associated downstream tasks. The findings demonstrate that
our approach achieves reduced memory usage compared to Adam, while maintaining
only a marginally higher memory footprint than Adafactor and CAME. Crucially,
our method not only maintains performance on par with Adam but also enables
accelerated convergence and improved generalization performance.

Our study’s key contributions are as follows:

  * •

In light of the singular value distribution of the second moment in Adam, we
present Adapprox, which utilizes randomized low-rank matrix approximation to
effectively approximate the second moment in Adam.

  * •

We develop an adaptive rank selection mechanism that balances precision and
memory savings by choosing an appropriate rank for approximating the target
matrix.

  * •

We integrate an optional cosine similarity guidance strategy into Adapprox,
which aims to expedite the convergence process and enhance the stability.

  * •

We showcase Adapprox’s efficacy on the GPT-2 pretraining and several
downstream tasks. Enabling the first moment, Adapprox achieves memory savings
ranging from 34.5% to 49.9% for the GPT-2 117M model and 33.8% to 49.9% for
the 345M model, compared to AdamW. Disabling the first moment elevates these
savings to 84.5% to 99.9% for the 117M model, and to 83.8% to 99.9% for the
345M model. Furthermore, our experimental results suggest that Adapprox may
offer faster convergence and potentially improved performance outcomes.

##  2 Related Work

Memory Efficient Optimizers. Memory efficient optimizers aim to reduce memory
usage by compressing optimizer states during training, ideally without
affecting the performance efficacy of their standard counterparts. Adafactor
(Shazeer & Stern, 2018) reduces memory usage by employing two main strategies:
the first is the optional omission of Adam’s first moment; and the second is
the compression of the second moment using a novel 1-rank matrix factorization
approach based on minimizing I-divergence (Lee & Seung, 1999). While Adafactor
achieves notable memory savings and computational efficiency, it is crucial to
recognize the potential compromise in accuracy that may arise from its
factorization approach (Anil et al., 2019; Luo et al., 2023). SM3 (Anil et
al., 2019) represents a memory efficient variation of Adagrad (Duchi et al.,
2011). Experimentally, it proves especially effective in situations where
gradients show natural activation patterns. CAME (Luo et al., 2023) builds
upon Adafactor’s framework by incorporating a confidence-guided strategy to
elevate approximation accuracy. However, this method also involves compressing
and storing confidence data using an identical matrix factorization technique,
resulting in similar accuracy-related challenges. Moreover, CAME’s dependency
on the first moment for its confidence strategy limits the potential for
omitting this data, thereby constraining further memory savings. 4-bit Adam
(Li et al., 2023) employs quantization techniques to compress Adam’s first and
second moments. Notably, this quantization is compatible with matrix
factorization methods.

Low-Rank Matrix Approximation. Low-rank matrix approximation seeks to
represent a matrix using lower-rank matrices, aiming for a more efficient data
representation while endeavoring to retain as much information as possible.
Low-rank matrix approximation is utilized across a wide spectrum of
applications, including principal component analysis (PCA) (Shen & Huang,
2008; Papailiopoulos et al., 2013), image processing (Haeffele et al., 2014;
Guo et al., 2017; Chen et al., 2017), and a variety of machine learning
scenarios (Paterek, 2007; Li et al., 2016).

##  3 Methodology

This section begins with a brief overview of Adam, followed by an examination
of the randomized low-rank approximation method. Subsequently, we introduce
the adaptive rank selection mechanism. We then provide a comprehensive
description of the Adapprox optimizer. Lastly, we delve into the proposed
cosine similarity guidance strategy.

###  3.1 Overview of the Adam Optimizer

Consider a function $f(W)$, where $W\in\mathbb{R}^{m\times n}$ denotes the
parameters of the neural network. The update rule for Adam (Kingma & Ba, 2014)
at the $t$-th iteration is defined as follows:

$\displaystyle\text{(Adam)}\left\\{\begin{array}[]{ll}G_{t}&=\nabla
f(W_{t-1}),\\\ M_{t}&=\beta_{1}M_{t-1}+(1-\beta_{1})G_{t},\\\
V_{t}&=\beta_{2}V_{t-1}+(1-\beta_{2})G_{t}^{2},\\\
\hat{M}_{t}&=M_{t}/(1-\beta_{1}^{t}),\\\
\hat{V}_{t}&=V_{t}/(1-\beta_{2}^{t}),\\\
W_{t}&=W_{t-1}-\alpha\hat{M}_{t}/(\sqrt{\hat{V}_{t}}+\epsilon).\end{array}\right.$
(7)

Here, all computations are element-wise. $G_{t}$ represents the gradient
arranged in matrix form. $M_{t}$ and $V_{t}$ are the exponential running
averages of the first and second moments, respectively. $\hat{M}_{t}$ and
$\hat{V}_{t}$ are the bias-corrected versions of $M_{t}$ and $V_{t}$.
$\beta_{1}$ and $\beta_{2}$ control these moment estimates. Additionally,
$\alpha$ denotes the learning rate, and $\epsilon$ is a small positive
constant introduced to prevent division by zero.

Building upon Adam, AdamW (Loshchilov & Hutter, 2018) decouples weight decay
from the gradient updates. With this change, the parameter update step in
AdamW is

$W_{t}=W_{t-1}-\alpha(\hat{M}_{t}/(\sqrt{\hat{V}_{t}}+\epsilon)+\lambda
W_{t-1}),$ (8)

where $\lambda$ is the rate of the weight decay.

As mentioned, Adam requires the storage of both $M_{t}$ and $V_{t}$ at each
step, which requires $O(mn)$ extra memory.

###  3.2 Low-Rank Approximation of the Second Moment

For a matrix $A\in\mathbb{R}^{m\times n}$, deriving its low-rank approximation
can be formulated as an optimization problem:

$\min_{Q,U}\|A-QU^{\top}\|_{F}^{2},$ (9)

where $\|\cdot\|_{F}$ is the Frobenius norm and $Q\in\mathbb{R}^{m\times k}$
and $U\in\mathbb{R}^{n\times k}$ are two feature matrices ($1\leq
k\ll\min\\{m,n\\}$). Then, $A_{k}=QU^{\top}$ is the $k$-rank approximation of
$A$. The optimal $A_{k}$ can be determined by performing a full Singular Value
Decomposition (SVD) and then truncating it to retain only the top $k$ singular
values and their corresponding singular vectors, resulting in the following
representation (see Theorem 2.4.8 in (Golub & Van Loan, 2013)):

$A_{k}=\sum_{i=1}^{k}\sigma_{i}u_{i}v_{i}^{T},$ (10)

where $\sigma_{1}\geq\sigma_{2}\geq\dots\geq\sigma_{k}\geq 0$ are the top $k$
singular values of $A$, and $u_{i}$ and $v_{i}$ $(1\leq i\leq k)$ are
corresponding left and right singular vectors. The approximation error can be
explicitly expressed as follows:

$\|A-A_{k}\|_{F}^{2}=\sum_{i=k+1}^{\min\\{m,n\\}}\sigma_{i}^{2}.$ (11)

Given Equation (11) and the singular value distribution of the second moment
matrix shown in Figure 1, utilizing a $k$-rank approximation to compress the
second moment matrix emerges as a rational choice. Nevertheless, the
computation of the full SVD for large matrices presents considerable
computational and memory challenges. We mitigate these issues by leveraging
randomized low-rank matrix approximation algorithms (Liberty et al., 2007;
Halko et al., 2011; Nakatsukasa, 2020), which provide a balance of
computational efficiency and memory economy, while still delivering high-
quality low-rank approximations.

Our implementation utilizes the Gaussian sampling variant of the randomized
SVD algorithm (Halko et al., 2011). In our approach, we bypass the SVD
estimation to streamline the process, concentrating solely on the extraction
of feature matrices without the need for singular values. Additionally, we
incorporate an oversampling mechanism, which further refines the algorithm’s
ability to capture more precise subspace representations. The comprehensive
procedure of this modified method is detailed in Algorithm 1, termed
Streamlined Randomized Subspace Iteration (S-RSI).

Algorithm 1 Streamlined Randomized Subspace Iteration

Inputs: Target matrix $A\in\mathbb{R}^{m\times n}$, target rank $k$, integer
$l$, and integer $p$ with $(k+p)\leq\min\\{m,n\\}$

$U\sim\mathcal{N}(0,1)$

$Q,R\leftarrow\mathbf{0}^{m\times(k+p)},\mathbf{0}^{(k+p)\times(k+p)}$

for $i\leftarrow 1,2,\dots,l$ do

$Q\leftarrow AU$

$Q,R\leftarrow\text{QR decomposition}(Q)$

$U\leftarrow A^{\top}Q$

end for

return $Q[:,:k],U^{\top}[:,:k]$

The S-RSI aims to compute an approximate basis $Q\in\mathbb{R}^{m\times k}$
which has orthonormal columns for the column space of the target matrix
$A\in\mathbb{R}^{m\times n}$ such that

$A_{k}=QQ^{\top}A.$ (12)

We form $U=Q^{\top}A$, and thus we obtain two feature matrices
$Q\in\mathbb{R}^{m\times k}$ and $U^{\top}\in\mathbb{R}^{n\times k}$.
Computing $Q$ is highly efficient using random sampling methods. Consider
drawing a random vector $u$, with each element independently and identically
distributed according to a standard Gaussian distribution. When we compute
$q=Au$, it effectively serves as a stochastic representation of the column
space of $A$, because $q$ represents a random linear combination of the
columns of $A$. By repeating this sampling process $k$ times, we get a set of
random vectors:

$\\{q_{i}\ |\ q_{i}=Au_{i},\ i=1,2,\dots,k\\}.$ (13)

Due to the inherent randomness in $u_{i}$’s generation, the set of vectors
$\\{u_{i}\\}_{i=1}^{k}$ are expected to occupy a general linear position,
which implies a high likelihood that any subset of these vectors is linearly
independent. Consequently, this observation leads us to propose the following:

######  Proposition 3.1.

Given a set of randomly generated vectors $\\{u_{i}\\}_{i=1}^{k}$ that are in
a general linear position, and a full rank matrix $A\in\mathbb{R}^{m\times
n}$, the set of vectors $\\{q_{i}\ |\ q_{i}=Au_{i}\\}_{i=1}^{k}$ are also
linearly independent.

###### Proof.

Linear independence of $\\{u_{i}\\}_{i=1}^{k}$ implies that
$\sum_{i=1}^{k}a_{i}u_{i}=0$ holds when all scalars $\\{a_{i}\\}_{i=1}^{k}$
are zero. We examine a linear combination of the vectors
$\\{q_{i}\\}_{i=1}^{k}$:

$\sum_{i=1}^{k}a_{i}q_{i}=\sum_{i=1}^{k}a_{i}Au_{i}=A\sum_{i=1}^{k}a_{i}u_{i}.$
(14)

As $A$ is full rank, $\sum_{i=1}^{k}a_{i}q_{i}\neq 0$ unless all $a_{i}$ are
zero, indicating that vectors $\\{q_{i}\\}_{i=1}^{k}$ are linearly
independent. ∎

Consequently, to derive an orthonormal basis for the column space of matrix
$A$, we first arrange the set $\\{q_{i}\\}_{i=1}^{k}$ as the columns of a
matrix $Q$, and then we apply an orthonormalization procedure, like QR
decomposition (Golub & Van Loan, 2013).

To enhance the precision of random sampling for basis determination in
scenarios where the input matrix exhibits a flat singular spectrum and is of
considerable size, we incorporate a power iteration technique (Rokhlin et al.,
2010; Halko et al., 2011; Golub & Van Loan, 2013). Specifically, this approach
applies the randomized sampling methodology to the modified matrix

$A^{\prime}=(AA^{\top})^{l}A,$ (15)

where $l$ is a small integer (for example, $l=5$). According to SVD, matrix
$A$ can be expressed as $A=Q\Sigma U^{T}$, where $Q$ and $U$ are orthogonal
matrices encapsulating the left and right singular vectors of $A$,
respectively, and $\Sigma$ is a diagonal matrix composed of the singular
values of $A$. Then, $A^{\prime}$ is derived as follows:

$A^{\prime}=(Q\Sigma U^{T}U\Sigma Q^{\top})^{l}Q\Sigma
U^{T}=Q(\Sigma)^{2l+1}U^{T}.$ (16)

Hence, while $A^{\prime}$ retains the same singular vectors as $A$, its
singular values exhibit an accelerated rate of decay:

$\sigma_{i}(A^{\prime})=\sigma_{i}(A)^{2l+1},$ (17)

This enhances the efficacy of distinguishing between more significant and less
significant singular vectors, thereby improving the overall accuracy of the
approximation. Furthermore, to augment the precision, we utilize a small
oversampling parameter $p$, such as $p=5$, to increase the sample size beyond
the target rank, thereby providing a buffer that mitigates the risk of
omitting significant components.

Through the S-RSI method, we can efficiently compress the storage of matrix
$A$ from $O(mn)$ to $O(k(m+n))$ by retaining only matrices $Q$ and $U^{\top}$
with time complexity of $O(lmn(k+p))$. Following the randomized SVD algorithm
(Halko et al., 2011), the approximation error bound is:

$\displaystyle\mathbb{E}\|A-QU^{\top}\|$
$\displaystyle\leq\left[\left(1+\sqrt{\frac{k}{p-1}}\right)^{2l+1}\sigma_{k+1}^{2l+1}\right.$
(18) $\displaystyle+\frac{e\sqrt{k+p}}{p}$
$\displaystyle\left.\sqrt{\sum_{j>k}\sigma_{j}^{2(2l+1)}}\right]^{1/(2l+1)}.$

According to Equation (18), the approximation error can be reduced by
increasing not just $k$, but also $p$ and $l$. We further demonstrate the
efficacy of the S-RSI through empirical comparisons. Specifically, we evaluate
its performance in comparison to Adafactor’s matrix factorization approach and
the SVD, focusing on all second-moment matrices obtained during the training
of a GPT-2 345M model using AdamW. Results are shown in Figure 2.

(a) Mean approximation error vs. rank.

(b) Mean computation time vs. rank.

Figure 2: Comparative analysis of the S-RSI ($l=5$ and $p=5$) against
Adafactor and SVD. All methods are applied to the second-moment matrices
derived from training a GPT-2 345M model using the AdamW, with results
captured at various stages of the training process.

Figure 2 compares the mean approximation error and computation time of varying
ranks, with the S-RSI employing hyperparameters $l=5$ and $p=5$. Adafactor’s
factorization method employs a fixed 1-rank approximation, leading to
consistent results regardless of rank changes. The SVD serves as a benchmark
for optimal low-rank approximation. The SVD and S-RSI show a substantial
decrease in approximation error with a minimal increase in rank. Besides, as
the rank increases, the S-RSI approaches the SVD benchmark bound. In terms of
computation time, Adafactor is the most efficient with a consistent duration.
The S-RSI not only significantly reduces computation time compared to SVD but
also approaches the efficiency of Adafactor. While the S-RSI shows a linear
relationship with rank $k$ theoretically, the actual rate of increase is
considerably less than $k$. This can be largely attributed to the multi-core
architecture of GPUs and parallel computing. Consequently, the S-RSI provides
an advantageous balance of approximation accuracy and computational
efficiency.

###  3.3 Adaptive Rank Selection

The results in Figure 2 underscore the importance of selecting the proper rank
$k$ for the S-RSI. Choosing a larger $k$ can result in increased computational
demands for minimal precision improvements, while a smaller $k$ risks
significantly compromising accuracy. To address this, we develop an adaptive
rank selection mechanism, which dynamically adjusts $k$ for each target
matrix.

The algorithm updates the rank $k_{t}$ at regular intervals of $\Delta s$
steps. Initially set as $k_{\textit{init}}$, $k_{t}$ is used to perform the
S-RSI. Then, we evaluate the approximation error rate:

$\xi=\frac{\|A-QU^{\top}\|_{F}}{\|A\|_{F}}.$ (19)

Should $\xi$ rise above a threshold $\xi_{\textit{thresh}}$, $k_{t}$ is
adjusted to $k_{t}+f(\xi)$, ensuring that it does not exceed the maximum rank
$k_{\textit{max}}$. We model $f$ as a variant of the sigmoid function
including four hyperparameters $\eta,\omega,\phi$ and $\tau$:

$f(\xi)=\left\lfloor\frac{\eta}{\exp(\omega\xi+\phi)+\tau}\right\rfloor,\
\xi>0,$ (20)

where $\eta>0$ sets the upper bound of $f$’s range, $\omega<0$ dictates the
rate of change, and $\phi$ and $\tau$ adjust the function’s offset. This
adheres to the heuristic principle that as $\xi$ increases, the expansion rate
of $k_{t}$ should initially be gradual, then accelerate, and ultimately
stabilize near an upper limit. By Equation (13), this extension requires
sampling $f(\xi)$ additional vectors from a standard Gaussian distribution and
then applying QR decomposition again. The revised $k_{t}$ is then maintained
for the next $\Delta s$ iterations. The pseudocode for this method is
summarized in Algorithm 2, which we designate as Adaptive S-RSI (AS-RSI).

Algorithm 2 Adaptive S-RSI

Inputs: Target matrix $A\in\mathbb{R}^{m\times n}$, rank $k_{t-1}$, initial
rank $k_{\textit{init}}$, the maximum rank $k_{\textit{max}}$, integer $l$,
integer $p$, threshold $\xi_{\textit{thresh}}$, step $t$, and adaptive
interval $\Delta s$

if $(t\mod\Delta s)=1$ then

$k_{t}\leftarrow k_{\textit{init}}$

repeat

$Q,U^{\top}\leftarrow$ S-RSI($A,k_{t},l,p$)

$\xi\leftarrow\|A-QU^{\top}\|_{F}/\|A\|_{F}$

$k_{t}\leftarrow\min\\{k_{t}+f(\xi),k_{\textit{max}}\\}${Equation (20)}

$p\leftarrow\min\\{p,k_{\textit{max}}-k_{t}\\}$

until $\xi\leq\xi_{\textit{thresh}}$

else

$k_{t}\leftarrow k_{t-1}$

$Q,U^{\top}\leftarrow$ S-RSI($A,k_{t},l,p$)

end if

return $Q[:,:k_{t}],U^{\top}[:,:k_{t}]$, $k_{t}$

###  3.4 Adapprox Algorithm

Algorithm 3 Adapprox

Inputs: Initial point $W_{0}\in\mathbb{R}^{m\times n}$,
$M_{0}={\mathbf{0}}^{m\times n}$ and $V_{0}=\mathbf{0}^{m\times n}$, learning
rates $\\{\alpha_{t}\\}_{t=1}^{T}$, moment decay $\beta_{1}$ and $\beta_{2}$,
small constant $\epsilon$, clipping threshold $d$, initial rank
$k_{\textit{init}}$, the maximum rank $k_{\textit{max}}$, integer $l$, integer
$p$ with $(k_{\textit{init}}+p)\leq k_{\textit{max}}$, threshold
$\xi_{\textit{thresh}}$, adaptive interval $\Delta s$, and weight decay rate
$\lambda$

Initialize $Q_{0}=\mathbf{0}^{m\times
k_{\textit{init}}},U_{0}^{\top}=\mathbf{0}^{n\times k_{\textit{init}}}$,
$k_{0}=k_{\textit{init}}$

for $t\leftarrow 1,2,\dots,T$ do

$G_{t}\leftarrow\nabla f_{t}(W_{t-1})$

$V_{t}\leftarrow\beta_{2}Q_{t-1}U_{t-1}^{\top}+(1-\beta_{2})G_{t}^{2}$

$Q_{t},U_{t}^{\top},k_{t}\leftarrow\text{AS-
RSI}(V_{t},k_{t-1},k_{\textit{init}},k_{\textit{max}},l,p,\xi_{\textit{thresh}},$

$\quad\quad\quad\quad\quad\quad t,\Delta s)$

$M_{t}\leftarrow G_{t}/(\sqrt{V_{t}}+\epsilon)$

$M_{t}\leftarrow M_{t}/\max(1,\text{RMS}(M_{t})/d)$

if $\beta_{1}>0$ then

$M_{t}\leftarrow\beta_{1}M_{t-1}+(1-\beta_{1})M_{t}$

end if

$W_{t}\leftarrow W_{t-1}-\alpha_{t}(M_{t}+\lambda W_{t-1})$

end for

The integration of the proposed methodologies culminates in the Adapprox
algorithm defined in Algorithm 3. At each step, $f_{t}$ represents a
stochastic realization of the objective function $f$, exemplified by the loss
function computed using a randomly selected mini-batch of data. We then
compute the gradient $G_{t}$ relative to the previous parameters and the
exponential running averages of the second moment $V_{t}$. Note that $V_{t-1}$
is reconstructed from $Q_{t-1}U_{t-1}^{\top}$, while $V_{t}$ is factored using
the AS-RSI. Subsequently, we calculate the update
$M_{t}=G_{t}/(\sqrt{V_{t}}+\epsilon)$ and incorporate the update clipping
mechanism as proposed in (Shazeer & Stern, 2018) to mitigate excessively large
updates:

$M_{t}\leftarrow\frac{M_{t}}{\max(1,{\text{RMS}(M_{t})}/{d})},\
\text{RMS}(M_{t})=\frac{\|M_{t}\|_{F}}{\sqrt{mn}},$

where $d$ is the clipping threshold. We also offer the option to omit the
first moment, depending on whether $\beta_{1}$ is set to zero. Finally,
parameter updates are executed in a decoupled weight decay fashion, as
delineated in Equation (8).

The proposed Adapprox, while rooted in Adam, diverges in three key aspects:
(1) it omits the bias correction steps; (2) it integrates an update clipping
mechanism that enhances performance (as detailed in Appendix A); and (3) it
modifies the first moment accumulator, replacing the running average of the
gradient with the running average of the update.

###  3.5 Cosine-Similarity Guidance Strategy

Inspired by CAME, we incorporate an optional mechanism aimed at enhancing the
update process and expediting the training phase when $\beta_{1}>0$.
Specifically, this involves a heuristic method that utilizes the cosine
similarity between the current update and the running average of updates to
assess the confidence level of the approximation.

To elaborate, for the given current update

$\hat{M}_{t}\leftarrow\frac{G_{t}}{\sqrt{V_{t}}+\epsilon},\
\hat{M}_{t}\leftarrow\frac{\hat{M}_{t}}{\max(1,{\text{RMS}(\hat{M}_{t})}/{d})},$
(21)

and the running average of updates

$M_{t}\leftarrow\beta_{1}M_{t-1}+(1-\beta_{1})\hat{M}_{t},$ (22)

we calculate the cosine similarity measure

$\theta_{\textit{cos}}=\frac{\sum_{i,j}(\hat{M}_{t}\odot
M_{t})_{ij}}{\|\hat{M}_{t}\|_{F}\|M_{t}\|_{F}}\in[-1,1],$ (23)

where $\odot$ denotes element-wise multiplication.

When $\theta_{\textit{cos}}$ approaches 1, it indicates a better alignment
between the directions of $\hat{M}_{t}$ and $M_{t}$. Conversely, a
$\theta_{\textit{cos}}$ value close to $-1$ suggests poor alignment. In the
case where $\theta_{\textit{cos}}=0$, $\hat{M}_{t}$ and $M_{t}$ are orthogonal
to each other. To modulate the update based on $\theta_{\textit{cos}}$, we
introduce a factor that either penalizes or accelerates the update
accordingly:

$M_{t}\leftarrow\frac{M_{t}}{1-\theta_{\textit{cos}}+\epsilon},$ (24)

which adjusts $M_{t}$ inversely with respect to $\theta_{\textit{cos}}$,
enhancing the update for alignment and penalizing it for divergence.

It is noteworthy that our proposed cosine-similarity guidance strategy is
effective exclusively when $\beta_{1}>0$, which is consistent with CAME.
However, our strategy is designed to utilize only the information available at
the current step, thereby obviating the need for extra memory.

##  4 Experiments

In this section, we present a comparative analysis of Adapprox against
existing counterparts in pretraining GPT-2 (Radford et al., 2019) and in
subsequent downstream tasks.

###  4.1 Setup

We investigate two GPT-2 configurations: 117M and 345M, as outlined in Table
1. Our pretraining experiments utilize The Pile dataset (Gao et al., 2020) and
the SentencePiece tokenizer (Kudo & Richardson, 2018). We evaluate the
pretrained models on several downstream tasks, including SQuAD v1.1 (Rajpurkar
et al., 2016), CoLA (Warstadt et al., 2018), MRPC (Dolan & Brockett, 2005),
SST-2 (Socher et al., 2013), and MNLI-m (Williams et al., 2017). Our primary
baselines include AdamW (Loshchilov & Hutter, 2018), Adafactor (Shazeer &
Stern, 2018), and CAME (Luo et al., 2023). We have implemented our
optimization algorithm using the PyTorch framework (Paszke et al., 2019).
Additionally, the pretraining of GPT-2 is conducted utilizing the Megatron-LM
framework (Shoeybi et al., 2019) and eight NVIDIA Tesla V-100 GPUs.

For GPT-2 pretraining, we set $\beta_{1}$ and $\beta_{2}$ at 0.9 and 0.999,
respectively, and maintain a consistent weight decay rate of 0.1 for all
compared algorithms. Adapprox’s additional parameters are specified as
follows: $\epsilon=1\times 10^{-8}$, $d=1$, $k_{\textit{init}}=1$,
$k_{\textit{max}}=0.25\min\\{m,n\\}$, $l=5$, $p=5$,
$\xi_{\textit{thresh}}=0.01$, $\Delta s=10$, and $\eta=200$, $\omega=-10$,
$\phi=-2.5$, $\tau=-9$. For Adafactor and CAME, the other parameters are set
to their respective default values. We adopt a linear warmup strategy followed
by a cosine-style learning rate decay, both integrated within the Megatron-LM
framework. To guarantee fair comparisons, all evaluated optimizers use uniform
training parameters for each model, selected through empirical testing and
established best practices. Specifically, the batch size, number of training
iterations, number of warmup iterations, peak learning rate, and minimum
learning rate are set as follows: for GPT-2 117M, they are 128, 100K, 1K,
$3\times 10^{-4}$, and $5\times 10^{-5}$; and for GPT-2 345M, 128, 100K, 1K,
$3\times 10^{-4}$, and $3\times 10^{-5}$, respectively.

For downstream tasks, we fine-tune GPT-2 models pretrained with each evaluated
optimizer for three epochs, adjusting learning rates individually per task.
Besides, cosine-similarity guidance is not employed in the fine-tuning process
of Adapprox.

Table 1: Model configurations.
Size | Layers | Hidden | Heads | Sequence
---|---|---|---|---
Length
117M | 12 | 768 | 12 | 1024
345M | 24 | 1024 | 16 | 1024

###  4.2 Memory Usage Comparison

(a) 117M, validation loss vs. iteration.

(b) 345M, validation loss vs. iteration.

(c) 117M, perplexity vs. iteration.

(d) 345M, perplexity vs. iteration.

Figure 3: Comparative analysis of Adapprox against AdamW, Adafactor, and CAME
on pretraining GPT-2 models. Table 2: Quantitative memory usage (MB)
comparison.
$\beta_{1}$ | Method | GPT-2 117M | GPT-2 345M
---|---|---|---
0.9 | AdamW | 949.7 (100.0%) | 2707.5 (100.0%)
Adafactor | 476.1 (50.1%) | 1356.7 (50.1%)
CAME | 476.8 (50.2%) | 1358.4 (50.2%)
Adapprox ($k_{\textit{init}}$) | 476.1 (50.1%) | 1356.7 (50.1%)
Adapprox ($k_{\textit{max}}$) | 622.0 (65.5%) | 1791.1 (65.5%)
0.0 | AdamW | 949.7 (100.0%) | 2707.5 (100.0%)
Adafactor | 1.2 (0.1%) | 2.9 (0.1%)
CAME | $-$ | $-$
Adapprox ($k_{\textit{init}}$) | 1.2 (0.1%) | 2.9 (0.1%)
Adapprox ($k_{\textit{max}}$) | 147.2 (15.5%) | 437.4 (16.2%)
Table 3: Results of fine-tuning performance on downstream tasks.
Model |  Method |  SQuAD v1.1 |  CoLA |  MRPC |  SST-2 |  MNLI-m |  Average
---|---|---|---|---|---|---|---
|  |  (F1) |  (Acc) |  (Acc) |  (Acc) |  (Acc) |
GPT-2 117M |  AdamW |  70.35 |  74.69 |  82.35 |  89.33 |  79.57 |  79.26
|  Adafactor |  72.57 |  75.55 |  79.41 |  89.33 |  80.13 |  79.40
|  CAME |  70.60 |  69.31 |  64.95 |  82.22 |  76.37 |  72.69
|  Adapprox |  72.89 |  75.83 |  80.63 |  90.94 |  80.15 |  80.09
GPT-2 345M |  AdamW |  75.06 |  79.39 |  82.84 |  91.17 |  82.33 |  82.16
|  Adafactor |  76.23 |  78.90 |  83.33 |  91.06 |  82.62 |  82.43
|  CAME |  75.50 |  69.22 |  71.81 |  82.22 |  81.72 |  76.09
|  Adapprox |  76.68 |  79.00 |  83.33 |  91.17 |  82.66 |  82.57

We assess the memory footprint associated with the state of each optimizer.
Adapprox utilizes an adaptive low-rank approximation, and we report its memory
usage based on predetermined $k_{\textit{init}}$ and $k_{\textit{max}}$
values. The actual memory usage falls between these boundaries. The results
are listed in Table 2. While Adapprox does increase memory usage compared to
Adafactor and CAME, it allows for a flexible trade-off between memory
efficiency and accuracy through the adjustment of $k_{\textit{init}}$ and
$k_{\textit{max}}$. In comparison to AdamW and with $\beta_{1}=0.9$, Adapprox
achieves memory reductions of 34.5% to 49.9% for GPT-2 117M and 33.8% to 49.9%
for GPT-2 345M. With $\beta_{1}=0$, CAME becomes non-viable, while AdamW still
allocates memory for the first moment. In contrast, both Adafactor and
Adapprox completely forego first moment memory allocation. In such scenarios,
Adapprox yields significant memory savings: 84.5% to 99.9% for GPT-2 117M and
83.8% to 99.9% for GPT-2 345M.

###  4.3 GPT-2 Training

Figure 3 presents the empirical comparison of Adapprox with AdamW, Adafactor,
and CAME in the pretraining of GPT-2 models. The top row shows the validation
loss across the 117M and 345M GPT-2 models, respectively, from left to right.
The bottom row displays the validation perplexity for these models, arranged
in the same sequence. Relative to AdamW, Adapprox generally exhibits better
validation loss and perplexity on the GPT-2 117M model, while delivering
comparable results on the GPT-2 345M. Adafactor consistently underperforms
compared to Adapprox, although the performance difference on the GPT-2 345M is
relatively modest. While CAME initially shows lower validation loss and
perplexity, it tends to converge to suboptimal outcomes over time. These
results indicate that Adapprox effectively balances accuracy with memory usage
and potentially offers faster convergence and superior performance compared to
Adafactor and CAME.

###  4.4 Downstream Tasks

We evaluate the downstream task performance of GPT-2 110M and 345M models,
each pretrained and fine-tuned with its corresponding optimizer. The empirical
results are presented in Table 3. Our experimental findings underscore the
superiority of Adapprox, as evidenced by better performance of GPT-2 models
trained and fine-tuned using it, compared to those using Adafactor and CAME.
Besides, Adapprox not only achieves performance comparable to AdamW but also
surpasses it in certain tasks. Additionally, we observe that Adapprox exhibits
stable performance across various learning rates, a conclusion drawn from
employing grid search to determine optimal learning rates for different
downstream tasks. In contrast, CAME shows sensitivity to learning rate
selection, frequently underperforming relative to other optimizers in these
sceniors. For detailed evidence, refer to Appendix B.

##  5 Discussion

While Adapprox and Adafactor achieve significant memory savings through the
low-rank approximation of the second moment and by omitting the first moment
(seen in Table 2), our experiments demonstrate that the absence of the first
moment has a notable impact on convergence speed and overall performance (see
Appendix C). Consequently, we recommend retaining the first moment unless
memory constraints are extremely prohibitive. Looking towards future research,
investigating techniques to compress the first moment could be a promising
direction. Furthermore, it is crucial to note that our approach is compatible
with other memory optimization techniques such as quantization and
recomputation. This compatibility opens the door for potential synergistic
enhancements in optimizer design.

##  6 Conclusion

In this paper, we introduce Adapprox, an innovative optimizer designed to
address the memory consumption challenges in training large-scale models. By
applying randomized low-rank matrix approximation to the second moment of
Adam, Adapprox achieves significant memory reduction with minimal impact on
model accuracy. This solution is further optimized with an adaptive rank
selection mechanism and an optional cosine similarity guidance strategy,
enhancing both stability and convergence speed. Our empirical evaluations,
encompassing GPT-2 117M and 345M models and their downstream tasks,
demonstrate Adapprox’s efficacy. It delivers considerable memory savings,
ranging from 33.8% to 49.9% over AdamW with the first moment intact, and
between 83.8% to 99.9% when the first moment is excluded. Despite a slight
trade-off in memory efficiency, Adapprox outperforms competitors like
Adafactor and CAME across several key metrics, including validation loss and
perplexity in pretraining, and F1 score and accuracy in downstream tasks.
These results establish Adapprox as a balanced approach for memory-efficient
training, effectively harmonizing efficiency with minimal accuracy trade-offs.

## Impact Statements

This paper presents work whose goal is to advance the field of Machine
Learning. There are many potential societal consequences of our work, none
which we feel must be specifically highlighted here.

## References

  * Anil et al. (2019) Anil, R., Gupta, V., Koren, T., and Singer, Y.  Memory efficient adaptive optimization.  _Advances in Neural Information Processing Systems_ , 32, 2019.
  * Batselier et al. (2018) Batselier, K., Yu, W., Daniel, L., and Wong, N.  Computing low-rank approximations of large-scale matrices with the tensor network randomized svd.  _SIAM Journal on Matrix Analysis and Applications_ , 39(3):1221–1244, 2018.
  * Brown et al. (2020) Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J. D., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., et al.  Language models are few-shot learners.  _Advances in neural information processing systems_ , 33:1877–1901, 2020.
  * Chen et al. (2017) Chen, Y., Guo, Y., Wang, Y., Wang, D., Peng, C., and He, G.  Denoising of hyperspectral images using nonconvex low rank matrix approximation.  _IEEE Transactions on Geoscience and Remote Sensing_ , 55(9):5366–5380, 2017.
  * Dolan & Brockett (2005) Dolan, B. and Brockett, C.  Automatically constructing a corpus of sentential paraphrases.  In _Third International Workshop on Paraphrasing (IWP2005)_ , 2005\.
  * Duchi et al. (2011) Duchi, J., Hazan, E., and Singer, Y.  Adaptive subgradient methods for online learning and stochastic optimization.  _Journal of machine learning research_ , 12(7), 2011.
  * Gao et al. (2020) Gao, L., Biderman, S., Black, S., Golding, L., Hoppe, T., Foster, C., Phang, J., He, H., Thite, A., Nabeshima, N., et al.  The pile: An 800gb dataset of diverse text for language modeling.  _arXiv preprint arXiv:2101.00027_ , 2020.
  * Golub & Van Loan (2013) Golub, G. H. and Van Loan, C. F.  _Matrix computations_.  JHU press, 2013.
  * Guo et al. (2017) Guo, Q., Gao, S., Zhang, X., Yin, Y., and Zhang, C.  Patch-based image inpainting via two-stage low rank approximation.  _IEEE transactions on visualization and computer graphics_ , 24(6):2023–2036, 2017.
  * Haeffele et al. (2014) Haeffele, B., Young, E., and Vidal, R.  Structured low-rank matrix factorization: Optimality, algorithm, and applications to image processing.  In _International conference on machine learning_ , pp. 2007–2015. PMLR, 2014.
  * Halko et al. (2011) Halko, N., Martinsson, P.-G., and Tropp, J. A.  Finding structure with randomness: Probabilistic algorithms for constructing approximate matrix decompositions.  _SIAM review_ , 53(2):217–288, 2011.
  * Kingma & Ba (2014) Kingma, D. P. and Ba, J.  Adam: A method for stochastic optimization.  _arXiv preprint arXiv:1412.6980_ , 2014.
  * Krizhevsky et al. (2017) Krizhevsky, A., Sutskever, I., and Hinton, G. E.  Imagenet classification with deep convolutional neural networks.  _Communications of the ACM_ , 60(6):84–90, 2017\.
  * Kudo & Richardson (2018) Kudo, T. and Richardson, J.  Sentencepiece: A simple and language independent subword tokenizer and detokenizer for neural text processing.  _arXiv preprint arXiv:1808.06226_ , 2018.
  * Lee & Seung (1999) Lee, D. D. and Seung, H. S.  Learning the parts of objects by non-negative matrix factorization.  _Nature_ , 401(6755):788–791, 1999.
  * Li et al. (2023) Li, B., Chen, J., and Zhu, J.  Memory efficient optimizers with 4-bit states.  _arXiv preprint arXiv:2309.01507_ , 2023.
  * Li et al. (2016) Li, D., Chen, C., Lv, Q., Yan, J., Shang, L., and Chu, S.  Low-rank matrix approximation with stability.  In _International Conference on Machine Learning_ , pp. 295–303. PMLR, 2016.
  * Li et al. (2014) Li, M., Bi, W., Kwok, J. T., and Lu, B.-L.  Large-scale nyström kernel matrix approximation using randomized svd.  _IEEE transactions on neural networks and learning systems_ , 26(1):152–164, 2014.
  * Liberty et al. (2007) Liberty, E., Woolfe, F., Martinsson, P.-G., Rokhlin, V., and Tygert, M.  Randomized algorithms for the low-rank approximation of matrices.  _Proceedings of the National Academy of Sciences_ , 104(51):20167–20172, 2007.
  * Loshchilov & Hutter (2018) Loshchilov, I. and Hutter, F.  Decoupled weight decay regularization.  In _International Conference on Learning Representations_ , 2018.
  * Luo et al. (2023) Luo, Y., Ren, X., Zheng, Z., Jiang, Z., Jiang, X., and You, Y.  Came: Confidence-guided adaptive memory efficient optimization.  In _Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)_ , pp. 4442–4453, 2023\.
  * Nakatsukasa (2020) Nakatsukasa, Y.  Fast and stable randomized low-rank matrix approximation.  _arXiv preprint arXiv:2009.11392_ , 2020.
  * Papailiopoulos et al. (2013) Papailiopoulos, D., Dimakis, A., and Korokythakis, S.  Sparse pca through low-rank approximations.  In _International Conference on Machine Learning_ , pp. 747–755. PMLR, 2013.
  * Paszke et al. (2019) Paszke, A., Gross, S., Massa, F., Lerer, A., Bradbury, J., Chanan, G., Killeen, T., Lin, Z., Gimelshein, N., Antiga, L., et al.  Pytorch: An imperative style, high-performance deep learning library.  In _Advances in Neural Information Processing Systems_ , volume 32, pp. 8024–8035, 2019.
  * Paterek (2007) Paterek, A.  Improving regularized singular value decomposition for collaborative filtering.  In _Proceedings of KDD cup and workshop_ , volume 2007, pp. 5–8, 2007.
  * Radford et al. (2019) Radford, A., Wu, J., Child, R., Luan, D., Amodei, D., Sutskever, I., et al.  Language models are unsupervised multitask learners.  _OpenAI blog_ , 1(8):9, 2019.
  * Rajpurkar et al. (2016) Rajpurkar, P., Zhang, J., Lopyrev, K., and Liang, P.  Squad: 100,000+ questions for machine comprehension of text.  _arXiv preprint arXiv:1606.05250_ , 2016.
  * Rokhlin et al. (2010) Rokhlin, V., Szlam, A., and Tygert, M.  A randomized algorithm for principal component analysis.  _SIAM Journal on Matrix Analysis and Applications_ , 31(3):1100–1124, 2010.
  * Shazeer & Stern (2018) Shazeer, N. and Stern, M.  Adafactor: Adaptive learning rates with sublinear memory cost.  In _International Conference on Machine Learning_ , pp. 4596–4604. PMLR, 2018.
  * Shen & Huang (2008) Shen, H. and Huang, J. Z.  Sparse principal component analysis via regularized low rank matrix approximation.  _Journal of multivariate analysis_ , 99(6):1015–1034, 2008.
  * Shoeybi et al. (2019) Shoeybi, M., Patwary, M., Puri, R., LeGresley, P., Casper, J., and Catanzaro, B.  Megatron-lm: Training multi-billion parameter language models using model parallelism, 2019.
  * Socher et al. (2013) Socher, R., Perelygin, A., Wu, J., Chuang, J., Manning, C. D., Ng, A. Y., and Potts, C.  Recursive deep models for semantic compositionality over a sentiment treebank.  In _Proceedings of the 2013 conference on empirical methods in natural language processing_ , pp. 1631–1642, 2013.
  * Steiner et al. (2023) Steiner, B., Elhoushi, M., Kahn, J., and Hegarty, J.  Model: memory optimizations for deep learning.  In _International Conference on Machine Learning_ , pp. 32618–32632. PMLR, 2023.
  * Warstadt et al. (2018) Warstadt, A., Singh, A., and Bowman, S. R.  Neural network acceptability judgments.  _arXiv preprint arXiv:1805.12471_ , 2018.
  * Williams et al. (2017) Williams, A., Nangia, N., and Bowman, S. R.  A broad-coverage challenge corpus for sentence understanding through inference.  _arXiv preprint arXiv:1704.05426_ , 2017.

##  Appendix A Ablation Experiment on Clipping Mechanism

Figure 4: Comparative analysis of training loss for the GPT-2 345M model
utilizing Adapprox with and without the clipping mechanism.

Figure 4 provides a comparative analysis of training loss convergence for the
GPT-2 345M model utilizing Adapprox with and without the implementation of the
clipping mechanism. The results indicate that the inclusion of a clipping
mechanism enhances Adapprox’s performance, resulting in lower training losses
at equivalent iterations.

##  Appendix B Learning Rate Sensitivity Analysis

Figure 5: Accuracy of the AdamW-pretrained GPT-2 345M model fine-tuned with
compared optimizers on the CoLA task across different learning rates.

In our experiments, Adapprox’s performance demonstrates stability across a
range of learning rates. This is illustrated in Figure 5, which compares the
accuracy of various optimizers in fine-tuning the AdamW-pretrained GPT-2 345M
model on the CoLA task under different learning rate conditions. Notably,
Adapprox shows minimal fluctuations, underlining its enhanced stability and
lower sensitivity to changes in learning rates.

##  Appendix C First Moment Efficacy Analysis

Figure 6: Training loss vs. iteration for AdamW, Adafactor, and Adapprox
optimizers, comparing scenarios with and without the first moment.

Figure 6 demonstrates that incorporating the first moment significantly
accelerates the convergence process, evidenced by achieving lower training
losses at the same iteration, for each optimizer examined, including AdamW,
Adafactor, and Adapprox. CAME is omitted from this analysis due to its
incompatibility with $\beta_{1}=0$. Furthermore, while AdamW exhibits
instability without the first moment, Adafactor and Adapprox mitigate this
through the use of a clipping mechanism, effectively reducing large,
unexpected updates and enhancing stability.