Adaptive Proximal Gradient Methods for Structured Neural Networks

1 Introduction . We study the regularized training of neural networks , which can be formulated as the following ( stochastic ) optimization problem minimize θ∈Rd F ( θ ) : = f ( θ ) ︷ ︸︸ ︷ Eξ∼P [ f ( θ ; ξ ) ] +R ( θ ) ( 1 ) where θ ∈ Rd represents the network parameter , ξ is the random variable representing mini-batch data samples , andR ( · ) is a regularizer encouraging low-dimensional structural constraints on θ . The technique of regularization is ubiquitous in machine learning as it can effectively prevent overfitting and yield better generalization . The ` 1-regularized training for Lasso estimators/sparse Gaussian graphical model ( GMRF ) estimation [ 2 , 3 ] and ` 2 weight decay [ 4 ] on parameters are prototypical examples . In the context of deep learning , important instances include network pruning [ 5 , 6 ] , which induces a sparse network structure , and network quantization [ 7 , 8 , 1 ] , which gives hard constraints so that parameters have only discrete values . For the unregularized case , i.e. , when R ( θ ) = 0 , stochastic gradient descent ( SGD ) has been a prevalent approach to solve the optimization problem stated in ( 1 ) . At each iteration , SGD evaluates the gradient on a randomly chosen subset of training samples ( mini-batch ) . While vanilla SGD employs a uniform learning rate for all coordinates , several adaptive variants have been proposed to 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . dynamically take advantage of the data geometry by scaling the learning rate for each coordinate by its gradient history . Prime examples of such approaches include ADAGRAD [ 9 ] , which adjusts the learning rate by the sum of all the past squared gradients , and exponential moving average ( EMA ) approaches such as RMSPROP [ 18 ] and ADAM [ 19 ] , which scale down the gradients by square roots of exponential moving averages of squared past gradients to essentially limit the scope of the adaptation to only a few recent gradients . In terms of theory , convergence analyses of these unregularized SGD methods , whether adaptive or not , have been well studied both for convex [ 19 , 20 ] and non-convex [ 21 , 22 ] loss f cases . For the regularized case , since the regularizer is often non-smooth around some region ( e.g . the ` 1 norm ) , modern machine learning libraries such as TensorFlow [ 23 ] and PyTorch [ 24 ] therefore resort to using the subgradient of the objective function F ( θ ) in ( 1 ) . Such a strategy is problematic as it may slow down convergence and result in oscillations . A simple idea to bypass the non-smoothness of a regularizer is via its proximal operator . This idea is the basis of proximal gradient descent ( PGD ) methods , which first update the parameter using the gradient of the loss function f ( θ ) and then perform a proximal mapping of R ( θ ) . In the non-stochastic case , PGD with both convex and non-convex regularizers has been extensively studied in the literature [ 25 , 26 , 11 , 12 , 27 ] . Another work , VMFB [ 28 ] , analyzes the preconditioned gradient descent on convex regularized problems with non-convex loss but does not consider the first-order momentum . In contrast , PGD in the stochastic setting has been little explored . [ 9 , 10 ] consider PGD to solve the stochastic objectives with convex regularizers . Recently , [ 15 ] studies non-convex and non-smooth regularized problems for DC ( difference of convex ) functions and [ 14 , 17 ] present non-asymptotic analysis for non-convex smooth loss and non-convex regularizers , which is the most general setting , but do not consider the preconditioner in the update rule . All the aforementioned studies of the stochastic case , however , focus either on limited settings ( e.g . [ 9 ] only covers the update rule of ADAGRAD ) with convex regularizers only , or on pure vanilla gradient descent for non-convex regularizers . Hence , they can not accommodate all advanced modern optimization algorithms with preconditioners , such as adaptive gradient methods . The only exception is PROX-SGD [ 16 ] , with the caveat that PROX-SGD update rule is not a pure PGD . Moreover , the theory in [ 16 ] only guarantees convergence , not how fast Prox-SGD converges , and the analysis is performed without considering preconditioners . In this paper , we propose an exact framework for stochastic proximal gradient methods with arbitrary positive preconditioners and lower semi-continuous ( possibly non-convex ) regularizers . With our framework , our goal is to provide theoretical and empirical understanding of stochastic proximal gradient methods for training structured neural networks . Our main contributions can be summarized as follows : • We propose the first general family of stochastic proximal gradient methods , which we term PROXGEN . We introduce two important instances stemming from our approach : ( i ) the first proximal version of ADAM [ 19 ] and ( ii ) a revised version of PROXQUANT [ 1 ] that improves upon the original approach for quantization-specific regularizers by incorporating the effect of preconditioners when computing proximal mappings . Algorithm 1 PROXGEN : A General Stochastic Proximal Gradient Method 1 : Input : Stepsize αt , { ρt } t=Tt=1 ∈ [ 0 , 1 ) , regularization parameter λ , and small constant 0 < δ < < 1 . 2 : Initialize : θ1 ∈ Rd , m0 = 0 ∈ Rd , and C0 = O ∈ Rd×d . 3 : for t = 1 , 2 , . . . , T do 4 : Draw a minibatch sample ξt from P 5 : gt ←∇f ( θt ; ξt ) . Stochastic gradient 6 : mt← ρtmt−1 + ( 1− ρt ) gt . 1st-order momentum 7 : Ct← Preconditioner construction 8 : θt+1 ∈ argmin θ∈Ω { 〈mt , θ〉+ λR ( θ ) + 1 2αt ( θ − θt ) T ( Ct + δI ) ( θ − θt ) } . Update rule 9 : end for • We analyze the convergence of the general PROXGEN family and identify essential conditions for convergence . We show that in general PROXGEN enjoys the same convergence rate as vanilla SGD , but more importantly that the adaptive methods can have faster convergence in terms of constant than vanilla SGD for sparse data . Our convergence guarantee encompasses several existing approaches as special cases . • In terms of practice , we demonstrate the superiority of proximal methods over subgradient-based methods with various non-convex regularizers which have not yet been studied in deep learning . Interestingly , our experiments indicate that the benefit of proximal methods over subgradient approaches is more pronounced with non-convex regularizers than with convex regularizers for learning sparse deep models . Table 1 summarizes the previous studies and our work in terms of stochastic PGD . 2 A Unified Framework of Adaptive Proximal Gradient Methods . In this section , we present PROXGEN , a general family of stochastic proximal gradient methods , and present both existing and novel instances as showcase examples in our family . Algorithm 1 describes the details of PROXGEN . The update rule on line 8 of Algorithm 1 can be written more compactly : θt+1 ∈ argmin θ∈Ω { 〈mt , θ〉+ λR ( θ ) + 1 2αt ( θ − θt ) T ( Ct + δI ) ( θ − θt ) } = proxCt+δIαtλR ( · ) ( θt − αt ( Ct + δI ) −1mt ) ( 2 ) where the proximal operator in ( 2 ) is defined as proxAh ( z ) = argminx { h ( x ) + 12‖x − z‖ 2 A } . In PROXGEN , we allow both the loss and the regularizer to be non-convex . Now , we introduce possible examples according to the proper combinations of preconditioners Ct and regularizersR ( · ) . Existing Instances of PROXGEN . We briefly recover some known examples in PROXGEN family . • ADAGRAD [ 9 ] is the first key instance of adaptive gradient methods where Ct = ( ∑t τ=1 gτg T τ ) 1/2 andR ( θ ) = ‖θ‖1 . Any convex regularizerR ( · ) is allowed . • The proximal Newton methods [ 29 ] employ the exact Hessian Ct = ∇2f ( θt ) andR ( θ ) = ‖θ‖1 . In addition , we can approximate the exact Hessian , which yield proximal Newton-type methods such as quasi-Newton approximation [ 30 ] , L-BFGS approximation [ 31 ] , and adding a multiple of the identity to the Hessian . Although the above examples enjoy good theoretical properties in convex settings , many of the modern practical optimization problems involve non-convex loss functions such as learning deep models . Moreover , it is known that non-convex regularizers yield better performance ( also in terms of theory ) than convex penalties in some applications ( see [ 32 , 33 , 34 , 35 ] and references therein ) . Considering this motivation and recent advanced optimizers , we arrive at the following new examples . Novel Instances of PROXGEN . Beyond the well-known methods above , PROXGEN naturally introduces proximal versions of standard SGD techniques developed for solving unregularized problems for deep learning . The following examples are just a few instances that have not been explored so far , and PROXGEN can cover a broader range of new examples depending on the combinations of preconditioners and regularizers . • The proximal version of ADAM [ 19 ] with ` q regularization is a possible example where Ct = √ βCt−1 + ( 1− β ) g2t with β ∈ [ 0 , 1 ) and R ( θ ) = ‖θ‖q for 0 ≤ q ≤ 1 . We validate empirically the superiority of our novel proximal version of ADAM over the usual subgradientbased counterpart in Section 4 . • We can also consider the proximal version of KFAC [ 36 ] . For an L-layer neural network , KFAC approximates the Fisher information matrix with layer-wise block diagonal structure where lth diagonal block Ct , [ l ] corresponds to Kronecker-factored approximation with respect to the parameters at l-th layer . The proximal version of KFAC corresponds to Ct , [ l ] = E [ δlδTl ] ⊗ E [ al−1aTl−1 ] and R ( θ ) = ‖θ‖q where δl is the gradient with respect to the outputs of l-th layer and al−1 is the activation of ( l − 1 ) -th layer . Examples of Proximal Mappings for PROXGEN . We provide update rules for PROXGEN with ` q regularization ( 0 ≤ q ≤ 1 ) and diagonal preconditioners , for which closed-form updates are available . Diagonal preconditioners are used by popular adaptive gradient methods such as ADAM . Note , however , that our framework and convergence analysis are not limited to diagonal preconditioners and apply to general positive preconditioners . Specifically , we consider regularizerR ( θ ) = λ ∑p j=1 |θj |q for θ ∈ Rp with diagonal preconditioner matrix Ct . Note that for Ct = I ( i.e . vanilla gradient descent ) , it is known that closed-form solutions exist for proximal mappings for q ∈ { 0 , 12 , 2 3 , 1 } [ 37 ] . We denote the i-th coordinate of the vector θt as θt , i and the diagonal entry [ Ct ] ii as Ct , i • ` 1 regularization . The proximal mappings for the case of ` 1 regularization with preconditioner can be computed efficiently via soft-thresholding as θ̂t , i = θt , i − αt mt , i Ct , i + δ , θt+1 , i = sign ( θ̂t , i ) ( ∣∣θ̂t , i∣∣− αtλ Ct , i + δ ) ( 3 ) • ` 0 regularization . In case of ` 0 regularization , we can compute the closed-form solutions via hard-thresholding as θ̂t , i = θt , i − αt mt , i Ct , i + δ , θt+1 , i = θ̂t , i , |θ̂t , i| > √ 2αtλ Ct , i+δ , 0 , |θ̂t , i| < √ 2αtλ Ct , i+δ { 0 , θ̂t , i } , |θ̂t , i| = √ 2αtλ Ct , i+δ ( 4 ) The closed-form proximal mappings for ` 1/2 and ` 2/3 regularization are provided in the Appendix . Table 2 : PROXQUANT versus revised PROXQUANT PROXQUANT proxαtλR ( · ) ( θt − αt ( Ct + δI ) −1mt ) Revised PROXQUANT prox Ct+δI αtλR ( · ) ( θt − αt ( Ct + δI ) −1mt ) Revised PROXQUANT [ 1 ] . The recently proposed PROXQUANT proposes novel regularizations for network quantization . Especially for binary quantization , a W-shaped regularizer is defined as Rbin ( θ ) = ‖θ − sign ( θ ) ‖1 where sign ( θ ) is applied on θ in an elementwise manner . Using this regularizer , the main difference between PROXQUANT and our PROXGEN approach is shown in Table 2 . Note that PROXQUANT ( top in Table 2 ) does not consider the effect of preconditioners when computing proximal mappings . Therefore , we revise the proximal update in PROXQUANT by considering preconditioners in proximal mappings with PROXGEN ( bottom in Table 2 ) . Moreover , we also propose generalized regularizers motivated by ` q regularization for 0 < q < 1 : Rqbin ( θ ) = ‖θ − sign ( θ ) ‖q . In terms of theory , [ 1 ] prove the convergence of PROXQUANT only for the full-batch gradient with differentiable regularizers , which is also guaranteed only for vanilla gradient descent . In contrast , using our revised PROXQUANT , we can completely bridge the gap in theory ( via Theorem 1 in Section 3 , which is stated for stochastic optimization ) , and we provide the exact update rule for solving the problem in ( 1 ) . We also investigate the empirical differences of PROXQUANT and our revised PROXQUANT in Section 4 .

Summary: This work introduces a proximal algorithm for training neural networks. Specifically, the method allows the use of non-convex regularization functions and general PSD pre-conditioners  in the objective, while guaranteeing convergence to a stationary point. Empirical results are provided for training sparse neural networks on CIFAR and ImageNet.