Adaptive Proximal Gradient Methods for Structured Neural Networks
Year: 2022

Abstract:
We consider the training of structured neural networks where the regularizer can be non-smooth and possibly non-convex. While popular machine learning libraries have resorted to stochastic (adaptive) subgradient approaches, the use of proximal gradient methods in the stochastic setting has been little explored and warrants further study, in particular regarding the incorporation of adaptivity. Towards this goal, we present a general framework of stochastic proximal gradient descent methods that allows for arbitrary positive preconditioners and lower semi-continuous regularizers. We derive two important instances of our framework: (i) the first proximal version of ADAM, one of the most popular adaptive SGD algorithm, and (ii) a revised version of PROXQUANT [1] for quantization-specific regularizers, which improves upon the original approach by incorporating the effect of preconditioners in the proximal mapping computations. We provide convergence guarantees for our framework and show that adaptive gradient methods can have faster convergence in terms of constant than vanilla SGD for sparse data. Lastly, we demonstrate the superiority of stochastic proximal methods compared to subgradient-based approaches via extensive experiments. Interestingly, our results indicate that the benefit of proximal approaches over sub-gradient counterparts is more pronounced for non-convex regularizers than for convex ones.

heading: 1 Introduction
text: 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 cannot 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.

heading: 2 A Unified Framework of Adaptive Proximal Gradient Methods
text: 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.

heading: 3 Convergence Analysis
text: In this section, we provide convergence guarantees for the PROXGEN family. Our goal is to find an -stationary point for the problem in (1) where is the required precision. For notational convenience, we assume that the regularization parameter λ is incorporated into R(θ) in (1). To guarantee the convergence under this setting, we should deal with the subdifferentials defined as: Definition 1 (Fréchet Subdifferential). Let ϕ be a real-valued function. The Fréchet subdifferential of ϕ at θ̄ with |ϕ(θ̄)| <∞ is defined by
∂̂ϕ(θ̄) := {θ∗ ∈ Ω ∣∣∣ lim inf
θ→θ̄ ϕ(θ)− ϕ(θ̄)− 〈θ∗, θ − θ̄〉 ‖θ − θ̄‖ ≥ 0}.
Definition 2 (Limiting Subdifferential). Let ∂̂ϕ(θ̄) be the Fréchet subdifferential in Definition 1. The limiting subdifferential of ϕ at θ̄ is defined by
∂ϕ(θ̄) := {u ∈ Rd : ∃θk ϕ→ θ̄, uk ∈ ∂̂ϕ(θk), uk → u}.
where θk ϕ→ θ̄ means θk → θ̄ with ϕ(θk)→ ϕ(θ̄).
To derive the convergence bound, we make the following mild conditions:
(C-1) (L-smoothness) The loss function f is differentiable, L-smooth, and lower-bounded: ∀x, y, ‖∇f(x)−∇f(y)‖ ≤ L‖x− y‖ and f(x∗) > −∞ for the optimal solution x∗.
(C-2) (Bounded variance) The stochastic gradient gt = ∇f(θt; ξ) is unbiased and has the bounded variance: Eξ [ ∇f(θt; ξ) ] = ∇f(θt), Eξ [ ‖gt −∇f(θt)‖2 ] ≤ σ2. (C-3) (i) final step-vector is finite, (ii) the stochastic gradient is bounded, and (iii) the momentum parameter should be exponentially decaying: (i) ‖θt+1 − θt‖ ≤ D, (ii) ‖gt‖ ≤ G, (iii) ρt = ρ0µ
t−1 with D,G > 0 and ρ0, µ ∈ [0, 1). (C-4) (Sufficiently positive-definite) The minimum eigenvalue of effective spectrums should be
uniformly lower bounded over all time t: ∀t, λmin ( αt(Ct + δI) −1) ≥ γ > 0. (C-1) and (C-2) are very standard in convergence analysis for optimization algorithms designed for deep learning such as ADAM, YOGI, and many others [21, 38, 39, 40, 41]. In addition, (C3) is extensively studied in previous literature for analysis of general non-convex optimization [19, 20, 42, 40, 41]. Lastly, a similar condition to (C-4) is also considered in [39, 42]. We note that (C-3) and (C-4) are reasonable conditions: It is well-known that the parameter of an overparametrized neural network hardly changes from the initial point during training [43, 44, 45], so one can expect that the diameter D of parameter space and the bound for the size of gradient G have very small values and can be understood as constants in rates of the results. To validate this for real cases, we train ResNet-34 on CIFAR-10 dataset. In Figure-1-(a), the difference of parameters ‖θt+1 − θt‖2 and the size of stochastic gradients ‖gt‖2 attain just 1 ∼ 3 while the parameter dimension d of ResNet-34 is about 107. Hence, the constants D and G in (C-3) are negligible compared to the problem dimension d in practice. The exponentially decaying momentum parameter assumption ρt = ρ0µ
t−1 could be relaxed to ρt = ρ0/t sacrificing the logarithmic factor in our analysis. Also, (C-4) is indeed easily satisfied both theoretically and empirically. This condition holds in theory for most of the popular optimization algorithms for deep learning such as ADAGRAD, ADAM, and KFAC (the constant γ is irrelevant to the problem dimension d for each algorithm, and we defer the derivations to Appendix D). In order to investigate whether these conditions could be satisfied in real problems, we revisit the experiments of training ResNet-34. In Figure 1-(b), we can see the minimum eigenvalue of αt(Ct + δI)−1 tends to increase, so the condition (C-4) is also satisfied empirically.
Since the loss function f is assumed to be differentiable as in (C-1) and it is known that ∂̂ϕ(θ) ⊆ ∂ϕ(θ), we have, at stationary points, 0 ∈ ∂̂F (θ) = ∇f(θ) + ∂̂R(θ), so the convergence criterion is slightly different from that of general non-convex optimization. Hence, we use the following convergence criterion E[dist(0, ∂̂F (θ))] ≤ for an -stationary point where dist(x,A) denotes the distance between a vector x and a set A. If no regularizer is considered (R = 0), this criterion boils down to the one usually used in non-convex optimization, E[‖∇f(θ)‖] ≤ . • Challenges specific to the analysis of PROXGEN. The most challenging issue in the analysis of PROXGEN compared to previous studies [14, 21] is that we should handle the momentum mt and
Diameter of Parameter Space
Gradient Norm
Condition (C-4)
non-trivial preconditioner Ct. In terms of adaptive gradient methods, [21] guarantees the convergence of a family of adaptive methods (but without proximal mapping) using the changes of effective learning rate (Γt := αt/ √ Vt − αt+1/ √ Vt+1 ≥ 0 where Vt is an adaptation matrix), which is a key quantity in their theory. [21] define a new sequence {zt} involving the quantity Γt and exploit the simple closed-form of the quantity zt+1 − zt to derive the convergence with coordinate-wise analysis. However, this proof technique is not available to the regularized problems since zt+1 − zt is not amenable anymore to compute in a simple closed-form due to the proximal mapping. On the other hand, our proof directly solves the quadratic subproblem w.r.t. θt at line 8 in Algorithm 1 to handle a regularizer term. It should also be emphasized that our proof skill can handle arbitrary positive curvatures (hence including more general non-diagonal one) that were not acceptable in [21]. In the context of proximal gradient descent, our proof is totally different from [14] which is only for vanilla SGD. Due to the existence of mt, it is highly non-trivial to bound the term ‖mt −∇f(θt)‖2 without suitable assumptions whereas ‖gt −∇f(θt)‖2 in [14] can be easily bounded using (C-2). Also, we need to deal with quadratic approximation term (θ − θt)T(Ct + δI)(θ − θt) in Algorithm 1 which is not problematic in [14] simply because Ct is trivially I . We could successfully bypass those difficulties using mild conditions (C-3) and (C-4), respectively.
We are ready to state our theorem for general convergence. Theorem 1. Let θa denote an iterate uniformly randomly chosen from {θ1, · · · , θT }. Under the conditions (C-1), (C-2), (C-3), (C-4) with the initial stepsize α0 ≤ δ3L and non-increasing stepsize αt, PROXGEN, Algorithm 1, is guaranteed to yield
Ea[dist(0, ∂̂F (θa))2] ≤ Q1σ
2
T T−1∑ t=0 1 bt + Q2∆ T + Q3 T
where ∆ = f(θ)− f(θ∗) with optimal point θ∗, and bt is the minibatch size at time t. The constants {Qi}3i=1 on the right-hand side depend on the constants {α0, δ, L,D,G, ρ0, µ, γ}, but not on T .
Note that the constants {Qi}3i=1 in Theorem 1 are completely independent of the problem dimension d. From Theorem 1, the appropriate minibatch size is important to ensure a good convergence. Various settings for the minibatch size could be employed for convergence guarantee, but considering practical cases, we provide the following important corollary for constant minibatch. Corollary 1 (Constant Mini-batch). Under the same assumptions as in Theorem 1 with sample size n and constant minibatch size bt = b = Θ(T ), we have Ea [ dist(0, ∂̂F (θa)) 2 ] ≤ O ( 1/T ) and the
total complexity is O(1/ 4) in order to have Ea [ dist ( 0, ∂̂F (θa) )] ≤ .
Here we make several remarks on our results and relationship with prior work.
• On Convergence Results. Note that our Corollary 1 achieves the optimal complexity O(1/ 4) of SGD to find -stationary points under the standard assumptions (C-1) ∼ (C-4). Recent studies [46, 47] show faster rate, but under additional stronger assumptions such as second-order smoothness (i.e., the smoothness of Hessian matrix). Also, we could relax the exponentially decaying momentum ρt = ρ0µ
t−1 in (C-3) to ρt = ρ0/t as mentioned in [48] with the logarithm factor as Q3 = O(log T ), which in result still ensures Õ(1/ 4).
• Advantages of using adaptive gradient methods in Theorem 1. We discuss how the constant γ in (C-4) affects the convergence in terms of theory. According to our proofs, γ depends on the algorithmic details and the constants Q1, Q2 and Q3 in Theorem 1 are proportional to 1/γ. The convergence rate depends on these constants and the benefit of preconditioners can be found here. To view this more clearly, we consider the diagonal matrix adaptation of ADAM [19], i.e. constant
stepsize αt = α and Ct = √ (1− β) ∑t τ=1 β
t−τgτ gτ , with β ∈ [0, 1) and the total iteration T . In this setting, the 1/γ can be computed as
Qi ∝ 1
γ =
√ (1− β) ∑t τ=1 β
t−τ‖gτ‖22 + δ α ≤ G+ δ α
where gτ is the gradient at time τ and δ is a small constant while the vanilla SGD (Ct = 0 and δ = 1) satisfies 1/γ = 1/α. Here, we can clearly see the advantages of adaptive methods (i.e., using preconditioners) since 1/γ could be dramatically smaller if ‖gτ‖2 << 1 holds roughly with small constant δ, which corresponds to sparse gradients ‖gτ‖2 (the data features are sparse). This coincides with the convex regret theory for adaptive gradient methods [9, 19, 48], which also holds in our theory with non-convex smooth loss and non-convex regularizers.
• Implications of condition (C-4) on theory. Our analysis relies on (C-4), the lower bound for the minimum eigenvalue of Γt := αt(Ct + δI)−1. This means that Theorem 1 guarantees Ea[dist(0, ∂̂F (θa)2] ≤ O(1/ √ T ) (in case of b = Θ(T ) as in Corollary 1) for any change of basis of Γt, so in that sense, we provide a worst-case analysis and there is room for more optimistic bounds.
• On minibatch in Corollary 1. The conditions b = Θ(T ) is considered as standard in many previous literature [38, 14] and is not stringent. In terms of stochastic optimization, it is natural in practice to choose the batch size b and the number of epochs e in advance. Then, the total number of iterations T satisfies the following relation: T = e× nb = e× n Θ(T ) . In this sense, the total iterations T should be an order of O( √ n) in practice. For example, this condition sets a minibatch size of approximately 200 and 1000 for CIFAR-10 and ImageNet dataset respectively, which is practical.
• Connections to second-order methods. Our analysis can provide guarantees for positive secondorder preconditioners as long as (C-4) is satisfied (The empirical Fisher information [36] is one example). Although second-order solvers generally enjoy very fast convergence under strongly convex loss [29, 49], it can be understood that our theory guarantees at least a sublinear rate for such second-order curvatures with less stringent conditions.

heading: 4 Experiments
text: We consider two important tasks for regularized training in deep learning communities: (i) training sparse neural networks and (ii) network quantization. Throughout our experiments, we consider ADAM as a representative of PROXGEN where mt = ρtmt−1 + (1− ρt)gt with constant decaying parameter ρt = 0.9 and Ct = √ βCt−1 + (1− β)g2t with β = 0.999 in Algorithm 1. The details on other hyperparameter/experiment settings are provided in the Appendix.
Training Sparse Neural Networks. Motivated by the lottery ticket hypothesis [50], we consider training VGG-16 [51] and ResNet-34 [52] on CIFAR-10 dataset using sparsity encouraging regularizers. Toward this, we consider the following objective function with possibly non-convex `q regularization: F (θ) := Eξ∼P[f(θ; ξ)] + λ ∑p j=1 |θj |q where 0 ≤ q ≤ 1. We train the network parameters with the closed-form proximal mappings introduced in Section 2. The results on VGG-16 are provided in Appendix.
We compare PROXGEN with subgradient methods and also include PROX-SGD [16] as a baseline especially for `1 regularization since PROX-SGD considers only convex regularizers. In PROX-SGD, the hand-crafted fine-tuned scheduling on αt and ρt is essential for fast convergence and good performance, but in our experiments we use standard settings ρt = 0.9. We first validate our theory in practice using constant stepsize in order to purely see the effect of proximal approaches (the results on this setting are provided in the Appendix). Then the step-decay learning rate scheduling is employed to consider standard training schemes for the state-of-the-art performance, which also satisfies the non-increasing stepsize condition in our Theorem 1. For `0 regularization, the problem in (1) cannot be optimized in a subgradient manner, so we compare PROXGEN with another popular baseline, `0hc [6] which approximates the `0-norm via hard-concrete distributions.
Figures 2 illustrates the results for ResNet-34. In terms of convergence, PROXGEN shows faster convergence than PROX-SGD for `1 case in Figure 2-(a), but there is no difference between PROXGEN and subgradient methods as in Figure 2-(a). However, there are notable differences in convergence for non-convex regularizers `1/2 and `2/3, which get bigger as q decreases. We believe this might be because the `q-norm derivative, q/|θ|1−q , is very large for non-zero tiny θ for q ∈ (0, 1). Meanwhile, ∂|θ|/∂θ is merely the sign value regardless of size of θ, so the large gradient of |θ|q may hinder convergence. The learning curves in Figure 2-(b,c) empirically corroborate this phenomenon.
In terms of performance, we can see that PROXGEN consistently achieves better performance than baselines for ResNet-34 with similar or even better sparsity level. Importantly, PROXGEN with `0 outperforms `0hc baseline by a great margin. This might be due to the design of `0hc , which approximates ‖θ‖0 = ∑p j=1 I{θj 6= 0} with binary mask zj parameterized by learnable probability πj for each coordinate. Thus, the number of parameters to be optimized is doubled, which might make optimization harder. In constrast, PROXGEN does not introduce additional parameters.
More results for other famous non-convex regularizers MCP [53] and SCAD [54] are in Appendix.
Training Group-Sparse Neural Networks. In the Appendix, we consider training Statistical Recurrent Units where `1,2 group-norm penalty is imposed on the input layer weights to detect non-linear Granger Causality [55]. As the proximal mappings for PROXGEN with group sparsity are not available in closed-form, we develop an efficient procedure for computation, whose derivations are also provided in the Appendix.
Training Binary Neural Networks. We consider the network quantization constraining the parameters to some set of discrete values which is a key approach for model compression. We evaluate our revised PROXQUANT in Table 2 with extended regularization Rqbin in Section 2. We consider the following objective function with quantization-specific regularizers: F (θ) := Eξ∼P[f(θ; ξ)] + λ ∑p j=1 |θj − sign(θj)|q where 0 ≤ q ≤ 1. For comparisons, we quantize ResNet weight parameters (except bias and activations) on CIFAR-10 and ImageNet dataset.
Table 3 presents the results. For all q values, revised PROXQUANT consistently outperforms the baselines except for ResNet-20, which implies PROXGEN may work better for larger networks. As such, our generalized regularizers Rqbin contribute to one of the state-of-the-art optimization-based methods in network quantization. Notably, revised PROXQUANT `1 greatly outperforms PROXQUANT baseline while these two approaches differ only in update rules (see Table 2). Hence, we can conclude that revised PROXQUANT based on PROXGEN provides an exact proximal update and also yields more generalizable solutions. In our experience, revised PROXQUANT `0 shows little degradation in performance, so we do not include this result. However, revised PROXQUANT `0 shows superiority to baselines for language modeling, whose preliminary results are in Appendix.
Table 4 illustrates Top-1/Top-5 error (%) for training ResNet on ImageNet with binary quantization. The most important thing is that our revised PROXQUANT shows great improvements in performance over the original PROXQUANT. Furthermore, PROXGEN shows superior performance to various baselines for weight quantization.

heading: 5 A Closer Look into Prox-SGD [16] vs. PROXGEN
text: Prox-SGD [16] is the approach closest to our PROXGEN method. However, PROX-SGD is not an exact proximal approach and is significantly different from PROXGEN. PROXGEN’s update rule
involves directly solving the quadratic subproblem (2). In contrast, PROX-SGD’s update rule consists of two stages: (i) solving the quadratic subproblem without learning rate (5), then (ii) updating the parameters with the computed direction (i.e. θ̂t − θt) by the learning rate αt (6).
θ̂t = prox Ct+δI λR(·) ( θt − (Ct + δI)−1mt ) ︸ ︷︷ ︸
no learning rate
, (5)
θt+1 = θt + αt(θ̂t − θt) (6)
To clearly see the differences between both approaches, we conduct two studies.
Study 1: Lasso Support Recovery. For this task, the two-stage update scheme of PROX-SGD might have some potential issues. For example, for `1-regularized problems, the updated parameter θt+1 (6) might not achieve exact zero (while θ̂t can) whereas θt+1 for PROXGEN (2) can attain exact zero value according to the update rule (3) in Section 2. Another potential caveat is that PROX-SGD might overestimate the sparsity level. In view of the above, we run Lasso simulations with different two initialization schemes: (i) random initialization and (ii) zero initialization. For random initialization, it can be seen in Figure 3-(a) that PROX-SGD could not achieve exact zero value, which corroborates our first observation. More interestingly, for zero initialization, we can see in Figure 3-(b) that the estimates using PROX-SGD are exactly zeros for all coordinates, which supports our second observation. This might be because θ̂t (5) is always zero since the quadratic subproblem does not consider the learning rate, which might overestimate the sparsity level. Hence, the subsequent iterate θt+1 would be always zero since we initialize the parameters with zeros, but PROXGEN recovers the correct support in both cases.
Study 2: DenseNet-201 on CIFAR-100 Dataset. To validate the superiority of PROXGEN upon PROX-SGD, we revisit the largest experiments in [16]. We train DenseNet-201 architecture on CIFAR-100 dataset with `1 regularization since PROX-SGD only consider convex regularizers. For both methods, we use the same hyperparameter settings for fair comparison. Figure 4 illustrates the training learning curves, and it can be seen that our PROXGEN achieves faster convergence as well as lower objective values. For our experience, the learning curves show the similar dynamics for different λ values.
Comparison of Theoretical Contributions. [16] guarantees the convergence of PROX-SGD, but not how fast it converges. Moreover, this is proved without considering preconditioners. In contrast, our analysis for
the PROXGEN framework appropriately incorporates the first-order momentum and arbitrary positive preconditioner with detailed non-asymptotic convergence.

heading: 6 Conclusion
text: In this work, we proposed PROXGEN, the first general family of stochastic proximal gradient methods. Within our framework, we presented novel examples of proximal versions of standard SGD approaches, including a proximal version of ADAM. We analyzed the convergence of the whole PROXGEN family and showed that PROXGEN can encompass the results of several previous studies. We also demonstrated that PROXGEN empirically outperforms subgradient-based methods for popular deep learning problems. As future work, we plan to further study efficient procedures to compute the proximal mappings for structured regularizers such as `1/`q-norms with preconditioners.

heading: Acknowledgement
text: This work was supported by the National Research Foundation of Korea (NRF) grants (2018R1A5A1059921, 2019R1C1C1009192) and Institute of Information & Communications Technology Planning & Evaluation (IITP) grants (No.2019-0-01371, Development of brain-inspired AI with human-like intelligence, No.2019-0-00075, Artificial Intelligence Graduate School Program (KAIST)) funded by the Korea government (MSIT).

Review summaries:
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.

The paper presents a framework for adaptive proximal stochastic methods. The general algorithm provided, ProxGen, is a variation on traditional proximal gradient descent in which the gradient is multiplied by the inverse of some $C_t+ \delta I$, and the norm in the prox operator is given by $\|x\| = x^{\top} (C_t+ \delta I) x$. They give a convergence result which matches the optimal rate for SGD, and they show experiments with sparse, group-sparse, and binary neural networks in which they show improvements over existing algorithms.