Understanding ELECTRA and Training an ELECTRA Language Model | by Thilina Rajapakse | Towards Data Science

Introduction
What Is Pre-Training?
Masked Language Modelling (MLM)
ELECTRA Pre-Training Approach
Efficiency Gains of the ELECTRA Approach
Training Your Own ELECTRA Model
Installation
Data preparation
Language Modeling Model
Training the model
Throwing out the generator and getting the discriminator
Fine-tuning the pre-trained model
Wrap-up
The process of training a Transformer model for use in a particular Natural Language Processing task is fairly simple, although it might not be easy. Start with a randomly initialized Transformer model, put together a huge (and I do mean huge) dataset containing text in the language or languages you are interested in, pre-train the Transformer on the huge dataset, and fine-tune the pre-trained Transformer on your particular task, using your task-specific dataset (which may be comparatively tiny).
The advantage of this method is that you only need labelled data for the final step (fine-tuning). Personally, if there’s anything I don’t like about my work with Deep Learning, it’s labelling the data! I’m sure any Deep Learning practitioner would agree that labelling is time-consuming, tedious, and potentially error-prone.
Pre-training is the process through which the Transformer model learns to model a language. In other words, the Transformer will learn good, context-dependent ways of representing text sequences. This knowledge can then be reused in downstream tasks, drastically cutting down the amount of required task-specific, labelled data as the model has already learned the language features and now only needs to fine-tune its representations to perform a specific task. For pre-training, the only requirement in terms of data is a lot of (hopefully) clean data. No labelling necessary! I shudder to think of the ordeal someone would have had to undergo if it had been necessary to label the gigabytes of text used to pre-train BERT.
That’s all well and good, but how exactly does a Transformer learn language representation from a huge dump of unlabelled text? Well, the original BERT model relied on two pre-training tasks, masked language modelling (MLM) and next sentence prediction. In next sentence prediction, the model is tasked with predicting whether two sequences of text naturally follow each other or not. This task was said to help with certain downstream tasks such as Question Answering and Natural Language Inference in the BERT paper although it was shown to be unnecessary in the later RoBERTa paper which only used masked language modelling. Regardless, we are more interested in the first approach, MLM, as this is what the ELECTRA pre-training method aims to improve upon.
In Masked Language Modelling, a certain percentage of the tokens (a token is a term for a single unit of a text sequence. It may be a word or a piece of a word) of an input sequence is masked, and the model is tasked with predicting the original token for the masked tokens. The masked tokens may be replaced with an actual mask token (e.g. [MASK]) or replaced with another random token from the vocabulary (the set of all tokens known to the model).
BERT, armed with this pre-training technique, was capable of shredding previously set benchmarks in many of the downstream NLP tasks. However, the authors of the ELECTRA paper note that MLM approaches only learns from the masked tokens (typically 15%) of any given example. This leads to a substantial increase in the computational resources needed to train a language model with MLM. Another downside of MLM is that the mask tokens only appear in the pre-training stage and never during fine-tuning or downstream usage. This discrepancy also contributes to a slight loss of performance in models trained using MLM.
ELECTRA (Efficiently Learning an Encoder that Classifies Token Replacements Accurately) is a new pre-training approach which aims to match or exceed the downstream performance of an MLM pre-trained model while using significantly less compute resources for the pre-training stage. The pre-training task in ELECTRA is based on detecting replaced tokens in the input sequence. This setup requires two Transformer models, a generator and a discriminator.
Let’s break down the pre-training process step-by-step.
For a given input sequence, randomly replace some tokens with a [MASK] token.The generator predicts the original tokens for all masked tokens.The input sequence to the discriminator is built by replacing [MASK] tokens with the generator predictions.For each token in the sequence, the discriminator predicts whether it is an original or whether it has been replaced by the generator.
For a given input sequence, randomly replace some tokens with a [MASK] token.
The generator predicts the original tokens for all masked tokens.
The input sequence to the discriminator is built by replacing [MASK] tokens with the generator predictions.
For each token in the sequence, the discriminator predicts whether it is an original or whether it has been replaced by the generator.
The generator model is trained to predict the original tokens for masked-out tokens while the discriminator model is trained to predict which tokens have been replaced given a corrupted sequence. This means that the discriminator loss can be calculated over all input tokens as it performs prediction on each token. With MLM, the model loss is only calculated over the masked tokens. This is shown to be a key difference between the two approaches and the primary reason behind ELECTRA’s greater efficiency.
This setup is similar to the training setup of a GAN (Generative Adversarial Network), except that the generator is not trained to try and deceive the discriminator (so it is not adversarial per se). Also, if the generator happens to predict the original token of a masked token correctly, that token is considered to be an original token (as the token has not been corrupted/changed).
The discriminator model is used for downstream tasks and the generator is thrown away after pre-training.
The ELECTRA paper highlights the major improvements of the ELECTRA pre-training approach and also shed light on why those improvements are seen.
All scores are GLUE benchmark scores given in the ELECTRA paper unless stated otherwise.
As mentioned earlier, the discriminator model’s loss is defined over all the tokens in the sequence as it must predict whether each token is an original or a replacement. To demonstrate the significance of this difference, the authors compare the ELECTRA model to an identically trained model (ELECTRA 15%) except that ELECTRA 15% only calculates the discriminator loss over the masked tokens. The GLUE benchmark is used for the comparison.
It is shown that the original ELECTRA approach yields a 85.0 score while ELECTRA 15% gets 82.4. (For comparison, BERT scored 82.2)
The ELECTRA is also compared to another model (All-Tokens MLM) where the masked tokens are replaced with generator predictions AND the model is tasked with predicting the original identity of all tokens in the input. The loss, too, is calculated over all input tokens. The All-Tokens MLM model scores 84.3, outperforming BERT’s 82.2, and nearly catching up to ELECTRA’s 85.0.
This clearly shows that the ability to calculate the loss over all input tokens significantly boosts the performance of a pre-trained model.
For this, the original BERT model is compared to a model (Replace MLM) trained using the MLM pre-training objective except that the masked tokens are replaced with a token from a generator rather than with an actual [MASK] token.
The replace MLM model scores 82.4, slightly outperforming BERT’s 82.2. This difference indicates that the pre-train/fine-tuning discrepancy does slightly harm BERT’s performance.
Comparatively, learning from all input tokens seems to have a much bigger impact than resolving the pre-train/fine-tune mismatch of masked tokens.
Summing up the comparisons in a table;
The performance gains of ELECTRA over BERT is shown to be larger at smaller model sizes.
One huge advantage of the ELECTRA pre-training approach is that it’s possible to train your own language models on a single GPU!
Below, I’ll show you how you can train your own language model with the Simple Transformers library.
Install Anaconda or Miniconda Package Manager from here.Create a new virtual environment and install packages.conda create -n simpletransformers python pandas tqdmconda activate simpletransformersconda install pytorch cudatoolkit=10.1 -c pytorchInstall Apex if you are using fp16 training. Please follow the instructions here. (Installing Apex from pip has caused issues for several people.)Install simpletransformers.pip install simpletransformers
Install Anaconda or Miniconda Package Manager from here.
Create a new virtual environment and install packages.conda create -n simpletransformers python pandas tqdmconda activate simpletransformersconda install pytorch cudatoolkit=10.1 -c pytorch
Install Apex if you are using fp16 training. Please follow the instructions here. (Installing Apex from pip has caused issues for several people.)
Install simpletransformers.pip install simpletransformers
We’ll be training our language model on Esperanto (inspired by Hugging Face’s tutorial here). Don’t worry if you don’t speak Esperanto, neither do I!
For pre-training a model we are going to need a (preferably large) corpus of text in Esperanto. I’ll be using the Esperanto text files from the Leipzig Corpora Collection. Specifically, I downloaded the following datasets;
2011 — Mixed (1M sentences)2012 — Newscrawl (1M sentences)2012 — Web (1M sentences)2016 — Wikipedia (300k sentences)
2011 — Mixed (1M sentences)
2012 — Newscrawl (1M sentences)
2012 — Web (1M sentences)
2016 — Wikipedia (300k sentences)
You should be able to improve the results by using a bigger dataset.
Download the datasets and extract the archives into a directory data/ .
Take all the “sentence” files and move them into the data/ directory.
Script for Linux/bash users. (Others: Why aren’t you on Linux? 😉)
for d in */;  do mv d/*sentences.txt .;done;
If you open one of the files, you’ll notice that they have two columns, with the first column containing the index and the second column containing the text. We just need the text, so we’ll drop the indexes, combine all the text, and split the text into train and test files.
Now we are ready to start training!
In Simple Transformers, all language modelling tasks are handled with the LanguageModelingModel class. You have tons of configuration options that you can use when performing any NLP task in Simple Transformers, although you don’t need to set each one (sensible defaults are used wherever possible).
List of common configuration options and their usage here.
Language modelling specific options and their usage here.
The gist above sets up a LanguageModelingModel which can be used to train our new model.
A single training epoch (with this configuration) takes a little under 2 hours on a Titan RTX GPU. To speed up training, you can increase evaluate_during_training_steps or turn off evaluate_during_training altogether.
When training a language model from scratch with Simple Transformers, it will automatically create a tokenizer for us from the specified train_file. You can configure the size of the trained tokenizer by setting a vocab_size in the train_args. In my case, I’m using a vocabulary of 52000 tokens.
You can also configure the architecture of the generator and the discriminator models as required. The configuration options are set in the two dictionaries generator_config and discriminator_config found in train_args. The ELECTRA paper recommends using a generator model that is 0.25–0.5 of the size of the discriminator. They also recommend decreasing the number of hidden layers and keeping the other parameters constant between the generator and the discriminator. With that in mind, I went with a small (12 layers) architecture for the discriminator and a similar generator albeit with 1/4th the number of hidden layers.
"generator_config": {"embedding_size": 128,"hidden_size": 256,"num_hidden_layers": 3,},"discriminator_config": {"embedding_size": 128,"hidden_size": 256,}
You can find all the architecture configuration options and their default values here.
Now that we’ve set up our models, all we need to do is initiate the training.
Running the above script will start training our language model!
You can take a look at all the training progress information (and a lot of charts!) here on Wandb.
Amazingly, the ELECTRA pre-training approach lets us train a brand-new language model, in a matter of hours, on a single GPU! The chart below tells the full story.
The flattish parts of the curve are where Simple Transformers did the evaluations. Sadly, I forgot to set a higher batch size for evaluation so the evaluations were a lot slower than they should have been.
When using the ELECTRA pre-training approach, the generator is typically discarded after training and only the discriminator is used. To facilitate this, Simple Transformers will save the generator and the discriminator separately once training completes.
However, if you terminate training before it completes, your model will be saved as a LanguageModelingModel which contains both the discriminator and the generator. In order to extract either the discriminator or the generator separately, you can use the methods save_discriminator() or save_generator().
The script below demonstrates how you can load a LanguageModelingModel and extract the generator model separately.
This will load the trained language model found in outputs/best_model into a LanguageModelingModel and then save only the discriminator into discriminator_trained/discriminator_model.
Now that we have a model that can “understand” Esperanto, we can go ahead and fine-tune it on a specific NLP task. For this, we can use the Part-of-Speech tagging dataset in Esperanto provided by Hugging Face (train and labels).
In my case, I downloaded the files and saved them to data/pos-tagging.
Using this we can train our ELECTRA pre-trained model (the discriminator) as a NERModel on the Part-of-Speech dataset. (For more info on Named Entity Recognition with Transformers, you can check out my article here)
Here are the results:
eval_loss = 0.18994220502207695f1_score = 0.9025234264922376precision = 0.9197880191456026recall = 0.8858950119055756
As you can see, the model is capable of performing impressively in the downstream task after only a few hours of training!
And that, folks, is how you can train a brand new language model from scratch with ELECTRA!
I am excited by the possibilities this opens up as ELECTRA should significantly lower the compute resource barrier to training your own language models. Hopefully, we will be seeing more and more language models being developed for all the languages spoken around the world!
As always, a huge shoutout to Hugging Face and their Transformers library, who’s work make all of this possible!