Pytorch implementation of paper: BERT (2018):”BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding”

The 2018 paper “BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding” by Devlin et al. (Google AI) introduced a shift from the standard “left-to-right” language modeling to a deeply bidirectional approach.

Unlike the original Transformer decoder (which masks future words), BERT uses a Masked Language Model (MLM) objective, allowing the representation of each word to fuse context from both its left and right simultaneously.

1. BERT Embeddings

BERT’s input is a summation of three embeddings: Token, Segment (to distinguish between two sentences), and Position (learned, unlike the sinusoidal ones in the original Transformer).

Python

import torch
import torch.nn as nn

class BERTEmbeddings(nn.Module):
    def __init__(self, vocab_size, hidden_size, max_len, n_segments):
        super().__init__()
        self.tok_embed = nn.Embedding(vocab_size, hidden_size)    # token embedding
        self.pos_embed = nn.Embedding(max_len, hidden_size)       # position embedding
        self.seg_embed = nn.Embedding(n_segments, hidden_size)    # segment(token_type) embedding
        self.norm = nn.LayerNorm(hidden_size)
        self.dropout = nn.Dropout(0.1)

    def forward(self, x, segment_label):
        seq_len = x.size(1)
        pos = torch.arange(seq_len, dtype=torch.long).unsqueeze(0).to(x.device)
        
        # BERT Input = Token + Position + Segment
        embedding = self.tok_embed(x) + self.pos_embed(pos) + self.seg_embed(segment_label)
        return self.dropout(self.norm(embedding))

---

2. The Multi-Head Attention (Bidirectional)

The core “Attention” is identical to the Transformer Encoder, but BERT explicitly removes the “causal mask” used in decoders, allowing every token to attend to every other token.

Python

class ScaledDotProductAttention(nn.Module):
    def forward(self, q, k, v, mask=None):
        d_k = q.size(-1)
        scores = torch.matmul(q, k.transpose(-2, -1)) / (d_k ** 0.5)
        
        if mask is not None:
            # Mask out padding tokens (0)
            scores = scores.masked_fill(mask == 0, -1e9)
            
        attn = torch.softmax(scores, dim=-1)
        return torch.matmul(attn, v)

---

3. The BERT Model & Pre-training Heads

BERT is pre-trained on two tasks:

1. Masked LM (MLM): Predict the [MASK] tokens.
2. Next Sentence Prediction (NSP): Predict if sentence B follows sentence A.

Python

class BERT(nn.Module):
    def __init__(self, vocab_size, hidden_size, n_layers, n_heads, d_ff):
        super().__init__()
        self.embedding = BERTEmbeddings(vocab_size, hidden_size, 512, 2)
        
        # Stack of Transformer Encoder Layers
        self.layers = nn.ModuleList([
            EncoderLayer(hidden_size, n_heads, d_ff, 0.1) for _ in range(n_layers)
        ])
        
        # NSP Head: Binary classification (isNext / NotNext)
        self.nsp_classifier = nn.Linear(hidden_size, 2)
        
        # MLM Head: Predict token from vocabulary
        self.mlm_classifier = nn.Linear(hidden_size, vocab_size)
        self.activation = nn.GELU() # BERT uses GELU instead of ReLU

    def forward(self, x, segment_label, mask):
        x = self.embedding(x, segment_label)
        
        for layer in self.layers:
            x = layer(x, mask)
            
        # Get the [CLS] token (first token) for NSP
        pooled_output = torch.tanh(nn.Linear(x.size(-1), x.size(-1))(x[:, 0]))
        nsp_logits = self.nsp_classifier(pooled_output)
        
        # Get all tokens for MLM
        mlm_logits = self.mlm_classifier(x)
        
        return mlm_logits, nsp_logits

---

Key Technical Details

• The [CLS] Token: BERT inserts a special [CLS] token at the start of every sequence. The final hidden state of this token is used as the aggregate sequence representation for classification tasks.
• GELU Activation: BERT popularized the Gaussian Error Linear Unit, which provides a smoother non-linearity than ReLU.
• Pre-training Strategy: 15% of tokens are chosen for masking. Of those, 80% are replaced with [MASK], 10% with a random token, and 10% are kept unchanged to ensure the model doesn’t only learn to focus on the mask.

Why it was a breakthrough

Before BERT, models like Word2Vec or GloVe gave a single static vector for a word (e.g., “bank” always had the same vector). BERT provides contextualized embeddings, meaning the vector for “bank” in “river bank” is completely different from “bank account” because the bidirectional attention sees the surrounding words before producing the vector.