Chapter 1: A Primer on Transformers

This chapter introduces the fundamental concepts and architecture of Transformer models, breaking down each essential component.

Introduction to the Transformer

The Transformer architecture, introduced in the paper "Attention Is All You Need," revolutionized sequence-to-sequence modeling by relying entirely on attention mechanisms, eschewing recurrent neural networks (RNNs) and convolutional neural networks (CNNs). This allows for greater parallelization during training and improved performance on tasks such as machine translation, text summarization, and question answering.

Understanding the Encoder of the Transformer

The encoder is responsible for processing the input sequence and generating a rich representation that captures the context of each element. It consists of a stack of identical layers.

Integrating All Encoder Components

An encoder block is typically composed of two main sub-layers: a multi-head self-attention mechanism and a position-wise feedforward network.

• Multi-Head Self-Attention: Allows the model to weigh the importance of different words in the input sequence when processing a particular word.
• Position-wise Feedforward Network: A simple, fully connected feedforward network applied independently to each position.
• Add and Norm Component: Residual connections and layer normalization are applied after each sub-layer to facilitate gradient flow and stabilize training.

Understanding the Decoder of a Transformer

The decoder is responsible for generating the output sequence, one element at a time, based on the encoder's output and the previously generated elements. It also consists of a stack of identical layers.

Integrating Encoder and Decoder

The decoder receives the output of the encoder and uses it to guide the generation process.

Combining All Decoder Components

A decoder block typically consists of three main sub-layers:

• Masked Multi-Head Self-Attention: Similar to the encoder's self-attention, but masked to prevent attending to future tokens in the output sequence, ensuring causality.
• Multi-Head Attention over Encoder Output: This layer allows the decoder to attend to relevant parts of the encoded input sequence.
• Position-wise Feedforward Network: Similar to the encoder's feedforward network.
• Add and Norm Component: Residual connections and layer normalization are applied after each sub-layer.

Detailed Understanding of the Key Mechanisms

Self-Attention Mechanism

The self-attention mechanism enables the model to calculate a weighted sum of input values, where the weights are determined by the relationships between different elements in the sequence.

Step 1 of Self-Attention

For each input element, three vectors are created: Query (Q), Key (K), and Value (V). These are derived by multiplying the input embedding with learnable weight matrices.

Step 2 of Self-Attention

The attention scores are calculated by taking the dot product of the Query vector of one element with the Key vectors of all elements in the sequence. This measures the compatibility or similarity.

$$ \text{Score}(Q, K) = Q \cdot K^T $$

Step 3 of Self-Attention

The scores are scaled by the square root of the dimension of the key vectors ($d_k$) to prevent very large dot products, which can lead to vanishing gradients in the softmax.

$$ \text{Scaled Score} = \frac{Q \cdot K^T}{\sqrt{d_k}} $$

Step 4 of Self-Attention

The scaled scores are then passed through a softmax function to obtain attention weights. These weights sum to 1 and represent the probability distribution of how much attention each element should pay to others.

$$ \text{Attention Weights} = \text{softmax}\left(\frac{Q \cdot K^T}{\sqrt{d_k}}\right) $$

The final output for an element is the weighted sum of the Value vectors, using the calculated attention weights.

$$ \text{Output} = \text{Attention Weights} \cdot V $$

Multi-Head Attention Mechanism

Multi-head attention allows the model to jointly attend to information from different representation subspaces at different positions. Instead of performing a single attention function, it performs several attention functions in parallel.

Multi-Head Attention in the Decoder

In the decoder, multi-head attention is used in two ways:

1. Masked Multi-Head Self-Attention: Applied to the decoder's input sequence to prevent attending to future tokens.
2. Multi-Head Attention over Encoder Output: Attends to the output of the encoder.

Feedforward Network

The position-wise feedforward network consists of two linear transformations with a ReLU activation in between. It is applied independently to each position in the sequence.

$$ \text{FFN}(x) = \max(0, xW_1 + b_1)W_2 + b_2 $$

Feedforward Network in the Decoder

The structure of the feedforward network in the decoder is the same as in the encoder.

Add and Norm Component

The "Add & Norm" component consists of a residual connection followed by layer normalization.

• Residual Connection (Add): The output of a sub-layer is added to the input of that sub-layer ($x + \text{Sublayer}(x)$). This helps in training deep networks by mitigating the vanishing gradient problem.
• Layer Normalization (Norm): Normalizes the activations across the features for each sample independently. This helps stabilize training and speeds up convergence.

Learning Position with Positional Encoding

Since the Transformer architecture does not inherently process sequences sequentially, positional information must be injected. Positional encodings are added to the input embeddings to provide information about the relative or absolute position of tokens in the sequence.

Commonly, sine and cosine functions of different frequencies are used:

$$ \text{PE}(pos, 2i) = \sin\left(\frac{pos}{10000^{2i/d_{model}}}\right) $$ $$ \text{PE}(pos, 2i+1) = \cos\left(\frac{pos}{10000^{2i/d_{model}}}\right) $$

where:

• $pos$ is the position of the token in the sequence.
• $i$ is the dimension of the embedding.
• $d_{model}$ is the dimension of the model's embeddings.

Linear and Softmax Layers

After the decoder stack, a linear layer projects the decoder's output to the vocabulary size. This is followed by a softmax layer, which converts these scores into probabilities for each word in the vocabulary, indicating the likelihood of each word being the next token in the output sequence.

Training the Transformer

Training a Transformer involves feeding it input sequences and target output sequences. The model learns by minimizing a loss function (e.g., cross-entropy) that measures the difference between the predicted output probabilities and the actual target tokens. Optimization techniques like Adam are commonly used.

Summary, Questions, and Further Reading

This chapter provided a foundational understanding of the Transformer architecture. Key components include the encoder and decoder stacks, self-attention, multi-head attention, feedforward networks, and positional encoding.

• Questions:
  • What are the benefits of using attention over recurrence?
  • How does masking in the decoder ensure causality?
  • What is the role of positional encoding?
• Further Reading:
  • "Attention Is All You Need" by Vaswani et al. (2017)
  • Illustrated Transformer (Jay Alammar's blog)