• on-device ai
  • quantization
  • transformer

Sub-3bit transformers for fast on-device translation

How binary-code-based quantization and BiQGEMM enable sub-3-bit Transformer inference on mobile devices. Originally written in 2023 based on a 2020 paper.

Code: github.com/insoochung/transformer_bcq

TL;DR

Binary-code-based quantization (BCQ) enables extremely low-bit quantization. Paired with BiQGEMM, it also supports fast on-device inference by multiplying quantized weights directly with floating-point activations.

For reference, Chung & Kim et al. (2020) demonstrated 3.5x speedup, 8.3x runtime memory reduction, and 11.8x model size compression with sub-3bit transformers on the on-device translation task while preserving BLEU scores.

Models are larger than ever

Cartoon showing a large neural network model outside a smartphone shop.
Figure 1. Large models can quickly exceed the constraints of personal devices.

As models grow, they become harder to run on personal devices.

Using cloud-hosted ML products can create privacy tradeoffs.

The cost of training large language models also creates barriers for small companies, concentrating access to frontier models within a few well-funded organizations.

That can limit who gets to experiment with these models and build new features around them.

Compression helps make models easier to deploy, but it is a multi-objective optimization problem: reduce inference latency, runtime memory usage, and model size while maintaining accuracy. Binary-code-based quantization offers one way to balance those goals.

This note looks at how BCQ supports transformer inference on personal devices.

Sub 3-bit transformer quantization

Chung & Kim et al. (2020) achieved sub-3-bit quantization for transformers with a BLEU score loss of at most 0.5 across three translation tasks. Furthermore, BCQ also delivers a 3.5x speedup and an 8.3x reduction in runtime memory footprint compared to full precision 32-bit models.

BCQ offers distinct advantages over the well-known uniform 8-bit quantization:

  1. Low bit weights can be efficiently multiplied with full precision activation, resulting in substantial speed gains with minimal impact on task performance. In contrast, uniform quantization requires activation quantization for quantized inference, often leading to higher quantization errors and latency overhead.
  2. BCQ enables extremely low bit quantization, as empirically demonstrated in Chung & Kim et al. (2020).

High inference overhead of uniform 8-bit quantization

To see why this helps, first consider the overhead of uniform quantization. Uniform 8-bit quantization maps floating-point weights to 8-bit integers. This process involves several steps for each weight matrix: determining the min-max value range, followed by mapping each FP value within the matrix to a corresponding value ranging from 0 to 255, or sometimes -127 to 127.

The objective is to ensure that the smallest value is mapped to 0 and the largest value is mapped to 255.

Uniform 8-bit quantization maps floating point values between min and max to integers from 0 to 255.
Figure 2. Uniform 8-bit quantization maps the floating-point range into 256 integer buckets.
Equation for uniform 8-bit quantization.
Figure 3. Quantization equation for the uniform 8-bit mapping.

For the inference paths considered here, the difficulty is multiplying INT8 weights with FP32 activations. The two options are:

  1. Runtime weight dequantization
  2. Quantized matmul by quantizing FP32 activation to INT8

Runtime weight dequantization converts the stored weights back to FP32. The matrix multiplication then uses FP32 weights and activations, so compression alone does not guarantee a speedup.

Quantized matmul by quantizing FP32 activation to INT8 introduces notable problems. First, quantizing the FP32 activation introduces additional noise to the overall quantization scheme, potentially leading to further degradation in task performance. Secondly, the matmul operation between INT8 weights and activation values results in an INT32 output matrix. To ensure compatibility, the output range should be pre-established, which often involves approximations that may not be fully accurate. See Bhandare et al. (2019) for an example of this direction in Transformer NMT.

Additionally, the INT32 result must be converted back to INT8 to be passed on to the next layer, typically involving a two-stage data conversion process: INT32 -> FP32, or FP16, -> INT8, further increasing overall latency.

Illustration of activation quantization, integer matrix multiplication, and output conversion.
Figure 4. Quantized matmul with uniform quantization requires activation quantization and output conversion.

The lack of a direct “as-is” multiplication scheme between INT8 weights and FP activation values forces uniform quantization’s inference to suffer from high memory and latency overhead and additional quantization noise.

Furthermore, studies have highlighted the varying significance of computations in full precision. As a consequence, achieving an optimal trade-off between accuracy and latency often requires a combination of full precision and quantized operations. Uniform quantization’s lack of flexible mixed precision matmul approach poses a challenge in this regard, limiting the ability to strike the desired balance.

Binary code based quantization

BCQ represents floating-point weights as 1-bit matrices B and floating-point scales grouped in alpha vectors. Each scale in an alpha vector corresponds to a row of the original weight matrix W. For instance, when applying a 2-bit BCQ to an FP matrix with shape m x n, we obtain two m x n 1-bit matrices B1 and B2, accompanied by two FP vectors alpha1 and alpha2 of length m.

Example of 2-bit BCQ weights represented as binary matrices and alpha scale vectors.
Figure 5. Example of 2-bit BCQ weights using binary matrices and floating-point scale vectors.
Binary tree showing values representable by 3-bit binary-code based quantization.
Figure 6. Values representable by 3-bit BCQ through binary basis combinations.

There are several ways to approximate W using B matrices and alpha vectors. A greedy approximation works as follows:

  1. Initialize the residual matrix R to be equal to the original FP weight matrix W.
  2. Determine the parameters of the Bk matrix by assigning either +1 or -1 based on the sign of the corresponding value in W.
  3. Compute the values in alphak such that it minimizes the difference between R and alphak * B, where each parameter in alpha is multiplied to a row in B.
  4. Update R by subtracting the product of B and alpha: R = R - alphak * Bk.
  5. Repeat steps 2-4 n times for n-bit BCQ.

Conversely, the dequantization process would be:

W=α1B1+α2B2++αkBkW' = \alpha_1 B_1 + \alpha_2 B_2 + \cdots + \alpha_k B_k

Efficient “as-is” matmul between binary weights and FP values

One advantage of BCQ is its efficient mixed-precision multiplication between quantized weights and FP activation values during inference. This process eliminates the need for runtime data conversion, thereby avoiding redundant overhead and the introduction of additional errors.

BiQGEMM uses dynamic programming to precompute and reuse partial sums. Consider the product of quantized weights WW' and a floating-point vector xx:

y=Wx=α1(B1x)+α2(B2x)++αk(Bkx)y = W'x = \alpha_1(B_1 \cdot x) + \alpha_2(B_2 \cdot x) + \cdots + \alpha_k(B_k \cdot x)

As discussed, the elements in BB are either 1 or -1, indicating that the output elements of BxB \cdot x are always derived from the set of all possible add-subtract combinations of xx‘s elements. Exploiting this property, we can 1. divide the BxB \cdot x workload into sub-regions, and 2. precompute all possible outcomes of BxB' \cdot x' per sub-region and reuse these values for efficient mixed precision matrix multiplication.

Splitting binary matrix B and input vector x into subregions for precomputation.
Figure 7. Splitting B and x into subregions enables efficient precomputation of reusable partial sums.

Consider the hypothetical scenario presented in Figure 7, where the B matrix has dimensions of 12 x 8, and x is a vector consisting of 12 elements. To obtain the resulting vector y, with a length of 8, each element in the i-th row of B is multiplied by the corresponding element in the x vector. To optimize the computation, we can divide x and B into three segments, as indicated by the white and gray highlights in the figure.

This allows us to express y as the sum of three sub-problems: B’ dot x’, B” dot x”, and B''' dot x'''. By addressing these sub-problems, each involving a subregion of x and B, we can efficiently perform the original matrix multiplication using 3 * 2^4 pre-computed intermediate values.

Once all the intermediate values are computed, the computation of B dot x simplifies to a sequence of indexing and addition operations, exploiting the reuse of the same intermediate values multiple times for improved efficiency. Furthermore, the subsequent multiplication between scale vectors and B dot x is straightforward.

The same approach extends to matrix-matrix multiplication. BiQGEMM also reuses negated pre-computation values to halve the precomputation work and uses bit-packed indices for faster indexing of pre-computed values. For more detailed explanations of these optimizations, refer to the original paper.

As demonstrated by Chung & Kim et al. (2020), BiQGEMM’s efficient matrix multiplication method provides fast and efficient inference, rendering BCQ a highly beneficial approach for on-device transformer use cases, such as translation. Also, the ability to perform matrix multiplication between quantized weights and FP activation values allows BCQ to have a flexible approach, such as only quantizing layers that are not subject to large loss in task performance when perturbed.

Recovering from high quantization error

When mapping FP values to extremely low-bit datatypes, such as sub-3 bits, significant quantization errors are likely to occur, resulting in a sharp decline in task performance. To address this issue, an effective approach is to incorporate a period of non-regularization, or pNR. By considering quantization as a form of weight regularization, we recognize that it limits the potential range of model weights within the loss surface.

Introducing pNR allows the target model to freely explore the loss space, occasionally enforcing the quantization constraint to strike a balance between quantization requirements and optimal task performance. For instance, in the case of transformers, a simple implementation could involve training the model with FP weights and applying quantization every 1000 batches to evaluate its performance. The figure below illustrates the training process with pNR.

Training diagram showing period of non-regularization with intermittent quantization constraints.
Figure 8. Applying a period of non-regularization during training.

Experiments show pNR constants ranging from 500 to 2000 yielded satisfactory results for transformers, with little variation between them. The use of a large pNR value implies that the quantization overhead is rarely incurred, resulting in only a marginal increase in overall training time.

Fine-grained compression approach for transformers

Different parts of a transformer tolerate different amounts of compression. Some layers are more sensitive to quantization errors than others. Embedding vectors also see very different amounts of use because word frequencies vary widely.

As a result, applying a uniform compression approach to all word vectors in an embedding may not yield optimal results. Compression rates and latency improvements also vary across different parts of transformers. For instance, the autoregressive inference scheme requires multiple decoding steps but only a single encoding step.

Addressing these challenges, Chung and Kim (2020) propose a quantization scheme that assigns varying bit precision to different parts of the transformer. This fine-grained quantization approach can be categorized based on the level of granularity in the quantization scheme.

Block-wise scheme

Table of FLOPs and Galaxy N10+ latency contribution for encoder, decoder, and embedding blocks.
Table 1. FLOPs and on-device Galaxy N10+ latency contribution of each Transformer block.

The contribution of each block in a transformer, encoder, decoder, and embedding, to the total latency is depicted in Table 1. While the encoder block demonstrates minimal overall overhead, the decoder and embedding-related functions account for the majority of the inference time. FLOPs alone do not predict latency: parameter movement and cache utilization also affect execution time.

Layer-wise scheme

Table showing each layer's sensitivity to quantization at different bit widths.
Table 2. Each layer’s sensitivity to quantization.

Table 2 illustrates the varying sensitivity of different layers to quantization. Certain layers exhibit higher sensitivity, resulting in a steep degradation in BLEU score when subjected to quantization errors. Based on this, Chung and Kim (2020) propose a quantization scheme that assigns more bits to sensitive layers and fewer to robust layers.

In general, higher bit precision is assigned to encoder sub-layers compared to decoder sub-layers, considering the block-wise contribution to the overall latency.

Vector-wise scheme

Word occurrence frequency in WMT19 translation datasets.
Figure 9. Frequency of word occurrence in WMT19 translation datasets.
BLEU score comparison between conventional method and fine-grained quantization method.
Figure 10. BLEU score comparison between same-bit assignment and fine-grained precision assignment.

In real-world text, words exhibit a power law frequency distribution, where certain words are highly frequent while others are rarely encountered. Building upon this observation, Chung and Kim (2020) propose a vector-wise scheme that assigns varying numbers of bits based on word frequency. Specifically, a group of the most frequent words is allocated a higher number of bits, such as 4 bits, while subsequent groups receive progressively fewer bits, such as 3, 2, and 1 bits.

This fine-grained approach enables efficient compression of the embedding block while preserving accuracy. Notably, Figure 10 demonstrates that employing this fine-grained approach maintains BLEU scores more effectively at lower average bit-precision levels.

Application: on-device translation with sub-3bit transformer

Table showing BLEU scores and model sizes for BCQ-ed Transformers.
Table 3. BCQ-ed Transformers show high compression while maintaining translation performance near the FP baseline.

Table 3 shows the model sizes and translation scores reported by Chung & Kim et al. (2020).

Table showing on-device Galaxy N10+ peak memory and average latency for BCQ models.
Table 4. On-device Galaxy N10+ latency and memory footprint of BCQ models.

The 2.7-bit transformer implementation brings substantial improvements. It accelerates inference by 3.5x, achieves an 8.3x runtime memory compression, and reduces model size by 11.8x, improving on-device translation efficiency.

Conclusion

BCQ supports extremely low-bit quantization while allowing binary weights to be multiplied directly with floating-point activations. Later work, such as Park et al. (2022) and Kwon et al. (2022), explores related BCQ-style ideas for larger generative models.

P.S. Implementation of BCQ is not as complex as it reads. This repository demonstrates the implementation of 3-bit BCQ on a transformer model for a translation task. It may give you some insight into how you can apply BCQ to your own project. Hope this helps.

References

* Equal contribution.