Transformers from Scratch: Self-Attention to Encoder–Decoder

A Transformer uses attention to mix information across token positions and a feed-forward network to transform each position. Removing recurrent state updates lets known training positions be processed in parallel within a layer. Autoregressive generation still produces successive tokens sequentially. This article builds the attention and residual blocks, adds position information, and connects an encoder to a causal decoder.

Queries, keys, and values

Queries and keys produce matching scores; values supply the vectors being combined. This is a useful retrieval analogy, but soft attention mixes contributions rather than selecting a single stored item. In self-attention, all three are learned projections of the same input \(X\). For one head and one sequence, \(X\) has shape \((T,d_{model})\), \(Q,K\) have shape \((T,d_k)\), and \(V\) has shape \((T,d_v)\). Rows of the score matrix correspond to queries; columns correspond to keys. Softmax normalizes each row across keys.

\[Q=XW^Q,\quad K=XW^K,\quad V=XW^V,\qquad\text{Attention}(Q,K,V)=\text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V\]

Independent query and key projections allow asymmetric raw scores. Even if raw scores are symmetric, row-wise softmax need not produce a symmetric weight matrix because row normalizers differ. For scores \((0,\log3)\), the weights are \((0.25,0.75)\); with scalar values 2 and 6, the output is 5. A value vector can carry many features, but one head uses the same positional weights across its value coordinates.

Attention was introduced to fix a specific encoder-decoder failure, described in Sequence-to-Sequence Models and Attention Explained.

Why divide by the square root of d

Under the simplifying assumption that all query and key components are mutually independent, mean zero, and variance one, \(q^Tk\) has variance \(d_k\). Dividing by \(\sqrt{d_k}\) makes that variance one. Learned projections need not satisfy these assumptions exactly. Large differences between logits can saturate softmax and attenuate gradients through it; the variance calculation motivates a scale, but does not prove that every unscaled layer stops learning.

import numpy as np

rng = np.random.default_rng(0)
d_k = 64
q = rng.standard_normal((1, d_k))
K = rng.standard_normal((10, d_k))

raw = (q @ K.T).ravel()
scaled = raw / np.sqrt(d_k)

def softmax(z):
    e = np.exp(z - z.max())
    return e / e.sum()

print(round(raw.std(), 2), round(scaled.std(), 2))
# 6.39 0.8
print(round(softmax(raw).max(), 4))
# 0.9297
print(round(softmax(scaled).max(), 4))
# 0.3075

In this draw, the unscaled peak is about 0.93 and the scaled peak about 0.31. These measure concentration, not training success. Results vary with the query, keys, their number, and their distribution. Scaling reduces the dimension-dependent spread under the assumptions above; other normalization or temperature choices are possible.

Masks

Two masks appear, and they serve different purposes.

A key-padding mask excludes padded key/value positions. It does not automatically zero outputs for padded queries or remove padding from a loss. A causal mask allows only current and earlier input positions. With shifted target inputs, later positions can contain the answer being predicted, so omitting causality leaks future information during training. That mismatch can make evaluation misleading; it does not imply one universal loss or generation outcome.

This implementation uses boolean masks with True meaning allowed. Disallowed scores become negative infinity before softmax, giving zero weight provided each row has at least one allowed key. Fully masked rows must be handled explicitly. Do not transfer mask polarity between APIs without checking: PyTorch MultiheadAttention boolean masks mark exclusions, while scaled_dot_product_attention uses True for allowed entries. Adding a negative value after softmax is not a substitute for masking logits.

import torch

def causal_mask(T):
    return torch.tril(torch.ones(T, T, dtype=torch.bool))

print(causal_mask(4).int())
# tensor([[1, 0, 0, 0],
#         [1, 1, 0, 0],
#         [1, 1, 1, 0],
#         [1, 1, 1, 1]], dtype=torch.int32)

Multi-head attention

Multi-head attention gives each position several separately normalized mixtures. With \(h\) heads, the implementation projects to width \(d_{model}\), then splits the projected features into heads of width \(d_k=d_{model}/h\). After attention, it concatenates the head outputs and applies an output projection. Each head learns projections of the full input, not merely a fixed slice of the original token features.

At fixed model width, these full-width Q, K, V and output projections contain \(4d_{model}^2+4d_{model}\) parameters including biases, independent of head count. Equal parameter counts do not imply identical expressive capacity: several attention distributions change the possible computation. Heads may learn different relationships, but distinct interpretable roles are not guaranteed.

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

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, n_heads, dropout=0.1):
        super().__init__()
        if n_heads < 1 or d_model < 1 or d_model % n_heads:
            raise ValueError("positive model width must be divisible by positive head count")
        self.h, self.d_k = n_heads, d_model // n_heads
        self.qkv = nn.Linear(d_model, 3 * d_model)
        self.out = nn.Linear(d_model, d_model)
        self.drop = nn.Dropout(dropout)

    def forward(self, x, mask=None, memory=None):
        if x.ndim != 3 or x.shape[-1] != self.out.in_features:
            raise ValueError("x must be (N,T,d_model)")
        N, T, D = x.shape
        if memory is None:
            q, k, v = self.qkv(x).chunk(3, dim=-1)
            S = T
        else:
            if memory.ndim != 3 or memory.shape[0] != N or memory.shape[2] != D:
                raise ValueError("memory must be (N,S,d_model)")
            S = memory.shape[1]
            wq, wk, wv = self.qkv.weight.chunk(3, dim=0)
            bq, bk, bv = self.qkv.bias.chunk(3)
            q = F.linear(x, wq, bq)
            k, v = F.linear(memory, wk, bk), F.linear(memory, wv, bv)
        if T == 0 or S == 0:
            raise ValueError("empty query or key sequence")
        q = q.reshape(N, T, self.h, self.d_k).transpose(1, 2)
        k, v = (t.reshape(N, S, self.h, self.d_k).transpose(1, 2) for t in (k, v))
        scores = q @ k.transpose(-2, -1) / self.d_k**0.5
        if mask is not None:
            if mask.dtype != torch.bool or mask.device != scores.device:
                raise ValueError("mask must be boolean on the input device")
            if mask.ndim not in (2, 4):
                raise ValueError("use (T,S) or explicit broadcastable (N,h,T,S) axes")
            allowed = torch.broadcast_to(mask, scores.shape)
            if not allowed.any(dim=-1).all():
                raise ValueError("each query must have at least one allowed key")
            scores = scores.masked_fill(~allowed, float("-inf"))
        weights = F.softmax(scores, dim=-1)
        y = (self.drop(weights) @ v).transpose(1, 2).reshape(N, T, D)
        return self.out(y)

mha = MultiHeadAttention(512, 8)
print(tuple(mha(torch.randn(2, 16, 512), causal_mask(16)).shape))
# (2, 16, 512)

The output transpose restores (batch, time, head, head-width) before concatenation; omitting it mixes time and head coordinates. The module also supports cross-attention: x supplies queries and memory supplies keys and values. Inputs and value vectors must be finite. A (N,S) padding mask must be expanded to (N,1,1,S); it is not accepted directly. Combine it with a (T,S) causal mask when needed. Attention dropout is applied after softmax, so its realized rows need not sum to one during training.

Position information

Without positional information, dropout, or an order-dependent mask, self-attention is permutation-equivariant: permuting input rows permutes output rows. It is not invariant. A symmetric pooling operation then loses order, while position-indexed outputs still move with their tokens. Causal masking itself introduces order-dependent structure, so the unmasked equivariance statement cannot be applied unchanged to causal decoders.

The original solution adds fixed sinusoids of geometrically spaced frequencies:

\[PE_{(pos,2i)}=\sin\left(\frac{pos}{10000^{2i/d}}\right),\qquad PE_{(pos,2i+1)}=\cos\left(\frac{pos}{10000^{2i/d}}\right)\]

Sinusoids need no learned parameters and can be evaluated at positions beyond the training range; that does not guarantee useful long-context generalization. Learned absolute tables have a configured maximum size and require extension or another adaptation outside it. Positions inside an allocated but untrained region are also not guaranteed to work.

Rotary position embedding (RoPE) rotates pairs of query and key coordinates using position-dependent angles. With the same rotation frequencies, \((R_mq)^T(R_nk)=q^TR_{n-m}k\): for fixed content vectors, the explicit positional factor depends on the relative offset. The score also depends on content, which may already encode context and positions. RoPE does not rotate values in its usual form, and evaluating its rotations at longer positions does not guarantee extrapolation quality.

The sinusoidal function below returns one row per position. The paired sine/cosine form here requires an even feature width; we check it explicitly. At position zero every sine coordinate is 0 and every cosine coordinate is 1.

def sinusoidal_positions(length, width, device=None):
    if length < 1 or width < 2 or width % 2:
        raise ValueError("positive length and positive even width required")
    position = torch.arange(length, device=device, dtype=torch.float32)[:, None]
    frequency = 10000 ** (-torch.arange(0, width, 2, device=device, dtype=torch.float32) / width)
    pe = torch.empty(length, width, device=device)
    pe[:, 0::2] = torch.sin(position * frequency)
    pe[:, 1::2] = torch.cos(position * frequency)
    return pe

print(sinusoidal_positions(3, 4)[0].tolist())
# [0.0, 1.0, 0.0, 1.0]

Adding these rows to token embeddings keeps the tensor shape unchanged. The complete example uses this function rather than the random position vectors used in the later symmetry exercise.

The block: LayerNorm, residual, feed-forward

A residual connection adds a sublayer output to its input, so both must have width \(d_{model}\). LayerNorm here normalizes each token’s features, with learned scale and bias. The position-wise feed-forward network applies the same two-layer MLP independently at every position; it mixes features, while attention mixes information across positions. A hidden width of \(4d_{model}\) is a common configuration, not a requirement.

The original Transformer used post-norm, \(\operatorname{LN}(x+\operatorname{Sublayer}(x))\). This example uses pre-norm, \(x+\operatorname{Sublayer}(\operatorname{LN}(x))\), which provides a direct identity residual path and can make optimization easier. Warmup and initialization still depend on the architecture and training setup; neither post-norm nor pre-norm has a universal convergence guarantee. We also use GELU instead of the original ReLU and a simplified dropout placement. A final normalization is added after a pre-norm stack in the complete example below.

class TransformerBlock(nn.Module):
    def __init__(self, d_model, n_heads, ff_mult=4, dropout=0.1):
        super().__init__()
        self.ln1 = nn.LayerNorm(d_model)
        self.attn = MultiHeadAttention(d_model, n_heads, dropout)
        self.ln2 = nn.LayerNorm(d_model)
        self.ff = nn.Sequential(
            nn.Linear(d_model, ff_mult * d_model),
            nn.GELU(),
            nn.Linear(ff_mult * d_model, d_model),
            nn.Dropout(dropout))

    def forward(self, x, mask=None):
        x = x + self.attn(self.ln1(x), mask)      # pre-norm
        x = x + self.ff(self.ln2(x))
        return x

block = TransformerBlock(512, 8)
print(sum(p.numel() for p in block.parameters()))
# 3152384

For this width-512 block, the feed-forward layers contain 2,099,712 of the 3,152,384 parameters, about 66.6%. That is a parameter allocation, not a measurement of how much each component contributes to the model’s capabilities.

Encoder, decoder, or both

An encoder layer has self-attention and a feed-forward sublayer. A decoder layer has causal self-attention, cross-attention to the encoder output, and a feed-forward sublayer. Cross-attention queries come from the decoder; keys and values come from the source representation. It uses the source–target weighted-sum idea of recurrent seq2seq attention with a different scoring implementation.

Decoder-only models use causal attention over a prompt and continuation. Encoder–decoder models instead build a separate source representation and condition target generation on it. Both can generate text; the choice affects computation, conditioning, and the training objective.

A language-model training implementation is in Building a Transformer Language Model in PyTorch. The compact example here also makes the encoder–decoder connection explicit.

This one-layer encoder–decoder reuses the attention implementation above. Its decoder adds cross-attention between causal self-attention and the MLP. All feature widths are 8, and each two-head attention splits them into heads of width 4. The source and target token tables and output head are learned separately. Run the Python blocks in order.

class DecoderBlock(nn.Module):
    def __init__(self, width, heads):
        super().__init__()
        self.norms = nn.ModuleList([nn.LayerNorm(width) for _ in range(3)])
        self.self_attn = MultiHeadAttention(width, heads, dropout=0)
        self.cross_attn = MultiHeadAttention(width, heads, dropout=0)
        self.ff = nn.Sequential(nn.Linear(width, 4*width), nn.GELU(), nn.Linear(4*width, width))
    def forward(self, x, memory, target_mask, source_mask):
        x = x + self.self_attn(self.norms[0](x), target_mask)
        x = x + self.cross_attn(self.norms[1](x), source_mask, memory=memory)
        return x + self.ff(self.norms[2](x))

torch.manual_seed(0)
PAD, BOS, EOS = 0, 1, 2
src = torch.tensor([[3, 4, 5], [4, 5, PAD]])
tgt_in = torch.tensor([[BOS, 4, 3], [BOS, 5, EOS]])
target = torch.tensor([[4, 3, EOS], [5, EOS, PAD]])
source_emb, target_emb = nn.Embedding(6, 8, padding_idx=PAD), nn.Embedding(6, 8, padding_idx=PAD)
encoder = TransformerBlock(8, 2, dropout=0)
decoder = DecoderBlock(8, 2)
enc_norm, dec_norm, head = nn.LayerNorm(8), nn.LayerNorm(8), nn.Linear(8, 6)
modules = [source_emb, target_emb, encoder, decoder, enc_norm, dec_norm, head]
parameters = [p for module in modules for p in module.parameters()]
optimizer = torch.optim.SGD(parameters, lr=0.01)
source_mask = (src != PAD)[:, None, None, :]
target_mask = causal_mask(tgt_in.shape[1])[None, None] & (tgt_in != PAD)[:, None, None, :]
source_x = source_emb(src) * 8**0.5 + sinusoidal_positions(src.shape[1], 8)
target_x = target_emb(tgt_in) * 8**0.5 + sinusoidal_positions(tgt_in.shape[1], 8)
memory = enc_norm(encoder(source_x, source_mask))
logits = head(dec_norm(decoder(target_x, memory, target_mask, source_mask)))
loss = F.cross_entropy(logits.reshape(-1, 6), target.reshape(-1), ignore_index=PAD)
optimizer.zero_grad()
loss.backward()
print(tuple(memory.shape), tuple(logits.shape), int((target != PAD).sum()))
# (2, 3, 8) (2, 3, 6) 5
print(all(p.grad is not None and bool(torch.isfinite(p.grad).all()) for p in parameters))
# True
optimizer.step()

Target inputs are shifted right: BOS predicts the first output, and EOS is a supervised target. The two masks hide source padding and future target keys; the loss separately ignores the padded target. Padded query states can remain nonzero after residuals and positional addition, but valid queries never use padded keys. This checks a complete synthetic training step. During generation, feed BOS, select the next token from the last-position logits, append it, and repeat until EOS or a length cap. Recomputing the prefix works with this code; efficient serving would cache keys and values. No translation quality is established by this example.

Complexity and practical checks

At fixed batch and head count, the dense score and value products cost \(O(T^2d_{model})\), while projections and a fixed-ratio MLP cost \(O(Td_{model}^2)\). This explicit implementation stores attention matrices with \(O(T^2)\) entries per head. Doubling length quadruples that term, not necessarily total runtime or memory. FlashAttention avoids storing the full matrix and uses memory scaling linearly with length at fixed width while retaining exact dense attention (up to floating-point differences) and quadratic arithmetic. Sparse and linear alternatives change computation or connectivity and need task-specific evaluation.

• Test causality in evaluation mode with dropout disabled: changing future inputs should leave earlier outputs equal within a suitable numerical tolerance. Also check that earlier-output gradients to future inputs are zero.
• Check that attention weights sum to 1 across allowed keys before attention dropout; fully masked rows need an explicit policy.
• Check batch and head broadcasting with unequal batch, head, and sequence sizes so an accidental axis match does not hide a mask error.
• Choose the learning-rate schedule, including any warmup, for the optimizer, initialization, and normalization scheme; monitor gradient norms and early loss.
• Scale positional encodings and token embeddings consistently; the original multiplies embeddings by \(\sqrt{d_{\text{model}}}\) before adding positions.
• Excluding normalization parameters and biases from weight decay is a common optimizer grouping; document and validate the choice.

The normalization layers referred to here are compared in Normalization Layers Explained: BatchNorm, LayerNorm, GroupNorm, RMSNorm.

Exercises

1. Scaling and concentration. Measure the average peak softmax weight over 200 random query-key draws at \(d = 16, 64, 256, 1024\), with and without the \(1/\sqrt{d}\) division. Report the trend in each.

You should get: increasing average peak weight without scaling and peaks around 0.05 with scaling in this experiment.

import torch
torch.manual_seed(0)
for d in (16, 64, 256, 1024):
    m_raw = m_sc = 0.0
    for _ in range(200):
        q, k = torch.randn(1, d), torch.randn(200, d)
        raw = q @ k.T
        m_raw += raw.softmax(-1).max().item()
        m_sc += (raw / d**0.5).softmax(-1).max().item()
    print(f"d={d:5d}  raw {m_raw/200:.4f}   scaled {m_sc/200:.4f}")
# d=   16  raw 0.4836   scaled 0.0489
# d=   64  raw 0.7491   scaled 0.0524
# d=  256  raw 0.8814   scaled 0.0546
# d= 1024  raw 0.9472   scaled 0.0479

2. Mask with a finite negative. Fill masked logits with float("-inf") and with -1e4, on a partially masked row and on a fully masked one. Then try masked_fill(mask, -1e9) on a float16 tensor. Report what each produces.

You should get: a NaN row from one fill value, and a hard error from one dtype-and-value combination.

import torch
logits = torch.zeros(1, 4)
part = torch.tensor([[False, True, True, True]])   # True = masked out
full = torch.tensor([[True, True, True, True]])

for name, m in (("partial", part), ("full", full)):
    for fill in (float("-inf"), -1e4):
        print(name, fill, logits.masked_fill(m, fill).softmax(-1).tolist())
# partial -inf [[1.0, 0.0, 0.0, 0.0]]
# partial -10000.0 [[1.0, 0.0, 0.0, 0.0]]
# full -inf [[nan, nan, nan, nan]]
# full -10000.0 [[0.25, 0.25, 0.25, 0.25]]

try:
    logits.half().masked_fill(part, -1e9)
except RuntimeError as e:
    print("fp16:", e)
# fp16: value cannot be converted to type at::Half without overflow

3. Attention is permutation-equivariant. Run self-attention on a sequence and on a shuffled copy of it, with and without positional encodings. Check whether the outputs agree after unshuffling, using the tolerance in the code.

You should get: an identical output in one case, to floating-point precision.

import torch
torch.manual_seed(0)
T, d = 6, 16
x = torch.randn(T, d)
perm = torch.randperm(T)
Wq, Wk, Wv = (torch.randn(d, d) for _ in range(3))

def attn(z):
    q, k, v = z @ Wq, z @ Wk, z @ Wv
    return (q @ k.T / d**0.5).softmax(-1) @ v

inv = perm.argsort()
print("no pos equal:", torch.allclose(attn(x[perm])[inv], attn(x), atol=1e-5, rtol=1e-5))
pos = torch.randn(T, d)
print("with pos equal:", torch.allclose(attn(x[perm] + pos)[inv], attn(x + pos), atol=1e-5, rtol=1e-5))
# no pos equal: True
# with pos equal: False

References

• Vaswani et al. (2017). Attention Is All You Need.
• Shaw, Uszkoreit, and Vaswani (2018). Self-Attention with Relative Position Representations. NAACL.
• Devlin, Chang, Lee, and Toutanova (2019). BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding. NAACL.
• Radford et al. (2019). Language Models Are Unsupervised Multitask Learners. OpenAI technical report.
• Raffel et al. (2020). Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer. JMLR.
• Xiong et al. (2020). On Layer Normalization in the Transformer Architecture. ICML.
• Su et al. (2024). RoFormer: Enhanced Transformer with Rotary Position Embedding. Neurocomputing.
• Dao et al. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness.