Unpacking Transformers: The Foundation of LLM Zero to Hero#

Why Did Transformers Become So Popular?#

1. Parallelization: Unlike RNNs, the Transformer architecture does not rely on sequential processing of input tokens. This results in higher computational efficiency and significantly faster training times, especially on large datasets.
2. Long-Range Dependencies: Transformers utilize attention mechanisms that help the model focus on different parts of the input sequence. This is in contrast to RNNs that tend to forget older tokens over longer sequences.
3. State-of-the-Art Results: Since their introduction, Transformers have consistently led to breakthroughs in NLP tasks such as machine translation, text generation, question answering, summarization, and much more.

Whether you are aiming to build a new chatbot, a machine translation tool, or a text summarization service, understanding Transformers is your key to unlocking the power behind state-of-the-art language models.

Attention Mechanism#

The concept of “attention” was introduced to let models focus on the most relevant parts of an input sequence when making a prediction. In a typical sequence-to-sequence task (like machine translation), the model attempts to align segments of the output with segments of the input.

An attention mechanism calculates a weighted sum of “value” vectors, where the weights derive from how well a “query” vector aligns with a series of “key” vectors:

1. Queries (Q): Represents the input that queries information from the entire sequence.
2. Keys (K): Acts like an address or a label to help match queries to relevant information.
3. Values (V): Contains the actual information.

The attention formula can be summarized as:

Attention(Q, K, V) = softmax( (QKᵀ) / √d_k ) × V

Where d_k is the dimension of the key vectors. The division by √d_k stabilizes gradients by mitigating large dot product values when dimensions are large.

Positional Encoding#

Since Transformers do not process tokens in a strictly sequential manner like RNNs, they need a way to incorporate positional information. This is where positional encoding comes in. Each token in the sequence receives a positional embedding, which describes its location within the sequence.

A common approach uses sinusoidal functions:

PE(pos, 2i) = sin(pos / 10000^(2i / d_model))
PE(pos, 2i+1) = cos(pos / 10000^(2i / d_model))

Here:

• pos is the position in the sequence (0, 1, 2, …),
• i is the dimension index,
• d_model is the embedding dimension.

This allows the model to learn relative position relationships between different tokens without losing the parallelization benefit.

Multi-Head Attention#

Instead of computing attention once, the Transformer splits the dimension of the embeddings into multiple heads. Each head performs an attention operation independently, and their outputs are concatenated, then linearly transformed. This approach allows the model to focus on different parts of the sequence from multiple perspectives.

Formally, multi-head attention can be written as:

MultiHead(Q, K, V) = Concat( head₁, head₂, …, headₕ ) × Wᴼ

where each headₖ is computed as:

headₖ = Attention(QWₖ^Q, KWₖ^K, VWₖ^V)

Residual Connections, Layer Normalization, and Feed-Forward Layers#

Each sub-layer in the Transformer (e.g., multi-head self-attention, feed-forward) is wrapped with residual connections and followed by layer normalization:

1. Residual Connections: Add the original input to the sub-layer’s output, which helps preserve gradients and mitigate the vanishing or exploding gradient problem.
2. Layer Normalization: Normalizes the summed output, smoothing optimization.
3. Position-wise Feed-Forward Layer: A fully connected feed-forward network applied to each attention output vector, typically featuring a non-linear activation like ReLU or GELU.

Prerequisites and Basic Setup#

We will assume you have:

• Python 3.7+
• PyTorch >= 1.10
• NumPy

Install the core dependencies (if you haven’t already):

1
pip install torch numpy

Defining the Model#

Let’s start by setting up a skeleton Transformer class. We’ll go step-by-step, focusing first on the building blocks.

1
import torch
2
import torch.nn as nn
3
import math
4

5
class Transformer(nn.Module):
6
    def __init__(self, d_model, n_heads, num_encoder_layers, num_decoder_layers, vocab_size, max_seq_len=512):
7
        super(Transformer, self).__init__()
8
        self.d_model = d_model
9
        self.embed = nn.Embedding(vocab_size, d_model)
10

11
        self.pos_encoding = PositionalEncoding(d_model, max_seq_len)
12

13
        encoder_layer = nn.TransformerEncoderLayer(d_model, n_heads)
14
        decoder_layer = nn.TransformerDecoderLayer(d_model, n_heads)
15

16
        self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=num_encoder_layers)
17
        self.decoder = nn.TransformerDecoder(decoder_layer, num_layers=num_decoder_layers)
18

19
        self.fc_out = nn.Linear(d_model, vocab_size)
20

21
    def forward(self, src, tgt):
22
        # src: [batch_size, src_seq_len]
23
        # tgt: [batch_size, tgt_seq_len]
24

25
        # Embedding + positional encoding
26
        src_emb = self.pos_encoding(self.embed(src) * math.sqrt(self.d_model))
27
        tgt_emb = self.pos_encoding(self.embed(tgt) * math.sqrt(self.d_model))
28

29
        # PyTorch transformer expects shape: [src_seq_len, batch_size, d_model]
30
        src_emb = src_emb.transpose(0, 1)
31
        tgt_emb = tgt_emb.transpose(0, 1)
32

33
        memory = self.encoder(src_emb)
34
        output = self.decoder(tgt_emb, memory)
35

36
        # Convert back to token predictions
37
        output = self.fc_out(output)
38

39
        # [tgt_seq_len, batch_size, vocab_size] => we can transpose if needed
40
        return output.transpose(0, 1)
41

42
class PositionalEncoding(nn.Module):
43
    def __init__(self, d_model, max_len=512):
44
        super(PositionalEncoding, self).__init__()
45
        position = torch.arange(0, max_len).unsqueeze(1)
46
        div_term = torch.exp(torch.arange(0, d_model, 2) * -(math.log(10000.0) / d_model))
47

48
        pe = torch.zeros(max_len, d_model)
49
        pe[:, 0::2] = torch.sin(position * div_term)
50
        pe[:, 1::2] = torch.cos(position * div_term)
51

52
        self.register_buffer('pe', pe.unsqueeze(0))
53

54
    def forward(self, x):
55
        # x: [batch_size, seq_len, d_model]
56
        seq_len = x.size(1)
57
        x = x + self.pe[:, :seq_len]
58
        return x

Here’s what’s happening in this snippet:

• We create a simple Transformer class that uses PyTorch’s built-in nn.TransformerEncoder and nn.TransformerDecoder for brevity.
• The PositionalEncoding class implements sinusoidal positional encodings.
• We multiply the embedding by math.sqrt(self.d_model) because in the original paper, the transformer’s embeddings are scaled by √d_model to account for the variance in the embeddings.

Implementing Multi-Head Self-Attention#

In practice, PyTorch has a ready-made attention mechanism, but let’s see how you might write a simplified version of the scaled dot-product attention:

1
def scaled_dot_product_attention(Q, K, V, mask=None):
2
    d_k = K.size(-1)
3
    scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(d_k)
4
    if mask is not None:
5
        scores = scores.masked_fill(mask == 0, float('-inf'))
6
    attn_weights = torch.softmax(scores, dim=-1)
7
    output = torch.matmul(attn_weights, V)
8
    return output, attn_weights

Now, for multi-head attention:

1
class MultiHeadAttention(nn.Module):
2
    def __init__(self, d_model, n_heads):
3
        super(MultiHeadAttention, self).__init__()
4
        self.d_model = d_model
5
        self.n_heads = n_heads
6

7
        assert d_model % n_heads == 0, "d_model must be divisible by n_heads"
8

9
        self.d_k = d_model // n_heads
10

11
        self.W_q = nn.Linear(d_model, d_model)
12
        self.W_k = nn.Linear(d_model, d_model)
13
        self.W_v = nn.Linear(d_model, d_model)
14

15
        self.out = nn.Linear(d_model, d_model)
16

17
    def forward(self, Q, K, V, mask=None):
18
        # Q, K, V: [batch_size, seq_len, d_model]
19
        batch_size = Q.size(0)
20

21
        Q = self.W_q(Q).view(batch_size, -1, self.n_heads, self.d_k)
22
        K = self.W_k(K).view(batch_size, -1, self.n_heads, self.d_k)
23
        V = self.W_v(V).view(batch_size, -1, self.n_heads, self.d_k)
24

25
        # transpose for attention calculation: [batch_size, n_heads, seq_len, d_k]
26
        Q = Q.transpose(1, 2)
27
        K = K.transpose(1, 2)
28
        V = V.transpose(1, 2)
29

30
        # scaled dot-product attention
31
        attn_output, attn_weights = scaled_dot_product_attention(Q, K, V, mask=mask)
32

33
        # concat heads
34
        attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model)
35

36
        # final linear layer
37
        output = self.out(attn_output)  # [batch_size, seq_len, d_model]
38

39
        return output, attn_weights

The above code should give you a sense for how attention is computed under the hood.

Putting It All Together#

In real-world scenarios, you’ll likely rely on pre-built modules (as with PyTorch’s nn.Transformer). However, looking under the hood helps clarify how these components interact. When training, you’ll need to:

1. Tokenize and numericalize your text.
2. Convert them to tensors and feed them into the model.
3. Apply masking (for example, to mask future tokens during training for a language modeling task).
4. Compute the loss (often cross-entropy for token classification).
5. Backpropagate and optimize.

Data Preparation#

For sequence-to-sequence tasks, you typically need parallel data: pairs of sentences in source and target languages. Assume you have a dataset of English-French sentence pairs:

1. Tokenization: Use a tokenizer (like Byte Pair Encoding - BPE, or a standard library such as Hugging Face’s tokenizers).
2. Convert tokens to indices: Build a vocabulary or use a pre-built vocabulary.
3. Create training pairs: src_sequences, tgt_sequences.
4. Add special tokens: Usually, <pad>, <sos>, <eos>.

Store these in a PyTorch Dataset and then wrap it with a DataLoader.

Hyperparameters and Optimization#

Key hyperparameters:

• d_model (embedding dimension): A typical starting point is 512.
• n_heads: 8 or 16 are common choices.
• number of layers (encoder/decoder): 6 each in the original paper, but can vary.
• learning rate: Often scheduled, e.g., using an Adam optimizer with “warmup steps.”

Example of an Adam optimizer with warmup scheduling:

1
import torch.optim as optim
2

3
model = Transformer(d_model=512, n_heads=8, num_encoder_layers=6, num_decoder_layers=6, vocab_size=10000)
4

5
optimizer = optim.Adam(model.parameters(), lr=1e-4, betas=(0.9, 0.98), eps=1e-9)
6
criterion = nn.CrossEntropyLoss(ignore_index=0)  # ignoring <pad> token
7

8
def train_batch(src_batch, tgt_batch):
9
    model.train()
10
    optimizer.zero_grad()
11

12
    # We assume tgt_batch has <sos> + ... + <eos>
13
    # Typically, you might split the target so the model predicts the next token
14
    # from the partial. For example, input to decoder is [<sos>, w1, w2, w3]
15
    # and we compare output with [w1, w2, w3, <eos>].
16

17
    output = model(src_batch, tgt_batch[:, :-1])  # all but the last token
18
    # output has shape [batch_size, tgt_seq_len-1, vocab_size]
19

20
    loss = criterion(output.reshape(-1, output.size(-1)), tgt_batch[:, 1:].reshape(-1))
21
    loss.backward()
22
    optimizer.step()
23

24
    return loss.item()

Evaluation Metrics#

For machine translation, common metrics include:

• BLEU (Bilingual Evaluation Understudy)
• ROUGE (Recall-Oriented Understudy for Gisting Evaluation)
• METEOR
• ChrF++ (especially for morphologically-rich languages)

You compute these metrics by comparing your model’s predictions with that of a reference translation across multiple test samples.

BERT, GPT, and T5#

• BERT (Bidirectional Encoder Representations from Transformers)

  • Uses only the encoder component and is pre-trained on tasks like masked language modeling (MLM) and next sentence prediction.
  • Excelled at tasks requiring a deep understanding of context (e.g., question answering, classification).

• GPT (Generative Pre-trained Transformer)

  • Based on the decoder portion, uses autoregressive language modeling.
  • GPT-2 and GPT-3 significantly scaled up parameters, establishing new benchmarks in generating coherent, contextually relevant text.

• T5 (Text-to-Text Transfer Transformer)

  • Treats every NLP task (classification, translation, summarization, etc.) as a text-to-text problem.
  • Uses an encoder-decoder architecture, pretrained on a massive “Fill-In-The-Blank” objective.

Pre-training and Fine-tuning#

Modern practice often revolves around:

1. Pre-training on large corpora with a generic task (e.g., language modeling, masked language modeling).
2. Fine-tuning on a specific dataset with a smaller, specialized objective.

This approach leverages learned language structures and dramatically reduces the need for large domain-specific datasets when building specialized models.

Scaling Laws and Transformer Variations#

• Scaling Laws: As you increase model size, dataset size, and compute, performance typically grows in a predictable way. Larger Transformers (e.g., GPT-3, GPT-4) are prime examples.
• Transformer Variants:
  • Longformer, BigBird: Adapted for very long sequences by introducing sparse attention.
  • Reformer: Reduces complexity using locality-sensitive hashing for attention.
  • Performers: Uses random feature maps for more efficient attention.

Combining Transformers with Other Architectures#

Transformers can be augmented or combined with other architectures for specialized tasks:

• ConvBERT mixes the strengths of CNN-like local feature extraction with Transformer context.
• Efficient Transformers: Weighted kernel-based methods and recurrence-based methods can be combined with self-attention for further memory or computational gains.