P ROX -SGD: T RAINING S TRUCTURED N EURAL N ETWORKS UNDER R EGULARIZATION AND C ONSTRAINTS
Anonymous authors
Paper under double-blind review
A BSTRACT
In this paper, we consider the problem of training structured neural networks (NN) with nonsmooth regularization (e.g. (cid:96) 1 -norm) and constraints (e.g. interval constraints).
We formulate training as a constrained nonsmooth nonconvex optimization problem, and propose a convergent proximal-type stochastic gradient descent (Prox-SGD) algorithm.
We show that under properly selected learning rates, with probability 1, every limit point of the sequence generated by the proposed ProxSGD algorithm is a stationary point.
Finally, to support the theoretical analysis and demonstrate the ﬂexibility of Prox-SGD, we show by extensive numerical tests how Prox-SGD can be used to train either sparse or binary neural networks through an adequate selection of the regularization function and constraint set.
I NTRODUCTION
In this paper, we consider the problem of training neural networks (NN) under constraints and regularization.
It is formulated as an optimization problem
[Equation]
where x is the parameter vector to optimize, y i is the i -th training example which consists of the training input and desired output, and m is the number of training examples.
The training loss f is assumed to be smooth (but nonconvex) with respect to x , the regularization r is assumed to be convex (but nonsmooth), proper and lower semicontinuous, and the constraint set X is convex and compact (closed and bounded).
When r ( x ) = 0 and X = R n , stochastic gradient descent (SGD) has been used to solve the optimization problem (1).
At each iteration, a minibatch of the m training examples are drawn randomly, and the obtained gradient is an unbiased estimate of the true gradient.
Therefore SGD generally moves along the descent direction, see Bertsekas & Tsitsiklis (2000).
SGD can be accelerated by replacing the instantaneous gradient estimates by a momentum aggregating all gradient in past iterations.
Despite the success and popularity of SGD with momentum, its convergence had been an open problem.
Assuming f is convex, analyzing the convergence was ﬁrst attempted in Kingma &
Ba (2015) and later concluded in Reddi et al.
The proof for a nonconvex f was later given in Chen et al.
; Lei et al.
In machine learning, the regularization function r is typically used to promote a certain structure in the optimal solution, for example sparsity as in, e.g., feature selection and compressed sensing, or a zero-mean-Gaussian prior on the parameters (Bach et al., 2011; Boyd et al., 2010).
It can be interpreted as a penalty function since at the optimal point x (cid:63) of problem (1), the value r ( x (cid:63) ) will be small.
One nominant example is the Tikhonov regularization r ( x ) = µ (cid:107) x (cid:107) 2 2 for some predeﬁned constant µ , and it can be used to alleviate the ill-conditioning and ensure that the magnitude of the weights will not become exceedingly large.
Another commonly used regularization, the (cid:96) 1 -norm where r ( x ) = µ (cid:107) x (cid:107) 1 = µ (cid:80) nj =1 | x j | (the convex surrogate of the (cid:96) 0 -norm), would encourage a sparse solution.
In the context of NN, it is used to (i) promote a sparse neural network (SNN) to alleviate overﬁtting and to allow a better generalization, (ii) accelerate the training process, and (iii) prune the network to reduce its complexity, see Louizos et al.
(2018) and Gale et al.
Technically, it is difﬁcult to analyze the regularizations as some commonly used convex regularizers are nonsmooth, for example, (cid:96) 1 -norm.
In current implementations of Tensorﬂow, the gradient of | x | is simply set to 0 when x = 0 .
This amounts to the stochastic subgradient descent method and usually exhibits slow convergence.
Other techniques to promote a SNN includes magnitude pruning and variational dropout, see Gale et al.
Although regularization can be interpreted as a constraint from the duality theory, sometimes it may still be more desirable to use explicit constraints, for example, (cid:80) x 2 j ≤ α , where the summation is over the weights on the same layer.
This is useful when we already know how to choose α .
Another example is the lower and upper bound on the weights, that is, l ≤ w ≤ u for some predeﬁned l and u .
Compared with regularization, constraints do not encourage the weights to stay in a small neighborhood of the initial weight, see Chapter 7.2 of Goodfellow et al.
(2016) for more details.
The set X models such explicit constraints, but it poses an additional challenge for stochastic gradient algorithms as the new weight obtained from the SGD method (with or without momentum) must be projected back to the set X to maintain its feasibility.
However, projection is a nonlinear operator, so the unbiasedness of the random gradient would be lost.
Therefore the convergence analysis for constrained problems is much more involved than unconstrained problems.
In this paper, we propose a convergent proximal-type stochastic gradient algorithm (Prox-SGD) to train neural networks under nonsmooth regularization and constraints.
It turns out momentum plays a central role in the convergence analysis.
We establish that with probability (w.p.) 1, every limit point of the sequence generated by Prox-SGD is a stationary point of the nonsmooth nonconvex problem (1).
This is in sharp contrast to unconstrained optimization, where the convergence of the vanilla SGD method has long been well understood while the convergence of the SGD method with momentum was only settled recently.
Nevertheless, the convergence rate of Prox-SGD is not derived in the current work and is worth further investigating.
To test the proposed algorithm, we consider two applications.
The ﬁrst application is to train a SNN, and we leverage (cid:96) 1 -regularization, that is,
[Equation]
The second application is to train a binary neural network (BNN) where the weights (and activations) are either 1 or -1 (see Courbariaux et al.
; Hou et al.
; Yin et al.
; Bai et al.
(2019) for more details).
To achieve this, we augment the loss function with a term that penalizes the weights if they are not +1 or -1: minimize
[Equation] [Equation]
subject to a j = 0 or 1 , j = 1 , . .
. , n, where µ is a given penalty parameter.
The binary variable a j can be interpreted as a switch for weight x j : when a j = 0 , (1 − a j )( x j − 1) 2 is activated, and there is a strong incentive for x j to be 1 (the analysis for a j = 1 is similar).
Since integer variables are difﬁcult to optimize, we relax a j to be a continuous variable between 0 and 1.
To summarize, a BNN can be obtained by solving the following regularized optimization problem under constraints with respect to x and a minimize
[Equation]
subject to 0 ≤ a j ≤ 1 , j [Equation] [Equation]
n.
[Equation]
If µ is properly selected (or sufﬁciently large), the optimal a j will be exactly or close to 0 or 1.
Consequently, regularization and constraints offer interpretability and ﬂexibility, which allows us to use more accurate models to promote structures in the neural networks, and the proposed convergent
Prox-SGD algorithm ensures efﬁcient training of such models.
T HE PROPOSED P ROX -SGD ALGORITHM
In this section, we describe the Prox-SGD algorithm to solve (1)
Background and setup.
We make the following blanket assumptions on problem (1).
• f i ( x , y ( i ) ) is smooth (continuously differentiable) but not necessarily convex.
• ∇ x f i ( x , y ( i ) ) is Lipschitz continuous with a ﬁnite constant L i for any y i . Thus ∇ f ( x ) is Lipschitz continuous with constant L (cid:44) 1 m (cid:80) mi =1 L i .
• r
( x ) is convex, proper and lower semicontinuous (not necessarily smooth).
• X is convex and compact.
We are interested in algorithms that can ﬁnd a stationary point of (1).
A stationary point x (cid:63) satisﬁes the optimality condition: at x = x (cid:63) ,
[Equation]
When r ( x ) = 0 and X = R n , the deterministic optimization problem 1 can be solved by the (batch) gradient descent method.
When m , the number of training examples, is large, it is computationally expensive to calculate the gradient.
Instead, we estimate the gradient by a minibatch of m ( t ) training examples.
We denote the minibatch by M ( t ) : its elements are drawn uniformly from { 1 , 2 , . . .
, m } and there are m ( t ) elements.
Then the estimated gradient is
[Equation]
and it is an unbiased estimate of the true gradient.
The proposed algorithm.
The instantaneous gradient g ( t ) is used to form an aggregate gradient (momentum) v ( t ) , which is updated recursively as follows
[Equation]
where ρ ( t ) is the stepsize (learning rate) for the momentum and ρ ( t ) ∈ (0 , 1] .
At iteration t , we propose to solve an approximation subproblem and denote its solution as (cid:98) x
( x ( t ) , v ( t ) , τ ( t )) , or simply (cid:98) x ( t [Equation] x [Equation]
argmin
[Equation]
A quadratic regularization term is incorporated so that the subproblem (7) is strongly convex and its modulus is the minimum element of the vector τ ( t ) , denoted as τ ( t ) and τ ( t ) =
min
j =1 ,...,n τ j ( t ) .
Note that τ ( t ) should be lower bounded by a positive constant that is strictly larger than 0, so that the quadratic regularization in (7) will not vanish.
The difference between two vectors (cid:98) x ( t ) and x ( t ) speciﬁes a direction starting at x ( t ) and ending at (cid:98) x ( t ) .
This update direction is used to reﬁne the weight vector
[Equation]
where (cid:15) ( t ) is a stepsize (learning rate) for the weight and (cid:15) ( t ) ∈ (0 , 1] .
Note that x ( t +1) is feasible as long as x ( t ) is feasible, as it is the convex combination of two feasible points x ( t ) and (cid:98) x ( t ) while the set X is convex.
The above steps (5)-(8) are summarized in Algorithm 1, which is termed proximal-type Stochastic
Gradient Descent (Prox-SGD), for the reason that the explicit constraint x ∈ X in (7) can also be formulated implicitly as a regularization function, more speciﬁcally, the indicator function δ X ( x ) .
If all elements of τ ( t ) are equal, then (cid:98) x ( t ) is exactly the proximal operator [Equation] argmin [Equation]
Input
[Equation]
[Equation]
[Equation]
[Equation]
[Equation]
[Equation] [Equation] [Equation] for t = 0 : 1 : T do
[Equation] [Equation]
end for
[Table]
[Table] [Equation] convex convex, compact
[Table] [Table]
) 1( ρ [Equation] (cid:15)
R n
AdaGrad (cid:15) (cid:112)
[Equation]
t
[Equation]
R n
RMSProp (cid:15) (cid:112) r
( t )
+ [Equation] [Equation]
R n ADAM ρ (cid:15) 1 −
[Equation]
t (cid:113) r ( t ) 1 − β t + δ
[Table]
Prox-SGD in Algorithm 1 bears a similar structure as several SGD algorithms, without and with momentum, see Table 1, and it allows to interpret some existing algorithms as special cases of the proposed framework.
For example, no momentum is used in SGD, and this amounts to setting ρ ( t )
= 1 in Algorithm 1.
In ADAM, the learning rate for momentum is a constant ρ and the learning rate for the weight vector is given by (cid:15)/ (1 − ρ t ) for some (cid:15) , and this simply amounts to setting ρ ( t )
=
ρ and (cid:15) ( t ) =
(cid:15)/ (1 − ρ t ) in Algorithm 1.
This interpretation also implies that the convergence conditions to be proposed shortly later are also suffcient for existing algorithms (although they are not meant to be the weakest conditions available in literature).
Solving the approximation subproblem (7).
Since (7) is strongly convex, (cid:98) x
( t ) is unique.
Generally (cid:98) x ( t ) in (7) does not admit a closed-form expression and should be solved by a generic solver.
However, some important special cases that are frequently used in practice can be solved efﬁciently.
• The trivial case is X = R n and r = 0 , where
[Equation]
where the vector division is understood to be element-wise.
When X = R n and r ( x ) = µ (cid:107) x (cid:107) 1 , (cid:98) x
( t ) has a closed-form expression that is known as the soft-thresholding operator
[Equation]
where S a ( b ) (cid:44) max( b − a , 0 ) − max( − b − a , 0 ) (Bach et al., 2011).
• If X = R n and r ( x ) = µ (cid:107) x (cid:107) 2 and τ ( t )
=
τ I for some τ , then (Parikh & Boyd, 2014)
[Equation]
If x is divided into blocks x 1 , x 2 , . .
, the (cid:96) 2 -regularization is commonly used to promote block sparsity (rather than element sparsity by (cid:96) 1 -regularization).
• When there is a bound constraint l ≤ x ≤ u , (cid:98) x ( t ) can simply be obtained by ﬁrst solving the approximation subproblem (7) without the bound constraint and then projecting the optimal point onto the interval [ l , u ] .
For example, when X = R n and r = 0 ,
[Equation]
x [Equation] with [ x ] ul = clip ( x , l , u ) (cid:44) min(max( x , l ) , u ) .
• If the constraint function is quadratic:
X = { x : (cid:107) x (cid:107) 22 ≤ 1 } , (cid:98) x ( t ) has a semi-analytical expression (up to a scalar Lagrange multiplier which can be found efﬁciently by the bisection method).
Approximation subproblem.
We explain why we update the weights by solving an approximation subproblem (7).
First, we denote (cid:101) f as the smooth part of the objective function in (7).
Clearly it depends on x ( t ) and v ( t ) (and thus M ( t ) ), while x ( t ) and v ( t ) depend on the old weights x (0) , . . .
, x ( t − 1) and momentum and v (0) , . .
, v ( t − 1) .
Deﬁne F ( t ) (cid:44) { x (0) , . . .
, x ( t ) , M (0) , . . .
, M ( t ) } as a shorthand notation for the trajectory generated by ProxSGD.
We formally write (cid:101) f as
[Equation]
It follows from the optimality of (cid:98) x
( t ) that
[Equation]
After inserting (13) and reorganizing the terms, the above inequality becomes
[Equation]
Since ∇ f ( x ) is Lipschitz continuous with constant L , we have
[Equation]
where the ﬁrst inequality follows from the descent lemma (applied to f ) and the second inequality follows from the Jensen’s inequality of the convex function r and the update rule (8).
If v ( t ) = ∇ f
( x ( t ))
(which is true asymptotically as we show shortly later), by replacing ∇ f ( x ( t )) in (16) by v ( t ) and inserting (14) into (16), we obtain
[Equation]
The right hand side (RHS) will be negative when (cid:15) ( t ) < 2 τ ( t )
L : this will eventually be satisﬁed as we shall use a decaying (cid:15) ( t ) .
This implies that the proposed update (8) will decrease the objective value of (1) after each iteration.
Momentum and algorithm convergence.
It turns out that the momentum (gradient averaging step) in (6) is essential for the convergence of Prox-SGD.
Under some mild technical assumptions we outline now, the aggregate gradient v ( t ) will converge to the true (unknown) gradient ∇ f ( x ( t )) .
This remark is made rigorous in the following theorem.
Theorem 1.
Assume that the unbiased gradient g ( t ) has a bounded second moment
[Equation]
(18) for some ﬁnite and positive constant C , and the sequence of stepsizes { ρ ( t ) } and { (cid:15) ( t ) } satisfy ∞ (cid:88) t =
[Equation]
t
[Equation]
Then lim t →∞ (cid:107) v ( t ) − ∇ f ( x ( t )) (cid:107) = 0 , and every limit point of the sequence { x ( t ) } is a stationary point of (1) w.p.1.
Proof.
Under the assumptions (18) and (19), it follows from Lemma 1 of Ruszczy´nski (1980) that v ( t ) → ∇ f ( x ( t )) .
Since the descent direction (cid:98) x ( t ) − x ( t ) is a descent direction in view of (14), the convergence of the Prox-SGD algorithm can be obtained by generalizing the line of analysis in Theorem 1 of Yang et al. (2016) for smooth optimization problems.
The detailed proof is included in the appendix to make the paper self-contained.
We draw some comments on the convergence analysis in Theorem 1. The bounded second moment assumption on the gradient g in (18) and decreasing stepsizes in (19) are standard assumptions in stochastic optimization and SGD.
What is noteworthy is that (cid:15) ( t ) should decrease faster than ρ ( t ) to ensure that v ( t ) → ∇ f ( x ( t )) .
But this is more of an interest from the theoretical perspective, and in practice, we observe that (cid:15) ( t ) /ρ
( t ) =
a for some constant a that is smaller than 1 usually yields satisfactory performance, as we show numerically in the next section.
According to Theorem 1, the momentum v ( t ) converges to the (unknown) true gradient ∇ f ( x ( t )) , so the Prox-SGD algorithm eventually behaves similar to the (deterministic) gradient descent algorithm.
This property is essential to guarantee the convergence of the Prox-SGD algorithm.
To guarantee the theoretical convergence, the quadratic gain τ ( t ) in the approximation subproblem (7) should be lower bounded by some positive constant (and it does not even have to be timevarying).
In practice, there are various rationales to deﬁne it (see Table 1), and they lead to different empirical convergence speed and generalization performance.
The technical assumptions in Theorem 1 may not always be fully satisﬁed by the neural networks deployed in practice, due to, e.g., the nonsmooth ReLU activation function, batch normalization and dropout.
Nevertheless, Theorem 1 still provides valuable guidance on the algorithm’s practical performance and the choice of the hyperparameters.
S IMULATION R ESULTS
In this section, we perform numerical experiments to test the proposed Prox-SGD algorithm. In particular, we ﬁrst train two SNN to compare Prox-SGD with ADAM (Kingma & Ba, 2015), AMSGrad (Reddi et al., 2018), ADABound (Luo et al., 2019) and SGD with momentum.
Then we train a BNN to illustrate the merit of regularization and constraints.
To ensure a fair comparison, the hyperparameters of all algorithms are chosen according to either the inventors’ recommendations or a hyperparameter search.
Furthermore, in all simulations, the quadratic gain τ ( t ) in Prox-SGD is updated in the same way as ADAM, with β = 0 .
999 (see Table 1).
3. S PARSE N EURAL N ETWORK : T RAINING C ONVOLUTION N EURAL N ETWORKS ON CIFAR-
We ﬁrst consider the multiclass classiﬁcation problem on CIFAR-10 dataset (Krizhevsky, 2009) with convolution neural network (CNN).
The network has 6 convolutional layers and each of them is followed by a batch normalization layer; the exact setting is shown in Table 2.
Following the parameter conﬁgurations of ADAM in Kingma & Ba (2015), AMSGrad in Reddi et al.
(2018), and ADABound in Luo et al.
(2019), we set ρ = 0 .
1 , β = 0 .
999 and (cid:15) = 0 .
001 (see Table 1), which are uniform for all the algorithms and commonly used in practice.
Note that we have also incorporated (cid:96) 1 -regularization in these algorithms by the built-in function in Tensorﬂow/Pytorch, [Table] which amounts to adding the subgradient of the (cid:96) 1 -norm to the gradient of the loss function.
For the proposed Prox-SGD, (cid:15) ( t ) and ρ ( t ) decrease over the iterations as follows,
[Equation]
Recall that the (cid:96) 1 -norm in the approximation subproblem naturally leads to the soft-thresholding proximal mapping [Table] see (10).
The regularization parameter µ in the soft-thresholding then permits controlling the sparsity of the parameter variable x ; in this experiment we set µ = 5 · 10 − 5 .
[Figure] [Table]
[Figure]
[Table]
[Figure] [Table]
[Figure] [Table]
[Figure] [Table] 3 Weights [Figure]
CD F
[-ADAMAMSAdaBoundProx-SGD -1 
-]ADAMAMSADABoundProx-SGD 0 0.02 0.04 0.
(c) CDF of Weights
In Figure 1, we compare the four algorithms (Prox-SGD, ADAM, AMSGrad, ADABound) in terms of three metrics, namely, the training loss, the test accuracy and the achieved sparsity.
On the one hand, Figure 1(a) shows that Prox-SGD outperforms ADAM, AMSGrad and ADABound in the achieved loss value. On the other hand, the accuracy achieved by these algorithms is comparable, see Figure 1(b).
The sparsity of the trained model is measured by the cumulative distribution function (CDF) of the weights’ value, which speciﬁes the percentage of weights before any given value.
For the proposed
Prox-SGD in Figure 1 (c), we can observe around 0 in the x-axis the abrupt change of the CDF in the y-axis, which implies that more than 90% of the weights are exactly zero.
By comparison, only 40%-50% are exactly zero by the other algorithms.
What is more, for this experiment, the soft-thresholding proximal operator in Prox-SGD does not increase the computation time: ADAM 17.24s (per epoch), AMSGrad 17.44s, ADABound 16.38s, Prox-SGD 16.04s.
Therefore, in this experiment, the proposed Prox-SGD with soft-thresholding proximal mapping has a clear and substantial advantage than other stochastic subgradient-based algorithms.
3. S PARSE N EURAL N ETWORK : T RAINING D ENSE N ET -201 ON CIFAR-
In this subsection, the performance of Prox-SGD is evaluated by a larger and more complex network and dataset. In particular, we train the DenseNet-201 network (Huang et al., 2017) for CIFAR- (Krizhevsky, 2009). DenseNet-201 is the deepest topology of the DenseNet family and belongs to the state of the art networks in image classiﬁcation tasks.
We train the network using ProxSGD, ADAM and SGD with momentum. To ensure a fair comparison among these algorithms, the learning rate is not explicitly decayed during training for all algorithms. The ideal hyperparameters for each algorithm were computed using a grid-search and the curves are averaged over ﬁve runs for each algorithm. All algorithms use a batch-size of 128.
For Prox-SGD, the regularization parameter is µ = 10 − 5 , the learning rate for the weight factor and momentum is, respectively, (cid:15) ( t )
= 0 .
( t + 4) 0 . 5 , ρ ( t ) = 0 .
( t + 4) 0 .
For ADAM, (cid:15) = 6 · 10 − 4 and ρ = 0 . SGD with momentum uses a learning rate of (cid:15) = 6 · 10 −and a momentum of 0.9 (so ρ = 0 . 1 ). The regularization parameter for both ADAM and SGD with momentum is µ = 10 − 4 .
[Figure]
Figure 2 shows the performance of Prox-SGD and other algorithms for DenseNet-201 trained on CIFAR-100.
In Figure 2(a) the training loss of the three algorithms is shown. Prox-SGD achieves the lowest training after Epoch 10.
The test set accuracy in Figure 2(b) shows that all algorithms achieve similar accuracy and Prox-SGD outperforms the other two during the earlier stages of training. We remark that this is achieved with a much sparser network as shown in Figure 2(c). In particular, we can see from the zoomed-in part of Figure 2(c) that SGD with momentum has approximately 70% of their weights at zero, while most weights learned by ADAM are not exactly zero (although they are very small). In contrast, Prox-SGD reaches the sparsity of 92-94%.
[Figure] [Figure] In Figure [Figure]
we demonstrate that Prox-SGD is much more efﬁcient in generating a SNN, irrespective of the hyperparameters (related to the learning rate). In particular, we try many different initial learning rate of the weight vector (cid:15) (0) for Prox-SGD and test their performance. From Figure 3(a)-(b) we easily see that, as expected [Figure]
the hyperparameters affect the achieved training loss and test accuracy [Figure] and many lead to a worse training loss and/or test accuracy than ADAM and SGD with momentum. However, Figure 3(c) shows that most of them (except when they are too small: 0.01 and 0.001) generate a much sparser NN than both ADAM and SGD with momentum.
These observations are also consistent with the theoretical framework in Section 2: interpretating ADAM and SGD with momentum as special cases of Prox-SGD implies that they have the same convergence rate, and the sparsity is due to the explicit use of the nonsmooth (cid:96) 1 -norm regularization. For this experiment, the soft-thresholding proximal operator in Prox-SGD increases the training time: the average time per epoch for Prox-SGD is 3.5 min, SGD with momentum 2.8 min and ADAM 2.9 min.
In view of the higher level of sparsity achieved by Prox-SGD, this increase in computation time is reasonable and affordable.
3. B INARY N EURAL N ETWORKS : T RAINING D EEP N EURAL N ETWORKS ON MNIST
In this subsection, we evaluate the proposed algorithm Prox-SGD in training the BNN by solving problem (3).
We train a 6-layer fully-connected deep neural network (DNN) for the MNIST dataset, and we use the tanh activation function to promote a binary activation output; see Table 3. The algorithm parameters are the same as Sec. 3.1, except that µ = 2 · 10 − 4 .
The chosen setup is particularly suited to evaluate the merit of the proposed method, since MNIST is a simple dataset and it allows us to investigate soly the effect of the proposed model and training algorithm.
[Table]
After customizing the general description in Algorithm 1 to problem (3), the approximation subproblem is
[Equation]
)( a − a ( t )) (cid:41) . Both (cid:98) x ( t ) and (cid:98) a ( t ) have a closed-form expression (cf.
(9) and (12))
(cid:98) x ( t ) =
x ( t ) − v x ( t )
τ x ( t ) , and (cid:98) a ( t ) =
(cid:20) a ( t ) − v a ( t )
τ a ( t ) (cid:21) , (21) where v x ( t ) and v a ( t ) are the momentum updated in the spirit of (6), with the gradients given by g x ( t )
= m ( t ) (cid:88) i
∈ M ( t ) ∇ f i ( x ( t ) , y ( i ) ) [Equation]
µ 2 ( x ( t ) [Equation] 2 a ( t ) − 1 ) , and g a ( t ) = µ x ( t ) .
[Figure]
Accuracy
[Figure]
) CDF [Figure] The training loss is shown in Figure 4 (a).
We remark that during the training process of Prox-SGD, the weights are not binarized, for the reason that the penalty should regularize the problem in a way such that the optimal weights (to which Prox-SGD converges) are exactly or close to 1 or -1.
After training is completed, the CDF of the learned weights is summarized in Figure 4(c), and then the learned weights are binarized to generate a full BNN whose test accuracy is in Figure 4 (b).
On the one hand, we see from Figure 4 (a)-(b) that the achieved training loss and test accuracy by BNN is worse than the standard full-precision DNN (possibly with soft-thresholding).
This is expected as BNN imposes regularization and constraints on the optimization problem and reduces the search space.
However, the difference in test accuracy is quite small.
On the other hand, we see from
Figure 4(c) that the regularization in the proposed formulation (3) is very effective in promoting binary weights: 15% of weights are in the range (-1,-0.5) and 15% of weights are in the range (0.5,1), and all the other weights are either -1 or 1.
As all weights are exactly or close to 1 or -1, we could just binarize the weights to exactly 1 or -1 only once by hard thresholding, after the training is completed, and thus the incurred performance loss is small (98% versus 95% for test accuracy).
In contrast, the weights generated by the full-precision DNN (that is, without regularization) are smoothly distributed in [ − 2 , 2] .
Even though the proposed formulation (3) doubles the number of parameters to optimize (from x in full-precision DNN to ( x , a ) in BNN Prox-SGD), the convergence speed is equally fast in terms of the number of iterations. The computation time is also roughly the same: full-precision DNN 13.06s (per epoch) and Prox-SGD 12.21s. We remark that g a ( t ) , the batch gradient w.r.t. a , has a closedform expression and it does not involve the back-propagation.
In comparison with the algorithm in Courbariaux et al.
(2016), the proposed Prox-SGD converges much faster and achieves a much better training loss and test accuracy (95% versus 89%, the computation time per epoch for Courbariaux et al. (2016) is 13.56s).
The notable performance improvement is due to the regularization and constraints.
Naturally we should make an effort of searching for a proper regularization parameter µ , but this effort is very well paid off.
Furthermore, we observe in the simulations that the performance is not sensitive to the exact value of µ , as long as it is in an appropriate range.
C ONCLUDING R EMARKS
In this paper, we proposed Prox-SGD, a proximal-type stochastic gradient descent algorithm with momentum, for constrained optimization problems where the smooth loss function is augmented by a nonsmooth and convex regularization.
We considered two applications, namely the stochastic training of SNN and BNN, to show that regularization and constraints can effectively promote structures in the learned network.
More generally, incorporating regularization and constraints allows us to use a more accurate and interpretable model for the problem at hand and the proposed convergent Prox-SGD algorithms ensures efﬁcient training.
Numerical tests showed that Prox-SGD outperforms state-of-the-art algorithms, in terms of convergence speed, achieved training loss and/or the desired structure in the learned neural networks. A A PPENDIX : P ROOF OF T HEOREM Proof. The claim lim t →∞ (cid:107) v ( t ) −∇ f ( x ( t ))
(cid:107) = 0 is a consequence of (Ruszczy´nski, 1980, Lemma 1).
To see this, we just need to verify that all the technical conditions therein are satisﬁed by the problem at hand.
Speciﬁcally, Condition (a) of (Ruszczy´nski, 1980, Lemma 1) is satisﬁed because X is closed and bounded.
Condition (b) of (Ruszczy´nski, 1980, Lemma 1) is exactly (18).
Conditions (c)-(d) of (Ruszczy´nski, 1980, Lemma 1) come from the stepsize rules in (19) of Theorem 1.
Condition (e) of (Ruszczy´nski, 1980, Lemma 1) comes from the Lipschitz property of ∇ f and stepsize rule in (19) of Theorem 1.
We need the following intermediate result to prove the limit point of the sequence x ( t ) is a stationary point of (1).
Lemma 1.
There exists a constant (cid:98) L such that
[Equation]
, and lim t 1 ,t 2 →∞ e ( t 1 , t 2 ) =
[Equation]
w.p.1.
Proof.
We assume without loss of generality (w.l.o.g.)
that τ ( t ) = τ 1 , and the approximation subproblem (7) reduces to
[Equation]
It is further equivalent to
[Equation]
[Equation]
where the (unique) optimal x and y is ( (cid:98) x ( t ) and r ( (cid:98) x ( t )) , respectively.
We assume w
[Equation]
l.o
[Equation]
g. that t 2 > t 1 .
It follows from ﬁrst-order optimality condition that
[Equation]
[Equation]
0 , ∀ x , y such that r ( x ) ≤ y. (23b)
Setting ( x , y )
= ( (cid:98) x
( t 2 ) , r ( (cid:98) x ( t 2 ))) in (23a) and ( x , y )
= ( (cid:98) x
( t 1 ) , r ( (cid:98) x ( t 1 ))) in (23b), and adding them up, we obtain
[Equation]
The term on the left hand side can be lower bounded as follows:
[Equation]
t 1 )
[Equation]
t
[Equation]
[Equation]
(25) where the inequality comes from the Lipschitz continuity of ∇ f ( x ) , with ε ( t ) (cid:44) (cid:107) v ( t ) − ∇ f ( x ( t ))
(cid:107) .
Combining the inequalities (24) and (25), we have
[Equation]
which leads to the desired (asymptotic) Lipschitz property:
[Equation]
[Table]
w
Deﬁne U ( x ) (cid:44) f
( x ) + r ( x ) .
Following the line of analysis from (15) to (16), we obtain
[Equation]
[Equation]
where in the last inequality we used (14) and the Cauchy-Schwarz inequality.
Let us show by contradiction that lim inf t
[Equation]
= 0 w.p.1.
Suppose lim inf t →∞ (cid:13)(cid:13)(cid:98) x
( t ) − x ( t ) (cid:13)(cid:13) ≥ χ > 0 with a positive probability.
Then we can ﬁnd a realization such that at the same time (cid:13)(cid:13)(cid:98) x ( t ) − x ( t ) (cid:13)(cid:13) ≥ χ > 0 for all t and lim t →∞ (cid:13)(cid:13) ∇ f ( x ( t )) − v ( t ) (cid:13)(cid:13)
= 0 ; we focus next on such a realization.
Using (cid:13)(cid:13)(cid:98) x
( t ) − x ( t ) (cid:13)(cid:13) ≥ χ > 0 , the inequality (27) is equivalent to
[Equation]
Since lim t →∞ (cid:13)(cid:13) ∇ f ( x ( t )) − v ( t ) (cid:13)(cid:13) = 0 , there exists a t 0 sufﬁciently large such that
[Equation]
Therefore, it follows from (28) and (29) that
[Equation]
which, in view of (cid:80) ∞ n = t 0 (cid:15) n +1 = ∞ , contradicts the boundedness of { U ( x ( t )) } .
Therefore it must be lim inf t →∞ (cid:107) (cid:98) x ( t ) − x ( t ) (cid:107) = 0 w.p.1.
Let us show by contradiction that lim sup t →∞ (cid:107) (cid:98) x ( t ) − x ( t ) (cid:107) = 0 w.p.1.
Suppose lim sup t →∞ (cid:107) (cid:98) x
( t ) − x ( t ) (cid:107) > 0 with a positive probability.
We focus next on a realization along with lim sup t →∞ (cid:107) (cid:98) x ( t ) − x ( t ) (cid:107) > 0 , lim t →∞ (cid:13)(cid:13) ∇ f ( x ( t )) − v ( t ) (cid:13)(cid:13)
= 0 , lim inf t →∞ (cid:13)(cid:13)(cid:98) x
( t ) − x ( t ) (cid:13)(cid:13)
= 0 , and lim t i ,t 2 →∞ e ( t 1 , t 2 )
= 0 , where e ( t 1 , t 2 ) is deﬁned in Lemma 1.
It follows from lim sup t →∞ (cid:107) (cid:98) x ( t ) − x ( t ) (cid:107) > 0 and lim inf t →∞ (cid:13)(cid:13)(cid:98) x
( t ) − x ( t ) (cid:13)(cid:13)
{+= 
+}
that there exists a δ > 0 such that (cid:107)(cid:52) x
( t ) (cid:107) ≥ 2 δ (with (cid:52) x ( t ) (cid:44) (cid:98) x ( t ) − x ( t ) ) for inﬁnitely many t and also (cid:107)(cid:52) x ( t ) (cid:107) <
δ for inﬁnitely many t .
Therefore, one can always ﬁnd an inﬁnite set of indexes, say T , having the following properties: for any t ∈ T , there exists an integer i t > t such that
[Equation]
Given the above bounds
[Equation]
the following holds: for all t ∈ T ,
[Equation]
32) implying that
[Equation]
Proceeding as in (32), we also have: for all t ∈ T ,
[Equation]
+
[Equation]
( t
[Equation]
t
[Equation]
[Equation]
which leads to
[Equation]
where the second inequality follows from (31).
It follows from (34) that there exists a ¯ δ 2 > 0 such that for sufﬁciently large t ∈ T ,
[Equation]
Here after we assume w.l.o.g. that (35) holds for all t ∈ T (in fact one can always restrict { x ( t ) } t ∈T to a proper subsequence).
We show now that (33) is in contradiction with the convergence of { U ( x ( t )) } .
Invoking (27), we have for all t ∈ T ,
[Equation]
and for t < n < i t ,
U
[Equation]
[Equation]
where the last inequality follows from (31).
Adding (36) and (37) over n = t + 1 , . .
, i t −and, for t ∈ T sufﬁciently large (so that τ − L(cid:15)
( t ) / 2 − δ − 1 (cid:13)(cid:13) ∇ f ( x ( n )) − v ( n ) (cid:13)(cid:13) ≥ (cid:98) τ > 0 and (cid:13)(cid:13) ∇ f ( x ( t )) − v ( t ) (cid:13)(cid:13) < (cid:98)
τ ¯ δ 22 /δ ), we have
U (
[Equation]
where (a) follows from τ − L(cid:15)
( t ) / 2 − δ − 1 (cid:13) (cid:13) ∇ f ( x ( n )) − v ( n ) (cid:13) (cid:13) ≥ (cid:98) τ > 0 ; (b) is due to (35); and in (c) we used (cid:13)(cid:13) ∇ f ( x ( t )) − v ( t ) (cid:13)(cid:13) < (cid:98)
τ ¯ δ 22 /δ .
Since { U ( x ( t )) } converges, it must be lim inf T (cid:51) t →∞ (cid:80) i t − 1 n = t +1 (cid:15) ( n ) = 0 , which contradicts (33).
Therefore, it must be lim sup t →∞ (cid:13)(cid:13)(cid:98) x ( t ) − x ( t ) (cid:13)(cid:13) = 0 w.p.1.
Finally, let us prove that every limit point of the sequence { x ( t ) } is a stationary solution of (1).
Let x (cid:63) be the limit point of the convergent subsequence { x ( t ) } t ∈T .
Taking the limit of (23a) over the index set T (and replacing w.l.o.g.
y by r ( x ) ), we have
[Equation]
where the last equality follows from: i)
lim t
[Equation]
, and ii)
[Equation]
x
[Equation]
− x ( t ) (cid:13)(cid:13) = 0 .
This is the desired ﬁrst-order optimality condition and x (cid:63) is a stationary point of (1).