U NDERSTANDING S TRAIGHT -T HROUGH E STIMATOR IN T RAINING A CTIVATION Q UANTIZED N EURAL N ETS
Anonymous authors
Paper under double-blind review
A BSTRACT
Training activation quantized neural networks involves minimizing a piecewise constant training loss whose gradient vanishes almost everywhere, which is undesirable for the standard back-propagation or chain rule.
An empirical way around this issue is to use a straight-through estimator (STE) (Bengio et al., 2013) in the backward pass, so that the “gradient” through the modiﬁed chain rule becomes non-trivial.
Since this unusual “gradient” is certainly not the gradient of training loss function, the following question arises naturally: why searching in its negative direction minimizes the training loss?
In this paper, we provide the theoretical justiﬁcation of the concept of STE by answering this question. We consider the problem of learning a two-linear-layer network with binarized ReLU activation and Gaussian input data.
We shall refer to the unusual “gradient” given by the
STE-modifed chain rule as coarse gradient.
Apparently, the choice of STE is not unique.
We prove that if the STE is properly chosen, the expected coarse gradient correlates positively with underlying true gradient (not available for the training), and its negation is a descent direction for minimizing the population loss. We further show the associated coarse gradient descent algorithm converges to a local minimum (more rigorously, a critical point) of the population loss minimization problem.
Moreover, we show that a relatively poor choice of STE may lead to instability of the training algorithm near certain local minima, which is also observed in our CIFAR-10 experiments.
I NTRODUCTION
Deep neural networks (DNN) have achieved the remarkable success in many machine learning applications such as computer vision (Krizhevsky et al., 2012;
Ren et al., 2015), natural language processing (Collobert & Weston, 2008) and reinforcement learning (Mnih et al., 2015; Silver et al., 2016).
However, the deployment of DNN typically require hundreds of megabytes of memory storage for the trainable full-precision ﬂoating-point parameters, and billions of ﬂoating-point operations to make a single inference.
To achieve substantial memory savings and energy efﬁciency at inference time, many recent efforts have been made to the training of coarsely quantized DNN, in the hope of maintaining the performance of their ﬂoat counterparts (Courbariaux et al., 2015; Rastegari et al., 2016;
Cai et al., 2017).
Training fully quantized DNN amounts to solving a very challenging optimization problem.
It calls for minimizing a piecewise constant and highly nonconvex empirical risk function f ( w ) subject to a discrete set-constraint w ∈ Q that characterizes the quantized weights.
In particular, weight quantization of DNN have been extensively studied in the literature; see for examples (Li et al., 2016; Zhu et al., 2016;
Zhou et al., 2017; Li et al., 2017; Hou & Kwok, 2018).
On the other hand, the gradient ∇ f ( w ) in training activation quantized DNN is almost everywhere (a.e.) zero, which makes the standard back-propagation inapplicable.
The arguably most effective way around this issue is nothing but to construct a non-trivial search direction by properly modifying the chain rule.
Speciﬁcally, one can replace the a.e. zero derivative of quantized activation function composited in the chain rule with a related surrogate.
This proxy derivative used in the backward pass only is referred as the straight-through estimator (STE) (Bengio et al., 2013). In the same paper, Bengio et al.
(2013) proposed an alternative approach based on stochastic neurons.
In addition, Friesen & Domingos (2017) proposed the feasible target propagation algorithm for learning hard-threshold (or binary activated) networks, which solves target propagation (Lee et al., 2015) via convex combinatorial optimization.
1. R ELATED W ORKS
The idea of STE dates back to the perceptron algorithm (Rosenblatt, 1957; 1962) proposed in 1950s for learning single-layer models with binary output, where the derivative of identity function was used as a proxy. Its convergence has been extensive discussed in the literature; see for examples, (Widrow & Lehr, 1990; Freund & Schapire, 1999) and the references therein. Same as the identity STE, Hinton (2012) proposed to back-propagate through the hard threshold function as if it had been the identity function, in training multi-layer networks with the hard threshold activation 1 { x> 0 } (a.k.a. binary neuron).
Bengio et al. (2013) proposed a STE variant which uses the derivative of the sigmoid function instead.
In the training of DNN with weights and activations constrained to ± 1 , (Hubara et al., 2016) substituted the derivative of the signum activation function with 1 {| x |≤ 1 } in the backward pass, known as the saturated STE.
Later the idea of STE was readily extended to the training of DNN with general quantized ReLU activations (Hubara et al., 2018;
Zhou et al., 2016; Cai et al., 2017; Choi et al., 2018), where some other proxies took place including the derivatives of vanilla ReLU and clipped ReLU.
Despite all the empirical success of STE, there is limited theoretical understanding of it in training DNN with step activations, besides the analysis of perceptron algorithms for single-layer models. Goel et al.
(2018) considers a similar setting with leaky ReLU activation of a one-hidden-layer network.
They showed the convergence of the proposed Convotron algorithm, which uses the identity
STE in the backward pass through the leaky ReLU layer. Other similar scenarios, where the derivative of certain layer composited in the loss function is not desirable for back-propagation, have also been brought up recently by (Wang et al., 2018) and (Athalye et al., 2018).
The former proposed an implicit weighted nonlocal Laplacian layer as the classiﬁer to improve the generalization accuracy of DNN.
In the backward pass, the derivative of a pre-trained fully-connected layer was used as a surrogate.
To circumvent the defense in adversarial attack (Szegedy et al., 2013), (Athalye et al., 2018) introduced the so-called backward pass differentiable approximation to deal with the obfuscated gradients, which shares the same spirit as STE.
Neither of these two papers theoretically justiﬁed the proposed training approach.
1. M OTIVATIONS AND C ONTRIBUTIONS
Throughout this paper, we shall refer to the “gradient” through the chain rule modiﬁed by STE as coarse gradient.
The coarse gradient is certainly not the gradient of training loss function, and it is generally not the gradient of any function.
Then it is natural to ask: why searching in its negative direction minimizes the training loss, since this is not the standard gradient descent? Moreover, the choice of STE is not unique, then what makes a good STE? From the optmization perspective, we make a deeper understanding towards STE and coarse gradient in training quantized ReLU nets, by presenting a principled framework to analyze general STEs including the simplest identity STE.
On the theoretical side, we consider three representative STEs for learning a two-linear-layer network with binary activation and Gaussian data: the derivatives of the identity function (Rosenblatt, 1957; Hinton, 2012; Goel et al., 2018), vanilla ReLU and the clipped ReLUs (Cai et al., 2017; Hubara et al., 2016).
We adopt the model of population loss minimization (Brutzkus & Globerson, 2017; Tian, 2017; Li & Yuan, 2017; Du et al., 2018). We show that the negative expected coarse gradients based on the vanilla and clipped ReLUs are descent directions for the minimizing the population loss, which yields monotonically descending energy in the training. In contrast, this is not true for the the identity STE (Lemma 10).
Moreover, we prove that the former guarantees the convergence to a critical point (a saddle point or a (local) minimizer) and the latter can be unstable sometimes near certain local minima (Theorem 1). Convergence to global minima can be further proved for properly initialized weights (Theorem 2 in Appendix F).
As complementary studies, we further compare the empirical performances of the STEs on MNIST and CIFAR-10 classiﬁcations with general quantized ReLU. Both vanilla and clipped ReLU work well on the relatively shallow LeNet-5. When it comes to the deeper VGG-11 and ResNet-20, the clipped ReLU STE is the best. In our CIFAR experiments in section 4.2, we observe that the training with identity STE is indeed unstable at good minima and can one converge to an inferior one with higher training loss and decreased generalization accuracy.
This is an implication of the poor performance of certain STEs, not because of the slow convergence of the corresponding coarse gradient descent, instead due to the fact that the algorithm can never reach a good local minimum. To our knowledge, existing convergence guarantees of training with the identity STE were established for the perceptron algorithm (Rosenblatt, 1957; 1962) and Convertron algorithm (Goel et al., 2018). It is worth noting that Convertron (Goel et al., 2018) assumes symmetric distribution of the data, weaker than in this paper.
These results, however, do not generalize to the network with two trainable layers studied in this paper. As aforementioned, the identity STE is actually a poor choice in this case. Moreover, it is not clear if their analyses can be applied to other STEs.
Similar to Convertron with leaky ReLU, the monotonicity of quantized activation function plays a role in coarse gradient descent. Indeed, all the STEs considered here exploit this property. But this is not the whole story. A good STE like the clipped ReLU also needs to match with quantized ReLU at the extrema, because otherwise the instability issue may arise (section 4.2).
Notations. (cid:107) · (cid:107) denotes the Euclidean norm of a vector or the spectral norm of a matrix.
0 n ∈ R
n represents the vector of all zeros, whereas 1 n ∈ R n the vector of all ones.
I n is the identity matrix of order n .
For any w , z ∈ R n , w (cid:62) z = (cid:104) w , z (cid:105) =
(cid:80)
i w i
z i is their inner product.
w (cid:12) z denotes the Hadamard product whose i th entry is given by ( w (cid:12) z ) i =
w i
z i .
L EARNING T WO -L INEAR -L AYER CNN WITH B INARY A CTIVATION
In this paper, we consider a two-linear-layer network model (Du et al., 2018) that outputs the prediction y
( Z , v , w ) :
= m (cid:88) i = v i σ ( Z (cid:62)
[Equation]
for some input Z ∈ R m × n .
Here w ∈ R n and v ∈ R m are the trainable weights in the ﬁrst and second linear layer, respectively; Z (cid:62) i denotes the i th row vector of Z ; the activation function σ acts component-wise on the vector Z w , i.e., σ ( Z w ) i
= σ (( Z w ) i )
=
σ ( Z (cid:62) i w ) .
The ﬁrst layer serves as a convolutional layer, where each row Z (cid:62) i can be viewed as a patch sampled from Z and the weight ﬁlter w is shared among all patches, and the second linear layer is the classiﬁer.
The label is generated according to y ∗ ( Z ) =
( v ∗ ) (cid:62) σ ( Z w ∗ ) for some true (non-zero) parameters v ∗ and w ∗ .
Moreover, we use the following squared sample loss
[Equation]
Unlike in (Du et al., 2018), the activation function σ here is not ReLU, but the binary function σ ( x ) = 1 { x> 0 } , same as the hard threshold activation (Hinton, 2012).
We assume that the entries of Z ∈ R m × n are i.i.d.
sampled from the Gaussian distribution N (0 , 1) (Zhong et al., 2017; Brutzkus & Globerson, 2017).
The legitimacy of this assumption comes from the use of batch normalization (Ioffe & Szegedy, 2015) in most architectures, which sends normalized inputs to the linear layers.
Since (cid:96) ( v , w ; Z ) = (cid:96) ( v , w /c ; Z ) for any scalar c > 0 , without loss of generality, we take (cid:107) w ∗ (cid:107) = 1 and cast the learning task as the following population loss minimization problem:
[Equation]
where the sample loss (cid:96) ( v , w ; Z ) is given by (1).
2. B ACK - PROPAGATION AND C OARSE G RADIENT D ESCENT
With the Gaussian assumption on Z , as will be shown in section 2.2, it is possible to ﬁnd the analytic expressions of f ( v , w ) and its gradient
[Equation]
The gradient of objective function, however, is not available for the network training.
In fact, we can only access the expected sample gradient, namely,
[Equation]
We remark that E Z (cid:2) ∂(cid:96)∂
w ( v , w ; Z ) (cid:3) is not the same as ∂f∂ [Equation] ( v , w ) =
[Equation] Z [
[Equation]
By the standard back-propagation or chain rule, we readily check that
[Equation]
and
[Equation]
Note that σ (cid:48) is zero a.e., which makes (4) inapplicable to the training.
The idea of STE is to simply replace the a.e. zero component σ (cid:48) in (4) with a related non-trivial function µ (cid:48)
(Hinton, 2012; Bengio et al., 2013; Hubara et al., 2016;
Cai et al., 2017), which is the derivative of some (sub)differentiable function µ .
More precisely, back-propagation using the STE µ (cid:48) gives the following non-trivial surrogate of ∂(cid:96)∂ w ( v , w ; Z ) , to which we refer as the coarse (partial) gradient
[Equation]
Using the STE µ (cid:48) to train the one-hidden-layer convolutional neural network (CNN) with binary activation gives rise to the (full-batch) coarse gradient descent described in Algorithm 1. Input : initialization v 0 ∈ R m , w 0 ∈ R n , learning rate η . for t = 0 , 1 , . . . do v t +1 = v t − η E Z (cid:2) ∂(cid:96)∂ v ( v t , w t ; Z ) (cid:3) w t +1 =
w t − η E Z [ g µ ( v t , w t ; Z )] end for
2. P RELIMINARIES
Let us present some preliminaries about the landscape of the population loss function f ( v , w ) .
To this end, we deﬁne the angle between w and w ∗ as θ ( w , w ∗ ) := arccos (cid:16) w (cid:62) w ∗ (cid:107) w (cid:107)(cid:107) w ∗ (cid:107) (cid:17) for any w (cid:54) = 0 n .
Recall that the label is given by y ∗ ( Z ) =
( v ∗ ) (cid:62) Z w ∗ from (1), we elaborate on the analytic expressions of f ( v , w ) and ∇ f ( v , w ) .
Lemma 1.
Suppose all entries of Z follow i.i.d.
N (0 , 1) .
If w (cid:54) = 0 n , the population loss f
( v , w ) is given by
[Equation]
In addition, f ( v , w ) =
18 ( v ∗ ) (cid:62) (cid:0) [Equation]
Due to space limitation, we defer all the technical proofs to the appendix.
Lemma 2.
If w (cid:54) = 0 n and θ ( w , w ∗ ) ∈ (0 , π ) , the partial gradients of f ( v , w ) w.r.t.
v and w are ∂f
[Equation]
and ∂f
[Equation]
respectively.
For any v ∈ R m , ( v , 0 m ) is impossible to be a local minimizer.
The only possible (local) minimizers of the model (2) are located at
[Equation]
) = [Equation] w , w [Equation]
=
[Equation]
Among them, { ( v , w ) : v = v ∗ , θ ( w , w ∗ )
= 0 } are obviously the global minimizers of (2).
We show that the stationary points, if exist, can only be saddle points, and { ( v , w ) : θ ( w , w ∗ ) =
π, v = (cid:0) I m + 1 m 1 (cid:62) m (cid:1) − 1 ( 1 m 1 (cid:62) m
− I m ) v ∗ } are the only potential spurious local minimizers.
Proposition 1.
If the true parameter v ∗ satisﬁes ( 1 (cid:62) m v ∗ ) 2 < m +12 (cid:107) v ∗ (cid:107) 2 , then
[Equation]
give the saddle points obeying (8), and { ( v , w ) : θ ( w , w ∗ )
=
π, v = (cid:0) I m + 1 m 1 (cid:62) m (cid:1) − 1 ( 1 m 1 (cid:62) m −
I m ) v ∗ } are the spurious local minimizers.
Otherwise, the model (2) has no saddle points or spurious local minimizers.
We further prove that the underlying true gradient ∇ f
( v , w ) given by (6) and (7), is Lipschitz continuous under a boundedness condition.
Lemma 3.
For any differentiable points ( v , w ) and (˜ v , ˜ w ) with min {(cid:107) w (cid:107) , (cid:107) ˜ w (cid:107)} =
c w > 0 and max {(cid:107) v (cid:107) , (cid:107) ˜
v (cid:107)} =
C v , there exists a Lipschitz constant L > 0 depending on C v and c w , such that
[Equation]
M AIN R ESULTS
We are most interested in the complex case where both the saddle points and spurious local minimizers are present.
Our main results are concerned with the behaviors of the coarse gradient descent summarized in Algorithm 1 when the derivatives of the vanilla ReLU and identity function serve as the STE, respectively.
Intuitively, the ReLU STE is supposed to outperform the identity STE, because ReLU is obviously a better approximation to the activation function σ ( x ) = 1 { x> 0 } .
Theorem 1 (Convergence) .
Let { ( v t , w t ) } be the sequence generated by Algorithm 1 with ReLU µ ( x ) = max { x, 0 } or clipped ReLU µ ( x ) = min { max { x, 0 } , 1 } .
Suppose (cid:107) v t (cid:107) ≤ C v
and (cid:107) w t (cid:107) ≥ c w for all t with some C v , c w > 0 .
Then if the learning rate η > 0 is sufﬁciently small, for any initialization ( v 0 , w 0 ) , the objective sequence { f ( v t , w t ) } is monotonically decreasing, and { ( v t , w t ) } converges to a saddle point or a (local) minimizer of the population loss minimization (2).
In addition, if 1 (cid:62) m v ∗ (cid:54) =
and m > 1 , the descent and convergence properties do not hold for Algorithm 1 with the identity function µ ( x ) = x near the local minimizers satisfying θ ( w , w ∗ )
=
π and v =
( I m + 1 m 1 (cid:62) m ) − 1 ( 1 m 1 (cid:62) m − I m ) v ∗ .
Remark 1.
The convergence guarantee for the coarse gradient descent is established under the assumption that there are inﬁnite training samples.
When there are only a few data, in a coarse scale, the empirical loss is roughly descending along the direction of negative coarse gradient, as illustrated by Figure 1.
As the sample size increases, the empirical loss gains monotonicity and smoothness.
This explains why STE works so well when there are a large number of training samples as in deep learning problems.
In the rest of this section, we sketch the mathematical analysis for the main results.
3. D ERIVATIVE OF THE V ANILLA R E LU AS STE
If we choose the derivative of ReLU µ ( x ) = max { x, 0 } as the STE in (5), it is easy to see µ (cid:48)
( x ) = σ ( x ) , and we have the following expressions of E Z (cid:2) ∂(cid:96)∂
v ( v , w ; Z ) (cid:3) and E Z (cid:104) g relu ( v , w ; Z ) (cid:105) for Algorithm 1. [Figure]
Lemma 4.
The expected partial gradient of (cid:96) ( v , w ; Z ) w.r.t.
v is
[Equation]
Let µ ( x ) = max { x, 0 } in (5).
The expected coarse gradient w.r.t.
w is
[Equation]
+
[Equation]
[Equation]
where h ( v , v ∗ ) = (cid:107) v (cid:107) {+
+}
[-[Equation] ∗ (cid:13)(cid:13)(cid:13) w (cid:107) w (cid:107) + w ∗ (cid:13)(cid:13)(cid:13) ,(11) where h ( v , v ∗ ) = (cid:107) v (cid:107) 
-] + ( 1 (cid:62) m v ) 2 − ( 1 (cid:62) m v )( 1 (cid:62) m v ∗ )
+ v (cid:62) v ∗ .
As stated in Lemma 5 below, the key observation is that the coarse partial gradient E Z (cid:104) g relu ( v , w ; Z ) (cid:105) has non-negative correlation with the true gradient ∂f∂ w ( v , w ) , and − E Z (cid:104)
g relu ( v , w ; Z ) (cid:105) together with − E Z (cid:2) ∂(cid:96)∂
v ( v , w ; Z ) (cid:3) form a descent direction for minimizing the population loss.
Lemma 5.
If w (cid:54) = 0 n and θ ( w , w ∗ ) ∈ (0 , π ) , then the inner product between the expected coarse and true gradients w.r.t.
w is
[Equation]
Moreover, if further (cid:107) v (cid:107) ≤ C v and (cid:107) w (cid:107) ≥ c w , there exists a constant A relu > 0 depending on C v and c w , such that
[Equation]
Clearly, when (cid:68) E Z (cid:104)
g relu ( v , w ; Z ) (cid:105) [Equation] ∂f∂ w ( v , w ) (cid:69) > 0 , E Z (cid:104)
g relu ( v , w ; Z ) (cid:105) is roughly in the same direction as ∂f∂ w ( v , w ) .
Moreover, since by Lemma 4, E Z (cid:2) ∂(cid:96)∂
v ( v , w ; Z ) (cid:3) =
∂f ∂ v ( v , w ) , we expect that the coarse gradient descent behaves like the gradient descent directly on f ( v , w ) .
Here we would like to highlight the signiﬁcance of the estimate (12) in guaranteeing the descent property of
[Equation]
a) is due to (12).
Therefore, if η is small enough, we have monotonically decreasing population loss until convergence.
Lemma 6.
When Algorithm 1 converges, E Z (cid:104) ∂(cid:96)∂
v ( v , w ; Z ) (cid:105) and E Z (cid:104)
g relu ( v , w ; Z ) (cid:105) vanish simultaneously, which only occurs at the Saddle points where (8) is satisﬁed according to Proposition 1. Minimizers of (2) where v
[Equation]
Lemma 6 states that when Algorithm 1 using ReLU STE converges, it can only converge to a critical point of the original loss function.
We emphasize that this is not necessarily true for any STE, because we are not using the true gradient of population loss in the training.
3. D ERIVATIVE OF THE C LIPPED R E LU AS STE
For the STE using clipped ReLU, µ ( x ) = min { max { x, 0 } , 1 } and µ (cid:48) ( x ) = 1 { 0 <x< 1 } ( x ) .
We have results similar to Lemmas 5 and 6.
That is, the coarse partial gradient using clipped ReLU
STE E Z (cid:104) g crelu ( v , w ; Z ) (cid:105) generally has positive correlation with the true partial gradient of the population loss ∂f∂ w ( v , w ) (Lemma 7)).
Moreover, the coarse gradient vanishes and only vanishes at the critical points (Lemma 8).
Lemma 7.
If w (cid:54) = 0 n and θ ( w , w ∗ ) ∈ (0 , π ) , then
[Equation]
v ∗ ) csc( θ/ 2) · q ( θ, w ) w (cid:107) w (cid:107) + w ∗ (cid:13)(cid:13)(cid:13) w (cid:107) w (cid:107) + w ∗ (cid:13)(cid:13)(cid:13) − ( v (cid:62) v ∗ ) ( p ( θ, w ) − cot( θ/ 2) · q ( θ, w ))
w (cid:107) w (cid:107) , where h ( v , v ∗ ) :
{+= (cid:107) v (cid:107) 
+}
+ ( 1 (cid:62) m v ) 2 − ( 1 (cid:62) m v )( 1 (cid:62) m v ∗ )
+ v (cid:62) v ∗ same as in Lemma 5, and p ( θ, w ) :
= 1 2 π (cid:90) π−
π 2 + θ cos( φ ) ξ (cid:18) sec( φ ) (cid:107) w (cid:107) (cid:19)
d φ, q ( θ, w ) :
= 1 2 π (cid:90) π−
π
2 + θ sin( φ ) ξ (cid:18) sec( φ ) (cid:107) w (cid:107) (cid:19)
d φ with ξ ( x ) :
= (cid:82) x 0 r 2 exp( − r 2 2 )d r .
The inner product between the expected coarse and true gradients w.r.t. w (cid:28) E Z (cid:104) g crelu ( v , w ; Z ) (cid:105) , ∂f ∂ w ( v , w ) (cid:29) ≥ 0 . Moreover, if further (cid:107) v (cid:107) ≤ C v and (cid:107) w (cid:107) ≥ c w , there exists a constant A crelu > 0 depending on C v and c w , such that
[Equation] Lemma 8. When Algorithm 1 converges, E Z (cid:104) ∂(cid:96)∂ v ( v , w ; Z ) (cid:105) and E Z (cid:104)
g crelu ( v , w ; Z ) (cid:105) vanish simultaneously, which only occurs at the 1. Saddle points where (8) is satisﬁed according to Proposition 1.
Minimizers of (2) where v = v ∗ , θ ( w , w ∗ ) = 0 , or v =
( I m + 1 m 1 (cid:62) m ) − 1 ( 1 m 1 (cid:62) m −
I m ) v ∗ , θ ( w , w ∗ ) = π .
3. D ERIVATIVE OF THE I DENTITY F UNCTION AS STE
Now we consider the derivative of identity function.
As opposed to the ReLU case, similar results to
Lemmas 5 and 6 are not valid anymore.
It happens that the coarse gradient derived from the identity
STE does not vanish at the local minimum, and Algorithm 1 may never converge there.
Lemma 9.
Let µ ( x ) =
x
in (5).
Then the expected coarse partial gradient w.r.t.
w is E
[Equation]
Z [Equation] If θ ( w , w ∗ ) = π and v = (cid:0) I m + 1 m 1 (cid:62)
m (cid:1) − 1 ( 1 m 1 (cid:62) m − I m ) v ∗ ,
[Equation]
{+i.e., E Z (cid:104) g id ( v , w ; Z ) (cid:105) does not vanish at the local minimizers if 1 (cid:62) m v ∗ (cid:54) = 
 and m > 1 .+}
Lemma 10.
If w (cid:54) = 0 n and θ ( w , w ∗ ) ∈ (0 , π ) , then the inner product between the expected coarse and true gradients w.r.t.
w is
[Equation]
When θ ( w , w ∗ ) → π , v → (cid:0) I m + 1 m 1 (cid:62)
m (cid:1) − 1 ( 1 m 1 (cid:62) m −
I m ) v ∗ , if 1 (cid:62) m v ∗ (cid:54) =
and m > 1 , we have
[Equation]
[Table] [Equation] → + ∞ . Lemma 9 suggests that if 1 (cid:62) m v ∗ (cid:54) = 0 , the coarse gradient descent will never converge near the spurious minimizers with θ ( w , w ∗ )
=
π and v = (cid:0) I m + 1 m 1 (cid:62) m (cid:1) − 1 ( 1 m 1 (cid:62) m − I m ) v ∗ , because
E Z (cid:104) g id ( v , w ; Z ) (cid:105) does not vanish there.
By the positive correlation implied by (15) of Lemma 10, for some proper ( v 0 , w 0 ) , the iterates { ( v t , w t ) } may move towards a spurious local minimizer in the beginning.
But when { ( v t , w t ) } approaches it, the descent property (13) does not hold for E Z [ g id ( v , w ; Z )] because of (16), hence the instability arises and the training loss begins to increase.
E XPERIMENTS
In practice, the vanilla ReLU may not deliver the best empirical performance.
Compared with the identity function and vanilla ReLU, the clipped ReLU proposed in (Cai et al., 2017) clearly better approximates the quantized ReLU.
As an example, we plot the 2-bit quantized ReLU and its associated clipped ReLU in Figure 3 in the appendix.
Besides the two STEs discussed above, we include the STE using derivative of the clipped ReLU for comparisons on training 2-bit and 4-bit activation networks for MNIST (LeCun et al., 1998) and CIFAR-10 (Krizhevsky, 2009) classiﬁcations.
Logtailed ReLU has also been proposed, we do not consider it here since it gives similar performance to the clipped ReLU as reported in (Cai et al., 2017).
We emphasize that we are not claiming the superiority of the quantization approach used here, as it is nothing but the HWGQ (Cai et al., 2017), except we consider the uniform quantization.
In all of our experiments, the weights are kept ﬂoat.
The optimizer we use is the stochastic (coarse) gradient descent with momentum = 0.9 for all experiments.
We train 50 epochs for LeNet-5 (LeCun et al., 1998) on MNIST, and 200 epochs for
VGG-11 (Simonyan & Zisserman, 2014) and ResNet-20 (He et al., 2016) on CIFAR-10.
The parameters/weights are initialized with those from their pre-trained full-precision counterparts.
The schedule of the learning rate is speciﬁed in Table 2 in the appendix.
The resolution α for the quantized ReLU needs to be carefully chosen to maintain the full-precision level accuracy.
To this end, we follow (Cai et al., 2017) and resort to a modiﬁed batch normalization layer (Ioffe & Szegedy, 2015) without the scale and shift, whose output components approximately follow a unit Gaussian distribution.
Then the α that ﬁts the input of activation layer the best can be pre-computed by a variant of Lloyd’s algorithm (Lloyd, 1982; Yin et al., 2018) applied to a set of simulated 1-D half-Gaussian data.
After determining the α , it will be ﬁxed during the whole training process.
Since the original LeNet-5 does not have batch normalization, we add one prior to each activation layer.
4. C OMPARISON R ESULTS
The experimental results are summarized in Table 1, where we record both the training losses and validation accuracies.
Among the three STEs, the derivative of clipped ReLU gives the best overall performance, followed by vanilla ReLU and then by the identity function.
For deeper networks, clipped ReLU is the best performer.
But on the shallow LeNet-5 network, vanilla ReLU exhibits comparable performance to the clipped ReLU, which is in line with our theoretical ﬁnding that ReLU is a superior STE for learning the one-hidden-layer (shallow) CNN. [Table]
4. I NSTABILITY AT G OOD M INIMA
We report the phenomenon of being repelled from a good minimum on ResNet-20 with 4-bit activations when using the identity STE to demonstrate the instability issue as predicited in Theorem 1.
By Table
[Table]
the coarse gradient descent algorithms using the vanilla and clipped ReLUs converge to the neighborhoods of the minima with validation accuracies (training losses) of 86 .
59% (0.25) and 91 .
24% (0.04), respectively, whereas that using the identity STE gives 54 .
Note that the landscape of the empirical loss function does not depend on which STE is used in the training.
Then we initialize training with the two improved minima and use the identity STE.
To see if the algorithm is stable there, we start the training with a tiny learning rate of 10 − 5 .
For both initializations, the training loss and validation error signiﬁcantly increase within the ﬁrst 20 epochs; see Figure 4.2.
To speedup training, at epoch 20, we switch to the normal schedule of learning rate speciﬁed in Table 2 and run 200 additional epochs.
The training using the identity STE ends up with a much worse minimum.
This is because the coarse gradient with identity STE does not vanish at the good minima in this case (Lemma 9).
Similarly, the poor performance of ReLU STE on 2-bit activated ResNet-  is also due to the instability of the corresponding training algorithm at good minima, as illustrated by Figure 5 in Appendix C.
C ONCLUDING R EMARKS
We provided the ﬁrst theoretical justiﬁcation for the concept of STE.
We considered two STEs: the derivatives of the identity function and the vanilla ReLU, in training one-hidden-layer CNN with binary activation.
We derived the explicit formulas of the expected coarse gradients corresponding to the two STEs, and showed that the negative expected coarse gradient based on vanilla ReLU is a descent direction for minimizing the population loss, whereas the identity STE is not.
Our experiments on MNIST and CIFAR-10 datasets veriﬁed the theoretical ﬁndings.
Looking ahead, we believe that the closeness between the quantized ReLU and the anti-derivative of STE plays an essential role in the performance of the coarse gradient descent.
Hence, we would like to explore and quantify this relationship in the future work.
A PPENDIX A. T HE P LOTS OF Q UANTIZED AND C LIPPED R E LU S
quantized ReLU clipped ReLU
B. T HE S CHEDULE OF L EARNING R ATE
[Table]
C. I NSTABILITY OF R E LU STE ON R ES N ET -20 WITH 2- BIT A CTIVATIONS
D. A DDITIONAL S UPPORTING L EMMAS
Lemma 11.
Let z ∈ R n be a Gaussian random vector with entries i.i.d.
sampled from N (0 , 1) .
Given nonzero vectors w , ˜ w ∈ R n with the angle θ , we have
[Equation]
and
[Equation]
Proof of Lemma 11.
The third identity was proved in Lemma A.1 of (Du et al., 2018).
To show the ﬁrst one, without loss of generality we assume w =
[ w 1 , 0 (cid:62) n − 1 ] (cid:62) with w 1 > 0 , then
E (cid:2) 1 { z (cid:62) w > 0 } (cid:3) =
P ( z 1 > 0)
= 12 .
We further assume ˜ w =
[ ˜ w 1 , ˜ w 2 , 0 (cid:62) n − 2 ] (cid:62) .
It is easy to see that
[Equation]
To prove the last identity, we use polar representation of two-dimensional Gaussian random variables, where r is the radius and φ is the angle with d P r = r exp( − r 2 / 2)d r and d P φ = 12 π d φ .
Then
E (cid:2)
[Equation]
= 0 for i ≥ 3 .
Moreover, E
[Equation]
and
[Equation]
Therefore,
[Equation]
, where the last equality holds because w (cid:107) w (cid:107) and w [Equation] w [Equation]
˜ w
[Equation]
˜ w (cid:107)
[Equation]
w (cid:107) w
[Equation]
+ ˜ w
[Equation] ˜ w [Equation] (cid:107) are two unit-normed vectors with angle θ/ 2 .
Lemma 12.
Let z ∈ R n be a Gaussian random vector with entries i.i.d.
sampled from N (0 , 1) .
Given nonzero vectors w , ˜ w ∈ R n with the angle θ , we have
[Equation]
Here p ( θ, w ) = 1 2 π (cid:90) π− π 2 + θ cos( φ ) ξ (cid:18) sec( φ ) (cid:107) w (cid:107) (cid:19) d φ, q ( θ, w ) = 1 2 π (cid:90) π− π
[Equation]
+ θ sin( φ ) ξ (cid:18) sec( φ ) (cid:107) w (cid:107) (cid:19)
d φ with ξ ( x ) =
(cid:82) x 0 r 2 exp( − r 2 2 )d r satisfying ξ (+ ∞ ) =
(cid:112)
π 2 .
In addition, p ( π, w ) =
q (0 , w ) =
q ( π, w ) = 0 .
2 Same as in Lemma 4, we redeﬁne E
[Equation]
˜
[Equation]
[Table]
[Table]
Proof of Lemma 12.
Following the proof of Lemma 11, we let w =
[ w 1 , 0 (cid:62) n − 1 ] (cid:62) with w 1 > 0 , ˜ w
=
[ ˜ w 1 , ˜ w 2 , 0 (cid:62) n − 2 ] (cid:62) , and use the polar representation of two-dimensional Gaussian distribution.
Then
[Equation]
and
[Equation]
since the integrand above is an odd function in φ .
Moreover, E (cid:2) z i 1 { 0 < z
[Equation]
<
[Equation]
=
[Equation]
for i ≥ 3 .
So the ﬁrst identity holds.
For the second one, we have
[Equation]
and similarly, E (cid:2) z 2 1 { 0 < z (cid:62) w < 1 , z (cid:62) ˜ w > 0 } (cid:3) =
q ( θ, w ) .
Therefore,
E [ z1 { 0 < z (cid:62) w < 1 , z (cid:62) ˜ w > 0 } ] =
[ p ( θ, w ) , q ( θ, w ) , 0 (cid:62) n − 2 ] (cid:62) .
Since w (cid:107) w (cid:107) =
[1 , 0 (cid:62) n − 1 ] (cid:62) and w (cid:107) w (cid:107) + ˜ w (cid:107) ˜
w (cid:107) (cid:107) w (cid:107) w (cid:107) + ˜ w (cid:107) ˜ w (cid:107) (cid:107) =
[cos( θ/ 2) , sin( θ/ 2) , 0 (cid:62) n − 2 ] (cid:62) , it is easy to check the second identity is also true.
Lemma 13.
Let p ( θ, w ) and q ( θ, w ) be deﬁned in Lemma 12.
Then for θ ∈ [ π 2 , π ] , we have 1.
p ( θ, w ) ≤ q ( θ, w ) .
(cid:0)
− θπ
(cid:1)
p
w ) ≤ q ( θ, w )
Proof of Lemma 13.
Let θ ∈
[ π 2 , π ] , since ξ ≥ 0 , p ( θ, w ) = 1 2 π (cid:90) πθ − πcos( φ ) ξ (cid:18) sec( φ ) (cid:107) w (cid:107) (cid:19)
d φ = 1 2 π (cid:90)
πθ − πsin (cid:16) π 2 − φ (cid:17) ξ (cid:18) sec( φ ) (cid:107) w (cid:107) (cid:19)
d φ ≤2 π (cid:90)
πθ − π
[Equation]
= q ( θ, w ) , where the last inequality is due to the rearrangement inequality since both sin( φ ) and ξ (cid:16) sec( φ ) (cid:107) w (cid:107) (cid:17) are increasing in φ on [0 , π 2 ] .
Since cos( φ ) ξ (cid:16) sec( φ ) (cid:107) w (cid:107) (cid:17) is even, we have
[Equation]
The ﬁrst inequality is due to part 1 which gives p ( π/ 2 , w ) ≤ q ( π/ 2 , w ) , whereas the second one holds because sin( φ ) ξ (cid:16) sec( φ ) (cid:107) w (cid:107) (cid:17) is increasing on [0 , π 2 ] .
Lemma 14.
For any nonzero vectors w and ˜ w with (cid:107) ˜ w (cid:107) ≥ (cid:107) w (cid:107) =
c > 0 , we have
[Table] [Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table] Proof of Lemma 14.
Since by Cauchy-Schwarz inequality,
[Equation]
we have
[Equation]
−
[Equation] [Equation]
Therefore,
[Equation]
where we used the fact sin( x ) ≥ 2 xπ for x ∈ [0 , π 2 ] and the estimate in (17). Since (cid:16) I n − ww (cid:62) [Equation]
w [Equation] 2 (cid:17) w ∗ is the projection of w ∗ onto the complement space of w , and likewise for (cid:16) I n − ˜ w ˜ w [Equation]
˜ w
[Equation] 2 (cid:17) w ∗ , the angle between (cid:16) I n − ww (cid:62) (cid:107) w (cid:107) 2 (cid:17) w ∗ and (cid:16) I n − ˜ w ˜ w (cid:62) (cid:107) ˜ w
[Equation] 2 (cid:17) w ∗ is equal to the angle between w and ˜ w .
Therefore,
[Equation]
(cid:29) , and thus (cid:13)(cid:13)(cid:13)(cid:13)(cid:13)(cid:13)(cid:107) w (cid:107) (cid:16) I n − ww (cid:62) (cid:107) w (cid:107) 2 (cid:17) w ∗ (cid:13)(cid:13)(cid:13)(cid:16)
I n − ww (cid:62) (cid:107) w (cid:107)
[Equation]
(cid:17) w ∗ (cid:13)(cid:13)(cid:13) − (cid:107) [Equation] (cid:107) (cid:16) [Equation]
(cid:107)
[Equation]
w (cid:107) [Equation] (cid:17) w ∗ (cid:13)(cid:13)(cid:13)(cid:16)
[Equation]
n − ˜ w ˜ w (cid:62) (cid:107) ˜ w (cid:107) 2 (cid:17) w ∗ (cid:13)(cid:13)(cid:13) (cid:13)(cid:13)(cid:13)(cid:13)(cid:13)(cid:13) =
(cid:13)(cid:13)(cid:13)(cid:13) w (cid:107) w (cid:107) 2 − ˜ w (cid:107) ˜ w (cid:107)(cid:13) (cid:13)(cid:13)(cid:13) = (cid:107) w − ˜ w (cid:107) (cid:107) w (cid:107)(cid:107) ˜ w (cid:107) ≤ c 2 (cid:107) w − ˜ w (cid:107) . The second equality above holds because (cid:13)(cid:13)(cid:13)(cid:13) w (cid:107) w (cid:107) 2 − ˜ w (cid:107) ˜ w (cid:107) 2 (cid:13)(cid:13)(cid:13)(cid:13)
[Equation]
E. M AIN P ROOFS
Lemma 1.
Suppose all entries of Z follow i.i.d.
N (0 , 1) .
If w (cid:54) = 0 n , the population loss f
( v , w ) is given by
[Equation]
(cid:35)
[Equation]
In addition, f ( v , w ) = 1 8 ( v ∗ ) (cid:62) (cid:0) I m + 1 m 1 (cid:62) m (cid:1) v ∗ for w = 0 n .
Proof of Lemma 1.
We notice that
[Equation]
Let Z (cid:62) i be the i th row vector of Z .
Since w (cid:54) = 0 n , using Lemma 11, we have E
[Equation]
= [Equation] 2 , and for i (cid:54)
=
j ,
E Z [Equation] σ [Equation]
Z w ) σ
[Equation]
Z w )
[Equation] = [Equation] σ
[Equation]
w )
[Equation]
( Z
[Equation]
j w ) [Equation] = E (cid:104) [Equation]
[Equation]
i w >
[Equation]
(cid:105)
[Equation] 0 } (cid:105) = 1 4 . Therefore, E Z (cid:2) σ ( Z w ) σ ( Z w ) (cid:62) (cid:3) = E Z (cid:2) σ ( Z w ∗ ) σ ( Z w ∗ ) (cid:62) (cid:3) = 14 (cid:0) I m + 1 m 1 (cid:62) m (cid:1) .
Furthermore,
E Z (cid:2) σ [Equation] Z w ) σ ( Z w [Equation]
[Equation]
(cid:3) ii = E (cid:104) [Equation] , and E Z
[Equation]
σ
[Equation]
Z w
[Equation]
Z w
[Equation]
So,
[Equation]
Then it is easy to validate the ﬁrst claim.
Moreover, if w = 0 n , then
[Equation]
Lemma 2.
If w (cid:54) = 0 n and θ ( w , w ∗ ) ∈ (0 , π ) , the partial gradients of f ( v , w ) w.r.t.
v and w are ∂f
[Equation]
and ∂f
[Equation]
w
[Equation]
, respectively.
Proof of Lemma 2.
The ﬁrst claim is trivial, and we only show the second one.
Since θ ( w , w ∗ ) =
arccos (cid:16) w (cid:62) w ∗ (cid:107) w (cid:107) (cid:17) is differentiable w.r.t.
w at θ ( w , w ∗ ) ∈ (0 , π ) , we have
[Equation]
Proposition 1.
If the true parameter v ∗ satisﬁes ( 1 (cid:62) m v ∗ [Equation] 2 < [Equation]
12 (cid:107) v ∗ (cid:107) 2 , then
[Equation]
give the saddle points obeying (8), and { ( v , w ) : θ ( w , w ∗ )
=
π, v = (cid:0) I m + [Equation] 1 (cid:62) m (cid:1) − 1 ( 1 m 1 (cid:62) m −
I m ) v ∗ } are the spurious local minimizers.
Otherwise, the model (2) has no saddle points or spurious local minimizers.
Proof of Proposition 1.
Suppose v (cid:62) v ∗ = 0 and ∂f∂ v ( v , w ) = 0 , then by Lemma 1,
[Equation] [Equation]
From (18) it follows that
[Equation]
On the other hand, from (18) it also follows that
[Equation]
where we used ( I m + 1 m 1 (cid:62) m ) 1 m = ( m [Equation] 1) 1 m .
Taking the difference of the two equalities above gives
[Equation]
By (19), we have θ ( w , w ∗ ) =
π
[Equation]
∗
[Equation]
v ∗ ) 2 , which requires π
[Equation]
( m + 1) (cid:107) v ∗ (cid:107) 2 − ( 1 (cid:62) m v ∗ ) 2 < π, or equivalently, ( 1 (cid:62) m v ∗ ) 2 <
[Equation]
Furthermore, since ∂f∂ v ( v , w ) = 0 , we have
[Equation]
Next
[Equation]
we check the local optimality of the staionary points.
By ignoring the scaling and constant terms
[Equation]
we rewrite the objective function as
[Equation]
It is easy to check that its Hessian matrix
[Equation]
is indeﬁnite.
Therefore, the stationary points are saddle points.
Moreover, if ( 1 (cid:62) m v ∗ ) 2 < m +1 2 (cid:107) v ∗ (cid:107) 2 , at the point ( v , θ ) =
(( I m + 1 m 1 (cid:62) m ) − 1 ( 1 m 1 (cid:62) m −
I m ) v ∗ , π ) , we have
[Equation]
where we used (20) in the last identity.
We consider an arbitrary point ( v + ∆ v , π + ∆ θ ) in the neighborhood of ( v , π ) with ∆ θ ≤ 0 .
The perturbed objective value is
[Equation]
On the right hand side, since v = ( I m + 1 m 1 (cid:62) m ) − 1 ( 1 m 1 (cid:62) m −
I m ) v ∗ is the unique minimizer to the quadratic function ˜ f
( v , π ) , we have if ∆ v (cid:54) =
m ,
[Equation]
Moreover, for sufﬁcently small (cid:107) ∆ v (cid:107) , it holds that ∆ θ · ( v + ∆ v ) (cid:62) v ∗ > 0 for ∆ θ < 0 because of (21).
Therefore, ˜ f ( v + ∆ v , π + ∆ θ ) >
˜ f ( v , π ) whenever (∆ v , ∆ θ ) is small and non-zero, and (( I m + 1 m 1 (cid:62) m ) − 1 ( 1 m 1 (cid:62) m −
I m ) v ∗ , π ) is a spurious local minimizer of ˜ f .
To prove the second claim, suppose ( 1 (cid:62) m v ∗ ) 2 ≥ m +12 (cid:107) v ∗ (cid:107) 2 , then either ∂f∂ w ( v , w ) does not exist, or ∂f∂ v ( v , w ) and ∂f∂ w ( v , w ) do not vanish simultaneously, and thus there is no stationary point.
At the point ( v , θ )
=
(( I m + 1 m 1 (cid:62) m ) − 1 ( 1 m 1 (cid:62) m −
I m ) v ∗ , π ) , we have
[Equation]
If v (cid:62) v ∗ > 0 , since ∇ ˜
f
( v , θ ) = 14 [ 0 (cid:62) m , 2 π v (cid:62) v ∗ ] (cid:62) , a small perturbation [ 0 (cid:62) m , ∆ θ ] (cid:62) with ∆ θ <will give a strictly decreased objective value, so ( v , θ ) =
(cid:0) ( I m + 1 m 1 (cid:62) m ) − 1 ( 1 m 1 (cid:62) m − I m ) v ∗ , π (cid:1) is not a local minimizer.
If v (cid:62) v ∗ = 0 , then ∇ ˜ f
( v , θ ) = 0 n +1 , the same conclusion can be reached by examining the second order necessary condition.
Lemma 3.
For any differentiable points ( v , w ) and (˜ v , ˜ w ) with min {(cid:107) w (cid:107) , (cid:107) ˜ w (cid:107)} =
c w > 0 and max {(cid:107) v (cid:107) , (cid:107) ˜
v (cid:107)} =
C v , there exists a Lipschitz constant L > 0 depending on C v and c w , such that
[Equation]
Proof of Lemma 3.
It is easy to check that (cid:107) I m + 1 m 1 (cid:62) m (cid:107) = m
+ 1 .
Then
[Equation]
where the last inequality is due to Lemma 14.1.
[Table]
[Equation]
where the second last inequality is to due to Lemma 14.2.
Combining the two inequalities above validates the claim.
Lemma 4.
The expected partial gradient of (cid:96) ( v , w ; Z ) w.r.t.
v is
[Equation]
Let µ ( x ) = max { x, 0 } in (5).
The expected coarse gradient w.r.t.
w is
E Z
[Equation]
g relu (
[Equation] [Equation]
[Equation]
where h ( v , v ∗ ) = (cid:107) v (cid:107)
+ ( 1 (cid:62) m v ) 2 − ( 1 (cid:62) m v )( 1 (cid:62) m v ∗ )
+ v (cid:62) v ∗ .
Proof of Lemma 4.
The ﬁrst claim is true because ∂(cid:96)∂ v ( v , w ; Z ) is linear in v .
By (5), g µ ( v , w ; Z ) =
Z (cid:62) (cid:0) µ (cid:48)
( Z w ) (cid:12) v (cid:1)(cid:16) v (cid:62) σ ( Z w ) − ( v ∗ ) (cid:62) σ ( Z w ∗ ) (cid:17) .
Using the fact that µ (cid:48)
=
σ = 1 { x> 0 } , we have
E Z [ g relu ( v , w ; Z )]
=
E Z (cid:34)(cid:32) m
(cid:88) i
= v i σ ( Z (cid:62) i w ) − m (cid:88) i = v ∗
i σ ( Z (cid:62) i w ∗ ) (cid:33) (cid:32) m (cid:88) i =
Z i v i σ ( Z (cid:62) i w ) (cid:33)(cid:35) =
E Z (cid:34)(cid:32) m
(cid:88) i
= v
i 1 { Z (cid:62) i w > 0 } − m (cid:88) i = v ∗
i 1 { Z (cid:62) i w ∗ > 0 } (cid:33) (cid:32) m (cid:88) i =
1 { Z (cid:62) i w > 0 } v i Z
i (cid:33)(cid:35) .
Invoking Lemma 11, we have
E (cid:104)
Z i 1 { Z (cid:62) i w > 0 , Z (cid:62) j w > 0 } (cid:105) =
(cid:40) 1 √ 2 π w (cid:107) w (cid:107) if i = j, 1 2 √ 2 π w (cid:107) w (cid:107) if i (cid:54) =
j, and
E (cid:104)
[Equation]
[Equation]
w
[Equation]
[Equation] (cid:105)
[Equation]
Therefore,
[Equation]
and the result follows
[Equation]
Lemma 5.
If w (cid:54) = 0 n and θ ( w , w ∗ ) ∈ (0 , π ) , then the inner product between the expected coarse and true gradients w.r.t.
w is
[Equation]
Moreover, if further (cid:107) v (cid:107) ≤ C v and (cid:107) w (cid:107) ≥ c w , there exists a constant A relu > 0 depending on C v and c w , such that
[Equation]
Proof of Lemma 5.
By Lemmas 2 and 5, we have
[Equation]
and
[Equation]
Notice that (cid:16) I n − ww (cid:62) (cid:107) w (cid:107) 2 (cid:17) w = 0 n and (cid:107) w ∗ (cid:107) =
1 , if θ ( w , w ∗ ) (cid:54) = 0 , π , then we have (cid:28) E Z (cid:104)
g relu ( v , w ; Z ) (cid:105) , ∂f ∂ w ( v , w ) (cid:29) =
cos (cid:18) θ ( w , w ∗ ) (cid:19) ( v (cid:62) v ∗ )( √ 2 π )(cid:42)(cid:107) w (cid:107) (cid:16) I n − ww (cid:62) (cid:107) w (cid:107) 2 (cid:17) w ∗ (cid:13)(cid:13)(cid:13)(cid:16) I n − ww (cid:62) (cid:107) w (cid:107) 2 (cid:17) w ∗ (cid:13)(cid:13)(cid:13) , w ∗ (cid:13)(cid:13)(cid:13) w (cid:107) w (cid:107) + w ∗ (cid:13)(cid:13)(cid:13) (cid:43) =
cos (cid:18) θ ( w , w ∗ ) (cid:19) ( v (cid:62) v ∗ )( √ 2 π )(cid:107) w (cid:107) 2 − ( w (cid:62) w ∗ )(cid:107)(cid:107) w (cid:107) 2 w ∗ − w ( w (cid:62) w ∗ ) (cid:107) (cid:107) w + (cid:107) w (cid:107) w ∗ (cid:107) = cos (cid:18) θ ( w , w ∗ ) (cid:19) ( v (cid:62) v
[Equation] [Equation]
) =
[Equation]
= cos
[Equation]
=
[Equation]
= sin
[Equation]
To show the second claim, without loss of generality, we assume (cid:107) w (cid:107)
= 1 .
Denote θ := θ ( w , w ∗ ) .
By Lemma 1, we have ∂f ∂
[Equation]
By Lemma
[Equation]
where
[Equation]
23) and by the ﬁrst claim,
[Equation]
Hence, for some A relu depending only on C v and c w , we have
[Equation]
[Table]
, where the equality is due to (22) and (23), the ﬁrst inequality is due to Cauchy-Schwarz inequality, the second inequality holds because the angle between w (cid:107) w (cid:107) and w (cid:107) w (cid:107) + w ∗ (cid:107) w (cid:107) w (cid:107) + w ∗ (cid:107) is θ 2 and (cid:13)(cid:13)(cid:13)(cid:13) w (cid:107) w (cid:107) − w (cid:107) w (cid:107) + w ∗
[Equation]
(cid:13)(cid:13)(cid:13)(cid:13) ≤ θ 2 , whereas the third inequality is due to sin( x ) ≥ 2 xπ , cos( x ) ≥ 1 − 2 xπ , and
[Equation]
for all x ∈ [0 , π 2 ] .
Lemma 6.
When Algorithm 1 converges, E Z (cid:104) ∂(cid:96)∂
v ( v , w ; Z ) (cid:105) and E Z (cid:104)
g relu ( v , w ; Z ) (cid:105) vanish simultaneously, which only occurs at the 1.
Saddle points where (8) is satisﬁed according to Proposition 1.
Minimizers of (2) where v = v ∗ , θ ( w , w ∗ ) = 0 , or v =
( I m + 1 m 1 (cid:62) m ) − 1 ( 1 m 1 (cid:62) m −
I m ) v ∗ , θ ( w , w ∗ ) = π .
Proof of Lemma 6.
By Lemma 4, suppose we have
E Z
[Equation]
24) and
[Equation] = 0 n , (25) where h ( v , v ∗ ) = (cid:107) v (cid:107) 2 +
[Equation] [Equation]
m v ) [-= (cid:107) v (cid:107) 
-]2 − [Equation]
[Equation] m v ) [Equation] [Equation]
m v ∗ ) + v (cid:62) v ∗ .
By (25), we must have θ ( w , w ∗ ) =
or θ ( w , w ∗ )
=
π or v (cid:62) v ∗ = 0 .
If θ ( w , w ∗ ) = 0 , then by (24), v = v ∗ , and (25) is satisﬁed.
If θ ( w , w ∗ ) =
π , then by (24), v = ( I m + 1 m 1 (cid:62) m ) − 1 ( 1 m 1 (cid:62) m − I m ) v ∗ , and (25) is satisﬁed.
If v (cid:62) v ∗ = 0 , then by (24), we have the expressions for v and θ ( w , w ∗ ) from Proposition 1, and (25) is satisﬁed.
Lemma 7.
If w (cid:54) = 0 n and θ ( w , w ∗ ) ∈ (0 , π ) , then [Equation] where h ( v , v ∗ ) : = {+(cid:107) v (cid:107) 
+}
+ ( 1 (cid:62)
[Equation]
m v )( 1 (cid:62) m v ∗ )
+ v (cid:62) v ∗ same as in Lemma 5, and p (
[Equation]
, w
[Equation]
:= 1 2 π
[Equation]
θ, w ) :
= 1 2 π (cid:90) π−
π
2 + θ sin( φ ) ξ (cid:18) sec( φ ) (cid:107) w (cid:107) (cid:19)
d φ with ξ ( x ) :
= (cid:82) x 0 r 2 exp( − r 2 2 )d r .
The inner product between the expected coarse and true gradients w.r.t.
w
[Equation]
Moreover, if further (cid:107) v (cid:107) ≤ C v and (cid:107) w (cid:107) ≥ c w , there exists a constant A crelu > 0 depending on C v and c w , such that
[Equation]
Proof of Lemma 7.
Denote θ := θ ( w , w ∗ ) .
We ﬁrst compute E Z (cid:104)
g crelu ( v , w ; Z ) (cid:105) .
By (5), g µ ( v , w ; Z ) =
Z (cid:62) (cid:0) µ (cid:48)
( Z w ) (cid:12) v (cid:1)(cid:16) v (cid:62) σ ( Z w ) − ( v ∗ ) (cid:62) σ ( Z w ∗ ) (cid:17) .
Since µ (cid:48)
= 1 { 0 <x< 1 } and σ = 1 { x> 0 } , we have
E
[Equation]
[Equation]
[Equation]
In the last equality above, we called Lemma 12.
Notice that (cid:16) I n − ww (cid:62) (cid:107) w (cid:107) 2 (cid:17) w = 0 n and (cid:107) w ∗ (cid:107) = 1 .
If θ ( w , w ∗ ) (cid:54) = 0 , π , then the inner product between E Z (cid:104)
g crelu ( v , w ; Z ) (cid:105) and ∂f∂ w ( v , w ) is given by (cid:28) E Z (cid:104)
g crelu ( v , w ; Z ) (cid:105) , ∂f
[Equation]
v ,
[Equation]
= csc (cid:18) θ (cid:19) q
( θ,
[Equation]
) 2 π (
[Equation]
[Equation]
(cid:13)(cid:13)(cid:13)(cid:16)
[Equation]
n
[Equation]
, w ∗ (cid:13)(cid:13)(cid:13)
[Equation]
(cid:13)(cid:13)(cid:13) (cid:43) =
q ( θ, w ) 2 π (cid:107) w (cid:107) ( v (cid:62) v ∗ ) 2 ≥ 0 .
In the last line, q ( θ, w ) ≥ 0 because sin( φ ) ξ (cid:16) sec( φ ) (cid:107) w (cid:107) (cid:17) is odd in φ and positive for φ ∈ (0 , π 2 ] .
Next, we bound (cid:13)(cid:13)(cid:13) E Z (cid:104)
g crelu ( v , w ; Z ) (cid:105)(cid:13)(cid:13)(cid:13) 2 .
Since (23) gives
[Equation]
where according to Lemma 1,
[Equation]
We rewrite the coarse partial gradient w.r.t.
w in (26) as
[Equation]
To prove the last claim
[Equation]
we notice that in the above equality,
[Equation]
Now, what is left is to bound (cid:0)(cid:0) 1 − θπ (cid:1) p
[Equation]
0 , w ) − p ( θ, w ) (cid:1) 2 , using a multiple of q ( θ, w ) .
Recall that
[Equation]
We ﬁrst show that both (cid:0)(cid:0) 1 − θπ (cid:1) p (0 , w ) − p ( θ, w ) (cid:1) 2 and q ( θ, w ) are symmetric with respect to θ
= π 2 on [0 , π ] .
This is because
[Equation]
and
[Equation]
Therefore, it sufﬁces to consider θ ∈ [ π 2 , π ] only.
Then calling Lemma 13 for θ ∈ [ π 2 , π ] , we have p ( θ, w ) ≤ q ( θ, w ) and (cid:0) 1 − θπ (cid:1) p (0 , w ) ≤ q ( θ, w ) .
Therefore,
[Equation]
Combining the above estimate together with (27), (28) and (29), and using Cauchy-Schwarz inequality, we have
[Equation]
[Equation]
q
[Equation]
, where p (0 , w ) and q ( θ, w ) are uniformly bounded.
This completes the proof.
Lemma 8.
When Algorithm 1 converges, E Z (cid:104) ∂(cid:96)∂
v ( v , w ; Z ) (cid:105) and E Z (cid:104)
g crelu ( v , w ; Z ) (cid:105) vanish simultaneously, which only occurs at the
Saddle points where
is
satisﬁed according to
Proposition 1.
Minimizers of (2)
where v = v ∗ , θ
( w
, w
∗ ) =
, or v =
( I
m
+ m
(cid:62)
m )
−
( 1 m 1 (cid:62) m − I m ) v ∗ , θ ( w , w ∗ )
= π .
Proof of Lemma 8.
The proof of Lemma 8 is similar to that of Lemma 6, and we omit it here.
The core part is that q ( θ, w ) deﬁned in Lemma 12 is non-negative and equals 0 only at θ = 0 , π , as well as p (0 , w ) ≥ p ( θ, w ) ≥ p ( π, w ) = 0 .
Lemma 9.
Let µ ( x ) =
x
in (5).
Then the expected coarse partial gradient w.r.t.
w is
[Equation]
w (cid:107) − ( v (cid:62) v ∗ ) w ∗ (cid:19) . If θ ( w , w ∗ ) = π and v = (cid:0) I m + 1 m 1 (cid:62)
m (cid:1) − 1 ( 1 m 1 (cid:62) m − I m ) v ∗ , (cid:13)(cid:13)(cid:13) E Z (cid:104)
g id ( v , w ; Z ) (cid:105)(cid:13)(cid:13)(cid:13) =
[Equation]
, i.e., E Z (cid:104) g id ( v , w ; Z ) (cid:105) does not vanish at the local minimizers if 1 (cid:62) m v ∗ (cid:54)
=
and m > 1 ,.
Proof of Lemma 9.
By (5),
[Equation]
Using the facts that µ (cid:48)
= 1 and σ = 1 { x> 0 } , we have
[Equation]
In the last equality above
[Equation]
we called the third identity in Lemma 11.
If θ ( w , w ∗ ) =
π and v = (cid:0)
I m
[Equation]
[Equation]
m − I m ) v ∗ , then
[Equation]
In the third equality, we used the identity ( I m + 1 m 1 (cid:62) m ) 1 m = ( m + 1) 1 m twice.
Lemma 10.
If w (cid:54) = 0 n and θ ( w , w ∗ ) ∈ (0 , π ) , then the inner product between the expected coarse and true gradients w.r.t. w is (cid:28)
[Equation]
When θ ( w , w ∗ ) → π , v → (cid:0) I
[Equation] + 1 m 1 (cid:62) m (cid:1) − [Equation]
m 1 (cid:62) m − I m ) v ∗ , if 1 (cid:62) m v ∗ (cid:54) =
and m > 1 , we have (cid:13)(cid:13)(cid:13) E Z (cid:104) g id ( v , w ; Z ) (cid:105)(cid:13)(cid:13)(cid:13)(cid:13)(cid:13)(cid:13) ∂f∂ v ( v , w ) (cid:13)(cid:13)(cid:13) 2 + (cid:68) E Z (cid:104) g id ( v , w ; Z ) (cid:105) , ∂f∂ w ( v , w ) (cid:69) →
+ ∞ .
Proof of Lemma 10.
By Lemmas 2 and 4, we have
[Equation]
(cid:62) [Equation] [Equation] (cid:107)
[Equation]
(cid:107) (cid:16) I n
[Equation]
ww (cid:62) (cid:107) w (cid:107) 2 (cid:17) w ∗ (cid:13)(cid:13)(cid:13)(cid:16) I n − ww (cid:62) (cid:107) w (cid:107) 2 (cid:17) w ∗ (cid:13)(cid:13)(cid:13) and
E Z (cid:104) g id ( v , w ; Z ) (cid:105)
= 1 √ 2 π (cid:18) (cid:107) v (cid:107) 2 w (cid:107) w (cid:107) − ( v (cid:62) v ∗ ) w ∗ (cid:19) . Since (cid:16) I n − ww (cid:62) (cid:107) w (cid:107) 2 (cid:17) w = 0 n and (cid:107) w ∗ (cid:107) = 1 , if θ ( w , w ∗ ) (cid:54) = 0 , π , then we have (cid:28) E Z (cid:104) g
id ( v , w
[Equation]
When θ ( w , w ∗ ) → π , v → (cid:0)
I
[Equation]
+ 1 m 1 (cid:62) m (cid:1)
[Equation]
1 ( 1 m 1 (cid:62) m − I m ) v ∗ , both (cid:13)(cid:13)(cid:13) ∂f∂ v ( v , w ) (cid:13)(cid:13)(cid:13) and (cid:68) E Z (cid:104) g
i
d ( v , w ; Z ) (cid:105) , ∂f∂ w ( v , w ) (cid:69) converge to 0. {+But if 1 (cid:62) m v ∗ (cid:54) = 
 and m > 1 , (cid:13)(cid:13)(cid:13) E Z (cid:104) g [Equation]+}
[-But if 1 (cid:62) m v ∗ (cid:54) = 
-][Equation]
(cid:105)(cid:13)(cid:13)(cid:13) → [Equation] 0 , which completes the proof.
Theorem 1.
Let { ( v t , w t ) } be the sequence generated by Algorithm 1 with ReLU µ ( x ) = max { x, 0 } .
Suppose (cid:107) v t (cid:107) ≤ C v
and (cid:107) w t (cid:107) ≥ c w for all t with some C v , c w > 0 .
Then if the learning rate η > 0 is sufﬁciently small, for any initialization ( v 0 , w 0 ) , the objective sequence { f ( v t , w t ) } is monotonically decreasing, and { ( v t , w t ) } converges to a saddle point or a (local) minimizer of the population loss minimization (2).
In addition, if 1 (cid:62) m v ∗ (cid:54) =
and m > 1 , the descent and convergence properties do not hold for Algorithm 1 with the identity function µ ( x ) = x near the local minimizers satisfying θ ( w , w ∗ )
=
π and v =
( I m + 1 m 1 (cid:62) m ) − 1 ( 1 m 1 (cid:62) m − I m ) v ∗ .
Proof of Theorem 1 .
If θ ( w t , w ∗ ) = 0 or π , then θ ( w T , w ∗ ) = 0 or π for all T ≥ t , and the original problem reduces to a quadratic program in terms of v .
So { v t } will converge to v ∗ or ( I m + 1 m 1 (cid:62) m ) − 1 ( 1 m 1 (cid:62) m − I m ) v ∗ by choosing a suitable step size η .
In either case, we have (cid:13)(cid:13)(cid:13) E Z [Equation]
v ( v t , w t ; Z ) (cid:105)(cid:13)(cid:13)(cid:13) and (cid:13)(cid:13)(cid:13) E Z (cid:104)
g relu ( v t , w t ; Z ) (cid:105)(cid:13)(cid:13)(cid:13) both converge to 0.
Else if θ ( w t , w ∗ ) ∈ (0 , π ) , we deﬁne for any a ∈ [0 , 1] that
[Equation]
and
[Equation]
which satisfy
[Equation]
Since by the assumption, (cid:107) v t (cid:107) ≤ C v and (cid:107) w t (cid:107) ≥ c w for all t , for sufﬁciently small η > 0 , it holds that (cid:107) v t ( a ) (cid:107) ≤ C v and (cid:107) w t ( a ) (cid:107) ≥ c w for all a ∈ [0 , 1] .
Possibly at some point awhere θ ( w t ( a 0 ) , w ∗ ) = 0 or π , the partial gradient ∂f∂ w ( v t ( a 0 ) , w t ( a 0 )) does not exist.
Otherwise, (cid:13)(cid:13)(cid:13) ∂f∂ w ( v t ( a ) , w t ( a )) (cid:13)(cid:13)(cid:13) is uniformly bounded for all a ∈ [0 , 1] / { a 0 } , which makes it integrable over the interval [0 , 1] .
Then for some constants L and A depending on C v and c w , we have
[Equation]
[Table]
[Equation]
The third equality is due to the fundamental theorem of calculus.
In the ﬁrst inequality, we called
Lemma 3 for ( v t , w t ) and ( v t ( a ) , w t ( a )) with a ∈ [0 , 1] / { a 0 } .
In the last inequality, we used
Lemma 5.
So when η < 2 (1+ A ) L , we have f ( v t +1 , w t +1 )
≤ f ( v t , w t ) .
Summing up the inequality (30) over t from 0 to ∞ and using f ≥ 0 , we have η ∞ (cid:88)
[Equation]
≤ [Equation] 0 , w 0 ) < ∞ .
Hence, lim t →∞
[Equation]
and
[Equation]
Invoking Lemma 5 again, we further have
[Equation]
Invoking Lemma 6, we have that coarse gradient descent with ReLU µ ( x ) only converges to a saddle point or a minimizer.
Using Lemmas 7, 8 and similar arguments, we can prove the convergence of coarse gradient descent with clipped ReLU STE.
The second claim follows from Lemmas 9 and 10.
F. C ONVERGENCE TO G LOBAL M INIMIZERS
We prove that if the initialization weights ( v 0 , w 0 ) satisfy ( v 0 ) (cid:62) v ∗ > 0 , θ ( w 0 , w ∗ ) <
π 2 and ( 1 (cid:62) m v ∗ )( 1 (cid:62) m v 0 ) ≤ ( 1 (cid:62) m v ∗ ) 2 , then we have convergence guarantee to global optima by using the vanilla or clipped ReLU STE.
Theorem 2.
Under the assumptions of Theorem 1, if further the initialization ( v 0 , w 0 ) satisﬁes ( v 0 ) (cid:62) v ∗ > 0 , θ ( w 0 , w ∗ ) <
π 2 and ( 1 (cid:62) m v ∗ )( 1 (cid:62) m v 0 ) ≤ ( 1 (cid:62) m v ∗ ) 2 , then by using the vanila or clipped
ReLU STE for sufﬁciently learning rate η > 0 , we have ( v t ) (cid:62) v ∗ > 0 and θ ( w t , w ∗ ) <
π 2 for all t > 0 , and { ( v t , w t ) } converges to a global minimizer.
Proof of Theorem 2.
Proof by induction.
Suppose ( v t ) (cid:62) v ∗ > 0 , θ ( w t , w ∗ ) <
π 2 and (
[Equation]
m v ∗ )( 1 (cid:62) m v t ) ≤ ( 1 (cid:62) m v ∗ ) 2 .
Then for small enough η , we have
[Equation]
and
[Equation]
Moreover, by Lemmas 4, 5 and 7, both E Z [ g relu ( v , w ; Z )] and E Z [ g crelu ( v , w ; Z )] can be written in the form of α 1 (cid:16) I m − ww
[Equation]
w
[Equation]
[Equation]
w ∗ + α 2 w , where α 1 ≤ 0 and α 2 is bounded by a constant depending on C v and c w .
Therefore,
[Equation]
and thus θ ( w t +1 , w ∗ ) <
π 2 .
Finally, since { ( v t , w t ) } converges, it can only converge to a global minimizer.