The Three Shapes
Encoder, decoder, and both, and why decoder only won
A transformer is a stack of the block week 8 finished with, plus an embedding at the front and a linear layer at the back. Everything else you have heard about them is a variation on where the blocks are and what they are allowed to look at.
The three shapes
| Shape | Blocks see | Trained to | Examples |
|---|---|---|---|
| Encoder only | The whole sequence, both directions | Fill in masked positions | BERT, sentence embeddings, classifiers |
| Decoder only | Earlier positions only | Predict the next token | GPT, Llama, almost every current model |
| Encoder and decoder | Encoder sees all of the input; decoder sees the input and its own past | Produce one sequence from another | The original transformer, T5, translation |
Decoder only won, and it is worth knowing why
Not because it is more expressive. An encoder can see both directions, which is strictly more information for understanding a sentence. Decoder only won because next-token prediction is a training task you can run on any text at all, with no labels, and because one model trained that way turns out to do translation, summarising and question answering without being built for any of them. The architecture is not the innovation. The training objective is.
An encoder, assembled
from torch import nn
from torch.utils.data import DataLoader, TensorDataset
PAD, START = 0, 10
VOCAB = 11 # digits 1 to 9, plus pad and start
def make_copy(n, length=6, seed=0):
"""Source is a sequence of digits; target is the same sequence.
Trivial for a person and a real test for a model: it has to align
output position i with input position i, which is exactly what
cross-attention is for."""
g = torch.Generator().manual_seed(seed)
src = torch.randint(1, 10, (n, length), generator=g)
start = torch.full((n, 1), START)
tgt_in = torch.cat([start, src[:, :-1]], dim=1)
return src, tgt_in, src
src_tr, tin_tr, tout_tr = make_copy(6000, seed=0)
src_va, tin_va, tout_va = make_copy(1000, seed=1)
train_loader = DataLoader(TensorDataset(src_tr, tin_tr, tout_tr),
batch_size=64, shuffle=True)
class Block(nn.Module):
"""Pre-normalisation transformer block, as week 8 day 7 built it."""
def __init__(self, width, heads, expansion=4, dropout=0.1):
super().__init__()
self.norm1 = nn.LayerNorm(width)
self.attn = nn.MultiheadAttention(width, heads, dropout=dropout,
batch_first=True)
self.norm2 = nn.LayerNorm(width)
self.ff = nn.Sequential(
nn.Linear(width, width * expansion), nn.GELU(),
nn.Linear(width * expansion, width), nn.Dropout(dropout))
def forward(self, x, mask=None):
h = self.norm1(x)
a, _ = self.attn(h, h, h, attn_mask=mask, need_weights=False)
x = x + a
return x + self.ff(self.norm2(x))
class Encoder(nn.Module):
def __init__(self, vocab=VOCAB, width=64, heads=4, blocks=2,
max_len=16):
super().__init__()
self.tok = nn.Embedding(vocab, width)
self.pos = nn.Embedding(max_len, width)
self.blocks = nn.ModuleList([Block(width, heads)
for _ in range(blocks)])
self.norm = nn.LayerNorm(width)
def forward(self, x):
h = self.tok(x) + self.pos(torch.arange(x.shape[1]))
for block in self.blocks:
h = block(h)
return self.norm(h)
torch.manual_seed(0)
enc = Encoder()
x = torch.randint(1, 10, (2, 6))
print('input ', tuple(x.shape))
print('output ', tuple(enc(x).shape), ' one vector per position')
print('parameters %d' % sum(p.numel() for p in enc.parameters()))
print('\nwhere the parameters are:')
for name, child in enc.named_children():
print(' %-8s %d' % (name, sum(p.numel() for p in child.parameters())))
output (2, 6, 64) one vector per position
parameters 101824
where the parameters are:
tok 704
pos 1024
blocks 99968
norm 128
Most of the weight is in the blocks, and inside a block most of it is in the feed forward network rather than the attention. That surprises people. Attention decides what to read; the feed forward layers are where most of the actual capacity lives, and current research on where a model stores facts keeps pointing at them.