BERT:Bidirectional Encoder Representations from Transformers

Input/Output Representations

Fig.1 Input Represention

For handling a variety of down-stream tasks, the input representation of BERT is able to unambiguously represent both a single sentence and a pair of sentences in one token sequence. Throughout this work, a “sentence” can be an arbitrary span of contiguous text, rather than an actual linguistic sentence. A “sequence” refers to the input token sequence to BERT, which may be a single sentence or two sentences packed together.

The first token of every sequence is always a special classification token [CLS]. The final hidden state corresponding to this token is used as the aggregate sequence representation for classification tasks. Sentence pairs are packed together into a single sequence and separated with a special token [SEP]. A learned embedding to every token is added to indicate whether it belongs to sentence A or sentence B. The input embeddings are the sum of the token embeddings, the segmentation embeddings and the position embeddings.. The model architecture of BERT is nearly the same as.

The Pytorch implementation is

import torch.nn.functional as F
import torch.nn as nn
import einops

class BERTEmbeddingLayer(nn.Module):
    def __init__(self, vocab_size, max_len, d_model):
        super(BERTEmbeddingLayer, self).__init__()
        self.t = nn.Embedding(vocab_size, d_model)
        self.p = nn.Embedding(max_len, d_model)
        # A: 1, B: 2
        self.s = nn.Embedding(3, d_model, padding_idx=0)

    def forward(self, input, segment):

        device = input.device
        b, seq_len = input.shape[:2]
        pos = einops.repeat(torch.arange(seq_len, dtype=torch.long, device=device), 's -> b s', b=b)
        output = self.t(input) + self.p(pos) + self.s(segment)

        return output

Masked Language Model

In order to train a deep bidirectional representation, BERT takes randomly masked sequences as input (those masked tokens are replaced by [MASK]) and then predicts those masked tokens. In this case, the final hidden vectors corresponding to the mask tokens are fed into an output softmax over the vocabulary.

Although this allows a bidirectional pre-trained model, a downside is that MLM creats a mismatch between pre-training and fine-tuning, since the [MASK] token does not appear during fine-tuning. To mitigate this, those “masked” words are not always replaced by the actual [MASK] token. The training data generator chooses 15% of the token positions at random for prediction.

Assuming the unlabeled sentence is my dog is hairy and during the random masking procedure hairy is choosen to be masked.

• 80% of the time: Replace the word with the [MASK] token, e.g., my dog is hairy -> my dog is [MASK]
• 10% of the time: Replace the word with a random word, e.g., my dog is hairy -> my dog is apple
• 10% of the time: Keep the word unchanged, e.g., my dog is hairy -> my dog is hairy. The purpose of this is to bias the representation towards the actual observed word.

The advantage of this procedure is that the Transformer encoder does not know which words it will be asked to predict or which have been replaced by random words, so it is forced to keep a distributional contextual representation of every input token. Additionally, because random replacement only occurs for 1.5% of all tokens, this does not seem to harm the model’s language understanding capability.

The Pytorch implementation is

class MaskedLanguageModel(nn.Module):
    def __init__(self, vocab_size, d_model):
        super(MaskedLanguageModel, self).__init__()
        self.fc = nn.Linear(d_model, vocab_size)

    def forward(self, input):
        return F.softmax(self.fc(input), dim=-1)

Next Sentence Prediction

Many important downstream tasks such as Question Answering (QA) and Natural Language Inference (NLI) are based on understanding the relationship between two sentences, which is not directly captured by language modeling. In order to train a model that understands sentence relationships, a binarized next sentence prediction task that can be trivially generated from any monolingual corpus is added. Specifically, when choosing the sentences A and B for each pre-training example, 50% of the time B is the actual next sentence that follows A (labeled as IsNext), and 50% of the time it is a random sentence from the corpus (labeled as NotNext)

The next sentence prediction task can be illustrated in the following examples.

Input   =   [CLS] the man went to [MASK] store [SEP]
                  he bought a gallon [MASK] milk [SEP]
Label   =   𝙸𝚜𝙽𝚎𝚡𝚝
Input	=   [CLS] the man [MASK] to the store [SEP]
                  penguin [MASK] are flight ##less birds [SEP]	
Label	=   𝙽𝚘𝚝𝙽𝚎𝚡𝚝

The Pytorch implementation is

class NextSentencePredict(nn.Module):
    def __init__(self, d_model):
        super(NextSentencePredict, self).__init__()
        self.fc = nn.Linear(d_model, 2)

    def forward(self, input):
        return F.softmax(self.fc(input), dim=-1)

BERT Pytorch Implementation

Code of other modules are from this blog NOTE: The activation function of Point-wise Feed-Forward Networks in BERT is GELU instead of ReLU .

class BERT(nn.Module):
    def __init__(self, num_blocks, vocab_size, max_len, d_model, n_heads, d_ff, dropout):
        super(BERT, self).__init__()
        self.ebd = BERTEmbeddingLayer(vocab_size, max_len, d_model)
        self.blks = nn.ModuleList([EncoderBlock(d_model, n_heads, d_ff, dropout) for _ in range(num_blocks)])
        self.mlm = MaskedLanguageModel(vocab_size, d_model)
        self.nsp = NextSentencePredict(d_model)


    def forward(self, input, segment, mask=None):

        if mask is None:
            mask = input > 0

        output = self.ebd(input, segment)
        for blk in self.blks:
            output = blk(output, mask)

        mlm = self.mlm(output)
        nsp = self.nsp(output[:, 0])

        return mlm, nsp

Reference

1. Devlin, J., Chang, M., Lee, K., & Toutanova, K. (2019). BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding. ArXiv, abs/1810.04805. ↩