BiT : Robustly Binarized Multi-distilled Transformer
Anonymous Author(s)
AffiliationAddress email
Abstract
Modern pre-trained transformers have rapidly advanced the state-of-the-art in machine learning, but have also grown in parameters and computational complexity,making them increasingly difficult to deploy in resource-constrained environments.
Binarization of the weights and activations of the network can significantly alleviatethese issues, however is technically challenging from an optimization perspective.
In this work, we identify a series of improvements which enables binary transformers at a much higher accuracy than what was possible previously.
These include atwo-set binarization scheme, a novel elastic binary activation function with learnedparameters, and a method to quantize a network to its limit by successively distilling higher precision models into lower precision students.
These approaches allowfor the first time, fully binarized transformer models that are at a practical levelof accuracy, approaching a full-precision BERT baseline on the GLUE languageunderstanding benchmark within as little as 5.9%.
Introduction
The past few years have witnessed tremendous advances in almost all applied fields of AI.
Itwould hardly be a simplification to say that the bulk of these advances was achieved by scalingthe transformer architecture (Vaswani et al., 2017) to ever larger sizes with increasing computationbudget (Devlin et al., 2019; Liu et al., 2019; Radford et al., 2018, 2019; Raffel et al., 2020; Brownet al., 2020).
On the other hand, mobile devices and wearables have proliferated and shrunk in size,with stringent requirements for storage, computation and energy consumption.
Consumers demandmore portability, while having access to all that current AI technology has to offer.
As a result, thegap between what is possible in AI, and what is deployable has never been wider.
While there is a variety of methods to increase inference efficiency in neural networks (e.g. knowledgedistillation, pruning), quantization has some attractive properties and has been widely successfulin practice (Gholami et al., 2021).
For one, storage and latency gains from quantization are deterministically defined for a given quantization level.
For instance, reducing the precision of modelparameters by a given factor immediately translates to an identical reduction in storage cost.
Similarly,reducing the precision of arithmetic operations results in a corresponding reduction in computationalcost.
Uniform quantization is hardware friendly, making it relatively simple to realize theoreticalimprovements in practice.
Binarization represents the extreme limit of quantization, promising a 32 × reduction in storage overfull-precision (32-bit) models.
Moreover, binary arithmetic completely eliminates multiplicationsin favor of bit-wise XNOR operations (Courbariaux et al., 2016; Rastegari et al., 2016), enablingeven further improvements when using special purpose hardware.
Energy efficiency improvementsbetween 100-1000x have been claimed to be possible with binary neural networks (BNNs), over theirfull-precision counterparts (Nurvitadhi et al., 2016).
The obvious challenge with binarization is the difficulty of optimization.
While all quantization isdiscontinuous, higher precisions allow approximating the full-precision network to a better extent,where with BNNs, this becomes much harder.
Surprisingly, researchers in computer vision have beenable to demonstrate BNNs with remarkable accuracy (Liu et al., 2018;
Qin et al., 2020; Martinez et al.
Unfortunately, while these works have mostly been developed on convolutional architecturesfor image tasks, they have not generalized well to transformer models.
For instance, recent work (Qinet al., 2021) has shown that a BERT (Devlin et al., 2019) model with binary weights and activationslags its full-precision counterpart by as much as 20 points on average on GLUE dataset.
Even forweights-only binarization, the loss landscape was shown to be too irregular, and recent work resortedto complex and specialized methods such as weight-splitting from half-width models to achieve areasonable accuracy (Bai et al., 2021).
With this background, we tackle the problem of fully binarizing transformer models to a high levelof accuracy.
With the expansion trend of transformers towards becoming the standard architecturechoice for all fields of AI, we believe a solution to this problem could be highly impactful.
Our approach follows the same paradigm as previous work, based on knowledge distillation (Hintonet al., 2015) from higher precision models using the straight-through estimator (STE) of Bengio et al.(2013).
In view of the optimization difficulties, we take the following steps to ensure that the studentand teacher models are well-matched:
• In Section 3 we describe a robust binarization framework, which allows the binary student
network to better match the output distribution of the teacher.
This allows us to achieveSoTA results for extreme activation quantization with BERT, producing models with little
loss in accuracy down to a quantization level of binary weights and 2-bit activations andimproves over previous setups by large margins in the fully binary (1-bit) setting.
It also
leads to competitive results for weight binarization with 4-bit activations using a single
knowledge distillation step.
• To further improve binary models, we propose a multi-distillation approach, describedin Section 4.
Instead of distilling directly into a fully binary model, we first distill anintermediate model of medium precision and acceptable accuracy.
This model then becomesthe teacher in the next round of distillation into increasingly quantized models.
Such amethod ensures that the student model doesn’t drift too far from the teacher, while alsoensuring as good an initialization as possible.
We call the resulting model BiT 1 .
In the vanilla setting without data augmentation, our approach reduces the accuracy gap to a fullprecision BERT-base model by half on the GLUE (Wang et al., 2019) benchmark compared to theprevious SoTA.
When using data augmentation, we are able to reduce the absolute accuracy gap toonly 5.9 points (from over 15 points previously).
In addition to the fully binary setting, we also
report
SoTA results with binary weights and 2-bit activations, where our models trail the full-precisionbaseline by only 3.5 points.
Our code with scripts to reproduce all of our results is packaged in thesupplementary material, and will be released publicly.
Background2.1 Transformer architecture
The transformer model of Vaswani et al.
(2017) is composed of an embedding layer, followed by
N transformer blocks and a linear output layer.
Each transformer block consists of a multi-headattention layer followed by a feed-forward network, as shown in Figure 1.
The multi-head attentionlayer is a concatenation of K scaled dot-product attention heads, defined by:
[Equation]
where d k is the dimension of each key, Q, K, V are weight matrices for the query, key and valuerespectively.
As such, the computation in a transformer model is limited to linear matrix multiplications and additions, pointwise non-linearities (most commonly Sigmoid (Han & Moraga,1995), GeLU (Hendrycks & Gimpel,
2.2 Quantization
A vector w is uniformly quantized to b -bit precision, if its entries are restricted to the set { 0 , 1 , . . .
, 2 b −1 } for asymmetric case or {− 2 b , − 2 b + 1 , . . .
, 2 b − 1 } for symmetric case, up to a real-valued scaleα .
This allows vector operations to utilize lower precision arithmetic, making them more efficient bya factor of Bb compared to full-precision calculation using B bits.
The scaling operation is still inhigher precision, but if the dimensionality of w ≫ Bb , then the extra computation is negligible.
A neural network with parameters quantized to b w bits takes up Bb w times less space.
However, to takeadvantage of lower-precision arithmetic, the input vectors (activations) to each vector/matrix operationalso need to be quantized.
A network which has weights quantized to b w bits and activations quantizedto b a bits is denoted as W b w A b a .
In this work, we’re specifically interested in W1A1 transformers.
Binary arithmetic is especially attractive, since multiplications reduce to XNOR operations, and canbe implemented orders of magnitude more efficiently using specialized hardware (Nurvitadhi et al.,2016).
2.3 Knowledge distillation
Knowledge distillation (KD) (Hinton et al., 2015) is a technique whereby a student network canbe trained to mimic the behavior of a teacher network.
This is especially useful when the studentnetwork is more efficient and easier to deploy than the more complex and cumbersome teacher.
Thebasic way of performing KD is by using the output distribution of the teacher model ( p ) as soft targetsfor training the student model.
If q is the student model’s output, then we have the loss term:
[Equation]
The advantage of KD over simple supervised training of a more efficient model is that the teachermodel provides a richer training signal including model confidence for each output class.
For computer vision tasks, distilling the final logits solely works well for binary neural networks (Liuet al., 2020).
If the student and teacher architectures are compatible, one can also distill intermediateactivations for faster convergence and better transfer and generalization (Aguilar et al., 2020):
[Equation]
where r si and r ti are the corresponding transformer block output activations from student and teacher.
Robust binarization setup
In this section we first bring together some best practices and minor improvements which we havefound helpful in simplifying previous work and building a strong baseline.
Then we present a novelactivation binarization scheme, which we will show to be critical to achieve good performance.
3.1 Two-set binarization scheme
In contrast to convolutional neural networks on images where activations exhibit comparable distributions, different activations in transformer blocks are performing different functionalities, and thus varyin their output distributions.
In particular, these activations can be divided into two categories: theactivations after Softmax/ReLU layer that contains positive values only and the remaining activationswith both positive and negative values (e.g., after matrix multiplication).
If we denote by X R thevector of activation values, then the two cases are X i R ∈ R
+
and X i R ∈ R respectively.
For the former set, mapping to the binary levels {− 1 , 1 } would result in a severe distribution mismatch.
Therefore we instead map non-negative activation layers to ˆX B ∈ { 0 , 1 } n and binarize activationlayers with X R ∈ R n to ˆX B ∈ {− 1 , 1 } n , shown in Figure 1.
A prior work BiBERT (Qin et al. , 2021)also suggests binarizing attention to { 0 , 1 } , but with bool function replacing SoftMax , while weempirically find that simply binarizing attentions after SoftMax to { 0 , 1 } works better and binarizing ReLU output to { 0 , 1 } instead of {− 1 , 1 } brings further improvements. (See Section A.3 for details).
Optimal scaling factor in two sets Additionally, we apply a layer-wise scaling factor to binarizedactivations to reduce the binarization error, i.e. , X B = α ˆX B .
The optimal values of α are different
[Figure]
for the ˆX B ∈ { 0 , 1 } n and ˆX B ∈ {− 1 , 1 } n cases and can be calculated by minimizing the l 2 error:
[Equation]
Following XNOR-Net (Rastegari et al., 2016), by expanding Eq.
3, we have
[Equation]
For the layers with X R ∈ R n we follow the traditional methods of binarizing activations (Rastegariet al., 2016; Liu et al., 2018) by taking the sign of real-valued activations:
[Equation]
In that case, ˆ X BT ˆX B = n X R , where n X R is number of elements in X R , and α ∗ can be solved as:
[Equation]
[Equation]
For the activations in attention layers or after the ReLU non-linearity layers with X R ∈ R n + , webinarize the activations to ˆX B ∈ { 0 , 1 } n by rounding the real-valued activations:
[Equation]
In that case, ˆ X BT ˆX B = n { X R ⩾ 0 .
5 } where n { X R ⩾ 0 .
5 } denotes the number of elements in X R that 137 are greater than or equal to 0 .
Then α ∗ can be solved as:
[Equation]
3.2 Best practices
We performed thorough experimentation and discovered the following modifications to be useful.
Simplified knowledge distillation Compared to previous BERT model binarization works (Bai et al.
Qin et al., 2021) which also attempt to distill the attention scores, we provide analysis and experimental results (Section 5.3) showing that using only L reps from transformer block outputs and
L logits is more effective while being simpler.
We also forego the two-step distillation scheme of Baiet al.
(2021) in favor of a single step, joint distillation, where our training loss is simply L logits + L reps .
Mean subtraction in weight binarization For weight binarization, centeralizing the real-valuedweights to be zero-mean before binarization can increase the information carrying capacity of thebinary weights.
Thus, for weight binarization, we have: W i B =
|| W R
||
l 1 n WR Sign( W i R − W R ) .
Gradient clipping Clipping gradients to 0 when X i R / ∈ [ − 1 , 1] (or X i R / ∈ [0 , 1] if ˆX B ∈ { 0 , 1 } n ) isa common technique for training binarized neural networks.
However, we find that clipping weightgradients is harmful for optimization.
Once a weight is outside of the clip range, the gradient is fixedto 0 , preventing further learning.
This is not so for activations, since the activation value changes foreach input.
As a result, we apply gradient clipping only to activations but not to weights.
Non-linearity We prefer ReLU activations whenever the output range is non-negative.
Combining these, we are able to build a strong baseline, which improves the accuracy by 9.6% overnaively binarized transformers.
Additionally, these techniques allow us to train a weight binarizedtransformer network in a single training step using knowledge distillation ( i.e. , without resorting toweight splitting as in BinaryBert (Bai et al., 2021))
(See Section 5.3 for details).
3.3 Elastic binarization function
The fixed scaling and threshold derived previously works reasonably well, but might not be optimalsince it ignores the distribution of the variable which is being binarized.
Ideally, these parameters canbe learned during training to minimize the target loss.
When using classical binarization methods, i.e. , ˆX i B =
Sign( X i R ) , the binary output is independentof the scale of the real-valued input.
However, in our case where ˆX i B =
⌊ Clip
( X i R , 0 , 1) ⌉ , thisindependence no longer holds.
Learning the scaling and threshold parameters, and how to approximatethe gradients precisely in the process becomes crucial for the final accuracy.
To handle this, inspired by the learnable threshold in ReActNet (Liu et al., 2020), we propose theelastic binarization function to learn both the scale α ∈ R + and the threshold β ∈ R :
[Equation]
In the function, we initialize α with α ∗ in Sec.
3.1 and β to be 0 , and train it with gradients fromthe final loss.
To back-propagate the gradients to α through the discretized binarization function,we follow the practice in Choi et al.
; Zhou et al.
(2016) to use straight-through estimator(STE)
(Bengio et al., 2013) to bypass the incoming gradients to the round function to be the outgoinggradients:
[Equation]
Then the gradients w.r.t.
β can be similarly calculates as:
[Equation]
For the layers that contain both positive and negative real-valued activations i.e. , X R ∈ R n , thebinarized values ˆX B ∈ {− 1 , 1 } n are indifferent to the scale inside the Sign function:
[Equation]
Require: D train , D dev ▷ Training Data Require: h▷
Full-precision Model Require: Q = { ( b 1 w , b 1 a ) , . . .
, ( b kw , b ka ) } ▷ Quantization Schedule 1: h teacher ← h 0 2: for b iw , b ia in Q do 3: h student ← Quantize ( h teacher , b iw , b ia ) 4:
KnowledgeDistill ( h student , h teacher , D train , D dev ) 5: h teacher ← h student 6: end for 7: return h student
In that case, since the effect of scaling factor α inside the Sign function can be ignored, the gradientw.r.t.
α can be simply calculated as:
[Equation]
In our ablations (Section 5.3 and A.2) we show that using this simple elastic binarization functioncan bring a 15.7% accuracy boost over our strong baseline on the GLUE benchmark.
Multi-distilled binary transformer
Classical knowledge distillation (KD) (Hinton et al., 2015) trains the outputs ( i.e. , logits) of astudent network to be close to those of a teacher, which is typically larger and more complex.
Thisapproach is quite general, and can work with any student-teacher pair which conforms to the sameoutput space.
However, in practice, knowledge transfer happens faster and more effectively if theintermediate representations are also distilled (Aguilar et al., 2020).
This approach has been founduseful when distilling to student models with similar architecture (Sanh et al., 2019), and in particularfor quantization (Bai et al., 2021;
Kim et al., 2019).
Note that having a similar student-teacher pair is a requirement for distilling representations.
Whilehow similar they need to be is an open question, intuitively a teacher which is architecturally closerto the student should make transfer of internal representations easier.
In the context of quantization, itis easy to see that lower precision students are progressively less similar to the full-precision teacher,which is one reason why binarization is difficult.
This suggests a multi-step approach, where instead of directly distilling from a full-precision teacherto the desired quantization level, we first distill into a model with sufficient precision in order topreserve quality.
This model can then be used as a teacher to distill into a further quantized student.
This process can be repeated multiple times, while at each step ensuring that the teacher and studentmodels are sufficiently similar, and the performance loss is limited.
This multi-distillation approachis sketched in Algorithm 1.
The multi-step distillation follows a quantization schedule , Q = { ( b 1 w , b 1 a ) , ( b 2 w , b 2 a ) , . . .
, ( b kw , b ka ) }with ( b 1 w , b 1 a ) >
( b 2 w , b 2 a ) > . . .
> ( b kw , b ka ) 2 .
( b kw , b ka ) is the target quantization level, which is inour case binary for both weights and activations.
In practice, we find that down to a quantization levelof W1A2, we can distill models of reasonable accuracy in single shot, following the best practicesoutlined in Section 3.2 (See our 1-1-2 baseline results in Table 1).
As a result, we follow a fixedquantization schedule, W32A32 → W1A2 → W1A1.
This is not necessarily optimal, and how toefficiently find the best quantization schedule is an interesting open problem.
We present our initialexplorations towards this direction in Section 5.5.
Combining the elastic binary activations with multi-distillation we obtain BiT , the robustly binarizedmulti-distilled transformer.
Note that BiT simultaneously ensures good initialization for the eventual(in our case binary) student model.
Since the binary loss landscape is highly irregular, good initialization is critical to aid optimization.
Previous work has proposed progressive distillation (Zhuanget al., 2018;
Yang et al., 2019) to tackle this problem, wherein the student network is quantized at
[Table]
increasing severity as the training progresses.
However, this method does not prevent the studentnetwork from drifting away from the teacher, which is always the full-precision model.
We compareto progressive distillation in Section A.1.
Main results
We follow recent work (Bai et al., 2021; Qin et al., 2021) in adopting the experimental setting of
Devlin et al.
(2018), and use the pre-trained BERT-base as our full-precision baseline.
We evaluateon GLUE (Wang et al., 2019), a varied set of language understanding tasks (see Section A.5 for a fulllist), as well as SQuAD (v1.1) (Rajpurkar et al., 2016), a popular machine reading comprehensiondataset.
5.1 GLUE results
Our main results on the GLUE benchmarks are presented in Table 1.
In the setting without dataaugmentation, where we only use the original training samples for knowledge distillation, we areable to reduce the gap to the full precision baseline by 49.8%, i.e. , from 20.7 in (Qin et al., 2021) to10.4 points.
We also see that our baseline models with elastic activation binarization already improveprevious SoTA by large margins.
In the binary weight setting (4-bit activations), we can match or outperform Bai et al.
(2021) withoutthe need for pre-training half-width models and subsequently splitting weights.
This result shouldmake binary weight models much easier to implement and deploy.
We also set a new state of the art for binary weight 2-bit activation (W1A2) models, with only a 3.5point degradation compared to the full-precision baseline (using data augmentation).
While not asefficient as binary, 2-bit arithmetic can also be performed without multiplications, making it a goodefficient alternative in applications where the performance cost of going to fully binary is significant.
5.1.1 Data augmentation
From Table 1, it can be observed that the datasets with small training sets still have a large gap fromthe full-precision baseline.
As a result, we employ data augmentation heuristics (following the exactsetup in Zhang et al.
on the datasets with small training sets (all except MNLI, QNLI) totake better advantage of our model’s strong representational capability.
This further reduces thequantization gap, with our models eventually trailing the full-precision model by only 5.9 points onaverage on the GLUE benchmark.
[Table]
5.2 SQuAD results
[Table]
We also evaluate on the popular machine readingcomprehension (MRC) dataset from Rajpurkaret al.
We compare to our own implementation of the MRC task on top of the BiBERTcodebase, since SQuAD results are not reportedin that work.
We also show results using the
BinaryBERT codebase, without using weightsplitting.
The results (Table 3) show that thistask is significantly harder than most documentclassification benchmarks in GLUE, and previous binarization methods fail to achieve any meaningful level of performance.
BiT does much better,but still trails the 32-bit baseline by 14.8 points in F1.
We conclude that despite the improvements wehave demonstrated on the GLUE benchmark, binarizing transformer models accurately is far from asolved problem in general.
5.3 Ablations
We start from the basic binarization implementation from Bai et al.
(2021), and add each of ourcontributions in sequence to get a better idea how each contributes to the performance.
The resultsare shown in Table 2.
We start by removing attention distillation (Section 3.2), which results in a 2.3% improvement (row 5vs. 4).
Then switching to our two-set binarization (Section 3.1), which binarizes the attention scoresdifferently than the feed-forward activations, which gives an additional 7.3% boost (row 6).
Thisresults in a much stronger baseline than what was used in prior works (row 2).
Moving from fixed to elastic binarization (Section 3.3) proves hugely important, pushing the averageaccuracy to 71.0% (row 7) from only 55.3% (row 6).
Note that this model already outperforms thecurrent state-of-the-art (row 3) by 7.8% points.
Finally, we add multi-step distillation (Section 4),which adds another 2.5 points, reaching the final accuracy of 73.5% on the GLUE benchmark.
5.4 Learned parameter visualization
We visualize the optimized α in the final BiT model.
As we can see from
Figure 2, the values of the α parametersvary significantly from layer to layer,and have apparent patterns according tolayer characteristics.
For example, theattention layers need to distribute the attention to different entries, thus the scaling factor for the attentions are learnedto be small, while the scaling factors for the query and key outputs are usually larger.
Note that thebiggest α value is 200 × of the smallest α , suggesting the importance of learning α dynamically.
5.5 Exploring multi-distillation paths
So far we have only considered the fixed quantization schedule, W32A32 → W1A2 → W1A1.
Thisis motivated by early experiments showing that one-step distillation to W1A2 works reasonably well.
We explored other optimal schedules, such as distilling to W1A8 resulted in a higher accuracy model,thus a better teacher to distill down to the eventual W1A1 student.
This suggests a trade-off betweenthe quality of the intermediate model, vs. the closeness to the target quantization level.
[Figure]
Figure 3, illustrates this trade-off.
We can seethat two-step distillation improves over one-stepin every case.
While higher precision intermediate models are better as expected, it is betterto use a lower precision teacher in the last stepsince it makes the learning task for the binarystudent model easier.
The closeness to the targetquantization is favored despite the lower accuracy of teacher model.
It is of course possible, though more cumbersome, to perform more than two distillationsteps.
We also experiment with a three-stepschedule, W32A32 → W1A4 → W1A2 →
W1A1, which is plotted in the same figure(dashed line).
We find that this particular 3step schedule does not improve over the 2-stepschedule W32A32 → W1A2 → W1A1.
While this result does not preclude existence of other moreoptimal schedules, we hypothesize that this is unlikely.
Related work
Convolutional neural network quantization Neural network quantization is a practical tool for
[Figure]
compressing model size and reduce storage (Hubara et al., 2017).
Quantization for convolutionalneural networks has been studied both in the uniform quantization (Choi et al., 2018;
Zhou et al.
Gong et al., 2019) and non-uniform quantization (Zhang et al., 2018; Miyashita et al., 2016;
Li et al., 2020) settings.
The quantization level has been progressively increased, from 8-bit (Wanget al., 2018;
Zhu et al., 2020) to 4-bit (Jung et al., 2019; Esser et al., 2019) and finally to the extreme1-bit case (Courbariaux et al., 2016; Rastegari et al., 2016; Liu et al., 2018; Martinez et al., 2020).
Transformer quantization Compared to the CNNs, transformers with attention layers are naturallymore challenging to quantize (Bondarenko et al., 2021).
Previous research mainly focused on 8-bitquantization (Zafrir et al., 2019; Fan et al., 2020) or 4-bit quantization (Shen et al., 2020; Zadehet al., 2020).
Extremely low-bit quantization for transformers has only been attempted very recently.
TernaryBERT (Zhang et al., 2020) proposed to ternarize the full-precision weights of a fine-tuned
BERT model.
As a follow-up to TernaryBERT, weight binarization was proposed in Bai et al.
Here, the network is trained by first training a ternary half-sized model, which is used as initialization.
Then a weight-splitting step results in a full-sized binarized model, which is further fine-tuned in asubsequent distillation step.
Binarizing both weight and activations in a transformer has proved to bechallenging.
BiBERT (Qin et al., 2021) made the first attempt in this direction with limited success.
Their model performed 20% worse than a real-valued baseline on the GLUE benchmark (Wang et al.
,2019), which even underperforms the original LSTM baselines.
Conclusion
Large pre-trained transformers have transformed NLP and are positioned to serve as the backbonefor all AI models.
In this work, we presented the first successful demonstration of a fully binarypre-trained transformer model.
While our approach is general and can be applied to any transformer,we have limited our evaluation to BERT-based models on the GLUE and SQuAD benchmarks.
Itremains to be seen how our conclusions will hold when applied to the wide variety of pre-trainedtransformer models which have gained popularity in recent years, from small mobile models, togigantic ones with hundreds of billions of parameters.
It will also be interesting to see the performanceof the approach on different domains (such as image and speech processing) and tasks (such as textand image generation).
Demonstrating the generality of this approach in a wider setting shouldsignificantly widen its impact, therefore we identify this as an important future direction.