Rotated Binary Neural Network

Given a CNN model, we denote $\mathbf{w}^{i}\in\mathbb{R}^{n^{i}}$ and $\mathbf{a}^{i}\in\mathbb{R}^{m^{i}}$ as its weights and feature maps in the $i$-th layer, where $n^{i}=c_{out}^{i}\cdot c_{in}^{i}\cdot w^{i}_{f}\cdot h^{i}_{f}$ and $m^{i}=c_{out}^{i}\cdot w^{i}_{a}\cdot h^{i}_{a}$. ($c^{i}_{out}$, $c^{i}_{in}$) represent the number of output and input channels, respectively. ($w_{f}^{i}$, $h_{f}^{i}$) and ($w_{a}^{i}$, $h_{a}^{i}$) are the width and height of filters and feature maps, respectively. We then have

$$\mathbf{a}^{i}=\mathbf{w}^{i}\circledast\mathbf{a}^{i-1},$$

where $\circledast$ is the standard convolution and we omit activation layers for simplicity. The BNN aims to convert $\mathbf{w}^{i}$ and $\mathbf{a}^{i}$ into $\mathbf{b}_{w}^{i}\in\{-1,+1\}^{n^{i}}$ and $\mathbf{b}_{a}^{i}\in\{-1,+1\}^{m^{i}}$, such that the convolution can be achieved by using the efficient xnor and bitcount logics. Following [23, 5], we binarize the activations with the sign function:

$$\mathbf{b}^{i}_{a}=\text{sign}(\mathbf{b}_{w}^{i}\odot\mathbf{b}^{i-1}_{a}),$$

where $\odot$ represents the xnor and bitcount logics, and sign($\cdot$) denotes the sign function which returns $1$ if the input is larger than zero, and $-1$ otherwise. Similar to [23, 45, 5], $\mathbf{b}_{w}^{i}$ can also be obtained by $\mathbf{b}_{w}^{i}=\text{sign}(\mathbf{w}^{i})$ and a scaling factor can be applied to compensate for the norm difference.

However, the existence of angular bias between $\mathbf{w}^{i}$ and $\mathbf{b}_{w}^{i}$ could lead to large quantization error as analyzed in Sec. 1. Besides, it results in consistent signs between $\mathbf{w}^{i}$ and $\mathbf{b}_{w}^{i}$ which lessens the information gain (see footnote 1). We aim to minimize this angular bias to reduce the quantization error, meanwhile increasing the probability of weight flips to increase the information gain.

As shown in Fig. 3, we consider applying a rotation matrix $\mathbf{R}^{i}\in\mathbb{R}^{n^{i}\times n^{i}}$ to $\mathbf{w}^{i}$ at the beginning of each training epoch, such that the angle ${\varphi}^{i}$ between the rotated weight vector ${(\mathbf{R}^{i})^{T}}\mathbf{w}^{i}$ and its binary vector $\text{sign}\big{(}{(\mathbf{R}^{i})^{T}}\mathbf{w}^{i}\big{)}$ should be minimized. To this end, we derive the following formulation:

$$\cos({\varphi}^{i})=\frac{\text{sign}\big{(}{(\mathbf{R}^{i})^{T}}\mathbf{w}^{i}\big{)}^{T}\big{(}{(\mathbf{R}^{i})^{T}}\mathbf{w}^{i}\big{)}}{\|\text{sign}\big{(}{(\mathbf{R}^{i})^{T}}\mathbf{w}^{i}\big{)}\|_{2}\|{(\mathbf{R}^{i})^{T}}\mathbf{w}^{i}\|_{2}},\quad s.t.\quad{(\mathbf{R}^{i})^{T}}\mathbf{R}^{i}=\mathbf{I}_{n^{i}},$$

(3)

where $\mathbf{I}_{n^{i}}\in\mathbb{R}^{n^{i}\times n^{i}}$ is an $n^{i}$-th order identity matrix. It is easy to know that $\|\text{sign}\big{(}{(\mathbf{R}^{i})^{T}}\mathbf{w}^{i}\big{)}\|_{2}=\sqrt{n^{i}}$, $\|{(\mathbf{R}^{i})^{T}}\mathbf{w}^{i}\|_{2}$ = $\|\mathbf{w}^{i}\|_{2}$, both of which are constant²²2To stress, the rotation is applied at the beginning of each training epoch instead of the training stage. Thus, $\|\mathbf{w}^{i}\|_{2}$ should be regarded as a constant.. Thus, Eq. (3) can be further simplified as

$$\begin{split}\cos({\varphi}^{i})&={\eta}^{i}\cdot{\text{sign}\big{(}{(\mathbf{R}^{i})^{T}}\mathbf{w}^{i}\big{)}^{T}}\big{(}{(\mathbf{R}^{i})^{T}}\mathbf{w}^{i}\big{)}={\eta}^{i}\cdot\text{tr}\big{(}\mathbf{b}_{w^{\prime}}^{i}(\mathbf{w}^{i})^{T}\mathbf{R}^{i}\big{)},\;s.t.\;{(\mathbf{R}^{i})^{T}}\mathbf{R}^{i}=\mathbf{I}_{n^{i}},\end{split}$$

(4)

where ${\eta}^{i}=1/\big{(}{\|\text{sign}\big{(}{(\mathbf{R}^{i})^{T}}\mathbf{w}^{i}\big{)}\|_{2}\|{(\mathbf{R}^{i})^{T}}\mathbf{w}^{i}\|_{2}}\big{)}=1/\big{(}{\sqrt{n^{i}}\|\mathbf{w}^{i}\|_{2}}\big{)}$, $\text{tr}(\cdot)$ returns the trace of the input matrix and $\mathbf{b}_{w^{\prime}}^{i}=\text{sign}\big{(}{(\mathbf{R}^{i})^{T}}\mathbf{w}^{i}\big{)}$. However, Eq. (4) involves a large rotation matrix, $n^{i}$ of which can be up to millions in a neural network. Direct optimization of $\mathbf{R}^{i}$ would consume massive memory and computation. Besides, performing a large rotation leads to $\mathcal{O}\big{(}(n^{i})^{2}\big{)}$ complexity in both space and time.

To deal with this, inspired by the properties of Kronecker product [27], we introduce a bi-rotation scenario where two smaller rotation matrices $\mathbf{R}_{1}^{i}\in\mathbb{R}^{n^{i}_{1}\times n^{i}_{1}}$ and $\mathbf{R}_{2}^{i}\in\mathbb{R}^{n^{i}_{2}\times n^{i}_{2}}$ are used to reconstruct the large rotation matrix $\mathbf{R}^{i}\in\mathbb{R}^{n^{i}\times n^{i}}$ with $n^{i}=n_{1}^{i}\cdot n_{2}^{i}$. One of the basic property of Kronecker product [27] is that if two matrices $\mathbf{R}_{1}^{i}\in\mathbb{R}^{n^{i}_{1}\times n^{i}_{1}}$ and $\mathbf{R}_{2}^{i}\in\mathbb{R}^{n^{i}_{2}\times n^{i}_{2}}$ are orthogonal, then $\mathbf{R}_{1}^{i}\otimes\mathbf{R}_{2}^{i}\in\mathbb{R}^{n_{1}^{i}n_{2}^{i}\times n_{1}^{i}n_{2}^{i}}$ is orthogonal as well, where $\otimes$ denotes the Kronecker product. Another basic property of Kronecker product comes at:

$$(\mathbf{w}^{i})^{T}{(\mathbf{R}_{1}^{i}\otimes\mathbf{R}_{2}^{i})}=\text{Vec}\big{(}{(\mathbf{R}_{2}^{i})^{T}}(\mathbf{W}^{i})^{T}{\mathbf{R}_{1}^{i}}\big{)},$$

(5)

where Vec($\cdot$) vectorizes its input and Vec($\mathbf{W}^{i}$) = $\mathbf{w}^{i}$. Thus, we can see that applying the bi-rotation to $\mathbf{W}^{i}$ is equivalent to applying an $\mathbf{R}^{i}=\mathbf{R}_{1}^{i}\otimes\mathbf{R}_{2}^{i}\in\mathbb{R}^{n_{1}^{i}n_{2}^{i}\times n_{1}^{i}n_{2}^{i}}$ rotation to $\mathbf{w}^{i}$. Learning two smaller matrices $\mathbf{R}_{1}^{i}$ and $\mathbf{R}_{2}^{i}$ can well reconstruct the large rotation matrix $\mathbf{R}^{i}$. Moreover, performing the bi-rotation consumes only $\mathcal{O}\big{(}(n_{1}^{i})^{2}+(n_{2}^{i})^{2}\big{)}$ space complexity and $\mathcal{O}\big{(}({n_{1}^{i}})^{2}n_{2}^{i}+n^{i}_{1}({n_{2}^{i}})^{2}\big{)}$ time complexity, respectively, leading to a significant complexity reduction compared to the large rotation³³3With $n^{i}_{1}\cdot n^{i}_{2}=n^{i}$, our RBNN achieves the least complexity when $n_{1}^{i}=n_{2}^{i}=\sqrt{n^{i}}$, which is also our experimental setting.. Accordingly, Eq. (4) can be reformulated as

$$\begin{split}\cos({\varphi}^{i})={\eta}^{i}\cdot\text{tr}&\Big{(}\mathbf{b}_{w^{\prime}}^{i}\text{Vec}\big{(}{(\mathbf{R}_{2}^{i})^{T}}(\mathbf{W}^{i})^{T}{\mathbf{R}_{1}^{i}}\big{)}\Big{)}={\eta}^{i}\cdot\text{tr}\big{(}\mathbf{B}_{W^{\prime}}^{i}{(\mathbf{R}_{2}^{i})^{T}}{(\mathbf{W}^{i})^{T}}\mathbf{R}^{i}_{1}\big{)},\\ &s.t.\qquad{(\mathbf{R}_{1}^{i})^{T}}\mathbf{R}_{1}^{i}=\mathbf{I}_{n_{1}^{i}},\quad{(\mathbf{R}_{2}^{i})^{T}}\mathbf{R}_{2}^{i}=\mathbf{I}_{n_{2}^{i}},\end{split}$$

(6)

where $\mathbf{B}^{i}_{W^{\prime}}=\text{sign}\big{(}{(\mathbf{R}_{1}^{i})^{T}}\mathbf{W}^{i}{\mathbf{R}_{2}^{i}}\big{)}$. Finally, we rewrite our optimization objective below:

$$\begin{split}&\mathop{\arg\max}_{\mathbf{B}_{W^{\prime}}^{i},\mathbf{R}^{i}_{1},\mathbf{R}^{i}_{2}}\text{tr}\big{(}\mathbf{B}_{W^{\prime}}^{i}{(\mathbf{R}_{2}^{i})^{T}}{(\mathbf{W}^{i})^{T}}\mathbf{R}^{i}_{1}\big{)},\\ s.t.\;\;\;\mathbf{B}_{W^{\prime}}^{i}\in\{-1,&+1\}^{n^{i}_{1}\times n^{i}_{2}},\;({\mathbf{R}_{1}^{i}})^{T}\mathbf{R}_{1}^{i}=\mathbf{I}_{n_{1}^{i}},\;({\mathbf{R}_{2}^{i}})^{T}\mathbf{R}_{2}^{i}=\mathbf{I}_{n_{2}^{i}}.\end{split}$$

(7)

Eq. (7) is non-convex w.r.t. $\mathbf{B}_{W^{\prime}}^{i}$, $\mathbf{R}^{i}_{1}$ and $\mathbf{R}^{i}_{2}$. To find a feasible solution, we adopt an alternating optimization approach, i.e., updating one variable with the rest two fixed until convergence.

1) $\mathbf{B}_{W^{\prime}}^{i}$-step: Fix $\mathbf{R}_{1}^{i}$ and $\mathbf{R}_{2}^{i}$, then learn the binarization $\mathbf{B}_{W^{\prime}}^{i}$. The sub-problem of Eq. (7) becomes:

$$\begin{split}\mathop{\arg\max}_{\mathbf{B}_{W^{\prime}}^{i}}\text{tr}\big{(}\mathbf{B}_{W^{\prime}}^{i}{(\mathbf{R}_{2}^{i})^{T}}{(\mathbf{W}^{i})^{T}}\mathbf{R}^{i}_{1}\big{)},\quad s.t.\quad\mathbf{B}_{W^{\prime}}^{i}\in\{-1,+1\}^{n^{i}_{1}\times n^{i}_{2}},\end{split}$$

(8)

which can be achieved by $\mathbf{B}_{W^{\prime}}^{i}=\text{sign}\big{(}{(\mathbf{R}^{i}_{1})^{T}}\mathbf{W}^{i}\mathbf{R}_{2}^{i}\big{)}$.

2) $\mathbf{R}_{1}^{i}$-step: Fix $\mathbf{B}_{W^{\prime}}^{i}$ and $\mathbf{R}_{2}^{i}$, then update $\mathbf{R}_{1}^{i}$. The corresponding sub-problem is:

$$\begin{split}\mathop{\arg\max}_{\mathbf{R}_{1}^{i}}\text{tr}\big{(}\mathbf{G}^{i}_{1}\mathbf{R}^{i}_{1}\big{)},\quad s.t.\quad({\mathbf{R}_{1}^{i}})^{T}\mathbf{R}_{1}^{i}=\mathbf{I}_{n_{1}^{i}},\end{split}$$

(9)

where $\mathbf{G}_{1}^{i}=\mathbf{B}_{W^{\prime}}^{i}{(\mathbf{R}_{2}^{i})^{T}}{(\mathbf{W}^{i})^{T}}$. The above maximum can be achieved by using the polar decomposition [34]: $\mathbf{R}_{1}^{i}=\mathbf{V}_{1}^{i}(\mathbf{U}_{1}^{i})^{T}$, where $\mathbf{G}_{1}^{i}=\mathbf{U}_{1}^{i}\mathbf{S}_{1}^{i}(\mathbf{V}_{1}^{i})^{T}$ is the SVD of $\mathbf{G}_{1}^{i}$.

3) $\mathbf{R}_{2}^{i}$-step: Fix $\mathbf{B}_{W^{\prime}}^{i}$ and $\mathbf{R}_{1}^{i}$, then update $\mathbf{R}_{2}^{i}$. The corresponding sub-problem becomes

$$\begin{split}\mathop{\arg\max}_{\mathbf{R}_{2}^{i}}\text{tr}\big{(}{(\mathbf{R}_{2}^{i})^{T}}\mathbf{G}_{2}^{i}\big{)},\quad s.t.\quad({\mathbf{R}_{2}^{i}})^{T}\mathbf{R}_{2}^{i}=\mathbf{I}_{n_{2}^{i}},\end{split}$$

(10)

where $\mathbf{G}_{2}^{i}={(\mathbf{W}^{i})^{T}}\mathbf{R}^{i}_{1}\mathbf{B}_{W^{\prime}}^{i}$. Similar to the updating rule for $\mathbf{R}_{1}^{i}$, the updating rule for $\mathbf{R}_{2}^{i}$ is $\mathbf{R}_{2}^{i}=\mathbf{U}_{2}^{i}(\mathbf{V}_{2}^{i})^{T}$, where $\mathbf{G}_{2}^{i}=\mathbf{U}_{2}^{i}\mathbf{S}_{2}^{i}(\mathbf{V}_{2}^{i})^{T}$ is the SVD of $\mathbf{G}_{2}^{i}$.

In the experiments, we iteratively update $\mathbf{B}_{W^{\prime}}^{i}$, $\mathbf{R}^{i}_{1}$ and $\mathbf{R}^{i}_{2}$, which can reach convergence after three cycles of updating. Therefore, the weight rotation can be efficiently implemented.

We narrow the angular bias between the full-precision weights and the binarization using our bi-rotation at the beginning of each training epoch. Then, we can set $\tilde{\mathbf{w}}^{i}=(\mathbf{R}^{i})^{T}\mathbf{w}^{i}$, which will be fed to the sign function and follow up the standard gradient update in the neural network. However, the alternating optimization may get trapped in a local optimum that either overshoots (Fig. 4(a)) or undershoots (Fig. 4(b)) the binarization $\text{sign}\big{(}{(\mathbf{R}^{i})^{T}}\mathbf{w}^{i}\big{)}$. To deal with this, we further propose to self-adjust the rotated weight vector as below:

$$\tilde{\mathbf{w}}^{i}=\mathbf{w}^{i}+\big{(}(\mathbf{R}^{i})^{T}\mathbf{w}^{i}-\mathbf{w}^{i}\big{)}\cdot{\alpha}^{i},$$

(11)

where ${\alpha}^{i}=\text{abs}\big{(}\sin({\beta}^{i})\big{)}\in[0,1]$ and ${\beta}^{i}\in\mathbb{R}$.

As can be seen from Fig. 4, Eq. (11) constrains that the final weight vector moves along the residual direction of $(\mathbf{R}^{i})^{T}\mathbf{w}^{i}-\mathbf{w}^{i}$ with ${\alpha}^{i}\geq 0$. It is intuitive that when overshooting, ${\alpha}^{i}\leq 1$; when undershooting, ${\alpha}^{i}\geq 1$. We empirically observe that overshooting is in a dominant position. Thus, we simply constrain ${\alpha}^{i}\in[0,1]$ to shrink the feasible region of ${\alpha}^{i}$, which we find can well further reduce the quantization error and boost the performance as demonstrated in Table 4. The final value of ${\alpha}^{i}$ varies across different layers. In Fig. 4(c), we show a toy example of how ${\alpha}^{i}$ updates during training in ResNet-20 (layer2.2.conv2).

At the beginning of each training epoch, with fixed $\mathbf{w}^{i}$, we learn the rotation matrix $\mathbf{R}^{i}$ ($\mathbf{R}^{i}_{1}$ and $\mathbf{R}^{i}_{2}$ actually). In the training stage, with fixed $\mathbf{R}^{i}$, we feed the sign function using Eq. (11) in the forward, and update $\mathbf{w}^{i}$ and ${\beta}^{i}$ in the backward.

The derivative of the sign function is almost zero everywhere, which makes the training unstable and degrades the accuracy performance. To solve it, various gradient approximations in the literature have been proposed to enable the gradient updating, e.g., straight through estimation [3], piecewise polynomial function [32], annealing hyperbolic tangent function [1], error decay estimator [37] and so on. Instead of simply using these existing approximations, in this paper, we further devise the following training-aware approximation function to replace the sign function:

$$\begin{split}F(x)=&\begin{cases}k\cdot\big{(}-\text{sign}(x)\frac{t^{2}x^{2}}{2}+\sqrt{2}tx\big{)},&|x|<\frac{\sqrt{2}}{t},\\ \text{sign}(x)\cdot k,&\mbox{else}.\end{cases}\end{split}$$

(12)

with

$$t={10}^{T_{\min}+\frac{e}{E}\cdot(T_{\max}-T_{\min})},\quad k=\max(\frac{1}{t},1),$$

where $T_{\min}=-2$, $T_{\max}=1$ in our implementation, $E$ is the number of training epochs and $e$ represents the current epoch. As can be seen, the shape of $F(x)$ relies on the value of $\frac{e}{E}$, which indicates the training progress. Then, the gradient of $F(x)$ w.r.t. the input $x$ can be obtained by

$$F^{\prime}(x)=\frac{\partial F(x)}{\partial x}=\max\big{(}k\cdot(\sqrt{2}t-|t^{2}x|),0\big{)}.$$

(13)

In Fig. 5, we visualize Eq. (12) and Eq. (13) w.r.t. varying values of $\frac{e}{E}$. In the early training stage, the gradient exists almost everywhere, which overcomes the drawback of sign function and enables the updating of the whole network. As the training proceeds, our approximation gradually becomes the sign-like function, which ensures the binary property. Thus, our approximation is training-aware. So far, we can have the gradient of loss function $\mathcal{L}$ w.r.t. the activation $\mathbf{a}^{i}$ and weight $\mathbf{w}^{i}$:

$${\sigma}_{\mathbf{a}^{i}}=\frac{\partial\mathcal{L}}{\partial F(\mathbf{a}^{i})}\cdot\frac{\partial F(\mathbf{a}^{i})}{\partial\mathbf{a}^{i}},\quad{\sigma}_{\mathbf{w}^{i}}=\frac{\partial\mathcal{L}}{\partial F(\tilde{\mathbf{w}}^{i})}\cdot\frac{\partial F(\tilde{\mathbf{w}}^{i})}{\partial\tilde{\mathbf{w}}^{i}}\cdot\frac{\partial\tilde{\mathbf{w}}^{i}}{\partial\mathbf{w}^{i}},$$

(14)

where

$$\frac{\partial\tilde{\mathbf{w}}^{i}}{\partial\mathbf{w}^{i}}=(1-{\alpha}^{i})\cdot\mathbf{I}_{n^{i}}+{\alpha}^{i}\cdot(\mathbf{R}^{i})^{T}.$$

(15)

Besides, the gradient of ${\alpha}^{i}$ in Eq. (11) can be obtained by

$${\sigma}_{{\alpha}^{i}}=\frac{\partial\mathcal{L}}{\partial{\alpha}^{i}}=\frac{\partial\mathcal{L}}{\partial\tilde{\mathbf{w}}^{i}}\frac{\partial\tilde{\mathbf{w}}^{i}}{\partial{\alpha}^{i}}=\mathop{\sum}_{j}\Big{[}{\sigma}_{\tilde{\mathbf{w}}^{i}}\big{(}(\mathbf{R}^{i})^{T}\mathbf{w}^{i}-\mathbf{w}^{i}\big{)}\Big{]}_{j},$$

(16)

where $[\cdot]_{j}$ denotes the $j$-th element of its input vector. Note that the error decay estimator in [37] can be also regarded as a training-aware approximation. However, our design in Eq. (12) is fundamentally different from that in [37] and its superiority is validated in Sec. 4.2.

In this section, we first show the benefit of our training-aware approximation over other recent advances [3, 32, 37]. And then, we show the effect of different components proposed in our RBNN. All the experiments are conducted on top of ResNet-20 with Bi-Real structure on CIFAR-10.

Table 4 compares the performances of RBNN based on various gradient approximations including straight through estimation (denoted by STE) [3], piecewise polynomial function (denoted by PPF) [32] and error decay estimator (denoted by EDE) [37]. As can be seen, STE shows the least accuracy. Though EDE also studies approximation with dynamic changes, its performance is even worse than PPF which is a fixed approximation. In contrast, our training-aware approximation achieves 1.8% improvements over EDE, which validates the effectiveness of our approximation.

To further understand the effect of each component in our RBNN, we conduct an ablation study by starting with the binarization using XNOR-Net [38] (denoted by B), and then gradually add different parts of training-aware approximation (denoted by T), weight rotation (denoted by R) and adjustable scheme (denoted by A). As shown in Table 4, the binarization using XNOR-Net suffers a great performance degradation of 8.0% compared with the full-precision model. By adding our weight rotation or training-aware approximation, the accuracy performance increases to 86.4% or 86.6%. Then, the collective effort of weight rotation and training-aware approximation further raises it to 87.1%. Lastly, by considering the adjustable weight vector in the training process, our RBNN achieves the highest accuracy of 87.8%. Therefore, each part of RBNN plays its unique role in improving the performance.

Fig. 6 shows the histograms of weights (before binarization) for XNOR-Net and our RBNN. It can be seen that the weight values for XNOR-Net are mixed up tightly around zero center and the value magnitude remains far less than 1. Thus it causes large quantization error when being pushed to the binary values of $-1$ and $+1$. On the contrary, our RBNN results in two-mode distributions, each of which is centered around $-1/+1$. Besides, there exist few weights around the zero, which creates a clear boundary between the two distributions. Thus, by the weight rotation, our RBNN effectively reduces quantization error as explained in Fig. 2.

As discussed in Sec. 2, the capacity of learning BNN is up to $2^{n}$ where $n$ is the total number of weight elements. When the probability of each element being $-1$ or $+1$ is equal during training, it reaches the maximum of $2^{n}$. We compare the initialization weights and binary weights, and then show the weight flipping rates of our RBNN and XNOR-Net across the layers of ResNet-20 in Fig. 7. As can be observed, XNOR-Net leads to a small flipping rate, i.e., most positive weights are directly quantized to $+1$, and vice versa. Differently, RBNN leads to around 50% weight flips each layer due to the introduced weight rotation as illustrated in Fig. 3, which thus maximizes the information gain during the training stage.