Training Structured Neural Networks Through Manifold Identification and Variance Reduction

1 Introduction . Training neural networks ( NNs ) with regularization to obtain a certain desired structure such as structured sparsity or discrete-valued parameters is a problem of increasing interest . Existing approaches either use stochastic subgradients of the regularized objective ( Wen et al. , 2016 ; 2018 ) or combine popular stochastic gradient algorithms for NNs , like SGD with momentum ( MSGD ) or Adam ( Kingma & Ba , 2015 ) , with the proximal operator associated with the regularizer to conduct proximal stochastic gradient updates to obtain a model with preferred structures ( Bai et al. , 2018 ; Yang et al. , 2019 ; Yun et al. , 2020 ; Deleu & Bengio , 2021 ) . Such methods come with proven convergence for certain measures of first-order optimality and have shown some empirical success in applications . However , we notice that an essential theoretical support lacking in existing methods is the guarantee for the output iterate to possess the same structure as that at the point of convergence . More specifically , often the imposed regularization is only known to induce a desired structure exactly at optimal or stationary points of the underlying optimization problem ( see for example , Zhao & Yu , 2006 ) , but training algorithms are only able to generate iterates asymptotically converging to a stationary point . Without further theoretical guarantees , it is unknown whether the output iterate , which is just an approximation of the stationary point , still has the same structure . For example , let us assume that sparsity is desired , the point of convergence is x∗ = ( 1 , 0 , 0 ) , and two algorithms respectively produce iterates { yt = ( 1 , t−1 , t−1 ) } and { zt = ( 1 + t−1 , 0 , 0 ) } . Clearly , both iterate sequences converge to x∗ , but only zt has the same desired structure as its limit point x∗ , while yt is not useful for sparsity despite that the point of convergence is . This work aims at filling this gap to propose an algorithm for training structured NNs that can provably make all its iterates after a finite number of iterations possess the desired structure of the stationary point to which the iterates converge . We term the structure at a stationary point a stationary structure , and it should be understood that for multiple stationary points , each might correspond to a different stationary structure , and we aim at identifying the one at the limit point of the iterates of an algorithm , instead of selecting the optimal one among all stationary structures . Although finding the structure at an inferior stationary point might seem not very meaningful , another reason for studying this identification property is that for the same point of convergence , the structure at the limit point is the most preferable one . Consider the same example above , we note that for any sequence { xt } converging to x∗ , xt1 6= 0 for all t large enough , for otherwise xt does not converge to x∗ . Therefore , xt can not be sparser than x∗ if xt → x∗.1 Identifying the structure of the point of convergence thus also amounts to finding the locally most ideal structure under the same convergence premise . It is well-known in the literature of nonlinear optimization that generating iterates consistently possessing the structure at the stationary point of convergence is possible if all points with the same structure near the stationary point can be presented locally as a manifold along which the regularizer is smooth . This manifold is often termed as the active manifold ( relative to the given stationary point ) , and the task of generating iterates staying in the active manifold relative to the point of convergence after finite iterations is called manifold identification ( Lewis , 2002 ; Hare & Lewis , 2004 ; Lewis & Zhang , 2013 ) . To identify the active manifold of a stationary point , we need the regularizer to be partly smooth ( Lewis , 2002 ; Hare & Lewis , 2004 ) at that point , roughly meaning that the regularizer is smooth along the active manifold around the point , while the change in its value is drastic along directions leaving the manifold . A more technical definition will be given in Section 3 . Fortunately , most regularizers used in machine learning are partly smooth , so stationary structure identification is possible , and various deterministic algorithms are known to achieve so ( Hare & Lewis , 2007 ; Hare , 2011 ; Wright , 2012 ; Liang et al. , 2017a ; b ; Li et al. , 2020 ; Lee , 2020 ; Bareilles et al. , 2020 ) . On the other hand , for stochastic gradient-related methods to identify a stationary structure , existing theory suggests that the variance of the gradient estimation needs to vanish as the iterates approach a stationary point ( Poon et al. , 2018 ) , and indeed , it is observed empirically that proximal stochastic gradient descent ( SGD ) is incapable of manifold identification due to the presence of the variance in the gradient estimation ( Lee & Wright , 2012 ; Sun et al. , 2019 ) .2 Poon et al . ( 2018 ) showed that variance-reduction methods such as SVRG ( Johnson & Zhang , 2013 ; Xiao & Zhang , 2014 ) and SAGA ( Defazio et al. , 2014 ) that utilize the finite-sum structure of empirical risk minimization to drive the variance of their gradient estimators to zero are suitable for this task . Unfortunately , with the standard practice of data augmentation in deep learning , training of deep learning models with a regularizer should be treated as the following stochastic optimization problem that minimizes the expected loss over a distribution , instead of the commonly seen finite-sum form : min W∈E F ( W ) : = Eξ∼D [ fξ ( W ) ] + ψ ( W ) , ( 1 ) where E is a Euclidean space with inner product 〈· , ·〉 and the associated norm ‖·‖ , D is a distribution over a space Ω , fξ is differentiable almost everywhere for all ξ ∈ Ω , and ψ ( W ) is a regularizer that might be nondifferentiable . We will also use the notation f ( W ) : = Eξ∼D [ fξ ( W ) ] . Without a finite-sum structure in ( 1 ) , Defazio & Bottou ( 2019 ) pointed out that classical variance-reduction methods are ineffective for deep learning , and one major reason is the necessity of periodically evaluating ∇f ( W ) ( or at least using a large batch from D to get a precise approximation of it ) in variance-reduction methods is intractable , hence manifold identification and therefore finding the stationary structure becomes an extremely tough task for deep learning . Although recently there are efforts in developing variancereduction methods for ( 1 ) inspired by online problems ( Wang et al. , 2019 ; Nguyen et al. , 2021 ; Pham et al. , 2020 ; Cutkosky & Orabona , 2019 ; Tran-Dinh et al. , 2019 ) , these methods all have multiple hyperparameters to tune and incur computational cost at least twice or 1See a more detailed discussion in Appendix B.1 . 2An exception is the interpolation case , in which the variance of plain SGD vanishes asymptot- ically . But data augmentation often fails this interpolation condition . thrice to that of ( proximal ) SGD . As the training of deep learning models is time- and resource-consuming , these drawbacks make such methods less ideal for deep learning . To tackle these difficulties , we extend the recently proposed modernized dual averaging framework ( Jelassi & Defazio , 2020 ) to the regularized setting by incorporating proximal operations , and obtain a new algorithm RMDA ( Regularized Modernized Dual Averaging ) for ( 1 ) . The proposed algorithm provably achieves variance reduction beyond finite-sum problems without any cost or hard-to-tune hyperparameters additional to those of proximal momentum SGD ( proxMSGD ) , and we provide theoretical guarantees for its convergence and ability for manifold identification . The key difference between RMDA and the original regularized dual averaging ( RDA ) of Xiao ( 2010 ) is that RMDA incorporates momentum and can achieve better performance for deep learning in terms of the generalization ability , and the new algorithm requires nontrivial proofs for its guarantees . We further conduct experiments on training deep learning models with a regularizer for structured-sparsity to demonstrate the ability of RMDA to identify the stationary structure without sacrificing the prediction accuracy . When the desired structure is ( unstructured ) sparsity , a popular approach is pruning that trims a given dense model to a specified level , and works like ( Gale et al. , 2019 ; Blalock et al. , 2020 ; Evci et al. , 2020 ; Verma & Pesquet , 2021 ) have shown promising results . However , as a post-processing approach , pruning is essentially different from structured training considered in this work , because pruning is mainly used when a model is available , while structured training combines training and structure inducing in one procedure to potentially reduce the computational cost and memory footprint when resources are scarce . We will also show in our experiment that RMDA can achieve better performance than a state-of-the-art pruning method , suggesting that structured training indeed has its merits for obtaining sparse NNs . The main contributions of this work are summarized as follows . • Principled analysis : We use the theory of manifold identification from nonlinear opti- mization to provide a unified way towards better understanding of algorithms for training structured neural networks . • Variance reduction beyond finite-sum with low cost : RMDA achieves variance reduction for problems that consist of an infinite-sum term plus a regularizer ( see Lemma 2 ) while incorporating momentum to improve the generalization performance . Its spatial and computational cost is almost the same as proxMSGD , and there is no additional hyperparameters to tune , making RMDA suitable for large-scale deep learning . • Structure identification : With the help of variance reduction , our theory shows that under suitable conditions , after a finite number of iterations , iterates of RMDA stay in the active manifold of its limit point . • Superior empirical performance : Experiments on neural networks with structured sparsity exemplify that RMDA can identify a stationary structure without reducing the validation accuracy , thus outperforming existing methods by achieving higher group sparsity . Another experiment on unstructured sparsity also shows RMDA outperforms a state-of-the-art pruning method . After this work is finished , we found a very recent paper Kungurtsev & Shikhman ( 2021 ) that proposed the same algorithm ( with slightly differences in the parameters setting in Line 5 of Algorithm 1 ) and analyzed the expected convergence of ( 1 ) under a specific scheduling of ct = st+1α−1t+1 when both terms are convex . In conrast , our work focuses on nonconvex deep learning problems , and especially on the manifold identification aspect . 2 Algorithm . Details of the proposed RMDA are in Algorithm 1 . At the t-th iteration with the iterate W t−1 , we draw an independent sample ξt ∼ D to compute the stochastic gradient ∇fξt ( W t−1 ) , decide a learning rate ηt , and update the weighted sum Vt of previous stochastic gradients using ηt and the scaling factor βt : = √ t : V0 : = 0 , Vt : = ∑t k=1 ηkβk∇fξk ( W k−1 ) = Vt−1 + ηtβt∇fξt ( W t−1 ) , ∀t > 0 . Algorithm 1 : RMDA ( W 0 , T , η ( · ) , c ( · ) ) input : Initial point W 0 , learning rate schedule η ( · ) , momentum schedule c ( · ) , number of epochs T 1 V0 ← 0 , α0 ← 0 2 for t = 1 , . . . , T do 3 βt ← √ t , st ← η ( t ) βt , αt ← αt−1 + st 4 Sample ξt ∼ D and compute V t ← V t−1 + st∇fξt ( W t−1 ) 5 W̃ t ← arg minW 〈V t , W 〉+ βt2 ∥∥W −W 0∥∥2 + αtψ ( W ) // ( 2 ) 6 W t ← ( 1− c ( t ) ) W t−1 + c ( t ) W̃ t output : The final model WT The tentative iterate W̃ t is then obtained by the proximal operation associated with ψ : W̃ t = proxαtβ−1t ψ ( W 0 − β−1t V t ) , αt : = ∑t k=1 βkηk , ( 2 ) where for any function g , proxg ( x ) : = arg miny ‖y − x‖ 2 /2 + g ( y ) is its proximal operator . The iterate is then updated along the direction W̃ t −W t−1 with a factor of ct ∈ [ 0 , 1 ] : W t = ( 1− ct ) W t−1 + ctW̃ t = W t−1 + ct ( W̃ t −W t−1 ) . ( 3 ) When ψ ≡ 0 , RMDA reduces to the modernized dual averaging algorithm of Jelassi & Defazio ( 2020 ) , in which case it has been shown that mixing W t−1 and W̃ t in ( 3 ) equals to introducing momentum ( Jelassi & Defazio , 2020 ; Tao et al. , 2018 ) . We found that this introduction of momentum greatly improves the performance of RMDA and is therefore essential for applying it on deep learning problems .

Summary: Regularization is generally used to impose desired structure on NN during training. The authors developed RMDA (Regularized Modernized Dual Averaging) which uses the weighted average of the previous SG to compute the tentative update via proximal operation associated with the regularizer. Then the parameters of the model are updated in the direction of this tentative update with a pre-defined factor. The authors theoretically analyzed the performance of their algorithm and showed that under some assumptions, RMDA can identify the structure of model in finite iterations.