Sockeye 2: A Toolkit for Neural Machine Translation (2020)

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 architectures, and distributed mixed precision training. These improvements result in faster training and inference, higher automatic metric scores, and a shorter path from research to production.

## introduction sockeye @xcite 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 winning submissions to wmt evaluations @xcite . based on the deep learning library mxnet @xcite , sockeye also powers amazon translate, showing industrialstrength performance in addition to the flexibility needed in academic environments. moreover, we are excited to see hardware manufacturers contributing optimizations to mxnet and sockeye. intel has demonstrated large performance gains for sockeye inference on intel skylake processors. 1 nvidia is working on significant performance improvements for sockeye's transformer @xcite 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. 1 https://www.intel.ai/amazing-inferenceperformance-with-intel-xeon-scalableprocessors/#gs.wrgsji paper discusses sockeye 2's streamlined gluon implementation ( §2), support for state of the art architectures ( §3), and improved model training ( §4). ## focus on state-of-the-art models due to the success of self-attentional models we concentrated development of sockeye 2 on the transformer @xcite , removing support for recurrent and convolutional models. using 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 encoder side. on wmt19 (en-fi, fi-en, en-de, de-en), a 20-encoder and 2-decoder layer configuration 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 @xcite 34.7 30h 20.1 14h plateau reduce 34.9 28h 20.7 12h ## training improvements sockeye 2 significantly accelerates training with horovod 2 integration @xcite ) and mxnet's automatic mixed precision (amp). horovod extends synchronous training to any number of gpus (including across nodes) while amp automatically detects and converts parts of the model that can run in reduced-precision mode (fp16) without loss of quality. sockeye also provides a data-driven alternative to the popular "inverse square root" learning schedule used by @xcite and @xcite : "plateau reduce" keeps the same learning rate until validation perplexity does not increase for several checkpoints, at which time it reduces the learning rate and rewinds all model and optimizer parameters to the best previous point. training concludes when validation perplexity reaches an extended plateau. in a wmt19 benchmark @xcite , plateau reduce training produces stronger models in slightly less time than the setup described by @xcite . results are presented in ## conclusion sockeye 2 provides out-of-the-box support for quickly training strong transformer models for re-2 https://github.com/horovod/horovod 3 https://github.com/awslabs/sockeye search or production. extensive configuration options and the simplified code base enable rapid development and experimentation. we invite the community to contribute their ideas to sockeye 2 and hope that the new programming model and performance improvements enable others to conduct effective and successful research.