Getting Started with Transformers in PyTorch

1. Why Transformers? The Fall of Recurrent Networks

To appreciate why Transformers conquered modern deep learning, we must first understand what made previous sequence models hit a performance brick wall.

For years, Natural Language Processing relied on Recurrent Neural Networks (RNNs), LSTMs (Long Short-Term Memory), and GRUs (Gated Recurrent Units). These architectures process inputs sequentially, updating a hidden state vector $h_t$ one step at a time:

This sequential formulation introduces two fatal limitations:

1. No GPU Parallelization: You cannot calculate the hidden state of token 50 until tokens 1 through 49 have finished calculating sequentially. Training speed is fundamentally throttled by sequence length ($O(N)$ sequential operations).
2. Vanishing Context & Information Bottleneck: The entire history of a sentence must be crammed into a fixed-size vector $h_t$. As sentences exceed 100 or 500 words, early context decays exponentially, causing models to "forget" distant antecedents and subject-verb agreements.

The Transformer eliminates recurrence completely. Instead of feeding hidden states along a temporal timeline, it uses Self-Attention to allow every single token to directly reference every other token in the sequence in $O(1)$ sequential operations. Modern GPUs can compute this through massive, highly optimized matrix multiplications.

2. The Transformer Blueprint: High-Level Architecture

The original Transformer follows an Encoder-Decoder structure designed for sequence-to-sequence translation (such as translating English to German):

Every Transformer block consists of the same core sub-modules:

• Input Embeddings & Positional Encodings: Convert discrete token IDs into dense geometric vectors infused with position.
• Multi-Head Self-Attention: Enable tokens to dynamically attend to relevant words across multiple representational subspaces.
• Position-Wise Feed-Forward Networks (FFN): Non-linear token-level feature transformations.
• Residual Connections & Layer Normalization: Preserve gradient flow and prevent internal covariate shift during deep backpropagation.

3. Step 1: Input Embeddings & Positional Encoding

Computers cannot read text natively; words are first split into numerical token IDs via tokenizers (like Byte-Pair Encoding). An embedding layer maps each integer ID to a continuous vector of dimension $d_{\text{model}}$ (e.g. 512 or 4096).

Why Positional Encoding is Mandatory

Self-attention is permutation-invariant. Without position information, the sentence "The dog bit the man" and "The man bit the dog" would yield identical attention representations! Since Transformers have no recurrence or convolution, we must explicitly inject order into the embeddings.

Vaswani et al. introduced sinusoidal positional encodings using combinations of sine and cosine functions of varying frequencies:

Here, $pos$ is the token's position in the sequence, and $i$ corresponds to the dimension index. Low dimensions alternate rapidly (high frequency), while high dimensions alternate slowly (low frequency), enabling the model to extrapolate to arbitrary sequence lengths.

Let's implement this cleanly in PyTorch:

import torch
import torch.nn as nn
import math

class PositionalEncoding(nn.Module):
    """
    Injects sinusoidal position embeddings into token embeddings.
    Shape: (batch_size, seq_len, d_model) -> (batch_size, seq_len, d_model)
    """
    def __init__(self, d_model: int, max_len: int = 5000, dropout: float = 0.1):
        super().__init__()
        self.dropout = nn.Dropout(p=dropout)

        # Initialize position matrix: (max_len, 1)
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        
        # Calculate frequency scaling: 10000^(2i / d_model)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))

        # Apply sine to even indices, cosine to odd indices
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)

        # Reshape to (1, max_len, d_model) for batch broadcasting
        pe = pe.unsqueeze(0)
        
        # Register as non-trainable persistent buffer
        self.register_buffer('pe', pe)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # x shape: (batch_size, seq_len, d_model)
        x = x + self.pe[:, :x.size(1)]
        return self.dropout(x)

4. Step 2: Demystifying Attention — Queries, Keys, and Values

The mathematical heart of the Transformer is Scaled Dot-Product Attention. It relies on three representations for each token:

The Formula & Why We Scale

The attention calculation is expressed as:

Let's dissect each component of this equation:

• $Q K^T$: Computes the raw dot-product similarity between every query and every key. For a sequence of length $N$, this produces an $(N \times N)$ attention matrix containing pairwise compatibility scores.
• Scaling by $\frac{1}{\sqrt{d_k}}$: As the key dimension $d_k$ grows large, dot products grow substantially in magnitude. Large inputs push the softmax function into regions with near-zero gradients (saturation), paralyzing model training. Dividing by $\sqrt{d_k}$ stabilizes variance to approximately 1.0.
• Softmax: Converts raw compatibility scores into a valid probability distribution where each row sums to 1.0.
• Multiplication with $V$: Computes a weighted sum of the values based on attention probabilities.

def scaled_dot_product_attention(Q, K, V, mask=None, dropout=None):
    """
    Computes Scaled Dot-Product Attention:
    Q, K, V: shape (batch_size, num_heads, seq_len, d_k)
    """
    d_k = Q.size(-1)
    
    # 1. Compute pairwise attention scores: (batch, heads, seq_len, seq_len)
    scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(d_k)
    
    # 2. Apply optional causal or padding mask
    if mask is not None:
        scores = scores.masked_fill(mask == 0, -1e9)
        
    # 3. Normalize scores across rows with Softmax
    attn_weights = torch.softmax(scores, dim=-1)
    
    if dropout is not None:
        attn_weights = dropout(attn_weights)
        
    # 4. Weighted combination of values
    output = torch.matmul(attn_weights, V)
    return output, attn_weights

5. Step 3: Multi-Head Attention (MHA) in PyTorch

If we only perform a single attention computation, the model can only focus on one dominant relationship at a time. For instance, in the sentence "The bank robber crossed the river bank", the word "bank" must attend to both "robber" (financial meaning) and "river" (geographical meaning).

Multi-Head Attention (MHA) solves this by projecting $Q, K, V$ into $h$ separate low-dimensional subspaces (e.g. $h = 8$ heads, each with $d_k = 512 / 8 = 64$). Each head independently learns different patterns (syntactic dependencies, semantic categories, punctuation cues, etc.).

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model: int = 512, num_heads: int = 8, dropout: float = 0.1):
        super().__init__()
        assert d_model % num_heads == 0, "d_model must be divisible by num_heads"
        
        self.d_model = d_model
        self.num_heads = num_heads
        self.d_k = d_model // num_heads

        # Linear projections for Query, Key, and Value
        self.W_q = nn.Linear(d_model, d_model)
        self.W_k = nn.Linear(d_model, d_model)
        self.W_v = nn.Linear(d_model, d_model)
        
        # Final output projection
        self.W_o = nn.Linear(d_model, d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, mask: torch.Tensor = None):
        batch_size = q.size(0)

        # 1. Project inputs: (batch, seq_len, d_model)
        # 2. Reshape and transpose to: (batch, num_heads, seq_len, d_k)
        q = self.W_q(q).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        k = self.W_k(k).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        v = self.W_v(v).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)

        # 3. Scaled dot-product attention per head in parallel
        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k)
        
        if mask is not None:
            scores = scores.masked_fill(mask == 0, -1e9)
            
        attn_weights = torch.softmax(scores, dim=-1)
        attn_weights = self.dropout(attn_weights)
        
        # (batch, num_heads, seq_len, d_k)
        context = torch.matmul(attn_weights, v)

        # 4. Concatenate all heads back together: (batch, seq_len, d_model)
        context = context.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model)

        # 5. Final linear projection
        return self.W_o(context), attn_weights

6. Step 4: Position-Wise Feed-Forward & Encoder Block

After attention captures relationships between different tokens, each token passes independently through a Position-Wise Feed-Forward Network (FFN). The FFN expands the dimension by $4\times$ ($d_{ff} = 2048$), applies a non-linearity (GELU or ReLU), and projects back down to $d_{\text{model}}$.

class PositionWiseFeedForward(nn.Module):
    def __init__(self, d_model: int = 512, d_ff: int = 2048, dropout: float = 0.1):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(d_ff, d_model),
            nn.Dropout(dropout)
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.net(x)

Pre-LN vs Post-LN Architecture

The original paper used Post-LayerNorm (LayerNorm applied after residual addition: $\text{LayerNorm}(x + \text{Sublayer}(x))$). However, virtually all modern large language models (such as GPT-3, LLaMA, and Mistral) utilize Pre-LayerNorm (LayerNorm applied before the sublayer: $x + \text{Sublayer}(\text{LayerNorm}(x))$), which ensures stable gradient propagation and eliminates the need for warm-up learning rate schedules.

class TransformerEncoderBlock(nn.Module):
    def __init__(self, d_model: int = 512, num_heads: int = 8, d_ff: int = 2048, dropout: float = 0.1):
        super().__init__()
        self.attn = MultiHeadAttention(d_model, num_heads, dropout)
        self.ffn = PositionWiseFeedForward(d_model, d_ff, dropout)
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x: torch.Tensor, mask: torch.Tensor = None):
        # Pre-LN Self-Attention with Residual Connection
        norm_x = self.norm1(x)
        attn_out, _ = self.attn(norm_x, norm_x, norm_x, mask)
        x = x + self.dropout(attn_out)

        # Pre-LN Feed-Forward with Residual Connection
        x = x + self.ffn(self.norm2(x))
        return x

7. Step 5: Assembling a Complete Transformer Classifier

Let's assemble our modular components into a fully functional end-to-end Transformer classifier in PyTorch — suitable for sentiment analysis, document classification, or intent detection.

class TransformerClassifier(nn.Module):
    def __init__(
        self,
        vocab_size: int,
        num_classes: int,
        d_model: int = 256,
        num_heads: int = 4,
        num_layers: int = 3,
        d_ff: int = 512,
        max_len: int = 256,
        dropout: float = 0.1
    ):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, d_model)
        self.pos_encoder = PositionalEncoding(d_model, max_len, dropout)
        
        # Stack multiple Transformer Encoder blocks
        self.layers = nn.ModuleList([
            TransformerEncoderBlock(d_model, num_heads, d_ff, dropout)
            for _ in range(num_layers)
        ])
        
        self.norm = nn.LayerNorm(d_model)
        self.classifier = nn.Sequential(
            nn.Linear(d_model, d_model // 2),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(d_model // 2, num_classes)
        )

    def forward(self, input_ids: torch.Tensor, mask: torch.Tensor = None) -> torch.Tensor:
        # 1. Embed tokens and add positional encoding: (batch, seq_len, d_model)
        x = self.embedding(input_ids)
        x = self.pos_encoder(x)

        # 2. Pass through stacked encoder layers
        for layer in self.layers:
            x = layer(x, mask)

        # 3. Final layer normalization
        x = self.norm(x)

        # 4. Global Mean Pooling across sequence dimension: (batch, d_model)
        pooled = x.mean(dim=1)

        # 5. Predict class logits: (batch, num_classes)
        logits = self.classifier(pooled)
        return logits

Running an End-to-End Test in PyTorch

We can verify our architecture with a live synthetic test pass:

# Hyperparameters
VOCAB_SIZE = 10000
NUM_CLASSES = 2  # Binary sentiment: Positive (1) vs Negative (0)
BATCH_SIZE = 8
SEQ_LEN = 32

# Instantiate model, optimizer and loss function
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = TransformerClassifier(
    vocab_size=VOCAB_SIZE,
    num_classes=NUM_CLASSES,
    d_model=128,
    num_heads=4,
    num_layers=2
).to(device)

criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=1e-2)

# Synthetic batch simulation
dummy_tokens = torch.randint(0, VOCAB_SIZE, (BATCH_SIZE, SEQ_LEN)).to(device)
dummy_labels = torch.randint(0, NUM_CLASSES, (BATCH_SIZE,)).to(device)

# Forward, backward, and optimization step
optimizer.zero_grad()
outputs = model(dummy_tokens)
loss = criterion(outputs, dummy_labels)
loss.backward()
optimizer.step()

print(f"✅ Forward pass successful!")
print(f"Logits shape: {outputs.shape}  |  Initial Loss: {loss.item():.4f}")
# Output: Logits shape: torch.Size([8, 2])  |  Initial Loss: 0.6931

8. Modern LLM Evolutions & Production Best Practices

While the original 2017 Transformer architecture remains the foundational blueprint, modern LLM research has introduced several critical hardware and algorithmic optimizations:

9. Transformer Family Comparison: BERT vs. GPT vs. T5

Depending on which sub-components of the original Transformer you utilize, you get distinct model families tailored to specific tasks:

10. Summary & Knowledge Check

Congratulations! You have demystified the inner mechanics of the most influential deep learning architecture of the 21st century and implemented every foundational layer from scratch in PyTorch.

Ready to dive deeper? Check out our interactive Deep Learning & AI Courses or explore our companion guide on What Is Generative AI?.