Lecture 2: Transformer Architecture Improvements and BERT Deep Dive

One. Course Details

This is the second lecture of Stanford University’s CME 295: Transformers and Large Language Models, taught by twin brothers Afshine and Shervine Amidi. The lecture opens with two critical logistical updates: first, the audio recording setup has been completely revised after suboptimal audio quality in the first lecture; second, the final exam is currently tentatively scheduled for Wednesday, December 10, but the instructors are working to move it to an earlier date that week and will announce the finalized schedule as soon as possible.

The lecture follows a clear two-part structure. In the first half, Afshine reviews core transformer fundamentals from Lecture 1 and breaks down the key architectural improvements to the original 2017 "Attention Is All You Need" model that have become standard in all modern transformers. In the second half, Shervine classifies all contemporary transformer-based models into three core families and delivers an in-depth deep dive into encoder-only architectures, with a primary focus on BERT and its industry-standard optimized variants.

---

Two. Key Learning Objectives

By the end of this lecture, students will be able to:

Compare and contrast four major position embedding methods (learned, sinusoidal, ALiBi, RoPE) and explain why Rotary Position Embeddings have become the universal industry standard.

Distinguish between post-norm, pre-norm, and RMSNorm layer normalization techniques and their respective impacts on training stability, convergence speed, and model performance.

Explain the motivation behind sparse sliding window attention and grouped query attention, and how these optimizations reduce the O(n²) computational complexity of vanilla self-attention.

Classify all modern transformer models into three core families (encoder-decoder, encoder-only, decoder-only) and identify their primary real-world use cases.

Describe the bidirectional attention mechanism of BERT, its two pre-training objectives, and its specialized input representation system.

Walk through the end-to-end fine-tuning process for BERT on both sequence-level classification tasks and token-level prediction tasks.

Explain the knowledge distillation technique used in DistilBERT and the key training optimizations introduced in RoBERTa.

---

Three. Memorable Course Quotes

"Each attention head is an opportunity for the model to learn a different way of projecting inputs into queries, keys, and values. They learn different types of relationships between words simultaneously."

"Self-attention is completely order-agnostic. It treats 'cat chases dog' exactly the same as 'dog chases cat'. We have to explicitly inject position information into the model."

"Rotary Position Embeddings are the standard today because they naturally encode relative distance between tokens and generalize perfectly to sequence lengths far longer than those seen during training."

"BERT is truly bidirectional: every token attends to every other token in the sequence. This is what makes it unbeatable for classification tasks, but also why it cannot generate text."

"The soft targets output by a model contain almost all the knowledge it has learned. This is the core insight behind distillation, which lets us build smaller models that retain 97% of the performance of larger ones."

---

Four. Detailed Lecture Notes

4.1 Transformer Fundamentals Recap

The lecture begins with a quick review of the scaled dot-product self-attention mechanism, the core building block of all transformers. Self-attention computes the representation of each token as a weighted sum of all other tokens in the sequence, using three learned projection matrices:

• Queries: What the current token is looking for in the sequence

• Keys: What each token in the sequence has to offer

• Values: The actual semantic content of each token

The standard formula for scaled dot-product attention is: Attention(Q, K, V) = softmax(QKᵀ/√dₖ)V

Multi-head attention runs this computation multiple times in parallel with separate projection matrices, allowing the model to learn different types of semantic and syntactic relationships between tokens simultaneously. The results from each head are concatenated and projected back to the original embedding dimension.

The instructors visualize multi-head attention using attention maps, which show the attention weights between pairs of tokens. For example, in the sentence "The application of the law is its most important feature", one attention head learns to associate the pronoun "its" with "law", while another head associates it with "application".

4.2 Key Improvements to the Transformer Architecture

While the core transformer architecture has remained remarkably stable since 2017, three critical components have seen significant optimizations that are now universal in all modern models.

4.2.1 Position Embeddings

Self-attention has no inherent sense of order, so position information must be explicitly added to token embeddings. Four major position embedding methods have been developed:

1. Learned position embeddings: Each position in the sequence has its own trainable embedding, which is added element-wise to the token embedding. This method is simple but cannot generalize to sequence lengths longer than those seen during training and can learn spurious biases from the training data.

2. Sinusoidal position embeddings: Used in the original transformer paper, these embeddings are generated using fixed sine and cosine functions of different frequencies. The dot product between two sinusoidal embeddings depends only on their relative distance, aligning with the intuition that closer tokens should be more similar. This method generalizes to arbitrary sequence lengths but is less flexible than learned embeddings.

3. ALiBi (Attention with Linear Biases): Instead of adding position information to the input embeddings, ALiBi adds a linear bias term directly to the attention scores, penalizing attention between tokens that are further apart. This method is simple, deterministic, and generalizes exceptionally well to long sequences.

4. Rotary Position Embeddings (RoPE): The current industry standard, RoPE rotates the query and key vectors by an angle proportional to their position in the sequence. The dot product between rotated queries and keys depends only on their relative distance, naturally encoding positional information. RoPE combines the best properties of all previous methods: it generalizes to arbitrary sequence lengths, encodes relative position naturally, and delivers state-of-the-art performance.

4.2.2 Layer Normalization

Layer normalization stabilizes training by normalizing the activations of each layer to have zero mean and unit variance. The original transformer used post-norm normalization, where normalization is applied after the residual connection. However, all modern models use pre-norm normalization, where normalization is applied before the sublayer (attention or feed-forward network), which dramatically improves training stability for deep models.

A further optimization is RMSNorm (Root Mean Square Normalization), which simplifies layer normalization by removing the mean subtraction step and only normalizing by the root mean square of the activations. RMSNorm reduces the number of learnable parameters and speeds up computation while maintaining performance comparable to full layer normalization.

4.2.3 Attention Mechanism Optimizations

Vanilla self-attention has O(n²) computational complexity, which becomes prohibitive for sequences longer than a few thousand tokens. Two major optimizations have addressed this limitation:

1. Sparse sliding window attention: Instead of allowing each token to attend to all other tokens, each token only attends to a fixed-size window of neighboring tokens. This reduces complexity to O(n) and is used in models like Longformer and Mistral. Multiple layers of sliding window attention create a large effective receptive field, similar to convolutions in computer vision.

2. Shared attention heads: To reduce memory usage during inference, modern models share projection matrices for keys and values across multiple attention heads. There are three standard variants:

   • Multi-Head Attention (MHA): Each head has its own query, key, and value projections (original transformer)

   • Multi-Query Attention (MQA): All heads share a single key and value projection

   • Group Query Attention (GQA): Heads are divided into groups, and each group shares a single key and value projection. This is the most common variant in modern LLMs, balancing performance and memory efficiency.

    

4.3 Modern Transformer Model Families

All contemporary transformer-based models fall into one of three core families, each optimized for different types of tasks:

1. Encoder-decoder models: Retain both the encoder and decoder components of the original transformer. They are primarily used for sequence-to-sequence tasks like machine translation and summarization. The most prominent example is the T5 (Text-to-Text Transfer Transformer) family, which frames all NLP tasks as text-to-text problems and uses a span corruption pre-training objective. Variants include mT5 (multilingual) and ByT5 (byte-level tokenization).

2. Encoder-only models: Remove the decoder component entirely. They use bidirectional self-attention, meaning every token attends to every other token in the sequence. This makes them ideal for classification and token-level tasks like sentiment analysis, named entity recognition, and question answering. The most famous example is BERT, covered in detail below.

3. Decoder-only models: Remove the encoder component entirely. They use causal (masked) self-attention, where each token only attends to itself and previous tokens. This makes them ideal for text generation and is the architecture used by all modern large language models (GPT, Llama, Mistral, etc.).

4.4 BERT Deep Dive

BERT (Bidirectional Encoder Representations from Transformers), introduced in 2018, revolutionized the field of NLP by demonstrating the power of bidirectional pre-trained representations. It remains the gold standard for most classification tasks today.

4.4.1 Core Architecture

BERT extracts only the encoder component from the original transformer and stacks multiple identical encoder layers. Its defining feature is bidirectional self-attention, which allows it to build context-aware representations that incorporate information from both the left and right of each token.

BERT introduces two special structural tokens:

• <CLS> token: A placeholder token added to the beginning of every sequence. The final embedding of the <CLS> token is used as the aggregate representation of the entire sequence for classification tasks.

• <SEP> token: A separator token used to distinguish between two sentences in paired tasks like next sentence prediction.

4.4.2 Input Representation

The input to BERT is the element-wise sum of three separate embeddings:

1. Token embeddings: Learned embeddings for each token in the vocabulary (BERT uses the WordPiece tokenizer with a vocabulary size of approximately 30,000).

2. Position embeddings: Learned embeddings that encode the position of each token in the sequence.

3. Segment embeddings: Learned embeddings that indicate which sentence a token belongs to (either segment A or segment B), used for paired tasks.

4.4.3 Pre-Training Objectives

BERT is pre-trained on two self-supervised objectives using massive amounts of unlabeled text:

1. Masked Language Model (MLM): 15% of tokens in the input are randomly selected and modified. 80% of the time they are replaced with a special <MASK> token, 10% of the time they are replaced with a random token, and 10% of the time they are left unchanged. The model is trained to predict the original token. This objective forces the model to learn deep bidirectional context.

2. Next Sentence Prediction (NSP): The model is given pairs of sentences and trained to predict whether the second sentence naturally follows the first. This objective was intended to teach the model sentence-level relationships, but later research showed it provides minimal practical benefit.

4.4.4 Fine-Tuning Process

After pre-training, BERT is fine-tuned on downstream tasks with minimal architectural modifications. A small task-specific linear layer is added on top of the pre-trained BERT encoder, and either the entire model or just the linear layer is fine-tuned on labeled task data. Common fine-tuning tasks include:

• Sentiment analysis: Linear layer on top of the <CLS> token to predict sentiment polarity.

• Named entity recognition: Linear layer on top of each token to predict entity types.

• Question answering: Two linear layers on top of each token to predict the start and end positions of the answer span.

4.5 BERT Variants

Two major variants of BERT have addressed its key limitations and become industry standards:

1. DistilBERT: Uses knowledge distillation to train a smaller, faster student model that mimics the output distribution of a larger teacher BERT model. DistilBERT has 50% fewer parameters and runs 60% faster while retaining 97% of the performance of the original BERT base model.

2. RoBERTa (Robustly Optimized BERT Pretraining Approach): Improves BERT by removing the NSP objective entirely, using dynamic masking (changing the masked tokens every epoch), training on 10x more data, and using longer sequence lengths. These changes deliver significant performance improvements on all downstream NLP benchmarks.