B I BERT: A CCURATE F ULLY B INARIZED BERT
Anonymous authors
Paper under double-blind review
A BSTRACT
The large pre-trained BERT has achieved remarkable performance on Natural
Language Processing (NLP) tasks but is also computation and memory expensive.
As one of the powerful compression approaches, binarization extremely reduces the computation and memory consumption by utilizing 1-bit parameters and bitwise operations.
Unfortunately, the full binarization of BERT ( i.e. , 1-bit weight, embedding, and activation) usually suffer a signiﬁcant performance drop, and there is rare study addressing this problem.
In this paper, with the theoretical justiﬁcation and empirical analysis, we identify that the severe performance drop can be mainly attributed to the information degradation and optimization direction mismatch respectively in the forward and backward propagation, and propose BiBERT , an accurate fully binarized BERT, to eliminate the performance bottlenecks.
Speciﬁcally, BiBERT introduces an efﬁcient Bi-Attention structure for maximizing representation information statistically and a DirectionMatching Distillation (DMD) scheme to optimize the full binarized BERT accurately.
Extensive experiments show that BiBERT outperforms both the straightforward baseline and existing state-of-the-art quantized BERTs with ultra-low bit activations by convincing margins on the NLP benchmark.
As the ﬁrst fully binarized BERT, our method yields impressive 56 .
3 × and 31 .
2 × saving on FLOPs and model size, demonstrating the vast advantages and potential of the fully binarized
BERT model in real-world resource-constrained scenarios.
I NTRODUCTION
[Figure]
FLOPs & size.
Our BiBERT enjoys most computation and storage savings while surpassing SOTA quantized
BERTs on GLUE benchmark with ultra-low bit activation.
Recently, the pre-trained language models have shown great power in various natural language processing (NLP) tasks (Wang et al., 2018a;
Qin et al., 2019; Rajpurkar et al., 2016).
In particular,
BERT (Devlin et al., 2018) signiﬁcantly improves the state-ofthe-art performance, while the massive parameters hinder their widespread deployment on edge devices in the real world.
Therefore, model compression has been actively studied to alleviate resource constraint issues, including quantization (Shen et al., 2020;
Zafrir et al., 2019), distillation (Jiao et al., 2020; Xu et al., 2020), pruning (McCarley et al., 2019; Gordon et al., 2020), parameter sharing (Lan et al., 2020), etc.
Among them, quantization emerges as an efﬁcient way to obtain the compact model by compressing the model parameters to lower bit-width representation, such as
Q-BERT (Shen et al., 2020), Q8BERT (Zafrir et al., 2019), and
GOBO (Zadeh et al., 2020).
However, the representation limitation and optimization difﬁculties come as a consequence of applying the discrete quantization, triggering severe performance drop in quantized BERT.
Fortunately, distillation becomes a common remedy in quantization as an auxiliary optimization approach to tackle the performance drop, which encourages the quantized BERT to mimic the full-precision model to exploit knowledge in the teacher’s representation (Bai et al., 2020; Zhang et al., 2020).
As the quantization scheme with the most aggressive bit-width (Zhou et al., 2016; Wang et al., 2018b; Xu et al., 2021), full binarization of BERT ( i.e. , 1-bit weight, embedding, and activation) further allows the model to utilize extremely compact 1-bit parameters and efﬁcient bitwise operations, which can largely promote the applications of BERT on edge devices in the real world (Wang
MHA
FFN
Add & Norm × 𝐿 Classifier
Embeddings
Add & Norm
Fully Binarized Transformer Layer ෍ 𝐅∈{𝐐,𝐊,𝐕} 𝐏 𝐅 −𝐏 𝐅𝑇 𝐏 𝐐 𝐏 𝐊 𝐏 𝐕
Direction-matching Distillation (DMD) 𝐁 𝐀 ⊠ bool(·) head 1 ,⋯,head 𝑁 𝐻 attentionscore 𝐀 𝐁 𝐐 𝐁 𝐊 ⨂ 𝐁 𝐕 query 𝐐 key 𝐊 value 𝐕
Bi-Attention 𝐏 𝐐𝑇 𝐏 𝐊𝑇 𝐏 𝐕𝑇
BAMM hiddenstates 𝐇
MHA
FFN
Add & Norm Add & Norm
Full-precision Transformer Layer (the teacher)
× 𝐿 et al., 2021; Liu et al., 2018).
However, although the BERT quantization equipped with distillation has been studied, it is still a signiﬁcant challenge to binarize BERT to extremely 1-bit, especially for its activation.
Compared with the weight and word embedding in BERT, the binarization of activation brings the most severe performance drop, and the model even crashes in NLP tasks.
So far, previous studies have pushed down the weight and embedding to be binarized (Bai et al., 2020), but none of them have ever achieved to binarize BERT with 1-bit activation accurately.
Therefore, we ﬁrst build a fully binarized BERT baseline, a straightforward yet effective solution based on common techniques.
Our study ﬁnds that the performance drop of BERT with binarized 1-bit weight, activation, and embedding (called fully binarized BERT) comes from the information degradation of the attention mechanism in the forward propagation and the optimization direction mismatch of distillation in the backward propagation.
First, the attention mechanism makes
BERT focus selectively on parts of the input and ignore the irrelevant content (Vaswani et al., 2017;
Chorowski et al., 2015; Wu et al., 2016).
While our analysis shows that direct binarization leads to the almost complete degradation of the information of attention weight (Figure 4), which results in the invalidation of the selection ability for attention mechanism.
Second, the distillation for the fully binarized BERT baseline utilizes the attention score, the direct binding product of two binarized activations.
However, we show that it causes severe optimization direction mismatch since the non-neglectable error between the defacto and expected optimization direction (Figure 5).
This paper provides empirical observations and theoretical formulations of the above-mentioned phenomena, and proposes a BiBERT to turn the full-precision BERT into the strong fully binarized model (see the overview in Figure 2).
To tackle the information degradation of the attention mechanism, we introduce an efﬁcient Bi-Attention structure based on the information theory.
Bi-Attention applies the binarized representation with maximized information entropy, allowing the binarized model to restore the perception of input contents.
Moreover, we developed the Direction-Matching
Distillation (DMD) scheme to eliminate the direction mismatch in the distillation.
DMD takes appropriate activation and utilizes knowledge from constructed similarity matrices in distillation to optimize accurately.
Our BiBERT, for the ﬁrst time, presents a promising route towards the accurate fully binarized BERT (with 1-bit weight, embedding, and activation).
The extensive experiments on the GLUE (Wang et al., 2018a) benchmark show that our BiBERT outperforms existing quantized BERT models with ultra-lower bit activation by convincing margins.
For example, the average accuracy of BiBERT exceeds 1-1-1 bit-width BinaryBERT ( 1 -bit weight, 1 -bit embedding and 1 -bit quantization) by 20.4% accuracy on average, and even better than 2-8-8 bit-width Q2BERT by 13.3%.
Besides, we highlight that our BiBERT gives impressive 56 .
3 × and 31 .
2 × saving on FLOPs and model size, respectively, which shows the vast advantages and potential of the fully binarized BERT model in terms of fast inference and ﬂexible deployment in real-world resource-constrained scenarios (Figure 1).
B UILDING A F ULLY B INARIZED BERT B ASELINE
First of all, we build a baseline to study the fully binarized BERT since it has never been proposed in previous works.
A straightforward solution is to binarize the representation in BERT architecture in the forward propagation and apply distillation to the optimization in the backward propagation.
2. B INARIZED BERT A RCHITECTURE
We give a brief introduction to the architecture of binarized BERT.
In general, the forward and backward propagation of sign function in binarized network can be formulated as:
Forward
[Equation]
where C is the cost function for the minibatch.
sign function is applied in the forward propagation while the straight-through estimator (STE) (Bengio et al., 2013) is used to obtain the derivative in the backward propagation.
As for the weight of binarized linear layers, the common practice is to redistribute the weight to zero-mean for retaining representation information (Rastegari et al., 2016;
Qin et al., 2020) and applies scaling factors to minimize quantization errors (Rastegari et al., 2016).
The activation is binarized by the sign without re-scaling for computational efﬁciency.
Thus, the computation can be expressed as
[Equation]
where W and X denote full-precision weight and activation, µ ( · ) denotes the mean value, α w is the scaling factors for weight, and ⊗ denotes the matrix multiplication with bitwise xnor and bitcount as presented in Appendix A.4.
The input data ﬁrst passes through a binarized embedding layer before being fed into the transformer blocks (Zhang et al., 2020; Bai et al., 2020).
And each transformer block consists of two main components: Multi-Head Attention (MHA) module and Feed-Forward Network (FFN).
The computation of MHA depends on queries Q , keys K and values V , which are derived from hidden states H ∈ R N × D .
N represents the length of the sequence, and D represents the dimension of features.
For a speciﬁc transformer layer, the computation in an attention head can be expressed as
[Equation]
where bi-linear Q , bi-linear K , bi-linear V represent three different binarized linear layers for
Q , K , V respectively.
Then we compute the attention score A as follow:
[Equation]
where B Q and B K are the binarized query and key, respectively.
Note that the obtained attention weight is then truncated by attention mask, and each row in A can be regarded as a k -dim vector, where k is the number of unmasked elements.
Then we binarize the attention weights B s A as
[Equation]
We follow original BERT architecture to carry on the rest of MHA and FFN in the binarized network.
2. D ISTILLATION FOR B INARIZED BERT
Distillation is a common and essential optimization approach to alleviate the performance drop of quantized BERT under ultra-low bit-width settings, which can be unobstructedly applied for any architectures to utilize the knowledge of a full-precision teacher model (Jiao et al., 2020; Bai et al., 2020; Zhang et al., 2020;
Wang et al., 2020).
The usual practice is to distill the attention score A Tl ,
MHA output M Tl , and hidden states H Tl in a layerwise manner from the full-precision teacher network, and transfer to the binarized student counterparts, i.e. , A l , M l , H l
( l = 1 , ..., L , where L represents the number of transformer layers), respectively.
We use the mean squared errors (MSE) as loss function to measure the difference between student and teacher networks for corresponding features:
[Equation]
Then the prediction-layer distillation loss is conducted by minimizing the soft cross-entropy (SCE) between teacher logits y T and student logits y .
The objective function is expressed as
[Equation]
[Figure]
T HE R ISE OF B I BERT
Although we can build a fully binarized BERT baseline and its training pipeline with common techniques, the performance is still of major concern.
Our study shows that the baseline suffers an immense information degradation of attention mechanism in the forward propagation and severe optimization direction mismatch in the backward propagation.
To solve these problems, in this section, we propose the BiBERT with theoretical and experimental justiﬁcations.
3. B OTTLENECKS OF F ULLY B INARIZED BERT B ASELINE
Intuitively, in the fully binarized BERT baseline, the information representation capability largely depends on structures of architecture, such as attention, which is severely limited due to the binarized parameters, and the discrete binarization also makes the optimization more difﬁcult.
This means the bottlenecks of the fully binarized BERT baseline come from architecture and optimization for the forward and backward propagation, respectively.
From the architecture perspective , we observe the accuracy drop caused by binarizing single structure by replacing it with a full-precision counterpart as Figure 3(a) shows.
We ﬁnd that binarizing MHA brings the most signiﬁcant drop of accuracy among all parts of the BERT, up to 11.49% (full-precision 89.09% vs. binarized 77.60%) and 11.72% on SST-2 and QQP, respectively.
While binarizing FFN and pooler layers brings less harm to the accuracy.
Binarization of intermediate or output layer only brings 2.19% and 5.52% drop on SST-2, and binarization of pooler layer only causes 0.82%.
Same results can be seen on QQP.
Thus, improving the attention structure is the highest priority to solve the accuracy drop of binarized BERT.
From the optimization perspective , we eliminate each distillation term to demonstrate the beneﬁts it brings in Figure 3(b).
The results show that for most distillation terms, solely removing them in the distillation will harm the performance, e.g. , removing the distillation of hidden states leads to 0.11% and 1.62% drop on SST-2 and QQP, and removing that of MHA outputs decrease the model accuracy to 73.24% (4.36% drop) on SST-2.
However, when the distillation loss of attention score is removed, the performance increases 0.81% and 1.10% on SST-2 and QQP respectively.
These observations inspire us to rethink the distillation of fully binarized BERT.
It requires a new design that could utilize the full-precision teacher’s knowledge better.
Based on the above experimental observations, we ﬁnd that (1) the existing attention structure should
NOT be directly applied in fully binarized BERT and (2) the distillation for the attention score in fully binarized BERT is actually harmful, which is contrary to the practice of many existing works.
Thus in this paper, we ﬁrst present the theoretical derivation of these two phenomena and then propose a well-designed attention structure and a novel distillation scheme for fully binarized BERT.
3. B I -A TTENTION
To address the information degradation of binarized representations in the forward propagation, we propose an efﬁcient Bi-Attention structure based on information theory, which statistically maximizes the entropy of representation and revives the attention mechanism in the fully binarized BERT.
3.2. I NFORMATION D EGRADATION IN A TTENTION S TRUCTURE
Since the representations (weight, activation, and embedding) with extremely compressed bit-width in fully binarized BERT have limited capabilities, the ideal binarized representation should preserve
[Figure]
the given full-precision counterparts as much as possible, which means the mutual information between binarized and full-precision representations should be maximized.
When the deterministic sign function is applied to binarize BERT, the goal is equivalent to maximizing the information entropy H ( B ) of binarized representation B (Messerschmitt, 1971), which is deﬁned as
[Equation]
where B ∈ {− 1 , 1 } is the random variable sampled from B with probability mass function p .
Therefore, the information entropy of binarized representation should be maximized to better preserve the full-precision counterparts and let the attention mechanism function well.
The application of zeromean pre-binarized weight in binarized linear layers is a representative practice that maximizes the information of binarized weight and activation (as in Section 2.1).
The related discussion and proof is shown in Appendix A.1.
As for the attention structure in full-precision BERT, the normalized attention weight obtained by softmax is essential.
But direct application of binarization function causes a complete information loss to binarized attention weight.
Speciﬁcally, since the softmax( A ) is regarded as following a probability distribution, the elements of B sA are all quantized to 1 (Figure 4(b)) and the information entropy H ( B sA ) degenerates to 0 .
A common measure to alleviate this information degradation is to shift the distribution of input tensors before applying the sign function, which is formulated as
[Equation]
where the shift parameter τ , also regarded as the threshold of binarization, is expected to maximize the entropy of the binarized ˆB
s A and is ﬁxed during the inference.
Theorem 1.
Given A ∈ R k with Gaussian distribution and the variable ˆB
s
A generated by ˆB A w = sign(softmax( A ) − τ ) , the threshold τ , which maximizes the information entropy H ( ˆB s A ) , is negatively correlated to the number of elements k .
However, Theorem 1 shows that it is hard to statistically determine τ to maximize the information entropy of binarized counterparts for the attention weight masked by the changeable length of attention mask.
This fact means that common measure for maximizing information (as in binarized linear layers) fails in the binarized attention structure.
The proof of Theorem 1 is shown in Appendix A.2.
Moreover, the attention weight obtained by the sign function is binarized to {− 1 , 1 } , while the original attention weight has a normalized value range [0 , 1] .
The negative value of attention weight in the binarized architecture is contrary to the intuition of the existing attention mechanism and is also empirically proved to be harmful to the attention structure in Appendix C.
3.2. B I -A TTENTION FOR M AXIMUM I NFORMATION E NTROPY
To mitigate the information degradation caused by binarization in the attention mechanism, we introduce an efﬁcient Bi-Attention structure for fully binarized BERT, which maximizes information entropy of binarized representations statistically and applies bitwise operations for fast inference.
We ﬁrst maximize the information entropy H ( ˆB s A ) .
The analysis in Section 3.2.1 shows that it is hard to statistically obtain a ﬁxed mean-shift τ for softmax( A ) that maximizes the entropy of the binarized parameters.
Fortunately, since both softmax and sign functions are order-preserving, there is a threshold φ ( τ, A ) to maximize Entropy sign( A − φ ( τ, A )) , which is equivalent to maximizing the information entropy of ˆB s A in Eq.
(9) (as the Proposition 2 proved in the Appendix A.6).
Theorem 2.
When the binarized query B Q = sign( Q ) ∈ {− 1 , 1 } N × D and key B K = sign( K ) ∈ {− 1 , 1 } N × D are entropy maximized in binarized attention, the probability mass function of each element A ij , i, j ∈ [1 , N ] sampled from attention score A = B Q ⊗ B K (cid:62) can be represented as p A (2 i − D )
= 0 .
5 D C iD , i ∈
[0 , D ] , which approximates the Gaussian distribution N (0 , D ) .
Theorem 2 shows that the distribution of A approximates the Gaussian distribution N (0 , D ) and is distributed symmetrically, we can thus trivially get that φ ( τ, A ) = 0 .
The detailed proof is given in Appendix A.3.
Therefore, simply binarizing the attention score A can maximize the information entropy of binarized representation.
Then, to revive the attention mechanism to capture crucial elements, here we are inspired by hard attention (Xu et al., 2015) to binarize the attention weight into the Boolean value, while our design is driven by information entropy maximization.
In Bi-Attention, we use bool function to binarize the attention score A , which is deﬁned as
[Equation]
By applying bool( · ) function, the elements in attention weight with lower value are binarized to 0 , and thus the obtained entropy-maximized attention weight can ﬁlter the crucial part of elements.
And the proposed Bi-Attention structure is ﬁnally expressed as
[Equation]
where B V is the binarized value obtained by sign( V ) , B A is the binarized attention weight, and (cid:2) is a well-designed Bitwise-Afﬁne Matrix Multiplication (BAMM) operator composed by ⊗ and bitshift to align training and inference representations and perform efﬁcient bitwise calculation.
The detailed of BAMM is illustrated as Figure 7 in Appendix A.4.
In a nutshell, in our Bi-Attention structure, the information entropy of binarized attention weight is maximized (as Figure 4(c) shows) to alleviate its immense information degradation and revive the attention mechanism.
Bi-Attention also achieves greater efﬁciency since the softmax is excluded.
3. D IRECTION -M ATCHING D ISTILLATION
To address the direction mismatch occurred in fully binarized BERT baseline in the backward propagation, we further propose a Direction-Matching Distillation (DMD) scheme with apposite distilled activations and the well-constructed similarity matrices to effectively utilize knowledge from the teacher, which optimizes the fully binarized BERT more accurately.
3.3. D IRECTION M ISMATCH
As an optimization technique based on element-level comparison of activation, distillation allows the binarized BERT to mimic the full-precision teacher model about intermediate activation.
However, we ﬁnd that the distillation causes direction mismatch for optimization in the fully binarized BERT baseline (Section 3.1), leading to insufﬁcient optimization and even harmful effects.
Eq.
(6) shows that, the distillation for attention score A in a speciﬁc layer can be expressed as
MSE( A , A T ) , where A and A T are attention scores in binarized student BERT and full-precision teacher BERT, respectively.
Since the attention score in fully binarized BERT is obtained by multiplying binarized query B Q and key B K , the loss (cid:96) att can be expressed as:
[Equation]
Theorem 3.
Given the variables X and X T follow N (0 , σ 1 ) , N (0 , σ 2 ) respectively, the proportion of optimization direction error is deﬁned as p error Q -bit =
p (sign( X − X T ) (cid:54) =
sign(quantize Q ( X ) −
X T )) , where quantize Q denotes the Q -bit symmetric quantization.
As Q reduces from 8 to 1 , p error Q -bit becomes larger.
[Figure]
The full binarized BERT baseline distills the attention score which has severe direction mismatch, while the knowledge (take queries as example in (b)) used in our BiBERT signiﬁcantly alleviates the mismatch.
Theorem 4 shows that the distillation for binarized activation based on the numerical comparison with full-precision activation causes the most severe optimization direction error among all bitwidths.
For example, for activation following the assumption of standard normal distribution, the error probability caused by 1bit binarization is approximately 4 .
4 × that of 4-bit quantization (14.4% vs. 3.3%).
The sudden increase in the probability of optimization direction mismatch makes it a critical issue in the fully binarized BERT while neglected in the previous BERT quantization studies.
The proof and discussion of Theorem 4 is found in
Appendix A.5.
Since the attention score is obtained by direct multiplication of two binarized activations (binarized query B Q and key B Q ) , its distillation is also misled by the direction mismatch.
In Figure 5(a), the direction mismatch in the distillation of attention score is common and severe, even causes up to higher degree optimization on the mismatch direction (mismatch 0.64 vs. match 0.45 on average).
The phenomenon in Figure 5 reveals that the element-wise gradient direction error for binarized activation found in Theorem 4 accumulates in the distillation of the attention score and eventually causes a signiﬁcant mismatch between the direction of overall optimization and that of the distillation loss.
Hence, the direction mismatch hinders the accurate optimization of the fully binarized BERT.
Second, the activation scales in the fully binarized student and full-precision teacher BERT are signiﬁcantly different, e.g. , A in fully binarized BERT approximately follows the ﬁxed distribution
N (0 , D ) while that of A T in full-precision teacher is ﬂexible.
Applying the discrete binarization function also makes numerical changes of activation more signiﬁcant.
These problems for activation also hinder effective distillation for the fully binarized BERT.
3.3. D IRECTION -M ATCHING D ISTILLATION FOR A CCURATE O PTIMIZATION
To solve the optimization direction mismatch in the distillation of the BERT full binarization, we propose the Direction-Matching Distillation (DMD) method in BiBERT.
We ﬁrst reselect the distilled activations for DMD.
As discussed in Section 3.3.1, severe direction mismatch is mainly caused by the distillation of the direct binding product of two binarized activations.
Thus, we distill the upstream query Q and key K instead of attention score in our DMD for distillation to utilize its knowledge while alleviating direction mismatch.
Moreover, inspired by an observation in Section 3.1 that the distillation of MHA output is of great help for improving performance, we also distill the value V to further cover all the inputs of MHA.
Then, we construct similarity pattern matrices for distilling activation, which can be expressed as
[Equation]
14) where (cid:107) · (cid:107) denotes (cid:96) 2 normalization.
Previous work shows that matrices constructed in this way are regarded as the speciﬁc patterns reﬂecting the semantic comprehension of network (Tung & Mori, 2019;
Martinez et al., 2020).
We further ﬁnd that the matrices are also scale-normalized and stable numerically since they focus more on endogenous relative relationships and thus are suitable for distillation between binarized and full-precision networks.
The corresponding P Q T , P K T , P V T are constructed in the same way by the full-precision teacher’s activation.
The distillation loss can be expressed as:
[Equation]
where L denotes the number of transformer layers, F DMD = { Q , K , V } , and loss terms (cid:96) hid and (cid:96) pred is constructed as in Eq.
(7) to supervise hidden states and predictions, respectively.
Our DMD scheme ﬁrst provides the matching optimization direction (Figure 5(b)) by reselecting appropriate distilled parameters and then constructs similarity matrices to eliminate scale differences and numerical instability, thereby improves fully binarized BERT by accurate optimization.
E XPERIMENTS
In this section, we conduct extensive experiments to validate the effectiveness of our proposed BiBERT for efﬁcient learning on the multiple architectures and the GLUE (Wang et al., 2018a) benchmark with diverse NLP tasks.
We ﬁrst conduct an ablation study on four tasks (SST-2, MRPC,
RTE, and QQP) for BiBERT on BERT BASE (12 hidden layers) architecture (Devlin et al., 2019) to showcase the beneﬁts of the Bi-Attention structure and DMD scheme separately.
Then we compare
BiBERT with the state-of-the-art (SOTA) quantized BERTs in terms of accuracy on BERT BASE .
Our designs stand out among fully binarized BERTs and even outperform some quantized models with more bit-width parameters.
We also evaluate our BiBERT on TinyBERT 6L (6 layers) and
TinyBERT 4L (Jiao et al., 2020), and BiBERT on these compact architectures even outperforms existing methods on BERT BASE .
In terms of efﬁciency, our BiBERT achieves an impressive 56 .
3 × and 31 .
2 × saving on FLOPs and model size.
The detailed experimental setup is given in Appendix B.
4. A BLATION S TUDY
[Table]
As shown in Table 1, the fully binarized BERT baseline suffers a severe performance drop on SST-2,
MRPC, RTE and QQP tasks.
Bi-Attention and DMD can improve the performance when used alone, and the two techniques further boost the performance considerably when combined together.
To conclude, the two techniques can promote each other to improve BiBERT and further close the performance gap between fully binarized BERT and full-precision counterpart.
4. C OMPARISON WITH SOTA M ETHODS
We compare our BiBERT with the SOTA BERT quantization methods under ultra-low bit-width in terms of accuracy and efﬁciency, to fully demonstrate the advantages of our design.
Accuracy Performance.
In Table 2 and Table 3, we show experiments on the BERT BASE (as default), TinyBERT 6L , and TinyBERT 4L architectures and the GLUE benchmark with or without data augmentation.
Since the MNLI and QQP have large data volume and it brings little beneﬁts for the performance to apply data augmentation, we do not apply augmentation for these two tasks, thus they are excluded in Table 3.
Results show that BiBERT outperforms other methods on the development set of GLUE benchmark, including TernaryBERT, BinaryBERT, Q-BERT, and Q2BERT.
Table 2 shows the results on the GLUE benchmark without data augmentation.
Our BiBERT surpasses existing methods on BERT BASE architecture by a wide margin in the average accuracy, and ﬁrst achieves convergence under ultra-low bit activation on some tasks, such as CoLA, MRPC, and
RTE.
While under ultra-low bit activation, the accuracy of TernaryBERT and BinaryBERT decreases severely which also does not improve under the 2-2-2 setting (about 4 × FLOPs and 2 × storage usage increase).
On some speciﬁc tasks like MRPC, higher bit-widths and consumption did not even make TernaryBERT and BinaryBERT converge.
With data augmentation, BiBERT achieves comparable performances with full-precision BERT on several tasks, e.g. , 90.9% accuracy (drop only 2.2%) on SST-2 (Table 3).
These results indicate that BiBERT makes full use of the limited representation capabilities by the well-designed structure and training scheme.
We noticed that BiBERT lags behind Q-BERT (2-8-8) and TernaryBERT (2-2-2) on MNLI and STS-B (with DA) while surpassing them on other tasks, indicating the fully binarized BERT may have greater improvement potential in these tasks.
BiBERT is also evaluated on compact TinyBERT 6L and TinyBERT 4L architectures, with about 2 .
× and 18 . 8 × fewer FLOPs, respectively.
The results show that BiBERT
[Table]
[Table] [Table] [Table] [Table]
[Table] [Table] [Table] on these compact architectures still outperforms existing quantization methods on BERT BASE , such as TernaryBERT and BinaryBERT.
It forcefully demonstrates that the targeted designs of BiBERT enable fully binarized BERTs to run on various architectures.
Efﬁciency Performance.
As shown in Table 2 and Table 3 above, our BiBERT achieves an impressive 56 .
3 × FLOPs and 31 .
2 × model size saving over the full-precision BERT.
Furthermore, beneﬁting from simple yet effective Bi-Attention structure which casts the expensive softmax operation into a well-engineered bit operation BAMM, our BiBERT surpasses other quantized BERTs in computation and storage saving while enjoying the best accuracy.
4. M ORE A NALYSIS
[Figure]
Information Performance.
To show the improvement of information performance by applying Bi-Attention, we compare the information entropy of binarized representations for baseline and BiBERT.
As shown in Figure 6(a), we take the ﬁrst heads in layerof each model, and the same phenomenon exists in all heads and layers.
During the training process, the information entropy of attention weight in BiBERT ﬂuctuates in a small range and is almost maximized, however, that of baseline is completely degraded to 0 .
Training Curves.
We plot training loss curves of fully binarized
BERT baseline and BiBERT on SST-2 without data augmentation in Figure 6(b).
Compared with the baseline model, our method has a faster convergence rate and achieves higher accuracy, suggesting ours advantages in terms of accurate optimization.
C ONCLUSION
We propose BiBERT towards the accurate fully binarized BERT.
We ﬁrst reveal the performance bottlenecks of the fully binarized BERT baseline and build a theoretical foundation for the impact of full binarization on BERT.
Then we propose Bi-Attention and DMD in BiBERT to improve performance.
BiBERT outperforms existing SOTA BERT quantization methods with ultra-low bit activation, giving an impressive 56 .
3 × FLOPs and 31 .
2 × model size saving.
Our work gives an insightful analysis and effective solution about the crucial issues in BERT full binarization, which blazes a promising path for the extreme compression of BERT.
We hope our work can provide directions for future research.
A PPENDIX FOR B I BERT A M AIN P ROOFS AND D ISCUSSION A. D ISCUSSION OF B INARIZED L INEAR L AYER
Proposition 1.
When a random variable X follows a zero-mean Gaussian distribution, the information entropy H ( B ) is maximized, where B X = sign( X ) .
Since the value of B w depends on the sign of W and the distribution of W is almost symmetric (He & Fan, 2019; Banner et al., 2018), the balanced operation can maximize the information entropy of binarized B w on the whole.
The balanced binarized weights with the zero-mean attribute can be obtained by subtracting the mean of full-precision weights.
As Proposition 1 shown, under the binomial distribution assumption and symmetric assumption of W , when the binarized weight is balanced, the information entropy of B w takes the maximum value, which means the binarized values should be evenly distributed.
Moreover, when the weight is zero-mean, the entropy of output Z (seen as the activation in the next binarized linear layer) in the network can also be maximized.
Supposing quantized activations B x have mean E [ B x ] = µ 1 , the mean of Z can be calculated by
[Equation]
Since the zero-mean weight is applied in each layer, we have Q w ⊗ 1 = 0 , and the mean of output is zero.
Therefore, the information entropy of activations in each layer can be maximized according to Proposition 1.
Therefore, a simple redistribution for the full-precision counterpart of binarized weights can simultaneously maximize the information entropy of binarized weights and activations.
The proof of Proposition 1 is presented as below:
Proof.
According to the deﬁnition of information entropy, we have
[Equation] [Equation]
[Equation]
Then we can get the derivative of H ( B X ) with respect to p B X ( − 1)
[Equation]
When we let d H ( B X ) d p BX ( − 1) = 0 to maximize the H ( B X ) , we have p B X ( − 1)
= 0 .
Since the deterministic sign function with the zero threshold is applied as the quantizer, the probability mass function of B X is represented as
[Equation]
where f X ( x ) is the probability density function of variable X .
Since X ∼ N (0 , σ ) , the f X ( x ) is deﬁned as
[Equation]
When the information entropy H ( B X ) is maximized, we have
[Equation]
A. P ROOF OF THE T HEOREM
Theorem 1.
Given A ∈ R k with Gaussian distribution and the variable ˆB
s
A generated by ˆB A w = sign(softmax( A ) − τ ) , the threshold τ , which maximizes the information entropy H ( ˆB s A ) , is negatively correlated to the number of elements k .
Proof.
Given the A
= { A 1 , A 2 , ..., A k } ∈ R k , each A i obeys Gaussian distribution N ( µ, σ ) .
Without loss of generality, we consider the threshold of the ﬁrst variable A 1 , which can maximize the information entropy.
Such threshold τ K satisﬁes
[Equation]
[Equation]
where [ · ] denotes the Iverson bracket that is deﬁned as [ P ] =
[Equation]
if P is true; 0 otherwise, (27) and the softmax function is
[Equation]
(28) p A i ( a i ) presents the probability that the i -th element is equal to a i .
Since the softmax k function is order-preserving, there exists exactly one threshold τ ( A 2 , A 3 , ..., A k ; τ k ) such that
[Equation]
In other words, we can convert the after-softmax threshold τ k to a before-softmax threshold τ ( A 2 , A 3 , ..., A k ; τ k ) for A 1 .
Then the Eq.
(26) can be express as
[Equation]
When we takes the { k + 1 } -th variable A k +1 into consider.
Since the (cid:80) ki =2 e A i < (cid:80) k +1 i =
2 e A i is always satisﬁed, τ
( A 2 , A 3 , ..., A k , A k +1 ; τ k ) >
τ ( A 2 , A 3 , ..., A k ; τ k ) is also always satisﬁed.
Consider the function F ( x ) , which is deﬁned below:
[Equation]
F ( x ) is a strictly monotone increasing function, which means F ( τ ( A 2 , A 3 , ..., A k , A k +1 ; τ k )) >
F ( τ ( A 2 , A 3 , ..., A k ; τ k )) .
Then we have
[Equation]
Since the information entropy of sign(softmax K +1 ( A 1 − τ K +1 )) is maximized, we have
[Equation]
∞
[Equation]
+
[Equation]
+
[Equation]
=
[Equation]
p
[Equation]
( a i ) da i = 0 .
5 , (35) thus,
[Equation]
Then we can get τ k +1 < τ k .
Therefore, the threshold τ which maximizes the information entropy
H ( ˆB s A ) is negatively correlated to the number of elements k .
A. P ROOF OF THE T HEOREM
Theorem 2.
When the binarized query B Q = sign( Q ) ∈ {− 1 , 1 } N × D and key B K = sign( K ) ∈ {− 1 , 1 } N × D are entropy maximized in binarized attention, the probability mass function of each element A ij ,
Proof.
First, we prove that the distribution of A ij can be approximated as a normal distribution.
Considering the deﬁnition of A , A ij can be expressed as
[Equation]
where B Q ,il represents the j -th element of i -th vector of B Q and B K ,il represents the j -th element of i -th vector of B K .
The value of the element B Q ,il × B K ,jl can be expressed as
[Equation]
The B Q ,il × B K ,jl only can take from two values and its value can be considered as the result of one
Bernoulli trial.
Thus for the random variable A ij sampled from the output tensor A , the probability mass function, p A can be expressed as
[Equation]
where p e denotes the probability that the element B Q ,il × B K ,jl takes 1 .
According to the De
Moivre–Laplace theorem (Walker & Helen, 1985), the normal distribution N ( µ, σ 2 ) can be used as an approximation of the binomial distribution under certain conditions, and the p A (2 i − D ) can be approximated as
[Equation]
and then, we can get the mean µ = 0 and variance σ = √ D of the approximated distribution N .
Now we give proof of this below.
According to Proposition 1, both B Q ,il and B K ,il have equal probability to be 1 or − 1 , which means p e = 0 .
Then we can rewrite the equation as
[Equation] [Equation]
Then we move to calculate the mean and standard variation of this distribution.
The mean of this distribution is deﬁned as
[Equation]
By the virtue of binomial coefﬁcient, we have
[Equation]
Besides, when D is an even number, we have (2 i − D )0 .
5 D C iD = 0 , i =
D 2 .
These equations prove the symmetry of function (2 i − D )0 .
5 D C iD .
Finally, we have
[Equation] [Equation]
[Equation]
The standard variation of p A is deﬁned as
[Equation]
To calculate the standard variation of p Z
[Equation]
we use Binomial Theorem and have several identical equations:
[Equation]
These identical equations help simplify Eq.
[Equation]
Now we proved that, the distribution of output is approximate normal distribution N (0 , D )
[Equation]
A. BAMM O PERATION IN B I -A TTENTION
In the Bi-Attention structure, we apply the bool function to obtain the binarized attention weights
B with values of 0 and 1, which maximizes the information entropy while restoring the perception of attention mechanism for the input.
However, we noticed that in the actual hardware deployment, the 1-bit matrix stored in the hardware is uniﬁed into the same form (with binary values of 1 and -1) and is supported by most existing hardware.
Therefore, we propose a new bitwise operation (cid:2) to support the computation between the binarized attention weight bool( A ) and the binarized value B V during inference, which is deﬁned as
[Equation]
where B A (cid:48) ∈ {− 1 , 1 } N × D is the representation of bool( A ) on hardware that B A (cid:48) ij = − 1 where bool( A ) ij = 0 , 1 ∈ { 1 } N × D is an all ones matrix, and ⊗ and (cid:28) is the bitwise matrix multiplication and bit-shift operation.
The (cid:2) can also be simulated as the common full-precision multiplication as the ⊗ during training.
We show the calculation process of the proposed (cid:2) operation in Figure 7, which also concludes the process of ⊗ .
-1𝐀
𝐁 𝐯 𝐕 -0.9 0.3 0.2 0.
-1𝐀
𝐁 𝐯 𝐕 × 2 0 2 1 -12 -2-1 1 1 1 -11 -1 1 𝐁 𝐀 ′ 1 -1 1 1 11 -1 1 𝐁 𝐯 ⨂ 1-bit representation on hardware -1 1 1 1 -11 -1 1 𝐁 𝐀 ′ 2 0 2 1 -12 -2⊠ 1 1 1 1 11 11 -1 1 1 11 -1⨂ + 𝐁 𝐯 ≫-1 1 1 1 -11 -1 1 𝐁 𝐀 ′ 1 -1 1 1 11 -1 1 𝐁 𝐯 ⨂ 1 1 1 -1 -
-3-1 11 1-1 1xnor popcount
Training
Inference
[Figure]
Theorem 4.
Given the variables X and X T follow N (0 , σ 1 ) , N (0 , σ 2 ) respectively, the proportion of optimization direction error is deﬁned as p error Q -bit =
p (sign( X − X T ) (cid:54) =
sign(quantize Q ( X ) −
X T )) , where quantize Q denotes the Q -bit symmetric quantization.
As Q reduces from 8 to 1 , p error Q -bit becomes larger.
Proof.
Given the random variables X ∼ N (0 , σ 1 ) and X T ∼ N (0 , σ 2 ) , the Q -bit symmetric quantization function quantize Q is expressed as (take X as an example)
[Equation]
where the (cid:98)·(cid:99) denotes the round down function, and the range [ − L, L ] is divided into 2 Q − 1 inter.
The optimization direction error occurs when sign( X − ˆ X ) =
sign(quantize Q ( X ) − X T ) , i.e.,
X > X T and quantize Q ( X ) <
X T or X < X T and quantize Q ( X ) > X T .
(1) When − L < X < L , (cid:98) (2 Q − 1)
X 2 L
[Equation]
quantize
[Equation]
X )
[Equation]
a) If X T < (cid:98) (2 Q − 1) X 2 L (cid:99)
L 2 Q − 1 , since (cid:98) (2 Q − 1)
X 2 L (cid:99) 2 L 2 Q − 1 <
X , we have X T <
X .
And since (cid:98) (2 Q − 1) X 2 L (cid:99) 2 L 2 Q − 1 < quantize Q ( X ) , X T < quantize Q ( X ) .
Thus, the optimization direction is always right in this case.
b) If X T > (cid:98) (2 Q − 1) X 2 L
+ 1 (cid:99)
L 2 Q − 1 , since (cid:98) (2 Q − 1)
X 2 L
+ 1 (cid:99)
L 2 Q − 1 > X , we have X T > X .
And since (cid:98) (2 Q − 1) X 2 L (cid:99) 2 L 2 Q − 1 > quantize Q ( X ) , X T > quantize Q ( X ) .
Thus, the optimization direction is always right in this case.
c)
If (cid:98) (2 Q − 1) X 2 L (cid:99)
L 2 Q −≤ X T ≤ (cid:98) (2 Q − 1) X
2 L + 1 (cid:99) 2 L 2 Q − 1 , ﬁrst, the probability of X >
X T and quantize Q ( X ) <
X T can be calculated as:
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Table] [Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
where f X,X T ( · , · ) is the probability density function of the joint probability distribution for { X, X T } , and [ · ] denotes the Iverson bracket as deﬁned in Eq.
Then we get the probability of X < X T and quantize Q ( X ) >
X T as p
[Equation]
[Equation]
Since f X,X T ( X, X T ) ≥ 0 is constant established, p error1 Q and p error2 Q increases as Q becomes smaller.
(2) When X > L , quantize Q ( X )
=
L .
X T > X > L = quantize Q ( X ) is constant established
when X T > X , and when X T <
X , the probability of X T > quantize Q ( X ) =
L is also constant based on the given distribution of X T .
Thus, the optimization direction is always right in this case.
(3) When X < − L , quantize Q ( X ) =
− L .
X T <
X < − L = quantize Q ( X ) is constant established when X T < X , and when X T > X , the probability of X T > quantize Q ( X ) = − L is also constant based on the given distribution of X T .
Thus, the optimization direction is always right in this case.
Although we have obtained the correlation between quantization bit-width and error of direction under symmetric quantization, it is difﬁcult to directly give an analytical representation of the error proportion of Gaussian distribution input under Q-bit quantization.
Therefore, we use the Monte
Carlo algorithm to simulate the probability of directional error caused by Q -bit by the error proportion of the pre-quantized data X ∈ R 10000 , where each element in X is sampled from the standard normal distribution.
When the quantization range is [ − 1 , 1] , the results are shown in Table 4.
This result experimentally proved our theorem, and shows that the probability of direction mismatch increases rapidly in 1-bit quantization.
A. P ROOF AND D ISCUSSION OF THE E QUIVALENCE
Considering the order-preserving characteristic of both sign and softmax , only the largest n elements (the value of n is related to the variable k for speciﬁc matrix A in Theorem 1) are binarized to 1 while others are binarized to − 1 .
Therefore, when sign(softmax( A ) − τ ) is optimized to have the maximized information entropy, there exists a corresponding threshold φ ( τ, A ) that maximizes the information entropy of sign( A − φ ( τ, A )) .
In addition, the conclusion of Theorem 2 that
A ∼ N (0 , D ) suggests that a ﬁxed threshold φ ( τ, A ) = 0 maximizes the information entropy of the binarized attention weight.
Proposition 2.
When the elements in tensor A are assumed as independent and identically distributed, there exists a φ ( τ, A )) making the entropy maximization of sign( A − φ ( τ, A )) be equivalent to that of sign(softmax( A ) − τ ) .
Proof.
As shown in Eq.
(8), the information entropy of ˆB A w = sign(softmax( A ) − τ ) can be expressed as:
[Equation]
where a s = softmax A ( a ) and a ∈ A .
Since softmax function is strictly order-preserving, there exists a φ ( τ, A ) makes ∀ a > τ , a s > φ ( τ, A ) , and it thereby makes following equations established:
[Equation]
where f and g are the probability density functions of variable a s and a .
And the information entropy H ( ˆB A w )
=
H ( ˆB A ) , where B A = sign( A − φ ( τ, A )) .
Therefore, when τ is optimized to maximize information, the information of B A is also maximized under the threshold φ ( τ, A ) .
B E XPERIMENTAL S ETUP B. D ATASET AND M ETRICS
We evaluate our method on the General Language Understanding Evaluation (Wang et al., 2018a) (GLUE) benchmark which consists of nine basic language tasks.
And we use the standard metrics for each GLUE task to measure the advantage of our method.
We use F1 score for QQP and MRPC tasks, Pearson Correlation and Spearman Correlation for STS-B, Mathews Correlation Coefﬁcient for CoLA and classiﬁcation accuracy for the rest tasks.
As for MNLI task, we report the accuracy on both in-domain evaluation MNLI-match (MNLI-m) and cross-domain evaluation MNLI-mismatch (MNLI-mm).
But we exclude WNLI task as previous studies do for its relatively small data volume and unstable behavior.
We also theoretically calculate the FLOPs at inference and also report the model size to give a comprehensive comparison on speed and storage.
We also provide a coarse estimation of model size and ﬂoating-point operations (FLOPs) at inference follow (Bai et al., 2020;
Zhou et al., 2016; Liu et al., 2018).
The matrix multiplication between an m -bit number and an n -bit number requires mn/ 64 FLOPs for a CPU with instruction size of 64bit.
Also, every 64 weight and embedding parameters are packed and stored together with a single instruction, so the model is signiﬁcantly compact.
B. I MPLEMENTATION
Backbone .
We follow (Bai et al., 2020) to take full-precision well-trained DynaBERT (Hou et al., 2020) as the teacher model to self-supervise the training of the binarized DynaBERT.
The number of transformer layers is 12 , the hidden state size is 768 and the number of heads in MHA is 12 , which is a typical setting of BERT-base.
Unlike (Bai et al., 2020), we do not pre-train an intermediate model as the initialization, but directly binarize the model from the full-precision one.
We use naive sign function to binarize parameters in the forward propagation and STE in the back-propagating, and also use MSE loss to distill the teacher model as described in Sec.
2.2 as the baseline binarized
BERT model.
Binarized Layers .
We follow the previous work to binarize the embedding layer, MHA and FFN in transformer layers, but leave full-precision classiﬁer.
That is because a common practice for BERT is to trivially cast the regression tasks to multiclass classiﬁcation tasks, such as STS-B with 5 classes and MNLI with 3 classes.
Therefore, the answers for above two tasks cannot be encoded by just 1-bit.
Settings .
The learning rate is initialized as 2 × 10 − 4 and decays to 0 with Adam optimizer.
We adopt data augmentation on GLUE tasks except MNLI and QQP for the little beneﬁt but it is timeconsuming.
It is noteworthy that we take more training epochs for every quantization method on each tasks to have a sufﬁcient training, which is 50 for CoLA, 20 for MRPC, STS-B and RTE,for SST-2 and QNLI, 5 for MNLI and QQP.
Other Quantized BERTs .
We implement TernaryBERT under 2-2-1 and 2-2-2 settings, and BinaryBERT under 1-1-1 and 1-1-2 settings for the compraisonal experiments.
For
BinaryBERT (Bai et al., 2020), except 4-bit and 8-bit activation settings mentioned in their paper, the ofﬁcial code ( https:// github.com/huawei-noah/Pretrained-Language-Model/blob/ 1dbaf58f17b8eea873d76aa388a6b0534b9ccdec/BinaryBERT/ ) provides the implementations for 1-bit or 2-bit activation settings, which are Binary-Weight-Network (BWN) and Ternary-Weight-Network (TWN), respectively, and we completely follow the released code.
Speciﬁcally, we apply Binary-Weight-Network (BWN) and Ternary-Weight-Network (TWN) methods to activation under 1-bit and 2-bit input settings (1-1-1 and 1-1-2), respectively.
The formulations and implementations are presented in (Bai et al., 2020) and ofﬁcial codes, respectively:
BWN:
[Equation]
TWN:
[Equation]
where sign( · ) is the sign function, ∆ = 0 .
7 n (cid:107) w t (cid:107) 1 and α = 1 | I | (cid:80) i ∈I | w ti
| with I =
{ i | ˆ w ti (cid:54) = 0 } .
For TernaryBERT, since there is no recommended setting under 1-bit and 2-bit input in its ofﬁcial code, we use the same implementation as BinaryBERT to quantize activation (BWN for 1-bit and
TWN for 2-bit) for fair comparison.
In addition, compared with TernaryBERT and BinaryBERT, under the 1-bit input setting, our BiBERT applies a simple sign function to binarize activation, while the former two apply real-time calculated scaling factors.
Therefore, under the same 1-1-1 setting, our BiBERT achieved better results with a smaller amount of computation.
C V ARIANTS OF B INARIZED A TTENTION M ECHANISM
In binarization literature, function sign is widely used to quantize a real number x ∈ R into y ∈ {− 1 , 1 } .
However, directly applying sign on attention weight is against the constraint that attention weight should be in [0 , 1] .
Inspired by this fact, we deﬁne function bool to map x into { 0 , 1 } .
We empirically evaluate these two binarization methods on several GLUE tasks in Table 5.
[Table]
Table 5 shows the performance of four different binarized attention mechanisms.
All the experiments use the DMD method and do not use data augmentation.
From these results, we can conclude that the tradition binarization method sign is not suitable to binarize attention weight since it can not selectively neglect the unrelated parts.
This operation not only obstacles the optimization process but also degrades the expressive ability of the whole attention mechanism.
It strongly shows the necessity of using a different binarization operator on attention weights.
D V ISUALIZATIONS
In Figure 8, we provide a coarse but bird-view of attention patterns across layers and heads in different model for same inputs on SST-2 task, including the full-precision, fully binarized BERT baseline and our BiBERT.
As can be seen in the ﬁgure, full-precision model always elicits certain tendency towards different contents, and each head captures a unique pattern with bias.
However, the attention perception is totally lost when fully binarized the BERT model with baseline methods.
The attention weight is identical regardless of inputs, and exactly degrades to a fully connection.
As for our BiBERT, it delivers various attention patterns among layers and heads.
We concede that some patterns seem to be completely different from the full-precision ones, some are even blank.
But we venture to attributes the phenomenon to the special design of our Bi-Attention structure.
We ensure the overall probability of attention weights to have a information entropy maximized distribution, but the function and quality of each head is special after an adaptive training.
It results in a selective attention according to the inputs, therefore we revive the perception of the model, and make it be able to elicit more diverse activation patterns.
[Figure]