Building BERT with PyTorch from scratch
BERT stands for Bidirectional Encoder Representation from Transformers. The original BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding, actually, explains everything you need to know about BERT.
Honestly saying, there are much better articles on the Internet explaining what BERT is, for example, BERT Explained: State of the art language model for NLP. After reading this article, you may have some questions about attention mechanisms; this article Illustrated: Self-Attention explains attentions.
In this paragraph I just want to run over the ideas of BERT and give more attention to the practical implementation.
BERT solves two tasks simultaneously:
- Next Sentence Prediction (NSP) ;
- Masked Language Model (MLM).
Building BERT
To build BERT we need to work out three steps:
- Prepare Dataset;
- Build a model;
- Build a trainer.
BERT
BERT module is a container that combines all the modules together and returns the output.
class BERT(nn.Module):
def __init__(self, vocab_size, dim_inp, dim_out, attention_heads=4):
super(BERT, self).__init__()
self.embedding = JointEmbedding(vocab_size, dim_inp)
self.encoder = Encoder(dim_inp, dim_out, attention_heads)
self.token_prediction_layer = nn.Linear(dim_inp, vocab_size)
self.softmax = nn.LogSoftmax(dim=-1)
self.classification_layer = nn.Linear(dim_inp, 2)
def forward(self, input_tensor: torch.Tensor, attention_mask: torch.Tensor):
embedded = self.embedding(input_tensor)
encoded = self.encoder(embedded, attention_mask)
token_predictions = self.token_prediction_layer(encoded)
first_word = encoded[:, 0, :]
return self.softmax(token_predictions), self.classification_layer(first_word)
We use linear layer (and softmax) with output that is equal to the vocabulary size for token prediction task.
self.token_prediction_layer = nn.Linear(dim_inp, vocab_size)
self.softmax = nn.LogSoftmax(dim=-1)
And linear layer with output of 2 for next sentence prediction task
self.classification_layer = nn.Linear(dim_inp, 2)
The output of the network
argmax(NSP output) = [1, 0] is NOT next sentence
argmax(NSP output) = [0, 1] is next sentence
Everything is simple in the forward. At first we calculate embedding, then pass embeddings to our encoder.
embedded = self.embedding(input_tensor)
encoded = self.encoder(embedded, attention_mask)
Secondly, we calculate the model output.
token_predictions = self.token_prediction_layer(encoded)
first_word = encoded[:, 0, :]
return self.softmax(token_predictions), self.classification_layer(first_word)
The full model graph is also available. To build the graph, run the script graph.py. It saves the graph to data/logs directory. Run tensorboard
tensorboard --logdir data/logs
Open http://localhost:6006 in the browser, go to Graph tab. You should see the graph of our BERT model.
Building BERT
To build BERT we need to work out three steps:
- Prepare Dataset;
- Build a model;
- Build a trainer.