Dynamic Multi-Branch Layers for On-Device Neural Machine Translation

Abstract

With the rapid development of artificial intelligence (AI), there is a trend in moving AI applications, such as neural machine translation (NMT), from cloud to mobile devices. Constrained by limited hardware resources and battery, the performance of on-device NMT systems is far from satisfactory. Inspired by conditional computation, we propose to improve the performance of on-device NMT systems with dynamic multi-branch layers. Specifically, we design a layer-wise dynamic multi-branch network with only one branch activated during training and inference. As not all branches are activated during training, we propose shared-private reparameterization to ensure sufficient training for each branch. At almost the same computational cost, our method achieves improvements of up to 1.7 BLEU points on the WMT14 English-German translation task and 1.8 BLEU points on the WMT20 Chinese-English translation task over the Transformer model, respectively. Compared with a strong baseline that also uses multiple branches, the proposed method is up to 1.5 times faster with the same number of parameters.

Source code is available at https://github.com/THUNLP-MT/Transformer-DMB

References in corpus (8)

• Sequence to Sequence Learning with Neural Networks
• To prune, or not to prune: exploring the efficacy of pruning for model compression
• Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer
• HAT: Hardware-Aware Transformers for Efficient Natural Language Processing
• The Evolved Transformer
• Lite Transformer with Long-Short Range Attention
• Exponentially Increasing the Capacity-to-Computation Ratio for Conditional Computation in Deep Learning
• Controlling Computation versus Quality for Neural Sequence Models

Cited by in corpus (1)

• MeMemo: On-device Retrieval Augmentation for Private and Personalized Text Generation