UNDERSTANDING STRAIGHT-THROUGH ESTIMATOR IN TRAINING ACTIVATION QUANTIZED NEURAL NETS
Year: 2019

Abstract:
Training activation quantized neural networks involves minimizing a piecewise constant function 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 only, so that the “gradient” through the modified chain rule becomes non-trivial. Since this unusual “gradient” is certainly not the gradient of loss function, the following question arises: why searching in its negative direction minimizes the training loss? In this paper, we provide the theoretical justification 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. The choice of STE is not unique. We prove that if the STE is properly chosen, the expected coarse gradient correlates positively with the population 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 critical point of the population loss minimization problem. Moreover, we show that a poor choice of STE leads to instability of the training algorithm near certain local minima, which is verified with CIFAR-10 experiments.

heading: 1 INTRODUCTION
text: 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 floating-point parameters, and billions of floating-point operations to make a single inference. To achieve substantial memory savings and energy efficiency at inference time, many recent efforts have been made to the training of coarsely quantized DNN, meanwhile maintaining the performance of their float counterparts (Courbariaux et al., 2015; Rastegari et al., 2016; Cai et al., 2017; Hubara et al., 2018; Yin et al., 2018b).
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; Li et al., 2017; Yin et al., 2016; 2018a; Hou & Kwok, 2018; He et al., 2018; Li & Hao, 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. Specifically, 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 (Lee et al., 2015) via convex combinatorial optimization.

heading: 1.1 RELATED WORKS
text: The idea of STE originates to the celebrated perceptron algorithm (Rosenblatt, 1957; 1962) in 1950s for learning single-layer perceptrons. The perceptron algorithm essentially does not calculate the “gradient” through the standard chain rule, but instead through a modified chain rule in which the derivative of identity function serves as the proxy of the original derivative of binary output function 1{x>0}. Its convergence has been extensive discussed in the literature; see for examples, (Widrow & Lehr, 1990; Freund & Schapire, 1999) and the references therein. Hinton (2012) extended this idea to train multi-layer networks with binary activations (a.k.a. binary neuron), namely, to backpropagate as if the activation had been the identity function. 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 employed 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; Yin et al., 2018b), where some other proxies took place including the derivatives of vanilla ReLU and clipped ReLU. Despite all the empirical success of STE, there is very limited theoretical understanding of it in training DNN with stair-case activations.
Goel et al. (2018) considers leaky ReLU activation of a one-hidden-layer network. They showed the convergence of the so-called Convertron algorithm, which uses the identity STE in the backward pass through the leaky ReLU layer. Other similar scenarios, where certain layers are not desirable for back-propagation, have 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 classifier to improve the generalization accuracy of DNN. In the backward pass, the derivative of a pre-trained fullyconnected layer was used as a surrogate. To circumvent adversarial defense (Szegedy et al., 2013), (Athalye et al., 2018) introduced the backward pass differentiable approximation, which shares the same spirit as STE, and successfully broke defenses at ICLR 2018 that rely on obfuscated gradients.

heading: 1.2 MAIN CONTRIBUTIONS
text: Throughout this paper, we shall refer to the “gradient” of loss function w.r.t. the weight variables through the STE-modified chain rule as coarse gradient. Since the backward and forward passes do not match, the coarse gradient is certainly not the gradient of loss function, and it is generally not the gradient of any function. Why searching in its negative direction minimizes the training loss, as this is not the standard gradient descent algorithm? Apparently, the choice of STE is non-unique, then what makes a good STE? From the optimization perspective, we take a step towards understanding STE in training quantized ReLU nets by attempting these questions.
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). For the first time, we prove that proper choices of STE give rise to training algorithms that are descent. Specifically, the negative expected coarse gradients based on STEs of the vanilla and clipped ReLUs are provably descent directions for the minimizing the population loss, which yield monotonically decreasing energy in the training. In contrast, this is not true for the identity STE. We further prove that the corresponding training algorithm can be unstable near certain local minima, because the coarse gradient may simply not vanish there.
Complementary to the analysis, we examine the empirical performances of the three STEs on MNIST and CIFAR-10 classifications with general quantized ReLU. While both vanilla and clipped ReLUs work very well on the relatively shallow LeNet-5, clipped ReLU STE is arguably the best for the deeper VGG-11 and ResNet-20. In our CIFAR experiments in section 4.2, we observe that the training using identity or ReLU STE can be unstable at good minima and repelled to an inferior one with substantially higher training loss and decreased generalization accuracy. This is an implication that poor STEs generate coarse gradients incompatible with the energy landscape, which is consistent with our theoretical finding about the identity STE.
To our knowledge, convergence guarantees of perceptron algorithm (Rosenblatt, 1957; 1962) and Convertron algorithm (Goel et al., 2018) were proved for the identity STE. It is worth noting that Convertron (Goel et al., 2018) makes weaker assumptions than in this paper. These results, however, do not generalize to the network with two trainable layers studied here. As aforementioned, the identity STE is actually a poor choice in our case. Moreover, it is not clear if their analyses can be extended to other STEs. Similar to Convertron with leaky ReLU, the monotonicity of quantized activation function plays a role in coarse gradient descent. Indeed, all three STEs considered here exploit this property. But this is not the whole story. A great STE like the clipped ReLU matches quantized ReLU at the extrema, otherwise the instability/incompatibility issue may arise.
Organization. In section 2, we study the energy landscape of a two-linear-layer network with binary activation and Gaussian data. We present the main results and sketch the mathematical analysis for STE in section 3. In section 4, we compare the empirical performances of different STEs in 2- bit and 4-bit activation quantization, and report the instability phenomena of the training algorithms associated with poor STEs observed in CIFAR experiments. Due to space limitation, all the technical proofs as well as some figures are deferred to the appendix.
Notations. ‖ · ‖ denotes the Euclidean norm of a vector or the spectral norm of a matrix. 0n ∈ Rn represents the vector of all zeros, whereas 1n ∈ Rn the vector of all ones. In is the identity matrix of order n. For any w, z ∈ Rn, w>z = 〈w, z〉 = ∑ i wizi is their inner product. w z denotes the Hadamard product whose ith entry is given by (w z)i = wizi.

heading: 2 LEARNING TWO-LINEAR-LAYER CNN WITH BINARY ACTIVATION
text: We consider a model similar to (Du et al., 2018) that outputs the prediction
y(Z,v,w) := m∑ i=1 viσ(Z > i w) = v >σ(Zw)
for some input Z ∈ Rm×n. Here w ∈ Rn and v ∈ Rm are the trainable weights in the first and second linear layer, respectively; Z>i denotes the ith row vector of Z; the activation function σ acts component-wise on the vector Zw, i.e., σ(Zw)i = σ((Zw)i) = σ(Z>i w). The first layer serves as a convolutional layer, where each row Z>i can be viewed as a patch sampled from Z and the weight filter w is shared among all patches, and the second linear layer is the classifier. The label is generated according to y∗(Z) = (v∗)>σ(Zw∗) for some true (non-zero) parameters v∗ and w∗. Moreover, we use the following squared sample loss
`(v,w;Z) := 1 2 (y(Z,v,w)− y∗(Z))2 = 1 2
( v>σ(Zw)− y∗(Z) )2 . (1)
Unlike in (Du et al., 2018), the activation function σ here is not ReLU, but the binary function σ(x) = 1{x>0}.
We assume that the entries of Z ∈ Rm×n are i.i.d. sampled from the Gaussian distribution N (0, 1) (Zhong et al., 2017; Brutzkus & Globerson, 2017). Since `(v,w;Z) = `(v,w/c;Z) for any scalar c > 0, without loss of generality, we take ‖w∗‖ = 1 and cast the learning task as the following population loss minimization problem:
min v∈Rm,w∈Rn
f(v,w) := EZ [`(v,w;Z)] , (2)
where the sample loss `(v,w;Z) is given by (1).

heading: 2.1 BACK-PROPAGATION AND COARSE GRADIENT DESCENT
text: With the Gaussian assumption on Z, as will be shown in section 2.2, it is possible to find the analytic expressions of f(v,w) and its gradient
∇f(v,w) := [ ∂f ∂v (v,w) ∂f ∂w (v,w) ] .
The gradient of objective function, however, is not available for the network training. In fact, we can only access the expected sample gradient, namely,
EZ [ ∂`
∂v (v,w;Z)
] and EZ [ ∂`
∂w (v,w;Z)
] .
We remark that EZ [ ∂` ∂w (v,w;Z) ] is not the same as ∂f∂w (v,w) = ∂EZ[`(v,w;Z)] ∂w . By the standard back-propagation or chain rule, we readily check that
∂` ∂v (v,w;Z) = σ(Zw)
( v>σ(Zw)− y∗(Z) ) (3)
and ∂`
∂w (v,w;Z) = Z>
( σ′(Zw) v )( v>σ(Zw)− y∗(Z) ) . (4)
Note that σ′ is zero a.e., which makes (4) inapplicable to the training. The idea of STE is to simply replace the a.e. zero component σ′ in (4) with a related non-trivial function µ′ (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 µ′ gives the following non-trivial surrogate of ∂`∂w (v,w;Z), to which we refer as the coarse (partial) gradient
gµ(v,w;Z) = Z >(µ′(Zw) v)(v>σ(Zw)− y∗(Z)). (5)
Using the STE µ′ to train the two-linear-layer convolutional neural network (CNN) with binary activation gives rise to the (full-batch) coarse gradient descent described in Algorithm 1.
Algorithm 1 Coarse gradient descent for learning two-linear-layer CNN with STE µ′. Input: initialization v0 ∈ Rm, w0 ∈ Rn, learning rate η.
for t = 0, 1, . . . do vt+1 = vt − η EZ [ ∂` ∂v (v t,wt;Z) ]
wt+1 = wt − η EZ [gµ(vt,wt;Z)] end for

heading: 2.2 PRELIMINARIES
text: Let us present some preliminaries about the landscape of the population loss function f(v,w). To this end, we define the angle between w and w∗ as θ(w,w∗) := arccos ( w>w∗
‖w‖‖w∗‖
) for any
w 6= 0n. Recall that the label is given by y∗(Z) = (v∗)>Zw∗ from (1), we elaborate on the analytic expressions of f(v,w) and ∇f(v,w). Lemma 1. If w 6= 0n, the population loss f(v,w) is given by
1
8
[ v> ( Im + 1m1 > m ) v − 2v> (( 1− 2
π θ(w,w∗)
) Im + 1m1 > m ) v∗ + (v∗)> ( Im + 1m1 > m ) v∗ ] .
In addition, f(v,w) = 18 (v ∗)> ( Im + 1m1 > m ) v∗ for w = 0n. Lemma 2. If w 6= 0n and θ(w,w∗) ∈ (0, π), the partial gradients of f(v,w) w.r.t. v and w are
∂f ∂v (v,w) = 1 4
( Im + 1m1 > m ) v − 1
4
(( 1− 2
π θ(w,w∗)
) Im + 1m1 > m ) v∗ (6)
and
∂f ∂w (v,w) = − v
>v∗
2π‖w‖
( In − ww >
‖w‖2 ) w∗∥∥∥(In − ww>‖w‖2 )w∗∥∥∥ , (7)
respectively.
For any v ∈ Rm, (v,0m) is impossible to be a local minimizer. The only possible (local) minimizers of the model (2) are located at
1. Stationary points where the gradients given by (6) and (7) vanish simultaneously (which may not be possible), i.e.,
v>v∗ = 0 and v = ( Im + 1m1 > m )−1(( 1− 2
π θ(w,w∗)
) Im + 1m1 > m ) v∗. (8)
2. Non-differentiable points where θ(w,w∗) = 0 and v = v∗, or θ(w,w∗) = π and v =( Im + 1m1 > m )−1 (1m1 > m − Im)v∗.
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 = ( Im + 1m1 > m )−1 (1m1 > m − Im)v∗} are the only potential spurious local minimizers. Proposition 1. If the true parameter v∗ satisfies (1>mv∗)2 < m+12 ‖v ∗‖2, then{
(v,w) : v = (Im + 1m1 > m) −1 (
−(1>mv∗)2
(m+ 1)‖v∗‖2 − (1>mv∗)2 Im + 1m1
> m
) v∗,
θ(w,w∗) = π
2
(m+ 1)‖v∗‖2
(m+ 1)‖v∗‖2 − (1>mv∗)2
} (9)
give the saddle points obeying (8), and {(v,w) : θ(w,w∗) = π, v = ( Im + 1m1 > m )−1 (1m1 > m − Im)v ∗} are the spurious local minimizers. Otherwise, the model (2) has no saddle points or spurious local minimizers.
We further prove that the population gradient∇f(v,w) given by (6) and (7), is Lipschitz continuous when restricted to bounded domains. Lemma 3. For any differentiable points (v,w) and (ṽ, w̃) with min{‖w‖, ‖w̃‖} = cw > 0 and max{‖v‖, ‖ṽ‖} = Cv , there exists a Lipschitz constant L > 0 depending on Cv and cw, such that
‖∇f(v,w)−∇f(ṽ, w̃)‖ ≤ L‖(v,w)− (ṽ, w̃)‖.

heading: 3 MAIN RESULTS
text: 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 and clipped ReLUs as well as the identity function serve as the STE, respectively. We shall prove that Algorithm 1 using the derivative of vanilla or clipped ReLU converges to a critical point, whereas that with the identity STE does not. Theorem 1 (Convergence). Let {(vt,wt)} be the sequence generated by Algorithm 1 with ReLU µ(x) = max{x, 0} or clipped ReLU µ(x) = min {max{x, 0}, 1}. Suppose ‖wt‖ ≥ cw for all t with some cw > 0. Then if the learning rate η > 0 is sufficiently small, for any initialization (v0,w0), the objective sequence {f(vt,wt)} is monotonically decreasing, and {(vt,wt)} converges to a saddle point or a (local) minimizer of the population loss minimization (2). In addition, if 1>mv
∗ 6= 0 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 = (Im + 1m1 > m) −1(1m1 > m − Im)v∗. Remark 1. The convergence guarantee for the coarse gradient descent is established under the assumption that there are infinite training samples. When there are only a few data, in a coarse scale, the empirical loss roughly descends 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 (proper) STE works so well with massive amounts of data as in deep learning.
Remark 2. The same results hold, if the Gaussian assumption on the input data is weakened to that their rows i.i.d. follow some rotation-invariant distribution. The proof will be substantially similar.
In the rest of this section, we sketch the mathematical analysis for the main results.

heading: 3.1 DERIVATIVE OF THE VANILLA RELU AS STE
text: If we choose the derivative of ReLU µ(x) = max{x, 0} as the STE in (5), it is easy to see µ′(x) = σ(x), and we have the following expressions of EZ [ ∂` ∂v (v,w;Z) ] and EZ [ grelu(v,w;Z) ] for Algorithm 1. Lemma 4. The expected partial gradient of `(v,w;Z) w.r.t. v is
EZ [ ∂`
∂v (v,w;Z)
] = ∂f
∂v (v,w). (10)
Let µ(x) = max{x, 0} in (5). The expected coarse gradient w.r.t. w is
EZ [ grelu(v,w;Z) ] = h(v,v∗)
2 √ 2π
w
‖w‖ − cos
( θ(w,w∗)
2
) v>v∗√
2π
w ‖w‖ + w ∗∥∥∥ w‖w‖ + w∗∥∥∥ , 1 (11) where h(v,v∗) = ‖v‖2 + (1>mv)2 − (1>mv)(1>mv∗) + v>v∗. As stated in Lemma 5 below, the key observation is that the coarse partial gradient EZ [ grelu(v,w;Z) ] has non-negative correlation with the population partial gradient ∂f∂w (v,w),
and −EZ [ grelu(v,w;Z) ] together with −EZ [ ∂` ∂v (v,w;Z) ] form a descent direction for minimizing the population loss. Lemma 5. If w 6= 0n and θ(w,w∗) ∈ (0, π), then the inner product between the expected coarse and population gradients w.r.t. w is〈
EZ [ grelu(v,w;Z) ] , ∂f
∂w (v,w)
〉 = sin (θ(w,w∗))
2( √ 2π)3‖w‖ (v>v∗)2 ≥ 0.
Moreover, if further ‖v‖ ≤ Cv and ‖w‖ ≥ cw, there exists a constant Arelu > 0 depending on Cv and cw, such that∥∥∥EZ[grelu(v,w;Z)]∥∥∥2 ≤ Arelu(∥∥∥∥∂f∂v (v,w) ∥∥∥∥2 +〈EZ[grelu(v,w;Z)], ∂f∂w (v,w) 〉) . (12) Clearly, when 〈 EZ [ grelu(v,w;Z) ] , ∂f∂w (v,w) 〉 > 0, EZ [ grelu(v,w;Z) ] is roughly in the same
direction as ∂f∂w (v,w). Moreover, since by Lemma 4, EZ [ ∂` ∂v (v,w;Z) ] = ∂f∂v (v,w), we expect
1We redefine the second term as 0n in the case θ(w,w∗) = π, or equivalently, w‖w‖ +w ∗ = 0n.
that the coarse gradient descent behaves like the gradient descent directly on f(v,w). Here we would like to highlight the significance of the estimate (12) in guaranteeing the descent property of Algorithm 1. By the Lipschitz continuity of∇f specified in Lemma 3, it holds that
f(vt+1,wt+1)− f(vt,wt) ≤ 〈 ∂f
∂v (vt,wt),vt+1 − vt
〉 + 〈 ∂f
∂w (vt,wt),wt+1 −wt 〉 + L
2 (‖vt+1 − vt‖2 + ‖wt+1 −wt‖2) = − ( η − Lη 2
2 )∥∥∥∥∂f∂v (vt,wt) ∥∥∥∥2 + Lη22 ∥∥∥EZ[grelu(vt,wt;Z)]∥∥∥2
− η 〈 ∂f
∂w (vt,wt),EZ
[ grelu(v t,wt;Z) ]〉
a) ≤ − ( η − (1 +Arelu) Lη2
2 )∥∥∥∥∂f∂v (vt,wt) ∥∥∥∥2
− ( η − AreluLη 2
2
)〈 ∂f
∂w (vt,wt),EZ
[ grelu(v t,wt;Z) ]〉 ,
(13)
where a) is due to (12). Therefore, if η is small enough, we have monotonically decreasing energy until convergence.
Lemma 6. When Algorithm 1 converges, EZ [ ∂` ∂v (v,w;Z) ] and EZ [ grelu(v,w;Z) ] vanish simultaneously, which only occurs at the
1. Saddle points where (8) is satisfied according to Proposition 1.
2. Minimizers of (2) where v = v∗, θ(w,w∗) = 0, or v = (Im+1m1>m) −1(1m1 > m−Im)v∗,
θ(w,w∗) = π.
Lemma 6 states that when Algorithm 1 using ReLU STE converges, it can only converge to a critical point of the population loss function.

heading: 3.2 DERIVATIVE OF THE CLIPPED RELU AS STE
text: For the STE using clipped ReLU, µ(x) = min {max{x, 0}, 1} and µ′(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 EZ [ gcrelu(v,w;Z) ] 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 6= 0n and θ(w,w∗) ∈ (0, π), then
EZ [gcrelu(v,w;Z)] = p(0,w)h(v,v∗)
2
w
‖w‖ − (v>v∗) csc(θ/2) · q(θ,w)
w ‖w‖ + w ∗∥∥∥ w‖w‖ + w∗∥∥∥ − (v>v∗) (p(θ,w)− cot(θ/2) · q(θ,w)) w
‖w‖ ,
where h(v,v∗) := ‖v‖2 + (1>mv)2 − (1>mv)(1>mv∗) + v>v∗ same as in Lemma 5, and
p(θ,w) := 1
2π
∫ π 2
−π2 +θ cos(φ)ξ
( sec(φ)
‖w‖
) dφ, q(θ,w) := 1
2π
∫ π 2
−π2 +θ sin(φ)ξ
( sec(φ)
‖w‖
) dφ
with ξ(x) := ∫ x 0 r2 exp(− r 2
2 )dr. The inner product between the expected coarse and true gradients w.r.t. w 〈 EZ [ gcrelu(v,w;Z) ] , ∂f
∂w (v,w)
〉 = q(θ,w)
2π‖w‖ (v>v∗)2 ≥ 0.
Moreover, if further ‖v‖ ≤ Cv and ‖w‖ ≥ cw, there exists a constant Acrelu > 0 depending on Cv and cw, such that∥∥∥EZ[gcrelu(v,w;Z)]∥∥∥2 ≤ Acrelu(∥∥∥∥∂f∂v (v,w) ∥∥∥∥2 +〈EZ[gcrelu(v,w;Z)], ∂f∂w (v,w) 〉) .
Lemma 8. When Algorithm 1 converges, EZ [ ∂` ∂v (v,w;Z) ] and EZ [ gcrelu(v,w;Z) ] vanish simultaneously, which only occurs at the
1. Saddle points where (8) is satisfied according to Proposition 1.
2. Minimizers of (2) where v = v∗, θ(w,w∗) = 0, or v = (Im+1m1>m) −1(1m1 > m−Im)v∗,
θ(w,w∗) = π.

heading: 3.3 DERIVATIVE OF THE IDENTITY FUNCTION AS STE
text: Now we consider the derivative of identity function. 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 local minima, 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
EZ [ gid(v,w;Z) ] =
1√ 2π
( ‖v‖2 w
‖w‖ − (v>v∗)w∗
) . (14)
If θ(w,w∗) = π and v = ( Im + 1m1 > m )−1 (1m1
> m − Im)v∗,∥∥∥EZ[gid(v,w;Z)]∥∥∥ = 2(m− 1)√ 2π(m+ 1)2 (1>mv ∗)2 ≥ 0,
i.e., EZ [ gid(v,w;Z) ] does not vanish at the local minimizers if 1>mv ∗ 6= 0 and m > 1.
Lemma 10. If w 6= 0n and θ(w,w∗) ∈ (0, π), then the inner product between the expected coarse and true gradients w.r.t. w is〈
EZ [ gid(v,w;Z) ] , ∂f
∂w (v,w)
〉 = sin (θ(w,w∗))
( √ 2π)3‖w‖ (v>v∗)2 ≥ 0. (15)
When θ(w,w∗)→ π, v → ( Im + 1m1 > m )−1 (1m1
> m − Im)v∗, if 1>mv∗ 6= 0 and m > 1, we have∥∥∥EZ[gid(v,w;Z)]∥∥∥2∥∥∥∂f∂v (v,w)∥∥∥2 + 〈EZ[gid(v,w;Z)], ∂f∂w (v,w)〉 → +∞. (16)
Lemma 9 suggests that if 1>mv ∗ 6= 0, the coarse gradient descent will never converge near the spurious minimizers with θ(w,w∗) = π and v = ( Im + 1m1 > m )−1 (1m1 > m − Im)v∗, because
EZ [ gid(v,w;Z) ] does not vanish there. By the positive correlation implied by (15) of Lemma
10, for some proper (v0,w0), the iterates {(vt,wt)} may move towards a local minimizer in the beginning. But when {(vt,wt)} approaches it, the descent property (13) does not hold for EZ [gid(v,w;Z)] because of (16), hence the training loss begins to increase and instability arises.

heading: 4 EXPERIMENTS
text: While our theory implies that both vanilla and clipped ReLUs learn a two-linear-layer CNN, their empirical performances on deeper nets are different. In this section, we compare the performances of the identity, ReLU and clipped ReLU STEs on MNIST (LeCun et al., 1998) and CIFAR-10 (Krizhevsky, 2009) benchmarks for 2-bit or 4-bit quantized activations. As an illustration, we plot
the 2-bit quantized ReLU and its associated clipped ReLU in Figure 3 in the appendix. Intuitively, the clipped ReLU should be the best performer, as it best approximates the original quantized ReLU. We also report the instability issue of the training algorithm when using an improper STE in section 4.2. In all experiments, the weights are kept float.
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 modified batch normalization layer (Ioffe & Szegedy, 2015) without the scale and shift, whose output components approximately follow a unit Gaussian distribution. Then the α that fits the input of activation layer the best can be pre-computed by a variant of Lloyd’s algorithm (Lloyd, 1982; Yin et al., 2018a) applied to a set of simulated 1-D half-Gaussian data. After determining the α, it will be fixed during the whole training process. Since the original LeNet-5 does not have batch normalization, we add one prior to each activation layer. 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.
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 specified in Table 2 in the appendix.

heading: 4.1 COMPARISON RESULTS
text: 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 relatively shallow LeNet-5 network, vanilla ReLU exhibits comparable performance to the clipped ReLU, which is somewhat in line with our theoretical finding that ReLU is a great STE for learning the two-linear-layer (shallow) CNN.

heading: 4.2 INSTABILITY
text: 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 predicted in Theorem 1. By Table 1, 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.16% (1.38). 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 significantly increase within the first 20 epochs; see Figure 4.2. To speedup training, at epoch 20, we switch to the normal schedule of learning rate specified 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-20
is also due to the instability of the corresponding training algorithm at good minima, as illustrated by Figure 4 in Appendix C, although it diverges much slower.

heading: 5 CONCLUDING REMARKS
text: We provided the first theoretical justification for the concept of STE that it gives rise to descent training algorithm. We considered three STEs: the derivatives of the identity function, vanilla ReLU and clipped ReLU, for learning a two-linear-layer CNN with binary activation. We derived the explicit formulas of the expected coarse gradients corresponding to the STEs, and showed that the negative expected coarse gradients based on vanilla and clipped ReLUs are descent directions for minimizing the population loss, whereas the identity STE is not since it generates a coarse gradient incompatible with the energy landscape. The instability/incompatibility issue was confirmed in CIFAR experiments for improper choices of STE. In the future work, we aim further understanding of coarse gradient descent for large-scale optimization problems with intractable gradients.

heading: ACKNOWLEDGMENTS
text: This work was partially supported by NSF grants DMS-1522383, IIS-1632935, ONR grant N0001418-1-2527, AFOSR grant FA9550-18-0167, DOE grant DE-SC0013839 and STROBE STC NSF grant DMR-1548924.

heading: A. THE PLOTS OF QUANTIZED AND CLIPPED RELUS

heading: B. THE SCHEDULE OF LEARNING RATE
text: C. INSTABILITY OF RELU STE ON RESNET-20 WITH 2-BIT ACTIVATIONS

heading: D. ADDITIONAL SUPPORTING LEMMAS
text: Lemma 11. Let z ∈ Rn be a Gaussian random vector with entries i.i.d. sampled from N (0, 1). Given nonzero vectors w, w̃ ∈ Rn with the angle θ, we have
E [ 1{z>w>0} ] = 1 2 , E [ 1{z>w>0, z>w̃>0} ] = π − θ 2π ,
and
E [ z1{z>w>0} ] =
1√ 2π w ‖w‖ , E [ z1{z>w>0, z>w̃>0} ] = cos(θ/2)√ 2π w ‖w‖ + w̃ ‖w̃‖∥∥∥ w‖w‖ + w̃‖w̃‖∥∥∥ , 2
Proof of Lemma 11. The third identity was proved in Lemma A.1 of (Du et al., 2018). To show the first one, without loss of generality we assume w = [w1,0>n−1] > with w1 > 0, then E [ 1{z>w>0} ] = P(z1 > 0) = 12 .
We further assume w̃ = [w̃1, w̃2,0>n−2] >. It is easy to see that
E [ 1{z>w>0, z>w̃>0} ] = P(z>w > 0, z>w̃ > 0) =
π − θ 2π .
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 dPr = r exp(−r2/2)dr and dPφ = 12πdφ. Then E [ zi1{z>w>0, z>w̃>0} ] = 0 for i ≥ 3. Moreover,
E [ z11{z>w>0, z>w̃>0} ] = 1
2π ∫ ∞ 0 r2 exp ( −r 2 2 ) dr ∫ π 2 −π2 +θ cos(φ)dφ = 1 + cos(θ) 2 √ 2π
and
E [ z21{z>w>0, z>w̃>0} ] = 1
2π ∫ ∞ 0 r2 exp ( −r 2 2 ) dr ∫ π 2 −π2 +θ sin(φ)dφ = sin(θ) 2 √ 2π .
Therefore,
E [ z1{z>w>0, z>w̃>0} ] =
cos(θ/2)√ 2π [cos(θ/2), sin(θ/2),0>n−2] > = cos(θ/2)√ 2π
w ‖w‖ + w̃ ‖w̃‖∥∥∥ w‖w‖ + w̃‖w̃‖∥∥∥ ,
where the last equality holds because w‖w‖ and w ‖w‖+ w̃ ‖w̃‖
‖ w‖w‖+ w̃‖w̃‖‖ are two unit-normed vectors with angle
θ/2.
Lemma 12. Let z ∈ Rn be a Gaussian random vector with entries i.i.d. sampled from N (0, 1). Given nonzero vectors w, w̃ ∈ Rn with the angle θ, we have
E [ z1{0<z>w<1} ] = p(0,w) w
‖w‖ ,
E [ z1{0<z>w<1, z>w̃>0} ] = ((p(θ,w))− cot(θ/2) · q(θ,w)) w
‖w‖
+ csc(θ/2) · q(θ,w) w ‖w‖ + w̃ ‖w̃‖∥∥∥ w‖w‖ + w̃‖w̃‖∥∥∥ .
Here
p(θ,w) = 1
2π
∫ π 2
−π2 +θ cos(φ)ξ
( sec(φ)
‖w‖
) dφ, q(θ,w) = 1
2π
∫ π 2
−π2 +θ sin(φ)ξ
( sec(φ)
‖w‖
) dφ
with ξ(x) = ∫ x 0 r2 exp(− r 2 2 )dr satisfying ξ(+∞) = √ π 2 . In addition,
p(π,w) = q(0,w) = q(π,w) = 0. 2Same as in Lemma 4, we redefine E [ z1{z>w>0, z>w̃>0} ] = 0n in the case θ(w, w̃) = π.
Proof of Lemma 12. Following the proof of Lemma 11, we let w = [w1,0>n−1]> with w1 > 0, w̃ = [w̃1, w̃2,0 > n−2]
>, and use the polar representation of two-dimensional Gaussian distribution. Then
E [ z11{0<z>w<1} ] = 1
2π
∫ π 2
−π2 cos(φ)
∫ sec(φ) ‖w‖
0
r2 exp ( −r 2
2
) drdφ = p(0,w)
and
E [ z21{0<z>w<1} ] = 1
2π
∫ π 2
−π2 sin(φ)
∫ sec(φ) ‖w‖
0
r2 exp ( −r 2
2
) drdφ = q(0,w) = 0,
since the integrand above is an odd function in φ. Moreover, E [ zi1{0<z>w<1} ] = 0 for i ≥ 3. So the first identity holds. For the second one, we have
E [ z11{0<z>w<1,z>w̃>0} ] = 1
2π
∫ π 2
−π2 +θ cos(φ)
∫ sec(φ) ‖w‖
0
r2 exp ( −r 2
2
) drdφ = p(θ,w),
and similarly, E [ z21{0<z>w<1,z>w̃>0} ] = q(θ,w). Therefore,
E[z1{0<z>w<1,z>w̃>0}] = [p(θ,w), q(θ,w),0>n−2]>.
Since w‖w‖ = [1,0 > n−1]
> and w ‖w‖+ w̃ ‖w̃‖
‖ w‖w‖+ w̃‖w̃‖‖ = [cos(θ/2), sin(θ/2),0>n−2] >, it is easy to check the
second identity is also true.
Lemma 13. Let p(θ, w) and q(θ,w) be defined in Lemma 12. Then for θ ∈ [π2 , π], we have
1. p(θ,w) ≤ q(θ,w). 2. ( 1− θπ ) p(0,w) ≤ q(θ,w).
Proof of Lemma 13. 1. Let θ ∈ [π2 , π], since ξ ≥ 0,
p(θ,w) = 1
2π
∫ π 2
θ−π2 cos(φ)ξ
( sec(φ)
‖w‖
) dφ = 1
2π
∫ π 2 θ−π2 sin (π 2 − φ ) ξ ( sec(φ) ‖w‖ ) dφ
≤ 1 2π
∫ π 2
θ−π2 sin(θ − φ)ξ
( sec(φ)
‖w‖
) dφ ≤ 1
2π
∫ π 2
θ−π2 sin(φ)ξ
( sec(φ)
‖w‖
) dφ = q(θ,w),
where the last inequality is due to the rearrangement inequality since both sin(φ) and ξ (
sec(φ) ‖w‖ ) are increasing in φ on [0, π2 ]. 2. Since cos(φ)ξ (
sec(φ) ‖w‖
) is even, we have
1
π
∫ π 2
−π2 cos(φ)ξ
( sec(φ)
‖w‖
) dφ = 2
π
∫ π 2
0
cos(φ)ξ
( sec(φ)
‖w‖
) dφ
≤ 2 π
∫ π 2
0
sin(φ)ξ
( sec(φ)
‖w‖
) dφ ≤ 1
π − θ
∫ π 2
θ−π2 sin(φ)ξ
( sec(φ)
‖w‖
) dφ.
The first inequality is due to part 1 which gives p(π/2,w) ≤ q(π/2,w), whereas the second one holds because sin(φ)ξ ( sec(φ) ‖w‖ ) is increasing on [0, π2 ].
Lemma 14. For any nonzero vectors w and w̃ with ‖w̃‖ ≥ ‖w‖ = c > 0, we have
1. |θ(w,w∗)− θ(w̃,w∗)| ≤ π2c‖w − w̃‖.
2. ∥∥∥∥∥∥ 1‖w‖ ( In−ww > ‖w‖2 ) w∗∥∥∥(In−ww>‖w‖2 )w∗∥∥∥ − 1‖w̃‖ ( In− w̃w̃ > ‖w̃‖2 ) w∗∥∥∥(In− w̃w̃>‖w̃‖2 )w∗∥∥∥ ∥∥∥∥∥∥ ≤ 1c2 ‖w − w̃‖. Proof of Lemma 14. 1. Since by Cauchy-Schwarz inequality,〈
w̃,w − cw̃ ‖w̃‖
〉 = w̃>w − c‖w̃‖ ≤ 0,
we have ‖w̃ −w‖2 = ∥∥∥∥(1− c‖w̃‖ ) w̃ − ( w − cw̃ ‖w̃‖ )∥∥∥∥2 ≥ ∥∥∥∥(1− c‖w̃‖ ) w̃ ∥∥∥∥2 + ∥∥∥∥w − cw̃‖w̃‖ ∥∥∥∥2
≥ ∥∥∥∥w − cw̃‖w̃‖ ∥∥∥∥2 = c2 ∥∥∥∥ w‖w‖ − w̃‖w̃‖ ∥∥∥∥2 . (17)
Therefore,
|θ(w,w∗)− θ(w̃,w∗)| ≤ θ(w, w̃) = θ ( w
‖w‖ ,
w̃
‖w̃‖
)
≤ π sin
θ ( w ‖w‖ , w̃ ‖w̃‖ ) 2  = π 2 ∥∥∥∥ w‖w‖ − w̃‖w̃‖ ∥∥∥∥ ≤ π2c‖w − w̃‖,
where we used the fact sin(x) ≥ 2xπ for x ∈ [0, π 2 ] and the estimate in (17). 2. Since ( In− ww >
‖w‖2 ) w∗ is the projection of w∗ onto the complement space of w, and likewise for(
In − w̃w̃ >
‖w̃‖2
) w∗, the angle between ( In − ww >
‖w‖2
) w∗ and ( In − w̃w̃ >
‖w̃‖2
) w∗ is equal to the angle
between w and w̃. Therefore,〈 ( In − ww >
‖w‖2 ) w∗∥∥∥(In − ww>‖w‖2 )w∗∥∥∥ , ( In − w̃w̃ > ‖w̃‖2 ) w∗∥∥∥(In − w̃w̃>‖w̃‖2 )w∗∥∥∥ 〉 = 〈 w ‖w‖ , w̃ ‖w̃‖ 〉 ,
and thus ∥∥∥∥∥∥ 1‖w‖ ( In − ww > ‖w‖2 ) w∗∥∥∥(In − ww>‖w‖2 )w∗∥∥∥ − 1 ‖w̃‖ ( In − w̃w̃ > ‖w̃‖2 ) w∗∥∥∥(In − w̃w̃>‖w̃‖2 )w∗∥∥∥ ∥∥∥∥∥∥ = ∥∥∥∥ w‖w‖2 − w̃‖w̃‖2 ∥∥∥∥ = ‖w − w̃‖‖w‖‖w̃‖ ≤ 1c2 ‖w − w̃‖.
The second equality above holds because∥∥∥∥ w‖w‖2 − w̃‖w̃‖2 ∥∥∥∥2 = 1‖w‖2 + 1‖w̃‖2 − 2〈w, w̃〉‖w‖2‖w̃‖2 = ‖w − w̃‖2‖w‖2‖w̃‖2 .

heading: E. MAIN PROOFS
text: Lemma 1. If w 6= 0n, the population loss f(v,w) is given by
1
8
[ v> ( Im + 1m1 > m ) v − 2v> (( 1− 2
π θ(w,w∗)
) Im + 1m1 > m ) v∗ + (v∗)> ( Im + 1m1 > m ) v∗ ] .
In addition, f(v,w) = 18 (v ∗)> ( Im + 1m1 > m ) v∗ for w = 0n.
Proof of Lemma 1. We notice that
f(v,w) = 1
2
( v>EZ[σ(Zw)σ(Zw)>]v − 2v>EZ[σ(Zw)σ(Zw∗)>]v∗
+ (v∗)>EZ[σ(Zw∗)σ(Zw∗)>]v∗ ) .
Let Z>i be the i th row vector of Z. Since w 6= 0n, using Lemma 11, we have
EZ [ σ(Zw)σ(Zw)> ] ii = E [ σ(Z>i w)σ(Z > i w) ] = E [ 1{Z>i w>0} ] = 1 2 ,
and for i 6= j,
EZ [ σ(Zw)σ(Zw)> ] ij = E [ σ(Z>i w)σ(Z > j w) ] = E [ 1{Z>i w>0} ] E [ 1{Z>j w>0} ] = 1 4 .
Therefore, EZ [ σ(Zw)σ(Zw)> ] = EZ [ σ(Zw∗)σ(Zw∗)> ] = 14 ( Im + 1m1 > m ) . Furthermore,
EZ [ σ(Zw)σ(Zw∗)> ] ii = E [ 1{Z>i w>0,Z>i w∗>0} ] = π − θ(w,w∗) 2π ,
and EZ [ σ(Zw)σ(Zw∗)> ] ij = 14 . So,
EZ [ σ(Zw)σ(Zw∗)> ] = 1
4
(( 1− 2θ(w,w ∗)
π
) Im + 1m1 > m ) .
Then it is easy to validate the first claim. Moreover, if w = 0n, then
f(v,w) = 1
2 (v∗)>EZ[σ(Zw∗)σ(Zw∗)>]v∗ =
1 8 (v∗)>
( Im + 1m1 > m ) v∗.
Lemma 2. If w 6= 0n and θ(w,w∗) ∈ (0, π), the partial gradients of f(v,w) w.r.t. v and w are
∂f ∂v (v,w) = 1 4
( Im + 1m1 > m ) v − 1
4
(( 1− 2
π θ(w,w∗)
) Im + 1m1 > m ) v∗
and
∂f ∂w (v,w) = − v
>v∗
2π‖w‖
( In − ww >
‖w‖2 ) w∗∥∥∥(In − ww>‖w‖2 )w∗∥∥∥ ,
respectively.
Proof of Lemma 2. The first claim is trivial, and we only show the second one. Since θ(w,w∗) = arccos ( w>w∗
‖w‖
) is differentiable w.r.t. w at θ(w,w∗) ∈ (0, π), we have
∂f ∂w (v,w) =
v>v∗
2π
∂θ ∂w (w,w∗) = −v
>v∗
2π
‖w‖2w∗ − (w>w∗)w ‖w‖3 √ 1− (w >w∗)2
‖w‖2
= − v >v∗
2π‖w‖
( In − ww >
‖w‖2 ) w∗∥∥∥(In − ww>‖w‖2 )w∗∥∥∥ .
Proposition 1. If the true parameter v∗ satisfies (1>mv∗)2 < m+12 ‖v ∗‖2, then{ (v,w) : v = (Im + 1m1 > m) −1 ( −(1>mv∗)2
(m+ 1)‖v∗‖2 − (1>mv∗)2 Im + 1m1
> m
) v∗,
θ(w,w∗) = π
2
(m+ 1)‖v∗‖2
(m+ 1)‖v∗‖2 − (1>mv∗)2 } give the saddle points obeying (8), and {(v,w) : θ(w,w∗) = π, v = ( Im + 1m1 > m )−1 (1m1 > m − Im)v ∗} are the spurious local minimizers. Otherwise, the model (2) has no saddle points or spurious local minimizers.
Proof of Proposition 1. Suppose v>v∗ = 0 and ∂f∂v (v,w) = 0, then by Lemma 1, 0 = v>v∗ = (v∗)>(Im + 1m1 > m) −1 ((
1− 2 π θ(w,w∗)
) Im + 1m1 > m ) v∗. (18)
From (18) it follows that
2 π θ(w,w∗)(v∗)>(Im+1m1 > m) −1v∗ = (v∗)>(Im+1m1
> m) −1 (Im + 1m1>m)v∗ = ‖v∗‖2. (19)
On the other hand, from (18) it also follows that( 2
π θ(w,w∗)− 1
) (v∗)>(Im + 1m1 > m) −1v∗ = (v∗)>(Im + 1m1 > m) −11m(1 > mv ∗) = (1>v∗)2
m+ 1 ,
where we used (Im + 1m1>m)1m = (m+ 1)1m. Taking the difference of the two equalities above gives
(v∗)>(Im + 1m1 > m) −1v∗ = ‖v∗‖2 − (1
> mv ∗)2
m+ 1 . (20)
By (19), we have θ(w,w∗) = π2 (m+1)‖v∗‖2 (m+1)‖v∗‖2−(1>mv∗)2 , which requires
π
2
(m+ 1)‖v∗‖2
(m+ 1)‖v∗‖2 − (1>mv∗)2 < π, or equivalently, (1>mv
∗)2 < m+ 1
2 ‖v∗‖2.
Furthermore, since ∂f∂v (v,w) = 0, we have v = (Im + 1m1 > m) −1 ((
1− 2 π θ(w,w∗)
) Im + 1m1 > m ) v∗
= (Im + 1m1 > m) −1 (
−(1>mv∗)2
(m+ 1)‖v∗‖2 − (1>mv∗)2 Im + 1m1
> m
) v∗.
Next, we check the local optimality of the stationary points. By ignoring the scaling and constant terms, we rewrite the objective function as
f̃(v, θ) := v> ( Im + 1m1 > m ) v − 2v> (( 1− 2
π θ
) Im + 1m1 > m ) v∗, for θ ∈ [0, π].
It is easy to check that its Hessian matrix ∇2f̃(v, θ) = [ 2(Im + 1m1 > m) 4 πv ∗
4 π (v ∗)> 0 ] is indefinite. Therefore, the stationary points are saddle points.
Moreover, if (1>mv ∗)2 < m+12 ‖v ∗‖2, at the point (v, θ) = ((Im + 1m1>m)−1(1m1>m − Im)v∗, π), we have
v>v∗ = (v∗)>(Im + 1m1 > m) −1(1m1 > m − Im)v∗
= ‖v∗‖2 − 2(v∗)>(Im + 1m1>m)−1v∗ = 2(1>mv ∗)2
m+ 1 − ‖v∗‖2 < 0, (21)
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
f̃(v + ∆v, π + ∆θ) = (v + ∆v)> ( Im + 1m1 > m ) (v + ∆v)− 2(v + ∆v)> ( 1m1 > m − Im ) v∗
+ 2∆θ
π (v + ∆v)>v∗.
On the right hand side, since v = (Im + 1m1>m) −1(1m1 > m − Im)v∗ is the unique minimizer to the quadratic function f̃(v, π), we have if ∆v 6= 0m,
(v + ∆v)> ( Im + 1m1 > m ) (v + ∆v)− 2(v + ∆v)> ( 1m1 > m − Im ) v∗ > f̃(v, π).
Moreover, for sufficiently small ‖∆v‖, it holds that ∆θ · (v + ∆v)>v∗ > 0 for ∆θ < 0 because of (21). Therefore, f̃(v + ∆v, π + ∆θ) > f̃(v, π) whenever (∆v,∆θ) is small and non-zero, and ((Im + 1m1 > m) −1(1m1 > m − Im)v∗, π) is a local minimizer of f̃ .
To prove the second claim, suppose (1>mv ∗)2 ≥ m+12 ‖v ∗‖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, θ) = ((Im + 1m1>m) −1(1m1 > m − Im)v∗, π), we have
v>v∗ = 2(1>mv ∗)2
m+ 1 − ‖v∗‖2 ≥ 0.
If v>v∗ > 0, since ∇f̃(v, θ) = 14 [0 > m, 2 πv >v∗]>, a small perturbation [0>m,∆θ] > with ∆θ < 0 will give a strictly decreased objective value, so (v, θ) = ( (Im + 1m1 > m) −1(1m1 > m − Im)v∗, π ) is not a local minimizer. If v>v∗ = 0, then∇f̃(v, θ) = 0n+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{‖w‖, ‖w̃‖} = cw > 0 and max{‖v‖, ‖ṽ‖} = Cv , there exists a Lipschitz constant L > 0 depending on Cv and cw, such that
‖∇f(v,w)−∇f(ṽ, w̃)‖ ≤ L‖(v,w)− (ṽ, w̃)‖.
Proof of Lemma 3. It is easy to check that ‖Im + 1m1>m‖ = m+ 1. Then∥∥∥∥∂f∂v (v,w)− ∂f∂v (ṽ, w̃) ∥∥∥∥ = 14 ∥∥∥∥(Im + 1m1>m)(v − ṽ) + 2π (θ(w,w∗)− θ(w̃,w∗))v∗ ∥∥∥∥
≤ m+ 1 4 ‖v − ṽ‖+ ‖v ∗‖ 2π |θ(w,w∗)− θ(w̃,w∗)| ≤ m+ 1 4 ‖v − ṽ‖+ ‖v ∗‖ 4cw ‖w − w̃‖
≤ 1 4
( m+ 1 +
‖v∗‖ cw
) ‖(v,w)− (ṽ, w̃)‖,
where the last inequality is due to Lemma 14.1.
We further have∥∥∥∥ ∂f∂w (v,w)− ∂f∂w (ṽ, w̃) ∥∥∥∥ = ∥∥∥∥∥∥ v >v∗ 2π‖w‖ ( In − ww > ‖w‖2 ) w∗∥∥∥(In − ww>‖w‖2 )w∗∥∥∥ − ṽ>v∗ 2π‖w̃‖ ( In − w̃w̃ > ‖w̃‖2 ) w∗∥∥∥(In − w̃w̃>‖w̃‖2 )w∗∥∥∥ ∥∥∥∥∥∥ ≤ ∥∥∥∥∥∥ v >v∗ 2π‖w‖ ( In − ww > ‖w‖2 ) w∗∥∥∥(In − ww>‖w‖2 )w∗∥∥∥ − v>v∗ 2π‖w̃‖ ( In − w̃w̃ > ‖w̃‖2 ) w∗∥∥∥(In − w̃w̃>‖w̃‖2 )w∗∥∥∥
∥∥∥∥∥∥ + ∥∥∥∥∥∥ v >v∗ 2π‖w̃‖ ( In − w̃w̃ > ‖w̃‖2 ) w∗∥∥∥(In − w̃w̃>‖w̃‖2 )w∗∥∥∥ − ṽ>v∗ 2π‖w̃‖ ( In − w̃w̃ > ‖w̃‖2 ) w∗∥∥∥(In − w̃w̃>‖w̃‖2 )w∗∥∥∥
∥∥∥∥∥∥ ≤ |v
>v∗| 2πc2w ‖w − w̃‖+ ‖v ∗‖ 2πcw ‖v − ṽ‖
≤ (Cv + cw)‖v ∗‖
2πc2w ‖(v,w)− (ṽ, w̃)‖,
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 `(v,w;Z) w.r.t. v is EZ [ ∂`
∂v (v,w;Z)
] = ∂f
∂v (v,w).
Let µ(x) = max{x, 0} in (5). The expected coarse gradient w.r.t. w is
EZ [ grelu(v,w;Z) ] = h(v,v∗)
2 √ 2π
w
‖w‖ − cos
( θ(w,w∗)
2
) v>v∗√
2π
w ‖w‖ + w ∗∥∥∥ w‖w‖ + w∗∥∥∥ , 3 where h(v,v∗) = ‖v‖2 + (1>mv)2 − (1>mv)(1>mv∗) + v>v∗.
Proof of Lemma 4. The first claim is true because ∂`∂v (v,w;Z) is linear in v. By (5), gµ(v,w;Z) = Z >(µ′(Zw) v)(v>σ(Zw)− (v∗)>σ(Zw∗)).
Using the fact that µ′ = σ = 1{x>0}, we have
EZ [grelu(v,w;Z)] = EZ [( m∑ i=1 viσ(Z > i w)− m∑ i=1 v∗i σ(Z > i w ∗) )( m∑ i=1 Ziviσ(Z > i w) )]
= EZ [( m∑ i=1 vi1{Z>i w>0} − m∑ i=1 v∗i 1{Z>i w∗>0} )( m∑ i=1 1{Z>i w>0}viZi )] .
Invoking Lemma 11, we have
E [ Zi1{Z>i w>0,Z>j w>0} ] = { 1√ 2π w ‖w‖ if i = j,
1 2 √ 2π w ‖w‖ if i 6= j,
and
E [ Zi1{Z>i w>0,Z>j w∗>0} ] =  cos(θ(w,w∗)/2)√ 2π w ‖w‖+w ∗ ‖ w‖w‖+w∗‖ if i = j,
1 2 √ 2π w ‖w‖ if i 6= j.
3We redefine the second term as 0n in the case θ(w,w∗) = π, or equivalently, w‖w‖ +w ∗ = 0n.
Therefore,
EZ [grelu(v,w;Z)] = m∑ i=1 v2i E [ Zi1{Z>i w>0} ] + m∑ i=1 m∑ j=1
j 6=i
vivjE [ Zi1{Z>i w>0,Z>j w>0} ]
− m∑ i=1 viv ∗ i E [ Zi1{Z>i w>0,Z>i w∗>0} ] −
m∑ i=1 m∑ j=1
j 6=i
viv ∗ jE [ Zi1{Z>i w>0,Z>j w∗>0} ]
= 1
2 √ 2π
( ‖v‖2 + (1>mv)2 ) w ‖w‖ − cos ( θ(w,w∗) 2 ) v>v∗√
2π
w ‖w‖ + w ∗∥∥∥ w‖w‖ + w∗∥∥∥ − 1
2 √ 2π
( (1>mv)(1 > mv ∗)− v>v∗ ) w ‖w‖ ,
and the result follows.
Lemma 5. If w 6= 0n and θ(w,w∗) ∈ (0, π), then the inner product between the expected coarse and true gradients w.r.t. w is
〈 EZ [ grelu(v,w;Z) ] , ∂f
∂w (v,w)
〉 = sin (θ(w,w∗))
2( √ 2π)3‖w‖ (v>v∗)2 ≥ 0.
Moreover, if further ‖v‖ ≤ Cv and ‖w‖ ≥ cw, there exists a constant Arelu > 0 depending on Cv and cw, such that
∥∥∥EZ[grelu(v,w;Z)]∥∥∥2 ≤ Arelu(∥∥∥∥∂f∂v (v,w) ∥∥∥∥2 +〈EZ[grelu(v,w;Z)], ∂f∂w (v,w) 〉) .
Proof of Lemma 5. By Lemmas 2 and 4, we have
∂f ∂w (v,w) = − v
>v∗
2π‖w‖
( In − ww >
‖w‖2 ) w∗∥∥∥(In − ww>‖w‖2 )w∗∥∥∥
and
EZ [ grelu(v,w;Z) ] = h(v,v∗)
2 √ 2π
w
‖w‖ − cos
( θ(w,w∗)
2
) v>v∗√
2π
w ‖w‖ + w ∗∥∥∥ w‖w‖ + w∗∥∥∥ .
Notice that ( In − ww >
‖w‖2
) w = 0n and ‖w∗‖ = 1, if θ(w,w∗) 6= 0, π, then we have
〈 EZ [ grelu(v,w;Z) ] , ∂f
∂w (v,w)
〉
= cos
( θ(w,w∗)
2
) (v>v∗)2
( √ 2π)3
〈 1
‖w‖
( In − ww >
‖w‖2 ) w∗∥∥∥(In − ww>‖w‖2 )w∗∥∥∥ , w∗∥∥∥ w‖w‖ + w∗∥∥∥ 〉
= cos
( θ(w,w∗)
2
) (v>v∗)2
( √ 2π)3 ‖w‖2 − (w>w∗)2 ‖‖w‖2w∗ −w(w>w∗)‖ ‖w + ‖w‖w∗‖
= cos
( θ(w,w∗)
2
) (v>v∗)2
( √ 2π)3 ‖w‖2 − (w>w∗)2√ ‖w‖4 − ‖w‖2(w>w∗)2 √ 2(‖w‖2 + ‖w‖(w>w∗))
= cos
( θ(w,w∗)
2
) (v>v∗)2
4( √ π‖w‖)3 ‖w‖2 − (w>w∗)2√ ‖w‖2 − (w>w∗)2 √ ‖w‖+ (w>w∗)
= cos
( θ(w,w∗)
2
) (v>v∗)2√1− w>w∗‖w‖ 4( √ π)3‖w‖
= cos
( θ(w,w∗)
2
) (v>v∗)2 √ 1− cos(θ(w,w∗))
4( √ π)3‖w‖
= sin (θ(w,w∗))
2( √ 2π)3‖w‖ (v>v∗)2.
To show the second claim, without loss of generality, we assume ‖w‖ = 1. Denote θ := θ(w,w∗). By Lemma 1, we have
∂f ∂v (v,w) = 1 4
( Im + 1m1 > m ) v − 1
4
(( 1− 2θ
π
) Im + 1m1 > m ) v∗.
By Lemma 4,
EZ [ grelu(v,w;Z) ] = h(v,v∗)
2 √ 2π
w
‖w‖ − cos
( θ
2
) v>v∗√
2π
w ‖w‖ + w ∗∥∥∥ w‖w‖ + w∗∥∥∥ , (22) where
h(v,v∗) = ‖v‖2 + (1>mv)2 − (1>mv)(1>mv∗) + v>v∗ = v> ( Im + 1m1 > m ) v − v>(1m1>m − Im)v∗
= v> ( Im + 1m1 > m ) v − v> ( 1m1 > m + ( 1− 2θ
π
) Im ) v∗ + 2 ( 1− θ
π
) v>v∗
= 4v> ∂f
∂v (v,w) + 2
( 1− θ
π
) v>v∗, (23)
and by the first claim,
〈 EZ [ grelu(v,w;Z) ] , ∂f
∂w (v,w)
〉 =
sin (θ)
2( √ 2π)3‖w‖ (v>v∗)2.
Hence, for some Arelu depending only on Cv and cw, we have∥∥∥EZ[grelu(v,w;Z)]∥∥∥2 = ∥∥∥∥∥2v> ∂f∂v (v,w)√2π w‖w‖ + cos ( θ 2 ) v>v∗√ 2π  w ‖w‖ − w ‖w‖ + w ∗∥∥∥ w‖w‖ + w∗∥∥∥ 
+ ( 1− θ
π − cos
( θ
2
)) v>v∗√
2π
w
‖w‖ ∥∥∥∥∥ 2
≤ 6C 2 v
π ∥∥∥∥∂f∂v (v,w) ∥∥∥∥2 + cos2(θ2 ) 3(v>v∗)2 2π ∥∥∥∥∥∥ w‖w‖ − w ‖w‖ + w ∗∥∥∥ w‖w‖ + w∗∥∥∥ ∥∥∥∥∥∥ 2
+ ( 1− θ
π − cos
( θ
2
))2 3(v>v∗)2
2π
≤ 6C 2 v
π ∥∥∥∥∂f∂v (v,w) ∥∥∥∥2 + cos2(θ2 ) 3θ2 8π (v>v∗)2 + ( 1− θ π − cos ( θ 2 ))2 3(v>v∗)2 2π
≤ 6C 2 v
π ∥∥∥∥∂f∂v (v,w) ∥∥∥∥2 + 3π8 cos2 ( θ 2 ) sin2 ( θ 2 ) (v>v∗)2 + 3 sin(θ) 2π (v>v∗)2
≤ Arelu (∥∥∥∥∂f∂v (v,w) ∥∥∥∥2 +〈EZ[grelu(v,w;Z)], ∂f∂w (v,w) 〉) ,
where the equality is due to (22) and (23), the first inequality is due to Cauchy-Schwarz inequality, the second inequality holds because the angle between w‖w‖ and w ‖w‖+w ∗
‖ w‖w‖+w∗‖ is θ2 and∥∥∥∥ w‖w‖ − w‖w‖+w∗‖ w‖w‖+w∗‖ ∥∥∥∥ ≤ θ2 , whereas the third inequality is due to sin(x) ≥ 2xπ , cos(x) ≥ 1 − 2xπ , and(
1− 2x π − cos(x)
)2 ≤ (
cos(x)− 1 + 2x π
)( cos(x) + 1− 2x
π
) ≤ sin(x)(2 cos(x)) = sin(2x)
for all x ∈ [0, π2 ].
Lemma 6. When Algorithm 1 converges, EZ [ ∂` ∂v (v,w;Z) ] and EZ [ grelu(v,w;Z) ] vanish simultaneously, which only occurs at the
1. Saddle points where (8) is satisfied according to Proposition 1.
2. Minimizers of (2) where v = v∗, θ(w,w∗) = 0, or v = (Im+1m1>m) −1(1m1 > m−Im)v∗,
θ(w,w∗) = π.
Proof of Lemma 6. By Lemma 4, suppose we have EZ [ ∂` ∂v (v,w;Z) ] = 1 4 ( Im + 1m1 > m ) v − 1 4 (( 1− 2 π θ(w,w∗) ) Im + 1m1 > m ) v∗ = 0m (24)
and EZ [ grelu(v,w;Z) ] = h(v,v∗)
2 √ 2π
w
‖w‖ − cos
( θ(w,w∗)
2
) v>v∗√
2π
w ‖w‖ + w ∗∥∥∥ w‖w‖ + w∗∥∥∥ = 0n, (25) where h(v,v∗) = ‖v‖2 + (1>mv)2− (1>mv)(1>mv∗) +v>v∗. By (25), we must have θ(w,w∗) = 0 or θ(w,w∗) = π or v>v∗ = 0.
If θ(w,w∗) = 0, then by (24), v = v∗, and (25) is satisfied.
If θ(w,w∗) = π, then by (24), v = (Im + 1m1>m) −1(1m1 > m − Im)v∗, and (25) is satisfied.
If v>v∗ = 0, then by (24), we have the expressions for v and θ(w,w∗) from Proposition 1, and (25) is satisfied.
Lemma 7. If w 6= 0n and θ(w,w∗) ∈ (0, π), then
EZ [gcrelu(v,w;Z)] = p(0,w)h(v,v∗)
2
w
‖w‖ − (v>v∗) csc(θ/2) · q(θ,w)
w ‖w‖ + w ∗∥∥∥ w‖w‖ + w∗∥∥∥ − (v>v∗) (p(θ,w)− cot(θ/2) · q(θ,w)) w
‖w‖ , (26)
where h(v,v∗) := ‖v‖2 + (1>mv)2 − (1>mv)(1>mv∗) + v>v∗ same as in Lemma 5, and
p(θ,w) := 1
2π
∫ π 2
−π2 +θ cos(φ)ξ
( sec(φ)
‖w‖
) dφ, q(θ,w) := 1
2π
∫ π 2
−π2 +θ sin(φ)ξ
( sec(φ)
‖w‖
) dφ
with ξ(x) := ∫ x 0 r2 exp(− r 2
2 )dr. The inner product between the expected coarse and true gradients w.r.t. w 〈 EZ [ gcrelu(v,w;Z) ] , ∂f
∂w (v,w)
〉 = q(θ,w)
2π‖w‖ (v>v∗)2 ≥ 0.
Moreover, if further ‖v‖ ≤ Cv and ‖w‖ ≥ cw, there exists a constant Acrelu > 0 depending on Cv and cw, such that∥∥∥EZ[gcrelu(v,w;Z)]∥∥∥2 ≤ Acrelu(∥∥∥∥∂f∂v (v,w) ∥∥∥∥2 +〈EZ[gcrelu(v,w;Z)], ∂f∂w (v,w) 〉) . Proof of Lemma 7. Denote θ := θ(w,w∗). We first compute EZ [ gcrelu(v,w;Z) ] . By (5),
gµ(v,w;Z) = Z >(µ′(Zw) v)(v>σ(Zw)− (v∗)>σ(Zw∗)).
Since µ′ = 1{0<x<1} and σ = 1{x>0}, we have
EZ [gcrelu(v,w;Z)] = EZ [( m∑ i=1 viµ ′(Z>i w)− m∑ i=1 v∗i µ ′(Z>i w ∗) )( m∑ i=1 Ziviσ(Z > i w) )]
= EZ [( m∑ i=1 vi1{0<Z>i w<1} − m∑ i=1 v∗i 1{0<Z>i w∗<1} )( m∑ i=1 1{Z>i w>0}viZi )]
= p(0,w)h(v,v∗)
2
w
‖w‖ − (v>v∗) csc(θ/2) · q(θ,w)
w ‖w‖ + w ∗∥∥∥ w‖w‖ + w∗∥∥∥ − (v>v∗) (p(θ, ‖w‖)− cot(θ/2) · q(θ,w)) w
‖w‖ .
In the last equality above, we called Lemma 12. Notice that ( In − ww >
‖w‖2
) w = 0n and ‖w∗‖ = 1. If θ(w,w∗) 6= 0, π, then the inner product
between EZ [ gcrelu(v,w;Z) ] and ∂f∂w (v,w) is given by〈
EZ [ gcrelu(v,w;Z) ] , ∂f
∂w (v,w)
〉
= csc
( θ
2
) q(θ,w)
2π (v>v∗)2
〈 1
‖w‖
( In − ww >
‖w‖2 ) w∗∥∥∥(In − ww>‖w‖2 )w∗∥∥∥ , w∗∥∥∥ w‖w‖ + w∗∥∥∥ 〉
= q(θ,w)
2π‖w‖ (v>v∗)2 ≥ 0.
In the last line, q(θ,w) ≥ 0 because sin(φ)ξ (
sec(φ) ‖w‖ ) is odd in φ and positive for φ ∈ (0, π2 ].
Next, we bound ∥∥∥EZ[gcrelu(v,w;Z)]∥∥∥2. Since (23) gives
h(v,v∗) = 4v> ∂f
∂v (v,w) + 2
( 1− θ
π
) v>v∗,
where according to Lemma 1,
∂f ∂v (v,w) = 1 4
( Im + 1m1 > m ) v − 1
4
(( 1− 2θ
π
) Im + 1m1 > m ) v∗.
We rewrite the coarse partial gradient w.r.t. w in (26) as EZ [gcrelu(v,w;Z)] = p(0,w) ( 2v> ∂f
∂v (v,w) +
( 1− θ
π
) v>v∗ ) w
‖w‖
− (v>v∗) csc(θ/2) · q(θ,w) w ‖w‖ + w ∗∥∥∥ w‖w‖ + w∗∥∥∥ − (v>v∗) (p(θ,w)− cot(θ/2) · q(θ,w)) w
‖w‖ .
= 2p(0,w)v> ∂f
∂v (v,w)
w
‖w‖
+ (v>v∗) (( 1− θ
π
) p(0,w)− p(θ,w) ) w
‖w‖
+ (v>v∗) csc(θ/2) · q(θ,w)  w ‖w‖ − w ‖w‖ + w ∗∥∥∥ w‖w‖ + w∗∥∥∥ 
− (v>v∗)(csc(θ/2)− cot(θ/2))q(θ,w) w ‖w‖
To prove the last claim, we notice that in the above equality,∥∥∥∥v> ∂f∂v (v,w) w‖w‖ ∥∥∥∥2 ≤ C2v ∥∥∥∥∂f∂v (v,w) ∥∥∥∥2 , (27) ∥∥∥∥(csc(θ/2)− cot(θ/2))q(θ,w) w‖w‖
∥∥∥∥2 = (csc(θ/2)− cot(θ/2))2q(θ,w)2 ≤ q(θ,w)2, (28) ∥∥∥∥∥∥csc(θ/2) · q(θ,w)  w ‖w‖ − w ‖w‖ + w ∗∥∥∥ w‖w‖ + w∗∥∥∥ ∥∥∥∥∥∥ 2 ≤ ( θ/2 sin(θ/2) )2 q(θ,w)2 ≤ π 2 4 q(θ,w)2. (29)
Now, what is left is to bound (( 1− θπ ) p(0,w)− p(θ,w) )2 , using a multiple of q(θ,w). Recall that
p(θ,w) = 1
2π
∫ π 2
−π2 +θ cos(φ)ξ
( sec(φ)
‖w‖
) dφ, q(θ,w) = 1
2π
∫ π 2
−π2 +θ sin(φ)ξ
( sec(φ)
‖w‖
) dφ.
We first show that both (( 1− θπ ) p(0,w)− p(θ,w) )2 and q(θ,w) are symmetric with respect to θ = π2 on [0, π]. This is because
q(π − θ,w) = 1 2π
∫ π 2
π 2−θ
sin(φ)ξ
( sec(φ)
‖w‖
) dφ
= q(θ,w) + 1
2π ∫ −π2 +θ π 2−θ sin(φ)ξ ( sec(φ) ‖w‖ ) dφ = q(θ,w).
and ( 1− π − θ
π
) p(0,w)− p(π − θ,w)
= θ
π
1
2π
∫ π 2
−π2 cos(φ)ξ
( sec(φ)
‖w‖
) dφ− 1
2π
∫ π 2
π 2−θ
cos(φ)ξ
( sec(φ)
‖w‖
) dφ
=
( θ π − 1 ) 1 2π ∫ π 2
−π2 cos(φ)ξ
( sec(φ)
‖w‖
) dφ+ 1
2π
∫ π 2−θ
−π2 cos(φ)ξ
( sec(φ)
‖w‖
) dφ
= − ((
1− θ π
) p(0,w)− p(θ,w) ) .
Therefore, it suffices to consider θ ∈ [π2 , π] only. Then calling Lemma 13 for θ ∈ [ π 2 , π], we have p(θ,w) ≤ q(θ,w) and ( 1− θπ ) p(0,w) ≤ q(θ,w). Therefore,(
p(θ,w)− (
1− θ π
) p(0,w) )2 ≤ q(θ,w)2.
Combining the above estimate together with (27), (28) and (29), and using Cauchy-Schwarz inequality, we have
‖EZ[gcrelu(v,w;Z)]‖2 ≤ 4 ( 4p(0,w)C2v ∥∥∥∥∂f∂v (v,w) ∥∥∥∥2 + (2 + π24 ) q(θ,w)2(v>v∗)2 ) ,
where p(0,w) and q(θ,w) are uniformly bounded. This completes the proof.
Lemma 8. When Algorithm 1 converges, EZ [ ∂` ∂v (v,w;Z) ] and EZ [ gcrelu(v,w;Z) ] vanish simultaneously, which only occurs at the
1. Saddle points where (8) is satisfied according to Proposition 1.
2. Minimizers of (2) where v = v∗, θ(w,w∗) = 0, or v = (Im+1m1>m) −1(1m1 > m−Im)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) defined 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
EZ [ gid(v,w;Z) ] =
1√ 2π
( ‖v‖2 w
‖w‖ − (v>v∗)w∗
) .
If θ(w,w∗) = π and v = ( Im + 1m1 > m )−1 (1m1
> m − Im)v∗,∥∥∥EZ[gid(v,w;Z)]∥∥∥ = 2(m− 1)√ 2π(m+ 1)2 (1>mv ∗)2 ≥ 0,
i.e., EZ [ gid(v,w;Z) ] does not vanish at the local minimizers if 1>mv ∗ 6= 0 and m > 1,.
Proof of Lemma 9. By (5), gµ(v,w;Z) = Z >(µ′(Zw) v)(v>σ(Zw)− (v∗)>σ(Zw∗)).
Using the facts that µ′ = 1 and σ = 1{x>0}, we have
EZ [gid(v,w;Z)] = EZ [( m∑ i=1 vi1{Z>i w>0} − m∑ i=1 v∗i 1{Z>i w∗>0} )( m∑ i=1 viZi )]
= m∑ i=1 m∑ j=1 vivjE [ Zi1{Z>j w>0} ] − m∑ i=1 m∑ j=1 v∗i vjE [ Zi1{Z>j w∗>0} ] =
1√ 2π
( ‖v‖2 w
‖w‖ − (v>v∗)w∗
) .
In the last equality above, we called the third identity in Lemma 11. If θ(w,w∗) = π and v =( Im + 1m1 > m )−1 (1m1 > m − Im)v∗, then
‖EZ [gid(v,w;Z)]‖ = 1√ 2π |v>(v + v∗)|
= 1√ 2π ∣∣∣(v∗)>(1m1>m − Im)(Im + 1m1>m)−1 ((Im + 1m1>m)−1(1m1>m − Im) + Im)v∗∣∣∣ =
2√ 2π ∣∣∣(v∗)>(1m1>m − Im)(Im + 1m1>m)−1(Im + 1m1>m)−11m(1>mv∗)∣∣∣ =
2√ 2π(m+ 1)2 ∣∣(v∗)>(1m1>m − Im)1m(1>mv∗)∣∣ =
2(m− 1)√ 2π(m+ 1)2 (1>mv ∗)2.
In the third equality, we used the identity (Im + 1m1>m)1m = (m+ 1)1m twice.
Lemma 10. If w 6= 0n and θ(w,w∗) ∈ (0, π), then the inner product between the expected coarse and true gradients w.r.t. w is〈
EZ [ gid(v,w;Z) ] , ∂f
∂w (v,w)
〉 = sin (θ(w,w∗))
( √ 2π)3‖w‖ (v>v∗)2 ≥ 0.
When θ(w,w∗)→ π, v → ( Im + 1m1 > m )−1 (1m1
> m − Im)v∗, if 1>mv∗ 6= 0 and m > 1, we have∥∥∥EZ[gid(v,w;Z)]∥∥∥2∥∥∥∂f∂v (v,w)∥∥∥2 + 〈EZ[gid(v,w;Z)], ∂f∂w (v,w)〉 → +∞.
Proof of Lemma 10. By Lemmas 2 and 4, we have
∂f ∂w (v,w) = − v
>v∗
2π‖w‖
( In − ww >
‖w‖2 ) w∗∥∥∥(In − ww>‖w‖2 )w∗∥∥∥
and
EZ [ gid(v,w;Z) ] =
1√ 2π
( ‖v‖2 w
‖w‖ − (v>v∗)w∗
) .
Since ( In − ww >
‖w‖2 ) w = 0n and ‖w∗‖ = 1, if θ(w,w∗) 6= 0, π, then we have〈
EZ [ gid(v,w;Z) ] , ∂f
∂w (v,w)
〉 = (v>v∗)2
( √ 2π)3‖w‖
(w∗)> ( In − ww >
‖w‖2 ) w∗∥∥∥(In − ww>‖w‖2 )w∗∥∥∥
= (v>v∗)2
( √ 2π)3‖w‖ 1− (w
>w∗)2
‖w‖2√ 1− (w >w∗)2
‖w‖2
= (v>v∗)2
( √ 2π)3‖w‖
√ 1− (w >w∗)2
‖w‖2
= (v>v∗)2
( √ 2π)3‖w‖ sin(θ(w,w∗)).
When θ(w,w∗) → π, v → ( Im + 1m1 > m )−1 (1m1 > m − Im)v∗, both ∥∥∥∂f∂v (v,w)∥∥∥ and 〈 EZ [ gid(v,w;Z) ] , ∂f∂w (v,w) 〉 converge to 0. But if 1>mv
∗ 6= 0 and m > 1,∥∥∥EZ[gid(v,w;Z)]∥∥∥→ 2(m−1)√2π(m+1)2 (1>mv∗)2 > 0, which completes the proof.
Theorem 1. Let {(vt,wt)} be the sequence generated by Algorithm 1 with ReLU µ(x) = max{x, 0} or clipped ReLU µ(x) = min{max{x, 0}, 1}. Suppose ‖wt‖ ≥ cw for all t with some cw > 0. Then if the learning rate η > 0 is sufficiently small, for any initialization (v0,w0), the objective sequence {f(vt,wt)} is monotonically decreasing, and {(vt,wt)} converges to a saddle point or a (local) minimizer of the population loss minimization (2). In addition, if 1>mv
∗ 6= 0 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 = (Im + 1m1 > m) −1(1m1 > m − Im)v∗.
Proof of Theorem 1 . We first prove the upper boundedness of {vt}. Due to the coerciveness of f(v,w) w.r.t v, there exists Cv > 0, such that ‖v‖ ≤ Cv for any v ∈ {v ∈ Rm : f(v,w) ≤ f(v0,w0) for some w}. In particular, ‖v0‖ ≤ Cv . Using induction, suppose we already have f(vt,wt) ≤ f(v0,w0) and ‖vt‖ ≤ Cv . If θ(wt,w∗) = 0 or π, then θ(wT ,w∗) = 0 or π for all T ≥ t, and the original problem reduces to a quadratic program in terms of v. So {vt}will converge to v∗ or (Im+1m1>m)
−1(1m1 > m−Im)v∗ by choosing a suitable step size η. In either case, we have∥∥∥EZ[ ∂`∂v (vt,wt;Z)]∥∥∥ and ∥∥∥EZ[grelu(vt,wt;Z)]∥∥∥ both converge to 0. Else if θ(wt,w∗) ∈ (0, π), we define for any a ∈ [0, 1] that
vt(a) := vt − a(vt+1 − vt) = vt − aηEZ [ ∂`
∂v (vt,wt;Z)
]
and
wt(a) := wt − a(wt+1 −wt) = wt − aηEZ [ grelu(v t,wt;Z) ] ,
which satisfy
vt(0) = vt, vt(1) = vt+1, wt(0) = wt, wt(1) = wt+1.
Let us fix 0 < c ≤ cw and C ≥ Cv . By the expressions of EZ [ ∂` ∂v (v t,wt;Z) ]
and EZ [grelu(vt,wt;Z)] given in Lemma 4, and since ‖wt‖ = 1, for sufficiently small η̃ depending on C0 and cw, with η ≤ η̃, it holds that ‖vt(a)‖ ≤ C and ‖wt(a)‖ ≥ c for all a ∈ [0, 1]. Possibly at some point a0 where θ(wt(a0),w∗) = 0 or π, the partial gradient ∂f∂w (v t(a0),w t(a0)) does not
exist. Otherwise, ∥∥∥ ∂f∂w (vt(a),wt(a))∥∥∥ is uniformly bounded for all a ∈ [0, 1]/{a0}, which makes
it integrable over the interval [0, 1]. Then for some constants L and Arelu depending on C and c, we
have f(vt+1,wt+1) = f(vt + (vt+1 − vt),wt + (wt+1 −wt))
= f(vt,wt) + ∫ 1 0 〈 ∂f ∂v (vt(a),wt(a)),vt+1 − vt 〉 da
+ ∫ 1 0 〈 ∂f ∂w (vt(a),wt(a)),wt+1 −wt 〉 da
= f(vt,wt) +
〈 ∂f
∂v (vt,wt),vt+1 − vt
〉 + 〈 ∂f
∂w (vt,wt),wt+1 −wt 〉 +
∫ 1 0 〈 ∂f ∂v (vt(a),wt(a))− ∂f ∂v (vt,wt),vt+1 − vt 〉 da
+ ∫ 1 0 〈 ∂f ∂w (vt(a),wt(a))− ∂f ∂w (vt,wt),wt+1 −wt 〉 da
≤ f(vt,wt)− ( η − Lη 2
2 )∥∥∥∥∂f∂v (vt,wt) ∥∥∥∥2
− η 〈 ∂f
∂w (vt,wt),EZ
[ grelu(v t,wt;Z) ]〉 + Lη2
2 ∥∥∥EZ[grelu(vt,wt;Z)]∥∥∥2 ≤ f(vt,wt)− ( η − (1 +Arelu) Lη2
2 )∥∥∥∥∂f∂v (vt,wt) ∥∥∥∥2
− ( η − AreluLη 2
2
)〈 ∂f
∂w (vt,wt),EZ
[ grelu(v t,wt;Z) ]〉 . (30)
The third equality is due to the fundamental theorem of calculus. In the first inequality, we called Lemma 3 for (vt,wt) and (vt(a),wt(a)) with a ∈ [0, 1]/{a0}. In the last inequality, we used Lemma 5. So when η < η0 := min { 2 (1+Arelu)L , η̃ }
, we have f(vt+1,wt+1) ≤ f(vt,wt) ≤ f(v0,w0), and thus ‖vt+1‖ ≤ Cv . Summing up the inequality (30) over t from 0 to∞ and using f ≥ 0, we have ∞∑ t=0 ( 1− (1 +Arelu) Lη 2 )∥∥∥∥∂f∂v (vt,wt) ∥∥∥∥2 + (1− AreluLη2 )〈 ∂f ∂w (vt,wt),EZ [ grelu(v t,wt;Z) ]〉
≤ f(v0,w0)/η <∞. Hence,
lim t→∞ ∥∥∥∥∂f∂v (vt,wt) ∥∥∥∥ = 0
and
lim t→∞
〈 ∂f
∂w (vt,wt),EZ
[ grelu(v t,wt;Z) ]〉 = 0.
Invoking Lemma 5 again, we further have
lim t→∞ ∥∥∥EZ[grelu(vt,wt;Z)]∥∥∥ = 0. Invoking Lemma 6, we have that coarse gradient descent with ReLU µ(x) (subsequentially) 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.

heading: F. CONVERGENCE TO GLOBAL MINIMIZERS
text: We prove that if the initialization weights (v0,w0) satisfy (v0)>v∗ > 0, θ(w0,w∗) < π2 and (1>mv ∗)(1>mv 0) ≤ (1>mv∗)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 (v0,w0) satisfies (v0)>v∗ > 0, θ(w0,w∗) < π2 and (1 > mv ∗)(1>mv
0) ≤ (1>mv∗)2, then by using the vanila or clipped ReLU STE for sufficiently learning rate η > 0, we have (vt)>v∗ > 0 and θ(wt,w∗) < π2 for all t > 0, and {(vt,wt)} converges to a global minimizer.
Proof of Theorem 2. Proof by induction. Suppose (vt)>v∗ > 0, θ(wt,w∗) < π2 and (1>mv ∗)(1>mv t) ≤ (1>mv∗)2. Then for small enough η, we have
(vt+1)>v∗ = ( vt − η
4
( (Im + 1m1 > m)v t − ((
1− 2 π θ(wt,w∗)
) Im + 1m1 > m ) v∗ ))> v∗
= (
1− η 4
) (vt)>v∗ + η
4
( (1>mv ∗)2 − (1>mv∗)(1>mv>) )
+ η
4
( 1− 2
π θ(wt,w∗)
) ‖v∗‖2 > 0.
and
(1>mv t+1)(1>mv ∗) =
( 1− (m+ 1)η
4
) (1>mv t)(1>mv ∗) + η
4
( m+ 1− 2
π θ(w,w∗)
) (1>mv ∗)2
≤ (
1− η 2π θ(w,w∗)
) (1>mv ∗)2 ≤ (1>mv∗)2.
Moreover, by Lemmas 4, 5 and 7, both EZ[grelu(v,w;Z)] and EZ[gcrelu(v,w;Z)] can be written in the form of α1 ( Im − ww >
‖w‖2
) w∗+α2w, where α1 ≤ 0 and α2 is bounded by a constant depending
on Cv and cw. Therefore, (wt+1)>w∗ = (1− ηα2)(wt)>w∗ − α1(w∗)> ( Im − ww>
‖w‖2
) w∗ ≥ (1− ηα2)(wt)>w∗ > 0,
and thus θ(wt+1,w∗) < π2 . Finally, since {(v t,wt)} converges, it can only converge to a global minimizer.

Review summaries:
The paper examines the use of STE for learning simple one-layer convolutional networks with binary activations and non overlapping patches. In this setting, the gradients are 0 almost everywhere (gradient of sign(x)) hence it is not clear how to use gradient descent. The approach studied here is to instead use the gradient of an alternative function such as ReLU or identity which is not always 0. The authors prove that if the ReLU's gradient is used then under gaussian distribution, the algorithm will converge to the local minimas/saddle points of the expected squared loss. they also show that the same does not hold for the identity's gradient.