Masked-Language Modeling With BERT | by James Briggs | Towards Data Science

BERT, everyone’s favorite transformer costs Google ~$7K to train [1] (and who knows how much in R&D costs). From there, we write a couple of lines of code to use the same model — all for free.
BERT has enjoyed unparalleled success in NLP thanks to two unique training approaches, masked-language modeling (MLM), and next sentence prediction (NSP).
In many cases, we might be able to take the pre-trained BERT model out-of-the-box and apply it successfully to our own language tasks.
But often, we might need to fine-tune the model.
Further training with MLM allows us to fine-tune BERT to better understand the particular use of language in a more specific domain.
Out-of-the-box BERT — great for general purpose use. Fine-tuned with MLM BERT — great for domain-specific use.
In this article, we’ll go into depth about what MLM is, how it works, and how we can improve our models with it.
MLM consists of giving BERT a sentence and optimizing the weights inside BERT to output the same sentence on the other side.
So we input a sentence and ask that BERT outputs the same sentence.
However, before we actually give BERT that input sentence — we mask a few tokens.
So we’re actually inputting an incomplete sentence and asking BERT to complete it for us.
What is the effect of this? Well, it’s like those questions many of us were given in school — where, given a sentence, we had to fill in the gaps.
In Autumn the ______ fall from the trees.
Do you know the answer? Most likely you do, and you do because you have considered the context of the sentence.
We see the words fall and trees — we know that the missing word is something that falls from trees.
A lot of things fall from trees, acorns, branches, leaves — but we have another condition, in Autumn — that narrows down our search, the most probable thing to fall from a tree in Autumn are leaves.
As humans, we use a mix of general world knowledge, and linguistic understanding to come to that conclusion. For BERT, this guess will come from reading a lot — and learning linguistic patterns incredibly well.
BERT may not know what Autumn, trees, and leaves are — but it does know that given linguistic patterns, and the context of these words, the answer is most likely to be leaves.
The outcome of this process — for BERT — is an improved comprehension of the style of language being used.
So we understand what MLM is doing, but how does this actually work? What are the logical steps that we’ll need to follow in code?
1.We tokenize our text. Just like we usually would with transformers, we begin with text tokenization.
From tokenization we will receive three different tensors:
input_ids
token_type_ids
attention_mask
We don’t need token_type_ids for MLM — and in this example attention_mask is not so important.
For us, the input_ids tensor is most important. Here, we will have a tokenized representation of our text — which is what we will be modifying moving forwards.
2.Create a labels tensor. We’re training our model here, so we need a labels tensor to calculate loss against — and optimize towards.
The labels tensor is simply input_ids — so all we need to do is make a copy.
3.Mask tokens in input_ids. Now that we’ve created a copy of input_ids for labels, we can go ahead and mask a random selection of tokens.
The BERT paper uses a 15% probability of masking each token during model pre-training, with a few additional rules — we’ll use a simplified version of this and assign a 15% probability of each word being masked.
4.Calculate loss. We process the input_ids and labels tensors through our BERT model and calculate the loss between them both.
Using this loss, we calculate the required gradient changes with BERT — and optimize our model weights.
The loss is calculated as the difference between the output probability distributions for each output ‘token’, and the true one-hot encoded labels.
Okay, that’s all great, but how can we demonstrate MLM in code?
We’ll be using HuggingFace’s transformers and PyTorch, alongside the bert-base-uncased model. So, let’s import and initialize everything first:
from transformers import BertTokenizer, BertForMaskedLM
import torch

tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
model = BertForMaskedLM.from_pretrained('bert-base-uncased')

text = ("After Abraham Lincoln won the November 1860 presidential "
        "election on an anti-slavery platform, an initial seven "
        "slave states declared their secession from the country "
        "to form the Confederacy. War broke out in April 1861 "
        "when secessionist forces attacked Fort Sumter in South "
        "Carolina, just over a month after Lincoln's "
        "inauguration.")

Some weights of the model checkpoint at bert-base-uncased were not used when initializing BertForMaskedLM: ['cls.seq_relationship.weight', 'cls.seq_relationship.bias']
- This IS expected if you are initializing BertForMaskedLM from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).
- This IS NOT expected if you are initializing BertForMaskedLM from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).

And now we move into each one of our logical steps, starting with:
1.Tokenization — tokenization is simple, we’ve already initialized a BertTokenizer, all we do now is tokenize our input text.
inputs = tokenizer(text, return_tensors='pt')

inputs.keys()

dict_keys(['input_ids', 'token_type_ids', 'attention_mask'])
inputs

{'input_ids': tensor([[  101,  2044,  8181,  5367,  2180,  1996,  2281,  7313,  4883,  2602,
          2006,  2019,  3424,  1011,  8864,  4132,  1010,  2019,  3988,  2698,
          6658,  2163,  4161,  2037, 22965,  2013,  1996,  2406,  2000,  2433,
          1996, 18179,  1012,  2162,  3631,  2041,  1999,  2258,  6863,  2043,
         22965,  2923,  2749,  4457,  3481,  7680,  3334,  1999,  2148,  3792,
          1010,  2074,  2058,  1037,  3204,  2044,  5367,  1005,  1055, 17331,
          1012,   102]]), 'token_type_ids': tensor([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
         0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
         0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]]), 'attention_mask': tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
         1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
         1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]])}
We won’t worry about padding/truncation for now. What we should be aware of are the three tensors described earlier — token_type_ids and attention_mask don’t need any attention from us — but input_ids does.
2.Create labels — The next step is easy, all we need to do here is clone our input_ids tensor into a new labels tensor. We’ll store this within the inputs variable too.
3.Masking — Now we need to mask a random selection of tokens in our input_ids tensor.
To create our 15% probability of masking any one token, we can use torch.rand alongside a condition of each value < 0.15. Together, these will produce our masking array mask_arr.
Now, we use mask_arr to select where to place our MASK tokens — but we don’t want to place a MASK token over other special tokens such as CLS or SEP tokens (101 and 102 respectively).
So, we need to add an additional condition. A check for positions containing the token ids 101 or 102.
Now that is our masking tensor, to apply it we will first extract the index positions where we find a True value — then use this selection to set values in these positions to 103 (the MASK token id).
Now we can see our MASK tokens represented by 103 in the input_ids tensor above.
4.Calculate Loss — Our final step here no different from the typical model training process.
With both input_ids and labels tensors inside our inputs dictionary, we can pass this to our model and return the model loss.
Cool, we’ve run through all of the essentials — but how would all of this look when fine-tuning a model?
There are two ways, (1) we implement our own version of the training function using everything we have learned so far, or (2) we use HuggingFace’s Trainer.
Trainer is clearly the way to go for an optimized, easy-to-use solution. And we’ll take a look at how we can use it — but first, let’s try implementing it ourselves.
After learning all of this, it would be a waste to not give implementing our own training function a go.
First, we need data. Because we’re just randomly masking a selection of tokens — we can use almost any text. We don’t need labeled or special data.
Here, we’ll use Meditations by Marcus Aurelius, sourced from here and preprocessed a little (clean version).
First, we’ll import/initialize and load our text data.
Then we tokenize — this time we do truncate and pad each sequence — as we have many sequences of different lengths.
Now we clone the input_ids to create our labels tensor.
Next is our masking code, which is a little different this time for two reasons:
Our mask should not include PAD tokens (as before with CLS and SEP).
We have many sequences — not just one.
We can see the value 103 assigned in the same position as the True value is found in the mask_arr tensor.
The inputs tensors are now ready — and we can begin setting them up to be fed into our model during training.
During training, we’ll be using a PyTorch DataLoader to load our data. To use this, we’ll need to format our data into a PyTorch Dataset object.
Now we’re ready to move onto our training loop. Before starting our loop we need to set up three things:
Move the model to GPU/CPU (GPU if available).
Activate model training mode.
Initialize an Adam with weighted decay optimizer.
Now we’re finally set up — we can begin training! We format this as a typical training loop in PyTorch.
And with that we’re done — we’ve implemented our own MLM fine-tuning script.
Switching across to our Trainer implementation — we’ll still need to do everything we did before up to to point that we created our dataset — as Trainer will be expecting this as input for training.
We’ll first define our training arguments, initialize the Trainer — then train!
So the Trainer approach is certainly much simpler — and allows us to implement checkpoints and other features by simply specifying them at initialization time.
Okay, so that’s everything we really need to know to get started with fine-tuning our models using MLM.
There’s a lot to MLM, but the concept and implementations are not too complex — and incredibly powerful.
Using what we’ve learned here, we can take the best models in NLP and fine-tune them to our more domain-specific language use-cases — needing nothing more than unlabelled text — often an easy data source to find.
I hope you enjoyed this article! If you have any questions, let me know via Twitter or in the comments below. If you’d like more content like this, I post on YouTube too.
Thanks for reading!
[1] The Staggering Cost of Training SOTA AI Models (2019), Synced Review
🤖 NLP With Transformers Course
*All images are by the author except where stated otherwise