Mastering Transformers: Key Extensions and Optimization Techniques Explained

1. Overall Structure of the Transformer

The Transformer is built from an Encoder and a Decoder, each typically composed of six identical blocks (the exact number is flexible). The Encoder processes the input sequence using self‑attention and a feed‑forward network, while the Decoder adds an extra attention layer that focuses on the Encoder’s output.

1.1 Encoder Structure

Each Encoder block contains a self‑attention layer followed by a feed‑forward network, both wrapped with residual connections (Add) and layer normalization (Norm). The same structure repeats across all blocks, producing an encoded representation matrix C .

1.2 Decoder Structure

The Decoder mirrors the Encoder but includes two attention layers: the first is a masked self‑attention to prevent attending to future tokens, and the second attends to the Encoder’s output C . A final softmax layer predicts the next token.

2. Processing Flow

Step 1: Each word in the input sentence is converted to a vector X by summing its word embedding and positional embedding.

Step 2: The matrix X is fed into the Encoder. After six Encoder blocks, the sentence is represented by matrix C .

Step 3: The Encoder output C is passed to the Decoder. The Decoder generates tokens sequentially, using a mask to hide future positions during training.

3. Input Representation

Each token’s representation x is the sum of its word embedding (e.g., Word2Vec or GloVe) and a positional embedding PE . The positional embedding uses sinusoidal functions to encode absolute or relative positions, allowing the model to handle sequences longer than those seen during training.

3.1 Word Embedding

Word embeddings can be pretrained with algorithms such as Word2Vec or GloVe, or learned jointly within the Transformer.

3.2 Positional Embedding

Because the Transformer lacks recurrence, positional embeddings inject order information. The sinusoidal formula ensures that embeddings for unseen positions can be computed and that relative distances are easy to derive.

It enables the model to handle sequences longer than any seen during training.

It allows simple computation of relative positions, e.g., PE(pos+k) from PE(pos) .

Adding word and positional embeddings yields the final input vector x .

4. Self‑Attention Mechanism

Self‑Attention computes three matrices: queries Q , keys K , and values V by linear projections of the input X . The attention scores are obtained by softmax((QKᵀ)/√d) , where d is the dimension of the vectors. The scores weight the values V to produce the output Z .

4.1 Multi‑Head Attention

Multiple self‑attention heads run in parallel, each with its own Q, K, V projections. Their outputs are concatenated and linearly transformed, allowing the model to capture diverse relational patterns.

5. Encoder Block Details

Each Encoder block consists of:

Multi‑Head Self‑Attention

Add & Norm (residual connection + layer normalization)

Feed‑Forward network (two linear layers with ReLU activation)

Another Add & Norm

6. Decoder Block Details

The Decoder block is similar but adds:

A masked Multi‑Head Self‑Attention (prevents attending to future tokens)

A second Multi‑Head Attention that queries the Encoder output C while using the Decoder’s previous output as queries

A final linear layer followed by softmax to predict the next token.

6.1 Masked Self‑Attention

During training, a mask matrix ensures that position i can only attend to positions ≤ i , preserving the autoregressive property.

7. Summary

Transformers replace recurrent structures with parallelizable self‑attention, enabling efficient training on large datasets. Positional embeddings inject order information, while multi‑head attention captures multiple relational aspects. The modular encoder‑decoder design, combined with residual connections and layer normalization, makes Transformers the foundation of modern NLP models.

8. PyTorch Implementations

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

# Define self‑attention module
class SelfAttention(nn.Module):
    def __init__(self, embed_dim):
        super(SelfAttention, self).__init__()
        self.query = nn.Linear(embed_dim, embed_dim)
        self.key = nn.Linear(embed_dim, embed_dim)
        self.value = nn.Linear(embed_dim, embed_dim)
    def forward(self, x):
        q = self.query(x)
        k = self.key(x)
        v = self.value(x)
        attn_weights = torch.matmul(q, k.transpose(1, 2))
        attn_weights = nn.functional.softmax(attn_weights, dim=-1)
        attended_values = torch.matmul(attn_weights, v)
        return attended_values

# Self‑attention classifier
class SelfAttentionClassifier(nn.Module):
    def __init__(self, embed_dim, hidden_dim, num_classes):
        super(SelfAttentionClassifier, self).__init__()
        self.attention = SelfAttention(embed_dim)
        self.fc1 = nn.Linear(embed_dim, hidden_dim)
        self.fc2 = nn.Linear(hidden_dim, num_classes)
    def forward(self, x):
        attended_values = self.attention(x)
        x = attended_values.mean(dim=1)  # average over sequence length
        x = self.fc1(x)
        x = torch.relu(x)
        x = self.fc2(x)
        return x

# Multi‑head self‑attention module
class MultiHeadSelfAttention(nn.Module):
    def __init__(self, embed_dim, num_heads):
        super(MultiHeadSelfAttention, self).__init__()
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads
        self.query = nn.Linear(embed_dim, embed_dim)
        self.key = nn.Linear(embed_dim, embed_dim)
        self.value = nn.Linear(embed_dim, embed_dim)
        self.fc = nn.Linear(embed_dim, embed_dim)
    def forward(self, x):
        batch_size, seq_len, embed_dim = x.size()
        q = self.query(x).view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        k = self.key(x).view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        v = self.value(x).view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        attn_weights = torch.matmul(q, k.transpose(-2, -1)) / torch.sqrt(torch.tensor(self.head_dim, dtype=torch.float))
        attn_weights = torch.softmax(attn_weights, dim=-1)
        attended_values = torch.matmul(attn_weights, v).transpose(1, 2).contiguous().view(batch_size, seq_len, embed_dim)
        x = self.fc(attended_values) + x
        return x

# Multi‑head self‑attention classifier
class MultiHeadSelfAttentionClassifier(nn.Module):
    def __init__(self, embed_dim, num_heads, hidden_dim, num_classes):
        super(MultiHeadSelfAttentionClassifier, self).__init__()
        self.attention = MultiHeadSelfAttention(embed_dim, num_heads)
        self.fc1 = nn.Linear(embed_dim, hidden_dim)
        self.fc2 = nn.Linear(hidden_dim, num_classes)
    def forward(self, x):
        x = self.attention(x)
        x = x.mean(dim=1)
        x = self.fc1(x)
        x = torch.relu(x)
        x = self.fc2(x)
        return x