Full Transformer Architecture — Encoder + Decoder

Why this session exists

Attention routes information, a position signal distinguishes locations, and a feed-forward sublayer transforms each token's state. A complete Transformer also needs explicit choices about residual paths, normalization, masks, output heads and training targets. Those choices determine what the model computes; none is adequately explained by saying that every modern model is the same block.

Here we assemble a small pre-LN decoder, train it on an inline character fixture, and separately exercise encoder-decoder cross-attention. We retain the original architectural breadth without pretending that a few CPU updates reproduce GPT pretraining.

Outcomes: draw both normalization orderings; count parameters; explain attention/FFN roles; choose among encoder, decoder and encoder-decoder interfaces; build and generate from TinyGPT; and trace a source/target pair through cross-attention.

Prerequisites

Sessions 106–110 cover tokenization, attention shapes, multi-head projections and position. Session 100 supplies modules/autograd; residual connections also appeared in the CNN lesson. Use (B,T,d) throughout and distinguish source length S from target length T.

(a) Intuition · 5 min

A tower repeats a floor plan, not the furniture. Transformer layers commonly repeat an architecture while retaining different parameters. “Residual stream as a shared bus” is useful: a branch reads a normalized view and adds a change to the stream. Addition can cancel or overwrite information algebraically; the metaphor does not guarantee permanent preservation of every earlier feature.

The five pieces of a Transformer block

In this pre-LN decoder block:

The functional pieces are attention, two residual additions, two normalizations and a nonlinear FFN. Counting them as “five pieces” is a teaching grouping, not a theorem that every Transformer has exactly five operations. Dropout, gates, cross-attention, parallel branches, expert routing and alternative norms change the design.

A residual sublayer y=x+f(x) has Jacobian I+J_f. The identity term provides a direct derivative path, but the total derivative can still cancel, grow or shrink. There is no universal 12-layer failure threshold without residuals and no guarantee of unattenuated gradients merely from adding them.

A short history

The 2017 Transformer was an encoder-decoder for translation, with post-LN and sinusoidal position [1]. GPT (2018) used a generative language-model pretraining objective; BERT (2018 preprint, 2019 proceedings) used bidirectional MLM plus next-sentence prediction. T5's 2019 preprint/2020 JMLR paper unified tasks into text-to-text and used relative bias. GPT-3 (2020) studied few-shot prompting at scale. LLaMA (2023) documented pre-RMSNorm, SwiGLU and RoPE [3]. These dated examples describe public designs, not hidden GPT-4/Claude internals or a universal architecture timeline.

The Transformer block

The same shape permits residual addition. Attention mixes across legal keys; the FFN shares weights across positions but processes each position separately. Attention is linear in V holding Q/K weights fixed; as a function of its input it is nonlinear because the attention weights depend on the input. Therefore a stack of attention layers does not collapse to a single linear map. The FFN supplies a separate, useful content transformation rather than repairing a supposedly wholly linear attention network.

The full architecture — decoder-only (GPT)

Token IDs index an embedding matrix E with shape (vocab,d). Add a learned position table in this fixture, pass through independent blocks, apply a final norm, then compute logits $hE^T$ with tied input/output weights. The result is (B,T,vocab), not probabilities until softmax. Cross-entropy consumes raw logits.

Teacher forcing pairs input tokens [BOS,a,b] with labels [a,b,EOS]. At index 1 the model sees BOS and a, and predicts b. Training all positions in parallel does not allow attention to future input positions. Feeding unshifted labels makes copying possible and teaches the wrong task. Autoregressive generation is sequential: the chosen token changes the next input.

Encoder-only vs decoder-only vs encoder-decoder

• Encoder: bidirectional self-attention over nonpadding source tokens. A task head may classify a sequence, tag tokens, predict masked tokens or pool an embedding. MLM is BERT's pretraining choice, not the definition of every encoder.
• Decoder-only: causal self-attention. A language-model head supplies next-token logits; hidden states can also support one-pass classification or embeddings. Generation is optional, not mandatory for every decoder application.
• Encoder-decoder: a bidirectional source encoder plus a causal target decoder. Cross-attention uses target states as Q and encoded source states as K/V. Its score matrix is (B,h,T,S), rectangular when T differs from S. It is a useful transduction design, not guaranteed to beat a decoder on translation.
• Prefix LM: a single sequence can have full visibility within a prefix and causal visibility in a suffix. This is a mask/objective choice; it is not synonymous with all PaLM or UL2 architectures.

For a three-token source and two-token target, source attention is 3×3; target self-attention is lower-triangular 2×2; target-to-source cross-attention is 2×3. Padding masks apply to legal keys in all three. Decoder label shifting remains necessary even with a correct target mask.

Anatomy: what each component adds

An ordinary bias-free MHA has 4d² weights. A two-matrix FFN with inner width f has 2df weights. At f=4d, that is 8d², twice attention, not four to eight times. With the fixture's biased FFN add f+d biases. Each affine LayerNorm has 2d parameters, so two norms add 4d. At d=32, f=128 the block total is 4096+8192+160+128=12,576. The program checks this exact count.

The 4d ratio is empirical parameter allocation, not a derived optimum. A gated SwiGLU FFN uses three matrices and approximately 3df weights; f near 8d/3 matches the 8d² budget before rounding. GELU, ReLU and SwiGLU differ mathematically and empirically; no fixed universal quality bump follows from replacing one with another. Sparse MoE often replaces FFNs with routed experts, increasing total capacity while activating a subset. MoE is not constrained by a theorem that it can never alter attention.

LayerNorm subtracts the feature mean and divides by the square root of biased feature variance plus epsilon, then optionally scales/shifts. RMSNorm divides by $\sqrt{\frac{1}{d}\sum_i x_i^2+\epsilon}$, with epsilon inside the square root, usually with a learned scale and no centering [6]. It need not produce zero mean. The residual stream itself is not bounded merely because branch inputs are normalized. Dropout is a tunable training regularizer, not something that must automatically be restored in every small-data fine-tune.

Pre-LN versus post-LN

Post-LN is $y=\operatorname{LN}(x+f(x))$, so its Jacobian is $J_{LN}(I+J_f)$. Pre-LN is $y=x+f(\operatorname{LN}(x))$, with Jacobian $I+J_fJ_{LN}$. This is the direct-path distinction. Xiong et al. analyze large expected gradients near the output at initialization in post-LN and the role of warmup [2]. They do not prove that nobody can train deep post-LN models, or that pre-LN always wins.

The original gradient experiment used squared output norm after a final LayerNorm. That objective is almost constant when the normalized variance is near one, making a small gradient unsurprising even for one layer. A better probe uses a nonconstant linear readout and matched weights/input. The CPU fixture compares both orderings this way; its gradient norms demonstrate computation, not a universal training ranking.

Choosing a model and reading an architecture paper

Start from the output contract, training data and serving budget, then compare matched candidates. A decoder can return a last-token vector in one forward pass; a bidirectional encoder does not automatically yield a useful sentence embedding without the appropriate objective/pooling. For a fixed label set at high volume, test an encoder baseline. For variable-length outputs compare decoder and encoder-decoder checkpoints with relevant training. Include quality, calibration, input lengths, batching and hardware, not unsourced 10–100× latency claims.

When reading a diagram, record: input/target IDs, position method, visibility, Q/K/V sources, normalization ordering, FFN activation/width, residual placement, tied weights, padding policy and loss positions. That checklist converts a figure into an implementable contract.

(c) Hands-on · 25 min

The first program trains a small tied-weight decoder on an inline character sequence and generates a short sample. No Shakespeare, pretrained weights or GPU is used. It also runs an encoder-decoder example on synthetic IDs, checking future-target and source-padding isolation. Its MLM extension demonstrates selected-position loss, not a pretrained BERT or a measured retrieval embedding.

import copy, math
import torch
from torch import nn
from torch.nn import functional as F
torch.set_num_threads(1); torch.manual_seed(111)
 
class Attention(nn.Module):
    def __init__(self,d,h):
        super().__init__(); assert d%h==0
        self.h=h; self.qkv=nn.Linear(d,3*d,bias=False); self.out=nn.Linear(d,d,bias=False)
    def forward(self,x,causal=True):
        b,t,d=x.shape
        q,k,v=[a.reshape(b,t,self.h,d//self.h).transpose(1,2) for a in self.qkv(x).chunk(3,-1)]
        y=F.scaled_dot_product_attention(q,k,v,is_causal=causal,dropout_p=0.0)
        return self.out(y.transpose(1,2).reshape(b,t,d))
class Block(nn.Module):
    def __init__(self,d=32,h=4,f=128):
        super().__init__();self.n1=nn.LayerNorm(d);self.n2=nn.LayerNorm(d)
        self.attn=Attention(d,h);self.ffn=nn.Sequential(nn.Linear(d,f),nn.GELU(),nn.Linear(f,d))
    def forward(self,x,causal=True):
        x=x+self.attn(self.n1(x),causal)
        return x+self.ffn(self.n2(x))
class TinyGPT(nn.Module):
    def __init__(self,vocab,length=16,d=32,layers=2):
        super().__init__();self.length=length
        self.emb=nn.Embedding(vocab,d);self.pos=nn.Embedding(length,d)
        self.blocks=nn.ModuleList([Block(d,4,4*d) for _ in range(layers)])
        self.norm=nn.LayerNorm(d);self.head=nn.Linear(d,vocab,bias=False)
        self.head.weight=self.emb.weight
    def forward(self,ids,causal=True):
        if ids.ndim!=2 or not 0<ids.shape[1]<=self.length:raise ValueError('invalid context length')
        x=self.emb(ids)+self.pos(torch.arange(ids.shape[1]))
        for block in self.blocks:x=block(x,causal)
        return self.head(self.norm(x))
    @torch.no_grad()
    def generate(self,ids,count,temperature=.8):
        if not math.isfinite(temperature) or temperature<=0:raise ValueError('positive temperature required')
        was=self.training;self.eval()
        try:
            for _ in range(count):
                logits=self(ids[:,-self.length:])[:,-1]
                ids=torch.cat((ids,torch.multinomial((logits/temperature).softmax(-1),1)),1)
            return ids
        finally:self.train(was)
 
text='red blue red green. '*12
chars=sorted(set(text)); encode={c:i for i,c in enumerate(chars)}
data=torch.tensor([encode[c] for c in text]);n=(len(data)-1)//16
assert n>0
x=data[:n*16].reshape(n,16);y=data[1:n*16+1].reshape(n,16)
model=TinyGPT(len(chars));assert model.head.weight is model.emb.weight
assert sum(p.numel() for p in model.blocks[0].parameters())==12576
opt=torch.optim.AdamW(model.parameters(),lr=.003)
losses=[]
for step in range(35):
    opt.zero_grad();logits=model(x);loss=F.cross_entropy(logits.flatten(0,1),y.flatten())
    assert torch.isfinite(loss);loss.backward()
    assert all(p.grad is None or torch.isfinite(p.grad).all() for p in model.parameters())
    nn.utils.clip_grad_norm_(model.parameters(),1.0);opt.step();losses.append(float(loss.detach()))
assert losses[-1]<losses[0]
model.eval();a=x[:1].clone();b=a.clone();b[:,8:]=(b[:,8:]+1)%len(chars)
assert torch.allclose(model(a)[:,:8],model(b)[:,:8],atol=1e-6)
assert not torch.allclose(model(a,False)[:,:8],model(b,False)[:,:8])
sample=model.generate(x[:1,:1],24)
print('synthetic training loss',losses[0],losses[-1]);print('fixture sample',''.join(chars[i] for i in sample[0]))
# MLM shape/selected-label extension (not pretrained BERT).
maskid=len(chars);mlm=TinyGPT(len(chars)+1)
selected=torch.zeros_like(x,dtype=torch.bool);selected[:,3]=True
corrupt=x.clone();corrupt[selected]=maskid
mloss=F.cross_entropy(mlm(corrupt,causal=False)[selected],x[selected]);mloss.backward()
assert selected.sum()==len(x) and torch.isfinite(mloss)
 
# Full source/target interface, no external model or dataset.
class Seq2Seq(nn.Module):
    def __init__(self):
        super().__init__();self.emb=nn.Embedding(20,16);self.pos=nn.Embedding(16,16)
        self.encoder=nn.TransformerEncoderLayer(16,2,32,dropout=0,batch_first=True,norm_first=True)
        self.decoder=nn.TransformerDecoderLayer(16,2,32,dropout=0,batch_first=True,norm_first=True)
        self.head=nn.Linear(16,20)
    def forward(self,src,tgt,pad):
        source=self.emb(src)+self.pos(torch.arange(src.shape[1]))
        target=self.emb(tgt)+self.pos(torch.arange(tgt.shape[1]))
        memory=self.encoder(source,src_key_padding_mask=pad)
        blocked=torch.ones(tgt.shape[1],tgt.shape[1],dtype=torch.bool).triu(1)
        out=self.decoder(target,memory,tgt_mask=blocked,memory_key_padding_mask=pad)
        return self.head(out)
s2s=Seq2Seq().eval();src=torch.tensor([[2,3,0]]);tgt=torch.tensor([[1,5]])
pad=torch.tensor([[False,False,True]])
z=s2s(src,tgt,pad);assert z.shape==(1,2,20)
other=tgt.clone();other[0,1]=7
assert torch.allclose(z[:,:1],s2s(src,other,pad)[:,:1],atol=1e-6)
other_src=src.clone();other_src[0,2]=9
assert torch.allclose(z,s2s(other_src,tgt,pad),atol=1e-6)
other_src[0,0]=8
assert not torch.allclose(z,s2s(other_src,tgt,pad))
 
# Matched pre/post norm probe: nonconstant linear readout, identical weights.
class Probe(nn.Module):
    def __init__(self,pre):
        super().__init__();self.pre=pre;self.n=nn.LayerNorm(8);self.f=nn.Linear(8,8)
    def forward(self,x):return x+self.f(self.n(x)) if self.pre else self.n(x+self.f(x))
pre=nn.Sequential(*[Probe(True) for _ in range(6)]).double();post=copy.deepcopy(pre)
for p in post:p.pre=False
inp=torch.randn(2,8,dtype=torch.float64);readout=torch.randn(2,8,dtype=torch.float64)
for net in (pre,post):
    u=inp.clone().requires_grad_();(net(u)*readout).sum().backward()
    assert torch.isfinite(u.grad).all();print('gradient probe pre',net[0].pre,float(u.grad.norm()))
print('decoder training, MLM labels, cross-attention masks and norm probes passed')

Anatomy of the script

A combined Linear(d,3d) is mathematically equivalent to three projections; whether it is faster depends on kernels and shapes. SDPA with is_causal=True supplies the square causal mask here and dropout_p=0.0 makes the fixture's attention deterministic in evaluation. This does not prove that a CUDA Flash kernel ran.

Weight tying shares the same parameter object. It saves Vd weights versus a separate head, not a fixed percentage of the whole model. A tied random embedding can yield highly nonuniform initial logits, so initial loss need not equal log(V). Cropping generated history resets positions in this simple learned-table recipe; this is a moving-window model, not a cache-equivalent continuation beyond the trained table.

The encoder-decoder uses independent source/target lengths and masks. PyTorch MultiheadAttention boolean masks use True for blocked positions, whereas functional SDPA boolean masks use True for allowed positions. The example uses the appropriate convention at each API boundary.

Practice: swap the objective, not just the mask

For an encoder MLM, reserve a mask ID, select positions, corrupt them, and compute loss against the original tokens only at selected positions. BERT's historical 15% selection and 80/10/10 replacement policy are developed in session 112. Removing a causal mask while retaining next-token targets can reveal the target directly; the code's future-token perturbation test detects this leakage.

For a fair head/depth/width experiment, keep data splits and token budget fixed, record parameter counts and repeat seeds. A better loss on a repeated synthetic string is not a neural scaling-law result. Inspect empty datasets and contexts longer than the position table explicitly instead of letting an undefined last batch counter fail later.

(d) Production reality · 15 min

LLaMA's documented norm, activation and position choices supply a concrete public variant [3]; fabricated ablation percentages and private GPT-4 parameter rumors do not. Some interpretability studies analyze FFNs as key-value memories, but facts are not known to live only there and fine-tuning changes are not universally concentrated in FFNs.

If loss diverges, inspect finite inputs/labels, target alignment, legal attention rows, actual loss reduction, learning rate/schedule, optimizer state, gradient norms and dtype overflow. Clipping is a mitigation, not a proof that the model or labels are correct. Reduced-precision failures require tracing actual tensors rather than assuming bf16 always fails or never fails.

Serving memory includes weights, K/V cache, activations and workspaces. A KV cache may dominate at long contexts or large batch, but weights can dominate for small batch/short contexts. Source encoding can be reused within one encoder-decoder generation; cross-attention K/V can also be cached. Compare complete request cost, not only the number of stacks.

Where this shows up in the rest of the plan

Session 112 compares objectives and model-loader interfaces; 113 handles sampling; 115 handles execution/cache efficiency; 117 handles retrieval evidence. LoRA may target attention, FFN or other linear maps. Vision Transformers reuse token-mixing blocks on patch sequences; diffusion Transformers change conditioning/objective. A common primitive is not evidence of identical whole-model behavior.

(e) Recall + stretch · 10 min

1. Block order? Pre-LN: norm, attention, add; norm, FFN, add. Post-LN: attention, add, norm; FFN, add, norm. A cross-attention decoder has another sublayer.
2. FFN role? Shared positionwise nonlinear content transformation, complementary to attention's input-dependent routing. Attention itself is not globally linear.
3. Why 4d? Historical empirical allocation: 8d² FFN versus 4d² attention. Three-matrix gating changes the equal-budget width to about 8d/3.
4. Three families? Full source visibility; causal prefix visibility; or full source plus causal target and cross-attention. Objectives and output heads must also be specified.
5. Pre/post-LN? The normalization Jacobian lies on the post-LN shortcut, but not on the explicit identity term of pre-LN. No depth cutoff or convergence guarantee follows.
6. Weight tying? One V×d parameter matrix is used by embedding and output projection; saves Vd versus an untied head. At V=128000 and d=4000, that is 512,000,000 weights; binary 128K/4K conventions give a different exact number.
7. Relative position? RoPE rotates projected Q/K; T5 adds a bucketed bias to scores. RoPE is not “bias-added.” Additive absolute position instead enters embeddings.
8. Divergence checks? Start with finite data and correct labels/masks, then optimizer/loss scale and gradients. Verify normalization placement and precision rather than blindly lowering a rate or adding clipping.
9. RMSNorm? Divide by the square root of the mean squared features plus epsilon inside the root, without centering; this differs from mean/variance LayerNorm. Any speed/quality comparison needs a named implementation and workload.
10. What changed since 2017? Public models vary norms, position, gates, attention sharing, routing, scale, data and post-training. Private GPT-4 design/size and optimizer details cannot be reconstructed from a generic block diagram.

Explain-out-loud test

Draw a two-token target attending to a three-token source; count a d=32/f=128 block; derive both residual Jacobians; and explain why an unmasked next-token training example can look successful while leaking its labels.

What comes next

Model families and objectives separates architecture, pretraining and output interface. Learning-plan hub.

Primary references

1. Attention Is All You Need, v7, §3.
2. Xiong et al., On Layer Normalization in the Transformer Architecture, ICML 2020.
3. LLaMA, 2023, §2.2.
4. PyTorch 2.9 SDPA contract.
5. OpenAI GPT-2 public implementation, dated source snapshot.
6. PyTorch 2.9 RMSNorm definition.