Sockeye 2: A Toolkit for Neural Machine Translation

SOCKEYE 2: A Toolkit for Neural Machine Translation Felix Hieber and Tobias Domhan and Michael Denkowski and David Vilar Amazon {fhieber,domhant,mdenkows,dvilar}@amazon.com Abstract We present SOCKEYE 2, a modernized and streamlined version of the SOCKEYE neural machine translation (NMT) toolkit. New features include a simplified code base through the use of MXNet’s Gluon API, a focus on state-of-the-art model ar- chitectures, and distributed mixed preci- sion training. These improvements result in faster training and inference, higher au- tomatic metric scores, and a shorter path from research to production. 1 Introduction SOCKEYE (Hieber et al., 2017) is a versatile toolkit for research in the fast-moving field of NMT. Since the initial release, it has been used in at least 25 scientific publications, including win- ning submissions to WMT evaluations (Scham- per et al., 2018). Based on the deep learning li- brary MXNet (Chen et al., 2015), SOCKEYE also powers Amazon Translate, showing industrial- strength performance in addition to the flexibil- ity needed in academic environments. Moreover, we are excited to see hardware manufacturers con- tributing optimizations to MXNet and SOCKEYE. Intel has demonstrated large performance gains for SOCKEYE inference on Intel Skylake proces- sors.1 NVIDIA is working on significant perfor- mance improvements for SOCKEYE’s transformer (Vaswani et al., 2017) implementation through fused operators and optimized beam search. This c⃝ 2020 The authors. This article is licensed under a Creative Commons 3.0 licence, no derivative works, attribution, CC- BY-ND. 1https://www.intel.ai/amazing-inference- performance-with-intel-xeon-scalable- processors/#gs.wrgsji paper discusses SOCKEYE 2’s streamlined Gluon implementation (§2), support for state of the art ar- chitectures (§3), and improved model training (§4). 2 Gluon Implementation SOCKEYE 2 adopts Gluon, the latest and preferred application programming interface (API) of the MXNet deep learning library. Gluon simplifies the code while improving overall performance. De- velopers can define building blocks of neural net- work architectures as Python classes and seam- lessly switch between eager execution for step-by- step debugging and cached computation graphs for maximum performance. Migration to Gluon sim- plifies training and inference code in SOCKEYE 2, reducing the overall number of lines of Python code by 25%. The hybridized Gluon transformer implementation in SOCKEYE 2 improves training speed by 14%, compared to SOCKEYE. 3 Focus on State-of-the-Art Models Due to the success of self-attentional models we concentrated development of SOCKEYE 2 on the transformer (Vaswani et al., 2017), removing sup- port for recurrent and convolutional models. Us- ing the pre-norm configuration by default allows for learning rate warm-up-free training. We found deep encoders and shallow decoders for transformers to be competitive in BLEU while significantly increasing decoding speed due to computational workload being shifted to the en- coder side. On WMT19 (EN–FI, FI–EN, EN–DE, DE–EN), a 20-encoder and 2-decoder layer config- uration improves on average by 0.2 BLEU over the baseline, while reducing decoding time by 50%. We also improved support for source factors by allowing to tie source factor and word embeddings, DE–EN EN–FI BLEU Time BLEU Time Ott et al. (2018) 34.7 30h 20.1 14h Plateau Reduce 34.9 28h 20.7 …

Content truncated to 500 words