Huggingface🤗Transformers: Retraining roberta-base using the RoBERTa MLM Procedure
By: Tanmay Garg
Tags: Pytorch, NLP, Bert, Transformers, Language Model

A hands-on guide

Transformers: Retraining roberta-base using the RoBERTa pre-training procedure

Photo by Alex Knight on Unsplash

Introduction

RoBERTa

Since BERT (Devlin et al., 2019) came out, the NLP community has been booming with the Transformer (Vaswani et al., 2017) encoder based Language Models enjoying state of the art (SOTA) results on a multitude of downstream tasks.

The RoBERTa model (Liu et al., 2019) introduces some key modifications above the BERT MLM (masked-language modeling) training procedure. The authors highlight “the importance of exploring previously unexplored design choices of BERT”. Details of these design choices can be found in the paper’s Experimental Setup section.

RoBERTa in the 🤗Transformers library

The 🤗Transformers library comes bundled with classes & utilities to apply various tasks on the RoBERTa model.

For example, it is extremely easy to get predictions for a sequence containing a mask token!

Yes! It is that simple!

‘score’: 0.167,‘sequence’: ‘
Send these pictures back!
’

'score': 0.108,'sequence': '
Send these photos back!
'

'score': 0.077,'sequence': '
Send these emails back!
'

...

This signifies what the “roberta-base” model predicts to be the best alternatives for the <mask> token.

Motivation: Beyond the pre-trained models

It is oftentimes desirable to re-train the LM to better capture the language characteristics of a downstream task.

For example, RoBERTa is trained on BookCorpus (Zhu et al., 2015), amongst other datasets.

A recently published work BerTweet (Nguyen et al., 2020) provides a pre-trained BERT model (using the RoBERTa procedure) on vast Twitter corpora in English. They argue that BerTweet better models the characteristic of language used on the Twitter subspace, outperforming previous SOTA models on Tweet NLP tasks.

Hence, it is a good indicator that the performance on downstream tasks is greatly influenced by what our LM captures!

Before we start

What this post is not about

This post does not delve into training the LM and tokenizer from scratch. Training from scratch is quite sufficiently covered in an official 🤗 post here.

Please note that these pre-trained models have been trained using a huge amount of data and computing resources. For example, according to this description, “roberta-base” was trained on 1024 V100 GPUs for 500K steps. Yes, you read that right. 1-0-2-4 GPUs 🤯. Such magnitudes of resources are often not available to us.

However, what we can do is retrain these available models for a few more epochs on a smaller dataset!

Retraining our very own RoBERTa LM 🤖🔥

Alright then, let’s get started.

We will be using transformers v3.5.1, however, this tutorial should work fine with the recently released v4.0.0 as well.

We will be using the Hate Speech Detection dataset (Basile et al., 2019) made available through TweetEval (Barbieri et al., 2020). Of course, you can use a dataset that better suits your downstream task! 🤩

Since our data is already present in a single file, we can go ahead and use the LineByLineTextDataset class.

The block_size argument gives the largest token length supported by the LM to be trained. “roberta-base” supports sequences of length 512 (including special tokens like 
 (start of sequence) and 
 (end of sequence).

For a finer control over the dataset, you can explore 🤗Datasets here.

The data collator object helps us to form input data batches in a form on which the LM can be trained. For example, it pads all examples of a batch to bring them to the same length.

As the names are quite self-explanatory, the TrainingArguments object holds some fields that help define the training process. The Trainer finally brings all of the objects that we have created till now together to facilitate the train process.

seed=1: seeds the RNG for the Trainer so that the results can be replicated when needed.

It took ~100mins for train to finish on Google Colab.

trainer.save_model(output_dir): helps us save the model to the output_dir so that we can load it using from_pretrained (or as done below).

Verify training

Go ahead and verify your LM! You’re in luck if your python environment has TensorBoard available, because the Trainer object logs the training in ‘./runs’ directory.

As before, we can get predictions for a sequence containing a mask token!

‘score’: 0.165, ‘sequence’: ‘
 Send these b*****s back! 
’

‘score’: 0.105, ‘sequence’: ‘
 Send these refugees back! 
’

‘score’: 0.096, ‘sequence’: ‘
 Send these people back! 
’

...

Oof, that’s hurtful! 🙉 I had to censor a word out for our PG-13 audiences 👶

Jokes aside, it also tells that our model has trained correctly and has started capturing the hateful language of the dataset.

Go ahead, tweak the hyper-parameters through the TrainerArguments object to see which setting gives the best results for your downstream task!