Regularized Adaptive Momentum Dual Averaging with an Efficient Inexact Subproblem Solver for Training Structured Neural Network
Authors: Zih-Syuan Huang
Ching-pei Lee
Venue: NeurIPS 2024 poster (2024)
Keywords: structured neural networks
variance reduction
manifold identification
proximal methods
adaptive methods
inexact subproblem solution
Primary area: optimization_for_deep_networks
TL;DR: A regularized adaptive momentum dual averaging algorithm for training structured neural networks with guarantees for finding the locally optimal structure, and exhibits outstanding performance in language, speech and image classification tasks.

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.

Regularized Adaptive Momentum Dual Averaging with an Efficient Inexact Subproblem Solver for Training Structured Neural Network
Zih-Syuan Huang
Department of Computer Science and Information Engineering National Taiwan University
Taipei 106, Taiwan r119222108ntu. edu. tw
Ching-pei Lee Department of Advanced Data Science Institute of Statistical Mathematics Tachikawa, Tokyo 190-8562, Japan chingpei Sorgpei@ism.ace ac.
Abstract
We propose Regularized Adaptíve Momentum Dual Averaging (RAMDA) algorithm for training structured neural networks. Similar to existing regular- ized adaptive methods, the subproblem for computing the update direction of RAMDA involves a nonsmooth regularizer and a diagonal preconditioner, and therefore does not possess closed-form solution in general. We thus also carefully devise an implementable inexactness condition that retains convergence guarantees similar the exact versions, and propose compan- ion efficient solver for the subproblems of both RAMDA and existing methods make them practically feasible. We leverage the theory of manifold iden- tification in variational analysis to show that, even in the presence of such inexactness, the iterates of RAMDA attain the idcal 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 samee point, making the first regularized adaptive method outputting models that possess outstanding predictive performance while being (lo- cally) 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::// h//ps://www.... .i..ggii rithub..o con//Lsmoptgroup/RAA....
1 Introduction
Since the recent emergence of ChatGPT, large language models (LLMs) and other huge deep learning models have garnered much attention and popularity, who are unfamiliar with machine learning. 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
among the public
38th Conference on Neural Information Processing Systems (NeurIPS 2024).

---

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, with or without a diagonal preconditioner for adaptiveness, to induce pre-specified type of desirable structure in the final model 40 51 g
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, [G] oberserd empirically that the structures obtained by those hods are highly suboptimal and unstable over iterations. They then proposed 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 H 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-scaat 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 stato-of-the-art prediction performance, including the transformer +d and the LSTM 15 The transformer is also
gaining prominence in computer vision for achieving exceptional performance [29]. Therefore, it is becoming increasingly important to devise methods that attain satisfactory performance for training these network archilectures with structure. For such modern architectures, adaptive methods like Adam 20 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-an [T0, I 52 28 22]. It is hence expected that the adaative RMDA of [6 might not lead promising results for such widely-used architectures and tasks.
This work aims to fill this gap to propose practical regularized adaptive method with guar- antees 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 adapliveness. However, this seemingly easy extension actually requires deliberation in two aspects. First, except for few exceptions, combination of even a simple diagonal precon- ditioner 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-wv wise divisions. Therefore, in the regularized case, the best we can hope for 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 bet the iexxctnesss and the preconditioner makes il particularly difficult to address the following three challenges simultaneously. 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 M. We are dealing with both a proconditioner that changes the whole algorithm, just like from SGD to Adagrad, and the inevitable incxact subproblem solutions that could nullify many useful properties (regard- ing the subdifferential) commonly used in convergence proofs. (ii) Structure: Thoorctical guarantees for structure identification is another critical aim of this work. Inexactness alone already makes this goal difficult; see Example 1 of 24 for a simple instance such that even
(i) Convergence: Proving convergence of a new algorithm
Sse the first paragraph in Section 1 and Appendix B of 16 for a discussion about why the structure at the point of convergence is locally optimal.
2

---

infinitesimal inexactness could hinder structure identification. Even without inexactness, finding a preconditioner that leads 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 loverage 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 37 [6, 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 finite approach
sum. We thus draw inspirations from |2 T6 to consider a stual-averagingtty A7 with momentum to attain variance reduction in this setting for the stochastic gradient estimation. However, we also need variance reduction for the preconditioner, 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-nont peeeanttoomm-baae 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 safe safeguard the preconditioner to keep the subproblems well-conditioned and the computation simple. We then propose to solve the subproblem by 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. An adaptive algorithm for finding locally optimal structures: RAMDA is the first rogularized 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 adaplive methods.
Efficientt subproblem solver for regularized adaptive methods: We propose an implementable inexactness condition and a companion 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 key step for realixing existing frameworks for regularized adaptive methods.
3. A method with outstanding empirical performance: Experiments on training mod- ern neural networks in computer vision (ImageNet), language modeling (Transformer-Xie and speech (Tacotron2) with structured sparsity show that RAMDA steadily outper- forms state of the art by achieving higher structured sparsity ratio and better prediction performance simultanconsly.
2 Related Work
Dual Averaging for Deep Learning. Our method is motivated by [6 that adapted 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 idcas from the empirically successful MADGRAD of 7 for training non-rogularized neural networks. RAMDA can thus also be scen as a gencralization 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, 7 only analyzed convergence rates of the objective value when the problem is convex. Our analysis of (i) variance reduction in the preconditioner, (ii) convergence the nonconvex nonsmooth regularized case, and (iii) structure identification guarantees are novel and closer
the famous regularized dual averaging 17 25
3

---

to properties desirable in practice. The first two items are also applicable when no regularizer present, so our theory also expands guarantees for MADGRAD.
Regularized Stochastic Algorithms for Deep Learning. Other than RMDA, there are several works training structured neural networks through regularization and its proximal operator, but none have structure guarantocs. 4 considered a simple regularized SGD method with momentum, but their convergence analysis is only for the nonadaptive case. 5 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 small class of problems. | proposed ProxSSI that extends ProxGen to the case of group-sparsity regularizers, whose corresponding subproblem indeed has no closed-form solution. They applied the subproblem solutions, and proposed a scemingly mild inexactness condition. 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 struct On the other hand, we will show that with our implement to able inexactness condition, ProxGen still possesses the same convergence guarantees in [5 without any additional requirement on the regularizer. Moreover, we will see in [Section 6 that the time cost of the subproblem solver of ProxSSI is prohibitively high.
method to obtain nearly-optimal
Structure and Manifold Identification. The major tool for our structure guarantees is 24 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 sequence of points converges to this stalionary point with their corresponding subgradients decreasing to zero, this sequence is guaranteed to eventually stay in this manifold, thus identifying the structure. 25
the thoory of manifold identification [2 13 27
37 16 have leveraged
this tool to show manifold identification for various stochastic algorithms, and the common key, as pointed out by [7, is variance reduction. Our analysis uses result given in 40 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.
minweg F (W) == E&-D [fs (W)) + + (W),
(1)
where E a Euclidean space with inner product () and its induced norm 11-11. D
distribution over space Ω representing all possible data modifications, f& is differentiable
almost everywhere for any E, and the possibly nonsmooth regularizer (WW) is for promoting desirable structure in the optimal solutions.
Our algorithm can be seen as a double-dual averaging method that incorporates momentum, a proximal operation for the regularization, and dual averaging for updating both the stochastic gradient estimation and the preconditioner. For ease of description, we assume without loss of generality that & = IR" in this section. At the tth iteration with learning rate T) and iterate W we first draw an independent and identically distributed sample & ~ DD compute the stochastic (sub)gradient G == Vfe, (WV-1) of the loss function at the current point W with respect to & and then update the weighted sum Vú of historical stochastic gradients and the weighted sum U, of their squared norms using the value sy
Vo =0 0.
V== s.G", U. U--++++ Gt G.
Vt > 0, Vt > 0,
8t = ===i
(2)
Uo = 0.
where o denotes the Hadamard (pointwise) product in E. We then construct the preconditioner pt and the weight. sum ae by
pt = Diag( V0 + +), a = ===
8k:
(3)
4

---

Algorithm 1 RAMDA (WO, T, e, (m)))a)))
y <- 0, U < 0, for t = = I...., do Sample & D, Compute V U by (2) and construct P by (3) and 8, t
0g 4-
B.V, a₁ t 00-1 St: G VJe. (W
8, t-
Compute W
by PG (WL-1 Wo, at VV. ar pt, ar0, (2)
Update W by
output:W" W
where >> O is a (usually small) constant for numerical stability and Diag(-) 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.
Wh arg min (Q+(W)= aw)(W) (V', w) + ( w P'(W Wo)),, W
(4)
where W/O is the initial point. Details regarding
(4) and how solve are deferred to
Section 4. The iterate then updated by averaging W and W with some o E [0, 1]:
W = (1- -1-e) w +
The choice of ps in (3) that uses the accumulated square the stochastic gradient norm as the preconditioner the key to adaptivity and is widely seen in adaptive methods such as Adagrad [1,, while the choice of the cubic root instead the square root is motivated by the impressive numerical performance of MADGRAD of Y for smooth problems without
regularization term. The averaging step in (5) with C * can be interpreted as incorporating
momentum term in the non-regularized non-adaptive case 1 19
4 Subproblem Solver
Given an ilerate W a momentum term my a preconditioner pt, and stepsize Whs existing regularized adaptive stochastic gradient algorithms for (1) can be summarized in the following form 51
1 ( - Wt-1 (W----1 P'(W - W ((W)), W = arg mig min (Q.(W)= = (me, W) + 2m,
(6)
W
whose form is similar to 4 When the preconditioner P is multiple of the identity matrix like in the case of T6 the exact subproblem solution of (4) can be efficiently computed through the proximal operator associated with the regularizer. However, major difficulty for realizing regularized adaptive methods, including the proposed RAMDA and the framework of 51 whose preconditioners are not 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., 33 to approxi- mately solve the subproblems in [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
we require
min sEOQ.(W*)
||||| S ft, Q((W))) & (W++-))).
for some pre-specified ét, where (W++1) is the (limiting) subdifferential [see, e.g., 38
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 (6] we enforce (7) but
with Qu replaced by Q, and W by Wt. We focus on the case of diagonal and positive pt
5

---

Algorithm 2 PG(Z", VV,0 (2
if v is nonconver then 0 - 8/2
for j j= ....., In do
Z & (2)-1 1 0(V + P(Z WD))
if ( holds with et = & and W = Z then 7 t- 22, and break
output: 2
and thus the largest eigenvalue max(diag(P")) where diag(-) is the vector formed by the diagonal entries of the input matrix, can be calculated easily and used to compute a step
size guaranleeing sufficient objective decrease. For cases in which this value is difficult to obtain, one can apply a simple racking linesearch for the subproblem to find a suitable step size efficiently. This PG subproblem solver is summarized in Algorithm To guarantee convergence for both our algorithm and the framework of 51 our analysis in Section requires that } satisfy
We will show in Section that
( holds after at most O(G²)) iterations of Algorithm 2
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 [6] for the framework of 5T. All proofs are in the appendices. Some of our results are inspired by [6] but with the added incxactness in (4) and the adaptiveness for the preconditioner, the analysis nontrivial. Recall that we assume that fe is differentiable only almost everywhere but not everywhere, which conforms with widely-used network structures like ReLU-type activations.
We first show that (7) can be attained by Algorithm and that the point of convergence of RAMDA is almost surely stationary point.
Theorem 1. Assume that (4) and
(6) has at least one optimal solution with a finite optimal objective value. Given e, 0, the number of iterations Algorithm 2 takes to satisfy [7) for
both (4) and (6) O(log(e,
when v/ is convEz and O ") when ib is nonconver.
Theorem 2. Consider [
generated by Algorithm for with (7) and {ct} and {a} Assume there is LL2 O such that for any 5. Je is almost surely --Lipschitzcommtinouuuyy dispaaooosssudfffeettttie freettabbl,,, so the expectation is also L-Ligyschitz-coniinnon
satisfying E = xo and (8)
differentiable, there is 2 0 such that Eg,-D|| (W+-1)]l* C for all t, and that the set of stationary points = {W 00 (F(W)} is nonempty. For any given W consider the event that {W*] converges to a point W (each event corresponds to a different W). If du is outer semicontinuous at W, this event has a nonzero probability, and {m satisfies
-sa
Σ(siar))* < oc, co,||wttt1wwi wI (s(a,)))) -
-e, 0.
then we have that W E with probability one conditional on this event. Moreover, (W') also converges this stationary point W.
Usually, convergence to a point requires some further regularity conditions like the Kurdyka- Lojasiewicz condition and boundedness of the iterates. However, existing frameworks regarding iterates convergence using such conditions also require the method analyzed to have mubgrradient-deecent-iikk behavior and to be descent algorithm. Neither of these hold true even for the basic stochastic gradient algorithm, and we leave the analysis for this part as 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 W. For this end, we first need introduce the notions of partial smoothness and prox-regularity, and impose these assumptions & at W.

---

Definition 1 (Partial Smoothness 26 12
A function vp is partly smooth ot point W
relative to set Mw 3 W if
1. Around W, .Mw is a C2 --manifold and llM is C².
2. vb regular (finite and the Fréchet subdifferential coincides with the limiting Préchet subdifferential) at all points W E MMw near W, with (bb(W) + 0. 3. The offine spon of @w((W is a translate of the normal space to Mw at W.
4. Ob is continuous at W relative to Mw
We often call My the active manifold W. Locally, this manifold represents all points near W that share the same structure induced by the regularized as W. Therefore, finding the active manifold is equivalent. to finding the locally optimal structure.
Definition 2 (Prox-regularity 36 A function vp is prox-regular al W for V E (WPWW) if db is locally lower semi-continuous around W, finäte at W, and there is > 0 such that ((W)) 2 +bW2)+++ (V, W - -2) & W2ll for every W1, W2 near W with (V22) close to (WW) and V (v2) close to V*. v/ is proz-regular at W if it is proz-regular for all V E (W).
Theorem 3. Consider Algorithm I with the conditions in Theorem 2 satisfied. Consider the event of (W') converging lo cerlain point W as in Theorem 2 If the probabilily of this event is nonzero; vb is prox-regular and subdifferentially continuous at W and partly smooth at W relative to the active C² manifold Mw; Ouj is outer semicontinuous at W; and the nondegeneracy condition 5 (W) E relative interior of Orb (W) holds at W, then conditional on this event, almost surely there is To 2 0 such that
W MW 22 To-
We note particularly that convex and weakly-convex regular, and subdifferentially continuous everywhere.
34 functions are all regular, prox-
We also show that our subproblem solver and condition can be effectively applied to the framework of [I while retaining the same convergence guarantees. As mentioned in ineetion 2 our result is much stronger than that of []] for having no unrealistic smoothness requirement on & and using implementable inexactness condition.
Theorem 4. For the framework in [oT) with the subproblem solved approrimately by Algo-
rillem 2) such that (7) holds with {ee] satisfying [8] Then Theorem of [5 still holds, bul
with the constunts {Q:) being also dependent
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 45 to consider structured sparsity as the representative structure our experiments. Particularly,
employ the group LASSO regularization bi to encourage group sparsity and disable weight decay in all experiments, except for the dense baselines. We begin from examining the efficiency and effectiveness of Pf for both RAMDA and existing regularized adaptive methods. We then consider tasks computer vision, language modeling, and speech to compare the following algorithms using Pytorch.
RAMDA: The proposed Algorithm
RMDA 16
ProxSGD
7U
ProxGen
We follow their experiments to use AdamW and apply our PG as the subproblem solver.
51 :
ProxSSI
These algorithms are introduced in Section 2 and also further sunmarized in Appendix A For each task, we also provide for relerence a baseline that does not include a group LASSO regularizer in the training (SGD with momentum (MSGD) for computer vision, and Adam

---

Table Weighted group sparsity and validation accuracy of different subproblem stopping
criteria.
No early stopping
Early stopping
/ata
Algorithm
Accuracy
Sparsity
Accuracy
Sparsity
VGG19 /
ProxGen
92.7 + 0.2%
88.8 + 0.0%
92.7 = 0.1%
86.9 _ 0.4%
CIFAR10
RAMDA
92.7 + 0.2%
86.7 + 0.3%
92.9 + 0.2%
86.3 + 0.4%
Re
ProxGen
73.6 + 0.1%
74.7 + 0.6%
74.0 = 0.1%
67.6 3.1%
CIFAR100
RAMDA
69.9 + 1.5%
69.5 + 2.1%
71.2 + 1.4%
67.5 + 1.6%
for the other two), but our comparison only amoug those for training structured models. Our code for reproducing the experiments and the hyperparameter settings are available at aattps://githut ctt/coommoo////e tessmoptgrouprranca
Additional details of the stability of the structure (level of structured sparsity here) over epochs of RAMDA are available in Appendix D
Model predictive ability, and 2. Structured sparsity level. The former is task-dependent and thus specified in each experiment. Regarding the latter, sparsifying neural networks while preserving its performance requires prior knowledge of model design. common approach is retaining certain parameters during the training process, and we adhere to this convention such that the bias, batch normalization [17], layer normalization [ and embedding layers do not have any sparsity-inducing regularization imposed on them p 35 For the rest, we adopt channel-wise grouping for convolutional layers, input-wise grouping for fully-connocted and LSTM layers during the training phase. For evaluation, our structured sparsity is caleulated using the weighted group sparsity with the weights proportional to the mumber of parameters in each group.
We use two criteria for comparison: 1.
We run each experiment with three different random initializations and show the mean and standard deviat ion of the validation predictive performance and the structured sparsity of the final model of all methods.
Subproblem We start from showing the effectiveness of our proposed subproblem solver for RAMDA and ProxGen. For both approaches, use Theorem 2 of g 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 (Q.(Zi-1)--Q) ((QQ(2))(((((((()) +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 carly stopping does not affect the outcome much. Given that early stopping more efficient, we will adopt it all subsequent experiments.
Next, 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 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.
solver is around 3X
Image Classification We conduct classical computer vision task of training ResNet50 14 with the ILSVRC 2012 ImageNet dataset By The result in Table 3 shows that RAMDA attains the best validation accuracy and structured sparsity simultancously.
Language Modeling For language modeling, we train Transformer-XL (base) G using the 100 dataset 3I. Transformer-XT. is comprised of embedding and non-embeddiigg layers, and in the PyTorrc implementation, the non-embedding layers are built using linear
8

---

Table 2: Weighted group sparsity, validation accuracy and time/epoch of ProxSSI and ProxGen for AAARIO/CIFARI00. We report the average time/epoch using one NVIDIA V100 GPU.
Algorithm
Accuracy
Sparsity
Time
Accuracy
Sparsity
Time
VGGT9/CIFARTO
VGGT9/CIFARI0
ProxSSI
92.8 + 0.1%
88.4 0.2%
79s
67.3 _ 0.1%
78.6 I 0.3%
79s
ProxGen
92.8 + 0.0%
86.6 = 0.1%
24s
68.1 + 0.4%
75.5 + 0.2%
26s
PesNet55//CFARRI
ResNet50/ t CIOOO
ProxSSI
94.0 = 0.1%
83.7 = 0.6%
260s
73.7 = 0.4%
70.4 E 0.7%
251s
ProxGen
94.1 + 0.1%
80.4 + 0.4%
70s
73.6 + 0.4%
65.5 + 3.6%
74s
Table 3: Weighted group sparsity and validation accuracy ImageNet/ReeNetta
Algorithm
Accuracy
Sparsity
MSGD
77.14 + 0.04%
RAMDA
74.53 + 0.10%
29.19 + 0.94%
ProxSGD
73.50 + 0.20%
17.54 + 1.26%
ProxGen
74.17 + 0.08%
20.29 + 0.22%
RMDA
74.47 + 0.08%
25.20 + 1.69%
and layer-normalization layers. We apply group LASSO regularization to the lincar layers, and present in Table 4 the perplexity and the weighted group sparsity of the models trained. We see that RAMDA gives the best perplexity and structured sparsity simultancously.
Speech Synthesis We consider Tacotron2 41 for speech synthesis on the LJSpeech dataset [8]. We apply regularization 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.
Time Efficiency In Tables 4 and we see that although RAMDA and ProxGen have more difficult subproblems without a closed-form solution to solve, our proposed PG solver highly efficient such that the running time of them is still close to other approaches, making these regularized adaptive approaches practically feasible.
Summary 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 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 proposed a regularized dual averaging method with adaptiveness, RAMDA, for training structured neural networks. Our method outperforms state of the art modern
Table 4: Weighted group sparsity and validation perplexity Transformer-X0 XL with WikiText- 103.
Alg.
Perplexity
Sparsity
Time/epode
Adam
23.00 + 0.05
6261 21s
RAMDA
26.97 + 0.10
36.2 + 0.3%
6954 + 30s
ProxSGD
27.42 0.02
33.1
1.5%
6167 12s
ProxGen
27.49 + 0.19
30.5 + 0.6%
6652 + 21s
RMDA
27.10 0.08
36.0
2.7%
6184 20s

---

Table 5: Weighted group sparsity and validation loss on Tacotron2 with LJSpoech.
Alg.
Loss
Sparsity
Time/epen /epcee
Adam
0.39 + 0.02
431
2s
RAMDA
0.44 + 0.01
52.9
1.6%
443 1s
ProxSGD
0.50 + 0.00
34.3 = 1.6%
431 + Os
ProxGen
0.45 = 0.01
45.6 = 0.9%
438 2s
RMDA
0.46 + 0.01
45.9 1.7%
431 + 2s
architectures including LSTM and transformers as well as the ImageNet problem. We also proposed subroutine with strong convergence guarantees approximately solve the regularized subproblem of both our method and an existing framework efficiently. Extensive experiments on group sparsity showed that our subproblem solver can greally 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///wwww../ gi github. com com/2hisyuan1214]
RAMDA/
Acknowledgement
Ching-pei's research is supported in part by the JSPS Grant-in-Aid for Research Activity Start-up 23K19981 and Grant-in-Aid for Early-Career Scientists 24K20845.
References
[1] Rohan Anil, Vincet Gupta, Tomer Koren, and Yoram Singer. Memory efficient adaptive optimization. In Advances in Neural Information Processing Systems, 2019. 2
2 Hedy Attouch, Jérôme Bolte, and Benar Fux Svaiter. Convergence of descent methods for semi-algebraic and tame problems: proximal algorithms, forward backward splitting, and regularized Gauss Seidel methods. Mathemotical programming, 137(1-2):91 129, 2013.
15
3] Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. Technical report, 2016. HrXiv:1607.06455
[4] Amir Beck. Applied Mathematics, Philadelphia, PA, United States, 2017. 15
Methods in Optimization. SIAM Society for Industrial and
5) Amir Beck and Mare Teboulle. A fast iterative thrinkaggethhreshokdnng algorithm for lincar inverse problems. SIAM Journal on Imaging Sciences, 2(1):183-202, 2009. 5
6) Zihang Dai, Zhilin Yang, Yiming Yang, Jaine Carbonell, Quoc Le, and Ruslan" Salakhut- dinov. Transformer-Xi XL: Attentive language models beyond a fixed-length context. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, 2019.
[7] Aaron Defazio and Samy Jelassi. Adaptivity without compromise: A momentumized, adaptive, dual averaged gradient method for stochastic optimization. Journal of Machine Learning Research, 23(144):11 1-44, 2022. 3 14
[8] Alexandre Défossez, Leon Bottou, Francis Bach, and Nicolas Usunier. A simple conver- gence proof of adam and adagrad. Transactions on Machine Learning Research, 2022.
ISSN 2835-8856. 2
9 Tristan Deleu and Yoshua Bengio. Structured sparsity inducing adaptive optimizers for deep learning. Technical report, 2021. arXiv:2102.03869, 2 7 S
10

---

[10] Michael Denkowski and Graham Neubig. Stronger baselines for trustable results in neural machine translation. Technical report, 2017. arXiv:1706.09733, 2
[11] John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of machine learning research, 12(7), 2011.
5
[12] Warren L. Hare and Adrian S. Lewis. Identifying active constraints via partial smoothness and prox-regularily. Journal of Conver Analysis, (122:251-2666 2004. 4 7
[13] Warren L. Hare and Adrian S. Lewis. Identifying active manifolds. Algorithmic Operations Research, 2(2))73-75, 2007.
[14] Kaiming Hc, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In IEEE conference on computer vision and pattern recognition, 2016.
20
[15] Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, 9(8):1735-1780, 1997.
[16] Zih-Syuan Huang and Ching-pei Lee. Training structured neural networks through manifold ddentification and variance reduction. In International Conference on Learning Representations, 2022. 3
14
[17] Sorgey Ioffe and Christian Szogedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In International conference machine learning, pages 448 456, 2015.
[18] Keith Ito and Linda Johnson. The LJ speech dataset, 2017.
[19] Samy Jelassi and Aaron Defazio. Dual averaging is surprisingly effective for deep learning optimization. Technical report, 2020. arXiv:2010. 10502. 5 14
[20] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In International Conference on Learning Representations, 2015. 2
[21] Alex Krizhevsky. Learning multiple layers of features from tiny images. Technical report,
2009.
20
[22] Frederik Kunstner, Jacques Chen, Jonathan Wilder Lavington, and Mark Schmidtt. 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.
[23] Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient- learning applied to document recognition. Proceedings of the IEEE, 86(11):2278 2324, 1998. 20
[24] Ching-pei Lee. Accelerating inexact successive quadratic approximation for regularized optimization through manifold identification. Mathematical Programming, 2023.
[25] Sangkyun Lee and Stephen J. Wright. Manifold identification in dual averaging for regularized stochastic online learning. Journal of Machine Learning Research, 13: 1705 1744, 2012. 3
[26] Adrian S. Lewis. Active sets, nonsmoothness, and sensitivity. STAM Journal on Optimization, 13(3):702-785 725, 2002. 7
[27] Adrian S. Lewis and Shanshan Zhang, Partial smoothness, tilt stability, and generalized hessians. SIAM Journal on Optimization, 23(1):74 94, 2013.
[28] Liyuan Liu, Xiaodong Liu, Jianfeng Gao, Weizhu Chen, and Jiawei Han. Understanding the difficulty of training transformers. In Proceedings of the 2020 Conference on Empirical Methods in Natural Longuage Processing (EMNLP), 2020.
11

---

[29] Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. Swin transformer: Hierarchical vision transformer using shifted windows. In Proceedings of the IEEE/CVF international conference on computer vision, 2021. 2
[30] Nya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In Interna- tional Conference on Learning Represenlalions, 2019. 14
[31] Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel
mixture models. In International Conference Learning Representations, 2017. g
[32] Yurii Nestcrov. Primal-dual subgradient methods for convex problems. Mathematical programming, 120(1):221-259, 2009,
[33] Yurii Nesterov. Gradient methods for minimizing composite functions. Mathematical Programming, 140(1):125-611, 2013. 5
(34] Evgeni Alekseevich Nurminskii. The quasigradient method for the solving of the nonlincar programming problems. Cybernetics, 911):145-150, 1973.
[35] Alexandra Peste, Eugenia Tofinova, Adrian Vladu, and Dan Alistarh. Ac/dc: Alternating compressed/deccome /ecompressed training of deep neural networks. In Advances in neural
information processing systems, 2021. g
[36] René Poliquin and R Rockafellar. Prox-regular functions in variational analysis. Thans- actions of the American Mathematical Society, 348(5):1805-15) 1996.
[37] Clarice Poon, Jingwei Liang, and Carola- Bibiane Schönlieb. Local convergence proper- ties of SAGA Prox-SVRG and acceleration. In International Conference on Machine Learning, 2018.
Carola-Bibiane
4
(38] R. Tyrrell Rockafellar and Roger J-B Wets. Variational analysis, volume 317. Springer Science & Business Media, 2009. 5
[39] Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, Alexander C. Berg, and Li Fei-Fei. Imagenet large scale visual recognition challenge. International journal of computer vision, 11):2): 252, 2015.
[40] Andrzej Ruszezyúski. Feasible direction methods for stochastic programming problems. Mathematical Programming, 19:220 229, 1980. 16
[41] Jonathan Shen, Ruoming Pang, Ron Weiss, Mike Schuster, Navdeep Jaitly, Zongheng Yang, Zhifeng Chen, Yu Zhang, Yuxuan Wang, Rj Skerrv-yaan Rif A. Saurous, Yannis Agiomyrgiannakii, and Yonghui Wu. Natural tts synthesis by conditioning wavenet mel spectrogram predictions. Tn 2018 IEFF international conference acoustics, speech and signal processing (ICASSP), 2018. g
[42) Karen Simonyan and Andrew Zisserman. Very deep convolutional networks for large- scale image recognition. In International Conference on Learning Representations, 2015.
20
[43] Wei Tao, Zhisong Pan, Gaowei Wu, and Qing Tao. Primal averaging: A new gradient evaluation step to attain the optimal individual convergence. IEEE transactions on cybermetics, 50(2):835 845, 2018. 5
[44] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasa Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in neural information processing systems, 2017. 2
[45] Wei Wen, Chunpeng Wu, Yandan Wang, Yiran Chen, and Hai Li. Learning st suctuueed sparsity deep neural networks. Advances neural information processing systems, pages 2074 2074-2082, 2016. /
12

---

[46] Han Xiao, Kashif Rasul, and Roland Vollgraf. Fashion-min mniss:: novel image dataset for benchmarking machine learning algorithms. Technical report, 2017. arXiv:170..07747,
22
[47] Lin Xiao. Dual averaging methods for regularized stochastic learning and online optimization. Journal of Machine Learning Research, 11(88):2543 2292523-666 2596, 2010.
3
[48) Zhenda Xie, Zheng Zhang, Yue Cao, Yutong Lin, Jianmin Bao, Zhuliang Yao, Qi Dai, and Han Hu. SimMIM: simple framework for masked image modeling. In Proceedings of the IEEE/CVE Conference on Computer Vision and Pattern Recognation, pages
9653-9663, 2022. 22
[49] Yang Yang, Yaxiong Yuan, Avraam Chatzimichailidis, Ruud JG van Sloun, Lei Lei, and Symcon Chatzinotas. ProxSGD: Training structured neural networks under regulariza- tion and constraints. In International Conference Learning Representations, 2019. 2
[50] Ming Yuan and Yi Lin. Model selection and estimation in regression with grouped variables. Journal of the Royal Statistical Society; Series B (Statistical Methodology), 68(1):49-67 2006.
[51] Jihun Yun, Aurélie C Lozano, and Eunho Yang. Adaptive proximal gradient methods for structured neural networks. In Advances in Neural Information Processing Systems, 2021.
[52] Jingzhao Zhang, Sai Praneeth Karimireddy, Andreas Veit, Seungyeon Kim, Sashank Reddi, Sanjiv Kumar, and Suvrit Sra. Why are adaptive methods good for attention models? In Advances in Neural Information Processing Systems, 2020. 2
Appendices
Table of Contents
More Experiment Details
13
A.1 Implementation and Hyperparameler Setting RAMDA
A.2 Details of the Algorithms Compared
A.3 Compulational Resource
B
Proofs
14
B.1 Convergence Result for ProxGen
18
Additional Experiments for Computer Vision
20
D Plots of Sparsity Level and Validation Accuracy Epochs
21
Experiment with Nuclear-norm Regularization
22
A More Experiment Details
This section describes details of our implementation of RAMDA and the setting of the experiments conducted in Section 6
13

---

A.1 Implementation and Hyperparameter Setting of RAMDA
Similar [6 we introduce restart strategy to the implementation of RAMDA. During each stage of the training, the learning rate π and the momentum factor are fixed. Once the epoch count enters the next stage, we reset the counter to 1 and use the output parameter W from the previous round as the new input parameter W to the same algorithm, set ac, V* and U 0, but keep the scheduling for 7 and c going without resetting them, and decrease by factor. We initialize a as either 0.1 or 0.01, depending on the problems, and use a constant a until the final stage, where we gradually increase by
miiin(cVi,))
where counts the training steps executed at this final stage. This momentum strategy is applied to both RAMDA and RMDA in our experiments.
A.2 Details of the Algorithms Compared
In
Table
we summarize details the algorithms compared in Section
Table 6: Algorithms used in the experiments.
Algorithm
Unregularized counterpart
Subproblem
RAMDA
MADGRAD
PG
RMDA
MDA 19
Closed-form solution
ProxSGD
MSGD
Closed-lorm solution
ProxGen
AdamW
30
PG
ProxSST
AdamW
30
Newton- Raphosn
A.3
Computational Resource
We conduct all experiments utilizing NVIDIA TESLA V100 (32 GB) GPUs. We employ eight V100 GPUs for each run of the ImageNet experiments. For all other experiments, utilize a single V100 GPU per run.
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 [ Assume that (4) and 66) has at least one optimal solution with finite optimal objective value. Given E > 0, the number of iterations of Algorithm 2) takes to satisfy (7) for
both and (6) O(ET)) when vb conver and O when vb is nonconvez.
Proof. We use the notation
f(() + (Z)
to unify the two objective function Qi/a, and Qu, where fe is the smooth part and we define the Lipschitz constant of Vf as L.
At each iteration,
PG solves the following subproblem
2 E arg min (Vf((Z',, Z - z²) + 2 20112 +22), 20,
Z
and thus from the first-order optimality conditions, clearly we have
V)((ztt)) ((12))))-AA(())----- - (gitz -2) (Q,(2)t1).
0,
We thus have from the Lipschitz continuity of V that
||s|| < [(((())))) - - - || z'll s (L+ +0- 1 -z'|.
1
min (20Q.(2:-1)
0,
99)
14

---

Note that Qt is lower bounded, say by Qi, and has at least one solution Z (unique when s) convex).
In the case that 1/ is convex, we know that 0, = 1/L,, and 3 clearly shows that the subproblem objoctive Q is c-strongly convex. Therefore, standard analysis of proximal gradient (sce, for example, 1 Lemma 10.4 and iThcorem 10.29] gives that
- 2 1 21 5 - (z))
Vj 0,
(10)
L (1 -1)) || z-1 Vj 1.
(11)
The combination of 99) 111 shows that it takes Ollog " iterations for PG to reach the required precision of Et. When & is nonconvex, we have that 02 = 1/((TT) and standard analysis 2 Section 5.1] gives
C
min k=
1z 7d
(12)
VJ
for some constant C' depending on L and (W?) - Qi- Therefore, (12) and (9) show that it takes O itcrations to reach the desired procision.
Theorem 2 Consider [W*) generated by Algorithm for with (7 and {ee} and { satisfying > c = oo and (8) Assume there is L 2 such that for any E. fe is almost surely c-iippchtt--oontimuuuuultyy
so the expectation is also L-Lipschitz-contimmoos differentiable, there is 2 0 such that (Eg.D//VVe, (W*=1)ll* S C for all i, and that the sel of stationary points == {W |0 } nonempty. For any given W consider the event that }W"} converges to a point W (each event corresponds to a different W.. If du outer semicontinuous at W, this event has a nonzero probability, and {nc} satisfies
Σ staT = =0, | (sta, a.e
(ssar < co,
(13)
(14)
then have that W 2 with probability conditional on this event. Moreover, {W*) also converges to this stationary point W.
Proof. First, we prove that when (W") converges to W, (W*) also converges to W. From
5
we have that
| W| s (1 - w|l t aa|lt w/|
(15)
Since W + W, for any > 0 we can find integer f, > 0 such that ll Wll < for all t 22 te. Therefore, by deducting e from both sides of (15) we get
||w - [w'-wWl-es wil-ss fia-as) (1IW---- w/l - S exp
Σ
(/[ -)), W2
k=tt
k=t.
By letting approach infinity and noting that Ea = 00, we see that
lim 4000
W W 5..e.
Because is arbitrary, we see that | W|l 0, and hence (W*) converges to W. Next, consider the update of a; Ut,
from (2) that
U a:-1
-
St fe. (W+-1)
St
Ut-1
St
+
(16)
a: 0t at-1
O2
at-1
an
15

---

Moreover, the assumptions on N satisfies all the required conditions of Lemma 1 of 40 We therefore apply Lemma 1 of 40 to conclude that
U
E
Ve (W*) o Vie (W")]
(17)
a4
The update for at 'v has a form analogous to (16) and we have from Jensen's inequality that
E ff. (w*=))]²² - V (WV-1)Il S <VC,
implying that the second moment also bounded expectation. We can therefore also apply Lemma of 40 to at av've and conclude that
V
V a.3.
VEEDD [ee (W*)]
(18)
at
We further notice that the union of two events that happens almost surely is still an event that happens almost surely.
From and
we can see that there sequence fet} such that
Vt
pd
(W* wo) OW("), | Sec
(19)
04
On
Our assumption (14)
implies that ar - xo, which together with
(8) leads to
2t
(20)
a2
From 188) that VE lff (W) Lipschitz continuous, and that W - W (which we have proven in the first part), we see that
V
VEEDD [F]]
(21)
For the third term, we have from
3 and
177
that
p
Ut
Diag
+
T.
a4
an
Again since dg - oo, the second term of the equation above converges to 0. Therefore, by we oblain that
17
pt 3
Diag /EEWD IV (W') Vife (W*))
Again from the continuity
VEEND life (W)l and that a + oo, we conclude that
pt
Diag (EED [W.e (W) o Ve (W)]
a.s. 0.
(22)
Finally, using the outer semicontinuity of (W) at W, we conclude from (19)
22 that
0 VEE-D [fe (W)] + lim ((V') S VIES-D [fe (W)] + (WW) OF(W)
with probability one, showing that W is stationary point almost surely.
LI
Theorem Consider Algorithm 1 with the conditions in Theorem 2 satisfied. Consider the
event of (W*} converying to certain point W as in Theorem 2 If the probability of this event is nonzero; & is prox-regular and subdifferentially continuous at W and partly smooth at W relative to the active C² manifold Mw; Oi is outer semicontinuous W; and the nondegeneracy condition
- -Vf (W) relative interior of db (W)
(23)
holds at W, then conditional this event, almost surely there is To such that
W Mws VI To-
(24)
In other words, the active manifold at W identified by the ilerates of Algorithm afler
finâte reumeer of iterations almost surely.
16

---

Proof. From
19
there exists sequence { such that
v
at + pt
Y' E (w(").
(W
WD + + =0= 4l.
(25)
at
at
at
For notational ease, we denote
S(W)= = BEWD [Se(W)]
(26)
From 25
we then get
VS(W')
V
2t
pt
( - wo) OF(Wy.
(27)
a0
a1
We aim to show that
dist(0, (F(())FF(W)))) = min
|||||
YESF(W!)
converges to 0 almost surely. From (7 have
dist(0, OFFt()))(((1)))
v
V
2x
pt
(
Wo)
0t
a2
at
V
V
2t
pt
(Vir Wo
+
+
01
02
at
VS(WT-
V
C
pt
( - (
(28)
04
02
where we get the first inequality from the triangle inequality and the second from (19)
According to (18) and
(22) there are {Ac) and (B,} such that
Vi
+ Ar. |ell a.s. O
(29)
at Diag
(VVSWW)) VS(WO) B,, ||Ba|| - 0.
Substituting the above two equations back to 28) we obtain (0st(O,FFFF)))
< VS(WW) V//WS) + ||
(ar
V Vf(W)) V/(W()) (F(W))))((WW
+1|Bel) |lire wo
12
< W - w
+ +1A+el + + (at VVF(W') VF(W?) +1|Bil) |lve moll
(30)
From Theorem
we know that : and W both converge to W,
|
-ie| W
[ | w|| || W|| 0.
From 8
and (14) we know that Ee/ar + 0. Because v¹: + W, we also have that
| w || w <c.
From W - W, we have that
- / VS(W) <00.
Vf(W")) Vf(W*)
Combining these results with (30] we conclude that
a ag L L| - w' + ||All
+ a VVS(W") V/(W))| + -lBall) | -wol a.a, o,
proving that
dist d0 (F(W)) an O.
On the other hand, since f is continuous and & is subdifferentially continuous at W (which implies F is also subdifferentially contnuous at W), W W, and that V.((V//)) Y. a* E (F)) (from Theorem 2), we know that F(W)) A.c, F(W) as well. Therefore, we can apply Lemma 1 of 24 to conclude that (24) indeed holds for some To oo almost surely. D
17

---

Algorithm 3 ProxGen (Wo, 7, T2, fmh (pe), foe) (c)) (4,),8))
mo +0
for t= -=,,..., do
Sample & D with batch size
G V-&,,, V[e,(W1))
my - P-m--1 (1 Gc)
Construct P satisfying P - 51
0, 9
Compute W by roughly solving (6) that satisfies 7
with Qe replacod by Q and Vi
replaced by W++,, using PG (W*, W my ni pt, 02 T2)
output: W
B.1 Convergence Result for ProxGen
We next discuss the convergence result for the framework of 51 with inexact added. For consistency, we first use our notations to introduce their framework, with our inexactness condition added, in Algorithm 3
In their analysis, 51 made the following four assumptions, and we will follow these assump- tions using the notation [26)
(C-1) The expected loss function f is Liipeediirccoottmoooyyyyyffeetee and lower- bounded for some L 2 O.
(C-2) The stochastic gradient G f'..(WW-)) an unbiased estimator of Vf(W(-1) with bounded variance.
a² b,"
REIND| = VV<D")"=V(((">1),
FEODD ||o V/(W">>> S
1 2 0,
where b, is the batch size of & and
constant.
(C-3) There some Po: f E [0, 1) and D,G such that [ - SD, ||'|| < C,
and Pt popt-1 for all t.
(C-4) There is some y >> 0 such that 1n 7 pd 1| S 1/7 0o for all
(C-5) There is 5 > 0 such that
p' p'20 8,
S
Vt 2 0.
(31)
3L
Theorem 4 For the framework in 15 with the subproblem solved upprocimalely by Algo-
rithm 2 such that (7) holds with {er} satisfying [8] Then Theorem of [5] still holds, but
with the constants {Q,} being also dependent on
Proof. The major flow of our proof follows that of 51 but with suitable modifications to accommodate the inexactness condition in the subproblem solving.
It is clear from
5l Lemma 1] that
|m.]l < G,
Vt 2 0.
(32)
By the update rule for m.
(6) and
we have that there is such that
0 2 +1 - pe)g. + p + (w") + 1 -oll--Waet W - m
|| S 6t
Vt 0,
leading to
v((W*) 2 11)g p.m
(p))(*--'*' OF(W").
(33)
7
18

---

We thus have from (33 and (C-4) that
dist(0, )F(
$ 2t + pt)g: vf(W**) + Pemp-1 (W* -1--) + 1 w(-1) (W - -w |
I 4 (1 pi)ge - Vf(W") + P=m + (W'- w|- -W-- + 46? + 4
1
(P')(w--W w--
171
+ 4 m - y
s 4 1-) v((W') + pam-- + (W* - b| +4 go + - 427.
111
1
1
(34)
We will separately bound the quantities Ti and || - w111 below.
From the subproblem objective requirement in we get
((1 ((1-Pe)ge Pc)ge + P.m4-1 -* -----> + (W"))
(W" w (W'----')>> s (W
21,.
(35)
From C-1)
we have
f(W*) < f(W (Vf(W( 1 -Wi1) + L
Wt - -
(36)
2
Summing 35 and
366
gives
((1-Pt)ge
p,m- W w - ||
W
2
F(W F(W*)
pt
+
(37)
Note that n/ 1 pt LI 2 from 31 so the sccond term in (37) is nonnegative.
(37)
together
with (C-3) then leads to
||W* 11
41
200 < F(W*-1)- - F(W") (g: (f( w - -W(-1) + (pigs, W -W(-1 - (pomm-1, wy--W--1
- - F(W") + 2T. 11:t V/(W"=1))1 + L - $1 1 11-1112
S F(W--1)-
|1 12 + 11-n-----11
2T.
+ |1|| W
< F(W1-1)- - -F(-)))-FWW))-- 2L
11 [I - -w +
+ Pojt--1D DG.
2L
Summing it over t = 1, 2,...... T and utilizing the assumption that the step sizes are increasing then give
8
L
W S 2L
11g - V1(
2rjo
where
C DgD
G 21(1 22)
A = F(W") min F(W),
From the inequality above, we obtain
- - H2
1g Vf(
(38)
1=1
t=1
for some constants H1, H2 depending on L, A, 8, no. and C₁. From (37)
we also have
(11 pt)g v.(W") + p,m,-1, W W' w'- w
< F(W(-1)- - F(W*)- - (VS(W)) - v/(W-(-1), W 1 -W | 11-1112
(11)-
201
< F(W(-!)- - F(W")) - (VS(W9- - V)(W"-1), w -WC-)).
19

---

Therefore, we obtain
Ti = |[| - pt)ge - V + + | - 1
+ (1 pt)ge - Vf(WW) + 1, W W1-1)
< [[11 -p) - V/(W--1)+ V/(W"-1) - vS(W") + + + ||** 11
-
- + =F(F(W(-)) - -(W"")((WW))))) -V -V/))--)) W - -
< 411 " - -1-
+
4p; 19 || - 1
+ 2 (F((W-1)) F(W") + -/|'" - 11-
< ) + 2p + 2L + | - 41| V/(W"=1)
(39)
Lel C2 2+ 21 4L2 7- and insert (39) into (34)
we gel
(et(0,0F(WW)))
< 4 4 (F(( 8p8 2t-1) G2 + C2 C2llW* 11-1112
+ 4119. -W-
vf(W )|
+
(40)
Therefore, we have from (8) and 40
and
C-2) that
E Bi,.. fo (dist (0, OF (F(w)))
(1 p-)ge
VS(W*) + + pm + me
(P*) -/'"--"-111f?
T 1=1
4
24
82
+4
Σ
T
Bl1g >)(
C2
_
BllW - 1- + Σ
7
1- 11
1=1
t=1
2A +
G2
1- 112 40
C2(H2 + H2o
Σ
T
) +
7
hn
Q T T
22
Q3 T
Σ be
+
+
T
for some constants Q1, Q2, Q3 dependent on (no, D. but not on T. This proves our theorem.
C Additional Experiments for Computer Vision
In this section, compare RAMDA with other methods on image classification with smaller datasets. They are:
1. Logistic regression (ncural network with no hidden layer) with the MNIST datasct
2. A modified VGG19 42] with the CIFAR10 dataset 21
3. The same VGG19 with the CIFAR100 dataset [2]]
4. A modified ResNet50 14 with the CIFAR10 dataset.
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 rum, with the initial point being the origin, as it is convex problem. Moreover, in this problem, when dealing with ProxSSI, ProxGen, and ProxSGD whose sparsity lovels are highly unstable over iterations, for all other problems, we report the group sparsity level of the final output.
report their highest weighted group sparsity over all epochs, but
Experiments in this subsection show that RAMDA might. sometimes perform worse than existing methods on smaller problems like 1IFAR10/100. Fortunately, for such smaller problems, the training cost is not very significant, and one can afford to try more algorithms.
20

---

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 + 0.14%
ProxSGD
92.82 + 0.09%
82.76 5.42%
ProxSSI
92.81 + 0.15%
88.40 + 0.23%
ProxGen
92.83 + 0.05%
86.64 = 0.12%
RMDA
93.13 + 0.10%
90.22 + 0.06%
RAMDA
92.89 + 0.13%
86.31 0.31%
VGG19/CIFARION
MSGD
74.07 E 0.05%
ProxSGD
71.96 + 0.15%
72.34 + 11.9%
ProxSSI
67.29 + 0.06%
78.58 + 0.34%
ProxGen
68.13 + 0.36%
75.46 + 0.17%
RMDA
71.96 + 0.31%
80.88 + 0.11%
RAMDA
70.47 + 0.25%
65.19 + 0.77%
ResSAE350/CFFARI
MSGD
95.54 + 0.19%
ProxSGD
92.36 + 0.05%
82.18 + 2.67%
ProxSSI
94.04 +: 0.12%
83.67 + 0.63%
ProxGen
94.07 + 0.12%
80.45 + 0.45%
RMDA
95.11 + 0.11%
85.64 = 0.12%
RAMDA
93.85 + 0.10%
81.99 + 1.26%
ResNee500/CIAARIO
MSGD
79.49 + 0.49%
ProxSGD
74.54 + 0.58%
49.29 + 5.91%
ProxSSI
73.65 + 0.39%
70.38 - 0.74%
ProxGen
73.63 + 0.43%
65.51 + 3.58%
RMDA
75.62 + 0.19%
79.97 0.27%
RAMDA
69.23 + 0.86%
68.65 + 1.83%
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. ResNet50 with the ILSVRC 2012 ImageNet dataset.
2. Transformer-LL with the WikiText-103 dataset.
3. Tacotron2 with the LJSpeech dataset.
4. Logistic Regression with the MNIST dataset.
5. A modified VGG19 with the CIFAR10 dalaset.
6. The same VGG19 with the CIFAR100 dataset.
7. A modified ResNet50 with the CIFAR10 dataset.
8. The same ResNet50 with the CIFAR100 dataset.
In the plot for Transformer-XI, 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
21

---

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, is possible that the structured sparsity level could be further improved.
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 muclear-norm regularizer per layer individually by treating each layer as a matrix. Given matrix X E Rm of rank r I min(m,n) with its singular value decomposition (SVD) X = UVVT, where U E IR"" V R are orthogonal and the positive definite diagonal matrix E grxr represents the nonzero singular values of X, the nuclear norm of X is computed by
XX|
E
and the corresponding proximal operator for >
proXIIII| (x) = UÊVT, where Σig = max Ei, - A}, i= I. Given point X with rank r*, the active manifold of the nuclear norm at X is M(X*) = (Y
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 lo exploit this structure for prediction acceleration and to make the training more efficient, possibly using iteralive 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 simple neural network with six fully-connected layers using the FashionMNIST datasct A6. Since this is 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 Lo 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 24 for masked image modeling 48 using the CIFAR10 datasct. 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 (able 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.
NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
22

---

ResNet50 ImageNet
Transformer-XL WikiText-100 103
Weigtted Studired Spersily Valiseion Aucurecy
0.8
Valdation Purplnelty
100
MSGD
0.4
ProdSOO
60
RMIA
PronCen
RAMDA
0.-
20
40
00
80 Epochs
100
120
140
100
20
30
50
70
Every 5000 Steps
0.4
0.2-
0.1
0.1
Weigited
0.0
0.0
100
Epochs
Every
(a) ResNet50 on ImageNet
(b) Transformer-XI. on WikiText-103
Tacotron2 on LJSpeech
Logistic Regression on MNIST
Weighted Stocture Sparsity Valdation Arturay
Welddiion Lows
Adam
1.00
Proiten
Fro
0.75
0.50
HAM
0.00
KNDA
Froa
0.86
MATOR
400
600
800
500
Weighted Stoctured Spersity
Epochs
0.6
0.4-
0.4
0.2-
0.2
200
400
600
800
100
200
300
400
500
Epochs
Epechs
(c) Tacotron2 on LJSpeech
(d) Logistic Regression on MNIST
VGG19 CIFAR10
VGG19 CIFAR100
Weighted Strictued Sporsiy Valtetion Noormn
1.00
Weighted Stractune Spansity Velidation Rectricy
0.8
0.75
HBGD
5500
Pro
Pru
0.50
Pro-
FrusS
ProsCen
0.25
0.2
0.00
0.0
200
400
600
800
1000
200
400
600
800
1000
Eporhss
Fpochs
0.75
0.75
0.50
0.50
0.25
0.25
0.00
0.00
800
Epochis
Epochss
(e) VGG19 CIFAR10
(l) VGG19 CIFAR100
ResNet50 on CIFAR10
ResNet50 on CIFAR100
Weighted Struittedddtuuuuueu stroctored Spersity Velddetim Accraey
Wiightdd Structured Sparsity vaidalion Accuary
0.75
0.75
0.50
0.50
0.25
RMDA
0.25
ANDA
AAMELA
RAMDA
0.00
0.00
600
Epochs
Epochs
0.75
0.50
0.25
0.00
0.00
200
400
600
800
1000
200
400
600
800
1000
Epochs
Epochs
(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.
23

---

ResNet50 ImageNet
Transformer-XL on WikiText-103
Weighted Structured Sparsity
0.35
Weighted Structured Sparsity
0.38
0.36
0.30
PraxSGD ProxGen RAMDA
0.34
ProkSCD
0.25
ProxGen
0.32
RAMDA
0.20
0.30
0.28
130
140
150
60
65
70
Epochs
Every 5000 Steps
(a) ResNet50 on ImageNet
(b)
Transformer-Xi on 1iki
Tacotron2 on LSppeehh
Logistic Regression on MNIST
Weighted Structured Sparsity
Weighted Structured Sparsity
0.6
0.5
0.50
0.4
ProxSGD
0.45
MAM M.M.w
Prox5SI
Prastin
0.3
ProxGen
RMDA
0.40
RMDA
0.2
RAMDA
pM / y yar
0.1
0.30
0.0
800
850
900
950
450
460
470
480
490
500
Epochs
Epochs
(c) Tacotron2 on LJSpeech
(d) Logistic Regression on MNIST
VGG19 CIFAR10
VGG19
CIFAR100
Weighted Structured Sparsity
Weighted Structured Sparsity
0.85
0.9
M
W/W
0.80
0.8
ProvSGO
ProxSGD
0.7
ProvSSi
Prox5S
PronGen HARDA BAPRIDA
ProxGen
0.6
0.70
RMDA RAMDA
0.5
0.65
0.4
920
980
920
980
1000
Epochs
Epochs
(e) VGG19 CIFAR10
(I) VGG19 CIFAR100
ResNet50 on CIFAR10
ResNet50 on CIFAR100
Weighted Structured Sparsity
Weighted Structured Sparsity
0.9
0.86
0.8
0.7
PrawGen
ProxGen RANDA
0.6
RAMION
0.82
0.5
0.80
0.4
900
920
940
960
980
1000
900
920
940
960
980
1000
Epochs
Epochs
(g) ResNet50 on CIFAR10
(h) ResNet50 on CIFAR100
Figure 2: Group sparsity level at the last epochs.
24

---

Table 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 = 0.29%
ProxSGD
87.54 0.52%
78.00 + 0.77%
ProxGen
86.66 = 0.33%
87.46 + 4.19%
RMDA
88.19 0.23%
91.88 = 0.12%
RAMDA
87.99 = 0.24%
89.59 = 0.42%
Table Low-rank level and validation loss of different. methods on pretraining modified vision transformer model using the CIFAR10 dataset for masked image modeling.
Algorithm
Validation loss
Low-rank level
AdamW
0.0865 0.0001
ProxSGD
0.1042 + 0.0003
82.60 + 0.34%
ProxGen
2.47% 0,1120 0.0019
82.64
RMDA
0.1054 + 0.0031
86.23
+ 0.41%
RAMDA
0.1035 + 0.0016
86.20 0.35%
Answer: [Yes]
Justification: Our claims accurately reflect the paper's contributions and scope.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: Limitations of our work discussed in Section (after and Appendix C
Theorem
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set assump- tions and a complete (and correct) proof?
Answer: [Yes]
Justification: All assumptions are stated clearly in each theorem statement, and all detailed proofs are provided in Appendix
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: Algorithm details are all given in the paper, and the parameter sel Lings all available in the supplementary materials.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: Our code is provided in the supplementary materials, our data are public data sets, and sufficient instructions are given in the README, in the supplementary materials.
25

---

6. Experimental Betting/Detail
Does the paper specify all the training and lest details (e.g. data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Question:
Answer: [Yes]
Justification: All details are given in either the main paper or the supplementary materials.
7. Experiment Statistical Significance
Question:
Does the paper report error bars suitably and correctly defined other appropriate information about the statistical significance the experiments?
Answer:
[Yes]
Justification: We report the mean and standard deviation of the comparison criteria over three different random initializations.
8. Experiments Compute Resources
Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Answer: [Yes]
Justification ion: Time of execution is reported in Section 6 and details of the computer
resources given in Appendix A.3
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code Ethics http://///aurrrr..///////////aaa
A Answer: [Yes]
Justification: The research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics.
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: This is fundamental research work and there is no foreseeable negative societal impact.
11. Safeguards
Question:
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Answer: [NA]
Justification:
The paper poses no such risks.
12. Licenses for existing assets
Question: Are the creators original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: (Yes]
Justification: We do cite all of the papers that proposed the models, datasets, and code we used.
13. New Assets
Question: Are new assets introduced the paper well documented and the documentation provided alongside the assets?
26

---

Answer: [Yes]
Justification: This paper will introduce new open-source software, and we have provided a README file for documentation of the package.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, well as details about compensation (if any)?
Answer: [NA]
Justification: The paper does not involve crowdsourcing nor research with human subjects.
15. Institutional Review Board (IRB) Approvals or Equivalent for Rescarch with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: The paper does not involve crowdsourcing nor rescarch with human subjects.
27


Reviews:
[{"review_id": "n2qbwOWnaT", "reviewer": "Reviewer_X8bY", "summary": "#### Summary\nThe paper introduces the Regularized Adaptive Momentum Dual Averaging (RAMDA) algorithm for training structured neural networks. RAMDA addresses the challenge of solving the subproblem involved in the regularized adaptive methods, which typically lacks a closed-form solution. The paper presents an inexactness condition that retains convergence guarantees and proposes an efficient subproblem solver. The algorithm leverages manifold identification theory to ensure that the iterates of RAMDA attain the ideal structure induced by the regularizer at convergence. Extensive experiments demonstrate the effectiveness of RAMDA in various tasks, including computer vision, language modeling, and speech synthesis.", "questions": "#### Questions\n1. **Computational Complexity**:\n    - Could you provide more details on the computational complexity of the proposed subproblem solver? How does it scale with increasing data size and model complexity?\n\n2. **Generality**:\n    - The paper focuses on structured neural networks with specific regularizers. Are there any challenges in extending RAMDA to other types of models or regularizers, such as those used in different machine learning tasks?\n\n3. **Comparison with Existing Methods**:\n    - How does RAMDA compare empirically with other state-of-the-art methods for training structured neural networks? Are there specific scenarios where RAMDA significantly outperforms these methods?\n\n4. **Implementation Guidelines**:\n    - Can you offer practical guidelines for implementing RAMDA in real-world scenarios? Specifically, how should practitioners tune the parameters, such as the learning rate and the inexactness threshold?\n\n5. **Assumptions and Limitations**:\n    - The paper discusses some assumptions and limitations. Could you elaborate on the key assumptions that are critical for the theoretical results, and how robust the method is to violations of these assumptions?", "limitations": "no", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "#### Strengths\n1. **Novel Algorithm**: RAMDA combines adaptive momentum dual averaging with efficient inexact subproblem solving, providing a practical and theoretically sound method for training structured neural networks.\n2. **Theoretical Guarantees**: The paper provides strong theoretical support, including convergence guarantees and structure identification, ensuring the algorithm's robustness.\n3. **Practical Efficiency**: The proposed inexact subproblem solver is efficient, making RAMDA feasible for large-scale applications.\n4. **Empirical Validation**: Extensive experiments across multiple domains demonstrate the superior performance of RAMDA in terms of both prediction accuracy and structured sparsity.", "weaknesses": "#### Weaknesses\n1. **Computational Complexity**: The computational complexity of the proposed subproblem solver, especially for high-dimensional data, needs more detailed discussion.\n2. **Generality**: While the paper focuses on specific types of structured neural networks, extending the methodology to other models and regularizers would enhance its generality.\n3. **Comparative Analysis**: More detailed comparisons with other state-of-the-art methods, beyond the provided benchmarks, would strengthen the empirical validation.\n4. **Implementation Details**: Practical guidelines for implementing RAMDA, including parameter tuning and handling different data distributions, are somewhat lacking.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the evaluation of our paper. Our reply is as follows.\n\n1. Computational Complexity: The complexity of the subproblem (Eq. 4) depends on the regularizer and especially its associated \n    proximal operation. Let the subproblem dimension, which is also the model size, be $n$. Most of the widely used regularizers come \n    with a proximal operator whose cost scales at $O(n)$, like the sparsity-inducing ones. However, there are also cases with higher \n    dependency on the problem dimension, such as the nuclear norm regularizer we adopted in appendix E, whose cost scales at \n    $O(n^3)$. Therefore, the subproblem complexity dependency on the model size would be dependent on the regularizer.\n    On the other hand, the data size does not affect the subproblem solver at all -- likely it affects the number of outer iterations and the \n    time cost of computing the stochastic gradient.\n\n2. Generality: We disagree that this paper focused on specific regularizers. Indeed for experimental purposes we have to select some \n    representative regularizers because there is no way to exhaust all possible regularizers in experiments. However, our discussion, \n    assumption on the regularizer, and analysis are all general such that a broad class of regularizers are included. On the other hand, \n    as we discussed in Appendix E and in the previous point, some regularizers might be useful but too costly for the subproblem solver, \n    and it might require more domain knowledge to design better structure grouping or subproblem solvers to tackle this cost issue in \n    these cases.\n\n\n    As for extending to other machine learning tasks, RAMDA is specifically designed for deep neural networks and probably less \n    suitable for other machine learning tasks. In particular, adaptive models are not that widely used in other machine learning tasks, up \n    to our knowledge, and therefore we do not expect our method to be as competitive on other machine learning tasks.\n\n\n3. Comparison with Existing Methods: We have compared RAMDA with other state-of-the-art methods for training structured neural networks in Section 6. Our numerical results have clearly indicated that RAMDA produces better prediction performance and higher structuredness simultaneously than state of the art, and with running time comparable to the most efficient algorithm. This shows that RAMDA indeed pushes forward the state of the art for training structured neural networks.\nIt is well-known that adaptive methods tend to be most successful in many state-of-the-art deep neural networks, especially the transformer (see Zhang et al. (2020)). It is also the case for RAMDA: our experiments have shown that it performs the best on such models. It also has outstanding performance on large-scale problems like the ImageNet problem.\n\n4. Implementation Guidelines:\n    There are four type of possible hyperparameters to tune:\n\n    - Learning rate: We do grid search to tune this hyperparameter in all of our experiments. Otherwise, $10^{-2}$ is a good guess if the \n        computational resources are limited. \n\n    - Learning rate scheduling: In all of the experiments, we use stage-wise learning rate scheduling for RAMDA. Each stage has an equal number of iterations, and the multiplicative factor of learning rate decay at each stage is $10^{-1}$.\n\n    - $\\lambda$ (weight of the regularization): We do grid search to tune this hyperparameter in all of our experiments.\n\n    - Momentum: $10^{-1}$ or $10^{-2}$ are good choices. We simply chose $10^{-2}$ in all the experiments of the Section 6.\n\n    Inexactness threshold: There are two hyperparemeters to controll the inexactness threshold (see Section 6):\n\n    - max_iters: the maximun number of the iterations. We set it to $100$ in all the experiments. \n\n    - rtol (: (previous_subproblem_objective-current_subproblem_objective)/(abs(current_subproblem_objective)+1.0): criterion for early \n        stopping. We set it to $10^{-8}$ in all the experiments except Table 1.\n\n    Implementation and experimental code for reproducibility has also been included in the supplementary materials.\n \n5. Assumptions and Limitations: Thms 1 & 4 have very little assumptions and limitations, and therefore they are applicable to almost all situations. For Thm 2, as discussed after the theorem statement, the key assumption is that the iterates converge to a point. Its satisfaction might require some additional regularity conditions that are usually satisfied by generic problems, but the analysis would be very involved and we will leave it as future work. In practice, we notice that with proper hyperparameter search, this assumption seems to hold true in general -- usually those hyperparameters that lead to divergent iterates are not adopted.\nFor Thm 3, the nondegeneracy condition is the key assumption, and again it holds for generic problems (since the relative boundary is a measure zero set, that the negative gradient falls on it is a probability zero event.)", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "XnawlCh1v9", "reviewer": "Reviewer_xuwn", "summary": "This paper develops regularized adaptive momentum dual averaging (RAMDA) for structured neural networks. The method uses the preconditioning matrix to accelerate the convergence of a regularized momentum dual averaging (RMDA) method at the price of requiring the local solver (e.g. standard proximal gradient methods) to solve the subproblem. By the preconditioning matrix inspired by the AdaGrad stepsizes in Eq. (2), RAMDA outperforms RMDA and other existing gradient-based methods for solving structured neural networks in various learning applications.", "questions": "I listed questions as part of weaknesses.", "limitations": "The limitations of current theoretical results are clearly and fairly discussed after Theorem 2.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Theoretical results suggest the convergence towards the solution to the subproblem when the proximal gradient methods are used as the local solvers, and almost surely convergence of RAMDA that derives from the manifold theory under the standard $L$-smoothness assumption on the objective functions $f$.   \n\n2. Empirical results illustrate the superior performance of RAMDA over RMDA and other existing gradient-based methods for various neural network tasks. Clear criteria, e.g. for solving the subproblems, are clearly stated in the numerical experiments.", "weaknesses": "1. I think there is an error in Eq. (3) where it should be the square root $\\sqrt{\\cdot}$ in the diagonal operator for $P^t$. This is because $P^t$ uses $U^t$ that is computed from the element-wise multiplicative product of the gradient $G^t$. Is $P^t$ inspired by the AdaGrad stepsizes? If so, then adding the justifications on using $P^t$ in RAMDA is worthwhile to better distinct RAMDA from RMDA. \n2. In the experiments, can you comment on the impact of the different $\\epsilon_t$ on the training performance of RAMDA? Because I believe that using $\\epsilon_t$ a bit higher than $10^{-8}$ set in your experiments RAMDA might achieve far lower training time than other methods while keeping still comparable perplexity to solve Transformer-XL with WikiText-103 in Table 4, or Tacotron2 with LJSpeech in Table 5.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for the careful evaluation of our paper and the invaluable suggestions. Our response is as follows.\n\n**Q1.**\n\n Equation 3 is indeed correct that we are using the cube root.\n  This choice of the preconditioner follows the empirical success of the MADGRAD algorithm of Defazio & Jelassi (2022) for smooth \n  optimization, and our proof also indicates that using the cube root still leads to convergence.\n  More discussion about the reasoning of using the cube root instead can be found in Sec 5.5.4 of Defazio & Jelassi (2022), but in \n  general we consider that the numerical results as more strong evidence than those verbal explanations.\n  On the other hand, even using the cube root, $P^t$ is indeed still inspired by AdaGrad to accumulate the previous gradient norm \n  squared. We will add discussion about AdaGrad in our revision\n\n**Q2.**\n\nIt is expected that with a looser subproblem stopping condition, we could reduce some running time of RAMDA. However, it is quite \n  unlikely that it can be faster than other methods: all methods share some common operations including calculation of the gradient \n  (which is usually the major bottleneck) and several vector additions (for updating the iterates and the accumulated gradient norm) per \n  iteration, and using a looser subproblem stopping condition would not affect these parts. The best we can hope for is that RAMDA and \n  ProxGen would be no slower than other methods that do not involve solving a subproblem with no closed-form solution.\n\n  With further investigation, we found that In Tacotron2 with LJSpeech, the time spent on solving the subproblem is very little, at 2.7% \n  and 1.6% of the total training time for RAMDA and ProxGen, respectively.\n  This suggests that loosening the subproblem stopping condition would provide only very limited benefit.\n  For Transformer-XL with WikiText-103, the time percentages spent on solving the subproblems are 10% for RAMDA and 5.9% for \n  ProxGen, indicating a rather high burden.\n  Consequently, for this problem, we tried to increase the stopping threshold from $10^{-8}$ to $10^{-6}$, which then resulted in a 20.2% \n  reduction in the subproblem time for RAMDA, while that of ProxGen almost remained unchanged (which suggests that probably there \n  is a sharp decrease in the objective improvement so that it reached both the old and the new stopping conditions at the same \n  iteration). As you conjectured, Both RAMDA and ProxGen with this looser stopping condition produced validation perplexity and \n  weighted structured sparsity comparable the original ones. However, as mentioned above, the running time difference is actually not that huge. (For ProxGen it is  likely that the algorithm behavior actually did not change given that the subproblem time remains unchanged.)", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "WjNxHPLIBa", "reviewer": "Reviewer_4z2u", "summary": "This article proposes an optimization algorithm RAMDA for training structured neural networks, which combines a number of optimization techniques including dual averaging, momentum, and coordinate-wise preconditioners. Similar to the existing RMDA algorithm, RAMDA also has the capacity to identify the local manifold structure of the solution. The author(s) provide theoretical analyses to justify the convergence property of RAMDA, and develop an inexact subproblem solver as required by RAMDA.", "questions": "See the \"Weaknesses\" section.", "limitations": "The authors have discussed the limitations of the proposed algorithm in Section 5 and Appendix C.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The proposed RAMDA algorithm extends the existing RMDA algorithm by adding a coordinate-wise preconditioner, and its theoretical analysis seems to be novel.", "weaknesses": "I think one major weakness of the current manuscript is the **correctness** of some theoretical results presented in the article.\n\n1. Theorem 1 suggests that the regularizer function $\\psi$ can be nonconvex. However, as the RAMDA algorithm heavily relies on the proximal operator of $\\psi$, how do you define the proximal operator when $\\psi$ is nonconvex? For example, equation (6) is used to define the new iterate $W^t$, but when $\\psi$ is nonconvex, it is likely that the \"argmin\" is a set and is not uniquely defined.\n\n2. Taking a closer look at the proof of Theorem 1, I feel that the author(s) may have a misunderstanding of an existing theorem. In Appendix B, equation (11) is obtained by citing Theorem 10.15 of [1]. However, Theorem 10.15 of [1] applies to functions of the form $F(x)=f(x)+\\psi(x)$, where $f$ is smooth and nonconvex, but $\\psi$ is convex. In other words, the non-convexity only applies to the smooth part, not the regularizer.\n\n3. If the findings above are valid, then the author(s) may need a thorough examination of the technical proofs to see if there is any error.\n\n4. If we assume $\\psi$ is convex, then there should be an Nesterov-accelerated version of Algorithm 2 that converges in $O(\\varepsilon_t^{-1/2})$ iterations, which is faster than the rate given in Theorem 1.\n\n[1] Beck, A. (2017). First-order methods in optimization. Society for Industrial and Applied Mathematics.\n\n\n===============================================================\n\nEdit: during the rebuttal the author(s) seem to have addressed the concerns above.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for the careful checking of our paper, including the proofs.\nIn current days, it has been increasingly rare to see a reviewer spending this much effort and time to review papers, and we truly appreciate it.\nAnd we thank the reviewer for pointing out our careless errors, which provides us a chance to correct them and obtain better guarantees.\nOur reply is as follows.\n\n**Q1**\n\nIt is indeed correct that when the regularizer is nonconvex, the argmin could be a set, and thus the proximal operator becomes a set operator.\nThis is already reflected in our proof of Thm 1 (see the eq between Lines 544 and 545 on P.14).\nIn this case, we can select any minimizer from the set of minima and the analysis would not be affected. (In the main paper, since we are only approximately solving the problem, whether the real solution is unique does not matter in the algorithmic description.)\nThis is the standard case in analysis involving the proximal operator of a nonconvex function. See, for example, the canonical work of Attouch et al. (2013).\n\n**Q2 & Q3**\n\nOur result is correct, but indeed the citation is incorrect and one minor assumption is missing, and thank you for pointing it out.\nA better reference (although without a formal theorem but just some discussion) would be section 5.1 of Attouch et al. (2013).\nAnd the additionally required assumption is that the objective function is lower-bounded by some finite value (which does not necessarily require convexity of the regularizer).\nFor completeness, we provide a proof for the claimed rate in (11) here.\n\nConsider applying proximal gradient with stepsize $\\theta_t$ to a function $\\bar Q_t(Z) := f_t(Z) + \\psi (Z)$, where $f_t$ is $L$-smooth and $\\psi$ comes with an easily computable proximal operator.\nLet the iterates be updated by\n\n$$Z^{j+1} \\in \\arg\\min_{Z} \\left(\\hat Q_t(Z; Z^j) := \\langle \\nabla f_t(Z^j), Z - Z^j \\rangle + \\frac{1}{2 \\theta_t} \\|Z - Z^j \\|^2 + \\psi(Z)\\right),$$\n\nthen from that $Z^{j+1}$ is a minimizer of $\\hat Q_t(Z; Z^j)$, we know that\n$\\hat Q_t(Z^{j+1}; Z^j) \\leq \\hat Q_t(Z^j; Z^j)$, which implies\n\n$$\\langle \\nabla f_t(Z^j), Z^{j+1} - Z^j \\rangle + \\frac{1}{2 \\theta_t} \\|Z^{j+1} - Z^j \\|^2 + \\psi(Z^{j+1}) \\leq \\psi(Z^j).$$\n\nOn the other hand, from the $L$-smoothness of $f_t$, we know that\n\n$$f_t(Z^{j+1}) - f_t(Z^j) \\leq \\langle \\nabla f_t(Z^j), Z - Z^j \\rangle + \\frac{L}{2} \\|Z - Z^j \\|^2.$$\n\nAdding these two inequalities together, we obtain that\n\n$$\\bar Q_t(Z^j) - \\bar Q_t(Z^{j+1}) \\geq (\\frac{1}{2 \\theta_t} - \\frac{L}{2}) \\|Z^{j+1} - Z^j\\|^2.$$\n\nBy taking $\\theta_t = 1/(2L)$ (as we did in the first line of Alg 2), summing the above inequality for $j=0,1,2,\\dotsc,T$, and noting that $\\bar Q_t(Z^0) - \\bar Q_t(Z^{T+1})$ is upper bounded by $\\bar Q_t(Z^0) - \\bar Q_t^*$, we obtain our claimed rate in equation (11).\n\n**Q4**\n\nFor the convex case, now we realize that the rate can actually be improved to $O(\\log \\epsilon_t^{-1})$. The idea is that when the regularizer is convex, $\\hat Q_t$ is actually strongly convex.\nIndeed, from our construction of $P^t$, we know that the smallest diagonal entry is at least $\\epsilon$, and therefore the subproblem is $\\epsilon$-strongly convex and the smooth part is $L$-smooth, where $L = \\max(diag(P^t)) + \\epsilon$.\nWe denote the global optimum by $\\bar Q^*$ and $\\kappa := L/\\epsilon$.\nWe know that $\\bar Q_t(Z^j) - \\bar Q_t(Z^{j+1})$ is upper bounded by $\\bar Q_t(Z^j) - \\bar Q_t^*$, which converges at a linear rate (see, for example, equation (3.4) of Drusvyatskiy & Lewis (2018)):\n\n$$ \\bar Q_t(Z^j) - \\bar Q_t (Z^{j+1}) \\leq \\bar Q_t(Z^j) - \\bar Q_t^* \\leq \\left(1 - \\frac{1}{2\\kappa}\\right)^j(\\bar Q_t(Z^0) - \\bar Q_t^*).$$\n\nOn the other hand, with the stepsize being $1/L$, we know from Theorem 1 of Nesterov (2013) that \n$$\\bar Q_t(Z^j) - \\bar Q_t(Z^{j+1}) \\geq \\frac{L}{2} \\|Z^j - Z^{j+1}\\|^2.$$\nCombination of these two results then gives that it takes at most $O(\\log \\epsilon_t^{-1})$ iterations to satisfy (7).\n\nOn the other hand, for a convex $\\psi$, accelerated proximal gradient has the same dependency on $\\epsilon_t$, although with a possibly faster linear convergence rate.\nHowever, accelerated proximal gradient does not guarantee that the objective value is always decreasing, so the second condition in (7) might be violated.\nThis is the why we did not apply it as our subproblem solver.\n\n\n**References:**\n\nAttouch, Hedy, Jérôme Bolte, and Benar Fux Svaiter. \"Convergence of descent methods for semi-algebraic and tame problems: proximal algorithms, forward–backward splitting, and regularized Gauss–Seidel methods.\" Mathematical Programming 137.1 (2013): 91-129.\n\nNesterov, Yu. \"Gradient methods for minimizing composite functions.\" Mathematical programming 140.1 (2013): 125-161.\n\nDrusvyatskiy, Dmitriy, and Adrian S. Lewis. \"Error bounds, quadratic growth, and linear convergence of proximal methods.\" Mathematics of Operations Research 43.3 (2018): 919-948.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We thank the reviewers for their careful evaluations of our paper.
We will individually reply to the reviewers to address their specific questions, and here we would like to highlight some changes that we will make in our revision.

- For Theorem 1, the rate when $\psi$ is convex can be improved to $O(\log \epsilon_t^{-1})$. This is from the unutilized observation that the objective function is actually strongly convex, and thus proximal gradient is expected to have a linear rate of convergence instead of the current slower sublinear one.

- In the proof of Theorem 1, reviewer 4z2u pointed out that there is a wrong citation for the case of nonconvex $\psi$. A better reference would be Section 5.1 of Attouch et al. (2013). We will also add the derivation to show that indeed this rate is correct, just the citation was wrong. More details are in our response to reviewer 4z2u.

- There is a typo in our proof of Theorem 1: in (16), (21), and the unnumbered
  equation before (21), 

$$\nabla E_{\xi \sim D} [ f_{\xi}( W^{t})] \circ \nabla E_{\xi\sim D} [ f_{\xi} ( W^{t}) ]$$

should instead be
  $$E_{\xi\sim D} t[ \nabla f_{\xi} ( W^{t} )   \circ \nabla f_{\xi}( W^{t} ) ].$$

**Reference:**

Attouch, Hedy, Jérôme Bolte, and Benar Fux Svaiter. "Convergence of descent methods for semi-algebraic and tame problems: proximal algorithms, forward–backward splitting, and regularized Gauss–Seidel methods." Mathematical Programming 137.1 (2013): 91-129.

Meta-review: All reviewers reach a consensus that the paper is worthy of publication. I suggest the authors to incorporate their general responses and necessary replies during the discussion to the revision of the paper.
Decision: Accept (poster)