# Regularized Adaptive Momentum Dual Averaging with an Efficient Inexact
Subproblem Solver for Training Structured Neural Network

Zih-Syuan Huang  Ching-pei Lee

###### Abstract

We propose a Regularized Adaptive Momentum Dual Averaging (RAMDA) algorithm
for training structured neural networks. Similar to existing regularized
adaptive methods, the subproblem for computing the update direction of RAMDA
involves a nonsmooth regularizer and a diagonal preconditioner, and therefore
does not possess a closed-form solution in general. We thus also carefully
devise an implementable inexactness condition that retains convergence
guarantees similar to the exact versions, and propose a companion efficient
solver for the subproblems of both RAMDA and existing methods to make them
practically feasible. We leverage the theory of manifold identification in
variational analysis to show that, even in the presence of such inexactness,
the iterates of RAMDA attain the ideal structure induced by the regularizer at
the stationary point of asymptotic convergence. This structure is locally
optimal near the point of convergence, so RAMDA is guaranteed to obtain the
best structure possible among all methods converging to the same point, making
it the first regularized adaptive method outputting models that possess
outstanding predictive performance while being (locally) optimally structured.
Extensive numerical experiments in large-scale modern computer vision,
language modeling, and speech tasks show that the proposed RAMDA is efficient
and consistently outperforms state of the art for training structured neural
network. Implementation of our algorithm is available at
https://www.github.com/ismoptgroup/RAMDA/.

Structured models, regularized optimization, adaptive method, manifold
identification, deep learning, inexact subproblem solution

###  1 Introduction

Since the recent emergence of ChatGPT, large language models (LLMs) and other
huge deep learning models have garnered much attention and popularity, even
among the public who are unfamiliar with machine learning. A challenge with
such gigantic neural network models is their vast number of model parameters,
reaching hundreds of billions, resulting in expensive storage and inference.
It thus becomes crucial to find ways to exploit structures in trained models
to reduce their spatial and prediction costs without degrading the prediction
performance. An active line of research is to explicitly add a nonsmooth
regularization term to the training objective function and apply proximal
stochastic (sub)gradient methods,111Note that nonsmoothness is often necessary
to induce desirable structures (Iutzeler & Malick, 2020). with or without a
diagonal preconditioner for adaptiveness, to induce a pre-specified type of
desirable structure in the final model (Yang et al., 2019; Yun et al., 2021;
Deleu & Bengio, 2021). Unfortunately, although the added regularizer indeed
induces some desirable structures at the stationary points of the training
objective function, the iterates of these methods only converge to those
stationary points asymptotically, but never really attain such a point at any
iteration. Therefore, whether the output model of these algorithms, which is
also an iterate that is only close enough to a stationary point, indeed
possesses the ideal structure at the nearby stationary point is unknown, and
theoretical analyses of these algorithms do not cover any guarantees regarding
the obtained structure. Indeed, Huang & Lee (2022) oberserd empirically that
the structures obtained by those methods are highly suboptimal and unstable
over iterations. They then proposed a regularized dual averaging method called
RMDA, and proved that after a finite number of steps, the iterates of RMDA can
stably identify the locally optimal structure induced by the regularizer at
the stationary point of asymptotic convergence.222See the first paragraph in
Section 1 and Appendix B of (Huang & Lee, 2022) for a discussion about why the
structure at the point of convergence is locally optimal. This is up to our
knowledge the only method with such structure guarantees for training
structured neural networks. With this property, their experiments demonstrated
that their method also empirically outperforms existing methods on modern
computer vision tasks. However, since RMDA does not incorporate adaptiveness
and their experiments are conducted only on medium-scale image classification
problems, its usefulness beyond computer vision is in doubt.

For a wide range of tasks in deep learning such as language modeling and
speech recognition, researchers have developed numerous architectures to
achieve state-of-the-art prediction performance, including the transformer
(Vaswani et al., 2017) and the LSTM (Hochreiter & Schmidhuber, 1997). The
transformer is also gaining prominence in computer vision for achieving
exceptional performance (Liu et al., 2021), so it is becoming increasingly
important to devise methods that attain satisfactory performance for training
these network architectures with structure. For such modern architectures,
adaptive algorithms like Adam (Kingma & Ba, 2015) that iteratively rescale the
stochastic gradient update directions via a coordinate-wise/diagonal
preconditioner are known to outperform their non-adaptive counterparts and
thus considered state-of-the-art (Denkowski & Neubig, 2017; Anil et al., 2019;
Zhang et al., 2020; Liu et al., 2020; Kunstner et al., 2023). It is hence
expected that the non-adaptive RMDA of Huang & Lee (2022) might not lead to
promising results for such widely-used architectures and tasks.

This work aims to fill this gap to propose a practical regularized adaptive
method with guarantees for both convergence and structure identification.
Since RMDA already has structure guarantees, it might look like we just need
to combine it with an arbitrary preconditioner for adaptiveness. However, this
seemingly easy extension actually requires deliberation in two aspects. First,
except for few exceptions, combination of even a simple diagonal
preconditioner and a nonsmooth regularizer makes the training subproblem
complicated with no closed-form solution. This is totally different from
adaptive methods with no regularization, whose subproblem optimal solution can
be easily computed by coordinate-wise divisions. Therefore, in the regularized
case, the best we can hope for is to apply an iterative approach to
approximately solve the subproblem. This calls for careful design and control
for the measure and the degree of the inexactness in the approximate
subproblem solution. The second aspect is the need of an appropriate
preconditioner that provides not only outstanding empirical performance but
also desirable theoretical properties. The interplay between the inexactness
and the preconditioner makes it particularly difficult to address the
following three challenges simultaneously. (i) _Convergence_ : Proving
convergence of a new algorithm with even just one added component is always a
nontrivial task. For example, although convergence of SGD has been well-
studied for decades, similar guarantees for its adaptive correspondence,
Adagrad, is not established until very recently (Défossez et al., 2022). We
are dealing with both a preconditioner that changes the whole algorithm, just
like from SGD to Adagrad, and the inevitable inexact subproblem solutions that
could nullify many useful properties (regarding the subdifferential) commonly
used in convergence proofs. (ii) _Structure_ : Theoretical guarantees for
structure identification is another critical aim of this work. Inexactness
alone already makes this goal difficult; see Example 1 of (Lee, 2023) for a
simple instance such that even infinitesimal inexactness could hinder
structure identification. Even without inexactness, finding a preconditioner
that leads to structure identification guarantees is already difficult because
no adaptive algorithm, even in the much simpler deterministic and exact
setting, is known to have such a guarantee. (iii) _Subproblem solver_ : Our
goal is a practical algorithm, so we need to solve the subproblem efficiently.
This requires the inexact measure be checkable and the degree quickly
attainable by a well-designed solver, and the preconditioner should make the
subproblem well-conditioned and cannot complicate the computation of the
solver.

To tackle these difficulties, we start from considering structure
identification. We leverage the theory of manifold identification in
variational analysis and nonlinear optimization to design a method that leads
to finite-iteration structure identification guarantees. As discussed by Poon
et al. (2018); Huang & Lee (2022), the key to such guarantees for stochastic
algorithms is to ensure the variance in the stochastic estimations decreases
to zero. Due to the standard practice of data augmentation in deep learning,
the training loss in the objective function is essentially the expected value
of the training loss over a certain probability distribution instead of a
finite sum. We thus draw inspirations from (Lee & Wright, 2012; Huang & Lee,
2022) to consider a dual-averaging-type approach (Nesterov, 2009; Xiao, 2010)
with momentum to attain variance reduction in this setting for the stochastic
gradient estimation. However, we also need variance reduction for the
preconditioner, so we carefully select a preconditioner whose update is in a
manner similar to dual averaging, and prove that its variance also decreases
to zero. We then conceive an implementable and practical subgradient-norm-
based inexactness measure compatible with the structure identification theory.
Further requirements are then added to the inexactness degree and the
preconditioner to ensure convergence, and we also safeguard the preconditioner
to keep the subproblems well-conditioned and the computation simple. We then
propose to solve the subproblem by a proximal gradient (PG) solver that
provably achieves our inexactness requirement efficiently. This leads to our
Regularized Adaptive Momentum Dual Averaging (RAMDA) algorithm.

We summarize our main contributions as follows.

  1. 1.

_An adaptive algorithm for finding locally optimal structures_ : RAMDA is the
first regularized _adaptive_ method guaranteed to find the locally optimal
structure possessed by the stationary point to which its iterates converge. It
thus produces models that are more structured while retaining the superb
prediction performance of adaptive methods.

  2. 2.

_Efficient subproblem solver for regularized adaptive methods_ : We propose an
implementable inexactness condition and a companion highly efficient
subproblem solver for regularized adaptive methods (including ours and
existing ones) whose subproblems have no closed-form solution. We show that
the induced inexactness does not affect convergence or structure
identification guarantees. This condition and subproblem solver thus also
serve as the key step for realizing existing regularized adaptive algorithmic
frameworks.

  3. 3.

_A method with outstanding empirical performance_ : Experiments on training
modern neural networks in computer vision (ImageNet), language modeling
(Transformer-XL), and speech (Tacotron2) with structured sparsity show that
RAMDA steadily outperforms state of the art by achieving higher structured
sparsity ratio and better prediction performance simultaneously.

###  2 Related Work

##### Dual Averaging for Deep Learning.

Our method is motivated by Huang & Lee (2022) that adapted the famous
regularized dual averaging (Xiao, 2010; Lee & Wright, 2012) approach with
momentum to train structured neural network models with data augmentation.
They selected dual averaging for the gradient estimation to achieve variance
reduction for structure guarantees, but their algorithm does not allow for
adaptiveness. Inspired by this approach, we also take dual-averaging-like
updates for the diagonal preconditioner in the subproblem for adaptiveness.
Our preconditioner design also borrows ideas from the empirically successful
MADGRAD of Defazio & Jelassi (2022) for training non-regularized neural
networks. RAMDA can thus also be seen as a generalization of MADGRAD to the
regularized setting. Since no regularizer is present, unlike RAMDA, the
subproblem of MADGRAD has a closed-form solution and no structure is expected.
Moreover, Defazio & Jelassi (2022) only analyzed convergence rates of the
objective value when the problem is convex. Our analysis of (i) variance
reduction in the preconditioner, (ii) convergence in the nonconvex nonsmooth
regularized case, and (iii) structure identification guarantees are novel and
closer to properties desirable in practice. The first two items are also
applicable when no regularizer is present, so our theory also expands
guarantees for MADGRAD.

##### Regularized Stochastic Algorithms for Deep Learning.

Other than RMDA, there are several works on training structured neural
networks through regularization and its proximal operator, but none have
structure guarantees. Yang et al. (2019) considered a simple regularized SGD
method with momentum, but their convergence analysis is only for the
nonadaptive case. Yun et al. (2021) studied a general regularized adaptive
framework ProxGen that incorporates diagonal preconditioners, and showed that
the subgradient of the objective function can decrease to the reciprocal of
the batch size, but their result does not guarantee further convergence to
stationary points. Moreover, they do not allow inexactness in the subproblem,
so their framework can be realized for only a small class of problems. Deleu &
Bengio (2021) extended ProxGen to the case of group-sparsity regularizers,
whose corresponding subproblem indeed has no closed-form solution. They
applied the Newton-Raphson method to obtain nearly-optimal subproblem
solutions, and proposed a seemingly mild inexactness condition to study
convergence. Unfortunately, their condition is not checkable, and their
corresponding convergence guarantee requires the regularizer to be locally
smooth around each iterate, which excludes most regularizers that induce
meaningful structures. On the other hand, we will show that with our
implementable inexactness condition, ProxGen still possesses the same
convergence guarantees in (Yun et al., 2021) without any additional
requirement on the regularizer. Moreover, we will see in Section 6 that the
time cost of the subproblem solver of Deleu & Bengio (2021) is prohibitively
high.

##### Structure and Manifold Identification.

The major tool for our structure guarantees is the theory of manifold
identification (Hare & Lewis, 2004, 2007; LewZ13a; Lee, 2023) in variational
analysis and nonlinear optimization. This theory shows that points possessing
the same structure induced by the regularizer at a stationary point form a
smooth manifold around this stationary point, and with properties from the
regularizer, if a sequence of points converges to this stationary point with
their corresponding subgradients decreasing to zero, this sequence is
guaranteed to eventually stay in this manifold, thus identifying the
structure. Lee & Wright (2012); Poon et al. (2018); Huang & Lee (2022) have
leveraged this tool to show manifold identification for various stochastic
algorithms, and the common key, as pointed out by Poon et al. (2018), is
variance reduction. Our analysis uses a result given in (Ruszczyński, 1980) to
prove so for both the gradient estimator and the preconditioner.

###  3 Problem Setting and Algorithm

As described in Section 1, we consider the case in which the training
objective function is the expectation over a probability distribution as
follows.

$\min\nolimits_{W\in{\mathcal{E}}}\quad
F\left(W\right)\coloneqq\mathbb{E}_{\xi\sim{\mathcal{D}}}\left[f_{\xi}\left(W\right)\right]+\psi\left(W\right),$
(1)

where ${\mathcal{E}}$ is a Euclidean space with inner product
$\langle\cdot,\,\cdot\rangle$ and its induced norm ${\left\|{\cdot}\right\|}$,
${\mathcal{D}}$ is a distribution over a space $\Omega$ representing all
possible data modifications, $f_{\xi}$ is differentiable almost everywhere for
any $\xi$, and the possibly nonsmooth regularizer $\psi(W)$ is for promoting a
desirable structure in the optimal solutions.333The weight for the
regularization term is embedded in $\psi(W)$ in our formulation.

Our algorithm can be seen as a double-dual averaging method that incorporates
a proximal operation for the regularization, momentum, and dual averaging for
updating both the stochastic gradient estimation and the preconditioner. For
the ease of description, we assume without loss of generality that
${\mathcal{E}}=\mathbb{R}^{n}$ in this section. At the $t$th iteration with
learning rate $\eta_{t}$ and iterate $W^{t-1}$, we first draw an independent
and identically distributed sample $\xi_{t}\sim{\mathcal{D}}$, calculate
$s_{t}\coloneqq\eta_{t}\sqrt{t}$, compute the stochastic (sub)gradient
$G^{t}\coloneqq\nabla f_{\xi_{t}}(W^{t-1})$ of the loss function at the
current point $W^{t-1}$ with respect to $\xi_{t}$, and then update the
weighted sum $V_{t}$ of historical stochastic gradients and the weighted sum
$U_{t}$ of their squared norms using the value $s_{t}$:

$\begin{cases}V_{0}\coloneqq 0,\quad V_{t}\coloneqq
V_{t-1}+s_{t}G^{t},\quad&\forall t>0,\\\ U_{0}\coloneqq 0,\quad U_{t}\coloneqq
U_{t-1}+s_{t}G^{t}\circ G^{t},\quad&\forall t>0,\end{cases}$ (2)

where $\circ$ denotes the Hadamard (pointwise) product in ${\mathcal{E}}$. We
then construct the preconditioner $P^{t}$ by

$P^{t}\coloneqq\operatorname{Diag}(\sqrt[3]{U^{t}}+\epsilon),$ (3)

where $\epsilon>0$ is a (usually small) constant for numerical stability and
$\operatorname{Diag}(\cdot)$ is the diagonal matrix whose diagonal entries are
the elements of the input vector. The update direction is then obtained by
(approximately) solving the following subproblem.

$\begin{gathered}\hat{W}^{t}\approx\operatorname*{arg\,min}\nolimits_{W}\,\big{(}Q_{t}(W)\coloneqq\alpha_{t}\psi(W)+\langle
V^{t},\,W\rangle\\\ +\frac{1}{2}\langle
W-W^{0},\,P^{t}(W-W^{0})\rangle\big{)},\quad\alpha_{t}\coloneqq\sum_{k=1}^{t}s_{k},\end{gathered}$
(4)

where $W^{0}$ is the initial point. Details regarding the subproblem Eq. 4 and
how to solve it are deferred to Section 4. The iterate is then updated by
averaging $\hat{W}^{t}$ and $W^{t-1}$ with some $c_{t}\in[0,1]$:

$\displaystyle W^{t}=\left(1-c_{t}\right)W^{t-1}+c_{t}\hat{W}^{t}.$ (5)

If one sets $P^{t}=\sqrt{t}I$, where $I$ is the identity matrix, this
framework recovers RMDA of (Huang & Lee, 2022). On the other hand, with Eq. 3,
adaptiveness is incorporated to obtain our RAMDA. The averaging step in Eq. 5
with $c_{t}\neq 1$ can be interpreted as incorporating a momentum term in the
non-regularized non-adaptive case (Tao et al., 2018; Jelassi & Defazio, 2020).

$V^{0}\leftarrow 0,\quad U^{0}\leftarrow 0,\quad\alpha_{0}\leftarrow 0$

for  _$t=1,\dotsc,T$_ do

Sample $\xi_{t}\sim{\mathcal{D}},\quad
s_{t}\leftarrow\eta_{t}\sqrt{t},\quad\alpha_{t}\leftarrow\alpha_{t-1}+s_{t}$

$G^{t}\leftarrow\nabla f_{\xi_{t}}(W^{t-1})$

Compute $V^{t},U^{t}$ by Eq. 2

Construct $P^{t}$ by Eq. 3, and
$\theta_{t}\leftarrow\max(\operatorname{diag}(P^{t}))^{-1}$

Compute $\hat{W}^{t}$ in Eq. 4 by PG
$(W^{t},W^{0},\alpha_{t}^{-1}V^{t},\alpha_{t}^{-1}P^{t},\alpha_{t}\theta_{t},T_{2},\epsilon_{t})$

Update $W^{t}$ by Eq. 5

output : $W^{T}$

Algorithm 1 RAMDA
$(W^{0},T,T_{2},\epsilon,\\{\eta_{t}\\},\\{c_{t}\\},\\{\epsilon_{t}\\})$

###  4 Subproblem Solver

Given an iterate $W^{t-1}$, a momentum term $m_{t}$, a preconditioner $P^{t}$,
and a stepsize $\eta_{t}$, existing regularized adaptive stochastic gradient
algorithms for Eq. 1 can be summarized in the following form (Yun et al.,
2021):

$\begin{gathered}W^{t}=\operatorname*{arg\,min}\nolimits_{W}\,\big{(}\hat{Q}_{t}(W)\coloneqq\langle
m_{t},\,W\rangle\\\ +\frac{1}{2\eta_{t}}\langle
W-W^{t-1},\,P^{t}(W-W^{t-1})\rangle+\psi(W)\big{)},\end{gathered}$ (6)

whose form is similar to Eq. 4. When the preconditioner $P^{t}$ is a multiple
of the identity matrix like in the case of (Huang & Lee, 2022), the exact
subproblem solution of Eq. 4 can be efficiently computed through the proximal
operator associated with the regularizer. However, a major difficulty for
realizing regularized adaptive methods, including the proposed RAMDA and the
framework of Yun et al. (2021) whose preconditioners are not a multiple of the
identity, is that except for few special regularizers, the subproblem usually
has no closed-form solution. We therefore consider using approximate solutions
of the subproblem.

We propose to apply a few iterations of proximal gradient (PG) (see, e.g.,
Beck & Teboulle, 2009; Nesterov, 2013) to approximately solve the subproblems
in Eqs. 4 and 6 when no closed-form solution is available, and we will show
theoretically and empirically in the following sections that the inexactness
of such approximate solutions has barely any effects on the theoretical
guarantees and the final model quality. For the inexactness of the approximate
solution in Eq. 4, we require

$\min_{s\in\partial
Q_{t}(\hat{W}^{t})}{\left\|{s}\right\|}\leq\epsilon_{t},\quad
Q_{t}(\hat{W}^{t})\leq Q_{t}(W^{t-1}),$ (7)

for some pre-specified $\epsilon_{t}$, where $\partial Q_{t}(W^{t+1})$ is the
(limiting) subdifferential (see, e.g., Rockafellar & Wets, 2009, Definition
8.3). This condition can be easily checked using information available in the
PG iterations. For the sake of time efficiency, we also impose an upper limit
for the number of PG iterations. Likewise, when applying our subproblem solver
to Eq. 6, we enforce Eq. 7 but with $Q_{t}$ replaced by $\hat{Q}_{t}$ and
$\hat{W}^{t}$ by $W^{t}$. We focus on the case of diagonal and positive
$P^{t}$, and thus the largest eigenvalue $\max(\operatorname{diag}(P^{t}))$,
where $\operatorname{diag}(\cdot)$ is the vector formed by the diagonal
entries of the input matrix, can be calculated easily and used to compute a
step size guaranteeing sufficient objective decrease. For cases in which this
value is difficult to obtain, one can apply a simple backtracking linesearch
for the subproblem to find a suitable step size efficiently. This PG
subproblem solver is summarized in Algorithm 2. To guarantee convergence for
both our algorithm and the framework of Yun et al. (2021), our analysis in
Section 5 requires that $\\{\epsilon_{t}\\}$ satisfy

$\bar{\epsilon}\coloneqq\sum\nolimits_{t=0}^{\infty}\epsilon_{t}^{2}<\infty.$
(8)

For any $\epsilon_{t}>0$, we will show in Section 5 that Eq. 7 holds after at
most $O(\epsilon_{t}^{-2})$ iterations of Algorithm 2.

if  _$\psi$ is nonconvex_ then $\theta\leftarrow\theta/2$

for  _$j=1,\dotsc,T_{2}$_ do

$Z^{j}\leftarrow\mbox{\rm prox}_{\psi}(Z^{j-1}-\theta(V+P(Z^{j-1}-W^{0})))$

if  _Eq. 7 holds with $\epsilon_{t}=\hat{\epsilon}$ and $\hat{W}^{t}=Z^{j}$_
then

$Z^{T_{2}}\leftarrow Z^{j}$, and break

output : $Z^{T_{2}}$

Algorithm 2 PG$(Z^{0},W^{0},V,P,\theta,T_{2},\hat{\epsilon})$

###  5 Analysis

This section discusses theoretical guarantees for RAMDA and the proposed
subproblem solver in Algorithm 2. We also prove convergence guarantees for
applying PG to approximately solve Eq. 6 for the framework of Yun et al.
(2021). Due to the space limit, all proofs are in the appendices. Some of our
results are inspired by (Huang & Lee, 2022), but with the added inexactness in
Eq. 4 and the adaptiveness for the preconditioner, the analysis is nontrivial.
Recall that we assume that $f_{\xi}$ is differentiable only almost everywhere
but not everywhere, which conforms with widely-used network structures like
ReLU-type activations.

We first show that Eq. 7 can be attained by Algorithm 2.

######  Theorem 1.

Assume that Eq. 4 and Eq. 6 has at least one optimal solution with a finite
optimal objective value. Given $\epsilon_{t}>0$, the number of iterations of
Algorithm 2 takes to satisfy Eq. 7 for both Eq. 4 and Eq. 6 is
$O(\epsilon_{t}^{-1})$ when $\psi$ is convex and $O(\epsilon_{t}^{-2})$ when
$\psi$ is nonconvex.

Next, we argue that the point of convergence $\bar{W}$ of RAMDA is almost
surely a stationary point.

######  Theorem 2.

Consider $\\{\hat{W}^{t}\\}$ generated by Algorithm 1 for Eq. 1, with Eq. 7
and $\\{c_{t}\\}$ and $\\{\epsilon_{t}\\}$ satisfying $\sum c_{t}=\infty$ and
Eq. 8. Assume there is $L\geq 0$ such that for any $\xi$, $f_{\xi}$ is almost
surely $L$-Lipschitz-continuously-differentiable, so the expectation is also
$L$-Lipschitz-continuously-differentiable, there is $C\geq 0$ such that
$\mathbb{E}_{\xi_{t}\sim{\mathcal{D}}}{\left\|{\nabla
f_{\xi_{t}}\left(W^{t-1}\right)}\right\|}^{4}\leq C$ for all $t$, and that the
set of stationary points ${\mathcal{Z}}\coloneqq\\{W\mid 0\in\partial F(W)\\}$
is nonempty. For any given $W^{0}$, consider the event that
$\\{\hat{W}^{t}\\}$ converges to a point $\bar{W}$ (each event corresponds to
a different $\bar{W}$). If $\partial\psi$ is outer semicontinuous at
$\bar{W}$, this event has a nonzero probability, and $\\{\eta_{t}\\}$
satisfies

$\displaystyle\sum
s_{t}\alpha_{t}^{-1}=\infty,\quad\sum\left(s_{t}\alpha_{t}^{-1}\right)^{2}<\infty,$
$\displaystyle{\left\|{W^{t+1}-W^{t}}\right\|}\left(s_{t}\alpha_{t}^{-1}\right)^{-1}\xrightarrow{\text{a.s.}}0,$

then we have that $\bar{W}\in{\mathcal{Z}}$ with probability one conditional
on this event. Moreover, $\\{W^{t}\\}$ also converges to this stationary point
$\bar{W}$.

Usually, convergence to a point requires some further regularity conditions
like the Kurdyka–Łojasiewicz condition and boundedness of the iterates.
However, existing frameworks regarding iterates convergence using such
conditions also require the method analyzed to have a subgradient-descent-like
behavior and to be a descent algorithm. Neither of these hold true even for
the basic stochastic gradient algorithm, and we leave the analysis for this
part as a challenging future work.

Our next key result shows that after a finite number of iterations, iterates
of RAMDA all possess the same structure as that of the point of convergence
$\bar{W}$. For this end, we first need to introduce the notions of partial
smoothness and prox-regularity, and impose these assumptions on $\psi$ at
$\bar{W}$.

######  Definition 1 (Partial Smoothness (Lewis, 2002; Hare & Lewis, 2004)).

A function $\psi$ is partly smooth at a point $\bar{W}$ relative to a set
${\mathcal{M}}_{\bar{W}}\ni\bar{W}$ if

  1. 1.

Around $\bar{W}$, ${\mathcal{M}}_{\bar{W}}$ is a $\mathcal{C}^{2}$-manifold
and $\psi|_{{\mathcal{M}}_{\bar{W}}}$ is $\mathcal{C}^{2}$.

  2. 2.

$\psi$ is regular (finite and the Fréchet subdifferential coincides with the
limiting Fréchet subdifferential) at all points $W\in{\mathcal{M}}_{\bar{W}}$
near $\bar{W}$, with $\partial\psi(W)\neq\emptyset$.

  3. 3.

The affine span of $\partial\psi(\bar{W})$ is a translate of the normal space
to ${\mathcal{M}}_{\bar{W}}$ at $\bar{W}$.

  4. 4.

$\partial\psi$ is continuous at $\bar{W}$ relative to
${\mathcal{M}}_{\bar{W}}$.

We often call ${\mathcal{M}}_{\bar{W}}$ the active manifold at $\bar{W}$.
Locally, this manifold represents all points near $\bar{W}$ that share the
same structure induced by the regularized as $\bar{W}$. Therefore, finding the
active manifold is equivalent to finding the locally optimal structure.

######  Definition 2 (Prox-regularity (Poliquin & Rockafellar, 1996)).

A function $\psi$ is prox-regular at $\bar{W}$ for
$V^{*}\in\partial\psi(\bar{W})$ if $\psi$ is locally lower semi-continuous
around $\bar{W}$, finite at $\bar{W}$, and there is $\rho>0$ such that
$\psi(W_{1})\geq\psi(W_{2})+\langle
V,\,W_{1}-W_{2}\rangle-\frac{\rho}{2}{\left\|{W_{1}-W_{2}}\right\|}^{2}$ for
every $W_{1},W_{2}$ near $\bar{W}$ with $\psi(W_{2})$ close to $\psi(\bar{W})$
and $V\in\partial\psi(W_{2})$ close to $V^{*}$. $\psi$ is prox-regular at
$\bar{W}$ if it is prox-regular for all $V\in\partial\psi(\bar{W})$.

######  Theorem 3.

Consider Algorithm 1 with the conditions in Theorem 2 satisfied. Consider the
event of $\\{\hat{W}^{t}\\}$ converging to a certain point $\bar{W}$ as in
Theorem 2. If the probability of this event is nonzero; $\psi$ is prox-regular
and subdifferentially continuous at $\bar{W}$ and partly smooth at $\bar{W}$
relative to the active $\mathcal{C}^{2}$ manifold ${\mathcal{M}}_{\bar{W}}$;
$\partial\psi$ is outer semicontinuous at $\bar{W}$; and the nondegeneracy
condition

$-\nabla f\left(\bar{W}\right)\in\textup{relative interior
of}\;\partial\psi\left(\bar{W}\right)$

holds at $\bar{W}$, then conditional on this event, almost surely there is
$T_{0}\geq 0$ such that

$\hat{W}^{t}\in{\mathcal{M}}_{\bar{W}},\quad\forall t\geq T_{0}.$

In other words, the active manifold at $\bar{W}$ is identified by the iterates
of Algorithm 1 after a finite number of iterations almost surely.

We note particularly that convex and weakly-convex (Nurminskii, 1973)
functions are all regular, prox-regular, and subdifferentially continuous
everywhere.

We also show that our subproblem solver and condition can be effectively
applied to the framework of Yun et al. (2021) while retaining the same
convergence guarantees. As mentioned in Section 2, our result is much stronger
than that of Deleu & Bengio (2021) for having no unrealistic smoothness
requirement on $\psi$ and using an implementable inexactness condition.

######  Theorem 4.

For the framework in (Yun et al., 2021) with the subproblem solved
approximately by Algorithm 2 such that Eq. 7 holds with $\\{\epsilon_{t}\\}$
satisfying Eq. 8. Then Theorem 1 of (Yun et al., 2021) still holds, but with
the constants $\\{Q_{i}\\}$ being also dependent on $\bar{\epsilon}$.

###  6 Experiments

This section examines the practical performance of RAMDA for training
structured neural networks. As sparsity is arguably one of the most widely
adopted structures in machine learning, we follow Wen et al. (2016) to
consider structured sparsity as the representative structure in our
experiments. Particularly, we employ the group LASSO regularization (Yuan &
Lin, 2006) to encourage group sparsity. We begin from examining the efficiency
and effectiveness of Algorithm 2 for both RAMDA and existing regularized
adaptive methods. We then consider tasks in computer vision, language
modeling, and speech to compare the following algorithms using Pytorch (Paszke
et al., 2019).

  * •

RAMDA: The proposed Algorithm 1.

  * •

RMDA (Huang & Lee, 2022): A regularized (non-adaptive) momentum dual-averaging
method.

  * •

ProxSGD (Yang et al., 2019): A non-adaptive proximal SGD method with momentum.

  * •

ProxGen (Yun et al., 2021): A general framework of stochastic proximal
methods. We follow their experiments to use AdamW with regularization and
apply our PG as the subproblem solver.

  * •

ProxSSI (Deleu & Bengio, 2021): A realization of ProxGen for group sparsity by
using the Newton-Raphson algorithm to solve the subproblem to near optimality.

These algorithms are summarized in Appendix A. For each task, we also provide
for reference a baseline that does not include a group LASSO regularizer in
the training (SGD with momentum (MSGD) for computer vision, and AdamW for the
other two), but our comparison is only among those for training structured
models. Our code for reproducing the experiments and the hyperparameter
settings are available at https://github.com/ismoptgroup/ramda_exp/.
Additional details of the stability of the structure (level of structured
sparsity here) over epochs of RAMDA are available in Appendix D.

We use two criteria for comparison: 1\. Model predictive ability, and 2.
Structured sparsity level. The former is task-dependent, so specifications for
each task are provided in their respective sections. Regarding the latter,
sparsifying neural networks while preserving its performance requires prior
knowledge of model design. A common approach is retaining certain parameters
during the training process, and we adhere to this convention such that the
bias, batch normalization (Ioffe & Szegedy, 2015), layer normalization (Ba et
al., 2016), and embedding layers do not have any sparsity-inducing
regularization imposed on them (Deleu & Bengio, 2021; Peste et al., 2021). For
the rest, we adopt channel-wise grouping for convolutional layers, input-wise
grouping for fully-connected and LSTM layers during the training phase. For
evaluation, our structured sparsity is calculated using the weighted group
sparsity with the weights proportional to the number of parameters in each
group.

Similar to Huang & Lee (2022), we introduce a restart strategy to the
implementation of RAMDA. During each stage of the training, the learning rate
$\eta_{t}$ and the momentum factor $c_{t}$ are fixed. Once the epoch count
enters the next stage, we reset the counter $t$ to $1$ and use the output
parameter $W^{T}$ from the previous round as the new input parameter $W^{0}$
to the same algorithm, set $\alpha_{t},V^{t}$ and $U^{t}$ to $0$, but keep the
scheduling for $\eta$ and $c$ going without resetting them. We initialize
$c_{t}$ as either $0.1$ or $0.01$, depending on the problems, and use a
constant $c_{t}$ until the final stage, where we gradually increase it by
$c_{t}=\min(c\sqrt{i},1)$, where $i$ counts the training steps executed at
this final stage. This momentum strategy is applied to both RAMDA and RMDA in
our experiments.

We run each experiment with three different random initializations and show
the mean and standard deviation of the validation predictive performance and
the structured sparsity of the final model of all methods.

####  6.1 Subproblem

Table 1: Group sparsity and validation accuracy of different subproblem stopping criteria. |  | No early stopping | Early stopping
---|---|---|---
Model/Data | Algorithm | Validation accuracy | Group sparsity | Validation accuracy | Group sparsity
VGG19 / | ProxGen | 92.70 $\pm$ 0.22% | 88.79 $\pm$ 0.03% | 92.67 $\pm$ 0.11% | 86.92 $\pm$ 0.38%
CIFAR10 | RAMDA | 92.70 $\pm$ 0.21% | 86.73 $\pm$ 0.30% | 92.92 $\pm$ 0.18% | 86.29 $\pm$ 0.40%
ResNet50 / | ProxGen | 73.63 $\pm$ 0.13% | 74.69 $\pm$ 0.58% | 74.03 $\pm$ 0.08% | 67.55 $\pm$ 3.05%
CIFAR100 | RAMDA | 69.90 $\pm$ 1.54% | 69.54 $\pm$ 2.07% | 71.23 $\pm$ 1.39% | 67.48 $\pm$ 1.59%

We start from showing the effectiveness of our proposed subproblem solver for
RAMDA and ProxGen. For both approaches, we use Theorem 2 of (Deleu & Bengio,
2021) to safely screen out a portion of groups that will be zero at the
optimal subproblem solution, and opt for the PG algorithm to solve the
remaining parts. We consider two practical stopping criteria for PG: 1.
Running until it reaches the maximum iterations (no early stopping), and 2.
Terminate when the subproblem objective improvement is small (early stopping).
For the former, we set the maximum to $100$. For the latter, we terminate PG
early if $(Q_{t}(Z^{j-1})-Q_{t}(Z^{j}))/(|Q_{t}(Z^{j}|+1)<10^{-8}$ is reached.
Moreover, to ensure incorporation of the preconditioner into ProxGen, we set
its minimum PG iterations to $2$. We examine how these stopping criteria
affect the final model of RAMDA and ProxGen using image classification
problems of a smaller scale. From Table 1, we see that early stopping does not
affect the outcome much. Given that early stopping is more efficient, we will
adopt it in all subsequent experiments.

Next, we compare ProxGen with ProxSSI (these two only differ in the subproblem
solver) to examine the efficiency and performance differences between solving
the subproblems approximately and (almost) exactly in Table 2. We see that our
solver is around 3X faster than ProxSSI, and the model qualities are similar.
We thus exclude ProxSSI from our comparisons in the following experiments due
to its excessively lengthy running time, especially for large-scale tasks.

Table 2: Weighted group sparsity, validation accuracy and time/epoch of ProxSSI and ProxGen for CIFAR10/CIFAR100. We report the average time per epoch using one NVIDIA V100 GPU. Algorithm | Accuracy | Sparsity | Time
---|---|---|---
VGG19/CIFAR10 |
ProxSSI | 92.8 $\pm$ 0.1% | 88.4 $\pm$ 0.2% | 79s
ProxGen | 92.8 $\pm$ 0.0% | 86.6 $\pm$ 0.1% | 24s
VGG19/CIFAR100 |
ProxSSI | 67.3 $\pm$ 0.1% | 78.6 $\pm$ 0.3% | 79s
ProxGen | 68.1 $\pm$ 0.4% | 75.5 $\pm$ 0.2% | 26s
ResNet50/CIFAR10 |
ProxSSI | 94.0 $\pm$ 0.1% | 83.7 $\pm$ 0.6% | 260s
ProxGen | 94.1 $\pm$ 0.1% | 80.4 $\pm$ 0.4% | 70s
ResNet50/CIFAR100 |
ProxSSI | 73.7 $\pm$ 0.4% | 70.4 $\pm$ 0.7% | 251s
ProxGen | 73.6 $\pm$ 0.4% | 65.5 $\pm$ 3.6% | 74s

####  6.2 Image Classification

We conduct a classical computer vision task of training ResNet50 (He et al.,
2016) with the ILSVRC 2012 ImageNet dataset (Russakovsky et al., 2015). The
result in Table 3 shows that RAMDA attains the best validation accuracy and
structured sparsity simultaneously.

Additional experiments in Appendix C show that RAMDA might sometimes perform
worse than existing methods on smaller problems like CIFAR10/100. But for such
smaller problems, the training cost is not very significant, and one can
afford to try more algorithms.

Table 3: Weighted group sparsity and validation accuracy of different methods for ImageNet on ResNet50. Algorithm | Accuracy | Sparsity
---|---|---
MSGD | 77.14 $\pm$ 0.04% | -
ProxSGD | 73.50 $\pm$ 0.20% | 17.54 $\pm$ 1.26%
ProxGen | 74.17 $\pm$ 0.08% | 20.29 $\pm$ 0.22%
RMDA | 74.47 $\pm$ 0.08% | 25.20 $\pm$ 1.69%
RAMDA | 74.53 $\pm$ 0.10% | 29.19 $\pm$ 0.94%

####  6.3 Language Modeling

For language modeling, we train Transformer-XL (base) (Dai et al., 2019) using
the WikiText-103 dataset (Merity et al., 2017). Transformer-XL is comprised of
embedding and non-embedding layers, and in the PyTorch implementation, the
non-embedding layers are built using linear and layer-normalization layers. We
apply group LASSO regularization to the linear layers, and present in Table 4
the perplexity and the weighted group sparsity of the models trained by
different methods. We see that RAMDA gives the best perplexity and structured
sparsity at the same time.

Table 4: Weighted group sparsity and validation perplexity on Transformer-XL with WikiText-103. Alg. | Perplexity | Sparsity | Time/epoch
---|---|---|---
AdamW | 23.00 $\pm$ 0.05 | - | 6261 $\pm$ 21s
ProxSGD | 27.42 $\pm$ 0.02 | 33.1 $\pm$ 1.5% | 6167 $\pm$ 12s
ProxGen | 27.49 $\pm$ 0.19 | 30.5 $\pm$ 0.6% | 6652 $\pm$ 21s
RMDA | 27.10 $\pm$ 0.08 | 36.0 $\pm$ 2.7% | 6184 $\pm$ 20s
RAMDA | 26.97 $\pm$ 0.10 | 36.2 $\pm$ 0.3% | 6954 $\pm$ 30s

####  6.4 Speech Synthesis

We consider Tacotron2 (Shen et al., 2018) for speech synthesis on the LJSpeech
dataset (Ito & Johnson, 2017). We apply regularization to the convolutional,
LSTM, and linear layers of Tacotron2 and show the results in Table 5. Clearly,
RAMDA gives the lowest validation loss and the highest group sparsity.

Table 5: Weighted group sparsity and validation loss on Tacotron2 with LJSpeech. Alg. | Loss | Sparsity | Time/epoch
---|---|---|---
AdamW | 0.39 $\pm$ 0.02 | - | 431 $\pm$ 2s
ProxSGD | 0.50 $\pm$ 0.00 | 34.3 $\pm$ 1.6% | 431 $\pm$ 0s
ProxGen | 0.45 $\pm$ 0.01 | 45.6 $\pm$ 0.9% | 438 $\pm$ 2s
RMDA | 0.46 $\pm$ 0.01 | 45.9 $\pm$ 1.7% | 431 $\pm$ 2s
RAMDA | 0.44 $\pm$ 0.01 | 52.9 $\pm$ 1.6% | 443 $\pm$ 1s

####  6.5 Time Efficiency and Summary

In Tables 4 and 5, we see that although RAMDA and ProxGen have more difficult
subproblems without a closed-form solution to solve, our proposed PG solver is
highly efficient such that the running time of them is still close to other
approaches, making these regularized adaptive approaches practically feasible.

In summary, thanks to its adaptive nature (for better predictive performance)
and its ability of manifold identification (for higher structured sparsity),
RAMDA is superior to state of the art on modern language modeling and speech
synthesis tasks as well as the ImageNet problem. We also observe from the
plots in the appendices that it is possible to further improve the sparsity
level of RAMDA if we run it for more epochs.

###  7 Conclusions

In this work, we propose a regularized dual averaging method with
adaptiveness, RAMDA, for training structured neural networks. Our method is
particularly useful for language models, speech synthesis, and the ImageNet
problem. It outperforms state of the art on modern architectures including
LSTM and transformers. We also propose a subroutine with strong convergence
guarantees to approximately solve the regularized subproblem of both our
method and an existing framework efficiently. Extensive experiments on group
sparsity show that our subproblem solver can greatly reduce the training time
for existing methods, and our proposed RAMDA achieves simultaneously higher
structured sparsity ratio and better prediction performance than existing
methods. Implementation of our method is available at
https://www.github.com/ismoptgroup/RAMDA/.

### References

  * Anil et al. (2019) Anil, R., Gupta, V., Koren, T., and Singer, Y.  Memory efficient adaptive optimization.  In _Advances in Neural Information Processing Systems_ , 2019.
  * Ba et al. (2016) Ba, J. L., Kiros, J. R., and Hinton, G. E.  Layer normalization.  Technical report, 2016.  arXiv:1607.06450.
  * Beck (2017) Beck, A.  _First-Order Methods in Optimization_.  SIAM - Society for Industrial and Applied Mathematics, Philadelphia, PA, United States, 2017.
  * Beck & Teboulle (2009) Beck, A. and Teboulle, M.  A fast iterative shrinkage-thresholding algorithm for linear inverse problems.  _SIAM Journal on Imaging Sciences_ , 2(1):183–202, 2009.
  * Dai et al. (2019) Dai, Z., Yang, Z., Yang, Y., Carbonell, J., Le, Q., and Salakhutdinov, R.  Transformer-XL: Attentive language models beyond a fixed-length context.  In _Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics_ , 2019.
  * Defazio & Jelassi (2022) Defazio, A. and Jelassi, S.  Adaptivity without compromise: A momentumized, adaptive, dual averaged gradient method for stochastic optimization.  _Journal of Machine Learning Research_ , 23(144):1–34, 2022.
  * Défossez et al. (2022) Défossez, A., Bottou, L., Bach, F., and Usunier, N.  A simple convergence proof of adam and adagrad.  _Transactions on Machine Learning Research_ , 2022.  ISSN 2835-8856.
  * Deleu & Bengio (2021) Deleu, T. and Bengio, Y.  Structured sparsity inducing adaptive optimizers for deep learning.  Technical report, 2021.  arXiv:2102.03869.
  * Denkowski & Neubig (2017) Denkowski, M. and Neubig, G.  Stronger baselines for trustable results in neural machine translation.  Technical report, 2017.  arXiv:1706.09733.
  * Hare & Lewis (2004) Hare, W. L. and Lewis, A. S.  Identifying active constraints via partial smoothness and prox-regularity.  _Journal of Convex Analysis_ , 11(2):251–266, 2004.
  * Hare & Lewis (2007) Hare, W. L. and Lewis, A. S.  Identifying active manifolds.  _Algorithmic Operations Research_ , 2(2):75–75, 2007.
  * He et al. (2016) He, K., Zhang, X., Ren, S., and Sun, J.  Deep residual learning for image recognition.  In _IEEE conference on computer vision and pattern recognition_ , 2016\.
  * Hochreiter & Schmidhuber (1997) Hochreiter, S. and Schmidhuber, J.  Long short-term memory.  _Neural computation_ , 9(8):1735–1780, 1997.
  * Huang & Lee (2022) Huang, Z.-S. and Lee, C.-p.  Training structured neural networks through manifold identification and variance reduction.  In _International Conference on Learning Representations_ , 2022.
  * Ioffe & Szegedy (2015) Ioffe, S. and Szegedy, C.  Batch normalization: Accelerating deep network training by reducing internal covariate shift.  In _International conference on machine learning_ , pp. 448–456, 2015.
  * Ito & Johnson (2017) Ito, K. and Johnson, L.  The LJ speech dataset, 2017.
  * Iutzeler & Malick (2020) Iutzeler, F. and Malick, J.  Nonsmoothness in machine learning: specific structure, proximal identification, and applications.  _Set-Valued and Variational Analysis_ , 28(4):661–678, 2020.
  * Jelassi & Defazio (2020) Jelassi, S. and Defazio, A.  Dual averaging is surprisingly effective for deep learning optimization.  Technical report, 2020.  arXiv:2010.10502.
  * Kingma & Ba (2015) Kingma, D. P. and Ba, J.  Adam: A method for stochastic optimization.  In _International Conference on Learning Representations_ , 2015.
  * Krizhevsky (2009) Krizhevsky, A.  Learning multiple layers of features from tiny images.  Technical report, 2009.
  * Kunstner et al. (2023) Kunstner, F., Chen, J., Lavington, J. W., and Schmidt, M.  Noise is not the main factor behind the gap between sgd and adam on transformers, but sign descent might be.  In _International Conference on Learning Representations_ , 2023.
  * LeCun et al. (1998) LeCun, Y., Bottou, L., Bengio, Y., and Haffner, P.  Gradient-based learning applied to document recognition.  _Proceedings of the IEEE_ , 86(11):2278–2324, 1998.
  * Lee (2023) Lee, C.-p.  Accelerating inexact successive quadratic approximation for regularized optimization through manifold identification.  _Mathematical Programming_ , 2023.
  * Lee & Wright (2012) Lee, S. and Wright, S. J.  Manifold identification in dual averaging for regularized stochastic online learning.  _Journal of Machine Learning Research_ , 13:1705–1744, 2012\.
  * Lewis (2002) Lewis, A. S.  Active sets, nonsmoothness, and sensitivity.  _SIAM Journal on Optimization_ , 13(3):702–725, 2002.
  * Liu et al. (2020) Liu, L., Liu, X., Gao, J., Chen, W., and Han, J.  Understanding the difficulty of training transformers.  In _Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP)_ , 2020.
  * Liu et al. (2021) Liu, Z., Lin, Y., Cao, Y., Hu, H., Wei, Y., Zhang, Z., Lin, S., and Guo, B.  Swin transformer: Hierarchical vision transformer using shifted windows.  In _Proceedings of the IEEE/CVF international conference on computer vision_ , 2021.
  * Loshchilov & Hutter (2019) Loshchilov, I. and Hutter, F.  Decoupled weight decay regularization.  In _International Conference on Learning Representations_ , 2019.
  * Merity et al. (2017) Merity, S., Xiong, C., Bradbury, J., and Socher, R.  Pointer sentinel mixture models.  In _International Conference on Learning Representations_ , 2017.
  * Nesterov (2009) Nesterov, Y.  Primal-dual subgradient methods for convex problems.  _Mathematical programming_ , 120(1):221–259, 2009\.
  * Nesterov (2013) Nesterov, Y.  Gradient methods for minimizing composite functions.  _Mathematical Programming_ , 140(1):125–161, 2013\.
  * Nurminskii (1973) Nurminskii, E. A.  The quasigradient method for the solving of the nonlinear programming problems.  _Cybernetics_ , 9(1):145–150, 1973.
  * Paszke et al. (2019) Paszke, A., Gross, S., Massa, F., Lerer, A., Bradbury, J., Chanan, G., Killeen, T., Lin, Z., Gimelshein, N., Antiga, L., et al.  Pytorch: An imperative style, high-performance deep learning library.  _Advances in neural information processing systems_ , pp. 8026–8037, 2019.
  * Peste et al. (2021) Peste, A., Iofinova, E., Vladu, A., and Alistarh, D.  Ac/dc: Alternating compressed/decompressed training of deep neural networks.  In _Advances in neural information processing systems_ , 2021.
  * Poliquin & Rockafellar (1996) Poliquin, R. and Rockafellar, R.  Prox-regular functions in variational analysis.  _Transactions of the American Mathematical Society_ , 348(5):1805–1838, 1996.
  * Poon et al. (2018) Poon, C., Liang, J., and Schönlieb, C.-B.  Local convergence properties of SAGA/prox-SVRG and acceleration.  In _International Conference on Machine Learning_ , 2018.
  * Rockafellar & Wets (2009) Rockafellar, R. T. and Wets, R. J.-B.  _Variational analysis_ , volume 317.  Springer Science & Business Media, 2009.
  * Russakovsky et al. (2015) Russakovsky, O., Deng, J., Su, H., Krause, J., Satheesh, S., Ma, S., Huang, Z., Karpathy, A., Khosla, A., Bernstein, M., Berg, A. C., and Fei-Fei, L.  Imagenet large scale visual recognition challenge.  _International journal of computer vision_ , 115(3):211–252, 2015.
  * Ruszczyński (1980) Ruszczyński, A.  Feasible direction methods for stochastic programming problems.  _Mathematical Programming_ , 19:220–229, 1980.
  * Shen et al. (2018) Shen, J., Pang, R., Weiss, R. J., Schuster, M., Jaitly, N., Yang, Z., Chen, Z., Zhang, Y., Wang, Y., Skerrv-Ryan, R., Saurous, R. A., Agiomyrgiannakis, Y., and Wu, Y.  Natural tts synthesis by conditioning wavenet on mel spectrogram predictions.  In _2018 IEEE international conference on acoustics, speech and signal processing (ICASSP)_ , 2018.
  * Simonyan & Zisserman (2015) Simonyan, K. and Zisserman, A.  Very deep convolutional networks for large-scale image recognition.  In _International Conference on Learning Representations_ , 2015.
  * Tao et al. (2018) Tao, W., Pan, Z., Wu, G., and Tao, Q.  Primal averaging: A new gradient evaluation step to attain the optimal individual convergence.  _IEEE transactions on cybernetics_ , 50(2):835–845, 2018.
  * Vaswani et al. (2017) Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, Ł., and Polosukhin, I.  Attention is all you need.  In _Advances in neural information processing systems_ , 2017.
  * Wen et al. (2016) Wen, W., Wu, C., Wang, Y., Chen, Y., and Li, H.  Learning structured sparsity in deep neural networks.  _Advances in neural information processing systems_ , pp. 2074–2082, 2016.
  * Xiao et al. (2017) Xiao, H., Rasul, K., and Vollgraf, R.  Fashion-mnist: a novel image dataset for benchmarking machine learning algorithms.  Technical report, 2017.  arXiv:1708.07747.
  * Xiao (2010) Xiao, L.  Dual averaging methods for regularized stochastic learning and online optimization.  _Journal of Machine Learning Research_ , 11(88):2543–2596, 2010.
  * Xie et al. (2022) Xie, Z., Zhang, Z., Cao, Y., Lin, Y., Bao, J., Yao, Z., Dai, Q., and Hu, H.  SimMIM: A simple framework for masked image modeling.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pp. 9653–9663, 2022.
  * Yang et al. (2019) Yang, Y., Yuan, Y., Chatzimichailidis, A., van Sloun, R. J., Lei, L., and Chatzinotas, S.  ProxSGD: Training structured neural networks under regularization and constraints.  In _International Conference on Learning Representations_ , 2019.
  * Yuan & Lin (2006) Yuan, M. and Lin, Y.  Model selection and estimation in regression with grouped variables.  _Journal of the Royal Statistical Society: Series B (Statistical Methodology)_ , 68(1):49–67, 2006.
  * Yun et al. (2021) Yun, J., Lozano, A. C., and Yang, E.  Adaptive proximal gradient methods for structured neural networks.  In _Advances in Neural Information Processing Systems_ , 2021.
  * Zhang et al. (2020) Zhang, J., Karimireddy, S. P., Veit, A., Kim, S., Reddi, S., Kumar, S., and Sra, S.  Why are adaptive methods good for attention models?  In _Advances in Neural Information Processing Systems_ , 2020.

## Appendices

###  Appendix A More Experiment Details

In Table 6, we summarize details of the algorithms compared in Section 6.

Table 6: Algorithms used in the experiments. Algorithm | Unregularized counterpart | Subproblem
---|---|---
RAMDA | MADGRAD (Defazio & Jelassi, 2022) | PG
RMDA | MDA (Jelassi & Defazio, 2020) | Closed-form solution
ProxSGD | MSGD | Closed-form solution
ProxGen | AdamW (Loshchilov & Hutter, 2019) | PG
ProxSSI | AdamW (Loshchilov & Hutter, 2019) | Newton-Raphson

###  Appendix B Proofs

This section provides proofs of the theoretical results stated in Section 5.
We restate these results and provide their corresponding proofs right after
each statement.

######  Theorem 1.

Assume that Eq. 4 and Eq. 6 has at least one optimal solution with a finite
optimal objective value. Given $\epsilon_{t}>0$, the number of iterations of
Algorithm 2 takes to satisfy Eq. 7 for both Eq. 4 and Eq. 6 is
$O(\epsilon_{t}^{-1})$ when $\psi$ is convex and $O(\epsilon_{t}^{-2})$ when
$\psi$ is nonconvex.

###### Proof.

We use the notation

$\bar{Q}_{t}(Z)=f_{t}(Z)+\psi(Z)$

to unify the two objective function $Q_{t}/\alpha_{t}$ and $\hat{Q}_{t}$,
where $f_{t}$ is the smooth part and we define the Lipschitz constant of
$\nabla f_{t}$ as $L$.

At each iteration, PG solves the following subproblem

$Z^{j+1}\in\operatorname*{arg\,min}_{Z}\quad\langle\nabla
f_{t}(Z^{j}),\,Z-Z^{j}\rangle+\frac{1}{2\theta_{t}}{\left\|{Z-Z^{j}}\right\|}^{2}+\psi(Z),$

and thus from the first-order optimality conditions, clearly we have

$\nabla f_{t}(Z^{j+1})-\nabla
f_{t}(Z^{j})-\frac{1}{\theta_{t}}\left(Z^{j+1}-Z^{j}\right)\in\partial\bar{Q}_{t}(Z^{j+1}).$

We thus have from the Lipschitz continuity of $\nabla f_{t}$ that

$\min_{s\in\partial\bar{Q}_{t}(Z^{j+1})}{\left\|{s}\right\|}\leq{\left\|{\nabla
f_{t}(Z^{j+1})-\nabla
f_{t}(Z^{j})}\right\|}+\frac{1}{\theta_{t}}{\left\|{Z^{j+1}-Z^{j}}\right\|}\leq\left(L+\theta_{t}^{-1}\right){\left\|{Z^{j+1}-Z^{j}}\right\|}.$
(9)

Note that $\bar{Q}_{t}$ is lower bounded, say by $\bar{Q}_{t}^{*}$, and has at
least one solution $Z^{*}$ (unique when $\psi$ is convex).

In the case that $\psi$ is convex, we know that $\theta_{t}=1/L$, and standard
analysis of proximal gradient (see, for example, Beck, 2017, Theorem 10.27)
gives that

$L{\left\|{Z^{j+1}-Z^{j}}\right\|}\leq\frac{2L{\left\|{Z^{0}-Z^{*}}\right\|}}{j},\quad\forall
j\geq 0.$ (10)

Therefore, the combination of Eq. 9 and Eq. 10 shows that it takes
$O(\epsilon_{t}^{-1})$ iterations for PG to reach the required precision of
$\epsilon_{t}$.

When $\psi$ is nonconvex, we have that $\theta_{t}=1/(2L)$ and standard
analysis (Beck, 2017, Theorem 10.15) gives

$\min_{k=0,1,\dots,j}{\left\|{Z^{j+1}-Z^{j}}\right\|}\leq\frac{C}{\sqrt{j}}$
(11)

for some constant $C$ depending on $L$ and
$\bar{Q}_{t}(W^{t})-\bar{Q}_{t}^{*}$. Therefore, Eq. 11 and Eq. 9 show that it
takes $O(\epsilon_{t}^{-2})$ iterations to reach the desired precision. ∎

######  Theorem 2.

Consider $\\{\hat{W}^{t}\\}$ generated by Algorithm 1 for Eq. 1, with Eq. 7
and $\\{c_{t}\\}$ and $\\{\epsilon_{t}\\}$ satisfying $\sum c_{t}=\infty$ and
Eq. 8. Assume there is $L\geq 0$ such that for any $\xi$, $f_{\xi}$ is almost
surely $L$-Lipschitz-continuously-differentiable, so the expectation is also
$L$-Lipschitz-continuously-differentiable, there is $C\geq 0$ such that
$\mathbb{E}_{\xi_{t}\sim{\mathcal{D}}}{\left\|{\nabla
f_{\xi_{t}}\left(W^{t-1}\right)}\right\|}^{4}\leq C$ for all $t$, and that the
set of stationary points ${\mathcal{Z}}\coloneqq\\{W\mid 0\in\partial F(W)\\}$
is nonempty. For any given $W^{0}$, consider the event that
$\\{\hat{W}^{t}\\}$ converges to a point $\bar{W}$ (each event corresponds to
a different $\bar{W}$). If $\partial\psi$ is outer semicontinuous at
$\bar{W}$, this event has a nonzero probability, and $\\{\eta_{t}\\}$
satisfies

$\displaystyle\sum
s_{t}\alpha_{t}^{-1}=\infty,\quad\sum\left(s_{t}\alpha_{t}^{-1}\right)^{2}<\infty,$
(12)
$\displaystyle{\left\|{W^{t+1}-W^{t}}\right\|}\left(s_{t}\alpha_{t}^{-1}\right)^{-1}\xrightarrow{\text{a.s.}}0,$
(13)

then we have that $\bar{W}\in{\mathcal{Z}}$ with probability one conditional
on this event. Moreover, $\\{W^{t}\\}$ also converges to this stationary point
$\bar{W}$.

###### Proof.

First, we prove that when $\\{\hat{W}^{t}\\}$ converges to $\bar{W}$,
$\\{W^{t}\\}$ also converges to $\bar{W}$. From Eq. 5, we have that

${\left\|{W^{t}-\bar{W}}\right\|}\leq(1-c_{t}){\left\|{W^{t-1}-\bar{W}}\right\|}+c_{t}{\left\|{\hat{W}^{t}-\bar{W}}\right\|}.$
(14)

Since $\hat{W}^{t}\rightarrow\bar{W}$, for any $\epsilon>0$ we can find an
integer $t_{\epsilon}\geq 0$ such that $\|{\hat{W}^{t}-\bar{W}}\|\leq\epsilon$
for all $t\geq t_{\epsilon}$. Therefore, by deducting $\epsilon$ from both
sides of Eq. 14, we get

${\left\|{W^{t}-\bar{W}}\right\|}-\epsilon\leq\left(\prod_{k=t_{\epsilon}}^{t}(1-c_{t})\right)\left({\left\|{W^{t_{\epsilon}-1}-\bar{W}}\right\|}-\epsilon\right)\leq\exp\left(-\sum_{k=t_{\epsilon}}^{t}c_{t}\right)\left({\left\|{W^{t_{\epsilon}-1}-\bar{W}}\right\|}-\epsilon\right),\quad\forall
t\geq t_{\epsilon}.$

By letting $t$ approach infinity and noting that $\sum c_{t}=\infty$, we see
that

$\lim_{t\rightarrow\infty}{\left\|{W^{t}-\bar{W}}\right\|}\leq\epsilon.$

Because $\epsilon$ is arbitrary, we see that $\|{W^{t}-\bar{W}}\|\rightarrow
0$, and hence $\\{W^{t}\\}$ converges to $\bar{W}$.

Next, consider the update of $\alpha_{t}^{-1}U^{t}$, we can see from Eq. 2
that

$\frac{U^{t}}{\alpha_{t}}=\frac{\alpha_{t-1}}{\alpha_{t}}\frac{U^{t-1}}{\alpha_{t-1}}+\frac{s_{t}\nabla
f_{\xi_{t}}(W^{t-1})}{\alpha_{t}}=\left(1-\frac{s_{t}}{\alpha_{t}}\right)\frac{U^{t-1}}{\alpha_{t-1}}+\frac{s_{t}}{\alpha_{t}}\nabla
f_{\xi_{t}}(W^{t-1}).$ (15)

Moreover, the assumptions on $\eta_{t}$ satisfies all the required conditions
of Lemma 1 of (Ruszczyński, 1980). We therefore apply Lemma 1 of (Ruszczyński,
1980) to conclude that

$\frac{U^{t}}{\alpha_{t}}\xrightarrow{\text{a.s.}}\nabla\mathbb{E}_{\xi\sim{\mathcal{D}}}\left[f_{\xi}\left(W^{t}\right)\right]\circ\nabla\mathbb{E}_{\xi\sim{\mathcal{D}}}\left[f_{\xi}\left(W^{t}\right)\right].$
(16)

The update for $\alpha_{t}^{-1}V^{t}$ has a form analogous to Eq. 15, and we
have from Jensen’s inequality that

$\mathbb{E}_{\xi_{t}\sim{\mathcal{D}}}{\left\|{\nabla
f_{\xi_{t}}\left(W^{t-1}\right)}\right\|}^{2}\leq\sqrt{\mathbb{E}_{\xi_{t}\sim{\mathcal{D}}}{\left\|{\nabla
f_{\xi_{t}}\left(W^{t-1}\right)}\right\|}^{4}}\leq\sqrt{C},$

implying that the second moment is also bounded in expectation. We can
therefore also apply Lemma 1 of (Ruszczyński, 1980) to $\alpha_{t}^{-1}V^{t}$
and conclude that

$\frac{V^{t}}{\alpha_{t}}\xrightarrow{\text{a.s.}}\nabla\mathbb{E}_{\xi\sim{\mathcal{D}}}\left[f_{\xi}\left(W^{t}\right)\right].$
(17)

We further notice that the union of two events that happens almost surely is
still an event that happens almost surely.

From Eqs. 4 and 7, we can see that there is a sequence $\\{z_{t}\\}$ such that

$-\left(\frac{V^{t}}{\alpha_{t}}+\frac{z_{t}}{\alpha_{t}}+\frac{P^{t}}{\alpha_{t}}(\hat{W}^{t}-W^{0})\right)\in\partial\psi(\hat{W}^{t}),\quad{\left\|{z_{t}}\right\|}\leq\epsilon_{t}.$
(18)

Our assumption in Eq. 13 implies that $\alpha_{t}\rightarrow\infty$, which
together with Eq. 8 leads to

$\frac{z_{t}}{\alpha_{t}}\rightarrow 0.$ (19)

From Eq. 17, that
$\nabla\mathbb{E}_{\xi\sim{\mathcal{D}}}\left[f_{\xi}\left(W\right)\right]$ is
Lipschitz continuous, and that $W^{t}\rightarrow\bar{W}$ (which we have proven
in the first part), we see that

$\frac{V^{t}}{\alpha_{t}}\xrightarrow{\text{a.s.}}\nabla\mathbb{E}_{\xi\sim{\mathcal{D}}}\left[f_{\xi}(\bar{W})\right].$
(20)

For the third term, we have from Eqs. 16 and 3 that

$\frac{P^{t}}{\alpha_{t}}=\alpha_{t}^{-\frac{2}{3}}\operatorname{Diag}\left(\sqrt[3]{\frac{U^{t}}{\alpha_{t}}}\right)+\frac{\epsilon}{\alpha_{t}}I.$

Again since $\alpha_{t}\rightarrow\infty$, the second term of the equation
above converges to $0$. Therefore, by Eq. 16, we obtain that

$\frac{P^{t}}{\alpha_{t}}\xrightarrow{\text{a.s.}}\alpha_{t}^{-\frac{2}{3}}\operatorname{Diag}\left(\sqrt[3]{\nabla\mathbb{E}_{\xi\sim{\mathcal{D}}}\left[f_{\xi}\left(W^{t}\right)\right]\circ\nabla\mathbb{E}_{\xi\sim{\mathcal{D}}}\left[f_{\xi}\left(W^{t}\right)\right]}\right).$

Again from the continuity of
$\nabla\mathbb{E}_{\xi\sim{\mathcal{D}}}\left[f_{\xi}\left(W\right)\right]$
and that $\alpha_{t}\rightarrow\infty$, we conclude that

$\frac{P^{t}}{\alpha_{t}}\xrightarrow{\text{a.s.}}\alpha_{t}^{-\frac{2}{3}}\operatorname{Diag}\left(\sqrt[3]{\nabla\mathbb{E}_{\xi\sim{\mathcal{D}}}\left[f_{\xi}\left(\bar{W}\right)\right]\circ\nabla\mathbb{E}_{\xi\sim{\mathcal{D}}}\left[f_{\xi}\left(\bar{W}\right)\right]}\right)\xrightarrow{\text{a.s.}}0.$
(21)

Finally, using the outer semicontinuity of $\partial\psi(W)$ at $\bar{W}$, we
conclude from Eqs. 18, 19, 20 and 21 that

$0\in\nabla\mathbb{E}_{\xi\sim{\mathcal{D}}}\left[f_{\xi}\left(\bar{W}\right)\right]+\lim_{t\rightarrow\infty}\psi(\hat{W}^{t})\subseteq\nabla\mathbb{E}_{\xi\sim{\mathcal{D}}}\left[f_{\xi}\left(\bar{W}\right)\right]+\psi(\bar{W})=\partial
F(\bar{W})$

with probability one, showing that $\bar{W}$ is a stationary point almost
surely. ∎

######  Theorem 3.

Consider Algorithm 1 with the conditions in Theorem 2 satisfied. Consider the
event of $\\{\hat{W}^{t}\\}$ converging to a certain point $\bar{W}$ as in
Theorem 2. If the probability of this event is nonzero; $\psi$ is prox-regular
and subdifferentially continuous at $\bar{W}$ and partly smooth at $\bar{W}$
relative to the active $\mathcal{C}^{2}$ manifold ${\mathcal{M}}_{\bar{W}}$;
$\partial\psi$ is outer semicontinuous at $\bar{W}$; and the nondegeneracy
condition

$-\nabla f\left(\bar{W}\right)\in\textup{relative interior
of}\;\partial\psi\left(\bar{W}\right)$ (22)

holds at $\bar{W}$, then conditional on this event, almost surely there is
$T_{0}\geq 0$ such that

$\hat{W}^{t}\in{\mathcal{M}}_{\bar{W}},\quad\forall t\geq T_{0}.$ (23)

In other words, the active manifold at $\bar{W}$ is identified by the iterates
of Algorithm 1 after a finite number of iterations almost surely.

###### Proof.

From Eq. 18, there exists a sequence $\\{Y^{t}\\}$ such that

$Y^{t}\in\partial\psi(\hat{W}^{t}),\quad\frac{V^{t}}{\alpha_{t}}+\frac{z_{t}}{\alpha_{t}}+\frac{P^{t}}{\alpha_{t}}(\hat{W}^{t}-W^{0})+Y^{t}=0,\quad\forall
t.$ (24)

For notational ease, we denote

$f(W)\coloneqq\mathbb{E}_{\xi\sim{\mathcal{D}}}\left[f_{\xi}(W)\right].$ (25)

From Eq. 24, we then get

$\nabla
f(\hat{W}^{t})-\frac{V^{t}}{\alpha_{t}}-\frac{z_{t}}{\alpha_{t}}-\frac{P^{t}}{\alpha_{t}}(\hat{W}^{t}-W^{0})\in\partial
F(\hat{W}^{t}).$ (26)

We aim to show that

$\mbox{\rm dist}(0,\partial F(\hat{W}^{t}))\coloneqq\min_{Y\in\partial
F(\hat{W}^{t})}{\left\|{Y}\right\|}$

converges to $0$ almost surely. From Eq. 26, we have

$\displaystyle\mbox{\rm dist}(0,\partial F(\hat{W}^{t}))$
$\displaystyle\leq{\left\|{\nabla
f(\hat{W}^{t})-\frac{V^{t}}{\alpha_{t}}-\frac{z_{t}}{\alpha_{t}}-\frac{P^{t}}{\alpha_{t}}(\hat{W}^{t}-W^{0})}\right\|}$
$\displaystyle\leq{\left\|{\nabla
f(\hat{W}^{t})-\frac{V^{t}}{\alpha_{t}}}\right\|}+{\left\|{\frac{z_{t}}{\alpha_{t}}}\right\|}+{\left\|{\frac{P^{t}}{\alpha_{t}}(\hat{W}^{t}-W^{0})}\right\|}$
$\displaystyle\leq{\left\|{\nabla
f(\hat{W}^{t})-\frac{V^{t}}{\alpha_{t}}}\right\|}+\frac{\epsilon_{t}}{\alpha_{t}}+{\left\|{\frac{P^{t}}{\alpha_{t}}(\hat{W}^{t}-W^{0})}\right\|},$
(27)

where we get the first inequality from the triangle inequality and the second
from Eq. 18. According to Eqs. 17 and 21, there are $\\{A_{t}\\}$ and
$\\{B_{t}\\}$ such that

$\begin{cases}\frac{V^{t}}{\alpha_{t}}&=\nabla
f(W^{t})+A_{t},\quad{\left\|{A_{t}}\right\|}\xrightarrow{\text{a.s.}}0\\\
\frac{P^{t}}{\alpha_{t}}&=\alpha_{t}^{-\frac{2}{3}}\operatorname{Diag}\left(\sqrt[3]{\nabla
f(W^{t})\circ\nabla
f(W^{t})}\right)+B_{t},\quad{\left\|{B_{t}}\right\|}\xrightarrow{\text{a.s.}}0.\end{cases}$
(28)

Substituting the above two equations back to Eq. 27, we obtain

$\displaystyle~{}\mbox{\rm dist}(0,\partial F(\hat{W}^{t}))$
$\displaystyle\leq$ $\displaystyle~{}{\left\|{\nabla f(\hat{W}^{t})-\nabla
f(W^{t})}\right\|}+{\left\|{A_{t}}\right\|}+\frac{\epsilon_{t}}{\alpha_{t}}+\left(\alpha_{t}^{-\frac{2}{3}}{\left\|{\sqrt[3]{\nabla
f(W^{t})\circ\nabla
f(W^{t})}}\right\|}_{\infty}+{\left\|{B_{t}}\right\|}\right){\left\|{\hat{W}^{t}-W^{0}}\right\|}$
$\displaystyle\leq$
$\displaystyle~{}L{\left\|{\hat{W}^{t}-W^{t}}\right\|}+{\left\|{A_{t}}\right\|}+\frac{\epsilon_{t}}{\alpha_{t}}+\left(\alpha_{t}^{-\frac{2}{3}}{\left\|{\sqrt[3]{\nabla
f(W^{t})\circ\nabla
f(W^{t})}}\right\|}_{\infty}+{\left\|{B_{t}}\right\|}\right){\left\|{\hat{W}^{t}-W^{0}}\right\|}.$
(29)

From Theorem 2, we know that $\hat{W}^{t}$ and $W^{t}$ both converge to
$\bar{W}$, so

${\left\|{\hat{W}^{t}-W^{t}}\right\|}\leq{\left\|{\hat{W}^{t}-\bar{W}}\right\|}+{\left\|{W^{t}-\bar{W}}\right\|}\rightarrow
0.$

From Eqs. 13 and 8, we know that $\epsilon_{t}/\alpha_{t}\rightarrow 0$.
Because $\hat{W}^{t}\rightarrow\bar{W}$, we also have that

${\left\|{\hat{W}^{t}-W^{0}}\right\|}\rightarrow{\left\|{\bar{W}-W^{0}}\right\|}<\infty.$

From $W^{t}\rightarrow\bar{W}$, we have that

${\left\|{\sqrt[3]{\nabla f(W^{t})\circ\nabla
f(W^{t})}}\right\|}_{\infty}\rightarrow{\left\|{\sqrt[3]{\nabla
f(\bar{W})\circ\nabla f(\bar{W})}}\right\|}_{\infty}<\infty.$

Combining these results with Eq. 29, we conclude that

$\displaystyle
L{\left\|{\hat{W}^{t}-W^{t}}\right\|}+{\left\|{A_{t}}\right\|}+\frac{\epsilon_{t}}{\alpha_{t}}+\left(\alpha_{t}^{-\frac{2}{3}}{\left\|{\sqrt[3]{\nabla
f(W^{t})\circ\nabla
f(W^{t})}}\right\|}_{\infty}+{\left\|{B_{t}}\right\|}\right){\left\|{\hat{W}^{t}-W^{0}}\right\|}\xrightarrow{\text{a.s.}}0,$

proving that

$\mbox{\rm dist}(0,\partial F(\hat{W}^{t}))\xrightarrow{\text{a.s.}}0.$

On the other hand, since $f$ is continuous and $\psi$ is subdifferentially
continuous at $\bar{W}$ (which implies $F$ is also subdifferentially contnuous
at $\bar{W}$), $\hat{W}^{t}\rightarrow\bar{W}$, and that $\nabla
f(\hat{W}^{t})+Y_{t}\xrightarrow{\text{a.s.}}0\in\partial F(\bar{W})$ (from
Theorem 2), we know that $F(\hat{W}^{t})\xrightarrow{\text{a.s.}}F(\bar{W})$
as well. Therefore, we can apply Lemma 1 of (Lee, 2023) to conclude that Eq.
23 indeed holds for some $T_{0}<\infty$ almost surely. ∎

####  B.1 Convergence Result for ProxGen

We next discuss the convergence result for the framework of (Yun et al., 2021)
with inexactness added. For consistency, we first use our notations to
introduce their framework, with our inexactness condition added, in Algorithm
3.

$m_{0}\leftarrow 0$

for  _$t=1,\dotsc,T$_ do

Sample $\xi_{t}\sim{\mathcal{D}}$ with batch size $b_{t}$

$G^{t}\leftarrow\nabla f_{\xi_{t}}(W^{t-1})$

$m_{t}\leftarrow\rho_{t}m_{t-1}+(1-\rho_{t})G^{t}$

Construct $P^{t}$ satisfying $P^{t}\succeq\delta I$

$\theta_{t}\leftarrow 1/{\left\|{P^{t}}\right\|}_{2}$

Compute $W^{t+1}$ by roughly solving Eq. 6 that satisfies Eq. 7 with $Q_{t}$
replaced by $\hat{Q}_{t}$ and $\hat{W}^{t}$ replaced by $W^{t+1}$, using PG
$(W^{t},W^{t},m_{t},\eta_{t}^{-1}P^{t},\theta_{t},T_{2},\epsilon_{t})$

output : $W^{T}$

Algorithm 3 ProxGen
$(W^{0},T,T_{2},\\{\eta_{t}\\},\\{\rho_{t}\\},\\{c_{t}\\},\\{\epsilon_{t}\\},\\{b_{t}\\},\delta
I)$

In their analysis, Yun et al. (2021) made the following four assumptions, and
we will follow these assumptions using the notation Eq. 25.

  1. (C-$\bf{1}$)

The expected loss function $f$ is $L$-Lipschitz-continuously-differentiable
and lower-bounded for some $L\geq 0$.

  2. (C-$\bf{2}$)

The stochastic gradient $G^{t}=\nabla f_{\xi_{t}}(W^{t-1})$ is an unbiased
estimator of $\nabla f(W^{t-1})$ with bounded variance.

$E_{\xi_{t}\sim{\mathcal{D}}}[G^{t}]=\nabla f(W^{t-1}),\quad E_{\xi_{t}\sim
D}\left[{\left\|{G^{t}-\nabla
f(W^{t-1})}\right\|}^{2}\right]\leq\frac{\sigma^{2}}{b_{t}},\quad\forall t\geq
0,$

where $b_{t}$ is the batch size of $\xi_{t}$ and $\sigma\geq 0$ is a constant.

  3. (C-$\bf{3}$)

There are some $\rho_{0},\mu\in[0,1)$ and $D,G>0$ such that
${\left\|{W^{t+1}-W^{t}}\right\|}\leq D$, ${\left\|{G^{t}}\right\|}\leq G$,
and $\rho_{t}=\rho_{0}\mu^{t-1}$ for all $t$.

  4. (C-$\bf{4}$)

There is some $\gamma>0$ such that
${\left\|{\eta_{t}^{-1}P^{t}}\right\|}_{2}\leq 1/\gamma<\infty$ for all $t$.

  5. (C-$\bf{5}$)

There is $\delta>0$ such that

$P^{t}\geq\delta,\quad\eta_{t}\leq\frac{\delta}{3L},\quad\forall t\geq 0.$
(30)

######  Theorem 4.

For the framework in (Yun et al., 2021) with the subproblem solved
approximately by Algorithm 2 such that Eq. 7 holds with $\\{\epsilon_{t}\\}$
satisfying Eq. 8. Then Theorem 1 of (Yun et al., 2021) still holds, but with
the constants $\\{Q_{i}\\}$ being also dependent on $\bar{\epsilon}$.

###### Proof.

The major flow of our proof follows that of (Yun et al., 2021) but with
suitable modifications to accommodate the inexactness condition in the
subproblem solving.

It is clear from (Yun et al., 2021, Lemma 1) that

${\left\|{m_{t}}\right\|}\leq G,\quad\forall t\geq 0.$ (31)

By the update rule for $m_{t}$, Eqs. 7 and 6, we have that there is $z_{t}$
such that

$\displaystyle 0\in
z_{t}+(1-\rho_{t})g_{t}+\rho_{t}m_{t-1}+{\partial}\psi(W^{t})+\frac{1}{\eta_{t}}(P^{t})(W^{t}-W^{t-1}),\quad{\left\|{z_{t}}\right\|}\leq\epsilon_{t},\quad\forall
t\geq 0,$

leading to

$\displaystyle\nabla
f(W^{t})-z_{t}-(1-\rho_{t})g_{t}-\rho_{t}m_{t-1}-\frac{1}{\eta_{t}}(P^{t})(W^{t}-W^{t-1})\in{\partial}F(W^{t}).$
(32)

We thus have from Eq. 32 and (C-$4$) that

$\displaystyle\mbox{\rm dist}(0,{\partial}F(W^{t}))^{2}$
$\displaystyle\leq~{}$ $\displaystyle\Big{\|}z_{t}+(1-\rho_{t})g_{t}-\nabla
f(W^{t})+\rho_{t}m_{t-1}+(W^{t}-W^{t-1})+\frac{1}{\eta_{t}}(P^{t})(W^{t}-W^{t-1})-(W^{t}-W^{t-1})\Big{\|}^{2}$
$\displaystyle\leq~{}$ $\displaystyle 4\Big{\|}(1-\rho_{t})g_{t}-\nabla
f(W^{t})+\rho_{t}m_{t-1}+(W^{t}-W^{t-1})\Big{\|}^{2}+4\epsilon_{t}^{2}+4\Big{\|}\frac{1}{\eta_{t}}(P^{t})(W^{t}-W^{t-1})\Big{\|}^{2}$
$\displaystyle+4\Big{\|}(W^{t}-W^{t-1})\Big{\|}^{2}$ $\displaystyle\leq~{}$
$\displaystyle 4\underbrace{\Big{\|}(1-\rho_{t})g_{t}-\nabla
f(W^{t})+\rho_{t}m_{t-1}+(W^{t}-W^{t-1})\Big{\|}^{2}}_{T_{1}}+4\Big{(}\frac{1}{\gamma^{2}}+1\Big{)}\|W^{t}-W^{t-1}\|^{2}+4\epsilon_{t}^{2}.$
(33)

We will separately bound the quantities $T_{1}$ and
${\left\|{W^{t}-W^{t-1}}\right\|}^{2}$ below.

From the subproblem objective requirement in Eq. 7, we get

$\big{\langle}(1-\rho_{t})g_{t}+\rho_{t}m_{t-1},W^{t}-W^{t-1}\big{\rangle}+\psi(W^{t})+\frac{1}{2\eta_{t}}\langle
W^{t}-W^{t-1},\,P^{t}(W^{t}-W^{t-1})\rangle\\\ \leq\psi(W^{t-1}).$ (34)

From (C-$1$), we have

$\displaystyle f(W^{t})\leq f(W^{t-1})+\langle\nabla
f(W^{t-1}),W^{t}-W^{t-1}\rangle+\frac{L}{2}\|W^{t}-W^{t-1}\|^{2}.$ (35)

Summing Eqs. 35 and 34 gives

$\big{\langle}(1-\rho_{t})g_{t}-\nabla
f(W^{t-1})+\rho_{t}m_{t-1},W^{t}-W^{t-1}\big{\rangle}+{\left\|{W^{t}-W^{t-1}}\right\|}^{2}_{\frac{P^{t}}{2\eta_{t}}-\frac{L}{2}I}\\\
\leq F(W^{t-1})-F(W^{t}).$ (36)

Note that $\eta_{t}^{-1}P^{t}-LI\succeq 0$ from Eq. 30, so the second term in
Eq. 36 is nonnegative. Eq. 36 together with (C-$3$) then leads to

$\displaystyle\|W^{t}-W^{t-1}\|_{\frac{P^{t}}{2\eta_{t}}-\frac{L}{2}I}^{2}$
$\displaystyle\leq~{}$ $\displaystyle
F(W^{t-1})-F(W^{t})-\big{\langle}g_{t}-\nabla
f(W^{t-1}),W^{t}-W^{t-1}\big{\rangle}+\langle\rho_{t}g_{t},W^{t}-W^{t-1}\rangle-\langle\rho_{t}m_{t-1},W^{t}-W^{t-1}\rangle$
$\displaystyle\leq~{}$ $\displaystyle
F(W^{t-1})-F(W^{t})+\frac{1}{2L}\|g_{t}-\nabla
f(W^{t-1})\|^{2}+\frac{L}{2}\|W^{t}-W^{t-1}\|^{2}+\frac{\rho_{t}^{2}}{2L}\|g_{t}\|^{2}+\frac{L}{2}\|W^{t}-W^{t-1}\|^{2}$
$\displaystyle\quad+\rho_{t}\|m_{t-1}\|\|W^{t}-W^{t-1}\|$
$\displaystyle\leq~{}$ $\displaystyle
F(W^{t-1})-F(W^{t})+\frac{1}{2L}\|g_{t}-\nabla
f(W^{t-1})\|^{2}++L\|W^{t}-W^{t-1}\|^{2}+\frac{\rho_{0}^{2}\mu^{2(t-1)}G^{2}}{2L}+\rho_{0}\mu^{t-1}DG.$

Summing it over $t=1,2,\dotsc,T$ and utilizing the assumption that the step
sizes are non-increasing then give

$\displaystyle\Big{(}\frac{\delta}{2\eta_{0}}-\frac{3}{2}L\Big{)}\sum\limits_{t=1}^{T}\|W^{t}-W^{t-1}\|^{2}\leq\Delta+C_{1}+\frac{1}{2L}\sum\limits_{t=1}^{T}\|g_{t}-\nabla
f(W^{t-1})\|^{2},$

where

$\Delta\coloneqq F(W^{0})-\min_{W}F(W),\quad
C_{1}\coloneqq\frac{\rho_{0}DG}{1-\mu}+\frac{\rho_{0}^{2}G^{2}}{2L(1-\mu^{2})}.$

From the inequality above, we obtain

$\displaystyle\sum\limits_{t=1}^{T}\|W^{t}-W^{t-1}\|^{2}\leq
H_{1}+H_{2}\sum\limits_{t=1}^{T}\|g_{t}-\nabla f(W^{t-1})\|^{2}$ (37)

for some constants $H_{1},H_{2}$ depending on $L,\Delta,\delta,\eta_{0}$, and
$C_{1}$. From Eq. 36, we also have

$\displaystyle\Big{\langle}(1-\rho_{t})g_{t}-\nabla
f(W^{t})+\rho_{t}m_{t-1},W^{t}-W^{t-1}\Big{\rangle}$ $\displaystyle\leq~{}$
$\displaystyle F(W^{t-1})-F(W^{t})-\big{\langle}\nabla f(W^{t})-\nabla
f(W^{t-1}),W^{t}-W^{t-1}\big{\rangle}-\big{\|}W^{t}-W^{t-1}\|_{\frac{1}{2\eta_{t}}(P^{t})-\frac{L}{2}I}^{2}$
$\displaystyle\leq~{}$ $\displaystyle F(W^{t-1})-F(W^{t})-\big{\langle}\nabla
f(W^{t})-\nabla f(W^{t-1}),W^{t}-W^{t-1}\big{\rangle}.$

Therefore, we obtain

$\displaystyle T_{1}$ $\displaystyle=\|(1-\rho_{t})g_{t}-\nabla
f(W^{t})+\rho_{t}m_{t-1}\|^{2}+\|W^{t}-W^{t-1}\|^{2}+2\Big{\langle}(1-\rho_{t})g_{t}-\nabla
f(W^{t})+\rho_{t}m_{t-1},W^{t}-W^{t-1}\Big{\rangle}$
$\displaystyle\leq\|(1-\rho_{t})g_{t}-\nabla f(W^{t-1})+\nabla
f(W^{t-1})-\nabla f(W^{t})+\rho_{t}m_{t-1}\|^{2}+\|W^{t}-W^{t-1}\|^{2}$
$\displaystyle\quad+2\left(F(W^{t-1})-F(W^{t})-\big{\langle}\nabla
f(W^{t})-\nabla f(W^{t-1}),W^{t}-W^{t-1}\big{\rangle}\right)$
$\displaystyle\leq 4\|g_{t}-\nabla
f(W^{t-1})\|^{2}+4L^{2}\|W^{t}-W^{t-1}\|^{2}+4\rho_{t}^{2}(\|m_{t-1}\|^{2}+\|g_{t}\|^{2})+\|W^{t}-W^{t-1}\|^{2}$
$\displaystyle\quad+2\left(F(W^{t-1})-F(W^{t})+L\|W^{t}-W^{t-1}\|^{2}\right)$
$\displaystyle\leq
2\left(F(W^{t-1})-F(W^{t})\right)+8\rho_{0}^{2}\mu^{2(t-1)}G^{2}+\Big{(}1+2L+4L^{2}\Big{)}\|W^{t}-W^{t-1}\|^{2}+4\|g_{t}-\nabla
f(W^{t-1})\|^{2}.$ (38)

Let $C_{2}\coloneqq 2+2L+4L^{2}+\gamma^{-2}$ and insert Eq. 38 into Eq. 33, we
get

$\displaystyle\mbox{\rm dist}(0,{\partial}F(W^{t}))^{2}$
$\displaystyle\leq~{}$ $\displaystyle
4\Bigg{(}2\left(F(W^{t-1})-F(W^{t})\right)+8\rho_{0}^{2}\mu^{2(t-1)}G^{2}+{C_{2}}\|W^{t}-W^{t-1}\|^{2}+4\|g_{t}-\nabla
f(W^{t-1})\|^{2}+\epsilon_{t}^{2}\Bigg{)}.$ (39)

Therefore, we have from Eqs. 39 and 8 and (C-$2$) that

$\displaystyle~{}\mathbb{E}_{a,\xi_{1},\dotsc,\xi_{T}}[\mbox{\rm
dist}(0,{\partial}F(W^{a}))^{2}]$ $\displaystyle\leq$
$\displaystyle~{}\frac{1}{T}\sum_{t=1}^{T}\mathbb{E}\Big{[}\big{\|}(1-\rho_{t})g_{t}-\nabla
f(W^{t})+z_{t}+\rho_{t}m_{t-1}+\frac{1}{\eta_{t}}(P^{t})(W^{t}-W^{t-1})\big{\|}^{2}\Big{]}$
$\displaystyle\leq$
$\displaystyle~{}\frac{4}{T}\Big{(}2\Delta+\frac{8\rho_{0}^{2}G^{2}}{1-\mu^{2}}+4\sum_{t=1}^{T}\mathbb{E}\|g_{t}-\nabla
f(W^{t-1})\|^{2}+C_{2}\sum_{t=1}^{T}\mathbb{E}\|W^{t}-W^{t-1}\|^{2}+\sum_{t=1}^{T}\epsilon_{t}^{2}\Big{)}$
$\displaystyle\leq$
$\displaystyle~{}\frac{4}{T}\Big{(}2\Delta+\frac{8\rho_{0}^{2}G^{2}}{1-\mu^{2}}+4\sigma^{2}\sum_{t=1}^{T}\frac{1}{b_{t}}+C_{2}(H_{1}+H_{2}\sigma^{2}\sum_{t=1}^{T}\frac{1}{b_{t}})+\bar{\epsilon}\Big{)}$
$\displaystyle\leq$
$\displaystyle~{}\frac{Q_{1}}{T}\sum_{t=1}^{T}\frac{1}{b_{t}}+\frac{Q_{2}\Delta}{T}+\frac{Q_{3}}{T},$

for some constants $Q_{1},Q_{2},Q_{3}$ dependent on
$\\{\eta_{0},\delta,\Delta,L,D,G,\rho_{0},\mu,\gamma,\bar{\epsilon}\\}$, but
not on $T$. This proves our theorem. ∎

###  Appendix C Additional Experiments for Computer Vision

In this section, we compare RAMDA with other methods on image classification
with smaller datasets. They are:

  1. 1.

Logistic regression (neural network with no hidden layer) with the MNIST
dataset (LeCun et al., 1998).

  2. 2.

A modified VGG19 (Simonyan & Zisserman, 2015) with the CIFAR10 dataset
(Krizhevsky, 2009).

  3. 3.

The same VGG19 with the CIFAR100 dataset (Krizhevsky, 2009).

  4. 4.

A modified ResNet50 (He et al., 2016) with the CIFAR10 dataset.

  5. 5.

The same ResNet50 with the CIFAR100 dataset.

The results are shown in Table 7. In the logistic regression problem, we only
perform a single run, with the initial point being the origin, as it is a
convex problem. Moreover, in this problem, when dealing with ProxSSI, ProxGen,
and ProxSGD whose sparsity levels are highly unstable over iterations, we
report their highest weighted group sparsity over all epochs, but for all
other problems, we report the group sparsity level of the final output.

Table 7: Group sparsity and validation accuracy of different methods on image classification with smaller datasets. Algorithm | Validation accuracy | Group sparsity
---|---|---
Logistic Regression/MNIST
ProxSGD | 91.31% | 39.29%
ProxSSI | 91.31% | 39.92%
ProxGen | 91.31% | 39.92%
RMDA | 91.34% | 57.02%
RAMDA | 91.35% | 57.40%
VGG19/CIFAR10
MSGD | 93.95 $\pm$ 0.14% | -
ProxSGD | 92.82 $\pm$ 0.09% | 82.76 $\pm$ 5.42%
ProxSSI | 92.81 $\pm$ 0.15% | 88.40 $\pm$ 0.23%
ProxGen | 92.83 $\pm$ 0.05% | 86.64 $\pm$ 0.12%
RMDA | 93.13 $\pm$ 0.10% | 90.22 $\pm$ 0.06%
RAMDA | 92.89 $\pm$ 0.13% | 86.31 $\pm$ 0.31%
VGG19/CIFAR100
MSGD | 74.07 $\pm$ 0.05% | -
ProxSGD | 71.96 $\pm$ 0.15% | 72.34 $\pm$ 11.9%
ProxSSI | 67.29 $\pm$ 0.06% | 78.58 $\pm$ 0.34%
ProxGen | 68.13 $\pm$ 0.36% | 75.46 $\pm$ 0.17%
RMDA | 71.96 $\pm$ 0.31% | 80.88 $\pm$ 0.11%
RAMDA | 70.47 $\pm$ 0.25% | 65.19 $\pm$ 0.77%
ResNet50/CIFAR10
MSGD | 95.54 $\pm$ 0.19% | -
ProxSGD | 92.36 $\pm$ 0.05% | 82.18 $\pm$ 2.67%
ProxSSI | 94.04 $\pm$ 0.12% | 83.67 $\pm$ 0.63%
ProxGen | 94.07 $\pm$ 0.12% | 80.45 $\pm$ 0.45%
RMDA | 95.11 $\pm$ 0.11% | 85.64 $\pm$ 0.12%
RAMDA | 93.85 $\pm$ 0.10% | 81.99 $\pm$ 1.26%
ResNet50/CIFAR100
MSGD | 79.49 $\pm$ 0.49% | -
ProxSGD | 74.54 $\pm$ 0.58% | 49.29 $\pm$ 5.91%
ProxSSI | 73.65 $\pm$ 0.39% | 70.38 $\pm$ 0.74%
ProxGen | 73.63 $\pm$ 0.43% | 65.51 $\pm$ 3.58%
RMDA | 75.62 $\pm$ 0.19% | 79.97 $\pm$ 0.27%
RAMDA | 69.23 $\pm$ 0.86% | 68.65 $\pm$ 1.83%

###  Appendix D Plots of Sparsity Level and Validation Accuracy over Epochs

We provide in Fig. 1 the plots of predictive ability and structured sparsity
over epochs of some representative experiments we have conducted. These
experiments are:

  1. 1.

ResNet50 with the ILSVRC 2012 ImageNet dataset.

  2. 2.

Transformer-XL with the WikiText-103 dataset.

  3. 3.

Tacotron2 with the LJSpeech dataset.

  4. 4.

Logistic Regression with the MNIST dataset.

  5. 5.

A modified VGG19 with the CIFAR10 dataset.

  6. 6.

The same VGG19 with the CIFAR100 dataset.

  7. 7.

A modified ResNet50 with the CIFAR10 dataset.

  8. 8.

The same ResNet50 with the CIFAR100 dataset.

In the plot for Transformer-XL, one step processes ten batches, and for our
batch size of 64, one epoch consists of 8,401 batches. We further observe in
the zoomed-in sparsity plots in Fig. 2 that the sparsity level of RAMDA is
stable at the final epochs. These plots corroborates our theory that RAMDA is
indeed capable of manifold identification while achieving competitive
prediction performance. On the other hand, in the absence of manifold
identification guarantees, the sparsity levels of ProxSGD, ProxSSI and ProxGen
exhibit oscillations that are sometimes drastic. We note that for the largest
problems Tacotron2 and Transformer-XL, the sparsity levels of RAMDA are still
gradually increasing even at the final epochs. This suggests that if we are
willing to run the algorithm for longer, it is possible that the structured
sparsity level could be further improved.

(a) ResNet50 on ImageNet

(b) Transformer-XL on WikiText-103

(c) Tacotron2 on LJSpeech

(d) Logistic Regression on MNIST

(e) VGG19 on CIFAR10

(f) VGG19 on CIFAR100

(g) ResNet50 on CIFAR10

(h) ResNet50 on CIFAR100

Figure 1: Group sparsity level and validation prediction performance v.s
epochs. In the plot for Transformer-XL, one step processes ten batches, and
for our batch size of 64, one epoch consists of 8,401 batches.

(a) ResNet50 on ImageNet

(b) Transformer-XL on WikiText-103

(c) Tacotron2 on LJSpeech

(d) Logistic Regression on MNIST

(e) VGG19 on CIFAR10

(f) VGG19 on CIFAR100

(g) ResNet50 on CIFAR10

(h) ResNet50 on CIFAR100

Figure 2: Group sparsity level at the last epochs.

###  Appendix E Experiment with Nuclear-norm Regularization

We further conduct some preliminary experiments with a different regularizer
to showcase that the proposed RAMDA can be applied to structures beyond
sparsity. We consider the structure such that each layer of the neural network
is low-rank, induced by imposing one nuclear-norm regularizer per layer
individually by treating each layer as a matrix. Given a matrix
$X\in\mathbb{R}^{m\times n}$ of rank $r\leq\min\\{m,n\\}$ with its singular
value decomposition (SVD) $X=U\Sigma V^{\top}$, where $U\in\mathbb{R}^{m\times
r}$, $V\in\mathbb{R}^{n\times r}$ are orthogonal and the positive definite
diagonal matrix $\Sigma\in\mathbb{R}^{r\times r}$ represents the nonzero
singular values of $X$, the nuclear norm of $X$ is computed by

${\left\|{X}\right\|}_{*}=\sum_{i=1}^{r}\Sigma_{i,i},$

and the corresponding proximal operator for $\lambda>0$ is

$\mbox{\rm
prox}_{\lambda{\left\|{\cdot}\right\|}_{*}}(X)=U\hat{\Sigma}V^{\top},\text{
where }\hat{\Sigma}_{i,i}=\max\\{0,\Sigma_{i,i}-\lambda\\},\quad
i=1,\dotsc,r.$

Given a point $X^{*}$ with rank $r^{*}$, the active manifold of the nuclear
norm at $X^{*}$ is

$\mathcal{M}(X^{*})=\\{Y\mid\text{rank}(Y)=r^{*}\\}.$

Using low-rankness to condense neural networks is itself an interesting
research topic, but conducting full SVDs could be rather time-consuming, so
applying this structure to larger problems is challenging but potentially
useful. How to exploit this structure for prediction acceleration and to make
the training more efficient, possibly using iterative methods to compute
approximate SVDs, is an interesting topic we plan to investigate in the near
future. Instead, the purpose of the preliminary experiment here is merely for
showing that our method is also applicable to other structures.

We first consider training a simple neural network with six fully-connected
layers using the FashionMNIST dataset (Xiao et al., 2017). Since this is a
rather small-scale problem and this is a image classification problem, we do
not expect RAMDA to outperform non-adaptive methods, especially the RMDA
method that is also able to identify the active manifold. The goal of this
experiment is just to demonstrate the possibilities of structures beyond
sparsity. The results are shown in Table 8. As we have anticipated, RAMDA is
indeed slightly worse than RMDA regarding the low-rank level and the
prediction accuracy, but it is still competitive and outperforms ProxGen and
ProxSGD. This exemplifies the potential of RAMDA as well as RMDA for training
neural networks with other useful structures.

We also conduct an experiment on pretraining a modified vision transformer
model (Liu et al., 2021) for masked image modeling (Xie et al., 2022) using
the CIFAR10 dataset. Following the standard practice of this task, we select
the model that gives the lowest validation loss among the last 50 epochs as
the final output. The results are shown in Table 9. We can see that RAMDA
attains the lowest validation loss and has a low-rank level almost identical
to that of RMDA. On the other hand, ProxSGD and ProxGen have worse low-rank
levels.

Table 8: Low-rank level and validation accuracy of different methods on training a six-layer fully-connected neural network with the FashionMNIST dataset for image classification. Algorithm | Validation accuracy | Low-rank level
---|---|---
MSGD | 89.95 $\pm$ 0.29% | -
ProxSGD | 87.54 $\pm$ 0.52% | 78.00 $\pm$ 0.77%
ProxGen | 86.66 $\pm$ 0.33% | 87.46 $\pm$ 4.19%
RMDA | 88.19 $\pm$ 0.23% | 91.88 $\pm$ 0.12%
RAMDA | 87.99 $\pm$ 0.24% | 89.59 $\pm$ 0.42%
Table 9: Low-rank level and validation loss of different methods on pretraining a modified vision transformer model using the CIFAR10 dataset for masked image modeling. Algorithm | Validation loss | Low-rank level
---|---|---
AdamW | 0.0865 $\pm$ 0.0001 | -
ProxSGD | 0.1042 $\pm$ 0.0003 | 82.60 $\pm$ 0.34%
ProxGen | 0.1120 $\pm$ 0.0019 | 82.64 $\pm$ 2.47%
RMDA | 0.1054 $\pm$ 0.0031 | 86.23 $\pm$ 0.41%
RAMDA | 0.1035 $\pm$ 0.0016 | 86.20 $\pm$ 0.35%