PROPER STRAIGHT THROUGH ESTIMATOR: BREAK-
Year: 2021

Abstract:
In the quantized network, its gradient shows vanishing except for nondifferentiable points. The network thus cannot be learned by the standard backpropagation, so that an alternative approach called Straight Through Estimator (STE), which replaces the part of the gradient with a simple differentiable function, is used. While STE is known to work well for learning the quantized network empirically, it has not been established theoretically. A recent study by Yin et al. (2019) has provided theoretical support for STE. However, its justification is still limited to the model in the one-hidden layer network with the binary activation where Gaussian generates the input data, and the true labels are output from the teacher network with the same binary network architecture. In this paper, we discuss the effectiveness of STEs in more general situations without assuming the shape of the input distribution and the labels. By considering the scale symmetry of the network and specific properties of the STEs, we find that STE with clipped Relu is superior to STEs with identity function and vanilla Relu. The clipped Relu STE, which breaks the scale symmetry, may pick up one of the local minima degenerated in scales, while the identity STE and vanilla Relu STE, which keep the scale symmetry, may not pick it up. To confirm this observation, we further present an analysis of a simple misspecified model as an example. We find that all the stationary points are identical with the vanishing points of the cRelu STE gradient, while some of them are not identical with the vanishing points of the identity and Relu STE. Finally we have numerically confirmed the observation for the mixture Gaussian model with various teacher network.

In the quantized network, its gradient shows vanishing except for nondifferentiable points. The network thus cannot be learned by the standard backpropagation, so that an alternative approach called Straight Through Estimator (STE), which replaces the part of the gradient with a simple differentiable function, is used. While STE is known to work well for learning the quantized network empirically, it has not been established theoretically. A recent study by Yin et al. (2019) has provided theoretical support for STE. However, its justification is still limited to the model in the one-hidden layer network with the binary activation where Gaussian generates the input data, and the true labels are output from the teacher network with the same binary network architecture. In this paper, we discuss the effectiveness of STEs in more general situations without assuming the shape of the input distribution and the labels. By considering the scale symmetry of the network and specific properties of the STEs, we find that STE with clipped Relu is superior to STEs with identity function and vanilla Relu. The clipped Relu STE, which breaks the scale symmetry, may pick up one of the local minima degenerated in scales, while the identity STE and vanilla Relu STE, which keep the scale symmetry, may not pick it up. To confirm this observation, we further present an analysis of a simple misspecified model as an example. We find that all the stationary points are identical with the vanishing points of the cRelu STE gradient, while some of them are not identical with the vanishing points of the identity and Relu STE. Finally we have numerically confirmed the observation for the mixture Gaussian model with various teacher network.
1 INTRODUCTION
Quantization of the weights and the activations is a promising technique to save the memory and accelerate the inference speed in deep neural networks (DNNs) which have been getting wider and deeper in recent years. There are two main approaches of the quantization for DNNs (Krishnamoorthi (2018)): Post-training quantization (PTQ) and quantization-aware training (QAT). In the PTQ, the pre-trained networks are simply quantized without re-training the model. This approach allows us to achieve the nearly floating point accuracy at 8-bits, while below 8-bits, this results in significant accuracy degradation. Although there have been recent attempts to alleviate the accuracy degradation in PTQ (Banner et al. (2018); Choukroun et al. (2019); Zhao et al. (2019); Kryzhanovskiy et al. (2021)), the QAT, where the quantized weights and activations are trained (e.g., Zhou et al. (2016); Hubara et al. (2017); Rastegari et al. (2016)), usually leads to better accuracy.
The difficulty in QAT is that the weights and the activations are discretized, and intrinsically nondifferentiable. If we take the derivatives forcibly, they either vanish or diverge. To avoid this problem, we replace them with the derivatives of some differentiable function in the backward pass only, called the Straight-Through Estimator (STE) (Hinton et al. (2012); Bengio et al. (2013)). Since the replacement leads to bias, it is not always possible to learn the network successfully. However this approach can be applied to the low bits below 8-bits with tolerable accuracy degradation (Zhou et al. (2016); Choi et al. (2018); Esser et al. (2019); Bhalgat et al. (2020)).
Originally, STE was introduced as the replacement with the derivative of the identity function by Hinton et al. (2012). Later, the term ”STE” has been extensively used as the replacement with various functions. In the binary network, Bengio et al. (2013) studied the replacement with the derivative of sigmoid function as well as the original STE, and Hubara et al. (2017) used the derivative of the identity function clipped in the region of |x| ≤ 1. In low-bits network, Zhou et al. (2016) used the quantized gradient after the replacement with the identity function, and Choi et al. (2018), Esser et al. (2019), and Bhalgat et al. (2020) used the the derivative of the clipped Relu (cRelu) for leaning the step size.
1.1 MAIN CONTRIBUTION
In this paper, we refer to the proxy of the gradient as the STE gradient 1. Although STE has been much success in training quantized DNNs empirically, the theoretical justification is very limited to the specific situations. It also remains unclear what differentiable functions should be used for the STE gradient. Our aim is to shed light on a new factor to determine the functions.
Without assuming the shape of the input distribution and the loss function, we first discuss the properties of the three STEs, identity STE, Relu STE, and cRelu STE in one-hidden-layer network with binary activation from the perspective of the intrinsic symmetry. We find that because the identity STE and Relu STE keep the scale symmetry, the STE gradients should be zero for any scale to converge to the true (local) minimum of the loss function, while in the case of the cRelu STE, which breaks the scale symmetry, even if its gradient is not zero at the minimum at some scale and it becomes zero at other scale, the cRelu STE converges to the minimum. Therefore, the backpropagation using the cRelu STE is most likely to converge to the minimum of the loss function among three STEs. This result reveals differences between Relu STE and cRelu STE, which could not be found in the models by Yin et al. (2019) and Long et al. (2021).
Recently, Kunin et al. (2020) discussed the symmetry embedded in non-quantized neural networks, and the effect of its breaking such as the discretization, the weight decay, and the stochasticity during training. We discuss the new effect of breaking symmetry by the proxy of the gradient in training the quantized neural networks.
Next, to confirm this observation, we have studied a similar model of the Gaussian input as discussed in Yin et al. (2019) as an example. We have employed a misspecified model suitable for most practical situations. Unlike the previous study, the true labels are presumed to be generated by the non-quantized network architecture. Consequently, we find that ignoring the scale degeneracy for the weight in front of the activation, all the stationary points are identical with the vanishing points in the cRelu STE gradient, while some of them are not identical with the vanishing points in the identity and the Relu STE gradients. In particular, at the global minimum point, both identity and Relu STE gradients do not vanish. Finally, to confirm it in more general case, we have numerically studied the model of the mixture Gaussian input.
1.2 RELATED WORKS
Recently, there have been a few theoretical studies on the justification of the STEs. Shekhovtsov & Yanush (2020) derived STE by the linearization of their proposed estimator in a stochastic binary deep network where the noises are injected before the binary activation to be a smooth network. Cheng et al. (2019) argued that STE can be interpreted as a projected Wasserstein gradient flow under certain conditions. Yin et al. (2019) that inspired our study examined which of the three STE’s, identity STE, Relu STE, and cRelu STE, converges to the true minimum in one-hiddenlayer network with binarized activation and Gaussian input data. They clarified that if the labels are obtained from a teacher network of the same architecture (Du et al. (2018)), all the three STEs show the non-negative correlation with the population gradient, but only the identity STE gradient does not become zero in a local minimum. This indicates that the back-propagation using either Relu STE or cRelu STE may convergence to the local minimum, while it is impossible to show the convergence to the local minimum using the identity STE. Furthermore, they showed that all the three STEs gradients vanish at the global minimum, indicating that they may achieve the global minimum independent of the choice of the STE, if we start with an appropriate initial value.
1It is called “coarse” gradient in Yin et al. (2019).
Very recently, Long et al. (2021) have discussed the justification of a class of STEs with certain monotonicity, which generalizes the Relu STE, in one-hidden-layer network for the hinge loss, claiming that the STE gradients vanish at the global minimum.
1.3 NOTATIONS
|·| denotes the Euclidean norm of a vector. A capital letter in bold denotes the matrix, e.g., Z, and its component is given by the letter with lower indices, e.g., Zij . A small letter in bold denotes the vector, e.g., g, and its component is given by the letter with lower indices, e.g., gj . Im×m is the m×m identity matrix, and 1m is the m-dimensional vector of all ones.
2 GENERAL DISCUSSION ON STES IN ONE-HIDDEN LAYER CNN WITH BINARY ACTIVATION
2.1 PRELIMINARIES
The one-hidden layer convolutional neural network with a binary activation function is realized as y(Z,v,w) = vTσ(Zw) = ∑ i viσ( ∑ j Zijwj), (1)
where σ(·) is the step function,
σ(x) = { 1 x > 0, 0 x ≤ 0, (2)
Z ∈ Rm×n is constructed by deividing the input data, e.g. an image, into m patches with size n, and w ∈ Rn and v ∈ Rm are the trainable weights, i.e. kernels, in the first and second layers, respectively (Du et al. (2018)). We assume that Z is generated by a continuous distribution P (Z). Due to the property of the step function, the network is invariant under the scale transformation for w, w → kw (k > 0). Using the loss for each sample, ` (y(v,w,Z); y∗) with y∗ being the labels generated by the true distribution q(y∗|Z), we can express the population loss function as
L(v,w; y∗) = EZ [`(y(v,w,Z); y∗)] . (3)
2.2 BACK-PROPAGATION AND STE GRADIENTS
The gradients of the loss for each sample w.r.t v and w can be formally expressed as ∂`
∂vi = ∂` ∂y ∂y ∂vi = ∂` ∂y σ( ∑ j Zijwj), (4)
∂` ∂wj = ∑ i ∂` ∂xi ∂xi ∂wj = ∑ i ∂` ∂y viδ ∑ j′ Zij′wj′ Zij , (5) where xi = σ( ∑ j Zijwj) and we used the property that the derivative of the step function is Dirac’s delta function, dσ(z)/dz = δ(z). The gradient w.r.t v can be utilized for learning the network, while it is impossible to use the gradient w.r.t w due to the presence of the delta function. Remarkably, both population gradients w.r.t v and w are smooth functions due to its average over the continuous distribution,
EZ [ ∂`
∂vi
] = ∫ ∏ i′,j′ dZ̃i′j′ P ({Z̃i′j′}) ∂` ∂y σ ( Z̃i1|w| ) , (6)
EZ [ ∂`
∂wj
] = 1 |w| ∑ i Fij(v,w, Z̃i1 = 0),
Fij(v,w, Z̃i1) ≡ ∫  ∏
(i′,j′) 6=(i,1)
dZ̃i′j′ P ({Z̃i′j′})[ ∂` ∂y vi ∑ k′ Oik′jZ̃ik′ ] , (7)
where we used the orthogonal transformation, Z̃ik′ = ∑ j′ Oik′j′Zij′ , Zij′ = ∑ k′ Oik′j′Z̃ik′ , with
Oi1j′ = wj′/|w| and ∑ j′ Oik′j′Oil′j′ = δk′,l′ . We can easily see that the gradient w.r.t. w does not
have the component in w-direction, ∑ j wjEZ [ ∂` ∂wj ] = 0. Its vanishing originally follows from the scale symmetry of the network: Because the loss has the scale symmetry, its gradient in the w-direction is obviously zero.
Instead of the population gradient, we use the STE in practice, which replaces the delta function with the derivative of some differentiable function µSTE, to train the kernel w,
gSTEj (v,w,Z) ≡ ∑ i ∂` ∂y viµ ′ STE ∑ j′ Zij′wj′ Zij . (8) We consider three types of µSTE(x): (i) the identity type µSTE(x) = x , (ii) vanilla ReLU type µSTE(x) = xσ(x) and (iii) the cReLU type µSTE(x) = xσ(x)σ(r − x) with the upper clipping value r. The surrogate back-propagation using the STE is described in Algorithm 1 (Yin et al. (2019)).
Algorithm 1 Surrogate back-propagation using STE for learning one-hidden-layer CNN. 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 [ gSTE(vt,wt;Z) ] end for
The average of the STE gradients over the input data is also expressed as
EZ [ gSTEj (v,w;Z) ] = ∑ i ∫ dZ̃i1µ ′ STE ( Z̃i1|w| ) Fij(v,w, Z̃i1). (9)
Note that in general, the STE gradients have the non-zero component in the w-direction due to their partial collapse of the form of the gradient, ∑ j wjEZ [ gSTEj (v,w;Z) ] 6= 0.
It is obvious from Eq.(9) that the only difference between the STE gradients with the identity function, Relu, and cRelu is the region of the integration over Z̃i1,
EZ [ gAj (v,w;Z) ] = ∑ i ∫ RA dZ̃i1Fij(v,w, Z̃i1), (10)
where A = id,Relu, cRelu, and the integration reigion RA are given as Rid = (−∞,+∞), RRelu = [0,+∞), andRcRelu = [0, r|w| ].
Following the scale invariance of Fij(v,w, Z̃i1) for w, Fi′j(v, kw, Z̃i′1) = Fi′j(v,w, Z̃i′1) (k > 0), the STE gradients have the following interesting features:
1. The identity and Relu STE gradients have scale invariance for w, EZ [ g id/Relu j (v, kw;Z) ] = EZ [ g id/Relu j (v,w;Z) ] . (11)
2. The cRelu STE gradient does not have scale invariance for w due to the clipping effect, EZ [ gcReluj (v, kw, r;Z) ] 6= EZ [ gcReluj (v,w, r;Z) ] , (12)
where we have explicitly shown the dependence of the upper clipping value r. 3. Instead, the scale transformation of w for the cRelu STE gradient can be compensated by
that of r:
EZ [ gcReluj (v,w, kr;Z) ] = EZ [ gcReluj (v,w/k, r;Z) ] . (13)
If we take kr with fixed r large enough to cover most of the distribution P (Z), the left-hand side is identical with the Relu STE gradient. Here we define the spread of the distribution as ρ. When we redefine w/k as w on the right hand, the cRelu STE gradient for |w| < r/ρ is approximately identical with the Relu STE gradient,
EZ [ gcReluj (v,w, r;Z) ] ' EZ [ gReluj (v,w;Z) ] for |w| < r/ρ. (14)
This result is intuitively obvious: when w is small enough to keep most of the pre-activation values in the clipped range, the cRelu gradient and the Relu gradient are approximately identical.
The important issue is that there is no guarantee that the weights obtained by the STE backpropagation are the ones at the (local) minimum of the loss function. In other words, at the stationary points, (v,w) = (vs,ws), defined as the vanishing points of the population gradient,
EZ [ ∂`
∂vi (v = vs,w = ws)
] = 0, EZ [ ∂`
∂wj (v = vs,w = ws)
] = 0, (15)
the STE gradient is not zero in general. Here we consider the local minimum point that the STE gradient does not vanish,
EZ [ gSTEj (vs,ws;Z) ] 6= 0. (16)
Note that (v,w) = (vs, kws) for any k > 0 is also the local minimum point due to the scale invariance of the loss. In the case of the Relu/identity STEs, their gradient has the same finite value for any scales due to the scale invariance shown in Eq.(11),
EZ [ g id/Relu j (vs, kws;Z) ] = EZ [ g id/Relu j (vs,ws;Z) ] 6= 0. (17)
In the case of the cRelu STE, however, the value of its gradient changes as the scale changes due to the breaking of the scale symmetry shown in Eq.(12). If it crosses zero as a function of scale,
EZ [ gcReluj (vs, k0ws, r;Z) ] = 0 for some k0, (18)
the back propagation using cRelu STE can converge at the local minimum as was done by using the population gradient. Therefore, we conjecture that due to breaking scale symmetry, the backpropagation using cRelu STE is the most likely to achieve the (local) minimum of the loss function in the three STEs. This also implies that the cRelu STE is less biased than the other. We confirm the conjecture analytically for the Gaussian input with Relu-type teacher network in Sec.3, and numerically for the mixture Gaussian input with various teacher network in Sec. 4.
Since the above discussion is focused only on symmetry, the conjecture can be generalized to networks with more symmetries in a straightforward way. In that case, if we employ an STE that breaks all of those symmetries instead of cRelu STE, it is more likely to achieve the minimum.
The similar mechanism can be also found in physical phenomena. For instance, this is the case of the ferromagnet, which has rotational symmetry in energy (or Hamiltonian). In the system, if we impose the magnetic field, which breaks the rotational symmetry, the spin states, corresponding to minimum points, become aligned with the same direction.
3 ONE-HIDDEN LAYER CNN WITH BINARY ACTIVATION: GAUSSIAN INPUT AND LABELS GENERATED BY NON-QUANTIZED RELU NETWORK
We consider the simple model similar to Yin et al. (2019) given by the Gaussian input with the variance σ̂2,Z ∼ P (Z) = (
1√ 2πσ̂2
)mn e− ∑ i,j Z2ij 2σ̂2 . We use the squared loss function,
`(y(v,w,Z); y∗(v∗,w∗,Z)) = 12 (y − y ∗) 2 . Unlike in Yin et al. (2019), the true labels are assumed to be generated by the non-quantized Relu network,
y∗(v∗,w∗,Z) = ∑ i v∗i fRelu ∑ j Zijw ∗ j  , (19)
with the Relu activation defined as fRelu(x) = xσ(x).
Note that the models are misspecified, i.e., the population loss cannnot be zero even at the global minimum. In fact, it is inferred from the accuracy degradation in most practical situations (Hubara et al. (2017)) that the binary networks are considered as misspecified models. On the other hand, if the true labels are generated by the same binarized architecture as was used in Yin et al. (2019), the network becomes a specified model.
As shown below, this misspecified model leads to a more striking difference between three STEs than the specified model. Particularly, the behaviors of three STE gradients at the global minimum are completely different from each other.
3.1 POPULATION LOSS, ITS GRADIENT AND STATIONARY POINTS
The population loss L(v,w;v∗,w∗) is derived in Appendix A.1 as
L(v,w;v∗,w∗) = EZ [`(v,w,Z;v∗,w∗)]
= 1 8 vT ( Im×m + 1m1 T m ) v − σ̂|w ∗| 2 √ 2π vT ( cosϕIm×m + 1m1 T m ) v∗
+ σ̂2
4π |w∗|2v∗T
( (π − 1)Im×m + 1m1Tm ) v∗ (20)
where ϕ is the angle between w∗ and w. Note that the loss is scale invariant for w, and thus its derivative satisfies ew · ∂L∂w = 0. The stationary points are given by
∂L ∂v = 0⇔ v = 2σ̂|w ∗|√ 2π
( cosϕIm×m +
1− cosϕ m+ 1 1m1 T m
) v∗ (21)
∂L ∂w = 0⇔ vTv∗ = 0 or ϕ = 0 or ϕ = π (22)
where we used ( Im×m + 1m1 T m )−1 = Im×m − 1m+11m1 T m.
There are three stationary points 2 as classified in Appendix A.2:
1. middle saddle point given by v = v̄ ≡ 2σ̂|w ∗|√
2π
( cos ϕ̄Im×m + 1−cos ϕ̄ m+1 1m1 T m ) v∗, ϕ =
ϕ̄ ≡ arccos [
(v∗T 1m) 2
−(m+1)(v∗)2+(v∗T 1m)2
] , which satisfies v̄Tv∗ = 0. This exists if and only
if (m+ 1) (v∗)2 ≥ 2 ( v∗T1m )2 .
2. local minimum if (m + 1) (v∗)2 ≥ 2 ( v∗T1m )2 , otherwise saddle point given by v =
vπ ≡ 2σ̂|w ∗|√
2π
( −Im×m + 2m+11m1 T m ) v∗, ϕ = π.
3. global minimum given by v = v0 ≡ 2σ̂|w ∗|√
2π v∗, ϕ = 0. At the global minimum, the
population loss does not become zero, L(v = v0,w = w∗) = σ̂2|w∗|2(π−2) 4π v ∗Tv.
3.2 STE GRADIENT
The STE gradient is given by
gSTEj = ∑ i viµ ′ STE ∑ j′ Zij′wj′ Zij ×
∑ i′ vi′σ ∑ j′ Zi′j′wj′ −∑ i′ v∗i′fRelu ∑ j′ Zi′j′w ∗ j′  . (23) 2The definition of the stationary points is not mathematically rigorous in our paper. Even if points are non-differentiable, but the gradient gives zero at the points in the one-side limit, we call them the stationary points.
The converging points by the back-propagation in Algorithm 1 are characterized by vanishing the gradients,
EZ [ ∂`
∂vi (v,w;Z)
] = 0, EZ [ gSTEj (v,w;Z) ] = 0. (24)
We discuss the explicit form of the three STE gradients and the associated vanishing points. The detailed derivation of the STE gradients is shown in Appendix A.3-A.5.
3.2.1 IDENTITY STE
The identity STE gradient is given by
EZ [ gid ]
= σ̂√ 2π w |w| ( vTv ) − σ̂ 2 2 w∗ ( vTv∗ ) (25)
Combined with Eq.(21), we find that the equation in Eq.(24) has no solution. Therefore, if we use the back-propagation by identity STE, it can not converge to any points.
3.2.2 RELU STE
The Relu STE gradient is expressed as
EZ [grelu] = σ̂
2 √ 2π
w |w| [ vT ( Im×m + 1m1 T m ) v ] − σ̂ 2 2π |w∗| |w| w [ vT ( −Im×m + 1m1Tm ) v∗ ]
− σ̂2 (
(π − ϕ)w∗ + |w ∗| sinϕ |w| w
) ( vTv∗ ) 2π . (26)
Combined with Eq.(21), the gradients vanish at the following points:
1. If v∗ satisfies (m+ 1) (v∗)2 ≥ 2 ( v∗T1m )2 , ϕ = ϕ̄ = arccos
[ (v∗T 1m) 2
−(m+1)(v∗)2+(v∗T 1m)2
] ,
v = v̄ = 2σ̂|w ∗|√
2π
( cosϕIm×m + 1−cosϕ m+1 1m1 T m ) v∗.
2. ϕ = π,v = v̄ = 2σ̂|w ∗|√
2π
( −Im×m + 2m+11m1 T m ) .
We find the solutions are identical with two stationary points shown in item 1 and 2 in Sec.3.1, leading to the the saddle point or local minimum, while the global minimum point at ϕ = 0,v = v0 cannot be obtained. Therefore, by using the back-propagation with Relu STE, it always converges to ϕ = π if (m+ 1) (v∗)2 ≥ 2 ( v∗T1m )2 , and it does not converge to any points, otherwise.
3.2.3 CLIPPED RELU STE
The cRelu STE is given by
EZ [gcrelu] = σ̂
2 √ 2π
w
|w|
( 1− e− 1 2 ( r σ̂|w| ) 2) vT ( Im×m + 1m1 T m ) v
− σ̂ 2
2π |w∗| |w|
( 1− e− 1 2 ( r σ̂|w| ) 2) w [ vT ( −Im×m + 1m1Tm ) v∗ ]
− σ̂2 { C(w,ϕ) w∗
|w∗| + S(w,ϕ)
( 1
sinϕ
w |w| − cotϕ w
∗
|w∗|
)}( vTv∗ ) , (27)
where C(|w|, ϕ) and S(|w|, ϕ) are given in Eq.(106). At ϕ = 0, π they are simplified as C(|w|, 0) = |w
∗| 2π ( π − rσ̂|w|e − 12 ( r σ̂|w| ) 2√ 2π − πerfc ( r√ 2σ̂|w| )) , C(|w|, π) = S(|w|, π) =
S(|w|, 0) = 0. Remarkably, the cRelu STE gradient is proportional to the Relu STE gradient in the case of vTv∗ = 0 or ϕ = π:
EZ [gcrelu] |vT v∗=0 or ϕ=π = ( 1− e− 1 2 ( r σ̂|w| ) 2) EZ [grelu] |vT v∗=0 or ϕ=π. (28)
As shown in Sec.3.2.2, all the vanishing points in the Relu STE can be found at vTv∗ = 0 or ϕ = π, leading to the saddle point or the local minimum, so that they are also found in the vanishing points of cRelu STE. On the other hand, the behavior of cRelu STE gradient at ϕ = 0 is not related to that of Relu STE gradient. The cRelu STE gradient at ϕ = 0 is written as
EZ [gcrelu] = σ̂
2 √ 2π w∗ ( 1− e− 1 2 ( r σ|w| ) 2) vT { 1 |w∗| ( Im×m + 1m1 T m ) v
− 2σ̂√ 2π
( (−1 + λ) Im×m + 1m1Tm ) v∗ } , (29)
where
λ = λ(|w|) = π − rσ̂|w|e
− 12 ( r σ̂|w| ) 2√ 2π − πerfc (
r√ 2σ̂|w| ) 1− e− 1 2 ( r σ̂|w| ) 2 . (30)
This becomes zero if
v = 2σ̂√ 2π |w∗|
( Im×m + 1m1 T m )−1 ( (−1 + λ) Im×m + 1m1Tm ) v∗
= 2σ̂√ 2π |w∗|
( (−1 + λ) Im×m +
2− λ m+ 1 1m1 T m
) v∗. (31)
To be consistent with Eq.(21) at ϕ = 0, only λ = 2 is allowed. The solution provides the global minimum shown in item 3 in Sec.3.1.
In fact, λ monotonically increases from 0 to π as rσ̂|w| increases, and thus the solution at λ = 2, corresponding to the global minimum point, can be obtained for any clipping value r and the variance σ̂2 by changing the scale |w|. Consequently, we get all the stationary points are found in cRelu STE:
1. If v∗ satisfies (m+ 1) (v∗)2 ≥ 2 ( v∗T1m )2 , ϕ = ϕ̄ = arccos
[ (v∗T 1m) 2
−(m+1)(v∗)2+(v∗T 1m)2
] ,
v = v̄ = 2σ̂|w ∗|√
2π
( cosϕIm×m + 1−cosϕ m+1 1m1 T m ) v∗.
2. ϕ = π,v = v̄ = 2σ̂|w ∗|√
2π
( −Im×m + 2m+11m1 T m ) .
3. ϕ = 0, |w| = w0 ≡ rσ̂ c0,v = v0 ≡ 2σ̂|w∗|√ 2π v∗, where w0 (or c0) is given by λ(w0) = 2.
Note that while the global minimum point is degenerated in scales for w, i.e. if (w,v) = (w0,v0) gives the global minimum, (w,v) = (kw0,v0) for any k > 0 also gives the global minimum, the cRelu STE picks up the point at the scale determined by λ = 2.
4 EXPERIMENTS
To confirm our conjecture in more general case, we have numerically studied the Gaussian mixture input with various mean values. As teacher networks, we have tested tanh-type and sin-type as well as Relu-type. For all the examined setups, cRelu STE behaves like the population gradient, while id/Relu STEs show qualitatively different behaviors as described below. To calculate the population loss and STE gradients, we have generated the mixture Gaussian samples and have taken their average. The population gradient has been obtained by calculating the finite difference of the population loss. We have demonstrated the back propagation with learning rate η = 0.01 given in Algorithm 1.
Shown in Fig. 1 are the results of ten mixture Gaussian input with random mean values for each components of Z. We employed m = 20, n = 25, and the tanh-type teacher network. We find that the population gradient and cRelu STE show similar results, while id/Relu STEs are completely different. This indicates cRelu STE is less biased than id/Relu STEs.
In the case of id/Relu STEs, at early steps up to around 500 step, |w| decreases around the magnitude of the update quantity, so that it begins to oscillate. Then it escapes the oscillation, and the loss function shows the convergence to a point different from the local minimum achieved by population gradient, while |w| becomes larger and larger due to their scale invariance. Interestingly, the values of the loss function are small compared to the one obtained by population gradient, which implies that id/Relu STEs avoid being trapped in the local solution due to their large bias. However, note that even at that point, the magnitude of w continues to grow and eventually become numerically unstable.
5 SUMMARY
We have found that breaking symmetry embedded in the network by STEs enhances the possibility of convergence to the true (local) minimum of the loss function. We have demonstrated that if an STE breaks the scale symmetry embedded in the one-hidden-layer network with a binary activation, it is more likely to achieve the local minimum than the one which keeps the symmetry. The discussion can be generalized to the network with more symmetries in a straightforward way. The more symmetries embedded in the network an STE breaks, the more likely it is to converge. To confirm the mechanism, we have studied three STEs, identity STE, Relu STE, and cRelu STE, in a simple misspecified model with Gaussian input. We have found that the back-propagation by the cRelu STE, which breaks the scale symmetry, can converge to the global minimum, while identy/Relu STE cannot. Finally we have numerically confirmed the mechanism for the mixture Gaussian model with various teacher network.
REFERENCES Ron Banner, Yury Nahshan, Elad Hoffer, and Daniel Soudry. Post-training 4-bit quantization of
convolution networks for rapid-deployment. arXiv preprint arXiv:1810.05723, 2018.
Yoshua Bengio, Nicholas Léonard, and Aaron Courville. Estimating or propagating gradients through stochastic neurons for conditional computation. arXiv preprint arXiv:1308.3432, 2013.
Yash Bhalgat, Jinwon Lee, Markus Nagel, Tijmen Blankevoort, and Nojun Kwak. Lsq+: Improving low-bit quantization through learnable offsets and better initialization. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition Workshops, pp. 696–697, 2020.
Pengyu Cheng, Chang Liu, Chunyuan Li, Dinghan Shen, Ricardo Henao, and Lawrence Carin. Straight-through estimator as projected wasserstein gradient flow. arXiv preprint arXiv:1910.02176, 2019.
Jungwook Choi, Zhuo Wang, Swagath Venkataramani, Pierce I-Jen Chuang, Vijayalakshmi Srinivasan, and Kailash Gopalakrishnan. Pact: Parameterized clipping activation for quantized neural networks. arXiv preprint arXiv:1805.06085, 2018.
Yoni Choukroun, Eli Kravchik, Fan Yang, and Pavel Kisilev. Low-bit quantization of neural networks for efficient inference. In ICCV Workshops, pp. 3009–3018, 2019.
Simon Du, Jason Lee, Yuandong Tian, Aarti Singh, and Barnabas Poczos. Gradient descent learns one-hidden-layer cnn: Don’t be afraid of spurious local minima. In International Conference on Machine Learning, pp. 1339–1348. PMLR, 2018.
Steven K Esser, Jeffrey L McKinstry, Deepika Bablani, Rathinakumar Appuswamy, and Dharmendra S Modha. Learned step size quantization. arXiv preprint arXiv:1902.08153, 2019.
Geoffrey Hinton, Nitsh Srivastava, and Kevin Swersky. Neural networks for machine learning. Coursera, video lectures, 264(1):2146–2153, 2012.
Itay Hubara, Matthieu Courbariaux, Daniel Soudry, Ran El-Yaniv, and Yoshua Bengio. Quantized neural networks: Training neural networks with low precision weights and activations. The Journal of Machine Learning Research, 18(1):6869–6898, 2017.
Raghuraman Krishnamoorthi. Quantizing deep convolutional networks for efficient inference: A whitepaper. arXiv preprint arXiv:1806.08342, 2018.
Vladimir Kryzhanovskiy, Gleb Balitskiy, Nikolay Kozyrskiy, and Aleksandr Zuruev. Qpp: Realtime quantization parameter prediction for deep neural networks. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 10684–10692, 2021.
Daniel Kunin, Javier Sagastuy-Brena, Surya Ganguli, Daniel LK Yamins, and Hidenori Tanaka. Neural mechanics: Symmetry and broken conservation laws in deep learning dynamics. arXiv preprint arXiv:2012.04728, 2020.
Ziang Long, Penghang Yin, and Jack Xin. Learning quantized neural nets by coarse gradient method for nonlinear classification. Research in the Mathematical Sciences, 8(3):1–19, 2021.
Mohammad Rastegari, Vicente Ordonez, Joseph Redmon, and Ali Farhadi. Xnor-net: Imagenet classification using binary convolutional neural networks. In European conference on computer vision, pp. 525–542. Springer, 2016.
Alexander Shekhovtsov and Viktor Yanush. Reintroducing straight-through estimators as principled methods for stochastic binary networks. arXiv preprint arXiv:2006.06880, 2020.
Penghang Yin, Jiancheng Lyu, Shuai Zhang, Stanley Osher, Yingyong Qi, and Jack Xin. Understanding straight-through estimator in training activation quantized neural nets. arXiv preprint arXiv:1903.05662, 2019.
Ritchie Zhao, Yuwei Hu, Jordan Dotzel, Chris De Sa, and Zhiru Zhang. Improving neural network quantization without retraining using outlier channel splitting. In International conference on machine learning, pp. 7543–7552. PMLR, 2019.
Shuchang Zhou, Yuxin Wu, Zekun Ni, Xinyu Zhou, He Wen, and Yuheng Zou. Dorefa-net: Training low bitwidth convolutional neural networks with low bitwidth gradients. arXiv preprint arXiv:1606.06160, 2016.
A DERIVATION OF SEC.3
In the Appendix, we provide the results of the Gaussian with variance 1,N (0, 1). For some function f(Z), the expectation value of the Gaussian with the generic variance σ̂,N (0, σ̂2), shown in Sec.3 is written as
EZ∼N (0,σ2)[f(Z)] =
√ 1
2πσ̂2 ∫ ∏ i,j dZijf(Z)e − ∑ i,j Z 2 ij 2σ̂2
=
√ 1
2π ∫ ∏ i,j dZ ′ijf(σ̂Z ′)e− ∑ i,j Z ′2 ij 2
= EZ∼N (0,1)[f(σ̂Z)] (32)
where we have changed the integral variables from Z to Z ′ = Z/σ. Therefore we get the results with the variance σ̂ by only changing from f(Z) to f(σ̂Z).
A.1 DERIVATION OF LOSS FUNCTION
The expectaion of the loss is divided into the following three terms:
L(v,w;v∗,w∗) = 1 2 EZ [ vTσ(Zw)vTσ(Zw) ] − EZ [ vTσ(Zw)v∗T fRelu(Zw ∗) ]
+ 1 2 EZ [ v∗T fRelu(Zw ∗)v∗T fRelu(Zw ∗) ]
(33)
The calculation of the first term and the third term are given in “Proof of Lemma 1” in Yin et al. (2019) and “Proof of Section 3” in Ref. Du et al. (2018), respectively:
1 2 EZ [ vTσ(Zw)vTσ(Zw) ] + 1 2 EZ [ v∗T fRelu(Zw ∗)v∗T fRelu(Zw ∗) ] = 1
8 vT ( Im×m + 1m1 T m ) v + 1 4π (w∗)2v∗T ( (π − 1)Im×m + 1m1Tm ) v∗ (34)
To discuss the second term, we evaluate the following quantity:
Fij = EZ [σ(Zw)fRelu(Zw∗)]ij
= ( 1√ 2π )mn ∫ ∏ k,l dZkl  e− 12 ∑k,l Z2klσ(∑ m′ Zim′wm′ )(∑ m′ Zjm′w ∗ m′ ) σ (∑ m′ Zjm′w ∗ m′ ) (35)
(i) i 6= j case:
Fij = ( 1√ 2π )2n ∫ (∏ l dZildZjl ) e− 1 2 ∑ l(Z 2 il+Z 2 jl)σ (∑ m′ Zim′wm′ )(∑ m′ Zjm′w ∗ m′ ) σ (∑ m′ Zjm′w ∗ m′ )
= ( 1√ 2π )2n ∫ (∏ l dZ̃ildZ̃jl ) e− 1 2 ∑ l(Z̃ 2 il+Z̃ 2 jl)σ ( Z̃i1|w| )( Z̃j1|w∗| ) σ ( Z̃j1|w∗| ) = ( 1√ 2π )2 ∫ dZ̃i1dZ̃j1e − 12 (Z̃ 2 i1+Z̃ 2 j1)σ ( Z̃i1 )( Z̃j1|w∗| ) σ ( Z̃j1
) = 1
2 √ 2π |w∗|, (36)
where we used the orthogonal transformation in the second line: Z̃il = ∑ m′ Oilm′Zim′ ,
Z̃jl = ∑ m′ Ojlm′Zjm′ (37)
with Oi1m′ = wm′/|w| and O j 1m′ = w ∗ m′/|w∗|.
(ii) i = j case:
Fii = ( 1√ 2π )n ∫ ∏ l dZile − 12 ∑ l Z 2 ilσ (∑ m Zimwm )(∑ m Zimw ∗ m ) σ (∑ m Zimw ∗ m ) . (38)
Without loss of generality, we can choose
w∗ = (w∗,0n−1),
w = (w1, w2,0n−2) (39)
with w∗ > 0, so that this can be written as
Fii = ( 1√ 2π )2 ∫ dZi1dZi2e − 12 (Z 2 i1+Z 2 i2)σ ( 2∑
m=1
Zimwm ) Zi1w ∗σ (Zi1)
= ( 1√ 2π )2 ∫ dZdθZe− 1 2Z 2 σ (cos(ϕ− θ))Z cos θw∗σ (cos θ)
= 1
2π ∫ ∞ 0 dZ ∫ π/2 ϕ−π/2 dθZ2e− 1 2Z 2 cos θw∗
= 1
2 √ 2π w∗(1 + cosϕ) (40)
where ϕ is the angle between w∗ and w. Eqs.(36) and (40) are combined as
Fij = 1
2 √ 2π |w∗|(1 + cosϕ)δij +
1 2 √ 2π |w∗|(1− δij), (41)
so that the second term in Eq. (33) is expressed as
EZ [ vTσ(Zw)v∗T fRelu(Zw ∗) ] = |w∗|
2 √ 2π vT ( cosϕIm×m + 1m1 T m ) v∗. (42)
Consequently, the loss function is summarized as
L(v,w;v∗,w∗) = 1 8 vT ( Im×m + 1m1 T m ) v − |w ∗| 2 √ 2π vT ( cosϕIm×m + 1m1 T m ) v∗
+ 1
4π |w∗|2v∗T
( (π − 1)Im×m + 1m1Tm ) v∗. (43)
The population gradient w.r.t v and w is thus given by
∂L ∂v = 1 4
( Im×m + 1m1 T m ) v − σ̂|w
∗| 2 √ 2π
( cosϕIm×m + 1m1 T m ) v∗, (44)
∂L ∂w = 1 |w| ∂L ∂ϕ eϕ = σ̂|w∗| 2 √ 2π|w| sinϕ
( vTv∗ ) eϕ, (45)
where without loss of generality, we can take ew = (cosϕ, sinϕ,0n−2), eϕ = (− sinϕ, cosϕ,0n−2). Note that ew · ∂L∂w = 0 is satisfied due to the scale symmetry for the loss.
A.2 CLASSIFICATION OF STATIONARY POINTS
Using Eq.(21), the loss at the three stationary points is rewritten as
L = −|w ∗|2
4π v∗T
( cos2 ϕIm×m + ( − (cosϕ− 1)2
m+ 1 + 1
) 1m1 T m ) v∗
+ |w∗|2
2π v∗T
( (π − 1) Im×m + 1m1Tm ) v∗. (46)
Obviously, this shows that ϕ = 0 is the global minimum point,
L ≥ L|ϕ=0 = |w∗|2
4π (π − 2)v∗Tv. (47)
At ϕ = π, this becomes
L|ϕ=π = − |w∗|2
4π v∗T
( Im×m + ( −4
m+ 1 + 1
) 1m1 T m ) v∗
+ |w∗|2
4π v∗T
( (π − 1) Im×m + 1m1Tm ) v∗
= |w∗|2
4π v∗T
( (π − 2) Im×m + 4
m+ 1 1m1
T m
) v∗ (48)
At ϕ = ϕ̄ = arccos [
(v∗T 1m) 2
−(m+1)(v∗)2+(v∗T 1m)2
] , this becomes
L|ϕ=ϕ̄ = − |w∗|2
4π v∗T { (v∗T1m)2 −(m+ 1) (v∗)2 + (v∗T1m)2 }2 Im×m
+ − (m+ 1){ (v∗)2 −(m+ 1) (v∗)2 + (v∗T1m)2 }2 + 1 1m1Tm v∗
+ |w∗|2
4π v∗T
( (π − 1) Im×m + 1m1Tm ) v∗
= −|w ∗|2
4π
{ ( v∗T1m )2 (v∗) 2
−(m+ 1) (v∗)2 + (v∗T1m)2 + ( v∗T1m
)2}
+ |w∗|2
4π v∗T
( (π − 1) Im×m + 1m1Tm ) v∗. (49)
Using these results, we obtain
L|ϕ=ϕ̄ − L|ϕ=π
= −|w ∗|2
4π
( v∗T1m )2 (v∗) 2
−(m+ 1) (v∗)2 + (v∗T1m)2 + |w∗|2 4π (v∗) 2 − |w ∗|2 4π
( v∗T1m )2 4 m+ 1
= |w∗|2
4π
{ (m+ 1) (v∗) 2 − 2 ( v∗T1m )2}2{ (m+ 1) (v∗) 2 − (v∗T1m)2 } (m+ 1) . (50)
Note that because (m + 1) (v∗)2 ≥ 2 ( v∗T1m )2 is satisfied to exist the stationary point at ϕ = ϕ̄, the loss at ϕ = π is smaller than the one at at ϕ = ϕ̄,
L|v=v̄,ϕ=ϕ̄ ≥ L|v=vπ,ϕ=π. (51) is satisfied. When (m + 1) (v∗)2 = 2 ( v∗T1m )2 , L|ϕ=ϕ̄ = L|ϕ=π, because the middle stationary point at ϕ = ϕ̄ is merged into the point at ϕ = π.
This implies that ϕ = 0 is the global minimum, ϕ = π is the local minimum, and ϕ = ϕ̄ is the saddle point or the local maximum if (m+ 1) (v∗)2 ≥ 2 ( v∗T1m )2 .
To fully clarify whether the stationary points are local minimum, maximum or saddle point, we have calculated the hessian matrix,
H =
( ∂2L ∂v2
∂2L ∂v∂ϕ
∂2L ∂ϕ∂v
∂2L ∂ϕ2
)
=
( 1 4 ( Im×m + 1m1 T m ) |w∗| 2 √ 2π sinϕv∗
|w∗| 2 √ 2π sinϕv∗T |w ∗| 2 √ 2π vTv∗ cosϕ
) . (52)
(i) v = v̄, ϕ = ϕ̄:
For z = (x, y)T (x ∈ Rm, y ∈ R),
zTHz = 1 4 xT ( Im×m + 1m1 T m ) x + |w∗|√ 2π y sin ϕ̄ ( xTv∗ ) = 1
4
( xT1m )2 +
( 1
2 x + |w∗|y sin ϕ̄√ 2π
v∗ )2 − |w ∗|2y2 sin2 ϕ̄ (v∗)2
2π , (53)
where we used v̄Tv∗ = 0. Therefore, depending on x and y, the Hessian can be either positive or negative, which means the stationary point is a saddle point.
(ii) v = vπ ≡ 2|w ∗|√
2π
( −Im×m + 2m+11m1 T m ) v∗, ϕ = π
H =
( 1 4 ( Im×m + 1m1 T m ) 0
0 − |w ∗|
2 √ 2π vTπ v
∗
) . (54)
Note that vTπ v ∗ = 2|w ∗|√ 2π v∗T ( −Im×m + 2m+11m1 T m ) v∗. If (m + 1) (v∗)2 ≥ 2 ( v∗T1m )2 , then vTπ v ∗ ≤ 0, so that the matrix becomes a positive definite matrix: For any x ∈ Rm, y ∈ R, this satisfies
1 4 xT ( Im×m + 1m1 T m ) x− |w ∗| 2 √ 2π vTπ v ∗y2 ≥ 0, (55)
which means the stationary point is the local minimum. On the other hand, if (m + 1) (v∗)2 ≤ 2 ( v∗T1m )2 , the matrix becomes either positive or negative, depending on x and y, so that this stationary point becomes a saddle point.
(iii) v = v0 ≡ 2|w ∗|√
2π v∗, ϕ = 0
H =
( 1 4 ( Im×m + 1m1 T m ) 0
0 |w ∗|
2 √ 2π vT0 v
∗
) . (56)
The Hessian becomes a positive definite matrix: For any x ∈ Rm, y ∈ R, this satisfies
1 4 xT ( Im×m + 1m1 T m ) x + |w∗| 2 √ 2π vT0 v ∗y2 ≥ 0. (57)
A.3 DERIVATION OF IDENTITY STE GRADIENT
By substituting µ′(x) = 1 into Eq.(23), we obtain
EZ [g]i = (
1√ 2π )mn ∫ ∏ k,l dZkl  e− 12 ∑k,l Z2kl ∑ j Zjivj  ×
(∑ k vkσ (∑ l Zklwl ) − ∑ k v∗kfRelu (∑ l Zklw ∗ l )) . (58)
We calculate each term as follows.
A.3.1 CALCULATION OF THE FIRST TERM
We calculate the first term as was done in Lemma 9 (and also Lemma 11) in Yin et al. (2019).
To discuss the first term, we evaluate
F idijk ≡ EZ [Zji [σ(Zw)]k] . (59)
Using F idijk, the first term can be expressed as
first term = ∑ j,k F idijkvjvk. (60)
(i) j 6= k case:
F idijk = 0, (61)
because of the odd symmetry for i-th and j-th components in the integrand.
(ii) j = k case:
F idijj = ( 1√ 2π )n ∫ (∏ l dZjl ) e− 1 2 ∑ l Z 2 jlZjiσ (∑ l Zjlwl )
= ( 1√ 2π )n ∫ (∏ l dZ̃jl ) e− 1 2 ∑ l Z̃ 2 jl (∑ m′ Oj Tim′Z̃jm′ ) σ ( Z̃j1 ) = ( 1√ 2π )n ∫ (∏ l dZ̃jl ) e− 1 2 ∑ l Z̃ 2 jlOj Ti1 Z̃j1σ ( Z̃j1
) = ( 1√ 2π )∫ dZ̃j1e − 12 Z̃ 2 j1Oj Ti1 Z̃j1σ ( Z̃j1
) = ( 1√ 2π ) Oj Ti1 = ( 1√ 2π ) Oj1i
= 1√ 2π wi |w|
(62)
where we used the orthogonal transformation in the second line:
Z̃jl = ∑ m′ Ojlm′Zjm′ (63)
with Oi1m′ = wm′/|w|. The first term is reduced to
first term = 1√ 2π wi |w|
( vTv ) . (64)
A.3.2 CALCULATION OF THE SECOND TERM
To discuss the second term, we evaluate
Gidijk = EZ [Zji [fRelu(Zw∗)]k] . (65)
Using Gidijk, the second term can be expressed as
second term = ∑ j,k Gidijkvjv ∗ k. (66)
(i) j 6= k case:
Gidijk = 0, (67)
because of the odd symmetry.
(ii) j = k case:
Gidijj = ( 1√ 2π )n ∫ (∏ l dZjl ) e− 1 2 ∑ l Z 2 jlZjifRelu (∑ l Zjlw ∗ l )
= ( 1√ 2π )n ∫ (∏ l dZ̃jl ) e− 1 2 ∑ l Z̃ 2 jl (∑ m′ Oj Tim′Z̃jm′ ) fRelu ( |w∗|Z̃j1 ) = ( 1√ 2π )n ∫ (∏ l dZ̃jl ) e− 1 2 ∑ l Z̃ 2 jlOj Ti1 Z̃j1fRelu ( |w∗|Z̃j1
) = ( 1√ 2π )∫ dZ̃j1e − 12 Z̃ 2 j1Oj Ti1 Z̃j1fRelu ( |w∗|Z̃j1
) = ( 1√ 2π )∫ dZ̃j1e − 12 Z̃ 2 j1Oj Ti1 Z̃ 2 j1|w∗|σ ( |w∗|Z̃j1
) = |w∗|
2 Oj Ti1 = |w∗| 2 Oj1i
= 1
2 w∗i (68)
where we have used the orthogonal transformation in the second line: Z̃jl = ∑ m′ Ojlm′Zjm′ (69)
with Oi1m′ = w∗m′/|w∗|, and also have used∫ ∞ 0 dxx2e−(1/2)x 2 = √ 2π 2 . (70)
The second term is written as
second term = 1 2 wi ( vTv∗ ) . (71)
A.4 DERIVATION OF RELU STE GRADIENT
By substituting µ′(x) = σ(x) into Eq.(23), we obtain
EZ [grelu]i = (
1√ 2π )mn ∫ ∏ k,l dZkl  e− 12 ∑k,l Z2kl ∑ j Zjivjσ (∑ l Zjlwl ) ×
(∑ k vkσ (∑ l Zklwl ) − ∑ k v∗kfRelu (∑ l Zklw ∗ l )) . (72)
We calculate each term as follows.
A.4.1 CALCULATION OF THE FIRST TERM
We calculate the first term as was done in Lemma9 (and also Lemma11) in Yin et al. (2019). We evaluate
F reluijk = EZ [ Zji [σ(Zw)]j [σ(Zw)]k ] . (73)
Using F reluijk , the first term can be expressed as first term = ∑ j,k F reluijk vjvk. (74)
(i) j 6= k case:
F reluijk = ( 1√ 2π )2n ∫ (∏ l dZjldZkl ) e− 1 2 ∑ l Z 2 jl+Z 2 klZjiσ (∑ l Zjlwl ) σ (∑ l Zklwl )
= ( 1√ 2π )2n ∫ (∏ l dZ̃jldZ̃kl ) e− 1 2 ∑ l Z̃ 2 jl+Z̃ 2 kl (∑ m′ Oj Tim′Z̃jm′ ) σ ( Z̃j1 ) σ ( Z̃k1 ) = ( 1√ 2π )2 ∫ dZ̃j1dZ̃k1e − 12 (Z̃ 2 j1+Z̃ 2 k1)Oj Ti1 Z̃j1σ ( Z̃j1 ) σ ( Z̃k1
) = 1
2 √ 2π Oj Ti1
= 1
2 √ 2π wi |w| , (75)
where we used the orthogonal transformation in the second line: Z̃jl = ∑ m′ Ojlm′Zjm′ (76)
with Oi1m′ = wm′/|w|. (ii) j = k case:
We get the same expression given in Eq.(62),
F reluijj = ( 1√ 2π )n ∫ (∏ l dZjl ) e− 1 2 ∑ l Z 2 jlZjiσ (∑ l Zjlwl )
= 1√ 2π wi |w| . (77)
The first term is written as
first term = 1
2 √ 2π wi |w|
vT ( Im×m + 1m1 T m ) v. (78)
A.4.2 CALCULATION OF SECOND TERM
To discuss the second term, we evaluate Greluijk = EZ [ Zji [ [σ(Zw)]j fRelu(Zw ∗) ] k ] . (79)
Using Greluijk , the second term can be expressed as second term = ∑ j,k Greluijk vjv ∗ k. (80)
(i) j 6= k case:
Greluijk = ( 1√ 2π )2n ∫ (∏ l dZjldZkl ) e− 1 2 ∑ l Z 2 jl+Z 2 klZjiσ (∑ l Zjlwl ) fRelu (∑ l Zklw ∗ l )
= ( 1√ 2π )2n ∫ (∏ l dZ̃jldZ̃kl ) e− 1 2 ∑ l Z̃ 2 jl+Z̃ 2 kl (∑ m′ Oj Tim′Z̃jm′ ) σ ( Z̃j1 ) fRelu ( |w∗|Z̃k1 ) = ( 1√ 2π )2 ∫ dZ̃j1dZ̃k1e − 12 (Z̃ 2 j1+Z̃ 2 k1)Oj Ti1 Z̃j1σ ( Z̃j1 ) |w∗|Z̃k1σ ( Z̃k1
) = |w∗| 2π Oj Ti1
= 1
2π |w∗| |w| wi, (81)
where we used the orthogonal transformation in the second line: Z̃jl = ∑ m′ Ojlm′Zjm′ ,
Z̃kl = ∑ m′ Oklm′Zkm′ (82)
with Oi1m′ = wm′/|w| and O j 1m′ = w ∗ m′/|w∗|.
(ii) j = k case:
Greluijj = ( 1√ 2π )m ∫ (∏ l dZjl ) e− 1 2 ∑ l Z 2 jlZjiσ (∑ l Zjlwl ) fRelu (∑ l Zjlw ∗ l )
= ( 1√ 2π )m ∫ (∏ l dZjl ) e− 1 2 ∑ l Z 2 jlZjiσ ( 2∑ l=1 Zjlwl ) fRelu (w ∗Zj1)
= 1
2π ∫ ( 2∏ l=1 dZjl ) e− 1 2 ∑2 l=1 Z 2 jlZji (δi1 + δi2)σ ( 2∑ l=1 Zjlwl ) fRelu (w ∗Zj1)
= 1
2π
∫ dZdθZe− 1 2Z 2 (Z cos θδi1 + Z sin θδi2)σ (Zw cos (θ − ϕ)) fRelu (w∗Z cos θ)
= 1
2π
∫ dZdθZe− 1 2Z 2 (Z cos θδi1 + Z sin θδi2)σ (Zw cos (θ − ϕ))w∗Z cos θσ (cos θ)
(83)
= w∗
2π ∫ ∞ 0 dZZ3e− 1 2Z 2 ∫ π/2 −π/2+ϕ dθ ( cos2 θδi1 + cos θ sin θδi2 ) = w∗
π
(( sin 2ϕ
4 +
1 2 (π − ϕ)
) δi1 +
1− cos 2ϕ 4 δi2 ) = w∗
π
( 1
2 (π − ϕ) δi1 +
sinϕ
2 (cosϕδi1 + sinϕδi2)
) (84)
where we have chosen w∗ = (w∗,0n−1),
w = (w1, w2,0n−2) (85) with w∗ > 0 and ϕ is defined as the angle between w∗ and w. We also have used the formula,∫ ∞
0
dxx3e−(1/2)x 2
= 2,∫ dθ cos2 θ = sin 2θ
4 +
1
2 θ,∫
dθ cos θ sin θ = −1 4 cos 2θ. (86)
Note that w∗i |w∗| = δi1,
wi |w| = cosϕδi1 + sinϕδi2. (87)
This reduces to
Greluijj = 1
π
( 1
2 (π − ϕ)w∗i +
sinϕ
2 |w∗| |w| wi
) (88)
In summary, the second term is written as
second term = 1
2π |w∗| |w| wi [ vT ( −Im×m + 1m1Tm ) v∗ ]
+ ( (π − ϕ)w∗i + |w∗| sinϕ |w| wi ) ( vTv∗ ) 2π . (89)
A.5 DERIVATION OF CLIPPED RELU STE GRADIENT
By substituting µ′(x) = σ(x)σ(r − x) into Eq.(23), we obtain
EZ [gcrelu]i = (
1√ 2π )mn ∫ ∏ k,l dZkl  e− 12 ∑k,l Z2kl ×
∑ j Zjivjσ (∑ l Zjlwl ) σ ( r − ∑ l Zjlwl ) ×
(∑ k vkσ (∑ l Zklwl ) − ∑ k v∗kfRelu (∑ l Zklw ∗ l )) . (90)
We calculate each term as follows.
A.5.1 CALCULATION OF THE FIRST TERM
To discuss the first term, we evaluate F creluijk = EZ [ Zji [σ(Zw) σ(r −Zw)]j [σ(Zw)]k ] . (91)
Using F creluijk , the first term can be expressed as
first term = ∑ j,k F creluijk vjvk. (92)
(i) j 6= k case:
F creluijk = ( 1√ 2π )2n ∫ (∏ l dZjldZkl ) e− 1 2 ∑ l Z 2 jl+Z 2 klZji
× σ (∑ l Zjlwl ) σ ( r − ∑ l Zjlwl ) σ (∑ l Zklwl )
= ( 1√ 2π )2n ∫ (∏ l dZ̃jldZ̃kl ) e− 1 2 ∑ l Z̃ 2 jl+Z̃ 2 kl (∑ m′ Oj Tim′Z̃jm′ )
× σ ( Z̃j1 ) σ ( r − |w|Z̃j1 ) σ ( Z̃k1 ) = ( 1√ 2π )2 ∫ dZ̃j1dZ̃k1e − 12 (Z̃ 2 j1+Z̃ 2 k1)Oj Ti1 Z̃j1σ ( Z̃j1 ) σ ( r − |w|Z̃j1 ) σ ( Z̃k1
) = 1
2 √ 2π ∫ r/|w| 0 dZ̃j1e − 12 Z̃ 2 j1Oj Ti1 Z̃j1
= 1
2 √ 2π Oj Ti1
( 1− e− 1 2 ( r |w| ) 2) = 1
2 √ 2π
( 1− e− 1 2 ( r |w| ) 2) wi |w| , (93)
where we used the orthogonal transformation in the second line: Z̃jl = ∑ m′ Ojlm′Zjm′ (94)
with Oi1m′ = wm′/|w|.
(ii) j = k case:
F creluijj = ( 1√ 2π )n ∫ (∏ l dZjl ) e− 1 2 ∑ l Z 2 jlZjiσ (∑ l Zjlwl ) σ ( r − ∑ l Zjlwl )
= ( 1√ 2π )n ∫ (∏ l dZ̃jl ) e− 1 2 ∑ l Z̃ 2 jl (∑ m′ Oj Tim′Z̃jm′ ) σ ( Z̃j1 ) σ ( r − |w|Z̃j1 ) =
1√ 2π
∫ dZ̃j1e − 12 Z̃ 2 j1Oj Ti1 Z̃j1σ ( Z̃j1 ) σ ( r − |w|Z̃j1 ) =
1√ 2π ∫ r/|w| 0 dZ̃j1e − 12 Z̃ 2 j1Oj Ti1 Z̃j1
= 1√ 2π Oj Ti1
( 1− e− 1 2 ( r |w| ) 2) =
1√ 2π
( 1− e− 1 2 ( r |w| ) 2) wi |w| , (95)
In summary, the first term is written as
first term = 1
2 √ 2π wi |w|
( 1− e− 1 2 ( r |w| ) 2) vT ( Im×m + 1m1 T m ) v. (96)
A.5.2 CALCULATION OF THE SECOND TERM
To discuss the second term, we evaluate Gcreluijk = EZ [ Zji [ [σ(Zw) σ(r −Zw)]j fRelu(Zw ∗) ] k ] . (97) Using Gcreluijk , the second term can be expressed as
second term = ∑ j,k Gcreluijk vjv ∗ k. (98)
(i) j 6= k case:
Gcreluijk = ( 1√ 2π )2n ∫ (∏ l dZjldZkl ) e− 1 2 ∑ l Z 2 jl+Z 2 klZji
× σ (∑ l Zjlwl ) σ ( r − ∑ l Zjlwl ) fRelu (∑ l Zklw ∗ l )
= ( 1√ 2π )2n ∫ (∏ l dZ̃jldZ̃kl ) e− 1 2 ∑ l Z̃ 2 jl+Z̃ 2 kl (∑ m′ Oj Tim′Z̃jm′ )
× σ ( Z̃j1 ) σ ( r − |w|Z̃j1 ) fRelu ( |w∗|Z̃k1 ) = ( 1√ 2π )2 ∫ dZ̃j1dZ̃k1e − 12 (Z̃ 2 j1+Z̃ 2 k1)Oj Ti1 Z̃j1
× σ ( Z̃j1 ) σ ( r − |w|Z̃j1 ) |w∗|Z̃k1σ ( Z̃k1 ) = ( 1√ 2π )2 ∫ dZ̃j1e − 12 Z̃ 2 j1Oj Ti1 Z̃j1σ ( Z̃j1 ) σ ( r − |w|Z̃j1 ) |w∗|
= ( 1√ 2π )2 ∫ r/|w| 0 dZ̃j1e − 12 Z̃ 2 j1Oj Ti1 Z̃j1|w ∗| = |w∗| 2π Oj Ti1 ( 1− e− 1 2 ( r |w| )
2) = |w∗| 2π ( 1− e− 1 2 ( r |w| ) 2) wi |w| , (99)
where we used the orthogonal transformation in the second line: Z̃il = ∑ m′ Oilm′Zim′ ,
Z̃kl = ∑ m′ Oklm′Zkm′ (100)
with Oi1m′ = wm′/|w| and O j 1m′ = w ∗ m′/|w∗|.
(ii) j = k case:
Gcreluijj = ( 1√ 2π )m ∫ (∏ l dZjl ) e− 1 2 ∑ l Z 2 jlZji
× σ (∑ l Zjlwl ) σ ( r − ∑ l Zjlwl ) fRelu (∑ l Zjlw ∗ l )
= ( 1√ 2π )m ∫ (∏ l dZjl ) e− 1 2 ∑ l Z 2 jlZji
× σ ( 2∑ l=1 Zjlwl ) σ ( r − 2∑ l=1 Zjlwl ) fRelu (w ∗Zj1)
= 1
2π ∫ ( 2∏ l=1 dZjl ) e− 1 2 ∑2 l=1 Z 2 jlZji (δi1 + δi2)
× σ ( 2∑ l=1 Zjlwl ) σ ( r − 2∑ l=1 Zjlwl ) fRelu (w ∗Zj1)
= 1
2π
∫ dZdθZe− 1 2Z 2
(Z cos θδi1 + Z sin θδi2)
× σ (Zw cos (θ − ϕ))σ ( r −
2∑ l=1 Zjlwl
) fRelu (w ∗Z cos θ)
= 1
2π
∫ dZdθZe− 1 2Z 2
(Z cos θδi1 + Z sin θδi2)
× σ (Zw cos (θ − ϕ))σ (r − Zw cos (θ − ϕ))w∗Z cos θσ (cos θ) (101)
= w∗
2π ∫ π/2 −π/2+ϕ dθ ( cos2 θδi1 + cos θ sin θδi2 ) ∫ rw cos(θ−ϕ) 0 dZZ3e− 1 2Z 2
= w∗
2π ∫ π/2 −π/2+ϕ dθ ( cos2 θδi1 + cos θ sin θδi2 ) × { 2− (( r
w cos (θ − ϕ)
)2 + 2 ) e− 1 2 ( r w cos(θ−ϕ) ) 2 } (102)
where we have chosen w∗ = (w∗,0n−1),
w = (w1, w2,0n−2) (103) with w∗ > 0 and ϕ is defined as the angle between w∗ and w. We also have used the formula,∫ t
0
dxx3e− 1 2x 2 = 2− (t2 + 2)e− 12 t 2 . (104)
Note that w∗i |w∗| = δi1,
wi |w| = cosϕδi1 + sinϕδi2. (105)
Using C(w,ϕ), S(w,ϕ) defined as
C(w,ϕ) ≡ w ∗
2π ∫ π/2 −π/2+ϕ dθ cos2 θ { 2− (( r w cos (θ − ϕ) )2 + 2 ) e− 1 2 ( r w cos(θ−ϕ) ) 2 } ,
S(w,ϕ) ≡ w ∗
2π ∫ π/2 −π/2+ϕ dθ sin θ cos θ { 2− (( r w cos (θ − ϕ) )2 + 2 ) e− 1 2 ( r w cos(θ−ϕ) ) 2 } ,
(106)
If ϕ 6= 0, π, Gcreluijj is formally written as
Gcreluijj = C(w,ϕ) w∗i |w∗| + S(w,ϕ)
( 1
sinϕ wi |w| − cotϕ w
∗ i
|w∗|
) (107)
If ϕ = 0 or π, Gcreluijj is formally written as
Gcreluijj = C(w, 0) w∗i |w∗| for ϕ = 0, (108)
Gcreluijj = 0 for ϕ = π (109)
where we used S(w, 0) = S(w, π) = C(w, π) = 0 and
C(w, 0) = w∗
2π ∫ π/2 −π/2 dθ cos2 θ { 2− (( r w cos θ )2 + 2 ) e− 1 2 ( r w cos θ ) 2 }
= w∗
2π
( π − r
w e− 1 2 ( r w )
2√ 2π − πerfc ( r√ 2w )) (110)
with the following formula∫ π/2 −π/2 dθe− a2 2 cos2 θ = πerfc ( a√ 2 ) , (111)∫ π/2
−π/2 dθ cos2 θe−
a2 2 cos2 θ = 1
2
( ae− a2 2 √ 2π + ( 1− a2 ) πerfc ( a√ 2 )) . (112)
In summary, the second term is written as
second term = 1
2π |w∗| |w|
( 1− e− 1 2 ( r |w| ) 2) wi [ vT ( −Im×m + 1m1Tm ) v∗ ]
+ { C(w,ϕ)
w∗i |w∗| + S(w,ϕ)
( 1
sinϕ wi |w| − cotϕ w
∗ i
|w∗|
)}( vTv∗ ) . (113)

Review summaries:
Training quantized neural network by the biased Straight-Through estimator (STE) is a common practice in the NN quantization literature. However, due to the biased nature of STE, there is no principled guidance in making algorithmic choices for such training methods. This work focuses on the choice of surrogate differentiable functions in STE, and looked at three common choices: identity, ReLU, and clipped ReLU. Through analysis of the scale invariance of gradients in STE, the paper suggests clipped ReLU might pick up local minima degenerated in scales, due to its property of breaking scale symmetry among the three. The authors further analyzed the stationary points of a simple misspecified model with Gaussian inputs to confirm their observations.

The paper does a theoretical analysis of various gradient approximations for the non differentiable sign function. They investigate 3 gradient estimators based on the STE assumption, namely (identity) STE, ReLU STE, and clipped STE (cReLU). They show theoretically that a gradient estimator which breaks the scale symmetry (clipped STE) in the network is more likely to achieve a local minimum than one that keeps the symmetry.

The paper considers a network with one hidden layer with binary valued neurons sharing the same weight vector but having non-overlapping receptive fields. The authors analyse the stationary points of a particular MSE loss when a straight through estimator (STE) is used for the loss gradient. They analyse three variants for the STE: (1) derivative of the identity function, (2) derivative of ReLU and (3) derivative of a piecewise linear squashing function. Assuming a simple independent distribution over the inputs and a particular loss, the authors show that the first two variants are biased and only the third STE variant can be unbiased. The authors conjecture that the reason for this is that the third STE variant breaks the scaling symmetry w.r.t. the weights of the first layer.

The paper investigates analytically the scale invariance properties of  different straight-through estimators used in training neural networks with binary activation and weights.   The main claim is that the cReLU STE, among the 3 considered, is the one that is more likely to give zero gradients in the minima of the population loss. This is (weakly) supported by theoretical analysis on a little variation of the shallow setting analyzed in Yin et al. (2019), where the teacher now contains ReLU activations.