# Implicit Bias of AdamW: $\ell_{\infty}$ Norm Constrained Optimization

Shuo Xie Zhiyuan Li
Toyota Technological Institute at Chicago
<EMAIL_ADDRESS>
###### Abstract

Adam with decoupled weight decay, also known as AdamW, is widely acclaimed for
its superior performance in language modeling tasks, surpassing Adam with
$\ell_{2}$ regularization in terms of generalization and optimization.
However, this advantage is not theoretically well-understood. One challenge
here is that though intuitively Adam with $\ell_{2}$ regularization optimizes
the $\ell_{2}$ regularized loss, it is not clear if AdamW optimizes a specific
objective. In this work, we make progress toward understanding the benefit of
AdamW by showing that it implicitly performs constrained optimization. More
concretely, we show in the full-batch setting, if AdamW converges with any
non-increasing learning rate schedule whose partial sum diverges, it must
converge to a KKT point of the original loss under the constraint that the
$\ell_{\infty}$ norm of the parameter is bounded by the inverse of the weight
decay factor. This result is built on the observation that Adam can be viewed
as a smoothed version of SignGD, which is the normalized steepest descent with
respect to $\ell_{\infty}$ norm, and a surprising connection between
normalized steepest descent with weight decay and Frank-Wolfe.

##  1 Introduction

$\mathtt{Adam}$ (Kingma & Ba, 2014) and its variant $\mathtt{AdamW}$
(Loshchilov & Hutter, 2018) have been the most successful and widely used
optimization algorithms in deep learning, especially for large language models
(LLMs), whose pre-training costs massively and cannot be done with SGD.
Despite its tremendous empirical success, we lack a good theoretical
understanding of $\mathtt{Adam}$’s underlying work and the roles of its
hyperparameters, in particular, _weight decay_. $\mathtt{AdamW}$ achieves
better optimization and generalization over $\mathtt{Adam}$ and decouples the
effect of the learning rate and the weight decay coefficient by using a
different implementation of weight decay (Loshchilov & Hutter, 2018). While
$\mathtt{Adam}$ implements weight decay as a $\ell_{2}$ regularizer of the
training objective, $\mathtt{AdamW}$ directly shrinks its weight per step,
known as the _decoupled weight decay_ (see Algorithm 1).

Algorithm 1 Adam with $\ell_{2}$ regularization and Adam with decoupled weight
decay ($\mathtt{AdamW}$)

$\beta_{1},\beta_{2}>0$, initialization ${\bm{x}}_{0}$, total steps $T$,
learning rate schedule $\\{\eta_{t}\\}_{t=1}^{T}$, weight decay coefficient
$\lambda$

${\bm{m}}_{0}\leftarrow{\bm{\mathbf{0}}},{\bm{v}}_{0}\leftarrow{\bm{\mathbf{0}}}$

for $t=1,2,\cdots,T$ :

${\bm{g}}_{t}\leftarrow\nabla L({\bm{x}}_{t-1})$$+\lambda{\bm{x}}_{t-1}$

${\bm{m}}_{t}\leftarrow\beta_{1}{\bm{m}}_{t-1}+(1-\beta_{1}){\bm{g}}_{t}$

${\bm{v}}_{t}\leftarrow\beta_{2}{\bm{v}}_{t-1}+(1-\beta_{2}){\bm{g}}_{t}^{2}$

${\bm{x}}_{t}\leftarrow{\bm{x}}_{t-1}-\eta_{t}\frac{{\bm{m}}_{t}}{\sqrt{{\bm{v}}_{t}}}$$-\lambda\eta_{t}{\bm{x}}_{t-1}$

return ${\bm{x}}_{T}$

However, the advantage of $\mathtt{AdamW}$ over $\mathtt{Adam}$ is mostly
empirical and our theoretical understanding is quite limited. Zhuang et al.
(2022) argues that one desirable property that $\mathtt{AdamW}$ has while
$\mathtt{Adam}$ does not is scale-freeness, meaning $\mathtt{AdamW}$ yields
the same optimization trajectory if loss is multiplied by any positive
constant. Yet this property does not give us enough information to understand
the difference regarding the optimization processes and the final learned
solutions between $\mathtt{AdamW}$ and $\mathtt{Adam}$ with $\ell_{2}$
regularization. Intuitively, if $\mathtt{Adam}$ with $\ell_{2}$ regularization
converges to some point, it converges to at least a stationary point of the
regularized loss function, if not a minimizer. But for $\mathtt{AdamW}$, it is
even not clear if it is optimizing any (regularized) loss function. Thus,
towards taking the first step of understanding the benefit of decoupled weight
decay in $\mathtt{AdamW}$, we ask the following question:

> Which solution does $\mathtt{AdamW}$ converge to, if it converges?

Our following main result Theorem 1.1 characterizes the implicit bias of
$\mathtt{AdamW}$ in the _deterministic_ case, where a full-batch loss is used:

######  Theorem 1.1.

For any continuously differentiable function $L:\mathbb{R}^{d}\to\mathbb{R}$,
$\beta_{1}\leq\beta_{2}<1$, initialization ${\bm{x}}_{0}$ and non-increasing
learning rate $\\{\eta_{t}\\}_{t=1}^{\infty}$ such that
$\sum_{t=1}^{\infty}\eta_{t}=\infty$, if the iterates of $\mathtt{AdamW}$
$\\{{\bm{x}}_{t}\\}_{t=0}^{\infty}$ on $L$ converges to some
${\bm{x}}_{\infty}$, then ${\bm{x}}_{\infty}$ is a KKT point (Definition 3.6)
of the constrained optimization problem
$\min_{\left\|{\bm{x}}\right\|_{\infty}\leq\frac{1}{\lambda}}L({\bm{x}})$.

If $L$ is additionally convex, then $\mathtt{AdamW}$ converges to the
constrained minimizer, _i.e._ ,
${\bm{x}}_{\infty}\in\operatorname*{arg\,min}_{\left\|{\bm{x}}\right\|_{\infty}\leq\frac{1}{\lambda}}L({\bm{x}})$.

Despite being simplistic, the full-batch setting is still a very interesting
and highly non-trivial regime, because the two main hypotheses of why
$\mathtt{Adam}$ outperforms $\mathtt{SGD}$ got challenged recently in the
deterministic regime (Kunstner et al., 2022). The first hypothesis is that
$\mathtt{Adam}$ outperforms $\mathtt{SGD}$ by better handling heavy-tailed
noise (Zhang et al., 2020). However, Kunstner et al. (2022) finds that
$\mathtt{Adam}$ still outperforms $\mathtt{GD}$ for optimizing language tasks
even in the full-batch setting. The second hypothesis is the smoothness of the
training loss landscape can linearly increase as the gradient norm increases
and thus clipping or normalization is necessary for gradient descent.
Intriguingly, Kunstner et al. (2022) finds that normalizing each update of GD
cannot close the gap towards $\mathtt{Adam}$ in the full-batch setting, but
normalizing _each coordinate_ to its sign (_i.e._ , $\mathtt{SignGD}$) closes
the gap.

The theoretical results and analysis in this work support the empirical
observation made by Kunstner et al. (2022). The way we prove Theorem 1.1 is
first to prove that normalized steepest descent with weight decay
($\mathtt{NSD}$-$\mathtt{WD}$) for any norm $\left\|\cdot\right\|$ must
converge to KKT points of the constrained optimization problem
$\min_{\left\|{\bm{x}}\right\|\leq\frac{1}{\lambda}}L({\bm{x}})$ (Theorem
3.7). Then we show that $\mathtt{AdamW}$ asymptotically behaves just like
$\mathtt{SignGD}$ with weight decay, which is the normalized steepest descent
w.r.t. $\ell_{\infty}$ norm with weight decay and the same proof framework
generalizes to $\mathtt{AdamW}$ if $\beta_{1}\leq\beta_{2}<1$. The condition
$\beta_{1}\leq\beta_{2}$ is crucial, and we provide a counter-example where
$1>\beta_{1}>\beta_{2}$ and $\mathtt{AdamW}$ converges somewhere else, instead
of a constrained minimizer.

It remains interesting why $\mathtt{SignGD}$ outperforms
$\mathtt{Normalized}$-$\mathtt{GD}$ in the full batch setting. They are both
normalized steepest descent but with respect to different norms –
$\mathtt{Normalized}$-$\mathtt{GD}$ picks the steepest direction under the
geometry of $\ell_{2}$ norm, while $\mathtt{SignGD}$ picks the steepest
direction under the geometry of the $\ell_{\infty}$ norm. It is natural to
make the following conjecture: _$\mathtt{Adam}$ outperforms $\mathtt{GD}$ due
to its utilization of $\ell_{\infty}$ geometry, under which the loss function
could have better properties, e.g., smaller smoothness_. Our main result
Theorem 1.1 provides positive evidence for this conjecture. We also provide a
convergence analysis for normalized steepest descent with weight decay for
convex loss, where the suboptimality against the constrained minimizer in norm
ball of radius $\frac{1}{\lambda}$ vanishes is $O(\frac{H}{T\lambda^{2}})$,
where $T$ is the total number of steps, $\lambda$ is the weight decay factor,
and $H$ is the smoothness of loss w.r.t. the particular norm used for picking
the steepest descent direction. Based on the convergence bound, we construct a
concrete $d$-dimensional loss function in Section 3.1 whose minimizer
${\bm{x}}^{*}$ satisfies
$\left\|{\bm{x}}^{*}\right\|_{2}\approx\sqrt{d}\left\|{\bm{x}}^{*}\right\|_{\infty}$
and $\mathtt{SignGD}$ with weight decay converges much faster than
$\mathtt{Normalized}$-$\mathtt{GD}$ with weight decay because
$\mathtt{SignGD}$ with weight decay can use a $\sqrt{d}$ times larger weight
decay factor $\lambda$ than $\mathtt{Normalized}$-$\mathtt{GD}$.

#### Contributions.

Below we summarize our contributions:

  1. 1.

In Section 3.1, we prove normalized steepest descent with weight decay
optimizes convex functions under norm constraints (Theorem 3.5). In Section
3.2, we prove it must converge to KKT points of the norm-constrained
optimization problem for general loss functions if it converges with a
learning rate schedule whose partial sum diverges (Theorem 3.7).

  2. 2.

In Section 4, we prove $\mathtt{AdamW}$ must converge to KKT points of the
norm-constrained optimization problem for general loss functions if it
converges with a non-increasing learning rate schedule whose partial sum
diverges (Theorem 1.1).

  3. 3.

Towards generalizing the proof of Theorem 3.7 to Theorem 1.1, we prove a novel
and tight upper bound on average update size of $\mathtt{Adam}$ (Lemma 4.2),
which holds even for non-deterministic settings as well and might be of
independent interest to the community. We test various predictions made by our
bound in experiments.

##  2 Preliminaries and Notations

#### Notations:

We use $\left\|\cdot\right\|$ to denote a general norm and
$\left\|\cdot\right\|_{*}$ to denote its dual norm. We say a function
$L:\mathbb{R}^{d}\to\mathbb{R}$ has $H$-lipschitz gradient w.r.t. norm
$\left\|\cdot\right\|$ for some $H>0$ iff for all
${\bm{x}},{\bm{y}}\in\mathbb{R}^{d}$, $\left\|\nabla L({\bm{x}})-\nabla
L({\bm{y}})\right\|_{*}\leq H\left\|{\bm{x}}-{\bm{y}}\right\|$. We define the
_smoothness_ of loss $L$ as the smallest positive $H$ w.r.t.
$\left\|\cdot\right\|$ such that $L$ has $H$-lipschitz gradient. We say a
function $L:\mathbb{R}^{d}\to\mathbb{R}$ is convex iff for any
${\bm{x}},{\bm{y}}\in\mathbb{R}^{d}$ and $\theta\in[0,1]$, it holds that
$L(\theta{\bm{x}}+(1-\theta){\bm{y}})\leq\theta
L({\bm{x}})+(1-\theta)L({\bm{y}})$. We define the _subgradients_ of convex
function $L$ at point ${\bm{x}}$ as $\\{{\bm{g}}\in\mathbb{R}^{d}\mid
L({\bm{y}})\geq
L({\bm{x}})+\left\langle{\bm{y}}-{\bm{x}},{\bm{g}}\right\rangle,\forall{\bm{y}}\in\mathbb{R}^{d}\\}$,
which is denoted by $\partial L({\bm{x}})$. When $L$ is differentiable at
${\bm{x}}$, $\partial L({\bm{x}})$ contains only one element, which is the
gradient $\nabla L({\bm{x}})$. In particular, all norms are convex functions
and we have the following standard lemma for the subgradients of norms:

######  Lemma 2.1.

For any norm $\left\|\cdot\right\|$ and ${\bm{x}}\in\mathbb{R}^{d}$,
$\partial\left\|{\bm{x}}\right\|=\\{{\bm{\Delta}}\in\mathbb{R}^{d}\mid\left\|{\bm{\Delta}}\right\|_{*}=1,\left\langle{\bm{\Delta}},{\bm{x}}\right\rangle=\left\|{\bm{x}}\right\|\\}$.

#### Steepest Descent:

We say ${\bm{v}}$ is a _steepest descent direction_ for objective function $L$
at current iterate ${\bm{x}}$ w.r.t. norm $\left\|\cdot\right\|$ iff
$\left\|{\bm{v}}\right\|=1$ and $\left\langle{\bm{v}},\nabla
L({\bm{x}})\right\rangle=\min_{\left\|{\bm{v}}^{\prime}\right\|\leq
1}\left\langle{\bm{v}}^{\prime},\nabla L({\bm{x}})\right\rangle$. Thus for all
steepest descent direction ${\bm{v}}$, we have that
$\left\langle{\bm{v}},\nabla L({\bm{x}})\right\rangle=-\left\|\nabla
L({\bm{x}})\right\|_{*}$.

Given initialization ${\bm{x}}_{0}$, learning rate schedule
$\\{\eta_{t}\\}_{t=0}^{\infty}$ and weight decay factor $\lambda$, the $t$th
iterate of _normalized steepest descent_ w.r.t. $\left\|\cdot\right\|$ with
decoupled weight decay is defined as

${\bm{x}}_{t}=(1-\lambda\eta_{t}){\bm{x}}_{t-1}-\eta_{t}{\bm{\Delta}}_{t}\text{,
where
}{\bm{\Delta}}_{t}\in\operatorname*{arg\,max}_{\left\|{\bm{\Delta}}\right\|\leq
1}\nabla L({\bm{x}}_{t-1})^{\top}{\bm{\Delta}}.$ (1)

Because the dual norm of the dual norm is always equal to the original norm,
by Lemma 2.1, we can also characterize the steepest descent directions as the
subgradient of its dual norm.

######  Lemma 2.2.

$\operatorname*{arg\,max}\limits_{\left\|{\bm{\Delta}}\right\|\leq 1}\nabla
L({\bm{x}})^{\top}{\bm{\Delta}}=\left.\partial\left\|{\bm{y}}\right\|_{*}\right|_{{\bm{y}}=\nabla
L({\bm{x}})}$.

For completeness, we also define the _steepest descent_ w.r.t.
$\left\|\cdot\right\|$ with decoupled weight decay below, though we will not
use it in our analysis. If we pick $\ell_{2}$ norm, Equation 2 becomes
standard gradient descent.

$\tilde{\bm{x}}_{t}=(1-\lambda\eta_{t})\tilde{\bm{x}}_{t-1}-\eta_{t}\tilde{\bm{\Delta}}_{t}\text{,
where
}\tilde{\bm{\Delta}}_{t}\in\operatorname*{arg\,max}_{\tilde{\bm{\Delta}}\in\mathbb{R}^{d}}\left(\nabla
L({\bm{x}}_{t-1})^{\top}\tilde{\bm{\Delta}}-\frac{1}{2}\left\|\tilde{\bm{\Delta}}\right\|^{2}\right).$
(2)

It can be shown that for each steepest descent update $\tilde{\bm{\Delta}}$
for objective $L$ at ${\bm{x}}$, there exists some normalized steepest descent
update ${\bm{\Delta}}$ satisfying $\tilde{\bm{\Delta}}=\left\|\nabla
L({\bm{x}})\right\|_{*}{\bm{\Delta}}$.

##  3 Warm Up: Implicit Bias of Normalized Steepest Descent w. Weight Decay

In this section, we aim to present some high-level intuition about the
constrained-minimization implicit bias of $\mathtt{AdamW}$ (Theorem 1.1), by
showing the same implicit bias for $\mathtt{SignGD}$ with weight decay, or
equivalently, normalized steepest descent w.r.t. $\ell_{\infty}$ norm.
$\mathtt{AdamW}$ is arguably a smoothed version of $\mathtt{SignGD}$, which
reduces the correlation between its numerator and denominator by using past
moving average and thus reduces the biasedness of the update direction in the
presence of noise. But intuitively, their behaviors are similar when there is
no noise and the learning rate is small.

Our analysis in this section holds for all norms, including the non-
differentiable ones, like $\left\|\cdot\right\|_{\infty}$.

###  3.1 Convex setting: constrained optimization

In this subsection, we give a simple non-asymptotic convergence analysis for
normalized Steepest descent w. weight decay ($\mathtt{NSD}$-$\mathtt{WD}$)
w.r.t. to general norms over smooth convex loss functions. If the norm of
initialization is no larger than $\frac{1}{\lambda}$ where $\lambda$ is the
weight decay factor then surprisingly $\mathtt{NSD}$-$\mathtt{WD}$ is exactly
equivalent to a well-known optimization algorithm in literature,
$\mathtt{Frank}$-$\mathtt{Wolfe}$ (Frank et al., 1956), where the constraint
set here is the norm ball with radius $\frac{1}{\lambda}$. If the norm of
initialization is larger than $\frac{1}{\lambda}$, then the analysis contains
an additional phase where the norm of iterates linearly converges to
$\frac{1}{\lambda}$. In this case, the iterate of $\mathtt{NSD}$-$\mathtt{WD}$
may always be outside the $\frac{1}{\lambda}$ norm ball, but still, the
convergence analysis of $\mathtt{Frank}$-$\mathtt{Wolfe}$ can be adopted
(_e.g._ , Jaggi (2013)). First, we show that the norm of the iterates will
shrink to $\frac{1}{\lambda}$ as long as the norm of each update is bounded by
$1$, i.e., $\left\|{\bm{\Delta}}_{t}\right\|\leq 1$. Note this conclusion
doesn’t use the convexity of the function $L({\bm{x}})$ nor the update
${\bm{\Delta}}_{t}$ being the steepest descent direction. It can hold under
non-deterministic settings.

######  Lemma 3.1.

For any learning rate schedule $\\{\eta_{t}\\}_{t=1}^{\infty}$ and update
$\\{{\bm{\Delta}}_{t}\\}_{t=1}^{\infty}$ such that $\lambda\eta_{t}<1$ and
$\left\|{\bm{\Delta}}_{t}\right\|\leq 1$,
$\left\|{\bm{x}}_{t}\right\|-\frac{1}{\lambda}\leq\max\left(e^{-\lambda\sum_{i=1}^{t}\eta_{i}}\left(\left\|{\bm{x}}_{0}\right\|-\frac{1}{\lambda}\right),0\right)$.

The proof is deferred to Section A.1. Lemma 3.1 shows that ${\bm{x}}_{t}$ is
either always inside the norm ball with radius $\frac{1}{\lambda}$, or their
distance shrinks exponentially as the sum of learning rates increases.
Whenever ${\bm{x}}_{t}$ gets into the norm ball with radius
$\frac{1}{\lambda}$, ${\bm{x}}_{t}$ will not leave it and the remaining
trajectory of $\mathtt{NSD}$-$\mathtt{WD}$ is exactly the same as
$\mathtt{Frank}$-$\mathtt{Wolfe}$, as shown in the following theorem. We note
the relationship between $\mathtt{Frank}$-$\mathtt{Wolfe}$ and steepest
descent algorithms is also observed very recently in the continuous case (Chen
et al., 2023).

######  Theorem 3.2.

For any norm $\left\|\cdot\right\|$, weight decay $\lambda$, and
$\left\|{\bm{x}}_{t-1}\right\|\leq\frac{1}{\lambda}$,
$\mathtt{NSD}$-$\mathtt{WD}$ with learning rate $\eta_{t}<\frac{1}{\lambda}$
and $\mathtt{Frank}$-$\mathtt{Wolfe}$ (Algorithm 2) with step size
$\gamma_{t}=\eta_{t}\lambda$ and convex set
$\mathcal{X}\triangleq\\{{\bm{y}}\mid\left\|{\bm{y}}\right\|\leq\frac{1}{\lambda}\\}$
generate the same next iterate ${\bm{x}}_{t}$.

Algorithm 2 $\mathtt{Frank}$-$\mathtt{Wolfe}$

convex set $\mathcal{X}$, ${\bm{x}}_{0}\in\mathcal{X}$, total steps $T$, step
sizes $\\{\gamma_{t}\\}_{t=1}^{T}$

for $t=1,2,\cdots,T$ :

${\bm{y}}_{t}\leftarrow\operatorname*{arg\,min}_{{\bm{y}}\in\mathcal{X}}\nabla
L({\bm{x}}_{t-1})^{\top}{\bm{y}}$

${\bm{x}}_{t}\leftarrow(1-\gamma_{t}){\bm{x}}_{t-1}+\gamma_{t}{\bm{y}}_{t}$

return ${\bm{x}}_{T}$

Define
${\bm{x}}^{*}=\operatorname*{arg\,min}_{\left\|{\bm{x}}\right\|\leq\frac{1}{\lambda}}L({\bm{x}})$
to be the constrained minimizer of convex function $L({\bm{x}})$. We first
compute how much the gap between $L({\bm{x}}_{t})$ and $L({\bm{x}}^{*})$ can
decrease in one normalized steepest descent step when the iterate
${\bm{x}}_{t}$ is bounded.

######  Lemma 3.3 (Descent Lemma for Smooth Convex Loss).

Suppose loss function $L$ is convex and has $H$-lipschitz gradient w.r.t. norm
$\left\|\cdot\right\|$. For iterates $\\{{\bm{x}}_{t}\\}$ in
$\mathtt{NSD}$-$\mathtt{WD}$ (Equation 1), we have that

$L({\bm{x}}_{t})-L({\bm{x}}^{*})\leq(1-\lambda\eta_{t})(L({\bm{x}}_{t-1})-L({\bm{x}}^{*}))+\frac{H}{2}\eta_{t}^{2}\left(1+\lambda\left\|{\bm{x}}_{t-1}\right\|\right)^{2}.$

The proof of Lemma 3.3 is deferred to Section A.2. With Lemma 3.3, we can
prove the convergence of $L({\bm{x}}_{t})$ for learning rate schedules with
certain conditions. The proof is also deferred to Section A.2.

######  Theorem 3.4.

Assume that $\eta_{t}\geq 0$, $\lim_{t\rightarrow\infty}\eta_{t}=0$ and
$\sum_{t=1}^{\infty}\eta_{t}=\infty$. For any convex loss $L$ with
$H$-lipschitz gradient,
$\lim_{t\rightarrow\infty}L({\bm{x}}_{t})=L({\bm{x}}^{*})$.

We also provide a specific example of learning rates
$\\{\eta_{t}\\}_{t=1}^{\infty}$ that can achieve $O(\frac{1}{t})$ convergence
of $f({\bm{x}}_{t})$, which is the same as $\mathtt{Frank}$-$\mathtt{Wolfe}$
over convex objectives (Jaggi, 2013) and the proof is standard. For
completeness, we provide a proof of Theorem 3.5 in Section A.2.

######  Theorem 3.5.

Define $B=\max{\\{\left\|{\bm{x}}_{0}\right\|,\frac{1}{\lambda}\\}}$. For
$\mathtt{NSD}$-$\mathtt{WD}$ with learning rate schedule
$\eta_{t}=\frac{2}{\lambda(t+1)}$, we have
$L({\bm{x}}_{t})-L({\bm{x}}^{*})\leq\frac{2H(1+\lambda
B)^{2}}{(t+2)\lambda^{2}}$ for $t\geq 1$.

Note that the descent rate highly depends on the smoothness coefficient $H$,
which is determined by the selected norm. Therefore, we provide a synthetic
example that may demonstrate the advantage of $\ell_{\infty}$ norm as
mentioned in Section 1. For some constant ${\bm{x}}^{*}\in\mathbb{R}^{d}$, the
loss function $g:\mathbb{R}^{d}\rightarrow\mathbb{R}$ is defined as

$g({\bm{x}})=\sum_{i=1}^{d}\frac{({\bm{x}}_{i}-{\bm{x}}^{*}_{i})^{2}}{i^{2}}.$
(3)

The Hessian matrix $\nabla^{2}g$ is a diagonal matrix with diagonal entries
$\\{\frac{2}{i^{2}}\\}_{i=1}^{d}$. For $\ell_{2}$ norm, the smoothness
coefficient is the largest eigenvalue of Hessian matrix, which is $2$. For
$\ell_{\infty}$ norm, the smoothness coefficient is the sum of the diagonal
entries because of the diagonality, which is $\sum_{i=1}^{d}\frac{2}{i^{2}}$.
It is upper bounded by $\sum_{i=1}^{\infty}\frac{2}{i^{2}}=\frac{\pi^{2}}{3}$
for all dimension $d$. However, if ${\bm{x}}^{*}$ is set to be in the unit
$\ell_{\infty}$ norm ball, its $\ell_{2}$ norm can be as large as $\sqrt{d}$,
which makes the suboptimality bound in Theorem 3.5 for $\ell_{2}$ norm
normalized steepest descent with weight decay $d$ times larger than its
$\ell_{\infty}$ counterpart. We implement steepest descent with
$\ell_{\infty}$ norm and $\ell_{2}$ norm on a $100$-dimension example in
Section 5.2 and find that $\ell_{\infty}$ norm can indeed work better.

###  3.2 Non-convex setting: convergence to KKT points

In this subsection, we study the implicit bias of $\mathtt{SignGD}$ (or more
generally, $\mathtt{NSD}$-$\mathtt{WD}$) when the loss is non-convex. In such
case, last-iterate parameter convergence is, in general, difficult to
show111Indeed, even for convex case, $\mathtt{Frank}$-$\mathtt{Wolfe}$ may not
converge in parameter.(Bolte et al., 2023), and thus we turn to study _what
parameters $\mathtt{SignGD}$ and $\mathtt{NSD}$-$\mathtt{WD}$ can converge
to_. Our main results Theorem 3.7 show that such parameters must be the KKT
points (see Definition 3.6) of the constrained optimization problems. In
particular, if the objective is convex, since the norm ball constraint is
always convex for all norm, all KKT points are constrained minimizers.

######  Definition 3.6 (KKT points).

We say ${\bm{x}}^{*}\in\mathbb{R}^{d}$ is a KKT point of the constrained
optimization problem
$\min_{\left\|{\bm{x}}\right\|\leq\frac{1}{\lambda}}L({\bm{x}})$ iff there
exists $s^{*}\geq 0$ such that
$\left\|{\bm{x}}^{*}\right\|\leq\frac{1}{\lambda}$, $0\in\nabla
L({\bm{x}}^{*})+\partial(s^{*}\left\|{\bm{x}}\right\|)|_{{\bm{x}}={\bm{x}}^{*}}$,
$s^{*}(\left\|{\bm{x}}^{*}\right\|-\frac{1}{\lambda})=0$.

For convex $L$, all KKT points ${\bm{x}}^{*}$ are optimal and the dual
variable $s^{*}\geq 0$ is the certificate for the optimality. To see that, for
any other $\left\|{\bm{y}}\right\|\leq\frac{1}{\lambda}$, it holds that
$L({\bm{y}})\geq
L({\bm{y}})+s^{*}(\left\|{\bm{y}}\right\|-\frac{1}{\lambda})\geq
L({\bm{x}}^{*})+s^{*}(\left\|{\bm{x}}^{*}\right\|-\frac{1}{\lambda})$, where
the second inequality is because $L({\bm{x}})+s^{*}\left\|{\bm{x}}\right\|$ is
also convex and $0$ is its subgradient at ${\bm{x}}^{*}$. Thus we conclude
$L({\bm{y}})\geq
L({\bm{x}}^{*})+s^{*}(\left\|{\bm{x}}^{*}\right\|-\frac{1}{\lambda})=L({\bm{x}}^{*})$.

Now we state the main result for this subsection.

######  Theorem 3.7 (Non-convex, KKT).

For any continuously differentiable function $L:\mathbb{R}^{d}\to\mathbb{R}$,
initialization ${\bm{x}}_{0}$, and learning rate schedule
$\\{\eta_{t}\\}_{t=1}^{\infty}$ such that
$\sum_{t=1}^{\infty}\eta_{t}=\infty$, if the iterates of
$\mathtt{NSD}$-$\mathtt{WD}$ $\\{{\bm{x}}_{t}\\}_{t=0}^{\infty}$ on $L$
converges to some ${\bm{x}}_{\infty}$, then ${\bm{x}}_{\infty}$ is a KKT point
(Definition 3.6) of the constrained optimization problem
$\min_{\left\|{\bm{x}}\right\|\leq\frac{1}{\lambda}}L({\bm{x}})$.

To prove Theorem 3.7, we use the following alternative characterization for
KKT points of $\min_{\left\|{\bm{x}}\right\|\leq\frac{1}{\lambda}}L({\bm{x}})$
below based on Lemma 2.1.

######  Lemma 3.8.

${\bm{x}}$ is a KKT point of
$\min_{\left\|{\bm{x}}\right\|\leq\frac{1}{\lambda}}L({\bm{x}})$ iff
$\left\|{\bm{x}}\right\|\leq\frac{1}{\lambda}$ and
$\left\langle-\lambda{\bm{x}},\nabla L({\bm{x}})\right\rangle=\left\|\nabla
L({\bm{x}})\right\|_{*}$.

With Lemma 3.8, next we illustrate the intuition for Theorem 3.7 for the case
where the dual norm $\left\|\cdot\right\|_{*}$ is continuously differentiable
at $\nabla L({\bm{x}}_{\infty})$ (like $\ell_{2}$-norm and $\nabla
L({\bm{x}}_{\infty})\neq\mathbf{0}$). For sufficiently large $t$, when $\nabla
L({\bm{x}}_{t})$ gets sufficiently close to $\nabla L({\bm{x}}_{\infty})$, the
descent direction $-{\bm{\Delta}}_{t}$ is unique, equal to
$-\nabla\left\|\nabla L({\bm{x}}_{t})\right\|_{*}$ by Lemma 2.2, and satisfies
$\left\langle{\bm{\Delta}}_{t},\nabla
L({\bm{x}}_{t})\right\rangle=\left\|\nabla L({\bm{x}}_{t})\right\|_{*}$.
Taking $t\to\infty$ we get $\left\langle-\nabla\left\|\nabla
L({\bm{x}}_{\infty})\right\|_{*},\nabla
L({\bm{x}}_{\infty})\right\rangle=\left\|\nabla
L({\bm{x}}_{\infty})\right\|_{*}$. Moreover, we must have $\nabla\left\|\nabla
L({\bm{x}}_{\infty})\right\|_{*}+\lambda{\bm{x}}_{\infty}=\lim_{t\to\infty}\left(\nabla\left\|\nabla
L({\bm{x}}_{t})\right\|_{*}+\lambda{\bm{x}}_{t}\right)=0$, otherwise
${\bm{x}}_{t}$ keeps moving towards
$\lim_{t\to\infty}\left(\nabla\left\|\nabla
L({\bm{x}}_{t})\right\|_{*}+\lambda{\bm{x}}_{t}\right)$ since
$\sum_{t=1}^{\infty}\eta_{t}=\infty$ and thus ${\bm{x}}_{\infty}$ cannot be
the limit. This implies the second condition in Lemma 3.8. The first condition
that $\left\|{\bm{x}}_{\infty}\right\|\leq\frac{1}{\lambda}$ is immediate from
Lemma 3.1 and that $\sum_{t=1}^{\infty}\eta_{t}=\infty$.

However, the above intuition no longer works when dual norm
$\left\|\cdot\right\|_{*}$ is not differentiable at $\nabla
L({\bm{x}}_{\infty})$. This could happen for $\ell_{\infty}$ norm where the
dual norm is $\ell_{1}$ norm and $\nabla L({\bm{x}}_{\infty})$ with
coordinates of value $0$, because the subgradient of absolute value function
at $0$, $\partial|0|$, could be anything between $-1$ and $1$. And more
generally, this could happen for any norm and $\nabla
L({\bm{x}}_{\infty})=\mathbf{0}$. If the limit point ${\bm{x}}_{\infty}$ has
zero gradient for $L$, then the steepest descent direction $-{\bm{\Delta}}$ is
provably not continuous around ${\bm{x}}_{\infty}$.

The following lemma (Lemma 3.9) circumvents the above issue by considering the
weighted average of past steepest descent directions, which provably
converges, given the iterates $\\{{\bm{x}}_{t}\\}_{t=1}^{\infty}$ converge.
Theorem 3.7 is a direct combination of Lemma 3.9 and Lemma 3.8 and we omit its
proof. The proof of Lemma 3.9 is deferred into Section A.3.

######  Lemma 3.9.

For any learning rate schedule $\\{\eta_{t}\\}_{t=1}^{\infty}$ satisfying
$\sum_{t=1}^{\infty}\eta_{t}=\infty$, if the iterates of
$\mathtt{NSD}$-$\mathtt{WD}$ $\\{{\bm{x}}_{t}\\}_{t=0}^{\infty}$ converges to
some ${\bm{x}}_{\infty}$, we have that

  1. 1.

${\bm{\Delta}}_{\infty}:=\lim\limits_{T\rightarrow\infty}\frac{\sum_{t=1}^{T}\eta_{t}{\bm{\Delta}}_{t}}{\sum_{t=1}^{T}\eta_{t}}$
exists and ${\bm{\Delta}}_{\infty}=-\lambda{\bm{x}}_{\infty}$.

  2. 2.

$\left\langle\nabla
L({\bm{x}}_{\infty}),{\bm{\Delta}}_{\infty}\right\rangle=\left\|\nabla
L({\bm{x}}_{\infty})\right\|_{*}$.

  3. 3.

$\left\|{\bm{\Delta}}_{\infty}\right\|\leq 1$.

##  4 Implicit Bias of AdamW

In this section, we extend the analysis on $\mathtt{NSD}$-$\mathtt{WD}$ in
Section 3 to $\mathtt{AdamW}$ to prove that the converged parameters of
$\mathtt{AdamW}$ is the KKT point of the constrained optimization problem. The
proof relies on an upper bound of average update size of $\mathtt{AdamW}$ and
we find that the bound can also be used to guide hyperparameter tuning in
empirical study.

We first state the analog of Lemma 3.9 for $\mathtt{AdamW}$ with the norm
being $\ell_{\infty}$ norm since we treat $\mathtt{AdamW}$ as a smoothed
version of $\mathtt{SignGD}$, which is Lemma 4.1. Here we additionally assume
that $\\{\eta_{t}\\}_{t=1}^{\infty}$ is non-increasing and ${\bm{\Delta}}_{t}$
is defined as $\frac{{\bm{m}}_{t}}{\sqrt{{\bm{v}}_{t}}}$ from Algorithm 1.
Theorem 1.1 is again a direct combination of Lemma 3.8 and Lemma 4.1.

######  Lemma 4.1.

For non-increasing learning rate schedule $\\{\eta_{t}\\}_{t=0}^{\infty}$
satisfying $\sum_{t=1}^{\infty}\eta_{t}=\infty$ and $\beta_{2}\geq\beta_{1}$,
we get $\\{{\bm{x}}_{t}\\}_{t=1}^{\infty}$ by running AdamW with weight decay
factor $\lambda$. If $\\{{\bm{x}}_{t}\\}_{t=0}^{\infty}$ converges to some
${\bm{x}}_{\infty}$, then it holds that

  1. 1.

${\bm{\Delta}}_{\infty}:=\lim\limits_{T\rightarrow\infty}\frac{\sum_{t=1}^{T}\eta_{t}{\bm{\Delta}}_{t}}{\sum_{t=1}^{T}\eta_{t}}$
exists and ${\bm{\Delta}}_{\infty}=-\lambda{\bm{x}}_{\infty}$.

  2. 2.

$\left\langle\nabla
L({\bm{x}}_{\infty}),{\bm{\Delta}}_{\infty}\right\rangle=\left\|\nabla
L({\bm{x}}_{\infty})\right\|_{1}$.

  3. 3.

$\left\|{\bm{\Delta}}_{\infty}\right\|_{\infty}\leq 1$.

The condition $\beta_{1}\leq\beta_{2}$ is necessary for the conclusion to
hold. Otherwise, the iterates can converge outside the $\ell_{\infty}$ norm
ball with radius $\frac{1}{\lambda}$ as shown in Section B.3.

The first two properties in Lemma 4.1 follow from a similar argument for Lemma
4.1, and the main technical difficulty here lies in the proof of the third
property. This is because for any single $t$,
$\left\|{\bm{\Delta}}_{t}\right\|$ could be larger than $1$, which is
different from the case of $\mathtt{NSD}$-$\mathtt{WD}$. To prove the third
property, we need a tight upper bound for the average update size of
$\mathtt{Adam}$-like update rule, which is Lemma 4.2. The proof of Lemma 4.1
is deferred to Appendix B.

###  4.1 Upper bound for average update size of $\mathtt{Adam}$

As mentioned earlier, $\mathtt{Adam}$ updates
$\left\|\frac{{\bm{m}}_{t}}{\sqrt{{\bm{v}}_{t}}}\right\|$ can easily go beyond
$1$ and thus we prove the following upper bound for the average update size of
$\mathtt{Adam}$ (Lemma 4.2). The proof of Lemma 4.2 is deferred to Section
B.1.

######  Lemma 4.2.

Given any $\beta_{1}\leq\beta_{2}<1$, suppose scalar sequences
$\\{v_{t}\\}_{t=0}^{\infty}$ and $\\{g_{t}\\}_{t=1}^{\infty}$ satisfy that
$v_{0}\geq 0,v_{1}>0$ and $v_{t}-\beta_{2}v_{t-1}\geq(1-\beta_{2})g_{t}^{2}$
for $t\geq 1$. Given initial value $|m_{0}|\leq\sqrt{v_{0}}$, define
$m_{t}=\beta_{1}m_{t-1}+(1-\beta_{1})g_{t}$ and
$\Delta_{t}=\frac{m_{t}}{\sqrt{v_{t}}}$ for $t\geq 1$. For any coefficients
$\\{\eta_{t}\\}_{t=1}^{\infty}$ and $T\in\mathbb{N}$, it always holds that

$\displaystyle\frac{|\sum_{t=1}^{T}\eta_{t}\Delta_{t}|}{\sum_{t=1}^{T}\eta_{t}}$
$\displaystyle\leq$
$\displaystyle\left(1+\frac{\beta_{2}-\beta_{1}}{1-\beta_{2}}\frac{\sum_{t=1}^{T}\eta_{t}\beta_{1}^{t-1}}{\sum_{t=1}^{T}\eta_{t}}+\frac{(\beta_{2}-\beta_{1})(1-\beta_{1})}{(1-\beta_{2})\sum_{t=1}^{T}\eta_{t}}\sum_{t=2}^{T}\left(\eta_{t}\frac{1-\beta_{1}^{t-1}}{1-\beta_{1}}-\sum_{i=1}^{T-t}\eta_{t+i}\beta_{1}^{i-1}\right)\ln{\frac{v_{t}}{v_{1}}}\right)^{\frac{1}{2}}.$
(4)

In particular, when $\beta_{1}=\beta_{2}$, it even holds that
$|\Delta_{t}|\leq 1$.

Note $\\{v_{t}\\}_{t=0}^{\infty}$ here only needs to satisfy a more general
condition rather than to be the exact moving average of $g_{t}^{2}$. It can be
applied to the practical scenario where a small positive constant $\epsilon$
is added to $\sqrt{{\bm{v}}_{t}}$ in the denominator to improve the numerical
stability of $\mathtt{Adam}$. It is easy to verify that for ${\bm{v}}_{t}$ in
Algorithm 1, we have that

$\displaystyle(\sqrt{{\bm{v}}_{t}}+\epsilon)^{2}-\beta_{2}(\sqrt{{\bm{v}}_{t-1}}+\epsilon)^{2}$
$\displaystyle\geq$
$\displaystyle({\bm{v}}_{t}-\beta_{2}{\bm{v}}_{t-1})+2\epsilon(\sqrt{{\bm{v}}_{t}}-\beta_{2}\sqrt{{\bm{v}}_{t-1}})$
$\displaystyle=$
$\displaystyle(1-\beta_{2}){\bm{g}}_{t}^{2}+2\epsilon(\sqrt{\beta_{2}^{2}{\bm{v}}_{t-1}+(1-\beta_{2}){\bm{g}}_{t}^{2}}-\sqrt{\beta_{2}^{2}{\bm{v}}_{t-1}})$
$\displaystyle\geq$ $\displaystyle(1-\beta_{2}){\bm{g}}_{t}^{2}.$

Therefore, for $\mathtt{Adam}$ with $\epsilon$, $v_{t}$ in Lemma 4.2 is always
lower bounded, and if we further have an upper bound for gradients, then we
can easily control the average update size of $\mathtt{Adam}$. One nice
property is that the upper bound only scales up logarithmically to
$1/\epsilon$, instead of linearly, as the naive upper bound scales.

#### Relationship of $\ell_{\infty}$ norm and hyperparameters

Another application of Lemma 4.2 is to provide a tight upper bound for the
norm of iterates for any setting, _e.g._ , before convergence or even when the
gradient is stochastic. In particular, when the learning rate does not change
over steps, we have the following upper bound whose proof is in Section B.4.

######  Lemma 4.3.

For any coordinate $j\in[d]$, for $\mathtt{AdamW}$ with constant learning rate
$\eta$ and weight decay factor $\lambda$, with $C\triangleq\max\limits_{1\leq
t\leq T}\left|\ln{\frac{{\bm{v}}_{t,j}}{{\bm{v}}_{1,j}}}\right|$, it holds
that

$\lambda\left|{\bm{x}}_{T,j}\right|-1\leq(1-\lambda\eta)^{T}\lambda\left|{\bm{x}}_{0,j}\right|+\frac{\lambda\eta(\beta_{2}-\beta_{1})\left[2C+\beta_{1}^{T}+(1-\lambda\eta)^{T}\right]}{2(1-\beta_{2})|1-\lambda\eta-\beta_{1}|}.$
(5)

When $\beta_{1}=\beta_{2}$, we only need
$T=\Omega\left(\frac{\log\left\|{\bm{x}}_{0}\right\|_{\infty}}{\lambda\eta}\right)$
to guarantee that $\left|{\bm{x}}_{T},j\right|$ is no larger than
$\frac{1}{\lambda}$ for any $\lambda\eta\leq 1$. However, when
$\beta_{1}<\beta_{2}$ and $\beta_{1}<1-\lambda\eta$, the dominating term on
the right-hand side is
$C\cdot\frac{\eta\lambda(\beta_{2}-\beta_{1})}{(1-\beta_{2})(1-\eta\lambda-\beta_{1})}$.
Assuming $C=O(1)$, it also requires $\lambda\eta\ll 1-\beta_{2}<1-\beta_{1}$
or $\lambda\eta<1-\beta_{2}\approx 1-\beta_{1}$ to ensure the remaining term
is small.

##  5 Experiments

In this section, we run experiments to verify the theoretical claims. In
Section 5.1, we show that the $\ell_{\infty}$ norm of iterates by
$\mathtt{AdamW}$ can converge below $\frac{1}{\lambda}$ as shown in Theorem
1.1 even when the function is non-convex. In Section 5.2, we show that
steepest descent w.r.t. $\ell_{\infty}$ norm works better than w.r.t.
$\ell_{2}$ norm for a specific function, which has better properties under
$\ell_{\infty}$ geometry.

###  5.1 Language modeling task on PTB

We train a small two-layer transformer for language modeling task on the Penn
Treebank dataset (PTB) (Marcus et al., 1993) based on the implementation
provided by Kunstner et al. (2022). We train the model in full batch without
dropout in order to get deterministic gradients and follow the constant
learning rate setting for the total $12800$ epochs. The learning rate $\eta$
is $\sqrt{10}\times 10^{-3}$222We follow the tuning process in Kunstner et al.
(2022) for $\mathtt{AdamW}$ and $\sqrt{10}\times 10^{-3}$ achieves the best
training performance. $\sqrt{10}\times 10^{-3}$ also achieves the best
training performance for full-batch $\mathtt{Adam}$ in Kunstner et al. (2022).
. For each setting of $\beta_{1}$, $\beta_{2}$, we use $\mathtt{Adam}$ and
$\mathtt{AdamW}$ with weight decay coefficient $\lambda=1,2$ to compare the
$\ell_{\infty}$ norm for iterates in each optimizer. We employ the standard
implementation in PyTorch but set $\epsilon$ to be $10^{-16}$ in
$\mathtt{Adam}$ and $\mathtt{AdamW}$ rather than $0$ to avoid division by zero
error because the gradient of some coordinates is $0$ in the first iteration.
Each run is repeated for $2$ random seeds to show the robustness of our claim.
More details can be found in Appendix C.

(a) $\beta_{1}=\beta_{2}=0.99$

(b) $\beta_{1}=\beta_{2}=0.999$

(c) $\beta_{1}=0.9$, $\beta_{2}=0.95$

(d) $\beta_{1}=0.9,\beta_{2}=0.999$

Figure 1: The $\ell_{\infty}$ norm of parameters during the training process
of language modeling task on PTB. The complete results for $\mathtt{Adam}$ are
in Figure 5. As predicted by Lemma 4.3, $\ell_{\infty}$ norm can be bounded by
$\frac{1}{\lambda}$ when $\beta_{1}=\beta_{2}$ or $\lambda\eta\ll
1-\beta_{2}<1-\beta_{1}$. However, for the default setting $\beta_{1}=0.9$ and
$\beta_{2}=0.999$, the $\ell_{\infty}$ norm of $\mathtt{AdamW}$ may not be
bounded by $\frac{1}{\lambda}$ because $1-\beta_{2}<\lambda\eta<1-\beta_{1}$.

From the discussion in Lemma 4.3, it requires either
$\beta_{1}\approx\beta_{2}$ or $\lambda\eta\ll 1-\beta_{2}$ when
$\beta_{1}<\beta_{2}$ in order for $\left\|{\bm{x}}_{T}\right\|_{\infty}$ to
be bounded by $\frac{1}{\lambda}$. To verify the first case, we employ the two
hyperparamter settings where $\beta_{1}=\beta_{2}=0.99$ and
$\beta_{1}=\beta_{2}=0.999$. To verify the second case, we employ the
hyperparameter setting where $\beta_{1}=0.9$ and $\beta_{2}=0.95$. The
$\ell_{\infty}$ norm of iterates for one random seed are shown in Figures
1(c), 1(a) and 1(b). In order to show the details around $0.5$ and $1$, we
truncate the range of y-axis and the full result is plotted in Figure 5 in
Appendix C. In all these three settings, the $\ell_{\infty}$ norm of iterates
in $\mathtt{Adam}$ keeps increasing while the $\ell_{\infty}$ norm of iterates
in $\mathtt{AdamW}$ is constrained below $1$ and $\frac{1}{2}$ for $\lambda=1$
and $2$ respectively. The results for another random seed show similar pattern
and are plotted in Figure 4 and Figure 6.

We also show that the condition is necessary for empirical study by training
with default $\beta_{1}=0.9$ and $\beta_{2}=0.999$. Now
$1-\beta_{2}<\lambda\eta$ and $\beta_{1}\neq\beta_{2}$, which breaks the
condition. The $\ell_{\infty}$ norm of iterates are shown in Figure 1(d). The
$\ell_{\infty}$ norm of $\mathtt{AdamW}$ can not be constrained by $1$ and
$\frac{1}{2}$ for $\lambda=1$ and $2$.

To study the phenomenon when the full batch assumption is removed, we repeat
the same experiment with the batch size being half of the training dataset
size. The results are shown in Figure 2. Note our norm upper bound Lemma 4.3
still holds in this case because it allows any sequence of gradients.
Interestingly, similar to the full batch case Figure 1, $\ell_{\infty}$ norm
still approaches $\frac{1}{\lambda}$, suggesting that our theory still
provides useful prediction even for stochastic $\mathtt{AdamW}$, if the noise
is small.

(a) $\beta_{1}=\beta_{2}=0.99$

(b) $\beta_{1}=\beta_{2}=0.999$

(c) $\beta_{1}=0.9$, $\beta_{2}=0.95$

(d) $\beta_{1}=0.9,\beta_{2}=0.999$

Figure 2: $\ell_{\infty}$ norm of parameters when the batch size is half of
the entire training set. Our norm upper bound Lemma 4.3 still holds in this
case.

###  5.2 A synthetic problem

As mentioned in Section 3.1, we construct the loss function Equation 3 for
some ${\bm{x}}^{*}\in\mathbb{R}^{100}$. The first $10$ coordinates of
${\bm{x}}^{*}$ is $1$ while the rest $90$ coordinates are uniformly sampled
between $[-1,1]$. With such initialization, the $\ell_{\infty}$ norm of
${\bm{x}}^{*}$ is always upper bounded by $1$ but the $\ell_{2}$ norm can be
as large as $10$.

We want to verify that an optimization algorithm can take advantage when it
employs the norm that is more suitable for the loss function. So we implement
normalized steepest descent and steepest descent with $\ell_{\infty}$ norm and
$\ell_{2}$ norm. In order to test the effect of weight decay, we also
implement normalized steepest descent with weight decay for both norms. In
order to be able to reach the global minimizer, the weight decay factor
$\lambda$ is set to be the lower bound of
$\frac{1}{\left\|{\bm{x}}^{*}\right\|}$, which is $1$ for $\ell_{\infty}$ norm
and $0.1$ for $\ell_{2}$ norm. The learning rate can be set according to
theoretical results. The learning rate for steepest descent is constant
$\frac{1}{H}$ in which $H$ is the smoothness coefficient. By Theorem 3.5, the
learning rate for normalized steepest descent with weight decay factor
$\lambda$ should be set as $\eta_{t}=\frac{2}{\lambda(t+1)}$. We also use the
same learning rate schedule for normalized steepest descent without weight
decay.

All the algorithms receive the same initialization ${\bm{x}}_{0}$, whose
coordinate is uniformly initialized from $[-5,5]$. The results for the first
$100$ iterations are shown in Figure 3. The steepest descent w.r.t.
$\ell_{\infty}$ norm always performs better than the steepest descent w.r.t.
$\ell_{2}$ norm no matter whether the update is normalized or not. For both
norms, the performance of the normalized steepest descent is improved when
weight decay is activated.

Figure 3: For both $\ell_{2}$ and $\ell_{\infty}$ norm, we plot training loss
of normalized steepest descent w. and w.o. weight decay and unnormalized
steepest descent over the quadratic loss
$g({\bm{x}})=\sum_{i=1}^{100}\frac{({\bm{x}}_{i}-{\bm{x}}^{*}_{i})^{2}}{i^{2}}$.
When weight decay is turned on, it is set as
$\frac{1}{\left\|{\bm{x}}^{*}\right\|}$ to preserve the optimal value even
with the norm constraints Theorem 3.4. We find that  $\ell_{\infty}$ norm
always outperforms $\ell_{2}$ norm regardless of the usage of weight decay and
irrespective of whether the steepest descent method is normalized. The usage
of weight decay accelerates the optimization for both $\ell_{\infty}$ norm and
$\ell_{2}$ norm.

##  6 Related Work

#### Adaptive Methods:

While Stochastic Gradient Descent (Robbins & Monro, 1951) remains popular for
optimizing deep learning models like ResNet (He et al., 2016), only adaptive
methods can efficiently train recently-emerged large language models (Zhang et
al., 2020). There has been a fruitful amount of research on adaptive gradient
method, including AdaGrad (Duchi et al., 2011), RMSProp (Tieleman & Hinton,
2012), AdaDelta (Zeiler, 2012), Adam (Kingma & Ba, 2014), AdaFactor (Shazeer &
Stern, 2018), AMSGrad (Reddi et al., 2018), AdaBound (Luo et al., 2018), Lion
(Chen et al., 2024), _etc_. Recently there have been also adaptive methods
attempting to accelerate by leveraging the second-order information, _e.g._ ,
AdaHessian (Yao et al., 2021) and Sophia (Liu et al., 2023). However, most
algorithms that are able to train large language models adopt coordinate-wise
adaptivity. In contrast, stochastic gradient descent, even equipped with
global gradient norm clipping, cannot match the performance of coordinate-wise
adaptive algorithms on language tasks (Li et al., 2022a). Previous work has
given convergence rate for RMSProp and Adam under different assumptions (Chen
et al., 2018; Zou et al., 2019; Shi & Li, 2021; Guo et al., 2021; Défossez et
al., 2022; Zhang et al., 2022).

Our work shows that $\mathtt{AdamW}$ and $\mathtt{SignGD}$ with weight decay
converge to the same point assuming convergence. Balles & Hennig (2018);
Kunstner et al. (2022) point out that the similarity with $\mathtt{SignGD}$
largely accounts for the advantage of $\mathtt{Adam}$ over $\mathtt{SGD}$.
Moreover, when $\mathtt{SignGD}$ is equipped with momentum which is one key
component of $\mathtt{Adam}$, it can achieve comparable empirical results with
$\mathtt{Adam}$ for various tasks (Balles & Hennig, 2018; Kunstner et al.,
2022; Bernstein et al., 2018; Crawshaw et al., 2022).

#### Role of Weight Decay:

The usage of weight decay, which refers to shrinking the parameter by a small
constant fraction, can be dated back to the 1980s (Rumelhart et al., 1986;
Hinton, 1987). It has been recognized as a standard trick to improve the
generalization performance of neural networks (Krogh & Hertz, 1991; Bos &
Chug, 1996) for a long time. Krizhevsky et al. (2012) first noticed that
weight decay can sometimes accelerate optimization in deep learning. For
modern architectures equipped with normalization layers, _e.g._ , BatchNorm
(Ioffe & Szegedy, 2015) and LayerNorm (Ba et al., 2016), only the direction of
the parameters before normalization layers matters, rather than their norms.
Turning on weight decay in such settings changes the effective learning rate
of the parameters (Hoffer et al., 2018; Arora et al., 2018; Zhang et al.,
2018; Li & Arora, 2019; Li et al., 2020).

Though weight decay is equivalent to $\ell_{2}$ regularization for SGD, for
steepest descent methods with general norms and adaptive methods like
$\mathtt{Adam}$, they lead to different optimization trajectories (Loshchilov
& Hutter, 2018; Zhang et al., 2018; Zhuang et al., 2022). The empirical
benefit of weight decay over $\ell_{2}$ regularization when they are different
is not well-understood in theory.

#### Implicit Regularization:

Our main result Theorem 1.1 shows that $\mathtt{AdamW}$ regularizes the
$\ell_{\infty}$ norm of the learned solution implicitly through modifying the
optimization dynamics, rather than directly modifying the objective, like
$\mathtt{Adam}$ with $\ell_{2}$ regularization. This kind of behavior is
termed _implicit regularization_ or _implicit bias_ of optimization
algorithms. Though there has been a large volume of works studying the
implicit bias of (Stochastic) GD and its non-adaptive variants, including
settings related to max margin (Soudry et al., 2018; Gunasekar et al., 2018;
Nacson et al., 2019a, b; Ji & Telgarsky, 2018; Lyu & Li, 2019), initialization
with norm (Gunasekar et al., 2017; Arora et al., 2019a; Li et al., 2019; Lyu
et al., 2021), kernel regime (Jacot et al., 2018; Arora et al., 2019b, c), and
flatness (Blanc et al., 2020; Damian et al., 2021; Li et al., 2021; Arora et
al., 2022; Li et al., 2022b; Wen et al., 2022; Damian et al., 2022), very few
results are known about the implicit bias of adaptive methods like
$\mathtt{Adam}$. Wilson et al. (2017) shows that for linear regression
problem, adaptive methods can converge to a solution whose elements have the
same magnitude under certain conditions. Their converged solution thus has
small $\ell_{\infty}$ norm while the converged solution of non-adaptive
methods is known to have the smallest $\ell_{2}$ norm among all the global
minimizers. Wang et al. (2021) shows that $\mathtt{Adam}$ behaves similarly to
non-adaptive methods like GD when the cross-entropy loss converges to $0$ due
to the positive numerical stability hyperparameter $\epsilon$ in the
denominator of $\mathtt{Adam}$’s update rule. The theoretical derivation by
Cattaneo et al. (2023) argues that $\mathtt{Adam}$ tends to find interpolating
solutions with small $\ell_{1}$ norm.

The concurrent work by Chen et al. (2023) is arguably the most related work to
us, where the recently discovered optimization algorithm by auto-search, Lion
(Chen et al., 2024), is elegantly generalized to a family of algorithms,
Lion-$\mathcal{K}$, where $\mathcal{K}$ is some convex function. When
$\mathcal{K}$ is chosen to be the dual norm and momentum in Lion-$\mathcal{K}$
is turned off, Lion-$\mathcal{K}$ becomes the normalized steepest descent.
Their analysis shows that even with momentum, the steepest normalized descent
with weight decay can be viewed as optimization under the original norm
constraint. However, in any Lion-$\mathcal{K}$ algorithm, the update at one
step $t$ only depends on past iterates through first-order momentum
${\bm{m}}_{t}$. Their analysis cannot be applied to $\mathtt{AdamW}$ because
$\mathtt{AdamW}$ cannot be written in the form of Lion-$\mathcal{K}$ for any
convex function $\mathcal{K}$. To see this, simply note that the update of
Lion-$\mathcal{K}$ for a fixed $\mathcal{K}$ is completely determined by
${\bm{g}}_{t},{\bm{m}}_{t}$ and ${\bm{x}}_{t}$ while the update of
$\mathtt{AdamW}$ can still be different if the second order momentum
${\bm{v}}_{t}$ is different. In terms of proof technique, Chen et al. (2023)
constructs the Lyapunov function while we directly characterize the KKT point
and connect the converged point to KKT point through the weighted average
update.

##  7 Discussion and Future Works

This work focuses on the implicit bias of $\mathtt{AdamW}$ in the
deterministic (or full-batch) case. Though our upper bound on the average
update size of $\mathtt{Adam}$ holds unconditionally on the input gradients,
regardless of stochasticity or not, it is unlikely that the
$\frac{1}{\lambda}$ upper bound can be reached when there is large gradient
noise, especially when $\beta_{2}$ is very close to $1$. In that case, the
denominator of the update of $\mathtt{AdamW}$ is roughly the square root of
the square of the expected gradient plus some additional gradient variance
term, which strictly dominates the expected gradient in the numerator. Malladi
et al. (2022) uses Stochastic Differential Equation (SDE) approximation to
model the trajectories of $\mathtt{Adam}$ in such regime and empirically tests
the implication of SDE approximation, namely the square root scaling rule.

The most interesting future direction is to understand in what sense the
optimization advantage of coordinate-wise adaptive methods like
$\mathtt{Adam}$ over standard SGD for language modeling tasks can be explained
by the conjecture implied by the findings of Kunstner et al. (2022), that _the
loss function for language modeling tasks has better properties under
$\ell_{\infty}$ geometry_. It would be interesting to understand if the loss
landscape in real-world applications shares common properties with our toy
quadratic example and induces similar results in Figure 3.

Another important future direction is to provide non-asymptotic convergence
rates for $\mathtt{AdamW}$ in both convex and non-convex settings.

##  8 Conclusions

We make the first step towards understanding the benefit of $\mathtt{AdamW}$
over $\mathtt{Adam}$ with $\ell_{2}$ regularization by characterizing the
implicit bias of $\mathtt{AdamW}$, _i.e._ , it can only converge to KKT points
of the $\ell_{\infty}$ norm constrained optimization problem. There are two
main insights behind this result: (1) $\mathtt{Adam}$ is a smoothed version of
$\mathtt{SignGD}$, which is the normalized steepest descent w.r.t.
$\ell_{\infty}$ norm; (2) for any norm, the corresponding normalized steepest
descent with weight decay is essentially Frank-Wolfe over the corresponding
norm ball, which is known to perform constrained optimization. Our main
technical contribution is a tight upper bound of the average update size of
$\mathtt{Adam}$ updates. We test its prediction on the relationship between
the $\ell_{\infty}$ norm of the parameters and the $\mathtt{AdamW}$
hyperparameters $\eta,\lambda$, $\beta_{1}$, $\beta_{2}$ on a language
modeling task.

## References

  * Arora et al. (2018) Sanjeev Arora, Zhiyuan Li, and Kaifeng Lyu.  Theoretical Analysis of Auto Rate-Tuning by Batch Normalization.  In _International Conference on Learning Representations_ , 2018.
  * Arora et al. (2019a) Sanjeev Arora, Nadav Cohen, Wei Hu, and Yuping Luo.  Implicit regularization in deep matrix factorization.  In _Advances in Neural Information Processing Systems_ , 2019a.
  * Arora et al. (2019b) Sanjeev Arora, Simon Du, Wei Hu, Zhiyuan Li, and Ruosong Wang.  Fine-grained analysis of optimization and generalization for overparameterized two-layer neural networks.  In _International Conference on Machine Learning_. PMLR, 2019b.
  * Arora et al. (2019c) Sanjeev Arora, Simon S Du, Wei Hu, Zhiyuan Li, Russ R Salakhutdinov, and Ruosong Wang.  On exact computation with an infinitely wide neural net.  _Advances in neural information processing systems_ , 32, 2019c.
  * Arora et al. (2022) Sanjeev Arora, Zhiyuan Li, and Abhishek Panigrahi.  Understanding gradient descent on the edge of stability in deep learning.  In _International Conference on Machine Learning_. PMLR, 2022.
  * Ba et al. (2016) Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton.  Layer normalization.  _arXiv preprint arXiv:1607.06450_ , 2016.
  * Balles & Hennig (2018) Lukas Balles and Philipp Hennig.  Dissecting adam: The sign, magnitude and variance of stochastic gradients.  In _International Conference on Machine Learning_ , 2018.
  * Bernstein et al. (2018) Jeremy Bernstein, Yu-Xiang Wang, Kamyar Azizzadenesheli, and Animashree Anandkumar.  signSGD: Compressed optimisation for non-convex problems.  In _International Conference on Machine Learning_ , 2018.
  * Blanc et al. (2020) Guy Blanc, Neha Gupta, Gregory Valiant, and Paul Valiant.  Implicit regularization for deep neural networks driven by an ornstein-uhlenbeck like process.  In _Conference on learning theory_. PMLR, 2020.
  * Bolte et al. (2023) Jérôme Bolte, Cyrille W Combettes, and Edouard Pauwels.  The iterates of the Frank–Wolfe algorithm may not converge.  _Mathematics of Operations Research_ , 2023.
  * Bos & Chug (1996) Siegfried Bos and E Chug.  Using weight decay to optimize the generalization ability of a perceptron.  In _Proceedings of International Conference on Neural Networks (ICNN’96)_. IEEE, 1996.
  * Cattaneo et al. (2023) Matias D Cattaneo, Jason M Klusowski, and Boris Shigida.  On the implicit bias of adam.  _arXiv preprint arXiv:2309.00079_ , 2023.
  * Chen et al. (2023) Lizhang Chen, Bo Liu, Kaizhao Liang, et al.  Lion Secretly Solves a Constrained Optimization: As Lyapunov Predicts.  In _The Twelfth International Conference on Learning Representations_ , 2023.
  * Chen et al. (2024) Xiangning Chen, Chen Liang, Da Huang, Esteban Real, Kaiyuan Wang, Hieu Pham, Xuanyi Dong, Thang Luong, Cho-Jui Hsieh, Yifeng Lu, et al.  Symbolic discovery of optimization algorithms.  _Advances in Neural Information Processing Systems_ , 2024.
  * Chen et al. (2018) Xiangyi Chen, Sijia Liu, Ruoyu Sun, and Mingyi Hong.  On the Convergence of A Class of Adam-Type Algorithms for Non-Convex Optimization.  In _International Conference on Learning Representations_ , 2018.
  * Crawshaw et al. (2022) Michael Crawshaw, Mingrui Liu, Francesco Orabona, Wei Zhang, and Zhenxun Zhuang.  Robustness to Unbounded Smoothness of Generalized SignSGD.  In _Advances in Neural Information Processing Systems_ , 2022.
  * Damian et al. (2021) Alex Damian, Tengyu Ma, and Jason D Lee.  Label Noise SGD Provably Prefers Flat Global Minimizers.  In _Advances in Neural Information Processing Systems_ , 2021.
  * Damian et al. (2022) Alex Damian, Eshaan Nichani, and Jason D Lee.  Self-Stabilization: The Implicit Bias of Gradient Descent at the Edge of Stability.  In _The Eleventh International Conference on Learning Representations_ , 2022.
  * Défossez et al. (2022) Alexandre Défossez, Leon Bottou, Francis Bach, and Nicolas Usunier.  A Simple Convergence Proof of Adam and Adagrad.  _Transactions on Machine Learning Research_ , 2022.
  * Duchi et al. (2011) John Duchi, Elad Hazan, and Yoram Singer.  Adaptive subgradient methods for online learning and stochastic optimization. _Journal of machine learning research_ , 2011.
  * Frank et al. (1956) Marguerite Frank, Philip Wolfe, et al.  An algorithm for quadratic programming.  _Naval research logistics quarterly_ , 1956.
  * Gunasekar et al. (2017) Suriya Gunasekar, Blake E Woodworth, Srinadh Bhojanapalli, Behnam Neyshabur, and Nati Srebro.  Implicit regularization in matrix factorization.  In _Advances in Neural Information Processing Systems_ , 2017.
  * Gunasekar et al. (2018) Suriya Gunasekar, Jason Lee, Daniel Soudry, and Nathan Srebro.  Characterizing implicit bias in terms of optimization geometry.  In _International Conference on Machine Learning_. PMLR, 2018.
  * Guo et al. (2021) Zhishuai Guo, Yi Xu, Wotao Yin, Rong Jin, and Tianbao Yang.  A novel convergence analysis for algorithms of the adam family.  _arXiv preprint arXiv:2112.03459_ , 2021.
  * He et al. (2016) Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun.  Deep residual learning for image recognition.  In _Proceedings of the IEEE conference on computer vision and pattern recognition_ , 2016.
  * Hinton (1987) Geoffrey E Hinton.  Learning translation invariant recognition in a massively parallel networks.  In _International conference on parallel architectures and languages Europe_ , 1987.
  * Hoffer et al. (2018) Elad Hoffer, Ron Banner, Itay Golan, and Daniel Soudry.  Norm matters: efficient and accurate normalization schemes in deep networks.  In _Advances in Neural Information Processing Systems_ , 2018.
  * Ioffe & Szegedy (2015) Sergey Ioffe and Christian Szegedy.  Batch normalization: Accelerating deep network training by reducing internal covariate shift.  In _International conference on machine learning_ , 2015.
  * Jacot et al. (2018) Arthur Jacot, Franck Gabriel, and Clément Hongler.  Neural tangent kernel: Convergence and generalization in neural networks.  In _Advances in neural information processing systems_ , 2018.
  * Jaggi (2013) Martin Jaggi.  Revisiting Frank-Wolfe: Projection-free sparse convex optimization.  In _International conference on machine learning_ , 2013.
  * Ji & Telgarsky (2018) Ziwei Ji and Matus Telgarsky.  Gradient descent aligns the layers of deep linear networks.  In _International Conference on Learning Representations_ , 2018.
  * Kingma & Ba (2014) Diederik P Kingma and Jimmy Ba.  Adam: A method for stochastic optimization.  _arXiv preprint arXiv:1412.6980_ , 2014.
  * Krizhevsky et al. (2012) Alex Krizhevsky, Ilya Sutskever, and Geoffrey E Hinton.  Imagenet classification with deep convolutional neural networks.  In _Advances in neural information processing systems_ , 2012.
  * Krogh & Hertz (1991) Anders Krogh and John Hertz.  A simple weight decay can improve generalization.  _Advances in neural information processing systems_ , 1991.
  * Kunstner et al. (2022) Frederik Kunstner, Jacques Chen, Jonathan Wilder Lavington, and Mark Schmidt.  Noise Is Not the Main Factor Behind the Gap Between Sgd and Adam on Transformers, But Sign Descent Might Be.  In _The Eleventh International Conference on Learning Representations_ , 2022.
  * Li et al. (2019) Yuanzhi Li, Colin Wei, and Tengyu Ma.  Towards explaining the regularization effect of initial large learning rate in training neural networks.  In _Advances in Neural Information Processing Systems_ , 2019.
  * Li & Arora (2019) Zhiyuan Li and Sanjeev Arora.  An Exponential Learning Rate Schedule for Deep Learning.  In _International Conference on Learning Representations_ , 2019.
  * Li et al. (2020) Zhiyuan Li, Kaifeng Lyu, and Sanjeev Arora.  Reconciling modern deep learning with traditional optimization analyses: The intrinsic learning rate.  _Advances in Neural Information Processing Systems_ , 2020.
  * Li et al. (2021) Zhiyuan Li, Tianhao Wang, and Sanjeev Arora.  What Happens after SGD Reaches Zero Loss?–A Mathematical Framework.  In _International Conference on Learning Representations_ , 2021.
  * Li et al. (2022a) Zhiyuan Li, Srinadh Bhojanapalli, Manzil Zaheer, Sashank Reddi, and Sanjiv Kumar.  Robust training of neural networks using scale invariant architectures.  In _International Conference on Machine Learning_ , 2022a.
  * Li et al. (2022b) Zhiyuan Li, Tianhao Wang, and Dingli Yu.  Fast Mixing of Stochastic Gradient Descent with Normalization and Weight Decay.  In _Advances in Neural Information Processing Systems_ , 2022b.
  * Liu et al. (2023) Hong Liu, Zhiyuan Li, David Leo Wright Hall, Percy Liang, and Tengyu Ma.  Sophia: A Scalable Stochastic Second-order Optimizer for Language Model Pre-training.  In _The Twelfth International Conference on Learning Representations_ , 2023.
  * Loshchilov & Hutter (2018) Ilya Loshchilov and Frank Hutter.  Decoupled Weight Decay Regularization.  In _International Conference on Learning Representations_ , 2018.
  * Luo et al. (2018) Liangchen Luo, Yuanhao Xiong, Yan Liu, and Xu Sun.  Adaptive Gradient Methods with Dynamic Bound of Learning Rate.  In _International Conference on Learning Representations_ , 2018.
  * Lyu & Li (2019) Kaifeng Lyu and Jian Li.  Gradient Descent Maximizes the Margin of Homogeneous Neural Networks.  In _International Conference on Learning Representations_ , 2019.
  * Lyu et al. (2021) Kaifeng Lyu, Zhiyuan Li, Runzhe Wang, and Sanjeev Arora.  Gradient Descent on Two-layer Nets: Margin Maximization and Simplicity Bias.  In _Advances in Neural Information Processing Systems_ , 2021.
  * Malladi et al. (2022) Sadhika Malladi, Kaifeng Lyu, Abhishek Panigrahi, and Sanjeev Arora.  On the SDEs and Scaling Rules for Adaptive Gradient Algorithms.  In _Advances in Neural Information Processing Systems_ , 2022.
  * Marcus et al. (1993) Mitch Marcus, Beatrice Santorini, and Mary Ann Marcinkiewicz.  Building a large annotated corpus of English: The Penn Treebank.  _Computational linguistics_ , 1993.
  * Nacson et al. (2019a) Mor Shpigel Nacson, Suriya Gunasekar, Jason Lee, Nathan Srebro, and Daniel Soudry.  Lexicographic and depth-sensitive margins in homogeneous and non-homogeneous deep models.  In _International Conference on Machine Learning_ , 2019a.
  * Nacson et al. (2019b) Mor Shpigel Nacson, Jason Lee, Suriya Gunasekar, Pedro Henrique Pamplona Savarese, Nathan Srebro, and Daniel Soudry.  Convergence of gradient descent on separable data.  In _The 22nd International Conference on Artificial Intelligence and Statistics_ , 2019b.
  * Reddi et al. (2018) Sashank J Reddi, Satyen Kale, and Sanjiv Kumar.  On the Convergence of Adam and Beyond.  In _International Conference on Learning Representations_ , 2018.
  * Robbins & Monro (1951) Herbert Robbins and Sutton Monro.  A stochastic approximation method.  _The annals of mathematical statistics_ , 1951.
  * Rumelhart et al. (1986) David E Rumelhart, Geoffrey E Hinton, and Ronald J Williams.  Learning representations by back-propagating errors.  _nature_ , 1986.
  * Shazeer & Stern (2018) Noam Shazeer and Mitchell Stern.  Adafactor: Adaptive learning rates with sublinear memory cost.  In _International Conference on Machine Learning_ , 2018.
  * Shi & Li (2021) Naichen Shi and Dawei Li.  Rmsprop converges with proper hyperparameter.  In _International conference on learning representation_ , 2021.
  * Soudry et al. (2018) Daniel Soudry, Elad Hoffer, Mor Shpigel Nacson, and Nathan Srebro.  The Implicit Bias of Gradient Descent on Separable Data.  In _International Conference on Learning Representations_ , 2018.
  * Tieleman & Hinton (2012) Tijmen Tieleman and Geoffrey Hinton.  Lecture 6.5-rmsprop, coursera: Neural networks for machine learning.  _University of Toronto, Technical Report_ , 2012.
  * Wang et al. (2021) Bohan Wang, Qi Meng, Wei Chen, and Tie-Yan Liu.  The implicit bias for adaptive optimization algorithms on homogeneous neural networks.  In _International Conference on Machine Learning_ , 2021.
  * Wen et al. (2022) Kaiyue Wen, Tengyu Ma, and Zhiyuan Li.  How Sharpness-Aware Minimization Minimizes Sharpness? In _The Eleventh International Conference on Learning Representations_ , 2022.
  * Wilson et al. (2017) Ashia C Wilson, Rebecca Roelofs, Mitchell Stern, Nati Srebro, and Benjamin Recht.  The marginal value of adaptive gradient methods in machine learning.  In _Advances in Neural Information Processing Systems_ , 2017.
  * Yao et al. (2021) Zhewei Yao, Amir Gholami, Sheng Shen, Mustafa Mustafa, Kurt Keutzer, and Michael Mahoney.  Adahessian: An adaptive second order optimizer for machine learning.  In _proceedings of the AAAI conference on artificial intelligence_ , 2021.
  * Zeiler (2012) Matthew D Zeiler.  ADADELTA: an adaptive learning rate method.  _arXiv preprint arXiv:1212.5701_ , 2012.
  * Zhang et al. (2018) Guodong Zhang, Chaoqi Wang, Bowen Xu, and Roger Grosse.  Three Mechanisms of Weight Decay Regularization.  In _International Conference on Learning Representations_ , 2018.
  * Zhang et al. (2020) Jingzhao Zhang, Sai Praneeth Karimireddy, Andreas Veit, Seungyeon Kim, Sashank Reddi, Sanjiv Kumar, and Suvrit Sra.  Why are adaptive methods good for attention models? _Advances in Neural Information Processing Systems_ , 2020.
  * Zhang et al. (2022) Yushun Zhang, Congliang Chen, Naichen Shi, Ruoyu Sun, and Zhi-Quan Luo.  Adam Can Converge Without Any Modification On Update Rules.  In _Advances in Neural Information Processing Systems_ , 2022.
  * Zhuang et al. (2022) Zhenxun Zhuang, Mingrui Liu, Ashok Cutkosky, and Francesco Orabona.  Understanding adamw through proximal methods and scale-freeness.  _Transactions on Machine Learning Research_ , 2022.
  * Zou et al. (2019) Fangyu Zou, Li Shen, Zequn Jie, Weizhong Zhang, and Wei Liu.  A sufficient condition for convergences of adam and rmsprop.  In _Proceedings of the IEEE/CVF Conference on computer vision and pattern recognition_ , 2019.

###### Contents

  1. 1 Introduction
  2. 2 Preliminaries and Notations
  3. 3 Warm Up: Implicit Bias of Normalized Steepest Descent w. Weight Decay
    1. 3.1 Convex setting: constrained optimization
    2. 3.2 Non-convex setting: convergence to KKT points
  4. 4 Implicit Bias of AdamW
    1. 4.1 Upper bound for average update size of $\mathtt{Adam}$
  5. 5 Experiments
    1. 5.1 Language modeling task on PTB
    2. 5.2 A synthetic problem
  6. 6 Related Work
  7. 7 Discussion and Future Works
  8. 8 Conclusions
  9. A Omitted Proofs in Section 3
    1. A.1 Omitted proofs for convergence into norm ball with bounded update
    2. A.2 Omitted proofs for convergence to constrained minimizer with proper learning rates
    3. A.3 Omitted Proofs for Lemma 3.9
  10. B Omitted Proofs in Section 4
    1. B.1 Omitted proofs for upper bound for average update size of $\mathtt{Adam}$
    2. B.2 Proof for Lemma 4.1
    3. B.3 A counter example when $\beta_{1}>\beta_{2}$
    4. B.4 Proof for upper bound for norm of iterates in $\mathtt{AdamW}$
  11. C Experimental Details and More Results

##  Appendix A Omitted Proofs in Section 3

In this section, we provide the omitted proofs in Section 3, which shows the
iterates and the converged solution by normalized steepest descent with
decoupled weight decay before diving into the analysis on $\mathtt{AdamW}$. In
Section A.1, we prove that the iterates will enter or stay in the norm ball
with radius $\frac{1}{\lambda}$ for any normalized update. In Section A.2, we
prove that the iterates of normalized steepest descent with weight decay will
converge to the constrained minimizer of $L({\bm{x}})$ in the same ball with
proper learning rates.

###  A.1 Omitted proofs for convergence into norm ball with bounded update

###### Proof.

Lemma 3.1 We prove by induction that
$\left\|{\bm{x}}_{t}\right\|\leq\frac{1}{\lambda}+\prod_{i=1}^{t}(1-\lambda\eta_{i})\left(\left\|{\bm{x}}_{0}\right\|-\frac{1}{\lambda}\right)$.

$\displaystyle\left\|{\bm{x}}_{t}\right\|-\frac{1}{\lambda}$
$\displaystyle=\left\|(1-\lambda\eta_{t}){\bm{x}}_{t-1}-\eta_{t}{\bm{\Delta}}_{t}\right\|-\frac{1}{\lambda}$
$\displaystyle\leq(1-\lambda\eta_{t})\left\|{\bm{x}}_{t-1}\right\|+\eta_{t}\left\|{\bm{\Delta}}_{t}\right\|-\frac{1}{\lambda}$
$\displaystyle\leq(1-\lambda\eta_{t})\left\|{\bm{x}}_{t-1}\right\|+\eta_{t}-\frac{1}{\lambda}$
$\displaystyle=(1-\lambda\eta_{t})\left(\left\|{\bm{x}}_{t-1}\right\|-\frac{1}{\lambda}\right)$
$\displaystyle\leq\prod_{i=1}^{t}(1-\lambda\eta_{i})\left(\left\|{\bm{x}}_{0}\right\|-\frac{1}{\lambda}\right).$

When $\left\|{\bm{x}}_{0}\right\|>\frac{1}{\lambda}$, we have that

$\displaystyle\left\|{\bm{x}}_{t}\right\|-\frac{1}{\lambda}$
$\displaystyle\leq\prod_{i=1}^{t}(1-\lambda\eta_{i})\left(\left\|{\bm{x}}_{0}\right\|-\frac{1}{\lambda}\right)$
$\displaystyle\leq\prod_{i=1}^{t}\exp(-\lambda\eta_{i})\left(\left\|{\bm{x}}_{0}\right\|-\frac{1}{\lambda}\right)$
$\displaystyle=\exp(-\lambda\sum_{i=1}^{t}\eta_{i})\left(\left\|{\bm{x}}_{0}\right\|-\frac{1}{\lambda}\right).$

When $\left\|{\bm{x}}_{0}\right\|\leq\frac{1}{\lambda}$,
$\left\|{\bm{x}}_{t}\right\|-\frac{1}{\lambda}\leq 0$. This completes the
proof. ∎

###  A.2 Omitted proofs for convergence to constrained minimizer with proper
learning rates

###### Proof of Lemma 3.3.

For normalized steepest descent update ${\bm{\Delta}}_{t}$ from Equation 1,

$\displaystyle\nabla L({\bm{x}}_{t-1})^{\top}({\bm{x}}_{t}-{\bm{x}}_{t-1})$
$\displaystyle=-\eta_{t}\nabla
L({\bm{x}}_{t-1})^{\top}{\bm{\Delta}}_{t}-\lambda\eta_{t}\nabla
L({\bm{x}}_{t-1})^{\top}{\bm{x}}_{t-1}$ $\displaystyle=-\eta_{t}\left\|\nabla
L({\bm{x}}_{t-1})\right\|_{*}-\lambda\eta_{t}\nabla
L({\bm{x}}_{t-1})^{\top}({\bm{x}}_{t-1}-{\bm{x}}^{*})-\lambda\eta_{t}\nabla
L({\bm{x}}_{t-1})^{\top}{\bm{x}}^{*}$ $\displaystyle\leq-\eta_{t}\left\|\nabla
L({\bm{x}}_{t-1})\right\|_{*}-\lambda\eta_{t}\left(L({\bm{x}}_{t-1})-L({\bm{x}}^{*})\right)+\lambda\eta_{t}\left\|\nabla
L({\bm{x}}_{t-1})\right\|_{*}\left\|{\bm{x}}^{*}\right\|$
$\displaystyle\leq-\lambda\eta_{t}\left(L({\bm{x}}_{t-1})-L({\bm{x}}^{*})\right),$

where the first inequality we use convexity of $L$ and the second inequality
uses $\left\|{\bm{x}}^{*}\right\|\leq 1$.

Since the gradient of $L$ is $H$-lipschitz, by Taylor expansion, we have that

$\displaystyle L({\bm{x}}_{t})-L({\bm{x}}_{t-1})$ $\displaystyle=$
$\displaystyle\int_{0}^{1}\nabla
L\left({\bm{x}}_{t-1}+\alpha({\bm{x}}_{t}-{\bm{x}}_{t-1})\right)^{\top}({\bm{x}}_{t}-{\bm{x}}_{t-1})\,d\alpha$
$\displaystyle=$ $\displaystyle\nabla
L({\bm{x}}_{t-1})^{\top}({\bm{x}}_{t}-{\bm{x}}_{t-1})+\int_{0}^{1}\left(\nabla
L\left({\bm{x}}_{t-1}+\alpha({\bm{x}}_{t}-{\bm{x}}_{t-1})\right)-\nabla
L({\bm{x}}_{t-1})\right)^{\top}({\bm{x}}_{t}-{\bm{x}}_{t-1})\,d\alpha$
$\displaystyle\leq$
$\displaystyle-\lambda\eta_{t}\left(L({\bm{x}}_{t-1})-L({\bm{x}}^{*})\right)+\int_{0}^{1}\left\|\nabla
L\left({\bm{x}}_{t-1}+\alpha({\bm{x}}_{t}-{\bm{x}}_{t-1})\right)-\nabla
L({\bm{x}}_{t-1})\right\|_{*}\left\|{\bm{x}}_{t}-{\bm{x}}_{t-1}\right\|\,d\alpha$
$\displaystyle\leq$
$\displaystyle-\lambda\eta_{t}\left(L({\bm{x}}_{t-1})-L({\bm{x}}^{*})\right)+\int_{0}^{1}H\alpha\left\|{\bm{x}}_{t}-{\bm{x}}_{t-1}\right\|^{2}\,d\alpha$
$\displaystyle=$
$\displaystyle-\lambda\eta_{t}\left(L({\bm{x}}_{t-1})-L({\bm{x}}^{*})\right)+\frac{H}{2}\left\|{\bm{x}}_{t}-{\bm{x}}_{t-1}\right\|^{2}.$

Because the update ${\bm{\Delta}}_{t}$ is normalized and thus have unit norm
by definition, it holds that

$\displaystyle\left\|{\bm{x}}_{t}-{\bm{x}}_{t-1}\right\|^{2}=\left\|-\eta_{t}{\bm{\Delta}}_{t}-\lambda\eta_{t}{\bm{x}}_{t-1}\right\|^{2}\leq\eta_{t}^{2}\left(\left\|{\bm{\Delta}}_{t}\right\|+\lambda\left\|{\bm{x}}_{t-1}\right\|\right)^{2}\leq\eta_{t}^{2}\left(1+\lambda\left\|{\bm{x}}_{t-1}\right\|\right)^{2}.$

Finally, we conclude that

$L({\bm{x}}_{t})-L({\bm{x}}^{*})\leq(1-\lambda\eta_{t})(L({\bm{x}}_{t-1})-L({\bm{x}}^{*}))+\frac{H}{2}\eta_{t}^{2}\left(1+\lambda\left\|{\bm{x}}_{t-1}\right\|\right)^{2}.$

∎

###### Proof of Theorem 3.4.

The proof of Theorem 3.4 is a direct application of Lemma A.1 on the one-step
descent lemma Lemma 3.3. ∎

######  Lemma A.1.

Assume that $\eta_{t}\geq 0$, $\lim_{t\rightarrow\infty}\eta_{t}=0$ and
$\sum_{t=1}^{\infty}\eta_{t}=\infty$. $C$ is any positive number and
$a_{0}\geq 0$. If the sequence $\\{a_{t}\\}_{t=0}^{\infty}$ satisfies that
$a_{t}\leq(1-\eta_{t})a_{t-1}+C\eta_{t}^{2}$, then
$\lim_{t\rightarrow\infty}a_{t}=0$.

###### Proof of Lemma A.1.

First we show by induction that $a_{t}\leq
a_{0}\exp\left(-\sum_{i=1}^{t}\eta_{i}\right)+C\sum_{i=1}^{t}\eta_{i}^{2}\exp\left(-\sum_{j=i+1}^{t}\eta_{j}\right)$.

$\displaystyle a_{t}$ $\displaystyle\leq(1-\eta_{t})a_{t-1}+C\eta_{t}^{2}$
$\displaystyle\leq\exp\left(-\eta_{t}\right)a_{t-1}+C\eta_{t}^{2}$
$\displaystyle\leq\exp\left(-\eta_{t}\right)\left[a_{0}\exp\left(-\sum_{i=1}^{t-1}\eta_{i}\right)+C\sum_{i=1}^{t-1}\eta_{i}^{2}\exp\left(-\sum_{j=i+1}^{t-1}\eta_{j}\right)\right]+C\eta_{t}^{2}$
$\displaystyle=a_{0}\exp\left(-\sum_{i=1}^{t}\eta_{i}\right)+C\sum_{i=1}^{t-1}\eta_{i}^{2}\exp\left(-\sum_{j=i+1}^{t}\eta_{j}\right)+C\eta_{t}^{2}$
$\displaystyle=a_{0}\exp\left(-\sum_{i=1}^{t}\eta_{i}\right)+C\sum_{i=1}^{t}\eta_{i}^{2}\exp\left(-\sum_{j=i+1}^{t}\eta_{j}\right).$

Because $\sum_{t=1}^{\infty}\eta_{t}=\infty$,
$\lim_{t\rightarrow\infty}a_{0}\exp\left(-\sum_{i=1}^{t}\eta_{i}\right)=0$. In
order to show $\lim_{t\rightarrow\infty}a_{t}=0$, it’s sufficient to show
$\lim_{t\rightarrow\infty}\sum_{i=1}^{t}\eta_{i}^{2}\exp\left(-\sum_{j=i+1}^{t}\eta_{j}\right)=0$.

For any $\epsilon>0$, $\eta^{*}$ is chosen such that
$\eta^{*}e^{\eta^{*}}=\frac{\epsilon}{2}$. There exists
$\tau\in\mathbb{N}^{+}$ such that $\eta_{i}\leq\eta^{*}$ for $i\geq\tau$. We
choose $T$ such that
$\exp(-\sum_{j=\tau}^{T}\eta_{j})\leq\frac{\epsilon}{2\sum_{i=1}^{\tau-1}\eta_{i}^{2}}$.
Then for any $t\geq T$, we have the following inequalities

$\displaystyle\sum_{i=1}^{\tau-1}\eta_{i}^{2}\exp\left(-\sum_{j=i+1}^{t}\eta_{j}\right)$
$\displaystyle\leq\sum_{i=0}^{\tau-1}\eta_{i}^{2}\exp\left(-\sum_{j=\tau}^{t}\eta_{j}\right)$
$\displaystyle\leq\left(\sum_{i=1}^{\tau-1}\eta_{i}^{2}\right)\exp\left(-\sum_{j=\tau}^{T}\eta_{j}\right)$
$\displaystyle\leq\frac{\epsilon}{2},$
$\displaystyle\sum_{i=\tau}^{t}\eta_{i}^{2}\exp\left(-\sum_{j=i+1}^{t}\eta_{j}\right)$
$\displaystyle\leq\eta^{*}\sum_{i=\tau}^{t}\eta_{i}\exp\left(-\sum_{j=i+1}^{t}\eta_{j}\right)$
$\displaystyle\leq\eta^{*}\sum_{i=\tau}^{t}\left(\exp(\eta_{i})-1\right)\exp\left(-\sum_{j=i+1}^{t}\eta_{j}\right)$
$\displaystyle=\eta^{*}\sum_{i=\tau}^{t}\exp(\eta_{i})\left(1-\exp(-\eta_{i})\right)\exp\left(-\sum_{j=i+1}^{t}\eta_{j}\right)$
$\displaystyle\leq\eta^{*}\sum_{i=\tau}^{t-1}\exp(\eta^{*})\left(1-\exp(-\eta_{i})\right)\exp\left(-\sum_{j=i+1}^{t}\eta_{j}\right)$
$\displaystyle=\eta^{*}\exp(\eta^{*})\sum_{i=\tau}^{t}\left[\exp\left(-\sum_{j=i+1}^{t}\eta_{j}\right)-\exp\left(-\sum_{j=i}^{t}\eta_{j}\right)\right]$
$\displaystyle=\eta^{*}\exp(\eta^{*})\left[1-\exp\left(-\sum_{j=\tau}^{t}\eta_{j}\right)\right]$
$\displaystyle\leq\eta^{*}\exp(\eta^{*})=\frac{\epsilon}{2},$
$\displaystyle\sum_{i=1}^{t}\eta_{i}^{2}\exp\left(-\sum_{j=i+1}^{t}\eta_{j}\right)$
$\displaystyle=\sum_{i=1}^{\tau-1}\eta_{i}^{2}\exp\left(-\sum_{j=i+1}^{t}\eta_{j}\right)+\sum_{i=\tau}^{t}\eta_{i}^{2}\exp\left(-\sum_{j=i+1}^{t}\eta_{j}\right)$
$\displaystyle\leq\frac{\epsilon}{2}+\frac{\epsilon}{2}=\epsilon.$

∎

###### Proof of Theorem 3.5.

From Lemma 3.1,
$\left\|{\bm{x}}_{t}\right\|\leq\max{\\{\left\|{\bm{x}}_{0}\right\|,\frac{1}{\lambda}\\}}=B$
for $t\geq 0$. Define $C\triangleq\frac{H(1+\lambda
B)^{2}}{2\lambda^{2}}\frac{4}{(t+1)^{2}}$.

We have that for $t=1$,

$\displaystyle L({\bm{x}}_{1})-L({\bm{x}}^{*})$
$\displaystyle\leq(1-1)(L({\bm{x}}_{0})-L({\bm{x}}^{*}))+\frac{H(1+\lambda
B)^{2}}{2\lambda^{2}}$ $\displaystyle=C\leq\frac{4C}{3}.$

Suppose $L({\bm{x}}_{t-1})-L({\bm{x}}^{*})\leq\frac{4C}{t+1}$, we have that

$\displaystyle L({\bm{x}}_{t})-L({\bm{x}}^{*})$
$\displaystyle\leq(1-\frac{2}{t+1})(L({\bm{x}}_{t-1})-L({\bm{x}}^{*}))+\frac{H(1+\lambda
B)^{2}}{2\lambda^{2}}\frac{4}{(t+1)^{2}}$
$\displaystyle\leq\frac{t-1}{t+1}\frac{4C}{t+1}+\frac{4C}{(t+1)^{2}}$
$\displaystyle=\frac{4Ct}{(t+1)^{2}}\leq\frac{4C}{t+2}.$

∎

###  A.3 Omitted Proofs for Lemma 3.9

###### Proof of Lemma 3.9.

  1. 1.

For any $\epsilon>0$, there exists $t^{\prime}$ such that
$\left\|{\bm{x}}_{t}-{\bm{x}}_{\infty}\right\|\leq\frac{\epsilon}{2\lambda}$
for any $t>t^{\prime}$. Because
$\eta_{t}{\bm{\Delta}}_{t}={\bm{x}}_{t-1}-{\bm{x}}_{t}-\lambda\eta_{t}{\bm{x}}_{t-1}$,
we have that

$\displaystyle\frac{\sum_{t=1}^{T}\eta_{t}{\bm{\Delta}}_{t}}{\sum_{t=1}^{T}\eta_{t}}$
$\displaystyle=\frac{{\bm{x}}_{0}-{\bm{x}}_{T}-\lambda\sum_{t=1}^{T}\eta_{t}{\bm{x}}_{t-1}}{\sum_{t=1}^{T}\eta_{t}}$
$\displaystyle=\frac{{\bm{x}}_{0}-{\bm{x}}_{T}-\lambda(\sum_{t=1}^{t^{\prime}}\eta_{t}{\bm{x}}_{t-1}-\sum_{t=1}^{t^{\prime}}\eta_{t}{\bm{x}}_{\infty})}{\sum_{t=1}^{T}\eta_{t}}-\lambda\frac{\sum_{t=1}^{t^{\prime}}\eta_{t}{\bm{x}}_{\infty}+\sum_{t=t^{\prime}+1}^{T}\eta_{t}{\bm{x}}_{t-1}}{\sum_{t=1}^{T}\eta_{t}}.$

There exists $T^{\prime}\geq t^{\prime}$ such that
$\sum_{t=1}^{T}\eta_{t}\geq\frac{2}{\epsilon}\left(\left\|{\bm{x}}_{0}-{\bm{x}}_{\infty}-\lambda\left(\sum_{t=1}^{t^{\prime}}\eta_{t}{\bm{x}}_{t-1}-\sum_{t=1}^{t^{\prime}}\eta_{t}{\bm{x}}_{\infty}\right)\right\|+\frac{\epsilon}{2}\right)$
for $T\geq T^{\prime}$. Then we have

$\displaystyle\left\|\frac{\sum_{t=1}^{T}\eta_{t}{\bm{\Delta}}_{t}}{\sum_{t=1}^{T}\eta_{t}}+\lambda{\bm{x}}_{\infty}\right\|$
$\displaystyle\leq\frac{\left\|{\bm{x}}_{0}-{\bm{x}}_{T}-\lambda\left(\sum_{t=1}^{t^{\prime}}\eta_{t}{\bm{x}}_{t-1}-\sum_{t=1}^{t^{\prime}}\eta_{t}{\bm{x}}_{\infty}\right)\right\|}{\sum_{t=1}^{T}\eta_{t}}+\lambda\frac{\sum_{t=t^{\prime}+1}^{T}\eta_{t}\left\|{\bm{x}}_{t-1}-{\bm{x}}_{\infty}\right\|}{\sum_{t=1}^{T}\eta_{t}}$
$\displaystyle\leq\frac{\left\|{\bm{x}}_{0}-{\bm{x}}_{\infty}-\lambda\left(\sum_{t=1}^{t^{\prime}}\eta_{t}{\bm{x}}_{t-1}-\sum_{t=1}^{t^{\prime}}\eta_{t}{\bm{x}}_{\infty}\right)\right\|+\left\|{\bm{x}}_{T}-{\bm{x}}_{\infty}\right\|}{\sum_{t=1}^{T}\eta_{t}}+\lambda\frac{\epsilon}{2\lambda}$
$\displaystyle\leq\frac{\epsilon}{2}+\frac{\epsilon}{2}=\epsilon.$

So
${\bm{\Delta}}_{\infty}:=\frac{\sum_{t=1}^{T}\eta_{t}{\bm{\Delta}}_{t}}{\sum_{t=1}^{T}\eta_{t}}$
exists and ${\bm{\Delta}}_{\infty}=-\lambda{\bm{x}}_{\infty}$.

  2. 2.

Because $\nabla L({\bm{x}})$ is a continuous function and
$\lim_{t\rightarrow\infty}{\bm{x}}_{t}={\bm{x}}_{\infty}$,
$\lim_{t\rightarrow\infty}\nabla L({\bm{x}}_{t})=\nabla L({\bm{x}}_{\infty})$.
For any $\epsilon>0$, there exists $T_{1}$ such that

$\left|\left\|\nabla L({\bm{x}}_{t})\right\|_{*}-\left\|\nabla
L({\bm{x}}_{\infty})\right\|_{*}\right|\leq\left\|\nabla
L({\bm{x}}_{t})-\nabla L({\bm{x}}_{\infty})\right\|_{*}\leq\frac{\epsilon}{3}$

for any $t\geq T_{1}$. It also holds that

$|\left\langle\nabla L({\bm{x}}_{t})-\nabla
L({\bm{x}}_{\infty}),{\bm{\Delta}}_{t}\right\rangle|\leq\left\|\nabla
L({\bm{x}}_{t})-\nabla
L({\bm{x}}_{\infty})\right\|_{*}\left\|{\bm{\Delta}}_{t}\right\|\leq\frac{\epsilon}{3}$

because $\left\|{\bm{\Delta}}_{t}\right\|\leq 1$. Because
$\sum_{t=1}^{\infty}\eta_{t}=\infty$, there exists $T_{2}\geq T_{1}$ such that

$\sum_{t=1}^{T_{2}}\eta_{t}\geq\frac{3}{\epsilon}\left|\sum_{t=1}^{T_{1}}\eta_{t}\left(\left\|\nabla
L({\bm{x}}_{t})\right\|_{*}-\left\|\nabla
L({\bm{x}}_{\infty})\right\|_{*}+\left\langle\nabla
L({\bm{x}}_{\infty})-\nabla
L({\bm{x}}_{t}),{\bm{\Delta}}_{t}\right\rangle\right)\right|.$

Then for any $T\geq T_{2}$, we have that

$\displaystyle\left|\frac{\sum_{t=1}^{T}\eta_{t}\left\langle\nabla
L({\bm{x}}_{\infty}),{\bm{\Delta}}_{t}\right\rangle}{\sum_{t=1}^{T}\eta_{t}}-\left\|\nabla
L({\bm{x}}_{\infty})\right\|_{*}\right|$ $\displaystyle=$
$\displaystyle\left|\frac{\sum_{t=1}^{T}\eta_{t}\left\langle\nabla
L({\bm{x}}_{t}),{\bm{\Delta}}_{t}\right\rangle}{\sum_{t=1}^{T}\eta_{t}}+\frac{\sum_{t=1}^{T}\eta_{t}\left\langle\nabla
L({\bm{x}}_{\infty})-\nabla
L({\bm{x}}_{t}),{\bm{\Delta}}_{t}\right\rangle}{\sum_{t=1}^{T}\eta_{t}}-\left\|\nabla
L({\bm{x}}_{\infty})\right\|_{*}\right|$ $\displaystyle=$
$\displaystyle\left|\frac{\sum_{t=1}^{T}\eta_{t}\left\|\nabla
L({\bm{x}}_{t})\right\|_{*}}{\sum_{t=1}^{T}\eta_{t}}+\frac{\sum_{t=1}^{T}\eta_{t}\left\langle\nabla
L({\bm{x}}_{\infty})-\nabla
L({\bm{x}}_{t}),{\bm{\Delta}}_{t}\right\rangle}{\sum_{t=1}^{T}\eta_{t}}-\left\|\nabla
L({\bm{x}}_{\infty})\right\|_{*}\right|$ $\displaystyle=$
$\displaystyle\left|\frac{\sum_{t=1}^{T}\eta_{t}(\left\|\nabla
L({\bm{x}}_{t})\right\|_{*}-\left\|\nabla
L({\bm{x}}_{\infty})\right\|_{*})}{\sum_{t=1}^{T}\eta_{t}}+\frac{\sum_{t=1}^{T}\eta_{t}\left\langle\nabla
L({\bm{x}}_{\infty})-\nabla
L({\bm{x}}_{t}),{\bm{\Delta}}_{t}\right\rangle}{\sum_{t=1}^{T}\eta_{t}}\right|$
$\displaystyle\leq$
$\displaystyle\frac{\left|\sum_{t=1}^{T_{1}}\eta_{t}\left(\left\|\nabla
L({\bm{x}}_{t})\right\|_{*}-\left\|\nabla
L({\bm{x}}_{\infty})\right\|_{*}+\left\langle\nabla
L({\bm{x}}_{\infty})-\nabla
L({\bm{x}}_{t}),{\bm{\Delta}}_{t}\right\rangle\right)\right|}{\sum_{t=1}^{T}\eta_{t}}$
$\displaystyle+$
$\displaystyle\frac{\sum_{t=T_{1}+1}^{T}\eta_{t}\left|\left\|\nabla
L({\bm{x}}_{t})\right\|_{*}-\left\|\nabla
L({\bm{x}}_{\infty})\right\|_{*}\right|}{\sum_{t=1}^{T}\eta_{t}}+\frac{\sum_{t=T_{1}+1}^{T}\eta_{t}\left|\left\langle\nabla
L({\bm{x}}_{\infty})-\nabla
L({\bm{x}}_{t}),{\bm{\Delta}}_{t}\right\rangle\right|}{\sum_{t=1}^{T}\eta_{t}}$
$\displaystyle\leq$
$\displaystyle\frac{\epsilon}{3}+\frac{\epsilon}{3}\frac{\sum_{t=T_{1}+1}^{T}\eta_{t}}{\sum_{t=1}^{T}\eta_{t}}+\frac{\epsilon}{3}\frac{\sum_{t=T_{1}+1}^{T}\eta_{t}}{\sum_{t=1}^{T}\eta_{t}}$
$\displaystyle\leq$ $\displaystyle\epsilon.$

Therefore, we prove that $\left\|\nabla
L({\bm{x}}_{\infty})\right\|_{*}=\lim_{T\to\infty}\frac{\sum_{t=1}^{T}\eta_{t}\left\langle\nabla
L({\bm{x}}_{\infty}),{\bm{\Delta}}_{t}\right\rangle}{\sum_{t=1}^{T}\eta_{t}}$.
On the other hand, we have that

$\displaystyle\left\langle\nabla
L({\bm{x}}_{\infty}),{\bm{\Delta}}_{\infty}\right\rangle$
$\displaystyle=\left\langle\nabla
L({\bm{x}}_{\infty}),\lim_{T\to\infty}\frac{\sum_{t=1}^{T}\eta_{t}{\bm{\Delta}}_{t}}{\sum_{t=1}^{T}\eta_{t}}\right\rangle$
$\displaystyle=\lim_{T\to\infty}\left\langle\nabla
L({\bm{x}}_{\infty}),\frac{\sum_{t=1}^{T}\eta_{t}{\bm{\Delta}}_{t}}{\sum_{t=1}^{T}\eta_{t}}\right\rangle$
$\displaystyle=\lim_{T\to\infty}\frac{\sum_{t=1}^{T}\eta_{t}\left\langle\nabla
L({\bm{x}}_{\infty}),{\bm{\Delta}}_{t}\right\rangle}{\sum_{t=1}^{T}\eta_{t}},$

which finishes the proof.

  3. 3.

For any $T$, we know
$\left\|\frac{\sum_{t=1}^{T}\eta_{t}{\bm{\Delta}}_{t}}{\sum_{t=1}^{T}\eta_{t}}\right\|\leq\frac{\sum_{t=1}^{T}\eta_{t}\left\|{\bm{\Delta}}_{t}\right\|}{\sum_{t=1}^{T}\eta_{t}}\leq\frac{\sum_{t=1}^{T}\eta_{t}}{\sum_{t=1}^{T}\eta_{t}}=1$.
By the continuity of $\left\|\cdot\right\|$,
$\left\|\Delta_{\infty}\right\|=\lim_{T\rightarrow\infty}\left\|\frac{\sum_{t=1}^{T}\eta_{t}{\bm{\Delta}}_{t}}{\sum_{t=1}^{T}\eta_{t}}\right\|\leq
1$.

∎

##  Appendix B Omitted Proofs in Section 4

###  B.1 Omitted proofs for upper bound for average update size of
$\mathtt{Adam}$

###### Proof of Lemma 4.2.

We first represent $m_{t}$ and $v_{t}$ as a weighted sum of $g_{t}$ and
$g_{t}^{2}$.

$\displaystyle m_{t}$
$\displaystyle=\beta_{1}m_{t-1}+(1-\beta_{1})g_{t}=\beta_{1}^{t}m_{0}+(1-\beta_{1})\sum_{i=0}^{t-1}\beta_{1}^{i}g_{t-i},$
$\displaystyle v_{t}$
$\displaystyle\geq\beta_{2}v_{t-1}+(1-\beta_{2})g_{t}^{2}\geq\beta_{2}^{t}v_{0}+(1-\beta_{2})\sum_{i=0}^{t-1}\beta_{2}^{i}g_{t-i}^{2}.$

By Cauchy–Schwarz inequality, we have that

$\displaystyle\left|\sum_{t=1}^{T}\eta_{t}\Delta_{t}\right|$
$\displaystyle=\left|\sum_{t=1}^{T}\eta_{t}\frac{\beta_{1}^{t}m_{0}+(1-\beta_{1})\sum_{i=0}^{t-1}\beta_{1}^{i}g_{t-i}}{\sqrt{v_{t}}}\right|$
$\displaystyle\leq\left[\sum_{t=1}^{T}\eta_{t}\left(\frac{\beta_{1}^{t}m_{0}^{2}}{v_{t}}+\sum_{i=0}^{t-1}(1-\beta_{1})\beta_{1}^{i}\frac{g_{t-i}^{2}}{v_{t}}\right)\right]^{\frac{1}{2}}\left[\sum_{t=1}^{T}\eta_{t}\left(\beta_{1}^{t}+\sum_{i=0}^{t-1}(1-\beta_{1})\beta_{1}^{i}\right)\right]^{\frac{1}{2}}$
$\displaystyle\leq\left(\sum_{t=1}^{T}\eta_{t}\frac{\beta_{1}^{t}v_{0}}{v_{t}}+\sum_{t=1}^{T}\eta_{t}\sum_{i=0}^{t-1}(1-\beta_{1})\beta_{1}^{i}\frac{v_{t-i}-\beta_{2}v_{t-i-1}}{(1-\beta_{2})v_{t}}\right)^{\frac{1}{2}}\left(\sum_{t=1}^{T}\eta_{t}\right)^{\frac{1}{2}}.$

We further analyze the first time in RHS and have that

$\displaystyle\sum_{t=1}^{T}\eta_{t}\frac{\beta_{1}^{t}v_{0}}{v_{t}}+\sum_{t=1}^{T}\eta_{t}\sum_{i=0}^{t-1}(1-\beta_{1})\beta_{1}^{i}\frac{v_{t-i}-\beta_{2}v_{t-i-1}}{(1-\beta_{2})v_{t}}$
$\displaystyle=$
$\displaystyle\sum_{t=1}^{T}\eta_{t}\frac{\beta_{1}^{t}v_{0}}{v_{t}}+\sum_{t=1}^{T}\eta_{t}\frac{1-\beta_{1}}{(1-\beta_{2})v_{t}}\left(v_{t}-\beta_{1}^{t-1}\beta_{2}v_{0}+\sum_{i=1}^{t-1}(\beta_{1}^{i}-\beta_{2}\beta_{1}^{i-1})v_{t-i}\right)$
$\displaystyle=$
$\displaystyle\sum_{t=1}^{T}\eta_{t}+\frac{\beta_{2}-\beta_{1}}{1-\beta_{2}}\sum_{t=1}^{T}\eta_{t}\left(1-\beta_{1}^{t-1}\frac{v_{0}}{v_{t}}-(1-\beta_{1})\sum_{i=1}^{t-1}\beta_{1}^{i-1}\frac{v_{t-i}}{v_{t}}\right)$
$\displaystyle\leq$
$\displaystyle\sum_{t=1}^{T}\eta_{t}+\frac{\beta_{2}-\beta_{1}}{1-\beta_{2}}\sum_{t=1}^{T}\eta_{t}\left(\beta_{1}^{t-1}+(1-\beta_{1})\sum_{i=1}^{t-1}\beta_{1}^{i-1}(1-\frac{v_{t-i}}{v_{t}})\right)$
$\displaystyle\leq$
$\displaystyle\sum_{t=1}^{T}\eta_{t}+\frac{\beta_{2}-\beta_{1}}{1-\beta_{2}}\sum_{t=1}^{T}\eta_{t}\left(\beta_{1}^{t-1}+(1-\beta_{1})\sum_{i=1}^{t-1}\beta_{1}^{i-1}\ln{(\frac{v_{t}}{v_{t-i}})}\right)$
$\displaystyle=$
$\displaystyle\sum_{t=1}^{T}\eta_{t}+\frac{\beta_{2}-\beta_{1}}{1-\beta_{2}}\sum_{t=1}^{T}\eta_{t}\beta_{1}^{t-1}+\frac{(\beta_{2}-\beta_{1})(1-\beta_{1})}{1-\beta_{2}}\sum_{t=1}^{T}\left(\eta_{t}\frac{1-\beta_{1}^{t-1}}{1-\beta_{1}}-\sum_{i=1}^{T-t}\eta_{t+i}\beta_{1}^{i-1}\right)\ln{v_{t}}$
$\displaystyle=$
$\displaystyle\sum_{t=1}^{T}\eta_{t}+\frac{\beta_{2}-\beta_{1}}{1-\beta_{2}}\sum_{t=1}^{T}\eta_{t}\beta_{1}^{t-1}+\frac{(\beta_{2}-\beta_{1})(1-\beta_{1})}{1-\beta_{2}}\sum_{t=2}^{T}\left(\eta_{t}\frac{1-\beta_{1}^{t-1}}{1-\beta_{1}}-\sum_{i=1}^{T-t}\eta_{t+i}\beta_{1}^{i-1}\right)\ln{\frac{v_{t}}{v_{1}}}.$

When $\beta_{1}=\beta_{2}$, we have that

$\displaystyle\left|\Delta_{t}\right|$
$\displaystyle=\left|\frac{\beta_{1}^{t}m_{0}+(1-\beta_{1})\sum_{i=0}^{t-1}\beta_{1}^{i}g_{t-i}}{\sqrt{v_{t}}}\right|$
$\displaystyle\leq\left(\frac{\beta_{1}^{t}m_{0}^{2}+(1-\beta_{1})\sum_{i=0}^{t-1}\beta_{1}^{i}g_{t-i}^{2}}{v_{t}}\right)^{\frac{1}{2}}\left(\beta_{1}^{t}+\sum_{i=0}^{t-1}(1-\beta_{1})\beta_{1}^{i}\right)^{\frac{1}{2}}$
$\displaystyle\leq\left(\frac{\beta_{1}^{t}v_{0}+(1-\beta_{1})\sum_{i=0}^{t-1}\beta_{1}^{i}g_{t-i}^{2}}{v_{t}}\right)^{\frac{1}{2}}$
$\displaystyle=1.$

∎

###  B.2 Proof for Lemma 4.1

###### Proof of Lemma 4.1.

  1. 1.

The proof for this part is the same as the proof of Lemma 3.9 in Section A.3.

  2. 2.

If $\nabla L({\bm{x}}_{\infty})={\bm{\mathbf{0}}}$, $\left\langle\nabla
L({\bm{x}}_{\infty}),{\bm{\Delta}}_{\infty}\right\rangle=0=\left\|\nabla
L({\bm{x}}_{\infty})\right\|_{1}$.

If $\nabla L({\bm{x}}_{\infty})\neq{\bm{\mathbf{0}}}$, we consider each
coordinate $j$ such that $\nabla L({\bm{x}}_{\infty})_{j}\neq 0$. Since we
have that $\lim_{t\rightarrow\infty}\nabla L({\bm{x}}_{t})_{j}=\nabla
L({\bm{x}}_{\infty})_{j}$, we can get the convergence for ${\bm{m}}_{t,j}$ and
${\bm{v}}_{t,j}$.

$\displaystyle{\bm{m}}_{t,j}$
$\displaystyle=(1-\beta_{1})\sum_{i=0}^{t}\beta_{1}^{i}{\bm{g}}_{t-i,j}\rightarrow\nabla
L({\bm{x}}_{\infty})_{j},$ $\displaystyle{\bm{v}}_{t,j}$
$\displaystyle=(1-\beta_{2})\sum_{i=0}^{t}\beta_{2}^{i}{\bm{g}}_{t-i,j}^{2}\rightarrow\nabla
L({\bm{x}}_{\infty})_{j}^{2}\neq 0.$

Then we have that
$\lim_{t\rightarrow\infty}{\bm{\Delta}}_{t,j}=\lim_{t\rightarrow\infty}\frac{{\bm{m}}_{t,j}}{\sqrt{{\bm{v}}_{t,j}}}=\text{sign}(\nabla
L({\bm{x}}_{\infty})_{j})$. For any $\epsilon>0$, there exists $t^{\prime}$
such that $\left\|{\bm{\Delta}}_{t,j}-\text{sign}\left(\nabla
L({\bm{x}}_{\infty})_{j}\right)\right\|\leq\frac{\epsilon}{2}$ for $t\geq
t^{\prime}$. And there exists $T^{\prime}\geq t^{\prime}$ such that
$\sum_{t=1}^{T}\eta_{t}\geq\frac{2}{\epsilon}\sum_{t=1}^{t^{\prime}}\eta_{t}\left({\bm{\Delta}}_{t,j}-\text{sign}\left(\nabla
L({\bm{x}}_{\infty})_{j}\right)\right)$ for any $T\geq T^{\prime}$. Then for
any $T\geq T^{\prime}$, we have that

$\displaystyle\left\|\frac{\sum_{t=1}^{T}\eta_{t}{\bm{\Delta}}_{t,j}}{\sum_{t=1}^{T}\eta_{t}}-\text{sign}(\nabla
L({\bm{x}}_{\infty})_{j})\right\|$ $\displaystyle\leq$
$\displaystyle\left\|\frac{\sum_{t=1}^{t^{\prime}}\eta_{t}\left({\bm{\Delta}}_{t,j}-\text{sign}(\nabla
L({\bm{x}}_{\infty})_{j})\right)}{\sum_{t=1}^{T}\eta_{t}}\right\|+\frac{\sum_{t=t^{\prime}+1}^{T}\eta_{t}\left\|{\bm{\Delta}}_{t,j}-\text{sign}(\nabla
L({\bm{x}}_{\infty})_{j})\right\|}{\sum_{t=1}^{T}\eta_{t}}$
$\displaystyle\leq$
$\displaystyle\frac{\epsilon}{2}+\frac{\epsilon}{2}=\epsilon.$

So
${\bm{\Delta}}_{\infty,j}=\lim_{T\rightarrow\infty}\frac{\sum_{t=1}^{T}\eta_{t}{\bm{\Delta}}_{t,j}}{\sum_{t=1}^{T}\eta_{t}}=\text{sign}(\nabla
L({\bm{x}}_{\infty})_{j})$ for $\nabla L({\bm{x}}_{\infty})_{j}\neq 0$. Then
we have that

$\displaystyle\left\langle\nabla
L({\bm{x}}_{\infty}),{\bm{\Delta}}_{\infty}\right\rangle$
$\displaystyle=\sum_{\nabla L({\bm{x}}_{\infty})_{j}\neq 0}\nabla
L({\bm{x}}_{\infty})_{j}{\bm{\Delta}}_{\infty,j}$ $\displaystyle=\sum_{\nabla
L({\bm{x}}_{\infty})_{j}\neq 0}|\nabla L({\bm{x}}_{\infty})_{j}|$
$\displaystyle=\left\|\nabla L({\bm{x}}_{\infty})\right\|_{1}.$

  3. 3.

For nonzero coordinate $j$ of $\nabla L({\bm{x}}_{\infty})$, from above we
have $|{\bm{\Delta}}_{\infty,j}|=|\text{sign}(\nabla
L({\bm{x}}_{\infty})_{j})|=1$.

For $j$ such that $\nabla L({\bm{x}}_{\infty})_{j}=0$, we know
$\lim_{t\rightarrow\infty}{\bm{g}}_{t,j}=\lim_{t\rightarrow\infty}{\bm{m}}_{t,j}=\lim_{t\rightarrow\infty}{\bm{v}}_{t,j}=0$.
We employ the upper bound for average update in Lemma 4.2 since
$\\{{\bm{g}}_{t,j}\\}_{t=1}^{\infty}$ and
$\\{{\bm{v}}_{t,j}\\}_{t=0}^{\infty}$ in Algorithm 1 satisfy the condition
that
${\bm{v}}_{t,j}-\beta_{2}{\bm{v}}_{t-1,j}\geq(1-\beta_{2}){\bm{g}}_{t,j}^{2}$
and ${\bm{m}}_{0,j}=0\leq\sqrt{{\bm{v}}_{0,j}}$. By Lemma 4.2 we have

$\displaystyle\left|\frac{\sum_{t=1}^{T}\eta_{t}{\bm{\Delta}}_{t,j}}{\sum_{t=1}^{T}\eta_{t}}\right|$
$\displaystyle\leq$
$\displaystyle\left(\frac{\sum_{t=1}^{T}\eta_{t}+\frac{\beta_{2}-\beta_{1}}{1-\beta_{2}}\sum_{t=1}^{T}\eta_{t}\beta_{1}^{t-1}+\frac{(\beta_{2}-\beta_{1})(1-\beta_{1})}{1-\beta_{2}}\sum_{t=2}^{T}\left(\eta_{t}\frac{1-\beta_{1}^{t-1}}{1-\beta_{1}}-\sum_{i=1}^{T-t}\eta_{t+i}\beta_{1}^{i-1}\right)\ln{\frac{{\bm{v}}_{t,j}}{{\bm{v}}_{1,j}}}}{\sum_{t=1}^{T}\eta_{t}}\right)^{\frac{1}{2}}.$

The denominator goes to $\infty$ when $T\rightarrow\infty$. So it suffices to
bound the last two terms in the numerator by constants in order to show
$\left\|{\bm{\Delta}}_{\infty}\right\|\leq 1$. Because $\eta_{t}$ is non-
increasing in $t$, it holds that

$\displaystyle\sum_{t=1}^{T}\eta_{t}\beta_{1}^{t}\leq\sum_{t=1}^{T}\eta_{1}\beta_{1}^{t}\leq\frac{\eta_{1}\beta_{1}}{1-\beta_{1}}.$

For the last term, we first analyze the coefficient between each
$\ln{{\bm{v}}_{t,j}}$. Define
$\alpha_{t}=\eta_{t}\frac{1-\beta_{1}^{t-1}}{1-\beta_{1}}-\sum_{i=1}^{T-t}\eta_{t+i}\beta_{1}^{i-1}$.
We claim that
$\left|\alpha_{t}\right|\leq\max{\\{\frac{\beta_{1}^{t-1}}{1-\beta_{1}}\eta_{t+1},\frac{\eta_{t}}{1-\beta_{1}}\\}}=\frac{\eta_{t}}{1-\beta_{1}}$.
This is because

$\displaystyle\alpha_{t}$
$\displaystyle\leq\eta_{t}\frac{1-\beta_{1}^{t-1}}{1-\beta_{1}}\leq\frac{\eta_{t}}{1-\beta_{1}},$

and again by monotonicity of learning rates $\eta_{t}$, we have that

$\displaystyle\alpha_{t}\geq\eta_{t+1}\frac{1-\beta_{1}^{t-1}}{1-\beta_{1}}-\sum_{i=1}^{T-t}\eta_{t+1}\beta_{1}^{i-1}\geq\eta_{t+1}\left(\frac{1-\beta_{1}^{t-1}}{1-\beta_{1}}-\sum_{i=1}^{\infty}\beta_{1}^{i-1}\right)=-\frac{\beta_{1}^{t-1}}{1-\beta_{1}}\eta_{t+1}.$

We can also have
$\ln{\frac{{\bm{v}}_{t,j}}{{\bm{v}}_{1,j}}}\geq(t-1)\ln{\beta_{2}}$ because

$\ln{{\bm{v}}_{t,j}}-\ln{{\bm{v}}_{t-1,j}}=\ln{\frac{{\bm{v}}_{t,j}}{{\bm{v}}_{t-1,j}}}=\ln{\frac{\beta_{2}{\bm{v}}_{t-1,j}+(1-\beta_{2}){\bm{g}}_{t,j}^{2}}{{\bm{v}}_{t-1,j}}}\geq\ln{\beta_{2}}.$

And there exists $t^{\prime}$ such that
$\ln{\frac{{\bm{v}}_{t,j}}{{\bm{v}}_{1,j}}}\leq 0$ for any $t\geq t^{\prime}$
because $\lim_{t\rightarrow\infty}{\bm{v}}_{t,j}=0$. Then

$\displaystyle\sum_{t=2}^{T}\alpha_{t}\ln{\frac{{\bm{v}}_{t,j}}{{\bm{v}}_{1,j}}}$
$\displaystyle\leq\sum_{t=2}^{t^{\prime}}\left|\alpha_{t}\right|\left|\ln{\frac{{\bm{v}}_{t,j}}{{\bm{v}}_{1,j}}}\right|+\sum_{t=t^{\prime}+1}^{T}\alpha_{t}\ln{\frac{{\bm{v}}_{t,j}}{{\bm{v}}_{1,j}}}$
$\displaystyle\leq\sum_{t=2}^{t^{\prime}}\frac{\eta_{t}}{1-\beta_{1}}\left|\ln{\frac{{\bm{v}}_{t,j}}{{\bm{v}}_{1,j}}}\right|+\sum_{\alpha_{t}\leq
0,\ t^{\prime}<t\leq T}\alpha_{t}\ln{\frac{{\bm{v}}_{t,j}}{{\bm{v}}_{1,j}}}$
$\displaystyle\leq\frac{\sum_{t=2}^{t^{\prime}}\eta_{t}\left|\ln{\frac{{\bm{v}}_{t,j}}{{\bm{v}}_{1,j}}}\right|}{1-\beta_{1}}+\sum_{\alpha_{t}\leq
0,\ t^{\prime}<t\leq
T}\left(-\frac{\beta_{1}^{t}}{1-\beta_{1}}\eta_{t+1}\right)(t-1)\ln{\beta_{2}}$
$\displaystyle\leq\frac{\sum_{t=2}^{t^{\prime}}\eta_{t}\left|\ln{\frac{{\bm{v}}_{t,j}}{{\bm{v}}_{1,j}}}\right|}{1-\beta_{1}}+(-\ln{\beta_{2}})\sum_{t=1}^{T}\frac{(t-1)\beta_{1}^{t}}{1-\beta_{1}}\eta_{t+1}$
$\displaystyle\leq\frac{\sum_{t=2}^{t^{\prime}}\eta_{t}\left|\ln{\frac{{\bm{v}}_{t,j}}{{\bm{v}}_{1,j}}}\right|}{1-\beta_{1}}+(-\ln{\beta_{2}})\eta_{1}\sum_{t=1}^{T}\frac{(t-1)\beta_{1}^{t}}{1-\beta_{1}}$
$\displaystyle\leq\frac{\sum_{t=2}^{t^{\prime}}\eta_{t}\left|\ln{\frac{{\bm{v}}_{t,j}}{{\bm{v}}_{1,j}}}\right|}{1-\beta_{1}}-\frac{\eta_{1}\beta_{1}^{2}\ln{\beta_{2}}}{(1-\beta_{1})^{3}}.$

Define
$C:=\frac{(\beta_{2}-\beta_{1})\eta_{1}\beta_{1}}{(1-\beta_{2})(1-\beta_{1})}+\frac{\beta_{2}-\beta_{1}}{1-\beta_{2}}\left(\sum_{t=2}^{t^{\prime}}\eta_{t}\left|\ln{{\bm{v}}_{t,j}}\right|-\frac{\eta_{1}\beta_{1}^{2}\ln{\beta_{2}}}{(1-\beta_{1})^{2}}\right)$,
we now have

$\left|\frac{\sum_{t=1}^{T}\eta_{t}{\bm{\Delta}}_{t,j}}{\sum_{t=1}^{T}\eta_{t}}\right|\leq\left(\frac{\sum_{t=1}^{T}\eta_{t}+C}{\sum_{t=1}^{T}\eta_{t}}\right)^{\frac{1}{2}}.$

Then
$\left|{\bm{\Delta}}_{\infty,j}\right|=\left|\lim\limits_{T\to\infty}\frac{\sum_{t=1}^{T}\eta_{t}{\bm{\Delta}}_{t,j}}{\sum_{t=1}^{T}\eta_{t}}\right|\leq\lim\limits_{T\to\infty}\left|\frac{\sum_{t=1}^{T}\eta_{t}{\bm{\Delta}}_{t,j}}{\sum_{t=1}^{T}\eta_{t}}\right|\leq
1$ because $\sum_{t=1}^{\infty}\eta_{T}=\infty$. This completes the proof.

∎

###  B.3 A counter example when $\beta_{1}>\beta_{2}$

For any $\lambda$ and $\beta_{1}>\beta_{2}$, we provide the following example
for which the iterates of $\mathtt{AdamW}$ will converge and the
$\ell_{\infty}$ norm of the converged solution is larger than
$\frac{1}{\lambda}$.

For some sufficiently small $\eta$, denote
$\tilde{x}=-\frac{1}{\lambda}\frac{\frac{1-\beta_{1}}{1-\lambda\eta-\beta_{1}}}{\sqrt{\frac{1-\beta_{2}}{(1-\lambda\eta)^{2}-\beta_{2}}}}$.
$L(x)$ is defined as $\frac{1}{2}(x-\tilde{x})^{2}$. For any starting point
$x_{0}>\tilde{x}$, $m_{0}$ is set as
$\frac{1-\beta_{1}}{1-\lambda\eta-\beta_{1}}g_{1}$ and $v_{0}$ is set as
$\frac{1-\beta_{2}}{(1-\lambda\eta)^{2}-\beta_{2}}g_{1}^{2}$ with
$g_{1}=\nabla L(x_{0})=x_{0}-\tilde{x}$. We show by induction that
$\frac{m_{t}}{\sqrt{v_{t}}}=-\lambda\tilde{x}$ and
$x_{t}-\tilde{x}=(1-\lambda\eta)(x_{t-1}-\tilde{x})$ for any $t\geq 1$.

When $t=1$, we have that

$\displaystyle m_{1}$
$\displaystyle=\beta_{1}m_{0}+(1-\beta_{1})g_{1}=\frac{(1-\beta_{1})(1-\lambda\eta)}{1-\lambda\eta-\beta_{1}}g_{1},$
$\displaystyle v_{1}$
$\displaystyle=\beta_{2}v_{0}+(1-\beta_{2})g_{1}^{2}=\frac{(1-\beta_{2})(1-\lambda\eta)^{2}}{(1-\lambda\eta)^{2}-\beta_{2}}g_{1}^{2}.$

Then
$\frac{m_{1}}{\sqrt{v_{1}}}=\text{sign}(g_{1})\frac{\frac{1-\beta_{1}}{1-\lambda\eta-\beta_{1}}}{\sqrt{\frac{1-\beta_{2}}{(1-\lambda\eta)^{2}-\beta_{2}}}}=\frac{\frac{1-\beta_{1}}{1-\lambda\eta-\beta_{1}}}{\sqrt{\frac{1-\beta_{2}}{(1-\lambda\eta)^{2}-\beta_{2}}}}=-\lambda\tilde{x}$,
which proves the first claim. For the second claim, we have that

$x_{1}-\tilde{x}=x_{0}-\eta\frac{m_{1}}{\sqrt{v_{1}}}-\lambda\eta
x_{0}-\tilde{x}=x_{0}+\lambda\eta\tilde{x}-\lambda\eta
x_{0}-\tilde{x}=(1-\lambda\eta)(x_{0}-\tilde{x}).$

Suppose the claims hold for any $0\leq i<t$. Then $g_{i+1}=\nabla
L(x_{i})=x_{i}-\tilde{x}=(1-\lambda\eta)^{i}g_{1}$. We have that

$\displaystyle m_{t}$
$\displaystyle=\beta_{1}^{t}m_{0}+(1-\beta_{1})\sum_{i=1}^{t}\beta_{1}^{t-i}g_{i}$
$\displaystyle=\beta_{1}^{t}\frac{1-\beta_{1}}{1-\lambda\eta-\beta_{1}}g_{1}+(1-\beta_{1})\sum_{i=1}^{t}\beta_{1}^{t-i}(1-\lambda\eta)^{i-1}g_{1}$
$\displaystyle=\left(\beta_{1}^{t}\frac{1-\beta_{1}}{1-\lambda\eta-\beta_{1}}+(1-\beta_{1})\frac{(1-\lambda\eta)^{t}-\beta_{1}^{t}}{1-\lambda\eta-\beta_{1}}\right)g_{1}$
$\displaystyle=\frac{1-\beta_{1}}{1-\lambda\eta-\beta_{1}}(1-\lambda\eta)^{t}g_{1},$

and

$\displaystyle v_{t}$
$\displaystyle=\beta_{2}^{t}v_{0}+(1-\beta_{2})\sum_{i=1}^{t}\beta_{2}^{t-i}g_{i}^{2}$
$\displaystyle=\beta_{2}^{t}\frac{1-\beta_{2}}{(1-\lambda\eta)^{2}-\beta_{2}}g_{1}^{2}+(1-\beta_{2})\sum_{i=1}^{t}\beta_{2}^{t-i}(1-\lambda\eta)^{2(i-1)}g_{1}^{2}$
$\displaystyle=\left(\beta_{2}^{t}\frac{1-\beta_{2}}{(1-\lambda\eta)^{2}-\beta_{2}}+(1-\beta_{2})\frac{(1-\lambda\eta)^{2t}-\beta_{2}^{t}}{(1-\lambda\eta)^{2}-\beta_{2}}\right)g_{1}^{2}$
$\displaystyle=\frac{1-\beta_{2}}{(1-\lambda\eta)^{2}-\beta_{2}}(1-\lambda\eta)^{2t}g_{1}^{2}.$

Then
$\frac{m_{t}}{\sqrt{v_{t}}}=\frac{\frac{1-\beta_{1}}{1-\lambda\eta-\beta_{1}}}{\sqrt{\frac{1-\beta_{2}}{(1-\lambda\eta)^{2}-\beta_{2}}}}=-\lambda\tilde{x}$.
We also have that

$x_{t}-\tilde{x}=x_{t-1}-\eta\frac{m_{t}}{\sqrt{v_{t}}}-\lambda\eta
x_{t-1}-\tilde{x}=x_{0}+\lambda\eta\tilde{x}-\lambda\eta
x_{0}-\tilde{x}=(1-\lambda\eta)(x_{0}-\tilde{x}).$

.

In this regime, $x_{t}$ will converge to $\tilde{x}$ because
$|x_{t}-\tilde{x}|=O((1-\lambda\eta)^{t})$. However, when
$\beta_{1}>\beta_{2}$ and $\lambda\eta$ is very small, $|\tilde{x}|$ can be
larger than $\frac{1}{\lambda}$. For example, when $\beta_{1}=0.99$,
$\beta_{2}=0.9$, $\lambda=0.1$ and $\eta=0.01$,
$|\tilde{x}|=10.999>\frac{1}{\lambda}$.

###  B.4 Proof for upper bound for norm of iterates in $\mathtt{AdamW}$

###### Proof of Lemma 4.3.

For $\mathtt{AdamW}$ with constant learning rate $\eta$ and each coordinate
$j$, ${\bm{x}}_{T,j}$ can be written as weighted average of past update

$\displaystyle{\bm{x}}_{T,j}$
$\displaystyle=(1-\lambda\eta)^{T}{\bm{x}}_{0,j}+\sum_{t=0}^{T-1}\eta(1-\lambda\eta)^{t}\frac{{\bm{m}}_{T-t,j}}{\sqrt{{\bm{v}}_{T-t,j}}}$
$\displaystyle=(1-\lambda\eta)^{T}{\bm{x}}_{0,j}+\sum_{t=1}^{T}\eta(1-\lambda\eta)^{T-t}\frac{{\bm{m}}_{t,j}}{\sqrt{{\bm{v}}_{t,j}}}.$

Define $\eta_{t}=\eta(1-\lambda\eta)^{T-t}$ for $1\leq t\leq T$. We apply
Lemma 4.2 on $\\{{\bm{v}}_{t,j}\\}_{t=1}^{T}$ and
$\\{{\bm{g}}_{t,j}\\}_{t=1}^{T}$ to bound
$\left|\sum_{t=1}^{T}\eta(1-\lambda\eta)^{T-t}\frac{{\bm{m}}_{t,j}}{\sqrt{{\bm{v}}_{t,j}}}\right|$.

We first compute
$\sum_{t=1}^{T}\eta_{t}=\frac{1-(1-\lambda\eta)^{T}}{\lambda}\leq\frac{1}{\lambda}$.
For the second term in Lemma 4.2, we have that

$\displaystyle\frac{\beta_{2}-\beta_{1}}{1-\beta_{2}}\sum_{t=1}^{T}\eta_{t}\beta_{1}^{t-1}$
$\displaystyle=\frac{\beta_{2}-\beta_{1}}{1-\beta_{2}}\sum_{t=1}^{T}\eta(1-\lambda\eta)^{T-t}\beta_{1}^{t-1}$
$\displaystyle=\frac{1}{\lambda}\frac{(\beta_{2}-\beta_{1})\lambda\eta[(1-\lambda\eta)^{T}-\beta_{1}^{T}]}{(1-\beta_{2})(1-\lambda\eta-\beta_{1})}$
$\displaystyle\leq\frac{(\beta_{2}-\beta_{1})\eta}{(1-\beta_{2})\left|1-\lambda\eta-\beta_{1}\right|}\left[\beta_{1}^{T}+(1-\lambda\eta)^{T}\right].$

For the last term, we define
$\alpha_{t}=\eta_{t}\frac{1-\beta_{1}^{t-1}}{1-\beta_{1}}-\sum_{i=1}^{T-t}\eta_{t+i}\beta_{1}^{i-1}$
and we can compute the exact form of $\alpha_{t}$ as following

$\displaystyle\alpha_{t}$
$\displaystyle=\eta(1-\lambda\eta)^{T-t}\frac{1-\beta_{1}^{t-1}}{1-\beta_{1}}-\sum_{i=1}^{T-t}\eta(1-\lambda\eta)^{T-t-i}\beta_{1}^{i-1}$
$\displaystyle=\eta(1-\lambda\eta)^{T-t}\frac{1-\beta_{1}^{t-1}}{1-\beta_{1}}-\frac{\eta[(1-\lambda\eta)^{T-t}-\beta_{1}^{T-t}]}{1-\lambda\eta-\beta_{1}}$
$\displaystyle=\frac{\eta(1-\lambda\eta)^{T-t}(-\lambda\eta)}{(1-\beta_{1})(1-\lambda\eta-\beta_{1})}-\frac{\eta(1-\lambda\eta)^{T-t}\beta_{1}^{t-1}}{1-\beta_{1}}+\frac{\eta\beta_{1}^{T-t}}{1-\lambda\eta-\beta_{1}}.$

Then we can bound the last term by showing that

$\displaystyle\frac{(\beta_{2}-\beta_{1})(1-\beta_{1})}{1-\beta_{2}}\sum_{t=2}^{T}\alpha_{t}\ln{\frac{{\bm{v}}_{t,j}}{{\bm{v}}_{1,j}}}$
$\displaystyle\leq$
$\displaystyle\frac{(\beta_{2}-\beta_{1})(1-\beta_{1})}{1-\beta_{2}}\sum_{t=2}^{T}|\alpha_{t}|\left|\ln{\frac{{\bm{v}}_{t,j}}{{\bm{v}}_{1,j}}}\right|$
$\displaystyle\leq$ $\displaystyle
C\frac{(\beta_{2}-\beta_{1})(1-\beta_{1})}{1-\beta_{2}}\left[\sum_{t=2}^{T}\frac{\lambda\eta^{2}(1-\lambda\eta)^{T-t}}{(1-\beta_{1})|1-\lambda\eta-\beta_{1}|}+\sum_{t=2}^{T}\frac{\eta(1-\lambda\eta)^{T-t}\beta_{1}^{t-1}}{1-\beta_{1}}+\sum_{t=2}^{T}\frac{\eta\beta_{1}^{T-t}}{|1-\lambda\eta-\beta_{1}|}\right]$
$\displaystyle=$ $\displaystyle
C\frac{(\beta_{2}-\beta_{1})(1-\beta_{1})}{1-\beta_{2}}\left[\frac{\eta[1-(1-\lambda\eta)^{T-1}]}{(1-\beta_{1})|1-\lambda\eta-\beta_{1}|}+\frac{\beta_{1}\eta|(1-\lambda\eta)^{T-1}-\beta_{1}^{T-1}|}{(1-\beta_{1})|1-\lambda\eta-\beta_{1}|}+\frac{\eta(1-\beta_{1}^{T-1})}{|1-\lambda\eta-\beta_{1}|(1-\beta_{1})}\right]$
$\displaystyle\leq$ $\displaystyle
C\frac{(\beta_{2}-\beta_{1})(1-\beta_{1})}{1-\beta_{2}}\left[\frac{\eta[1-(1-\lambda\eta)^{T-1}]}{(1-\beta_{1})|1-\lambda\eta-\beta_{1}|}+\frac{\beta_{1}\eta\left((1-\lambda\eta)^{T-1}+\beta_{1}^{T-1}\right)}{(1-\beta_{1})|1-\lambda\eta-\beta_{1}|}+\frac{\eta(1-\beta_{1}^{T-1})}{|1-\lambda\eta-\beta_{1}|(1-\beta_{1})}\right]$
$\displaystyle=$
$\displaystyle\frac{2C\eta(\beta_{2}-\beta_{1})}{(1-\beta_{2})|1-\lambda\eta-\beta_{1}|}+\frac{C\eta(\beta_{2}-\beta_{1})(\beta_{1}-1)}{(1-\beta_{2})|1-\lambda\eta-\beta_{1}|}\left[\beta_{1}^{T-1}+(1-\lambda\eta)^{T-1}\right]$
$\displaystyle\leq$
$\displaystyle\frac{2C\eta(\beta_{2}-\beta_{1})}{(1-\beta_{2})|1-\lambda\eta-\beta_{1}|}.$

Therefore, we have the following bound

$\displaystyle\lambda|{\bm{x}}_{T,j}|$
$\displaystyle\leq\lambda(1-\lambda\eta)^{T}|{\bm{x}}_{0,j}|+\lambda\left|\sum_{t=1}^{T}\eta_{t}\frac{{\bm{m}}_{t,j}}{\sqrt{{\bm{v}}_{t,j}}}\right|$
$\displaystyle\leq\lambda(1-\lambda\eta)^{T}|{\bm{x}}_{0,j}|+\lambda\sum_{t=1}^{T}\eta_{t}\biggl{[}1+\frac{(\beta_{2}-\beta_{1})\eta\left[\beta_{1}^{T}+(1-\lambda\eta)^{T}\right]}{\sum_{t=1}^{T}\eta_{t}(1-\beta_{2})\left|1-\lambda\eta-\beta_{1}\right|}+\frac{2C\eta(\beta_{2}-\beta_{1})}{\sum_{t=1}^{T}\eta_{t}(1-\beta_{2})|1-\lambda\eta-\beta_{1}|}\biggr{]}^{\frac{1}{2}}$
$\displaystyle\leq\lambda(1-\lambda\eta)^{T}|{\bm{x}}_{0,j}|+\left[(\lambda\sum_{t=1}^{T}\eta_{t})^{2}+\lambda^{2}\sum_{t=1}^{T}\eta_{t}\left[\frac{(\beta_{2}-\beta_{1})\eta\left[\beta_{1}^{T}+(1-\lambda\eta)^{T}\right]+2C\eta(\beta_{2}-\beta_{1})}{(1-\beta_{2})\left|1-\lambda\eta-\beta_{1}\right|}\right]\right]^{\frac{1}{2}}$
$\displaystyle\leq\lambda(1-\lambda\eta)^{T}|{\bm{x}}_{0,j}|+\left[1+\lambda\frac{(\beta_{2}-\beta_{1})\eta\left[\beta_{1}^{T}+(1-\lambda\eta)^{T}\right]+2C\eta(\beta_{2}-\beta_{1})}{(1-\beta_{2})\left|1-\lambda\eta-\beta_{1}\right|}\right]^{\frac{1}{2}}$
$\displaystyle\leq\lambda(1-\lambda\eta)^{T}|{\bm{x}}_{0,j}|+\left[1+\frac{\lambda\eta(\beta_{2}-\beta_{1})\left[\beta_{1}^{T}+(1-\lambda\eta)^{T}\right]}{2(1-\beta_{2})|1-\lambda\eta-\beta_{1}|}+C\frac{\lambda\eta(\beta_{2}-\beta_{1})}{(1-\beta_{2})|1-\lambda\eta-\beta_{1}|}\right]$
$\displaystyle\leq
1+\lambda(1-\lambda\eta)^{T}\left\|{\bm{x}}_{0}\right\|_{\infty}+\frac{\lambda\eta(\beta_{2}-\beta_{1})\left[\beta_{1}^{T}+(1-\lambda\eta)^{T}\right]}{2(1-\beta_{2})|1-\lambda\eta-\beta_{1}|}+C\frac{\lambda\eta(\beta_{2}-\beta_{1})}{(1-\beta_{2})|1-\lambda\eta-\beta_{1}|}.$

This completes the proof. ∎

##  Appendix C Experimental Details and More Results

The architecture of the two-layer transformer is the same as in Kunstner et
al. (2022), which is also used as a tutorial example in PyTorch. It consists
of a 200-dimensional embedding layer, $2$ transformer layers and a linear
layer. Each transformer layer consists of a $2$-head self-attention and an MLP
with a hidden dimension $200$. The experiments are run on a single A4000 or
A6000.

As mentioned in Section 5, we present the results for another random seed in
Figure 4. We also plot the results in full range for the two random seeds in
Figure 5 and Figure 6 to show that the $\ell_{\infty}$ norm of parameters for
$\mathtt{Adam}$ keeps increasing.

(a) $\beta_{1}=\beta_{2}=0.99$

(b) $\beta_{1}=\beta_{2}=0.999$

(c) $\beta_{1}=0.9$, $\beta_{2}=0.95$

(d) $\beta_{1}=0.9,\beta_{2}=0.999$

Figure 4: $\ell_{\infty}$ norm of parameters for $\mathtt{Adam}$ and
$\mathtt{AdamW}$ with different $\beta_{1},\beta_{2}$ for seed $1$

(a) $\beta_{1}=\beta_{2}=0.99$

(b) $\beta_{1}=\beta_{2}=0.999$

(c) $\beta_{1}=0.9$, $\beta_{2}=0.95$

(d) $\beta_{1}=0.9,\beta_{2}=0.999$

Figure 5: $\ell_{\infty}$ norm of parameters for $\mathtt{Adam}$ and
$\mathtt{AdamW}$ with different $\beta_{1},\beta_{2}$ for seed $0$. The range
of y-axis is extended to show the full result of $\mathtt{Adam}$.

(a) $\beta_{1}=\beta_{2}=0.99$

(b) $\beta_{1}=\beta_{2}=0.999$

(c) $\beta_{1}=0.9$, $\beta_{2}=0.95$

(d) $\beta_{1}=0.9,\beta_{2}=0.999$

Figure 6: $\ell_{\infty}$ norm of parameters for $\mathtt{Adam}$ and
$\mathtt{AdamW}$ with different $\beta_{1},\beta_{2}$ for seed $1$. The range
of y-axis is extended to show the full result of $\mathtt{Adam}$.