MobileNMT: Enabling Translation in 15MB and 30ms
Year: 2023

Abstract:
Deploying NMT models on mobile devices is essential for privacy, low latency, and offline scenarios. For high model capacity, NMT models are rather large. Running these models on devices is challenging with limited storage, memory, computation, and power consumption. Existing work either only focuses on a single metric such as FLOPs or general engine which is not good at auto-regressive decoding. In this paper, we present MobileNMT, a system that can translate in 15MB and 30ms on devices. We propose a series of principles for model compression when combined with quantization. Further, we implement an engine that is friendly to INT8 and decoding. With the co-design of model and engine, compared with the existing system, we speed up 47.0× and save 99.5% of memory with only 11.6% loss of BLEU. The code is publicly available at https://github.com/zjersey/Lightseq-ARM.

heading: 1 Introduction
text: As a classic subfield of natural language processing, neural machine translation (NMT) has achieved great success in recent years. Most of the studies focus on improving the accuracy of large machine translation systems, ignoring whether such models are easy to be deployed in real-world scenarios.
Here we adopt four metrics to evaluate whether an NMT model is deployment-friendly. (1) Model size is the most important metric in model compression (Han et al., 2016). (2) Floating-point operations (FLOPs) is commonly used to evaluate computational complexity in neural architecture design. (3) Memory or Memory mapped I/O (MMI/O) reflects the memory requirements of the real running system. (4) Decoding speed depends on many realistic factors such as engine implementation and the power of avaliable processors.
∗This work is done during the internship at ByteDance. †Corresponding author.
In this paper, we propose MobileNMT, a Transformer-based machine translation system that can translate in 15MB and 30ms. First, we propose three principles for designing parameter-limited MT models: 1) To compress embedding, reducing vocabulary size is simple and effective compared to embedding factorization; 2) To compress the encoder and decoder, reducing the model width is much more efficient in computation and memory than cross-layer parameter sharing; 3) Encoder depth is very important to ensure accuracy. To achieve higher accuracy, we adjust the training hyperparameters according to the newly designed structure, and adopt sequence-level knowledge distillation. For industrial deployment, we optimize general matrix multiplication (GEMM) and memory in our own inference engine and use the 8-bit integer for storage and computation. As shown in Table 1, the 10MB MobileNMT achieves 88.4%
368
Scaling :: V :::::::: performs ::::: better :::: than :::::: Scaling ::: E. ::::: Width :::::::: performs ::::: nearly ::: the ::::: same :::: with ::::::: Sharing.
performance of Transformer-big with only 1.1% size and runs 47.0× faster on decoding, which can be easily deployed and used.
Our contributions are summarized as follows:
• We propose three principles for parameterlimited MT models to make more efficient use of computation and memory resources.
• We adjust training strategies according to the newly designed structure to achieve higher translation accuracy.
• We develop a mobile inference engine to bridge the gap between industrial practice and theoretical research.

heading: 2 Architecture Design Principles
text: For model compression and acceleration, most studies focus on a single metric such as model size or FLOPs, without considering the real-world applications. In this section, we consider four metrics including model size, FLOPs, memory usage, and decoding speed, and then propose three design principles for parameter-limited MT models. We choose Transformer (Appendix A) as our baseline because of its great success in machine translation.

heading: 2.1 Embedding Compression
text: The vocabulary size V usually reaches tens of thousands in NMT models (Akhbardeh et al., 2021). The parameters can reach tens of millions and greatly affect the overall parameter efficiency.
Embedding Factorization (Scaling E). For model compression, embedding factorization has been widely studied (Lan et al., 2020; Grave et al., 2017; Baevski and Auli, 2019). To decouple the
embedding dimension E and hidden dimension H , it additionally introduces a trainable transformation weight W T ∈ RE×H , where E ≤ H . After factorization, the embedding parameters will be decreased from O(V ×H) to O(V ×E+E×H).
Reducing Vocabulary Size (Scaling V). A more direct way to compress embedding is to reduce the vocabulary size V . To reduce the risk of out-ofvocabulary words, here we adopt Byte-Pair Encoding (BPE) (Sennrich et al., 2016; Ott et al., 2018; Ding et al., 2019; Liu et al., 2020). For most studies on machine translation, the adopted BPE merge operations range from 30∼40K (Ding et al., 2019). Volt proves that we can find a well-performing vocabulary with higher BLEU and smaller BPE merge operations (Xu et al., 2021). Experiments in Lin et al. (2021)’work also show that smaller vocabularies may be better.
Reducing Vocabulary Size Performs Better. To compare the two embedding compression methods, here we select three baseline models of different sizes. The model settings are shown in Table 1. As shown in Table 2, the parameters and FLOPs are almost the same in these two methods. As shown
::::: figures ::::: show :::: that ::::::: reducing :::::: model ::::: width :: is :::: more ::::::::::::::::: quantization-friendly :::: than :::::::::: cross-layer :::::::: parameter ::::::: sharing.
For :::::::::: embedding
::::::::::: compression, :::::::: reducing :::::::::: vocabulary :::: size :::::::: (Scaling ::: V) : is ::::: more ::::::: simple :::: and :::::::: effective. ::::: For :::::::::::::: encoder/decoder ::::::::::: compression, ::::::: reducing :::::: model :::::: width ::::::: (Width) :: is ::::: more ::::::: efficient :: in :::::::::: computation ::: and :::::::: memory.
in the first row of Fig. 2, compared to reducing vocabulary size, the model with embedding factorization performs poorly in most cases, especially when the parameters are limited.

heading: 2.2 Encoder/Decoder Compression
text: For encoder and decoder compression, here we compare models with cross-layer parameter sharing and model width reduction.
Cross-Layer Parameter Sharing (Sharing). The most widespread use of parameter sharing is in convolutional neural networks (Long et al., 2015). In recent years, it has also been investigated on NLP and NLU tasks. Among them, cross-layer parameter sharing can provide stronger nonlinearity along the model depth while keeping the parameters unchanged (Dehghani et al., 2019; Takase and Kiyono, 2021; Lan et al., 2020).
Reducing Model Width (Width). Since model depth has been proven to be important in natural language processing tasks such as machine translation (Devlin et al., 2019; Liu et al., 2020; Wang et al., 2022; Liu et al., 2020), here we keep the
Encoder ::::: depth
::: can :: be ::::::::: considered :: as ::: the ::::: most :::::::: important ::::::::: dimension.
depth unchanged and reduce the model width. Reducing Model Width is More Efficient and Quantization-Friendly. In the second row of Fig. 2, these two methods perform nearly the same. However, Table 2 shows that there is a large difference in FLOPs and MMI/O, which means reducing model width is much more efficient in computation and memory. Since it is necessary to quantize these models for greater compression, we further compare the weights and output ranges of the two methods in Fig. 3. It can obviously be observed that models with parameter sharing have larger ranges of values for both weight and output, which is not quantization-friendly. The right figure also verifies this: when we apply post-training quantization (PTQ) (Sung et al., 2015; Banner et al., 2019; Choukroun et al., 2019) to these two methods, cross-layer parameter sharing performs poorly.

heading: 2.3 Deep Encoder and Shallow Decoder
text: Fig. 4 studies how different dimensions affect the Transformer performance. In order to analyze the impact of each dimension separately, here we only change one specific dimension and keep the others
unchanged. The point on the left of the Small Baseline represents scaling one dimension down, while the point on the right represents scaling one dimension up. We can see that Encoder Depth is more important than other dimensions, which is consistent with the related work on large-scale models (Wang et al., 2019, 2022). Based on the above discussion, we finally build a deep encoder and a shallow decoder, while reducing the vocab size and model width. Two MobileNMT models of different sizes are built here and the detailed settings are shown in Table 3.

heading: 3 Training Strategies

heading: 3.1 Pre-Training with Knowledge Distillation
text: In order to improve the performance of compressed models, recent studies distill knowledge from a well-trained full-precision teacher network to a student network (Mishra and Marr, 2018) or directly use a quantized teacher network (Kim et al., 2019). Here we adopt sequence-level knowledge distilla-
Removing ::::::: dropout ::::: from ::::::::::: MobileNMT :::: can :::: lead
:: to :::::::: significant ::::::::::: performance :::::::::::: improvement. :::::: While ::::: larger ::::::: learning :::: rates ::: can :::: also ::::::: improve ::::: model ::::::::::: performance, ::: the ::::: model :::: will :::::: become :::::::::::::::::::: quantization-unfriendly.
::::: Larger ::::::: learning :::: rate ::: will ::::: result :: in ::::: larger ::::: range :: of :::::: values.
tion because it has shown to be effective for NMT tasks. The most basic full-precision Transformerbase model is adopted as the teacher.

heading: 3.2 Quantization
text: The process of quantizing a transformer model can be divided into two steps: 1) constructing quantizers; 2) applying the quantization-aware training (QAT) (Courbariaux et al., 2015) based on the pretrained model we have obtained in Section 3.1.
FFN and Attention Quantizers. The original Transformer layer includes two types of sublayers: the attention sublayer and feed-forward network (FFN) (Vaswani et al., 2017). Here we construct the quantizer for each linear in the attention and FFN, and quantize both the weights and activations as shown in Fig. 5. Since most computations are spent on matrix multiplication, all biases and residuals are kept in full precision for accuracy preservation. Since quantization will change the range of network outputs, here we add a learnable weight γi to the i-th sublayer to learn how to combine the output and the residual surrounding it.
Quantization-Aware Training. Since MobileNMT only has 10M/20M parameters, quantizing such a small model inevitably results in performance loss, so we perform QAT after constructing
:::: more :::::::::::::::::: quantization-friendly.
the quantizers. Before QAT, we pre-compute all scaling parameters based on a forward running on the pre-trained distillation model obtained in Section 3.1. It takes nearly no additional costs, but provides a good initialization. For engineering development, we choose the uniform quantization scheme because of it is hardware-friendly (Liu et al., 2022). For 8-bit quantization, we use the element-wise quantization (Lee et al., 2021). For lower-bit quantization, such as 4-bit integer, we use the row-wise quantization (Faraone et al., 2018).

heading: 3.3 Training Hyperparameters
text: Compared to the original Transformer model, MobileNMT introduced in Section 2 has fewer parameters and different architectures, so different training hyperparameters are needed.
Removing Dropout. Since our models have fewer parameters, we do not need to impose strong regularizations on them and we remove dropout from the entire model. The left part of Fig. 6 shows that removing dropout will lead to an improvement of almost two BLEU points.
Larger Learning Rate. Here we follow the configuration provided in Wang et al. (2019) with a larger learning rate (0.01→ 0.02), a larger training batch (4096 → 8192), and more warmup steps (4000 → 8000). As shown in the right part of Fig. 6, it can improve model performance by more than 0.5 BLEU points (red bars). However, after PTQ, the model with 0.02 learning rate performs significantly worse than 0.01 (blue bars). As shown in Fig. 7, the network weights and outputs become larger when using a larger learning rate, which is not quantization-friendly. L2 Regularization. To solve the above problem, this paper adopts L2 regularization applied to weight (also called weight decay). It adds the squared magnitude of the network weights as the penalty term to the original loss function and en-
courage the weights to be smaller. As shown in the left two parts of Fig. 8, with L2 regularization, both the network weights and output values will become significantly smaller. The right part of Fig. 8 shows the performance of PTQ when applying different degrees of L2 regularization. The red and blue bars represent the model performance before and after PTQ. We can see that L2 regularization does improve the model performance after PTQ.

heading: 4 The Engine
text: This section introduces the detailed implementations of our inference engine.

heading: 4.1 GEMM Optimization
text: According to statistics on the ONNX Runtime platform, general matrix multiplication (GEMM) accounts for 80.44% of the overall decoding time, demonstrating that optimizing GEMM is the key to decoding speed up. We optimize GEMM from three aspects: (1) Replacing 32-bit floating points
with 8-bit integers in GEMM for model quantization. (2) The Arm instruction set we use allows multiple integers to be processed in parallel in a single instruction, which takes full advantage of the processor throughput. (3) To improve the cache hit and the register usage, we adjust the layout of the tensor in memory to ensure that the instruction reads data from continuous space. Specifically, we convert each 4× 4 block in the original layout into a contiguous vector of size 16. An example can be seen in Fig. 9.

heading: 4.2 Memory Optimization
text: As shown in Fig. 10 in the appendix C, except for GEMM, other operations account for only 19.56% of the decoding time but will be frequently performed, resulting in a large amount of temporary memory. To improve memory efficiency, we take two strategies: (1) To avoid frequent memorymapped I/O and footprint, our engine integrates all adjacent fine-grained operations between two GEMM operations into one fused operation. (2) To save temporary memory, different operations are allowed to share the same space, provided that these operations do not interfere with each other at the same time. Through memory sharing, only two 8-bit memory buffers, and one 32-bit buffer need to be pre-allocated in the Transformer encoder to hold intermediate results.

heading: 5 Experiments

heading: 5.1 Setups
text: We evaluate our methods on two WMT benchmarks. For the WMT14 En-De task (4.5M pairs), we choose newstest-2013 as the validation set and
newstest-2014 as the test set. For the WMT14 EnFr task (35M pairs), we validate the system on the combination of newstest-2012 and newstest2013, and test it on newstest-2014. Details of the architecture were introduced in Section 2, and training hyperparameters were introduced in Section 3. For model compression ratio and decoding speed up, we choose Transformer-big as the benchmark (1.0×). Other details of experimental setups are introduced in Appendix D.

heading: 5.2 Results
text: Table 4 shows the results of different systems on WMT14 En-De and En-Fr. Table 5 shows the comparison of MobileNMT with other parameter-efficient methods based on Transformer. MobileNMT-10MB and MobileNMT-20MB are
two models we have built with different sizes, which are introduced in Table 3.
On WMT14 En-De, our MobileNMT-10MB requires only 4.6% of the parameters to maintain 88.4% performance of Transformer-big, while it achieves 87.2× compression ratio and 47.0× speed up. Our MobileNMT-20MB can maintain 95.5% performance of Transformer-big with only 9.2% parameters, while it achieves 43.6× compression ratio and 27.7× speed up. Experiments on En-Fr show similar results. In addition, thanks to the memory optimization strategies adopted in our engine, MobileNMT requires significantly less running memory than other models (0.5%∼0.9% of Transformer-big). All these experiments demonstrate that MobileNMT is efficient in terms of parameters, computation, and memory, and can be easily deployed on mobile devices.

heading: 6 Conclusion
text: We propose MobileNMT, a Transformer-based machine translation system that can translate in 15MB and 30ms. It uses existing resources efficiently and can be easily deployed in real-world scenarios. We develop a mobile inference engine with GEMM and memory optimization, hoping that it can bridge the gap between theoretical research and real-world applications on efficient machine translation.

heading: Acknowledgments
text: This work was supported in part by the National Science Foundation of China (No. 62276056), the National Key R&D Program of China, the China HTRD Center Project (No. 2020AAA0107904), the Natural Science Foundation of Liaoning Province of China (2022-KF-16-01), the Yunnan Provincial Major Science and Technology Special Plan Projects (No. 202103AA080015), the Fundamental Research Funds for the Central Universities (Nos. N2216016, N2216001, and N2216002), and the Program of Introducing Talents of Discipline to Universities, Plan 111 (No. B16009).
Limitations
Multilingual Translation. Here we mainly discuss the design principles of efficient architectures for bilingual machine translation. Compared with bilingual translation, multilingual translation tasks require significantly more parameters and computations to perform well, and different model scales
may lead to different design considerations. We will leave this for future exploration. Knowledge Distillation. As a small model that requires only 10MB/20MB of storage, MobileNMT will inevitably suffer from performance loss compared to other Transformer-based models. To reduce performance loss, here we adopt knowledge distillation and choose the Transformer-base model as the teacher. From a training efficiency perspective, although the teacher model can help MobileNMT improve performance, it also introduces additional training costs. Compatibility. Here our inference engine only provides implementation for the ARM CPU. We will make it available for other AI accelerator (such as NPU) on mobile devices in the future.

heading: A Transformer Architecture
text: We chose Transformer for study because it is one of the most successful neural models for machine translation. It consists of a N -layer encoder and a M -layer decoder, where N=M=6 in the original Transformer-base and Transformer-big. Each encoder layer consists of two sublayers, including the self-attention and feed-forward network (FFN). Each decoder layer has an additional crossattention sublayer to bridge the encoder and decoder.
The self-attention takes the output X of the previous sublayer as its input. The cross-attention is similar to the self-attention, except that it takes the encoder output as an additional input. Both types of attention first compute the attention distribution Ax and then averageX byAx. We denote the transformation matrices ofQ,K, V asWq,Wk,Wv, the subsequent transformation matrices as Wo, and the attention as Ya = Attn(X), then:
Ax = SoftMax( XWqW
T k X T
√ d
) (1)
Ya = AxXWvWo (2)
The FFN applies non-linear transformation to its input X . We denote the FFN as Yf = FFN(X):
Yf = ReLU(XW1 + b1)W2 + b2 (3)
where W1 and b1 denote the weight and bias of the first linear transformation, W2 and b2 are parameters of the second transformation.
Here we preprocess each sublayer input by the layer normalization (Ba et al., 2016). All sublayers are coupled with the residual connection (He et al., 2016a).

heading: B PTQ and QAT
text: As an appealing solution to model compression, quantization enables the model to use lower-bit values (such as 8-bit integer) to compute faster and consume less storage space (Hubara et al., 2016; Micikevicius et al., 2018; Quinn and Ballesteros, 2018; Jacob et al., 2018).
Post-Training Quantization (PTQ) can be seen as the basis for Quantization Aware Training (QAT), it adds quantization nodes to a well-trained floatingpoint model. To quantize a floating-point tensor r to a tensor with n bits, a scale s is introduced to map these two types of values (Wu, 2020):
s = max(r)−min(r)
2n − 1 (4)
colors refer to Model Architecture in Section 2, Training Strategies in Section 3 and Greedy Search.
To get a faster computation speed, both weights and activations will be quantized to n-bit. Suppose rm = min(r), the quantization function is:
Q(r) = b(r − rm)/se × s+ rm (5)
where b·e represents rounding to the nearest integer. However, in PTQ, applying quantization directly to the floating-point network will result in significant performance losses. Based on PTQ, QAT simulates the behavior of n-bit computation by minimizing quantization errors during training, which helps the model achieve higher accuracy. In addition to the learnable weights of the model itself, s is also learnable.

heading: C Operations except GEMM
text: Since general matrix multiplication (GEMM) accounts for 80.44% of the overall decoding time, we have concluded that optimizing GEMM is the key to decoding speed up in Section 4. As for operations except GEMM, Fig. 10 shows the proportion of running time in the decoding process. The corresponding data is measured in 32-bit floating point format on the ONNX Runtime.

heading: D Setups
text: All sentences were segmented into sequences of sub-word units (Sennrich et al., 2016). In the implementation, we adopt the normalization before layers (Baevski and Auli, 2019; Xiong et al., 2020; Nguyen and Salazar, 2019). Most previous work only shared source and target vocabularies on the En-De task. In our MobileNMT, both En-De and En-Fr adopt shared vocabularies for efficiency reasons, which leads to a larger compression gain at the expense of performance. We test on the model ensemble by averaging the last 5 checkpoints and report SacreBLEU scores (Post, 2018).
0 1 2 3 4 5 6 7
·10−2
IDENTITY SLICE
CUMSUM EXPAND WHERE
MULTI_HEAD HOST_CONVERTER
SHAPE SPLIT
UNSQUEEZE CONCAT
CONSTANT_OF_SHAPE MULTI_HEAD_TRANS
BINARY_MUL REDUCE
BINARY_UNKNOWN GATHER
CONVERTER BINARY_ADD
LAYER_NORM ARGMAX
DECODER_ATTENTION DECODER_SELF_ATTENTION ENCODER_SELF_ATTENTION
2.09 · 10−6 2.55 · 10−6 4.79 · 10−6 1.22 · 10−5 2.45 · 10−5 5.52 · 10−5 6.28 · 10−5 7.1 · 10−5 9.49 · 10−5 1.13 · 10−4 1.52 · 10−4 1.91 · 10−4 2 · 10−4 2.01 · 10−4 3.57 · 10−4 4.13 · 10−4 4.63 · 10−4
2.99 · 10−3 4.88 · 10−3 5.36 · 10−3
1.88 · 10−2 1.96 · 10−2
3.76 · 10−2 5.89 · 10−2
Percentage [%]
For the experiments of MobileNMT in Table 4, we use the greedy search algorithm in our engine. Compared with beam search, greedy search can lead to more efficient decoding. For the experiments of TFLite in Table 4, since TFLite will expand all loop subgraphs, it is hard to support the entire decoding process (30 steps) of the Transformerbig/base model with limited memory (6GB in Google Pixel 4). For the memory of these two models, we only record the running memory of 1 step. For the corresponding latencies, we estimate the 30-step latency according to the 1-step and 5- step latencies. It is worth noting that except for the memory and latency on Transformer-big/base, all other data statistics are measured in real-world.
E Analysis
E.1 Ablation Study Table 6 summarizes how each part of Section 2 and Section 3 affects the overall performance. Each row in Table 6 represents the result of applying the current part to the system obtained in the previous row.
To reduce the model parameters from 65M to 10M, the model performance decreased from 27.40 to 22.54, which illustrates the importance of network parameters on model capacity. We observe that both knowledge distillation and tuning hyperparameters can bring significant performance improvements (from 22.54 to 25.48), which effectively compensate for the performance loss caused by parameter reduction.
E.2 Quantization Study Table 7 studies how performance changes when quantizing the model to lower bits (i.e., 4-bit, 3-bit, and 2-bit). As introduced in Section 3.2, for 8-bit quantization, we use the element-wise quantization method (Lee et al., 2021). For lower-bit quantization, we use the row-wise quantization for accuracy preservation (Faraone et al., 2018).
As shown in Table 7, 8-bit and 4-bit quantization have almost no negative effect on model performance. When quantizing the model to lower bits, such as 3-bit and 2-bit integers, model performance will drop dramatically.

heading: Future Work
text: We will leave this for future exploration. Knowledge Distillation. As a small model that requires only 10MB/20MB of storage, MobileNMT will inevitably suffer from performance loss compared to other Transformer-based models. To reduce performance loss, here we adopt knowledge distillation and choose the Transformer-base model as the teacher. From a training efficiency perspective, although the teacher model can help MobileNMT improve performance, it also introduces additional training costs. Compatibility. Here our inference engine only provides implementation for the ARM CPU. We will make it available for other AI accelerator (such as NPU) on mobile devices in the future