Transformer Architecture Visual Guide
Self-Attention Explained Step by Step

What Is Self-Attention?

Self-attention allows every position in a sequence to attend to every other position in the same sequence. Unlike recurrent networks, all positions are processed in parallel — giving transformers their speed advantage and ability to capture long-range dependencies.

The Query, Key, Value Framework

Each input token is projected into three vectors using learned weight matrices:

• Query (Q) — What this token is looking for
• Key (K) — What this token offers to other tokens
• Value (V) — The actual content to be aggregated

The scaling factor sqrt(d_k) prevents dot products from growing too large in high-dimensional spaces, which would push softmax into saturation zones with near-zero gradients.

PyTorch Implementation

import torch
import torch.nn as nn
import torch.nn.functional as F
import math

class ScaledDotProductAttention(nn.Module):
    def forward(self, Q, K, V, mask=None):
        d_k = Q.size(-1)
        # Attention scores: (batch, heads, seq, seq)
        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(d_k)
        if mask is not None:
            scores = scores.masked_fill(mask == 0, float('-inf'))
        attn_weights = F.softmax(scores, dim=-1)
        return torch.matmul(attn_weights, V), attn_weights


class MultiHeadAttention(nn.Module):
    def __init__(self, d_model=512, num_heads=8):
        super().__init__()
        assert d_model % num_heads == 0
        self.d_k = d_model // num_heads
        self.num_heads = num_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)
        self.attention = ScaledDotProductAttention()

    def split_heads(self, x):
        B, S, D = x.shape
        return x.view(B, S, self.num_heads, self.d_k).transpose(1, 2)

    def forward(self, Q, K, V, mask=None):
        Q = self.split_heads(self.W_q(Q))  # (B, H, S, d_k)
        K = self.split_heads(self.W_k(K))
        V = self.split_heads(self.W_v(V))
        x, weights = self.attention(Q, K, V, mask)
        B, H, S, dk = x.shape
        x = x.transpose(1, 2).contiguous().view(B, S, H * dk)
        return self.W_o(x), weights  # (B, S, d_model)


# Usage
d_model, num_heads, seq_len, batch = 512, 8, 64, 2
x = torch.randn(batch, seq_len, d_model)
mha = MultiHeadAttention(d_model, num_heads)
out, attn = mha(x, x, x)  # Self-attention: Q=K=V=x
print(out.shape)   # (2, 64, 512)
print(attn.shape)  # (2, 8, 64, 64) — attention maps per head

Key Dimensions to Know

For a standard BERT-base transformer (d_model=768, num_heads=12):

• d_k = d_model / num_heads = 768 / 12 = 64
• Q, K, V matrices: (seq_len, 64) per head
• Attention matrix: (seq_len, seq_len) — quadratic in sequence length
• Multi-head output: (seq_len, 768) after concatenation + projection

Why Multi-Head?

Multiple attention heads allow the model to jointly attend to information from different representation subspaces. Head 1 might capture syntactic dependencies, head 5 might track coreference, and head 11 might focus on positional proximity — all in parallel within a single layer.