Transformer Architecture Explained: A Visual Guide

Transformer Architecture Explained: A Visual Guide

The Transformer is arguably the most important neural network architecture of the last decade. Introduced in the 2017 paper "Attention Is All You Need" by Vaswani et al., it replaced recurrent networks (RNNs/LSTMs) as the dominant approach for sequence modeling — and its influence now extends to vision, audio, biology, and robotics.

Every major language model today — GPT-4, Claude, Llama, Gemini — is built on the Transformer. Understanding this architecture is essential for anyone working in ML.

The Big Picture

At its core, the Transformer is an encoder-decoder model designed for sequence-to-sequence tasks (e.g., translation). But modern variants often use only one half:

• Encoder-only (e.g., BERT) — great for classification, embeddings, understanding
• Decoder-only (e.g., GPT, Llama) — great for text generation
• Encoder-decoder (e.g., T5, BART) — great for translation, summarization

Input Tokens → [Encoder] → Latent Representation → [Decoder] → Output Tokens

Step 1: Input Embedding + Positional Encoding

Transformers process all tokens in parallel (unlike RNNs, which process sequentially). But this means they have no inherent sense of word order. Positional encodings are added to token embeddings to inject position information:

import torch
import torch.nn as nn
import math

class PositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=5000):
        super().__init__()
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len).unsqueeze(1).float()
        div_term = torch.exp(
            torch.arange(0, d_model, 2).float()
            * (-math.log(10000.0) / d_model)
        )
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        self.register_buffer("pe", pe.unsqueeze(0))

    def forward(self, x):
        return x + self.pe[:, :x.size(1)]

The original paper uses sinusoidal functions (shown above), while modern models often use learned positional embeddings or Rotary Position Embeddings (RoPE).

Step 2: Multi-Head Self-Attention

This is the heart of the Transformer. Self-attention lets every token attend to every other token in the sequence, capturing dependencies regardless of distance.

For each token, we compute three vectors:

• Query (Q) — "What am I looking for?"
• Key (K) — "What do I contain?"
• Value (V) — "What information do I provide?"

The attention scores are computed as:

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

Multi-head attention runs this computation multiple times in parallel with different learned projections, allowing the model to attend to different types of relationships simultaneously:

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, n_heads):
        super().__init__()
        self.n_heads = n_heads
        self.d_k = d_model // n_heads
        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)
        self.W_o = nn.Linear(d_model, d_model)

    def forward(self, x):
        B, L, D = x.shape
        Q = self.W_q(x).view(B, L, self.n_heads, self.d_k).transpose(1, 2)
        K = self.W_k(x).view(B, L, self.n_heads, self.d_k).transpose(1, 2)
        V = self.W_v(x).view(B, L, self.n_heads, self.d_k).transpose(1, 2)

        scores = (Q @ K.transpose(-2, -1)) / math.sqrt(self.d_k)
        attn = torch.softmax(scores, dim=-1)
        out = (attn @ V).transpose(1, 2).contiguous().view(B, L, D)
        return self.W_o(out)

Step 3: Feed-Forward Network

After attention, each token passes through a position-wise feed-forward network — two linear layers with a non-linearity in between:

FFN(x) = ReLU(xW₁ + b₁)W₂ + b₂

Modern Transformers often replace ReLU with GeLU or SwiGLU activations for better performance.

Step 4: Layer Norm + Residual Connections

Each sub-layer (attention and FFN) is wrapped with a residual connection and layer normalization:

output = LayerNorm(x + SubLayer(x))

Residual connections prevent the vanishing gradient problem in deep networks. Layer normalization stabilizes training. Modern architectures often use Pre-Norm (normalize before the sub-layer) instead of the original Post-Norm.

The Decoder Difference

The decoder adds one critical modification: causal (masked) self-attention. When generating text left-to-right, each token should only attend to previous tokens, not future ones. This is achieved by masking future positions with negative infinity before the softmax:

mask = torch.triu(torch.ones(L, L), diagonal=1).bool()
scores = scores.masked_fill(mask, float("-inf"))

Why Transformers Won

The Transformer's dominance comes down to three key advantages:

1. Parallelism — Unlike RNNs, all tokens are processed simultaneously during training, enabling massive GPU utilization
2. Long-range dependencies — Self-attention connects any two tokens in O(1) layers, while RNNs require O(n) steps
3. Scalability — The architecture scales predictably with more parameters, data, and compute (the "scaling laws" that drive modern AI)

Beyond Text

The Transformer architecture has been successfully adapted to nearly every domain: Vision Transformers (ViT) for images, Audio Spectrogram Transformers for sound, AlphaFold for protein structure prediction, and Decision Transformers for reinforcement learning. Its generality is one of its greatest strengths.

Understanding the Transformer is not just an academic exercise — it's the foundation for virtually all modern AI systems.

---

Want to go deeper? Check out our ML Basic Course →