Train Neural Machine Translation Models with Sockeye
Have you ever wondered how you can use machine learning (ML) for translation? With our new framework, Sockeye, you can model machine translation (MT) and other sequence-to-sequence tasks. Sockeye, which is built on Apache MXNet, does most of the heavy lifting for building, training, and running state-of-the-art sequence-to-sequence models.
In natural language processing (NLP), many tasks revolve around solving sequence prediction problems. For example, in MT, the task is predicting a sequence of translated words, given a sequence of input words. Models that perform this kind of task are often called sequence-to-sequence models. Lately, deep neural networks (DNNs) have significantly advanced the performance of these models. Sockeye provides both a state-of-the-art implementation of neural machine translation (NMT) models and a platform to conduct NMT research.
Sockeye is built on Apache MXNet, a fast and scalable deep learning library. The Sockeye codebase leverages unique features from MXNet. For example, it mixes declarative and imperative programming styles through the symbolic and imperative MXNet APIs. It also uses data parallelism to train models on multiple GPUs.
In this post, we provide an overview of NMT, and then show how to use Sockeye to train a minimal NMT model with attention.
To understand what’s going on under the hood in Sockeye, let’s take a look at the neural network architecture that many academic groups and industry commonly use.
The network has three major components: the encoder, the decoder, and the attention mechanism. The encoder reads the source sentence one word at a time until the end of sentence () and produces a hidden representation of the sentence. The encoder is often implemented as a recurrent neural network (RNN), such as a long short-term memory (LSTM) network.
The decoder, which is also implemented as an RNN, produces the target sentence one word at a time, starting with a beginning-of-sentence symbol (). It has access to the source sentence through an attention mechanism that generated a context vector. Using the attention mechanism, the decoder can decide which words are most relevant for generating the next target word. This way, the decoder has access to the entire input sentence at all times.
The next word that the network generates becomes an input to the decoder. The decoder produces the subsequent word based on the generated word and its hidden representation. The network continues generating words until it produces a special end-of-sentence symbol,
Sockeye implements state-of-the-art sequence-to-sequence models in MXNet. It also provides appropriate default values for all of the hyperparameters for sequence-to-sequence models. For optimization, you don’t need to worry about stopping criteria, metric tracking, or weight initialization. You can simply run the provided training command line interface (CLI).
You can easily change the basic model architecture, including the following elements:
Sockeye also supports more advanced features, such as:
For training, Sockeye gives you full control over important optimization parameters. For example, you can set the optimizer types, learning rate, momentum, weight decay, and early-stopping conditions. Sockeye tracks multiple metrics (including MT-specific metrics, like BLEU) on the training and validation data.
We plan to continuously extend the Sockeye feature set to provide researchers a platform for experimenting with new ideas for NMT.
Now, let’s train our first NMT model. We expect that the following commands are run on the shell of any Unix like operating system like Linux or Mac OS X.
First, acquire a parallel corpus.
