Title: Regularized Adaptive Momentum Dual Averaging with an Efficient Inexact Subproblem Solver for Training Structured Neural Network
Venue: NeurIPS 2024
Accepted: yes
Keywords: structured neural networks

variance reduction

manifold identification

proximal methods

adaptive methods

inexact subproblem solution
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

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/ .

1 Introduction

Since the recent emergence of ChatGPT, large language models (LLMs) and other huge deep learning models have garnered much attention and popularity, even among the public who are unfamiliar with machine learning. A challenge with such gigantic neural network models is their vast number of model parameters, reaching hundreds of billions, resulting in expensive storage and inference. It thus becomes crucial to find ways to exploit structures in

trained models to reduce their spatial and prediction costs without degrading the prediction performance. An active line of research is to explicitly add a nonsmooth regularization term to the training objective function and apply proximal stochastic (sub)gradient methods, with or without a diagonal preconditioner for adaptiveness, to induce a pre-specified type of desirable structure in the final model [49, 51, 9]. 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, [16] oberserd empirically that the structures obtained by those methods are highly suboptimal and unstable over iterations. They then proposed a regularized dual averaging method called RMDA , and proved that after a finite number of steps, the iterates of RMDA can stably identify the locally optimal structure induced by the regularizer at the stationary point of asymptotic convergence. 1 This is up to our knowledge the only method with such structure guarantees for training structured neural networks. With this property, their experiments demonstrated that their method also empirically outperforms existing methods on modern computer vision tasks. However, since RMDA does not incorporate adaptiveness and their experiments are conducted only on medium-scale image classification problems, its usefulness beyond computer vision is in doubt.

For a wide range of tasks in deep learning such as language modeling and speech recognition, researchers have developed numerous architectures to achieve state-of-the-art prediction performance, including the transformer [44] 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 architectures 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-art [10, 1, 52, 28, 22]. It is hence expected that the non-adaptive RMDA of [16] might not lead to promising results for such widely-used architectures and tasks.

infinitesimal inexactness could hinder structure identification. Even without inexactness, finding a preconditioner that leads to structure identification guarantees is already difficult because no adaptive algorithm, even in the much simpler deterministic and exact setting, is known to have such a guarantee. (iii) Subproblem solver : Our goal is a practical algorithm, so we need to solve the subproblem efficiently. This requires the inexact measure be checkable and the degree quickly attainable by a well-designed solver, and the preconditioner should make the subproblem well-conditioned and cannot complicate the computation of the solver.

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

We summarize our main contributions as follows.

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

Efficient 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 a key step for realizing existing frameworks for regularized adaptive methods.

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

2 Related Work

Dual Averaging for Deep Learning. Our method is motivated by [16] that adapted the famous regularized dual averaging [47, 25] approach with momentum to train structured neural network models with data augmentation. They selected dual averaging for the gradient estimation to achieve variance reduction for structure guarantees, but their algorithm does not allow for adaptiveness. Inspired by this approach, we also take dual-averaging-like updates for the diagonal preconditioner in the subproblem for adaptiveness. Our preconditioner design also borrows ideas from the empirically successful MADGRAD of [7] for training non-regularized neural networks. RAMDA can thus also be seen as a generalization of MADGRAD to the regularized setting. Since no regularizer is present, unlike RAMDA , the subproblem of MADGRAD has a closed-form solution and no structure is expected. Moreover, [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 in the nonconvex nonsmooth regularized case, and (iii) structure identification guarantees are novel and closer

to properties desirable in practice. The first two items are also applicable when no regularizer is present, so our theory also expands guarantees for MADGRAD .

Regularized Stochastic Algorithms for Deep Learning. Other than RMDA , there are several works on training structured neural networks through regularization and its proximal operator, but none have structure guarantees. [49] considered a simple regularized SGD method with momentum, but their convergence analysis is only for the nonadaptive case. [51] studied a general regularized adaptive framework ProxGen that incorporates diagonal preconditioners, and showed that the subgradient of the objective function can decrease to the reciprocal of the batch size, but their result does not guarantee further convergence to stationary points. Moreover, they do not allow inexactness in the subproblem, so their framework can be realized for only a small class of problems. [9] proposed ProxSSI that extends ProxGen to the case of group-sparsity regularizers, whose corresponding subproblem indeed has no closed-form solution. They applied the Newton-Raphson method to obtain nearly-optimal subproblem solutions, and proposed a seemingly mild inexactness condition. Unfortunately, their condition is not checkable, and their corresponding convergence guarantee requires the regularizer to be locally smooth around each iterate, which excludes most regularizers that induce meaningful structures. On the other hand, we will show that with our implementable inexactness condition, ProxGen still possesses the same convergence guarantees in [51] 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.

Structure and Manifold Identification. The major tool for our structure guarantees is the theory of manifold identification [12, 13, 27, 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 a sequence of points converges to this stationary point with their corresponding subgradients decreasing to zero, this sequence is guaranteed to eventually stay in this manifold, thus identifying the structure. [25, 37, 16] have leveraged this tool to show manifold identification for various stochastic algorithms, and the common key, as pointed out by [37], is variance reduction. Our analysis uses a 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.

where E is a Euclidean space with inner product 〈· , ·〉 and its induced norm ‖·‖ , D is a distribution over a space Ω representing all possible data modifications, f ξ is differentiable almost everywhere for any ξ , and the possibly nonsmooth regularizer ψ ( W ) is for promoting a 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 E = R n in this section. At the t th iteration with learning rate η t and iterate W t -1 , we first draw an independent and identically distributed sample ξ t ∼ D , compute the stochastic (sub)gradient G t := ∇ f ξ t ( W t -1 ) of the loss function at the current point W t -1 with respect to ξ t , and then update the weighted sum V t of historical stochastic gradients and the weighted sum U t of their squared norms using the value s t :

where ◦ denotes the Hadamard (pointwise) product in E . We then construct the preconditioner P t and the weight sum α t by

where /epsilon1 > 0 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.

where W 0 is the initial point. Details regarding (4) and how to solve it are deferred to Section 4. The iterate is then updated by averaging ˆ W t and W t -1 with some c t ∈ [0 , 1]:

/negationslash

The choice of P t in (3) that uses the accumulated square of the stochastic gradient norm as the preconditioner is the key to adaptivity and is widely seen in adaptive methods such as Adagrad [11], while the choice of the cubic root instead of the square root is motivated by the impressive numerical performance of MADGRAD of [7] for smooth problems without a regularization term. The averaging step in (5) with c t = 1 can be interpreted as incorporating a momentum term in the non-regularized non-adaptive case [43, 19].

4 Subproblem Solver

Given an iterate W t -1 , a momentum term m t , a preconditioner P t , and a stepsize η t , existing regularized adaptive stochastic gradient algorithms for (1) can be summarized in the following form [51]:

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

We propose to apply a few iterations of proximal gradient (PG) [see, e.g. , 5, 33] to approximately 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 in (4), we require

for some pre-specified /epsilon1 t , where ∂Q t ( W t +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 Q t replaced by ˆ Q t and ˆ W t by W t . We focus on the case of diagonal and positive P t ,

and thus the largest eigenvalue max(diag( P t )), 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 guaranteeing sufficient objective decrease. For cases in which this value is difficult to obtain, one can apply a simple backtracking linesearch for the subproblem to find a suitable step size efficiently. This PG subproblem solver is summarized in Algorithm 2. To guarantee convergence for both our algorithm and the framework of [51], our analysis in Section 5 requires that { /epsilon1 t } satisfy

We will show in Section 5 that (7) holds after at most O ( /epsilon1 -2 t ) iterations of Algorithm 2.

5 Analysis

This section discusses theoretical guarantees for RAMDA and the proposed subproblem solver in Algorithm 2. We also prove convergence guarantees for applying PG to approximately solve (6) for the framework of [51]. All proofs are in the appendices. Some of our results are inspired by [16], but with the added inexactness in (4) and the adaptiveness for the preconditioner, the analysis is nontrivial. Recall that we assume that f ξ 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 2 and that the point of convergence of RAMDA is almost surely a stationary point.

Theorem 2. Consider { ˆ W t } generated by Algorithm 1 for (1) , with (7) and { c t } and { /epsilon1 t } satisfying ∑ c t = ∞ and (8) . Assume there is L ≥ 0 such that for any ξ , f ξ is almost surely L -Lipschitz-continuously-differentiable, so the expectation is also L -Lipschitz-continuouslydifferentiable, there is C ≥ 0 such that E ξ t ∼D ∥ ∥ ∇ f ξ t ( W t -1 )∥ ∥ 4 ≤ C for all t , and that the set of stationary points Z := { W | 0 ∈ ∂F ( W ) } is nonempty. For any given W 0 , consider the event that { ˆ W t } converges to a point ¯ W (each event corresponds to a different ¯ W ). If ∂ψ is outer semicontinuous at ¯ W , this event has a nonzero probability, and { η t } satisfies

Theorem 1. Assume that (4) and (6) has at least one optimal solution with a finite optimal objective value. Given /epsilon1 t > 0 , the number of iterations Algorithm 2 takes to satisfy (7) for both (4) and (6) is O (log( /epsilon1 -1 t )) when ψ is convex and O ( /epsilon1 -2 t ) when ψ is nonconvex.

then we have that ¯ W ∈ Z with probability one conditional on this event. Moreover, { W t } also converges to this stationary point ¯ W .

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

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

Definition 1 (Partial Smoothness [26, 12]) . A function ψ is partly smooth at a point ¯ W relative to a set M ¯ W /owner ¯ W if

/negationslash

ψ is regular (finite and the Fréchet subdifferential coincides with the limiting Fréchet subdifferential) at all points W ∈ M ¯ W near ¯ W , with ∂ψ ( W ) = ∅ .

Around ¯ W , M ¯ W is a C 2 -manifold and ψ | M ¯ W is C 2 .

The affine span of ∂ψ ( ¯ W ) is a translate of the normal space to M ¯ W at ¯ W .

∂ψ is continuous at ¯ W relative to M ¯ W .

We often call M ¯ W the active manifold at ¯ 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 ψ is prox-regular at ¯ W for V ∗ ∈ ∂ψ ( ¯ W ) if ψ is locally lower semi-continuous around ¯ W , finite at ¯ W , and there is ρ > 0 such that ψ ( W 1 ) ≥ ψ ( W 2 )+ 〈 V, W 1 -W 2 〉-ρ 2 ‖ W 1 -W 2 ‖ 2 for every W 1 , W 2 near ¯ W with ψ ( W 2 ) close to ψ ( ¯ W ) and V ∈ ∂ψ ( W 2 ) close to V ∗ . ψ is prox-regular at ¯ W if it is prox-regular for all V ∈ ∂ψ ( ¯ W ) .

Theorem 3. Consider Algorithm 1 with the conditions in Theorem 2 satisfied. Consider the event of { ˆ W t } converging to a 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 2 manifold M ¯ W ; ∂ψ is outer semicontinuous at ¯ W ; and the nondegeneracy condition -∇ f ( ¯ W ) ∈ relative interior of ∂ψ ( ¯ W ) holds at ¯ W , then conditional on this event, almost surely there is T 0 ≥ 0 such that

We note particularly that convex and weakly-convex [34] functions are all regular, proxregular, and subdifferentially continuous everywhere.

We also show that our subproblem solver and condition can be effectively applied to the framework of [51] while retaining the same convergence guarantees. As mentioned in Section 2, our result is much stronger than that of [9] for having no unrealistic smoothness requirement on ψ and using an implementable inexactness condition.

Theorem 4. For the framework in [51] with the subproblem solved approximately by Algorithm 2 such that (7) holds with { /epsilon1 t } satisfying (8) . Then Theorem 1 of [51] still holds, but with the constants { Q i } being also dependent on ¯ /epsilon1 .

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 in our experiments. Particularly, we employ the group LASSO regularization [50] 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 PG for both RAMDA and existing regularized adaptive methods. We then consider tasks in computer vision, language modeling, and speech to compare the following algorithms using Pytorch.

RAMDA : The proposed Algorithm 1.

ProxSGD [49]

RMDA [16]

ProxGen [51]: We follow their experiments to use AdamW and apply our PG as the subproblem solver.

ProxSSI [9]

These algorithms are introduced in Section 2 and also further summarized in Appendix A. For each task, we also provide for reference a baseline that does not include a group LASSO regularizer in the training (SGD with momentum ( MSGD ) for computer vision, and Adam

No early stopping

Early stopping

for the other two), but our comparison is only among those for training structured models. Our code for reproducing the experiments and the hyperparameter settings are available at https://github.com/ismoptgroup/ramda_exp/ . Additional details of the stability of the structure (level of structured sparsity here) over epochs of RAMDA are available in Appendix D.

We use two criteria for comparison: 1. Model predictive ability, and 2. Structured sparsity level. The former is task-dependent and thus specified in each experiment. Regarding the latter, sparsifying neural networks while preserving its performance requires prior knowledge of model design. A common approach is retaining certain parameters during the training process, and we adhere to this convention such that the bias, batch normalization [17], layer normalization [3], and embedding layers do not have any sparsity-inducing regularization imposed on them [9, 35]. For the rest, we adopt channel-wise grouping for convolutional layers, input-wise grouping for fully-connected and LSTM layers during the training phase. For evaluation, our structured sparsity is calculated using the weighted group sparsity with the weights proportional to the number of parameters in each group.

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

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

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

Image Classification We conduct a classical computer vision task of training ResNet50 [14] with the ILSVRC 2012 ImageNet dataset [39]. The result in Table 3 shows that RAMDA attains the best validation accuracy and structured sparsity simultaneously.

Language Modeling For language modeling, we train Transformer-XL (base) [6] using the WikiText-103 dataset [31]. Transformer-XL is comprised of embedding and non-embedding layers, and in the PyTorch implementation, the non-embedding layers are built using linear

Algorithm

Sparsity

Time

Sparsity

Time

Algorithm

Sparsity

and layer-normalization layers. We apply group LASSO regularization to the linear layers, and present in Table 4 the perplexity and the weighted group sparsity of the models trained. We see that RAMDA gives the best perplexity and structured sparsity simultaneously.

Speech Synthesis We consider Tacotron2 [41] for speech synthesis on the LJSpeech dataset [18]. We apply regularization to the convolutional, LSTM, and linear layers of Tacotron2 and show the results in Table 5. Clearly, RAMDA gives the lowest validation loss and the highest group sparsity.

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

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 on modern language modeling and speech synthesis tasks as well as the ImageNet problem. We also observe from the plots in the appendices that it is possible to further improve the sparsity level of RAMDA if we run it for more epochs.

7 Conclusions

In this work, we proposed a regularized dual averaging method with adaptiveness, RAMDA , for training structured neural networks. Our method outperforms state of the art on modern

Alg.

Perplexity

Sparsity

Time/epoch

Alg.

Loss

Sparsity

Time/epoch

architectures including LSTM and transformers as well as the ImageNet problem. We also proposed a subroutine with strong convergence guarantees to approximately solve the regularized subproblem of both our method and an existing framework efficiently. Extensive experiments on group sparsity showed that our subproblem solver can greatly reduce the training time for existing methods, and our proposed RAMDA achieves simultaneously higher structured sparsity ratio and better prediction performance than existing methods. Implementation of our method is available at https://www.github.com/zhisyuan1214/ RAMDA/ .

References

Rohan Anil, Vineet Gupta, Tomer Koren, and Yoram Singer. Memory efficient adaptive optimization. In Advances in Neural Information Processing Systems , 2019. 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. Mathematical programming , 137(1-2):91-129, 2013. 15

Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. Technical report, 2016. arXiv:1607.06450. 8

Amir Beck. First-Order Methods in Optimization . SIAM - Society for Industrial and Applied Mathematics, Philadelphia, PA, United States, 2017. 15

Amir Beck and Marc Teboulle. A fast iterative shrinkage-thresholding algorithm for linear inverse problems. SIAM Journal on Imaging Sciences , 2(1):183-202, 2009. 5

Zihang Dai, Zhilin Yang, Yiming Yang, Jaime Carbonell, Quoc Le, and Ruslan" Salakhutdinov. Transformer-XL: Attentive language models beyond a fixed-length context. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics , 2019. 8

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):1-34, 2022. 3, 5, 14

Alexandre Défossez, Leon Bottou, Francis Bach, and Nicolas Usunier. A simple convergence proof of adam and adagrad. Transactions on Machine Learning Research , 2022. ISSN 2835-8856. 2

Tristan Deleu and Yoshua Bengio. Structured sparsity inducing adaptive optimizers for deep learning. Technical report, 2021. arXiv:2102.03869. 2, 4, 7, 8

Michael Denkowski and Graham Neubig. Stronger baselines for trustable results in neural machine translation. Technical report, 2017. arXiv:1706.09733. 2

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

Warren L. Hare and Adrian S. Lewis. Identifying active constraints via partial smoothness and prox-regularity. Journal of Convex Analysis , 11(2):251-266, 2004. 4, 7

Warren L. Hare and Adrian S. Lewis. Identifying active manifolds. Algorithmic Operations Research , 2(2):75-75, 2007. 4

Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In IEEE conference on computer vision and pattern recognition , 2016. 8, 20

Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation , 9(8):1735-1780, 1997. 2

Zih-Syuan Huang and Ching-pei Lee. Training structured neural networks through manifold identification and variance reduction. In International Conference on Learning Representations , 2022. 2, 3, 4, 5, 6, 7, 14

Sergey Ioffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In International conference on machine learning , pages 448-456, 2015. 8

Keith Ito and Linda Johnson. The LJ speech dataset, 2017. 9

Samy Jelassi and Aaron Defazio. Dual averaging is surprisingly effective for deep learning optimization. Technical report, 2020. arXiv:2010.10502. 5, 14

Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In International Conference on Learning Representations , 2015. 2

Alex Krizhevsky. Learning multiple layers of features from tiny images. Technical report, 2009. 20

Frederik Kunstner, Jacques Chen, Jonathan Wilder Lavington, and Mark Schmidt. Noise is not the main factor behind the gap between sgd and adam on transformers, but sign descent might be. In International Conference on Learning Representations , 2023. 2

Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient-based learning applied to document recognition. Proceedings of the IEEE , 86(11):2278-2324, 1998. 20

Ching-pei Lee. Accelerating inexact successive quadratic approximation for regularized optimization through manifold identification. Mathematical Programming , 2023. 2, 4, 17

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, 4

Adrian S. Lewis. Active sets, nonsmoothness, and sensitivity. SIAM Journal on Optimization , 13(3):702-725, 2002. 7

Adrian S. Lewis and Shanshan Zhang. Partial smoothness, tilt stability, and generalized hessians. SIAM Journal on Optimization , 23(1):74-94, 2013. 4

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 Language Processing (EMNLP) , 2020. 2

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, 22

Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Conference on Learning Representations , 2019. 14

Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models. In International Conference on Learning Representations , 2017. 8

Yurii Nesterov. Primal-dual subgradient methods for convex problems. Mathematical programming , 120(1):221-259, 2009. 3

Yurii Nesterov. Gradient methods for minimizing composite functions. Mathematical Programming , 140(1):125-161, 2013. 5

Evgeni Alekseevich Nurminskii. The quasigradient method for the solving of the nonlinear programming problems. Cybernetics , 9(1):145-150, 1973. 7

Alexandra Peste, Eugenia Iofinova, Adrian Vladu, and Dan Alistarh. Ac/dc: Alternating compressed/decompressed training of deep neural networks. In Advances in neural information processing systems , 2021. 8

René Poliquin and R Rockafellar. Prox-regular functions in variational analysis. Transactions of the American Mathematical Society , 348(5):1805-1838, 1996. 7

Clarice Poon, Jingwei Liang, and Carola-Bibiane Schönlieb. Local convergence properties of SAGA/prox-SVRG and acceleration. In International Conference on Machine Learning , 2018. 3, 4

R. Tyrrell Rockafellar and Roger J-B Wets. Variational analysis , volume 317. Springer Science & Business Media, 2009. 5

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 , 115(3):211-252, 2015. 8

Andrzej Ruszczyński. Feasible direction methods for stochastic programming problems. Mathematical Programming , 19:220-229, 1980. 4, 16

Jonathan Shen, Ruoming Pang, Ron J Weiss, Mike Schuster, Navdeep Jaitly, Zongheng Yang, Zhifeng Chen, Yu Zhang, Yuxuan Wang, Rj Skerrv-Ryan, Rif A. Saurous, Yannis Agiomyrgiannakis, and Yonghui Wu. Natural tts synthesis by conditioning wavenet on mel spectrogram predictions. In 2018 IEEE international conference on acoustics, speech and signal processing (ICASSP) , 2018. 9

Karen Simonyan and Andrew Zisserman. Very deep convolutional networks for largescale image recognition. In International Conference on Learning Representations , 2015. 20

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 cybernetics , 50(2):835-845, 2018. 5

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in neural information processing systems , 2017. 2

Wei Wen, Chunpeng Wu, Yandan Wang, Yiran Chen, and Hai Li. Learning structured sparsity in deep neural networks. Advances in neural information processing systems , pages 2074-2082, 2016. 7

Appendices

Table of Contents

A More Experiment Details

This section describes details of our implementation of RAMDA and the setting of the experiments conducted in Section 6.

A.1 Implementation and Hyperparameter Setting of RAMDA

Similar to [16], we introduce a restart strategy to the implementation of RAMDA . During each stage of the training, the learning rate η t and the momentum factor c t are fixed. Once the epoch count enters the next stage, we reset the counter t to 1 and use the output parameter W T from the previous round as the new input parameter W 0 to the same algorithm, set α t , V t and U t to 0, but keep the scheduling for η and c going without resetting them, and decrease /epsilon1 by a factor. We initialize c t as either 0 . 1 or 0 . 01, depending on the problems, and use a constant c t until the final stage, where we gradually increase it by

where i 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 6, we summarize details of the algorithms compared in Section 6.

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, we 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 1. Assume that (4) and (6) has at least one optimal solution with a finite optimal objective value. Given /epsilon1 t > 0 , the number of iterations of Algorithm 2 takes to satisfy (7) for both (4) and (6) is O ( /epsilon1 -1 t ) when ψ is convex and O ( /epsilon1 -2 t ) when ψ is nonconvex.

Proof. We use the notation to unify the two objective function Q t /α t and ˆ Q t , where f t is the smooth part and we define the Lipschitz constant of ∇ f t as L .

At each iteration, PG solves the following subproblem

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

We thus have from the Lipschitz continuity of ∇ f t that

Note that ¯ Q t is lower bounded, say by ¯ Q ∗ t , and has at least one solution Z ∗ (unique when ψ is convex).

In the case that ψ is convex, we know that θ t = 1 /L , and (3) clearly shows that the subproblem objective ¯ Q is /epsilon1 -strongly convex. Therefore, standard analysis of proximal gradient [see, for example, 4, Lemma 10.4 and iTheorem 10.29] gives that

The combination of (9)-(11) shows that it takes O (log /epsilon1 -1 t ) iterations for PG to reach the required precision of /epsilon1 t .

When ψ is nonconvex, we have that θ t = 1 / (2 L ) and standard analysis [2, Section 5.1] gives

for some constant C depending on L and ¯ Q t ( W t ) -¯ Q ∗ t . Therefore, (12) and (9) show that it takes O ( /epsilon1 -2 t ) iterations to reach the desired precision.

Theorem 2. Consider { ˆ W t } generated by Algorithm 1 for (1) , with (7) and { c t } and { /epsilon1 t } satisfying ∑ c t = ∞ and (8) . Assume there is L ≥ 0 such that for any ξ , f ξ is almost surely L -Lipschitz-continuously-differentiable, so the expectation is also L -Lipschitz-continuouslydifferentiable, there is C ≥ 0 such that E ξ t ∼D ∥ ∥ ∇ f ξ t ( W t -1 )∥ ∥ 4 ≤ C for all t , and that the set of stationary points Z := { W | 0 ∈ ∂F ( W ) } is nonempty. For any given W 0 , consider the event that { ˆ W t } converges to a point ¯ W (each event corresponds to a different ¯ W ). If ∂ψ is outer semicontinuous at ¯ W , this event has a nonzero probability, and { η t } satisfies

then we have that ¯ W ∈ Z with probability one conditional on this event. Moreover, { W t } also converges to this stationary point ¯ W .

Proof. First, we prove that when { ˆ W t } converges to ¯ W , { W t } also converges to ¯ W . From (5), we have that

Since ˆ W t → ¯ W , for any /epsilon1 > 0 we can find an integer t /epsilon1 ≥ 0 such that ‖ ˆ W t -¯ W ‖ ≤ /epsilon1 for all t ≥ t /epsilon1 . Therefore, by deducting /epsilon1 from both sides of (15), we get

By letting t approach infinity and noting that ∑ c t = ∞ , we see that

Because /epsilon1 is arbitrary, we see that ‖ W t -¯ W ‖ → 0, and hence { W t } converges to ¯ W .

Next, consider the update of α -1 t U t , we can see from (2) that

Moreover, the assumptions on η t satisfies all the required conditions of Lemma 1 of [40]. We therefore apply Lemma 1 of [40] to conclude that

The update for α -1 t V t has a form analogous to (16), and we have from Jensen's inequality that implying that the second moment is also bounded in expectation. We can therefore also apply Lemma 1 of [40] to α -1 t V t and conclude that

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

From (4) and (7), we can see that there is a sequence { z t } such that

Our assumption in (14) implies that α t →∞ , which together with (8) leads to

From (18), that ∇ E ξ ∼D [ f ξ ( W )] is Lipschitz continuous, and that W t → ¯ W (which we have proven in the first part), we see that

For the third term, we have from (3) and (17) that

Again since α t →∞ , the second term of the equation above converges to 0. Therefore, by (17), we obtain that

Again from the continuity of ∇ E ξ ∼D [ f ξ ( W )] and that α t →∞ , we conclude that

Finally, using the outer semicontinuity of ∂ψ ( W ) at ¯ W , we conclude from (19)-(22) that with probability one, showing that ¯ W is a stationary point almost surely.

Theorem 3. Consider Algorithm 1 with the conditions in Theorem 2 satisfied. Consider the event of { ˆ W t } converging to a 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 2 manifold M ¯ W ; ∂ψ is outer semicontinuous at ¯ W ; and the nondegeneracy condition

holds at ¯ W , then conditional on this event, almost surely there is T 0 ≥ 0 such that

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

Proof. From (19), there exists a sequence { Y t } such that

For notational ease, we denote

From (25), we then get

We aim to show that

converges to 0 almost surely. From (27), we have

where we get the first inequality from the triangle inequality and the second from (19). According to (18) and (22), there are { A t } and { B t } such that

Substituting the above two equations back to (28), we obtain dist(0 , ∂F ( ˆ W t ))

From Theorem 2, we know that ˆ W t and W t both converge to ¯ W , so

From (8) and (14), we know that /epsilon1 t /α t → 0. Because ˆ W t → ¯ W , we also have that

From W t → ¯ W , we have that

Combining these results with (30), we conclude that

proving that

On the other hand, since f is continuous and ψ is subdifferentially continuous at ¯ W (which implies F is also subdifferentially contnuous at ¯ W ), ˆ W t → ¯ W , and that ∇ f ( ˆ W t ) + Y t a.s. - - → 0 ∈ ∂F ( ¯ W ) (from Theorem 2), we know that F ( ˆ W t ) a.s. - - → F ( ¯ W ) as well. Therefore, we can apply Lemma 1 of [24] to conclude that (24) indeed holds for some T 0 < ∞ almost surely.

Algorithm 3 ProxGen ( W 0 , T, T 2 , { η t } , { ρ t } , { c t } , { /epsilon1 t } , { b t } , δI )

for t = 1 , . . . , T do m 0 ← 0

output: W T

t

Compute W t +1 by roughly solving (6) that satisfies (7) with Q t replaced by ˆ Q t and ˆ W t replaced by W t +1 , using PG ( W t , W t , m t , η -1 t P t , θ t , T 2 , /epsilon1 t )

B.1 Convergence Result for ProxGen

We next discuss the convergence result for the framework of [51] with inexactness 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 assumptions using the notation (26).

(C-1) The expected loss function f is L -Lipschitz-continuously-differentiable and lowerbounded for some L ≥ 0.

(C-2) The stochastic gradient G t = ∇ f ξ t ( W t -1 ) is an unbiased estimator of ∇ f ( W t -1 ) with bounded variance.

where b t is the batch size of ξ t and σ ≥ 0 is a constant.

(C-3) There are some ρ 0 , µ ∈ [0 , 1) and D,G > 0 such that ∥ ∥ W t +1 -W t ∥ ∥ ≤ D , ‖ G t ‖ ≤ G , and ρ t = ρ 0 µ t -1 for all t .

(C-4) There is some γ > 0 such that ∥ ∥ η -1 t P t ∥ ∥ 2 ≤ 1 /γ < ∞ for all t .

(C-5) There is δ > 0 such that

Theorem 4. For the framework in [51] with the subproblem solved approximately by Algorithm 2 such that (7) holds with { /epsilon1 t } satisfying (8) . Then Theorem 1 of [51] still holds, but with the constants { Q i } being also dependent on ¯ /epsilon1 .

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 [51, Lemma 1] that

By the update rule for m t , (6) and (7), we have that there is z t such that

leading to

We thus have from (33) and (C-4) that dist(0 , ∂F ( W t )) 2

We will separately bound the quantities T 1 and ∥ ∥ W t -W t -1 ∥ ∥ 2 below.

From the subproblem objective requirement in (7), we get

From (C-1), we have

Summing (35) and (36) gives

Note that η -1 t P t -LI /followsequal 0 from (31), so the second term in (37) is nonnegative. (37) together with (C-3) then leads to

Summing it over t = 1 , 2 , . . . , T and utilizing the assumption that the step sizes are nonincreasing then give where

From the inequality above, we obtain

for some constants H 1 , H 2 depending on L, ∆ , δ, η 0 , and C 1 . From (37), we also have

Therefore, we obtain

Therefore, we have from (8) and (40) and (C-2) that

for some constants Q 1 , Q 2 , Q 3 dependent on { η 0 , δ, ∆ , L, D, G, ρ 0 , µ, γ, ¯ /epsilon1 } , but not on T . This proves our theorem.

C Additional Experiments for Computer Vision

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

Logistic regression (neural network with no hidden layer) with the MNIST dataset [23].

A modified VGG19 [42] with the CIFAR10 dataset [21].

The same VGG19 with the CIFAR100 dataset [21].

A modified ResNet50 [14] with the CIFAR10 dataset.

The same ResNet50 with the CIFAR100 dataset.

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

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

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:

ResNet50 with the ILSVRC 2012 ImageNet dataset.

Transformer-XL with the WikiText-103 dataset.

Tacotron2 with the LJSpeech dataset.

Logistic Regression with the MNIST dataset.

A modified VGG19 with the CIFAR10 dataset.

The same VGG19 with the CIFAR100 dataset.

A modified ResNet50 with the CIFAR10 dataset.

The same ResNet50 with the CIFAR100 dataset.

In the plot for Transformer-XL, one step processes ten batches, and for our batch size of 64, one epoch consists of 8,401 batches. We further observe in the zoomed-in sparsity

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

E Experiment with Nuclear-norm Regularization

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

and the corresponding proximal operator for λ > 0 is

Given a point X ∗ with rank r ∗ , the active manifold of the nuclear norm at X ∗ is

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

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

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

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

0

20

40

60

80

100

120

140

160

Epochs

0.0

0.2

0.4

0.6

0.8

Validation Accuracy

ResNet50 on ImageNet

MSGD

ProxSGD

ProxGen

RMDA

RAMDA

0

20

40

60

80

100

120

140

160

Epochs

0.0

0.1

0.2

0.3

Weighted Structured Sparsity

(a) ResNet50 on ImageNet

0

10

20

30

40

50

60

70

Every 5000 Steps

20

40

60

80

100

Validation Perplexity

Transformer-XL on WikiText-103

Adam

ProxSGD

ProxGen

RMDA

RAMDA

0

10

20

30

40

50

60

70

Every 5000 Steps

0.0

0.1

0.2

0.3

0.4

Weighted Structured Sparsity

(b) Transformer-XL on WikiText-103

0

200

400

600

800

Epochs

0.50

0.75

1.00

Validation Loss

Tacotron2 on LJSpeech

Adam

ProxSGD

ProxGen

RMDA

RAMDA

0

200

400

600

800

Epochs

0.0

0.2

0.4

0.6

Weighted Structured Sparsity

(c) Tacotron2 on LJSpeech

0

100

200

300

400

500

Epochs

0.86

0.88

0.90

0.92

Validation Accuracy

Logistic Regression on MNIST

ProxSGD

ProxSSI

ProxGen

RMDA

RAMDA

0

100

200

300

400

500

Epochs

0.0

0.2

0.4

0.6

Weighted Structured Sparsity

(d) Logistic Regression on MNIST

0

200

400

600

800

1000

Epochs

0.00

0.25

0.50

0.75

1.00

Validation Accuracy

VGG19 on CIFAR10

MSGD

ProxSGD

ProxSSI

ProxGen

RMDA

RAMDA

0

200

400

600

800

1000

Epochs

0.00

0.25

0.50

0.75

Weighted Structured Sparsity

(e) VGG19 on CIFAR10

0

200

400

600

800

1000

Epochs

0.0

0.2

0.4

0.6

0.8

Validation Accuracy

VGG19 on CIFAR100

MSGD

ProxSGD

ProxSSI

ProxGen

RMDA

RAMDA

0

200

400

600

800

1000

Epochs

0.00

0.25

0.50

0.75

Weighted Structured Sparsity

(f) VGG19 on CIFAR100

0

200

400

600

800

1000

Epochs

0.00

0.25

0.50

0.75

1.00

Validation Accuracy

ResNet50 on CIFAR10

MSGD

ProxSGD

ProxSSI

ProxGen

RMDA

RAMDA

0

200

400

600

800

1000

Epochs

0.00

0.25

0.50

0.75

Weighted Structured Sparsity

(g) ResNet50 on CIFAR10

0

200

400

600

800

1000

Epochs

0.00

0.25

0.50

0.75

Validation Accuracy

ResNet50 on CIFAR100

MSGD

ProxSGD

ProxSSI

ProxGen

RMDA

RAMDA

0

200

400

600

800

1000

Epochs

0.00

0.25

0.50

0.75

Weighted Structured Sparsity

(h) ResNet50 on CIFAR100

130

140

150

Epochs

0.15

0.20

0.25

0.30

0.35

Weighted Structured Sparsity

ResNet50 on ImageNet

ProxSGD

ProxGen

RMDA

RAMDA

(a) ResNet50 on ImageNet

60

65

70

75

Every 5000 Steps

0.28

0.30

0.32

0.34

0.36

0.38

Weighted Structured Sparsity

Transformer-XL on WikiText-103

ProxSGD

ProxGen

RMDA

RAMDA

(b) Transformer-XL on WikiText-103

800

850

900

950

Epochs

0.30

0.35

0.40

0.45

0.50

0.55

Weighted Structured Sparsity

Tacotron2 on LJSpeech

ProxSGD

ProxGen

RMDA

RAMDA

(c) Tacotron2 on LJSpeech

450

460

470

480

490

500

Epochs

0.0

0.1

0.2

0.3

0.4

0.5

0.6

Weighted Structured Sparsity

Logistic Regression on MNIST

ProxSGD

ProxSSI

ProxGen

RMDA

RAMDA

(d) Logistic Regression on MNIST

900

920

940

960

980

1000

Epochs

0.4

0.5

0.6

0.7

0.8

0.9

Weighted Structured Sparsity

VGG19 on CIFAR10

ProxSGD

ProxSSI

ProxGen

RMDA

RAMDA

(e) VGG19 on CIFAR10

900

920

940

960

980

1000

Epochs

0.60

0.65

0.70

0.75

0.80

0.85

Weighted Structured Sparsity

VGG19 on CIFAR100

ProxSGD

ProxSSI

ProxGen

RMDA

RAMDA

(f) VGG19 on CIFAR100

900

920

940

960

980

1000

Epochs

0.80

0.82

0.84

0.86

Weighted Structured Sparsity

ResNet50 on CIFAR10

ProxSGD

ProxSSI

ProxGen

RMDA

RAMDA

(g) ResNet50 on CIFAR10

900

920

940

960

980

1000

Epochs

0.4

0.5

0.6

0.7

0.8

0.9

Weighted Structured Sparsity

ResNet50 on CIFAR100

ProxSGD

ProxSSI

ProxGen

RMDA

RAMDA

(h) ResNet50 on CIFAR100

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 are discussed in Section 5 (after Theorem 2) and Appendix C.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions 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 B.

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 settings are 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.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

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 or other appropriate information about the statistical significance of 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: Time of execution is reported in Section 6, and details of the computer resources are 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 of Ethics https://neurips.cc/public/EthicsGuidelines ?

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 a 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 or models that have a 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 or 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 in the paper well documented and is the documentation provided alongside the assets?

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, as 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 Research 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 on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: The paper does not involve crowdsourcing nor research with human subjects.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nThe proposed RAMDA algorithm extends the existing RMDA algorithm by adding a coordinate-wise preconditioner, and its theoretical analysis seems to be novel.\n\nWEAKNESSES:\nI 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.\n\nQUESTIONS:\nSee the \"Weaknesses\" section.\n\nLIMITATIONS:\nThe authors have discussed the limitations of the proposed algorithm in Section 5 and Appendix C.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-11"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\n1. 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.\n\nQUESTIONS:\nI listed questions as part of weaknesses.\n\nLIMITATIONS:\nThe limitations of current theoretical results are clearly and fairly discussed after Theorem 2.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\n#### 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.\n\nSTRENGTHS:\n#### 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.\n\nWEAKNESSES:\n#### 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.\n\nQUESTIONS:\n#### 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?\n\nLIMITATIONS:\nno", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-13"}