BERT: Bidirectional Encoder Representations

Chapter Overview

BERT (Bidirectional Encoder Representations from Transformers) revolutionized NLP by introducing effective bidirectional pre-training. This chapter covers BERT's architecture, pre-training objectives (masked language modeling and next sentence prediction), fine-tuning strategies, and variants (RoBERTa, ALBERT, DistilBERT).

Learning Objectives

1. Understand BERT's encoder-only architecture
2. Implement masked language modeling (MLM)
3. Apply BERT to downstream tasks via fine-tuning
4. Compare BERT variants and their improvements
5. Analyze BERT's learned representations
6. Understand limitations and failure modes

BERT Architecture

Model Specification

BERT represents a fundamental shift in how we approach natural language understanding by using bidirectional context throughout the entire model. Unlike autoregressive language models such as GPT that can only attend to previous tokens, BERT's encoder-only architecture allows each token to attend to all other tokens in the sequence simultaneously. This bidirectional attention enables BERT to build richer contextual representations that capture both left and right context, making it particularly effective for understanding tasks like question answering, named entity recognition, and text classification.

The architecture comes in two primary configurations that trade off between computational efficiency and model capacity. BERT-base uses 12 transformer encoder layers with hidden dimension $d_{\text{model}} = 768$, employing 12 attention heads where each head operates on dimension $d_k = d_v = 64$. The feed-forward network in each layer expands to dimension $d_{ff} = 3072$, following the standard 4× expansion ratio. This configuration contains approximately 110 million parameters and was designed to be trainable on a modest cluster of TPUs while still achieving strong performance across a wide range of tasks.

BERT-large doubles the model depth to 24 layers and increases the hidden dimension to $d_{\text{model}} = 1024$ with 16 attention heads of dimension $d_k = d_v = 64$ each. The feed-forward dimension scales proportionally to $d_{ff} = 4096$, maintaining the 4× expansion ratio. With approximately 340 million parameters, BERT-large achieves substantially better performance on challenging tasks but requires significantly more computational resources for both training and inference. The parameter count scales roughly quadratically with hidden dimension due to the $d_{\text{model}}^2$ terms in attention projections and feed-forward networks.

BERT-base specifications:

• Layers: $L = 12$
• Hidden size: $d_{\text{model}} = 768$
• Attention heads: $h = 12$, head dimension: $d_k = d_v = 64$
• Feed-forward size: $d_{ff} = 3072$
• Maximum sequence length: $n_{\max} = 512$
• Vocabulary size: $V \approx 30{,}000$ (WordPiece)
• Total parameters: $\approx 110$M

BERT-large specifications:

• Layers: $L = 24$
• Hidden size: $d_{\text{model}} = 1024$
• Attention heads: $h = 16$, head dimension: $d_k = d_v = 64$
• Feed-forward size: $d_{ff} = 4096$
• Maximum sequence length: $n_{\max} = 512$
• Vocabulary size: $V \approx 30{,}000$ (WordPiece)
• Total parameters: $\approx 340$M

Parameter Breakdown and Memory Requirements

Understanding BERT's parameter distribution is essential for memory planning and optimization. BERT-base totals $\approx$110M parameters: embeddings account for $\sim$21\% ($\sim$23M for token embeddings alone), while the 12 encoder layers contribute $\sim$78\% at 7.1M parameters per layer (see Section~[ref] for the complete worked analysis including FLOPs, activation memory, and hardware timing).

The memory footprint depends critically on numerical precision. In standard FP32 (32-bit floating point), each parameter requires 4 bytes, so BERT-base occupies $110{,}000{,}000 \times 4 = 440$ MB. Modern training typically uses mixed precision with FP16 or BF16 (16-bit formats) for activations and gradients while maintaining FP32 master weights for numerical stability. This reduces the working memory for forward and backward passes to $110{,}000{,}000 \times 2 = 220$ MB for the model parameters, though the optimizer still maintains FP32 copies. For inference, pure FP16 weights require only 220 MB, enabling BERT-base to run comfortably on consumer GPUs with 8-16 GB of memory.

BERT-large's parameter distribution follows the same structure but scales significantly. Token embeddings remain at 23 million parameters (vocabulary size unchanged), but each layer now contains $4 \times 1024^2 = 4{,}194{,}304$ attention parameters and $2 \times 1024 \times 4096 = 8{,}388{,}608$ feed-forward parameters, totaling approximately 12.6 million parameters per layer. With 24 layers, the transformer stack contributes 302 million parameters. The total of 340 million parameters requires 1.36 GB in FP32 or 680 MB in FP16. This larger footprint means BERT-large training typically requires GPUs with at least 16 GB of memory (such as NVIDIA V100 or A100), and inference benefits from GPUs with 12+ GB to accommodate reasonable batch sizes.

Input Representation

Token Embeddings: WordPiece tokenization, vocabulary $\approx 30{,}000$

Segment Embeddings: Distinguish sentence A vs B (for sentence-pair tasks)

Position Embeddings: Learned absolute positions (not sinusoidal)

Special tokens:

• [CLS]: Start of sequence, used for classification
• [SEP]: Separate sentences
• [MASK]: Masked token for MLM
• [PAD]: Padding

Pre-Training Objectives

Masked Language Modeling (MLM)

Masked Language Modeling represents BERT's core pre-training objective and the key innovation that enables bidirectional pre-training. Unlike traditional left-to-right language modeling where the model can only condition on previous tokens, MLM randomly masks a subset of input tokens and trains the model to predict the original tokens based on bidirectional context. This approach allows BERT to learn deep bidirectional representations that capture both left and right context simultaneously, which proves crucial for understanding tasks.

The masking strategy employs a carefully designed 80-10-10 split that prevents the model from simply memorizing the training data or developing trivial solutions. When a token is selected for masking (15\% of all tokens), it is replaced with the special [MASK] token 80\% of the time, creating the primary training signal. However, if the model only ever saw [MASK] tokens during training, it would learn to rely exclusively on this special token and fail to generalize during fine-tuning, where [MASK] tokens never appear. To address this mismatch, 10\% of selected tokens are replaced with random tokens from the vocabulary, forcing the model to maintain robust representations even when the input contains noise. The remaining 10\% of selected tokens are kept unchanged, encouraging the model to preserve meaningful representations for all tokens rather than only attending to masked positions.

Objective:

The computational cost of MLM is substantial but manageable. For each training example, only 15\% of tokens contribute to the loss, meaning the model must process approximately 6.7 times as many tokens as a standard language model to see the same number of training signals. However, this cost is offset by the ability to process entire sequences in parallel rather than autoregressively. For a sequence of length 512 with 15\% masking, approximately 77 tokens are masked per example. The prediction head for each masked token performs a matrix multiplication with the vocabulary matrix $\mW_{\text{vocab}} \in \R^{d_{\text{model}} \times V}$ followed by softmax over $V \approx 30{,}000$ tokens, requiring $2 \times 77 \times 768 \times 30{,}000 \approx 3.6$ billion FLOPs per example—a small fraction of the 96.6 billion FLOPs required for the BERT-base forward pass itself.

Why this masking strategy?

• 80\% [MASK]: Standard masking
• 10\% random: Prevents over-reliance on [MASK] token
• 10\% original: Encourages model to maintain representations

Next Sentence Prediction (NSP)

Training data:

• 50\%: B actually follows A (label: IsNext)
• 50\%: B is random sentence (label: NotNext)

NSP Loss:

Total pre-training loss:

Next Sentence Prediction was introduced as a secondary pre-training objective to help BERT learn relationships between sentence pairs, which is crucial for tasks like question answering and natural language inference. The task takes two sentences A and B, where 50\% of the time B is the actual next sentence that follows A in the corpus, and 50\% of the time B is a random sentence from elsewhere. The model must predict whether B genuinely follows A by using the representation of the special [CLS] token, which is prepended to every input sequence and serves as an aggregate sequence representation.

The motivation for NSP was that many downstream tasks require understanding relationships between sentence pairs, and pre-training on this objective should provide useful inductive biases. However, subsequent research, particularly the RoBERTa paper, demonstrated that NSP provides minimal or even negative benefit to downstream task performance. The likely explanation is that NSP is too easy a task—the model can often distinguish random sentences from consecutive sentences based on topic coherence alone, without learning the deeper inter-sentence relationships that would transfer to downstream tasks. RoBERTa removed NSP entirely and instead trained with longer sequences and dynamic masking, achieving substantially better performance. Modern BERT-style models typically omit NSP or replace it with more challenging objectives like sentence order prediction (used in ALBERT) or span corruption (used in T5).

Training Details and Computational Cost

Hardware and Training Configuration

BERT's original training represents a landmark in demonstrating that large-scale pre-training on commodity hardware clusters could produce models with broad applicability. BERT-base was trained on 16 Cloud TPU v3 chips, where each chip contains 2 cores for a total of 32 TPU cores. However, the paper reports using 4 Cloud TPU v3 Pods with 16 chips each, giving 64 TPU cores total. Each TPU v3 core provides approximately 123 TFLOPS of bfloat16 performance, yielding a combined peak performance of roughly 7.9 PFLOPS for the full training setup. In practice, achieving even 40-50\% utilization on such distributed training is considered excellent, suggesting an effective compute rate of 3-4 PFLOPS during training.

The training configuration uses a batch size of 256 sequences, each of maximum length 512 tokens, for a total of 131,072 tokens per batch. This large batch size is essential for stable training with the Adam optimizer and enables efficient utilization of the TPU hardware, which achieves best performance with large matrix operations. The learning rate schedule employs a warmup phase over the first 10,000 steps where the learning rate increases linearly from 0 to the peak learning rate of $1 \times 10^{-4}$, followed by linear decay to 0 over the remaining training steps. This warmup prevents the large initial gradients from destabilizing training, while the decay helps the model converge to a better optimum.

BERT-base training runs for 1 million steps with this configuration, requiring approximately 4 days of continuous training on the 64 TPU cores. Each training step processes 256 sequences of 512 tokens, so the total training data comprises $1{,}000{,}}000 \times 256 \times 512 = 131$ billion tokens. The training corpus consists of BooksCorpus (800 million words) and English Wikipedia (2.5 billion words), totaling approximately 3.3 billion words or roughly 4.4 billion tokens after WordPiece tokenization. This means the model sees each token approximately 30 times during training, providing sufficient repetition for the model to learn robust representations while maintaining diversity through the random masking strategy.

Computational Cost Analysis

The computational cost of BERT pre-training can be estimated from the FLOPs per training step and the total number of steps. As analyzed in Chapter 12, BERT-base requires approximately 96.6 billion FLOPs for a forward pass on a single sequence of length 512. The backward pass requires roughly twice the FLOPs of the forward pass, giving approximately 290 billion FLOPs per training step per sequence. With batch size 256, each training step requires $290 \times 256 = 74.2$ trillion FLOPs. Over 1 million training steps, the total computation is approximately $74.2 \times 10^{15} \times 10^6 = 7.42 \times 10^{22}$ FLOPs, or 74.2 zettaFLOPs.

At an effective compute rate of 3.5 PFLOPS (assuming 45\% utilization of the 7.9 PFLOPS peak), each training step takes approximately $74.2 \times 10^{12} / (3.5 \times 10^{15}) = 21.2$ milliseconds. Over 1 million steps, this yields $21.2 \times 10^6 = 21.2$ million milliseconds, or approximately 5.9 hours of pure compute time. However, the reported 4-day training time includes data loading, checkpointing, and distributed communication overhead, which can easily account for a factor of 16× slowdown in practice. This discrepancy highlights the importance of efficient data pipelines and communication strategies in distributed training.

The estimated cost of BERT-base pre-training in 2018 was approximately \$7,000 using Google Cloud TPU pricing. This relatively modest cost (compared to later models like GPT-3, which cost millions of dollars to train) made BERT accessible to academic research groups and smaller companies, democratizing access to large-scale pre-trained models. The cost breaks down to roughly \$1.75 per hour for the TPU Pod configuration, and 4 days of training equals 96 hours, yielding $96 \times \$1.75 \approx \$168$. However, the \$7,000 figure likely includes multiple training runs, hyperparameter tuning, and ablation studies rather than a single successful training run.

BERT-large requires substantially more computation due to its larger size. With 340 million parameters compared to BERT-base's 110 million, and 24 layers instead of 12, BERT-large requires approximately 3-4× the FLOPs per training step. The original paper reports training BERT-large for 1 million steps on 64 TPU cores, taking approximately 4 days as well, though this likely involved more aggressive optimization or different batch sizes. The estimated cost for BERT-large pre-training was around \$20,000-\$25,000, reflecting the increased computational requirements.

Fine-Tuning BERT

Fine-Tuning Procedure and Memory Requirements

Fine-tuning BERT for downstream tasks represents one of the model's key advantages: the pre-trained representations can be adapted to specific tasks with relatively small amounts of labeled data and modest computational resources. The fine-tuning process adds a task-specific head on top of the pre-trained BERT model and trains the entire model end-to-end on the target task. This approach typically requires only 2-4 epochs of training on task-specific data, taking minutes to hours on a single GPU depending on dataset size.

The memory requirements for fine-tuning are substantially lower than pre-training because fine-tuning typically uses smaller batch sizes and shorter sequences. For BERT-base fine-tuning on a classification task with batch size 16 and sequence length 128, the memory breakdown is approximately: 440 MB for model parameters (FP32), 440 MB for gradients, 880 MB for Adam optimizer states, and roughly 2 GB for activations. This totals approximately 3.8 GB, comfortably fitting on consumer GPUs with 8 GB of memory like the RTX 2080 or RTX 3070. Using mixed precision training reduces this to approximately 2.5 GB, enabling batch sizes of 32-48 on the same hardware.

The typical hyperparameter configuration for BERT fine-tuning uses a much smaller learning rate than pre-training to avoid catastrophically forgetting the pre-trained representations. Learning rates in the range of $2 \times 10^{-5}$ to $5 \times 10^{-5}$ work well for most tasks, with a linear warmup over 10\% of training steps followed by linear decay. The batch size typically ranges from 16 to 32 for most tasks, though larger batch sizes (64-128) can improve performance on tasks with abundant training data. Fine-tuning for 2-4 epochs is usually sufficient, as longer training often leads to overfitting on the task-specific data.

Classification Tasks

For sequence classification (sentiment, topic, etc.):

1. Add classification head on [CLS] token
   $$ \text{logits} = \mW_{\text{cls}} \vh_{[CLS]} + \vb_{\text{cls}} $$
2. Fine-tune entire model end-to-end

Token-Level Tasks

For named entity recognition (NER), POS tagging:

1. Add classification head on each token
   $$ \text{logits}_i = \mW_{\text{token}} \vh_i + \vb_{\text{token}} $$
2. Predict label for each token independently

Question Answering (SQuAD)

For span-based QA:

1. Input: [CLS] Question [SEP] Context [SEP]
2. Predict start and end positions in context
   $$\begin{align} P_{\text{start}}(i) &= \text{softmax}(\vh_i\transpose \mathbf{s}) \\ P_{\text{end}}(i) &= \text{softmax}(\vh_i\transpose \mathbf{e}) \end{align}$$
   where $\mathbf{s}, \mathbf{e} \in \R^{768}$ are learned vectors.

BERT Variants

RoBERTa (Robustly Optimized BERT)

RoBERTa, introduced by Facebook AI Research in 2019, demonstrated that BERT was significantly undertrained and that careful attention to training procedures could yield substantial improvements without architectural changes. The key insight was that many of BERT's design choices were not carefully ablated, and several modifications to the training procedure could dramatically improve performance. RoBERTa achieves state-of-the-art results on GLUE, SQuAD, and RACE benchmarks by training the same architecture as BERT-base and BERT-large with improved training procedures.

The first major change removes the Next Sentence Prediction objective entirely, training only with masked language modeling. Ablation studies showed that NSP actually hurt performance on downstream tasks, likely because the task is too simple and doesn't provide useful training signal. Instead, RoBERTa uses full-length sequences of up to 512 tokens sampled contiguously from documents, allowing the model to learn longer-range dependencies without the artificial sentence-pair structure. This change alone improves downstream task performance by 1-2\% across most benchmarks.

Dynamic masking represents another crucial improvement over BERT's static masking. BERT generates masked training examples once during data preprocessing and uses the same masked version throughout all epochs of training. This means the model sees the exact same masked examples multiple times, potentially leading to memorization. RoBERTa instead generates new masking patterns every time a sequence is fed to the model, ensuring that the model never sees the exact same masked example twice across the entire training run. This dynamic masking provides more diverse training signal and improves generalization, particularly for longer training runs.

The training scale increases dramatically compared to BERT. RoBERTa uses batch sizes of 8,192 sequences (32× larger than BERT's 256), enabled by gradient accumulation across multiple steps. Larger batch sizes improve training stability and allow for higher learning rates, accelerating convergence. The training data expands from BERT's 16 GB of text (BooksCorpus + Wikipedia) to 160 GB, incorporating CC-News (76 GB), OpenWebText (38 GB), and Stories (31 GB) in addition to the original sources. This 10× increase in training data provides substantially more diverse examples for the model to learn from.

Perhaps most importantly, RoBERTa trains for much longer than BERT. While BERT-base trains for 1 million steps, RoBERTa trains for 500,000 steps with the larger batch size, corresponding to processing 4× more tokens overall. Some RoBERTa variants train for even longer, up to 1 million steps with the large batch size. This extended training allows the model to better learn the training objective and develop more robust representations. The computational cost of RoBERTa training is approximately 10-15× higher than BERT due to the combination of more data, larger batches, and longer training, estimated at \$50,000-\$100,000 for the full training run.

The results demonstrate the value of these improvements. RoBERTa-base matches or exceeds BERT-large performance on most tasks despite having the same architecture as BERT-base, and RoBERTa-large achieves new state-of-the-art results across GLUE, SQuAD 2.0, and RACE. The improvements are particularly pronounced on tasks requiring deep language understanding, such as natural language inference and reading comprehension, where RoBERTa-large improves over BERT-large by 2-4\% absolute.

ALBERT (A Lite BERT)

ALBERT addresses BERT's memory consumption and training time through parameter sharing and factorized embeddings, achieving comparable performance with far fewer parameters. The key insight is that many of BERT's parameters are redundant, and careful parameter sharing can maintain model capacity while dramatically reducing memory requirements. ALBERT-xxlarge achieves similar performance to BERT-large with only 235 million parameters compared to BERT-large's 340 million, and the parameter sharing enables training on hardware that couldn't accommodate the full BERT-large model.

Factorized embedding parameterization represents the first major innovation. In BERT, the token embedding dimension equals the hidden dimension ($V \times d_{\text{model}}$), meaning the vocabulary matrix for BERT-base contains $30{,}000 \times 768 = 23$ million parameters. ALBERT observes that token embeddings are meant to learn context-independent representations, while hidden layers learn context-dependent representations, so these don't need to share the same dimension. ALBERT instead projects the vocabulary to a smaller embedding dimension $E$ (typically 128), then projects from $E$ to $d_{\text{model}}$. This factorization reduces embedding parameters from $V \times d_{\text{model}}$ to $V \times E + E \times d_{\text{model}}$. For ALBERT-base with $E=128$: $30{,}000 \times 128 + 128 \times 768 = 3{,}938{,}304$ parameters, an 83\% reduction from BERT's 23 million embedding parameters.

Cross-layer parameter sharing takes the reduction further by using the same parameters for all transformer layers. Instead of having unique parameters for each of the 12 or 24 layers, ALBERT uses a single set of layer parameters that is applied repeatedly. This reduces the transformer stack parameters by a factor of $L$ (the number of layers). For ALBERT-base, the 12-layer transformer stack requires only 7.1 million parameters (one layer's worth) instead of BERT-base's 85 million parameters. The memory savings are dramatic: ALBERT-base contains only 12 million parameters total compared to BERT-base's 110 million, an 89\% reduction.

The performance impact of parameter sharing is surprisingly modest. Cross-layer parameter sharing causes a small performance degradation (typically 1-2\% on downstream tasks) compared to models with unique layer parameters, but this is far less than the 89\% parameter reduction would suggest. The likely explanation is that the transformer layers learn similar functions across depths, so sharing parameters doesn't severely limit model capacity. The factorized embeddings actually improve performance slightly by preventing the model from overfitting to the vocabulary and encouraging more robust token representations.

ALBERT also replaces Next Sentence Prediction with Sentence Order Prediction (SOP), a more challenging task. Instead of distinguishing consecutive sentences from random sentences (which can often be solved by topic matching), SOP requires the model to distinguish the correct sentence order from the reversed order. Given sentences A and B that appear consecutively in the corpus, 50\% of examples present them as A-B (positive) and 50\% as B-A (negative). This task requires understanding inter-sentence coherence and cannot be solved by topic matching alone, providing a more useful training signal than NSP.

The memory and speed implications are substantial. ALBERT-base requires only 48 MB in FP32 (12 million parameters × 4 bytes) compared to BERT-base's 440 MB, enabling training with much larger batch sizes on the same hardware. However, parameter sharing doesn't reduce computation—each layer still performs the same number of FLOPs, just with shared parameters. This means ALBERT-base has similar training time per step as BERT-base despite the parameter reduction. The primary benefit is memory efficiency, not computational efficiency. ALBERT-xxlarge, with $d_{\text{model}}=4096$ and 12 layers, contains 235 million parameters but achieves performance comparable to BERT-large (340 million parameters) on most benchmarks.

DistilBERT: Knowledge Distillation for Compression

DistilBERT demonstrates that knowledge distillation can compress BERT to a fraction of its size while retaining most of its performance, making deployment feasible on resource-constrained devices. The core idea is to train a smaller "student" model to mimic the behavior of the larger "teacher" BERT model, transferring the knowledge encoded in the teacher's parameters to the more compact student architecture. DistilBERT achieves 97\% of BERT-base's performance on GLUE while being 40\% smaller and 60\% faster, making it practical for production deployment where latency and memory are critical.

The DistilBERT architecture uses 6 transformer layers instead of BERT-base's 12, with the same hidden dimension of 768 and 12 attention heads. This halving of depth reduces parameters from 110 million to approximately 66 million, a 40\% reduction. The parameter savings come entirely from the transformer layers (6 layers × 7.1M parameters = 42.6M instead of 85M), while the embedding layer remains unchanged at 23 million parameters. The reduced depth means each forward pass requires only 6 layers of computation instead of 12, directly translating to a 2× speedup in the ideal case. In practice, the speedup is approximately 1.6-1.7× due to overhead from embedding lookups and the final prediction layer, which don't scale with depth.

The distillation loss combines three components to transfer knowledge from teacher to student. The first component is the standard cross-entropy loss between the student's predictions and the true labels, ensuring the student learns the correct task. The second component is the distillation loss, which minimizes the KL divergence between the student's output distribution and the teacher's output distribution. The teacher's "soft" predictions (probability distributions over the vocabulary) contain more information than the hard labels alone—for example, the teacher might assign 0.7 probability to the correct token, 0.2 to a synonym, and 0.1 to other tokens, revealing semantic relationships. The student learns from this richer signal by matching the teacher's distribution.

Distillation loss:

The temperature parameter $T$ (typically 2-4) softens the probability distributions, amplifying the differences between high-probability and low-probability tokens. Higher temperatures make the distributions more uniform, providing more information about the teacher's uncertainty. The coefficient $\alpha$ (typically 0.5) balances the two objectives. The third component (not shown in the simplified equation) is a cosine embedding loss that encourages the student's hidden states to align with the teacher's hidden states, providing additional supervision beyond the output distributions.

The training procedure initializes the student by taking every other layer from the pre-trained BERT teacher, providing a warm start that accelerates convergence. The student is then trained on the same masked language modeling task as BERT, but with the teacher's soft targets providing additional supervision. Training DistilBERT requires approximately 90 hours on 8 NVIDIA V100 GPUs, substantially less than the 4 days on 64 TPU cores required for BERT-base pre-training. The reduced training time reflects both the smaller model size and the benefit of learning from the teacher's predictions rather than learning from scratch.

The performance-efficiency trade-off is compelling for production deployment. DistilBERT retains 97\% of BERT-base's performance on GLUE tasks, with the largest drops on tasks requiring deep reasoning (2-3\% degradation on MNLI and QQP) and smaller drops on simpler tasks (0-1\% on SST-2 and MRPC). The 40\% parameter reduction translates directly to memory savings: DistilBERT requires 264 MB in FP32 compared to BERT-base's 440 MB, enabling deployment on mobile devices and edge hardware. The 60\% speedup (1.6-1.7× faster) reduces inference latency from approximately 14 ms to 8-9 ms per sequence on an NVIDIA V100, crucial for real-time applications.

Memory and Speed Comparisons

Comparing BERT variants across memory footprint and inference speed reveals clear trade-offs between model capacity and deployment efficiency. The following analysis uses BERT-base as the baseline and measures relative performance on GLUE benchmark tasks, memory consumption in FP32, and inference throughput on an NVIDIA V100 GPU with batch size 1 and sequence length 128.

Parameter counts and memory:

• BERT-base: 110M parameters, 440 MB (FP32), 220 MB (FP16)
• BERT-large: 340M parameters, 1.36 GB (FP32), 680 MB (FP16)
• RoBERTa-base: 125M parameters, 500 MB (FP32), 250 MB (FP16) — slightly larger due to different vocabulary
• RoBERTa-large: 355M parameters, 1.42 GB (FP32), 710 MB (FP16)
• ALBERT-base: 12M parameters, 48 MB (FP32), 24 MB (FP16) — 89\% reduction from BERT-base
• ALBERT-xxlarge: 235M parameters, 940 MB (FP32), 470 MB (FP16)
• DistilBERT: 66M parameters, 264 MB (FP32), 132 MB (FP16) — 40\% reduction from BERT-base

Inference speed (sequences/second on V100, batch size 1, sequence length 128):

• BERT-base: $\sim$140 sequences/sec (baseline)
• BERT-large: $\sim$50 sequences/sec (2.8× slower due to larger dimensions and more layers)
• RoBERTa-base: $\sim$135 sequences/sec (similar to BERT-base)
• ALBERT-base: $\sim$140 sequences/sec (same speed as BERT-base despite fewer parameters—computation unchanged)
• DistilBERT: $\sim$230 sequences/sec (1.6× faster due to half the layers)

Performance on GLUE (average score):

• BERT-base: 78.5 (baseline)
• BERT-large: 82.1 (+3.6 points)
• RoBERTa-base: 80.4 (+1.9 points — matches BERT-large with base architecture)
• RoBERTa-large: 84.3 (+5.8 points — new state-of-the-art)
• ALBERT-base: 77.2 (-1.3 points — slight degradation from parameter sharing)
• ALBERT-xxlarge: 82.3 (+3.8 points — matches BERT-large with fewer parameters)
• DistilBERT: 76.1 (-2.4 points — retains 97\% of performance)

These comparisons reveal distinct use cases for each variant. RoBERTa offers the best performance when computational resources are available, making it ideal for offline processing and high-accuracy applications. ALBERT provides excellent memory efficiency for training and deployment on memory-constrained hardware, though it doesn't reduce inference time. DistilBERT offers the best balance of speed and performance for latency-sensitive applications like real-time search and interactive systems. BERT-base remains a strong baseline that balances performance, speed, and memory for most applications.

Hardware Requirements and Deployment

GPU Memory Requirements

Understanding GPU memory requirements is essential for selecting appropriate hardware for BERT training and inference. The memory footprint depends on whether the model is being trained or used for inference, the batch size, sequence length, and numerical precision. For training, memory must accommodate model parameters, gradients, optimizer states, and activations, while inference requires only parameters and a single forward pass of activations.

BERT-base training in FP32 with batch size 32 and sequence length 512 requires approximately 13.8 GB of memory, as detailed in Chapter 12. This breaks down to 440 MB for parameters, 440 MB for gradients, 880 MB for Adam optimizer states (first and second moments), and approximately 12 GB for activations across the 12 layers. This memory requirement fits comfortably on NVIDIA V100 GPUs with 16 GB of memory, RTX 3090 GPUs with 24 GB, or A100 GPUs with 40-80 GB. However, the batch size must be reduced for GPUs with less memory—an RTX 3080 with 10 GB can train BERT-base with batch size 16-20, while an RTX 3070 with 8 GB is limited to batch size 8-12.

Mixed precision training with FP16 or BF16 dramatically reduces memory consumption by storing activations in 16-bit format while maintaining FP32 master weights for numerical stability. For BERT-base, mixed precision reduces total memory from 13.8 GB to approximately 8 GB, enabling batch size 32 on GPUs with 12 GB of memory or batch size 64 on GPUs with 24 GB. The memory savings come primarily from activations, which are halved from 12 GB to 6 GB, while the parameter and optimizer memory increases slightly from 1.76 GB to 1.98 GB due to maintaining both FP16 and FP32 copies of parameters. Gradient checkpointing can further reduce memory by recomputing activations during the backward pass rather than storing them, reducing activation memory by approximately 80\% at the cost of 20-30\% longer training time.

BERT-large training requires substantially more memory due to its larger hidden dimension and greater depth. In FP32 with batch size 32 and sequence length 512, BERT-large requires approximately 32 GB of memory: 1.36 GB for parameters, 1.36 GB for gradients, 2.72 GB for optimizer states, and approximately 26 GB for activations. This necessitates GPUs with at least 32 GB of memory, such as the NVIDIA V100 (32 GB variant) or A100 (40-80 GB). Mixed precision reduces this to approximately 18 GB, enabling training on A100 40 GB GPUs with batch size 16-24. Consumer GPUs like the RTX 3090 (24 GB) can train BERT-large with mixed precision and gradient checkpointing at batch size 8-12, though training time increases significantly.

Inference memory requirements are far more modest because they don't include gradients, optimizer states, or stored activations for backpropagation. BERT-base inference in FP32 requires only 440 MB for parameters plus approximately 200-300 MB for a single forward pass of activations with batch size 1, totaling under 1 GB. In FP16, this drops to approximately 400 MB total, enabling BERT-base inference on virtually any GPU, including mobile GPUs and edge devices. BERT-large inference requires approximately 1.5 GB in FP32 or 800 MB in FP16, still easily fitting on consumer GPUs. The primary consideration for inference is batch size: larger batch sizes improve throughput but increase activation memory linearly. A V100 with 16 GB can run BERT-base inference with batch size 128-256 in FP16, achieving throughput of 15,000-20,000 sequences per second.

Batch Size Limits by GPU Type

The maximum batch size for BERT training varies significantly across GPU types, directly impacting training throughput and efficiency. Larger batch sizes improve GPU utilization by amortizing memory bandwidth costs across more computation, but they're limited by available memory. The following analysis assumes mixed precision training with sequence length 512 for BERT-base and BERT-large.

BERT-base maximum batch sizes (mixed precision, sequence length 512):

• RTX 3060 (12 GB): Batch size 24-28 without gradient checkpointing, 48-56 with checkpointing
• RTX 3070 (8 GB): Batch size 14-18 without gradient checkpointing, 28-36 with checkpointing
• RTX 3080 (10 GB): Batch size 18-22 without gradient checkpointing, 36-44 with checkpointing
• RTX 3090 (24 GB): Batch size 48-56 without gradient checkpointing, 96-112 with checkpointing
• V100 (16 GB): Batch size 28-32 without gradient checkpointing, 56-64 with checkpointing
• V100 (32 GB): Batch size 64-72 without gradient checkpointing, 128-144 with checkpointing
• A100 (40 GB): Batch size 80-96 without gradient checkpointing, 160-192 with checkpointing
• A100 (80 GB): Batch size 160-192 without gradient checkpointing, 320-384 with checkpointing

BERT-large maximum batch sizes (mixed precision, sequence length 512):

• RTX 3090 (24 GB): Batch size 12-16 without gradient checkpointing, 24-32 with checkpointing
• V100 (16 GB): Batch size 6-8 without gradient checkpointing, 12-16 with checkpointing
• V100 (32 GB): Batch size 16-20 without gradient checkpointing, 32-40 with checkpointing
• A100 (40 GB): Batch size 20-24 without gradient checkpointing, 40-48 with checkpointing
• A100 (80 GB): Batch size 48-56 without gradient checkpointing, 96-112 with checkpointing

These batch size limits have direct implications for training efficiency. Smaller batch sizes reduce GPU utilization because the model spends more time on memory transfers relative to computation. For BERT-base on an RTX 3070 with batch size 16, GPU utilization typically reaches only 50-60\% of peak FLOPS, while an A100 with batch size 96 can achieve 70-80\% utilization. Gradient accumulation can simulate larger batch sizes by accumulating gradients over multiple forward-backward passes before updating parameters, enabling effective batch sizes of 128-256 even on GPUs limited to batch size 16-32 per step. However, gradient accumulation increases training time proportionally to the accumulation steps.

Inference Speed Analysis

Inference speed determines the feasibility of deploying BERT in production systems where latency and throughput are critical. We measure inference speed in two ways: latency (time per sequence for batch size 1, important for interactive applications) and throughput (sequences per second for large batches, important for offline processing). The following measurements use sequence length 128, which is typical for many classification and NER tasks, and FP16 precision on NVIDIA GPUs.

BERT-base inference latency (batch size 1, sequence length 128):

• V100: 7.2 ms per sequence (139 sequences/sec)
• A100: 3.8 ms per sequence (263 sequences/sec) — 1.9× faster than V100
• RTX 3090: 5.1 ms per sequence (196 sequences/sec)
• CPU (Intel Xeon Gold 6248): 45-60 ms per sequence (17-22 sequences/sec) — 10-15× slower than GPU

BERT-base inference throughput (batch size 128, sequence length 128):

• V100: 18,000-20,000 sequences/sec
• A100: 35,000-40,000 sequences/sec — 2× faster than V100
• RTX 3090: 25,000-28,000 sequences/sec

BERT-large inference latency (batch size 1, sequence length 128):

• V100: 20.5 ms per sequence (49 sequences/sec)
• A100: 10.8 ms per sequence (93 sequences/sec) — 1.9× faster than V100
• RTX 3090: 14.2 ms per sequence (70 sequences/sec)

BERT-large inference throughput (batch size 64, sequence length 128):

• V100: 6,500-7,500 sequences/sec
• A100: 12,000-14,000 sequences/sec — 1.9× faster than V100
• RTX 3090: 9,000-10,500 sequences/sec

The A100's superior performance comes from its higher memory bandwidth (1.6 TB/s vs V100's 900 GB/s) and more powerful Tensor Cores (312 TFLOPS FP16 vs V100's 125 TFLOPS). For BERT inference, which is often memory-bandwidth bound due to loading model parameters, the A100's bandwidth advantage is particularly valuable. The approximately 2× speedup of A100 over V100 holds across different batch sizes and model sizes, making the A100 the preferred choice for production BERT deployment when latency is critical.

Sequence length significantly impacts inference speed due to the quadratic scaling of attention computation. For BERT-base on a V100, increasing sequence length from 128 to 512 (4× longer) increases latency from 7.2 ms to approximately 18 ms (2.5× slower), less than the 4× that pure quadratic scaling would suggest because the feed-forward network and embedding layers don't scale quadratically. For very long sequences approaching the 512 token maximum, attention computation dominates and the scaling approaches quadratic. This explains why efficient attention mechanisms (Chapter 16) focus on reducing the $O(n^2)$ attention complexity for long-context applications.

Analysis and Interpretability

What BERT Learns

Lower layers: Syntactic information (POS tags, parse trees)

Middle layers: Semantic information (word sense, entity types)

Upper layers: Task-specific information

Attention patterns:

• Some heads attend to next token (language modeling pattern)
• Some heads attend to syntactic relations (e.g., verbs to subjects)
• Some heads attend broadly (averaging)

Probing Tasks

Test what linguistic information is encoded:

• Surface: Sentence length, word order
• Syntactic: POS tags, dependency labels, constituency trees
• Semantic: Named entities, semantic roles, coreference

Method: Train linear classifier on frozen BERT representations

Result: BERT captures surprisingly rich linguistic structure!

Exercises

Solutions

Full solutions for all exercises are available at \url{https://deeplearning.hofkensvermeulen.be}.

import torch
import torch.nn as nn
from transformers import BertTokenizer, BertForMaskedLM
import numpy as np

def create_mlm_data(sentence, tokenizer, mask_prob=0.15):
    """Create masked language modeling training data"""
    # Tokenize
    tokens = tokenizer.tokenize(sentence)
    token_ids = tokenizer.convert_tokens_to_ids(tokens)
    
    # Create labels (copy of original)
    labels = token_ids.copy()
    
    # Mask tokens
    masked_indices = []
    for i in range(len(token_ids)):
        if np.random.random() < mask_prob:
            masked_indices.append(i)
            
            # 80
            if np.random.random() < 0.8:
                token_ids[i] = tokenizer.mask_token_id
            # 10
            elif np.random.random() < 0.5:
                token_ids[i] = np.random.randint(
                    0, tokenizer.vocab_size
                )
            # 10
    
    # Set non-masked positions to -100 (ignored in loss)
    for i in range(len(labels)):
        if i not in masked_indices:
            labels[i] = -100
    
    return token_ids, labels, masked_indices

# Example
sentence = "The quick brown fox jumps"
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
model = BertForMaskedLM.from_pretrained('bert-base-uncased')

# Create masked input
input_ids, labels, masked_idx = create_mlm_data(sentence, tokenizer)

# Convert to tensors
input_tensor = torch.tensor([input_ids])
label_tensor = torch.tensor([labels])

# Forward pass
outputs = model(input_tensor, labels=label_tensor)
loss = outputs.loss
logits = outputs.logits

print(f"Original: {sentence}")
print(f"Tokens: {tokenizer.convert_ids_to_tokens(input_ids)}")
print(f"Masked positions: {masked_idx}")
print(f"MLM Loss: {loss.item():.4f}")

# Get predictions for masked positions
with torch.no_grad():
    predictions = torch.softmax(logits[0], dim=-1)

print("\nPredictions for masked positions:")
for idx in masked_idx:
    # Get top 5 predictions
    probs, indices = torch.topk(predictions[idx], k=5)
    predicted_tokens = tokenizer.convert_ids_to_tokens(
        indices.tolist()
    )
    
    original_token = tokenizer.convert_ids_to_tokens([labels[idx]])[0]
    
    print(f"\nPosition {idx} (original: '{original_token}'):")
    for token, prob in zip(predicted_tokens, probs):
        print(f"  {token}: {prob.item():.4f}")

Original: The quick brown fox jumps
Tokens: ['the', '[MASK]', 'brown', '[MASK]', 'jumps']
Masked positions: [1, 3]
MLM Loss: 2.3456

Predictions for masked positions:

Position 1 (original: 'quick'):
  quick: 0.8234
  fast: 0.0892
  little: 0.0234
  small: 0.0156
  big: 0.0089

Position 3 (original: 'fox'):
  fox: 0.7123
  dog: 0.1234
  cat: 0.0567
  animal: 0.0234
  wolf: 0.0189

from transformers import BertForSequenceClassification, AdamW
from torch.utils.data import DataLoader, TensorDataset
import matplotlib.pyplot as plt

def finetune_bert(train_loader, val_loader, freeze_bert=False, 
                  epochs=10):
    model = BertForSequenceClassification.from_pretrained(
        'bert-base-uncased', 
        num_labels=2
    )
    
    # Option 1: Freeze BERT layers, train only classifier
    if freeze_bert:
        for param in model.bert.parameters():
            param.requires_grad = False
    
    # Optimizer
    optimizer = AdamW(
        filter(lambda p: p.requires_grad, model.parameters()),
        lr=2e-5 if not freeze_bert else 1e-3
    )
    
    train_losses, val_losses = [], []
    train_accs, val_accs = [], []
    
    for epoch in range(epochs):
        # Training
        model.train()
        train_loss, train_correct = 0, 0
        
        for batch in train_loader:
            input_ids, attention_mask, labels = batch
            
            optimizer.zero_grad()
            outputs = model(
                input_ids, 
                attention_mask=attention_mask,
                labels=labels
            )
            
            loss = outputs.loss
            loss.backward()
            optimizer.step()
            
            train_loss += loss.item()
            preds = outputs.logits.argmax(dim=-1)
            train_correct += (preds == labels).sum().item()
        
        # Validation
        model.eval()
        val_loss, val_correct = 0, 0
        
        with torch.no_grad():
            for batch in val_loader:
                input_ids, attention_mask, labels = batch
                outputs = model(
                    input_ids,
                    attention_mask=attention_mask,
                    labels=labels
                )
                val_loss += outputs.loss.item()
                preds = outputs.logits.argmax(dim=-1)
                val_correct += (preds == labels).sum().item()
        
        # Record metrics
        train_losses.append(train_loss / len(train_loader))
        val_losses.append(val_loss / len(val_loader))
        train_accs.append(train_correct / len(train_loader.dataset))
        val_accs.append(val_correct / len(val_loader.dataset))
        
        print(f"Epoch {epoch+1}: "
              f"Train Loss={train_losses[-1]:.4f}, "
              f"Val Acc={val_accs[-1]:.4f}")
    
    return train_losses, val_losses, train_accs, val_accs

# Run both approaches
results_frozen = finetune_bert(
    train_loader, val_loader, 
    freeze_bert=True, epochs=10
)

results_full = finetune_bert(
    train_loader, val_loader,
    freeze_bert=False, epochs=10
)

# Plot learning curves
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 5))

# Loss curves
ax1.plot(results_frozen[1], label='Frozen BERT', marker='o')
ax1.plot(results_full[1], label='Full fine-tuning', marker='s')
ax1.set_xlabel('Epoch')
ax1.set_ylabel('Validation Loss')
ax1.set_title('Learning Curves: Loss')
ax1.legend()
ax1.grid(True)

# Accuracy curves
ax2.plot(results_frozen[3], label='Frozen BERT', marker='o')
ax2.plot(results_full[3], label='Full fine-tuning', marker='s')
ax2.set_xlabel('Epoch')
ax2.set_ylabel('Validation Accuracy')
ax2.set_title('Learning Curves: Accuracy')
ax2.legend()
ax2.grid(True)

from transformers import BertModel, BertTokenizer
import torch
import matplotlib.pyplot as plt
import seaborn as sns

def visualize_attention(sentence, model, tokenizer, layer=0):
    """Visualize attention patterns for all heads in a layer"""
    # Tokenize
    inputs = tokenizer(sentence, return_tensors='pt')
    tokens = tokenizer.convert_ids_to_tokens(inputs['input_ids'][0])
    
    # Get attention weights
    with torch.no_grad():
        outputs = model(**inputs, output_attentions=True)
        attentions = outputs.attentions  # Tuple of (layer, batch, head, seq, seq)
    
    # Extract attention for specified layer
    attn = attentions[layer][0]  # Shape: (num_heads, seq_len, seq_len)
    
    # Plot all heads
    num_heads = attn.shape[0]
    fig, axes = plt.subplots(3, 4, figsize=(20, 15))
    
    for head_idx in range(num_heads):
        ax = axes[head_idx // 4, head_idx 
        
        # Get attention matrix for this head
        attn_matrix = attn[head_idx].numpy()
        
        # Plot heatmap
        sns.heatmap(
            attn_matrix,
            xticklabels=tokens,
            yticklabels=tokens,
            cmap='viridis',
            ax=ax,
            cbar=True,
            square=True
        )
        ax.set_title(f'Head {head_idx}')
        ax.set_xlabel('Key')
        ax.set_ylabel('Query')
    
    plt.tight_layout()
    plt.savefig(f'attention_layer_{layer}.png', dpi=150)
    
    return attn, tokens

# Example
sentence = "The cat that chased the mouse ran away"
model = BertModel.from_pretrained('bert-base-uncased')
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')

# Visualize different layers
for layer in [0, 5, 11]:
    attn, tokens = visualize_attention(sentence, model, tokenizer, layer)
    print(f"\nLayer {layer} attention patterns analyzed")

def analyze_attention_patterns(attn, tokens):
    """Identify heads with specific patterns"""
    num_heads = attn.shape[0]
    
    # (1) Adjacent word attention
    adjacent_heads = []
    for head in range(num_heads):
        # Check diagonal attention (adjacent words)
        diagonal_score = 0
        for i in range(len(tokens) - 1):
            diagonal_score += attn[head, i, i+1] + attn[head, i+1, i]
        diagonal_score /= (2 * (len(tokens) - 1))
        
        if diagonal_score > 0.3:  # Threshold
            adjacent_heads.append((head, diagonal_score))
    
    # (2) Subject-verb relations
    # "cat" (position 2) -> "chased" (position 4)
    # "cat" (position 2) -> "ran" (position 8)
    subj_verb_heads = []
    cat_idx = tokens.index('cat')
    chased_idx = tokens.index('chased')
    ran_idx = tokens.index('ran')
    
    for head in range(num_heads):
        score = (attn[head, cat_idx, chased_idx] + 
                attn[head, cat_idx, ran_idx]) / 2
        if score > 0.2:
            subj_verb_heads.append((head, score))
    
    # (3) Long-range dependencies
    # "cat" (position 2) -> "away" (position 9)
    long_range_heads = []
    away_idx = tokens.index('away')
    
    for head in range(num_heads):
        score = attn[head, cat_idx, away_idx]
        if score > 0.15:
            long_range_heads.append((head, score))
    
    return adjacent_heads, subj_verb_heads, long_range_heads

# Analyze patterns
adjacent, subj_verb, long_range = analyze_attention_patterns(attn, tokens)

print("(1) Heads capturing adjacent words:")
for head, score in sorted(adjacent, key=lambda x: x[1], reverse=True):
    print(f"  Head {head}: {score:.3f}")

print("\n(2) Heads capturing subject-verb relations:")
for head, score in sorted(subj_verb, key=lambda x: x[1], reverse=True):
    print(f"  Head {head}: {score:.3f}")

print("\n(3) Heads capturing long-range dependencies:")
for head, score in sorted(long_range, key=lambda x: x[1], reverse=True):
    print(f"  Head {head}: {score:.3f}")

(1) Heads capturing adjacent words:
  Head 2: 0.456
  Head 7: 0.389
  Head 10: 0.342

(2) Heads capturing subject-verb relations:
  Head 5: 0.678
  Head 8: 0.534
  Head 11: 0.423

(3) Heads capturing long-range dependencies:
  Head 3: 0.289
  Head 9: 0.234
  Head 6: 0.198