BERT: Bidirectional Encoder Representations from Transformers

We will implement a simple BERT model from scratch and train it on a small dataset.

Data Preprocessing

Like GPT, BERT trains on raw, unlabeled text — no human-written pairs needed. We reuse the same cached Arabic Wikipedia corpus (1,500 articles) from the GPT experiment, but shape it very differently, since Masked Language Modeling is a fundamentally different task from next-token prediction.

Tokenization: AraBERT’s pretrained WordPiece tokenizer AraBERT’s tokenizer was trained specifically for Arabic. It also comes with [CLS], [SEP], [MASK], and [PAD] already built into its 64,000-token vocabulary.

Shaping the data: masked sequences, not shifted ones. FOr each sequence, we do the following:

80% of the time → replace with MASK_ID
10% of the time → replace with a random token id from the vocab
10% of the time → leave it exactly as it was (no change to input_ids)

Example:

input_ids:  [CLS, A, [MASK], C, D, E, X, G, H, SEP]     ← B became [MASK], F became a random token X

labels:     [-100, -100, B, -100, -100, -100, F, -100, -100, -100]  ← original B and F saved here

Show Code

import torch
import random
import torch.nn as nn

from bert_data_pipeline import train_loader, tokenizer, VOCAB_SIZE, SEQ_LEN, CLS_ID, SEP_ID, MASK_ID, PAD_ID

print("vocab size:", VOCAB_SIZE)
print("seq len:", SEQ_LEN)

input_ids, labels, _ = next(iter(train_loader))
print("input_ids:", input_ids.shape)
print("labels:", labels.shape)

vocab size: 64000
seq len: 128
input_ids: torch.Size([32, 128])
labels: torch.Size([32, 128])

Show Code

from torch.utils.data import DataLoader, random_split
from bert_data_pipeline import train_dataset, BATCH_SIZE

val_size = max(1, int(0.1 * len(train_dataset)))
train_size = len(train_dataset) - val_size

train_subset, val_subset = random_split(
    train_dataset,
    [train_size, val_size],
    generator=torch.Generator().manual_seed(42),
)

# Your dataset generates fresh masks every time an item is fetched.
# Save validation examples once to keep their masks fixed.
random_state = random.getstate()
random.seed(42)
val_examples = [val_subset[i] for i in range(len(val_subset))]
random.setstate(random_state)

train_loader = DataLoader(
    train_subset,
    batch_size=BATCH_SIZE,
    shuffle=True,
)

val_loader = DataLoader(
    val_examples,
    batch_size=BATCH_SIZE,
    shuffle=False,
)

THe above is containing these:

input_ids — the actual token ids fed into the embedding layer. Each row is [CLS] token token [MASK] token ... [SEP]
labels — same shape, but almost entirely -100. Only at the positions that were selected for masking does labels hold the original, uncorrupted token id — the answer the model is being trained to recover.

Model Architecture

BERT is your existing Encoder from the Attention notebook, almost unchanged — same self-attention, residuals, LayerNorm, feed-forward blocks, stacked N times. There are two main differences from the GPT-1 model:

1- No causal mask. The GPT’s causal_mask blocked each position from seeing anything ahead of it. BERT’s self-attention gets no such restriction — every position attends to every other position, both directions, every layer.

2- A new prediction head for MLM. Instead of GPT’s output_layer predicting “what’s the next token,” BERT’s output layer predicts “what was the original token at this position,” but only evaluated at the masked positions.

The flow of data through the BERT model is as follows:

input_ids → embeddings → stack of encoder blocks (bidirectional self-attention) → linear layer projecting to vocab size at every position → loss computed only where labels != -100

Show Code

class BERTEmbedding(nn.Module):
    def __init__(self, VOCAB_SIZE, SEQ_LEN, d_model):
        super().__init__()
        self.token_embedding =  nn.Embedding(VOCAB_SIZE, d_model)
        self.position_embedding = nn.Embedding(SEQ_LEN, d_model)
        self.segment_embedding = nn.Embedding(2, d_model) # to which segment (A/B) the sentance belongs
        
    def forward(self, input_ids, segment_ids):
        batch_size, seq_len = input_ids.shape
        
        positions = torch.arange(seq_len, device=input_ids.device) # [0, 1, 2, ..., 127]
        # This will broadcasts the above row across all 32 sequences in the batch
        positions = positions.unsqueeze(0).expand(batch_size, seq_len)
    
        
        token_embeds = self.token_embedding(input_ids) # (batch_size, seq_len, d_model)
        position_emb = self.position_embedding(positions) # (batch_size, seq_len, d_model)
        segment_emb = self.segment_embedding(segment_ids) # (batch_size, seq_len, d_model)
        
        
        return token_embeds + position_emb + segment_emb

Below is the class definition for a single BERT encoder layer, in which the self-attention is bidirectional (no causal mask) followed by a feed-forward block, with residual connections and layer normalization.

Show Code

class BERTEncoderLayer(nn.Module):
    def __init__(self, d_model, num_heads, ff_hidden):
        super().__init__()
        self.self_attention = nn.MultiheadAttention(d_model, num_heads, batch_first=True)
        self.layer_norm1 = nn.LayerNorm(d_model)
        self.feed_forward = nn.Sequential(
            nn.Linear(d_model, ff_hidden),
            nn.ReLU(),
            nn.Linear(ff_hidden, d_model)
        )
        self.layer_norm2 = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(0.1)
        
    def forward(self, x, attention_mask=None):
        # Self-attention
        attn_output, _ = self.self_attention(x, x, x, key_padding_mask=attention_mask)
        attn_output = self.dropout(attn_output)
        x = self.layer_norm1(x + attn_output)  # Add & Norm
        
        # Feed-forward
        ff_output = self.feed_forward(x)
        ff_output = self.dropout(ff_output)
        x = self.layer_norm2(x + ff_output)  # Add & Norm
        
        return x

THen we stack N of these encoder layers to form the full BERT encoder block. The final output is a sequence of hidden states, one for each input token.

Show Code

class BERTEncoder(nn.Module):
    def __init__(self, d_model, num_heads, ff_hidden, num_layers):
        super().__init__()
        #initialize the encoder layers then wrap them in a nn.ModuleList for N layers
        self.encoder_layers = nn.ModuleList([BERTEncoderLayer(d_model, num_heads, ff_hidden) for _ in range(num_layers)])
        
    def forward(self, x, attention_mask=None):
        for layer in self.encoder_layers:
            x = layer(x, attention_mask=attention_mask)
        return x

Finally, this class is wrapping the embedding layer, the encoder stack, and the output layer (The MLM head) into a single model. The forward pass returns the logits for each position in the sequence, which are then used to compute the loss against the labels.

Show Code

class BERT(nn.Module):
    def __init__(self, vocab_size, seq_len, d_model, num_heads, ff_hidden, num_layers):
        super().__init__()
        self.embedding = BERTEmbedding(vocab_size, seq_len, d_model)
        self.encoder = BERTEncoder(d_model, num_heads, ff_hidden, num_layers)
        self.mlm_head = nn.Linear(d_model, vocab_size)

    def forward(self, input_ids, attention_mask=None):
        segment_ids = torch.zeros_like(input_ids)
        x = self.embedding(input_ids, segment_ids)
        x = self.encoder(x, attention_mask=attention_mask)
        logits = self.mlm_head(x)
        return logits

Show Code

d_model = 128
num_heads = 4
ff_hidden = 512
num_layers = 4

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print("Training on:", device)

model = BERT(
    VOCAB_SIZE, SEQ_LEN, d_model, num_heads, ff_hidden, num_layers
).to(device)

input_ids, labels, attention_mask = next(iter(train_loader))
input_ids = input_ids.to(device)
labels = labels.to(device)
attention_mask = attention_mask.to(device)
with torch.no_grad():
    logits = model(input_ids, attention_mask=(attention_mask == 0))
print("logits:", logits.shape)   # expect (32, 128, 64000)

Training Loop

Show Code

# evaluation function

@torch.no_grad()
def evaluate(model, loader, loss_fn):
    model.eval()

    total_loss = 0.0
    total_targets = 0

    for input_ids, labels, attention_mask in loader:
        input_ids = input_ids.to(device)
        labels = labels.to(device)
        attention_mask = attention_mask.to(device)

        logits = model(
            input_ids,
            attention_mask=(attention_mask == 0),
        )

        num_targets = (labels != -100).sum().item()
        if num_targets == 0:
            continue

        loss = loss_fn(
            logits.reshape(-1, VOCAB_SIZE),
            labels.reshape(-1),
        )

        total_loss += loss.item() * num_targets
        total_targets += num_targets

    return total_loss / total_targets if total_targets else float("nan")

Show Code

import math

loss_fn = nn.CrossEntropyLoss(ignore_index=-100)
optimizer = torch.optim.Adam(model.parameters(), betas=(0.9, 0.98), eps=1e-9)


class NoamScheduler:
    def __init__(self, optimizer, d_model, warmup_steps=1000):
        self.optimizer = optimizer
        self.d_model = d_model
        self.warmup_steps = warmup_steps
        self.step_num = 0

    def step(self):
        self.step_num += 1
        lr = (self.d_model ** -0.5) * min(self.step_num ** -
                                          0.5, self.step_num * self.warmup_steps ** -1.5)
        for group in self.optimizer.param_groups:
            group['lr'] = lr
        self.optimizer.step()


scheduler = NoamScheduler(optimizer, d_model, warmup_steps=1000)

Show Code

num_epochs = 5
loss_history = []
train_epoch_losses = []
val_epoch_losses = []

for epoch in range(num_epochs):
    model.train()  # Re-enable dropout after validation.

    epoch_loss = 0.0
    epoch_targets = 0
    window_loss = 0.0
    window_targets = 0

    for batch_idx, (input_ids, labels, attention_mask) in enumerate(
        train_loader, start=1
    ):
        input_ids = input_ids.to(device)
        labels = labels.to(device)
        attention_mask = attention_mask.to(device)

        num_targets = (labels != -100).sum().item()
        if num_targets == 0:
            continue

        logits = model(
            input_ids,
            attention_mask=(attention_mask == 0),
        )

        loss = loss_fn(
            logits.reshape(-1, VOCAB_SIZE),
            labels.reshape(-1),
        )

        optimizer.zero_grad()
        loss.backward()
        scheduler.step()  # Your scheduler also calls optimizer.step().

        batch_loss = loss.item()
        loss_history.append(batch_loss)

        epoch_loss += batch_loss * num_targets
        epoch_targets += num_targets
        window_loss += batch_loss * num_targets
        window_targets += num_targets

        if batch_idx % 200 == 0 or batch_idx == len(train_loader):
            print(
                f"epoch {epoch + 1} "
                f"batch {batch_idx}/{len(train_loader)} "
                f"average train loss {window_loss / window_targets:.4f}"
            )
            window_loss = 0.0
            window_targets = 0

    train_loss = epoch_loss / epoch_targets
    val_loss = evaluate(model, val_loader, loss_fn)

    train_epoch_losses.append(train_loss)
    val_epoch_losses.append(val_loss)

    print(
        f"epoch {epoch + 1} finished | "
        f"train loss {train_loss:.4f} | "
        f"validation loss {val_loss:.4f}"
    )

Show Code

num_epochs = 1
loss_history = []

for epoch in range(num_epochs):
    model.train()
    for batch_idx, (input_ids, labels, attention_mask) in enumerate(train_loader):
        input_ids = input_ids.to(device)
        labels = labels.to(device)
        attention_mask = attention_mask.to(device)
        logits = model(input_ids, attention_mask=(attention_mask == 0))

        loss = loss_fn(logits.view(-1, VOCAB_SIZE), labels.view(-1))

        optimizer.zero_grad()
        loss.backward()
        scheduler.step()

        loss_history.append(loss.item())

        if batch_idx % 200 == 0:
            print(
                f"epoch {epoch+1} batch {batch_idx}/{len(train_loader)} loss {loss.item():.4f}")

Data Preprocessing

Model Architecture

Training Loop

Comments