BiT : Robustly Binarized Multi-distilled Transformer
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 alleviate these issues, however, is technically challenging from an optimization perspective.
In this work, we identify a series of improvements that enables binary transformers at a much higher accuracy than what was possible previously.
These include a two-set binarization scheme, a novel elastic binary activation function with learned parameters, and a method to quantize a network to its limit by successively distilling higher precision models into lower precision students.
These approaches allow for the first time, fully binarized transformer models that are at a practical level of accuracy, approaching a full-precision BERT baseline on the GLUE language understanding benchmark within as little as 5.9%.
Code and models are available at: https://github.com/facebookresearch/bit .
Introduction
The past few years have witnessed tremendous advances in almost all applied fields of AI.
It would hardly be a simplification to say that the bulk of these advances was achieved by scaling the transformer architecture (Vaswani et al., 2017) to ever larger sizes with increasing computation budget (Devlin et al., 2019; Liu et al., 2019; Radford et al., 2018, 2019; Raffel et al., 2020; Brown et 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 demand more portability, while having access to all that current AI technology has to offer.
As a result, the gap 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. knowledge distillation, pruning), quantization has some attractive properties and has been widely successful in 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 model parameters 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 computational cost.
Uniform quantization is hardware friendly, making it relatively simple to realize theoretical improvements in practice.
Binarization represents the extreme limit of quantization, promising a 32 × reduction in storage over full-precision (32-bit) models.
Moreover, binary arithmetic completely eliminates multiplications in favor of bit-wise XNOR operations (Courbariaux et al., 2016; Rastegari et al., 2016), enabling even further improvements when using special purpose hardware.
Energy efficiency improvements between 100-1000x have been claimed to be possible with binary neural networks (BNNs), over their full-precision counterparts (Nurvitadhi et al., 2016).
The obvious challenge with binarization is the difficulty of optimization.
While all quantization is discontinuous, 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 been able to demonstrate BNNs with remarkable accuracy (Liu et al., 2018; Qin et al., 2020; Martinez et al., 2020).
Unfortunately, while these works have mostly been developed on convolutional architectures for image tasks, they have not generalized well to transformer models.
For instance, recent work (Qin et al., 2021) has shown that a BERT (Devlin et al., 2019) model with binary weights and activations lags its full-precision counterpart by as much as 20 points on average on GLUE dataset.
Even for weights-only binarization, the loss landscape was shown to be too irregular, and recent work resorted to complex and specialized methods such as weight-splitting from half-width models to achieve a reasonable accuracy (Bai et al., 2021).
With this background, we tackle the problem of fully binarizing transformer models to a high level of accuracy.
With the expansion trend of transformers towards becoming the standard architecture choice 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 (Hinton et al., 2015) from higher precision models using the straight-through estimator (STE) of Bengio et al.
In view of the optimization difficulties, we take the following steps to ensure that the student and 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 achieve SoTA 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 and improves 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, described in Section 4.
Instead of distilling directly into a fully binary model, we first distill an intermediate model of medium precision and acceptable accuracy.
This model then becomes the teacher in the next round of distillation into increasingly quantized models.
Such a method ensures that the student model doesn’t drift too far from the teacher, while also ensuring as good an initialization as possible.
We call the resulting model BiT 2 .
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 the previous SoTA.
When using data augmentation, we are able to reduce the absolute accuracy gap to only 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-precision baseline by only 3.5 points.
Background 2.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-head attention layer followed by a feed-forward network, as shown in Figure 1.
The multi-head attention layer 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 value respectively.
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, 2016) or ReLU (Nair & Hinton, 2010) ) and the Softmax operation (Bridle, 1989).
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 by a factor of Bb compared to full-precision calculation using B bits.
The scaling operation is still in higher 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 take advantage of lower-precision arithmetic, the input vectors (activations) to each vector/matrix operation also need to be quantized.
A network which has weights quantized to b w bits and activations quantized to 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 can be 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 can be trained to mimic the behavior of a teacher network.
This is especially useful when the student network is more efficient and easier to deploy than the more complex and cumbersome teacher.
The basic way of performing KD is by using the output distribution of the teacher model ( p ) as soft targets for 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 teacher model 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 (Liu et al., 2020).
If the student and teacher architectures are compatible, one can also distill intermediate activations 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 have found helpful in simplifying previous work and building a strong baseline.
Then we present a novel activation 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 vary in their output distributions.
In particular, these activations can be divided into two categories: the [Figure] activations after Softmax/ReLU layer that contains positive values only and the remaining activations with both positive and negative values (e.g., after matrix multiplication).
If we denote by X R the vector 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 activation layers 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 we empirically 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 binarized activations to reduce the binarization error, i.e. , X B = α ˆX B .
The optimal values of α are different 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 (Rastegari et 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]
For the activations in attention layers or after the ReLU non-linearity layers with X R ∈ R n + , we binarize 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 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., 2021; 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 Bai et 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-valued weights to be zero-mean before binarization can increase the information carrying capacity of the binary weights.
Thus, for weight binarization, we have: W i B =
|| W R
||
[Equation]
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 ) is a common technique for training binarized neural networks.
However, we find that clipping weight gradients is harmful for optimization.
Once a weight is outside of the clip range, the gradient is fixed to 0 , preventing further learning.
This is not so for activations, since the activation value changes for each 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% over naively binarized transformers.
Additionally, these techniques allow us to train a weight binarized transformer network in a single training step using knowledge distillation ( i.e. , without resorting to weight 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 optimal since it ignores the distribution of the variable which is being binarized.
Ideally, these parameters can be 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 independent of the scale of the real-valued input.
However, in our case where ˆX i B =
⌊ Clip
( X i R , 0 , 1) ⌉ , this independence no longer holds.
Learning the scaling and threshold parameters, and how to approximate the 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 the elastic 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 from the final loss.
To back-propagate the gradients to α through the discretized binarization function, we follow the practice in Choi et al.
; Zhou et al.
; Esser et al. (2019) to use straight-through estimator (STE) (Bengio et al., 2013) to bypass the incoming gradients to the round function to be
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 the outgoing gradients:
[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 , the binarized values ˆX B ∈ {− 1 , 1 } n are indifferent to the scale inside the Sign function:
[Equation]
[Equation]
In that case, since the effect of scaling factor α inside the Sign function can be ignored, the gradient w.r.t.
α can be simply calculated as:
[Equation]
[Equation]
In our ablations (Section 5.3 and A.2) we show that using this simple elastic binarization function can 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 a student network to be close to those of a teacher, which is typically larger and more complex.
This approach is quite general, and can work with any student-teacher pair which conforms to the same output space.
However, in practice, knowledge transfer happens faster and more effectively if the intermediate representations are also distilled (Aguilar et al., 2020).
This approach has been found useful when distilling to student models with similar architecture (Sanh et al., 2019), and in particular for quantization (Bai et al., 2021;
Kim et al., 2019).
Note that having a similar student-teacher pair is a requirement for distilling representations.
While how similar they need to be is an open question, intuitively a teacher which is architecturally closer to the student should make transfer of internal representations easier.
In the context of quantization, it is easy to see that lower precision students are progressively less similar to the full-precision teacher, which is one reason why binarization is difficult.
[Table]
This suggests a multi-step approach, where instead of directly distilling from a full-precision teacher to the desired quantization level, we first distill into a model with sufficient precision in order to preserve 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 student models are sufficiently similar, and the performance loss is limited.
This multi-distillation approach is 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 ) 3 .
( b kw , b ka ) is the target quantization level, which is in our case binary for both weights and activations.
In practice, we find that down to a quantization level of W1A2, we can distill models of reasonable accuracy in single shot, following the best practices outlined in Section 3.2 (See our 1-1-2 baseline results in Table 1).
As a result, we follow a fixed quantization schedule, W32A32 → W1A2 → W1A1.
This is not necessarily optimal, and how to efficiently find the best quantization schedule is an interesting open problem.
We present our initial explorations towards this direction in Section 5.5.
Combining the elastic binary activations with multi-distillation we obtain BiT , the robustly binarized multi-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 (Zhuang et al., 2018;
Yang et al., 2019) to tackle this problem, wherein the student network is quantized at increasing severity as the training progresses.
However, this method does not prevent the student network from drifting away from the teacher, which is always the full-precision model.
We compare to 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.
and use the pre-trained BERT-base as our full-precision baseline.
We evaluate on GLUE (Wang et al., 2019), a varied set of language understanding tasks (see Section A.5 for a full list), as well as SQuAD (v1.1) (Rajpurkar et al., 2016), a popular machine reading comprehension dataset. 3 ( a, b ) > ( c, d ) if a > c and b ≥ d or a ≥ c and b > d . [Table]
5.1 GLUE results
Our main results on the GLUE benchmarks are presented in Table 1.
In the setting without data augmentation, where we only use the original training samples for knowledge distillation, we are able to reduce the gap to the full precision baseline by 49.8%, i.e. , from 20.7 in (Qin et al., 2021) to 10.4 points.
We also see that our baseline models with elastic activation binarization already improve previous SoTA by large margins.
In the binary weight setting (4-bit activations), we can match or outperform Bai et al.
(2021) without the need for pre-training half-width models and subsequently splitting weights.
This result should make 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. point degradation compared to the full-precision baseline (using data augmentation).
While not as efficient as binary, 2-bit arithmetic can also be performed without multiplications, making it a good efficient 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 from the full-precision baseline.
As a result, we employ data augmentation heuristics (following the exact setup in Zhang et al.
on the datasets with small training sets (all except MNLI, QNLI) to take better advantage of our model’s strong representational capability.
This further reduces the quantization gap, with our models eventually trailing the full-precision model by only 5.9 points on average on the GLUE benchmark.
5.2 SQuAD results
[Table]
We also evaluate on the popular machine reading comprehension (MRC) dataset from Rajpurkar et al.
We compare to our own implementation of the MRC task on top of the BiBERT codebase, since SQuAD results are not reported in that work.
We also show results using the
BinaryBERT codebase, without using weight splitting.
The results (Table 3) show that this task is significantly harder than most document classification 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 we have demonstrated on the GLUE benchmark, binarizing transformer models accurately is far from a solved problem in general.
5.3 Ablations
We start from the basic binarization implementation from Bai et al.
(2021), and add each of our contributions in sequence to get a better idea how each contributes to the performance.
The results are shown in Table 2.
We start by removing attention distillation (Section 3.2), which results in a 2.3% improvement (rowvs. 4).
Then switching to our two-set binarization (Section 3.1), which binarizes the attention scores differently than the feed-forward activations, which gives an additional 7.3% boost (row 6).
This results 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 average accuracy to 71.0% (row 7) from only 55.3% (row 6).
Note that this model already outperforms the current 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 α parameters vary significantly from layer to layer, and have apparent patterns according to layer characteristics.
For example, the attention layers need to distribute the attention to different entries, thus the scaling factor for the attentions are learned to be small, while the scaling factors for the query and key outputs are usually larger.
Note that the biggest α 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.
This is 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 between the quality of the intermediate model, vs. the closeness to the target quantization level.
[Figure]
Figure 3, illustrates this trade-off.
We can see that two-step distillation improves over one-step in every case.
While higher precision intermediate models are better as expected, it is better to use a lower precision teacher in the last step since it makes the learning task for the binary student model easier.
The closeness to the target quantization is favored despite the lower accuracy of teacher model.
It is of course possible, though more cumbersome, to perform more than two distillation steps.
We also experiment with a three-step schedule, 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-step schedule W32A32 → W1A2 → W1A1.
While this result does not preclude existence of other more optimal schedules, we hypothesize that this is unlikely.
Related work
Convolutional neural network quantization Neural network quantization is a practical tool for compressing model size and reduce storage (Hubara et al., 2017).
Quantization for convolutional neural networks has been studied both in the uniform quantization (Choi et al., 2018;
Zhou et al., 2016;
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 (Wang et al., 2018;
Zhu et al., 2020) to 4-bit (Jung et al., 2019; Esser et al., 2019; Liu et al., 2022) and finally to the extreme 1-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 naturally more challenging to quantize (Bondarenko et al., 2021).
Previous research mainly focused on 8-bit quantization (Zafrir et al., 2019; Fan et al., 2020) or 4-bit quantization (Shen et al., 2020; Zadeh et 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 a subsequent distillation step.
Binarizing both weight and activations in a transformer has proved to be challenging.
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 backbone for all AI models.
In this work, we presented the first successful demonstration of a fully binary pre-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.
It remains to be seen how our conclusions will hold when applied to the wide variety of pre-trained transformer models which have gained popularity in recent years, from small mobile models, to gigantic ones with hundreds of billions of parameters.
It will also be interesting to see the performance of the approach on different domains (such as image and speech processing) and tasks (such as text and image generation).
Demonstrating the generality of this approach in a wider setting should significantly widen its impact, therefore we identify this as an important future direction.