How to Train an mT5 Model for Translation With Simple Transformers
By: Thilina Rajapakse
Tags: Machine Learning, Artificial Intelligence, Naturallanguageprocessing, Translation

mT5 is a multilingual Transformer model pre-trained on a dataset (mC4) containing text from 101 different languages. The architecture of the mT5 model (based on T5) is designed to support any Natural Language Processing task (classification, NER, question answering, etc.) by reframing the required task as a sequence-to-sequence task.

In other words — text goes in, and text comes out. For example, in a classification task, the input to the model can be the text sequence to be classified, and the output from the model will be the class label for the sequence. For translation, this is even more straight forward. The text that goes in is in one language, and the text that comes out is in another.

Considering the multilingual capabilities of mT5 and the suitability of the sequence-to-sequence format for language translation, let’s see how we can fine-tune an mT5 model for machine translation. For this article, we’ll be training a translation model to translate between Sinhalese (my native language!) and English. It’s quite challenging to train good translation models for low-resource languages like Sinhalese because of, well, the low availability of resources (training data). Hopefully, the multilingual pre-training on a huge dataset (which includes Sinhalese, although not a lot of it) will help the mT5 model compensate for inadequate training data in the form of direct Sinhalese-English (and vice versa) sequences.

We’ll be using the Simple Transformers library (built on the Huggingface Transformers library) to train the mT5 model. The training and testing data will be obtained from The Tatoeba Translation Challenge. The graphs and charts are generated from Weights & Biases, which is supported natively in Simple Transformers for experiment tracking and hyperparameter optimization.

Note: You can find all the code in this article in the examples/t5/mt5_translation directory (link) of the Simple Transformers Github repo.

Outline

Installing Simple Transformers Downloading a dataset for translation Training the model Evaluating the model — Calculating the BLEU score Wrap up

Setup

You can find the most up-to-date installation instructions in the Simple Transformers documentation.

1. Install Anaconda or Miniconda Package Manager from here.

2. Create a new virtual environment and install packages.

conda create -n st python pandas tqdm sacrebleu

conda activate st

3. If using CUDA:

conda install pytorch>=1.6 cudatoolkit=10.2 -c pytorch

else:

conda install pytorch cpuonly -c pytorch

4. Install simple transformers.

pip install simpletransformers

Data Preparation

The training and test can be obtained from the Tatoeba Translation Challenge data page. You’ll also find datasets for a whole host of other languages there (including Sindarin, a language spoken by Elves in The Lord of the Rings 😀).

If you want to try training a translation model for another language, you can download the dataset for that language instead of Sinhalese. All the other steps in this article apply to any language dataset.

If you are too lazy to search for the dataset, here’s the direct link. 😜

Download the (compressed) translation dataset and extract the archive. Extract the train.trg.gz and train.src.gz archives (yes, the training data is in archives inside archives). Check whether you have the train.trg, train.src, test.trg, test.src files in the data/eng-sin/ directory. (If not, it’s probably easiest to move the files to this location as the code examples will assume that these files can be found here)

Now, let’s build the tsv files that we will use to train and test our mT5 model.

Running the code above will write the two files, train.tsv and eval.tsv , to the data/ directory.

Model Training

Once we have the data files, we are ready to start training the model.

First, we’ll import the necessary stuff and set up logging.

Next, we set up our training and evaluation data.

Here, we remove the prefix values from the datasets because we expect the model to infer the required task based on the input. If the input is in English, then it should be translated to Sinhalese. If it’s in Sinhalese, then it should be translated to English. The model shouldn’t need a prefix to figure this out after the training!

You can use a prefix value to tell an mT5 (or T5) to perform a specific task. This is quite useful to train a model which can perform multiple tasks, as shown in the article below.

Sidenote on GPU memory usage

The amount of GPU memory required to train a Transformer model depends on many different factors (maximum sequence length, number of layers, number of attention heads, size of the hidden dimensions, size of the vocabulary, etc.). Out of these, the maximum sequence length of the model is one of the most significant.

For the self-attention mechanisms used in the mT5 model, the memory requirement grows quadratically with the input sequence length (O(n²) space complexity). I.e., When the sequence length doubles, the memory required quadruples.

Also, mT5 has a much larger vocabulary than T5 (~250,000 tokens to ~32,000 tokens), contributing to mT5 being quite punishing in terms of GPU memory required.

The takeaway from all this is that the number of tokens we can input to the model (the maximum sequence length) comes at a hefty premium. Based on this, it’s wasteful to use even a small number of tokens on the prefix if the model can do without.

Now, let’s get back to training the model!

Here, we specify how we want the model to be set up and initialize the pre-trained mT5 model according to model_args .

I’m using a maximum sequence length ( max_seq_length ) of 96 and a train/eval batch sizes of 20. Generally, larger batch sizes mean better GPU utilization, and therefore, shorter training times. As mentioned earlier, longer sequences require more GPU memory, which means smaller batch sizes and longer training times. The maximum sequence length of 96 allows the model to work with reasonably long text (typically a few sentences) while also keeping the training time practical.

Note that you may need to tweak these values to train the model on your own GPU. If you run out of GPU memory (CUDA memory error), try reducing the batch sizes and/or the maximum sequence length.

If you want to try the fine-tuned model, you can find it here on the Huggingface model hub.

Now, to run the training, we just need to call the train_model() method.

As easy as that! The fine-tuned model will be saved to the outputs directory at the end of the training (see docs for more info on model saving).

With these settings, the model took a little over 10 hours to complete the training on an RTX 3090 (24 GB VRAM).

Time taken for training — by author

I probably could have gotten away with slightly larger batch sizes (as you can see below), but I didn’t want to run the risk of the training crashing as I was running this overnight!

GPU memory usage during training — by author

Playing it safe with a batch size of 20 meant that the GPU was not fully utilized, but, 80%-ish is not bad!

GPU utilization (cores) — by author

Visualizing Training Progress

Setting the wandb_project value in model_args tells Simple Transformers to log the training progress to Weights & Biases automatically. You can find all the logged data for this experiment here.

Training loss

Graph of training loss — by author

Evaluation loss

Graph of evaluation loss — by author

The actual loss values here don’t tell us too much, but the fact that they are decreasing does mean that the model is learning!

In fact, it appears as though the model has not yet converged as the evaluation loss is still decreasing. Training for another epoch or two might very well improve the model performance, but, that would take another 10 or 20 hours!

To try out the fine-tuned model in a web-based GUI (Streamlit app), use the terminal command simple-viewer .

Evaluating the Model

The standard metric used to evaluate and compare machine translation models is the BLEU score, specifically, the BLEU scheme used by the annual Conference on Machine Translation (WMT). The SacreBLEU library can be used to calculate this score.

For more information on the BLEU score, please refer to this paper by Matt Post.

Since The Tatoeba Challenge also provides the BLEU scores for the benchmark translation models, we can easily compare our model to the benchmark model.

Now, let’s load our fine-tuned model and see how it stacks up!

We import the necessary stuff (note the sacrebleu library) and initialize the model just as we did for training, except that we load the fine-tuned model from outputs/ rather than the pre-trained model.

We also set some parameters for generating text (decoding) with the model. Here, the max_length is the maximum length for the output from the model rather than the input.

If you’d like to learn more about the decoding process, please refer to the decoding algorithms section in this article and this excellent notebook by Huggingface.

Next, we’ll prepare the data for evaluation.

Here, we load the evaluation data and prepare separate lists of inputs and true translations (for English to Sinhalese and vice versa).

With the model and evaluation data loaded and ready, we can go ahead and do the translations. With Simple Transformers, we just call model.predict() with the input data. Then, we use the sacrebleu tool to calculate the BLEU score.

The sacrebleu library should be installed in your virtual environment if you followed the setup instructions. If not, you can install it with pip install sacrebleu .

Running this gives us the following scores (rounded off):

English to Sinhalese: 10.3

Sinhalese to English: 24.4

Both these scores improve upon the scores posted by the translation model in the Tatoeba Challenge!

Wrapping Up

The mT5 model does an excellent job of translating between Sinhalese and English, despite the limited training data availability.

mT5 outperforms the scores posted in the Tatoeba Challange. However, it should be noted that the benchmark model in the challenge is trained on several other languages in addition to Sinhalese and English. In addition, the mT5 model requires significantly more compute resources to train and to use. On the other hand, the mT5 model has the potential to improve upon the current scores with more training.

Finally, you can find the fine-tuned model on the Huggingface model hub here. You can use it directly with Simple Transformers as shown below.