SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Yichi Zhang*
Cornell University  
yz2499@cornell. edu
Ankush Garg*
Google DeepMind  
ankugarg@google.com
Google DeepMind  
yuancao@google.com
Lukasz Lew
Google Research lew@google.com
Behrooz Ghorbani†
OpenAI ghorbani@openai.com
Zhiru Zhang
Cornell University zhiruz@cornell.edu
Orhan First
Google DeepMind orhanf@google.com
The rapid scaling of language models is motivating research using low-bitwidth quantization. In this work, we propose a novel binarization technique for Transformers applied to machine translation (BMT), the first of its kind. We identify and address the problem of inflated dot-product variance when using one-bit weights and activations. Specifically, BMT leverages additional LayerNorms and residual connections to improve binarization quality. Experiments on the WMT dataset show that a one-bit weight-only Transformer can achieve the same quality as a float one, while being  $16 \times$  smaller in size. One-bit activations incur varying degrees of quality drop, but mitigated by the proposed architectural changes. We further conduct a scaling law study using production-scale translation datasets, which shows that one-bit weight Transformers scale and generalize well in both in-domain and out-of-domain settings<sup>3</sup>.
Neural language models are scaling, with the parameter count of recent models, such as the GPT family, roughly increased by  $10 \times$  per year [29]. A scaling law study by Kaplan et al. [21] suggests that the continuous increase in model parameters is strongly correlated with performance improvement. This trend has been validated by recent successes in large-scale models, such as the 540-billion parameter Pathways Language Model (PaLM), which achieves breakthrough performance on language understanding and generation [11]. The 540-billion parameter Minerva [25] also exceeded the national average on the National Math Exam in Poland in 2021, where language models were previously far from human-level. Similarly, in the field of neural machine translation (MT), the scaling law holds, as reported by Ghorbani et al. [16], with the translation quality improving as the model size increases.
The aggressive scaling trend resulted in unprecedented challenges in model serving. In particular:
The inference cost grows exponentially. The size and computational complexity of language models are increasing rapidly, with roughly a  $10 \times$  increase in model size and a  $100 \times$  increase in operation
$^{\dagger}$ Work done while at Google.
<sup>3</sup>Source code is available in the init2winit library: https://github.com/google/init2winit/blob/master/init2winit/model_lib/xformerTranslate_binary.py
count per year [18]. However, the energy efficiency of hardware used to run these models is not keeping pace. Specifically, the energy required for FP32 operations has improved by only  $2.5 \times$  over the past 11 years (2007-2018), from  $45\mathrm{nm}$  to  $7\mathrm{nm}$  process nodes. Over the same period, DRAM access energy has only improved by  $6.3 \times$  [20]. The ever-growing gap between model size inflation and inefficiency in hardware energy utility is causing inference energy to grow exponentially, which is becoming a major cost of running language models in datacenters.
The inter-chip communication overhead becomes non-negligible. Data parallelism alone is no longer sufficient for models at such a large scale since one matrix multiplication cannot fit on a single accelerator chip. Each weight tensor in PaLM [11], for example, is partitioned across 3072 TPUv4 chips in a pod. This leads to a huge overhead on transferring the weights and intermediate activations across the datacenter networks.
Latency-critical applications can now hardly benefit from parameter caching. Loading model parameters from DRAM to on-chip accelerator memory often takes a lot of time during inference. In the past, parameter caching was an effective optimization for latency because it reused model weights and avoided off-chip memory transfers. However, evaluations on edge TPUs reported that this method works best for models with fewer than 30 million parameters [36]. For larger models, parameter caching even becomes harmful. Benefits from compiler optimizations are diminishing, and the serving latency becomes almost proportional to the model parameter count. In our case, the smallest translation model has about 50 million parameters. Improving latency thus boils down to increasing memory bandwidth alone.
Quantization can significantly reduce inference cost. Binarization is an extreme case where both the weights and activations of a matrix multiplication (matmul) are quantized to a single bit. Compared to the Brain floating-point format (bfloat16) [1]  $^4$ , binarization reduces the weight size by  $16\times$ , thus significantly lowering the memory and communication overhead. Moreover, a binarized matmul can be carried out by XNOR operations followed by a population count, which is estimated to be  $256\times$  more energy-efficient than the bfloat16 counterpart [39].
Prior work shows that BERT can be binarized for pretraining [5, 33, 28]; however, it is important to note that the BERT and MT models, which both use Transformer as their core [37], are very different. One key difference is the architecture: while an MT model has both an encoder and a decoder, BERT only has an encoder. This difference can impact the quality of encoder quantization because every cross attention layer in the decoder requires outputs from the encoder. Another difference is that MT model inference produces a sequence of text, while BERT performs a single text classification. This is critical because each word in the output translation sequence affects the generation of the next word. The sampling distribution of a word is therefore crucial and should be preserved after binarization, but for BERT, only the peak of the logits needs to be preserved. Due to these differences, directly applying BERT binarization techniques to MT can easily result in a lower quality model.
In this work, we investigate binarized Transformer for neural machine translation, which, to our knowledge, is the first study on this topic. Each Transformer block contains an attention layer and a feed-forward network (FFN). We binarize the weights and activations separately so we can study how each one affects the quality of the model. We found that binarizing weights did not significantly affect accuracy, but that traditional methods for binarizing activations led to poor performance due to activation magnitude explosion. Then, we propose a new method for activation binarization that uses a simple scaling factor and additional residual connections.
To understand the scaling behavior of the proposed 1-bit Transformer in practice, we further evaluate it on our in-house production-scale translation dataset that contains three billion sentence pairs. We for the first time demonstrate that the 1-bit weight Transformer scales and generalizes similarly well as the float one, even on the out-of-domain data. We also analyze sentences sampled from both models' outputs and find that the 1-bit Transformer generates a similar translation quality as its float counterpart. Binarization can therefore be a potential candidate for future MT model serving.
The success of Transformer has spurred an active body of work to quantize it to lower precision. In this section, we review a subset of these efforts that inspired our approach.
4 In the remaining paper, "float" refers to bfloat16.
Transformer quantization. Much of the prior effort focused on 8-bit Transformer. Bhandare et al. [8] reported a less than 0.5 BLEU drop on the WMT14 En-De translation task with 8 bits. Prato et al. [32] showed an 8-bit Transformer preserved the translation quality. For non-generative tasks, Zafrir et al. [38] quantized BERT to 8-bit with marginal quality loss. When pushed down to 4 bits, though Prato et al. [32] reported an 8 BLEU degradation for MT, Aji and Heafield [2] reported almost no BLEU loss by using a logarithmic quantization scheme.
The exploration on 1-bit Transformers centered around BERT. Usually binarization is directly applied and the focus is on improving the training recipe. Bai et al. [5] initiated the attempt by splitting a ternary BERT into a binary one, then fine-tuning. It achieved  $41\%$  average accuracy on the GLUE benchmarks. Qin et al. [33] proposed to distill each intermediate layer outputs from a floating-point model. Recently, Liu et al. [28] proposed to incrementally quantize the model, e.g., from 32-bit to 4-bit to 2-bit, finally to 1-bit, and it improved the GLUE accuracy to  $73.5\%$ .
Binarized vision models. Courbariaux et al. [12] pioneered the investigation on binarized deep neural nets. Recently, PokeBNN [39] established a pareto SOTA on the ImageNet recognition task. We inherit the binarization functions and training recipes from PokeBNN.
Generalizability. Hooker et al. [19] show that compressed models do not generalize well on out-of-domain (OOD) data. We are particularly interested in evaluating BMT under OOD settings and analyze its generalizability.
3 Algorithm and Model Architecture
In this section, we introduce the methodology of binarizing a Transformer-based MT model. We first define the binarization equations, then show that directly applying the equations to Transformer will produce an inferior model quality because of the dot-product variance inflation. A scaling factor is then proposed as a solution to this problem, and we discuss using LayerNorm [4] to replace fixed scaling factors. Finally, we combine and present the architectural changes that are necessary to improve the binarized model quality.
3.1 Binarization Equations
We follow the approach defined in PokeBNN [39] and AQT [24], which includes an important hyperparameter “ $B$ ”. The function of casting floating-point values into binary values is
$$
\operatorname {c l i p} (x, x _ {\min }, x _ {\max }) := \min  (x _ {\max }, \max  (x _ {\min }, x))
$$
$$
x _ {b} := \left(\operatorname {f l o o r} \left(\operatorname {c l i p} \left(\frac {x}{B}, - 1 + \epsilon , 1 - \epsilon\right)\right) + 0. 5\right) \times B
$$
where  $x$  is the input tensor,  $\epsilon$  is a small floating-point number that prevents overflow when taking the floor, and  $B$  is the binarization bound. In the backward propagation, the floor function is ignored, i.e.,  $\frac{\partial \text{floor}(x)}{\partial x} \coloneqq 1$ , known as the straight-through estimator [12]. The gradient of the entire binarization function is then  $\frac{\partial x_b}{x} = 1_{x \in [-B, B]}$ , otherwise zero. The bound  $B$  therefore serves as a hyperparameter that controls the range of the input values that will have non-zero gradients. Note that  $B$  also serves as a scaling factor for the outputs since the binarization function maps  $x \to \left\{-\frac{B}{2}, +\frac{B}{2}\right\}$ . The bound  $B$  can also generalize to a vector, depending on the granularity of binarization. The finest granularity, however, is one bound value for each dot product, i.e., per contraction dimension, so that the binarized matrix multiplication can be accelerated.
For a dense layer in Transformer of the form  $A \cdot W$ , where  $A^{N \times d_{\mathrm{model}}}$  is the input activations and  $W^{d_{\mathrm{model}} \times d_k}$  is the model weights, we instead compute a binarized matmul  $A_b \cdot \bar{W}_b$ . Throughout the experiments we apply binarization bound  $B_W$  and  $B_A$  for weights and activations, respectively.
$$
B _ {W} = \max  (\operatorname {a b s} (W), \operatorname {a x i s} = d _ {\text {m o d e l}}), B _ {A} = \max  (\operatorname {a b s} (A), \operatorname {a x i s} = d _ {\text {m o d e l}})
$$
where axis is the dimension along which max is taken. Using one axis means the bound is per channel and per example [24]. Both  $B_A^{N \times 1}$  and  $B_W^{1 \times d_k}$  are vectors that contain maximum absolute values along the contraction dimension. Note that the weight binarization bound  $B_W$  is static in inference though it is updated in every training iteration. The activation bound  $B_A$  is dynamic.
3.2 Variance Inflation in Binarization
We start by applying the binarization function to feed-forward networks (FFNs), leaving other modules as float. We observe that directly binarizing the weights preserves the model quality, but binarizing the input activations causes the training to not converge in the context of machine translation. To understand the reason of this behavior, we analyze the variance of the dot product magnitude with and without binarization. Our analysis reveals that binarizing both weights and activations will statistically inflate the magnitude, leading to abnormal signal propagation within the neural network [10]. We present the details of this analysis as follows.
Let each weight of a dense layer be randomly initialized and sampled from a zero-mean normal distribution,  $w \sim \mathcal{N}(0, \sigma_w^2)$ . Assume each input activation is independent of the weights and identically distributed as  $a \sim \mathcal{N}(0, \sigma_a^2)$ . After applying the binarization function, both  $w_b$  and  $a_b$  are still centered at zero and have an equal probability of being either  $-\frac{B}{2}$  or  $+\frac{B}{2}$ , namely, they follow the probability mass function defined as  $\operatorname*{Pr}\left(x_b\right) = \begin{cases} \frac{1}{2} & x_b = -\frac{B}{2} \\ \frac{1}{2} & x_b = +\frac{B}{2} \end{cases}$ . Hence the variance of a binarized multiplication is
$$
\operatorname {V a r} \left(a _ {b} \cdot w _ {b}\right) = \mathbb {E} \left[ a _ {b} ^ {2} \right] \cdot \mathbb {E} \left[ w _ {b} ^ {2} \right] - \mathbb {E} ^ {2} \left[ a _ {b} \right] \cdot \mathbb {E} ^ {2} \left[ w _ {b} \right] = \sum_ {a _ {b}} a _ {b} ^ {2} \cdot \Pr \left(a _ {b}\right) \cdot \sum_ {w _ {b}} w _ {b} ^ {2} \cdot \Pr \left(w _ {b}\right) - 0 = \frac {B ^ {4}}{1 6}
$$
The variance of a binarized dot product is then  $\mathrm{Var}(A_b \cdot W_b) = \sum_{n=0}^{D-1} \mathrm{Var}_n(a_b \cdot w_b) = \frac{B^4}{16} \cdot D$ , where  $D$  is the dimensionality of the dot product, i.e., the hidden projection dimension in an FFN, and  $n$  is the index of each entry in the vector.
Following the same analysis, the variance of a floating-point dot-product is  $\operatorname{Var}(A \cdot W) = \sigma_a^2 \cdot \sigma_w^2 \cdot D$ . Note that the commonly used Xavier initializer [17] equalizes the variance of the activations across layers.  $\sigma_w^2$  will therefore be initialized as  $\frac{1}{D}$ , so  $\operatorname{Var}(A \cdot W) = \sigma_a^2$ , which is usually at the scale of 1.
Meanwhile, the common binarization bound is  $B \in [1,3]$  [12, 39, 7]. Our Transformer FFN employs a hidden projection dimension  $D = 4096$  throughout the experiments. Therefore,  $\operatorname{Var}(A_b \cdot W_b) \gg \operatorname{Var}(A \cdot W)$ . Binarization heavily inflates the dot product variance by at least  $256 \times$ , which will be reflected in the magnitude of the dense layer outputs. Also note that  $\operatorname{Var}(A_b \cdot W_b) \propto D$ , indicating that Transformer with a larger width will potentially suffer more from the convergence issue.
3.3 A Scaling Factor as the Solution
Inspired by the scaling factor  $\sqrt{d_k}$  in the scaled dot-product attention Attention  $(Q,K,V) = \mathrm{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V$  in the original Transformer [37], we propose a scaling factor for each binarized dense layer, i.e.,
$$
\mathrm {D e n s e} \left(A _ {b}\right) = \frac {A _ {b} \cdot W _ {b}}{s}
$$
The scaling factor  $s$  is a hyperparameter that suppresses dot-product variance inflation, while in the attention layer  $\sqrt{d_k}$  prevents the dot products from entering small-gradient regions of the softmax function. According to the analysis in Section 3.2, its value is estimated to be  $s \propto \sqrt{D}$  in order to cancel the multiplicative effect from  $D$  on the variance.
To verify how the magnitude of the scaling factor affects the training loss, we sweep  $s$  in Section 5. In practice,  $s \geq 64$  can make the training converge.
3.4 Replacement of Scaling Factor with LayerNorm
While the scaling factor  $s$  enables the binarization of FFNs, it requires hyperparameter tuning, which can be challenging for
billion-parameter translation models. To address this deficiency, we propose using layer normalization
Figure 1: BMT Multi-Head Attention — Differences from the original Transformer are highlighted (in yellow). All linear projections and einsums can be binarized.
(LayerNorm) [4] as a drop-in replacement for the scaling factor, which has the form of  $\mathrm{LN}(x) = \frac{x - \mathbb{E}[x]}{\sqrt{\operatorname{Var}(x) + \epsilon}}\cdot \gamma +\beta$ , where  $\gamma$  and  $\beta$  are learnable parameters. Besides the fact that  $\gamma$  can incorporate the scaling factor  $s$ , LayerNorm also has the following advantages.
The scaling factor is now dynamic and adaptive during training. The binarization function employs a dynamic bound  $B$ , so  $\operatorname{Var}(A_b \cdot W_b)$  varies. The learnable parameter  $\gamma$  in LayerNorm can better capture the changes in the dot product variance and hence properly normalize it.
LayerNorm also redistributes the input activations. It enables the binarization of a tensor with all positive values. A directly binarized FFN is  $\mathrm{FFN}(A) = \max \left(0, A_b W_{1b} + b_1\right)_b W_{2b} + b_2$ , where  $W_1, b_1$  and  $W_2, b_2$  are the weights and biases for the first and second dense layer, respectively. One may note that the activations  $\max \left(0, A_b W_{1b} + b_1\right)$  are all positive. The binarization function will then map the entire tensor to a constant  $+\frac{B}{2}$ , which undermines the model training. With the help LayerNorm, however, the activations are redistributed and more balanced in terms of the number of positive and negative values. This enables the normal  $\{-1, +1\}$  (bipolar) binarization of the second dense layer. Qin et al. [33], Liu et al. [28] used  $\{0, 1\}$  binarization instead in binarized BERT to overcome the issue of constant positive values. It yields a ternary matrix multiplication since  $A \in \{0, 1\}^{N \times D}$  and  $W \in \{-1, +1\}^{D \times K}$ , which incurs nontrivial additional overhead if computed on binary hardware accelerator. The complete proposed 1-bit FFN has the structure of
$$
\operatorname {F F N} (A) = \operatorname {L N} (\operatorname {L N} (\max  (0, A _ {b} W _ {1 b} + b _ {1})) _ {b} \cdot W _ {2 b} + b _ {2})
$$
When proceeding to the attention binarization, we add a LayerNorm to the output of each linear projection layer for the same reasons. We verified in Section 5 that a dynamic and adaptive scaling factor in LayerNorm indeed outperformed a fixed one.
3.5 Residual Connection in Attention Layers
In attention layers, we also add a shortcut connection to the output linear projection layer. Combined with the additional LayerNorm, the output projection then becomes Out  $(A) = \mathrm{LN}(A \cdot W) + A$ . In BNNs, gradients of a binarized layer are approximated due to the straight-through estimator. This will eventually lead the optimization into a different direction as we stack more binarized layers. Liu et al. [26] proposed adding additional residual connections in BNNs, which became a useful method for partially addressing this issue. We therefore adopt it in our model. Note that this modification is unnecessary for QKV (query, key, value) linear projections. The shortcut around the entire attention layer in the original Transform serves the same purpose. We will also demonstrate the effectiveness of the shortcut connection in the ablation study in Section 5.
The complete modified attention architecture is shown in Figure 1, where we highlight the differences from the original one. The extra layer normalization and shortcut connection are both elementwise. Their overhead is small, especially comparing to the benefits of binarization.
In this section, we empirically evaluate our proposed binarized Transformer on MT tasks at difference scales. To investigate the impact of binarizing different layers, we first train a standard 6-layer encoder-decoder (6L6L) Transformer on the WMT2017 De-En translation dataset [9] and evaluate it on the WMT2014 De-En dataset. We then choose the 1-bit weight model variant and study its practical scaling law on in-house translation datasets. We also compare the translation qualities of both 1-bit and float models. Throughout the experiments, the embedding table and the prediction head layer are not binarized.
4.1 WMT Results
We binarize five different matmuls in a Transformer. In an attention layer there are (1) QKV linear projections; (2) activation-activation matmul between queries and keys (QK Einsum); (3) activation-activation matmul between attention scores and values (Score-V Einsum); (4) output linear projection. In an FFN there are two dense layers of the same type. To study their individual impact, we binarize their weights and activations separately. In our experiments we use the following training details.
Model. We use a 6L6L Transformer as the base model. Embedding dimension is 1024. Each multi-head attention layer has 16 heads, with a dimension of 1024 for QKV if combining all the heads. The hidden projection dimension in FFNs is 4096. Dropout layers have a dropout rate of 0.1.
Scheduler. We adopt a three-stage training scheme, where the learning rate (LR) of each stage decreases from base to zero following a cosine decay. A quantization event starts at the beginning of each stage. We first train the model in float. In the second stage, all weights will be binarized. In the last stage, both weights and activations will be binarized.
Table 1: BMT results on the WMT dataset. Training uses WMT2017 De-En and evaluation uses WMT2014 De-En. Binarized activations or weights are labeled by checkmarks. Unlabeled tensors remains bfloat16. 1-bit weights models have 25MB of weight storage size while the float one has 399MB,  $\sim$ $16\times$  compression. As a comparison, the baseline model (last row) directly applies XNOR-Net style [34] binarization used in previous works [33, 28], sign function followed by a normalization. BLEU evaluation employs a beam size of 4.
ATTENTION 1-BIT	FFN 1-BIT	METRICS
AQKV	WQKV	AOUT	WOUT	QK	SCORE-V	AFFN	WFFN	VAL LOSS	BLEU
FLOAT	1.39	26.35
BMT-1	✓	✓	✓	1.38	25.93
BMT-2	✓	✓	1.40	25.44
BMT-3	✓	✓	✓	✓	1.51	24.11
BMT-4	✓	✓	✓	✓	✓	1.72	21.55
BMT-5	✓	✓	✓	✓	✓	1.60	21.06
BMT-6	✓	✓	✓	✓	✓	✓	1.89	17.87
BMT-7	✓	✓	✓	✓	1.76	18.27
BMT-8	✓	✓	✓	✓	✓	2.81	9.42
BASE [33, 28]	✓	✓	8.07	0.21
Training. We apply knowledge distillation (KD) during training. KD replaces the ground truth label in the cross-entropy loss function with the softmaxed logits from the teacher model, so it is optional for users. Adam optimizer [22] is used with  $\beta_{1} = 0.9$  and  $\beta_{2} = 0.98$ . No weight decay is applied. Batch size is 1024. Base learning rate is 0.001. The first LR cycle has 50000 steps, others have 88339 steps. We train the model with a  $4\times 8$  TPU topology.
Observations. The evaluation results on WMT2014 De-En translation dataset is shown in Table 1. We mainly rely on the validation loss for comparing the model quality since BLEU score has a higher variation [16]. From the table we have the following key observations.
Weight-only binarization preserves the model loss. The float 6L6L Transformer baseline has a 1.39 validation loss. In contrast, binarizing all dense layer weights (in both attention layers and FFNs) produces an even lower loss (1.38, BMT-1), though the BLEU score slightly drops by about 0.4. Both metrics indicate that the 1-bit weight model has a similar translation quality to the float baseline. Binarization has the potential to compress the model storage size by  $16 \times$  while preserving the quality.