Attention

Week 8 of 18 · Sequences · 7 days

Week 08 · Sequences

Attention

Week 08 · Day 1 of 7

The Bottleneck

Why one fixed-size state is not enough, and the idea that replaced it

A recurrent layer compresses everything it has read into one fixed-size vector. For a short sequence that is fine. For a long one it is a bottleneck, and the fix is to stop compressing and start looking back.

The bottleneck, measured

import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset

def make_sequences(n, length=8, seed=0):
    """Week 7's task: is the digit at position 0 repeated later?
    Balanced by construction, and unanswerable from a bag of counts."""
    g = torch.Generator().manual_seed(seed)
    first = torch.randint(1, 10, (n, 1), generator=g)
    rest = torch.randint(1, 9, (n, length - 1), generator=g)
    rest = rest + (rest >= first).long()
    y = torch.randint(0, 2, (n,), generator=g)
    where = torch.randint(0, length - 1, (n, 1), generator=g)
    idx = torch.arange(length - 1).unsqueeze(0)
    plant = (y.unsqueeze(1) == 1) & (idx == where)
    rest = torch.where(plant, first.expand_as(rest), rest)
    return torch.cat([first, rest], dim=1), y

X_tr, y_tr = make_sequences(8000, seed=0)
X_va, y_va = make_sequences(2000, seed=1)
train_loader = DataLoader(TensorDataset(X_tr, y_tr), batch_size=64,
                          shuffle=True)
val_loader = DataLoader(TensorDataset(X_va, y_va), batch_size=512)
def fit(model, epochs=10, lr=3e-3, loader=None, val=None):
    torch.manual_seed(0)
    opt = torch.optim.AdamW(model.parameters(), lr=lr)
    loss_fn = nn.CrossEntropyLoss()
    for _ in range(epochs):
        model.train()
        for xb, yb in (loader or train_loader):
            opt.zero_grad()
            loss_fn(model(xb), yb).backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            opt.step()
    model.eval()
    right = seen = 0
    with torch.no_grad():
        for xb, yb in (val or val_loader):
            right += (model(xb).argmax(1) == yb).sum().item()
            seen += yb.numel()
    return right / seen
class LastState(nn.Module):
    def __init__(self, hidden):
        super().__init__()
        self.emb = nn.Embedding(10, 32)
        self.rnn = nn.LSTM(32, hidden, batch_first=True)
        self.head = nn.Linear(hidden, 2)

    def forward(self, x):
        out, _ = self.rnn(self.emb(x))
        return self.head(out[:, -1, :])

print('%14s %12s' % ('hidden size', 'accuracy'))
for hidden in [2, 4, 8, 64]:
    torch.manual_seed(0)
    print('%14d %12.4f' % (hidden, fit(LastState(hidden))))

   hidden size     accuracy
             2       0.5000
             4       0.5025
             8       0.5050
            64       0.9860

Everything the model knows about eight digits has to fit through that vector. Squeeze it and performance falls, because the state is genuinely carrying information rather than acting as a formality. Now imagine a sentence of two hundred words compressed into the same place, which is what machine translation was doing in 2014.

The idea

Attention: Instead of compressing a sequence into one vector, keep every position and let the model decide, per query, how much of each position to read. The output is a weighted average of the positions, and the weights are computed from the content rather than fixed in advance.

import torch

# Four positions, each with a value we might want to read.
values = torch.tensor([[1.0, 0.0],
                       [0.0, 1.0],
                       [2.0, 2.0],
                       [-1.0, 0.5]])

for name, weights in [('all of position 2', torch.tensor([0., 0., 1., 0.])),
                      ('an even average', torch.tensor([.25, .25, .25, .25])),
                      ('mostly 0 and 1', torch.tensor([.45, .45, .05, .05]))]:
    print('%-20s -> %s' % (name, (weights @ values).tolist()))
print('\nattention is this, with the weights computed rather than chosen.')

all of position 2    -> [2.0, 2.0]
an even average      -> [0.5, 0.875]
mostly 0 and 1       -> [0.5, 0.574999988079071]

attention is this, with the weights computed rather than chosen.

Day 1 takeaway

A single fixed-size state is a bottleneck, and shrinking it measurably costs accuracy. Attention removes the bottleneck by keeping every position and computing, on demand, a weighted average over them.

Week 08 · Day 2 of 7

Scaled Dot-Product Attention

Queries, keys and values, and why the square root is there

By Pawan Rama Mali645 words

The weights come from asking every position how well it matches what you are looking for. That is a dot product, and the rest is bookkeeping.

Queries, keys and values

Query, key and value: Three projections of the same input. The query is what this position is looking for, the key is what each position offers, and the value is what it passes on if selected. Query against key gives a score; the scores become weights; the weights average the values.

import torch

def softmax(z, dim=-1):
    z = z - z.max(dim=dim, keepdim=True).values
    e = z.exp()
    return e / e.sum(dim=dim, keepdim=True)

def attention(Q, K, V):
    d = K.shape[-1]
    scores = Q @ K.transpose(-2, -1) / d ** 0.5
    weights = softmax(scores)
    return weights @ V, weights

torch.manual_seed(0)
x = torch.randn(5, 8)                    # 5 positions, width 8
Wq, Wk, Wv = (torch.randn(8, 8) * 0.4 for _ in range(3))
out, w = attention(x @ Wq, x @ Wk, x @ Wv)

print('input  ', tuple(x.shape))
print('output ', tuple(out.shape), ' one vector per position')
print('weights', tuple(w.shape), ' one row per query')
print('\nevery row sums to 1:', bool(torch.allclose(w.sum(-1),
                                                    torch.ones(5))))
print('\nweights (row = the position doing the looking):')
for i, row in enumerate(w):
    print('  %d  %s' % (i, ' '.join('%.3f' % v for v in row.tolist())))

input   (5, 8)
output  (5, 8)  one vector per position
weights (5, 5)  one row per query

every row sums to 1: True

weights (row = the position doing the looking):
  0  0.623 0.181 0.063 0.108 0.026
  1  0.277 0.304 0.105 0.175 0.138
  2  0.030 0.058 0.091 0.060 0.760
  3  0.061 0.174 0.110 0.162 0.493
  4  0.045 0.039 0.425 0.070 0.420

Why the square root

import torch

torch.manual_seed(0)
print('%8s %16s %16s %16s'
      % ('width', 'score std', 'max weight', 'max weight scaled'))
for d in [8, 64, 512]:
    q = torch.randn(1, d)
    k = torch.randn(20, d)
    raw = q @ k.T
    scaled = raw / d ** 0.5
    print('%8d %16.3f %16.4f %16.4f'
          % (d, raw.std().item(), raw.softmax(-1).max().item(),
             scaled.softmax(-1).max().item()))

   width        score std       max weight max weight scaled
       8            3.579           0.7797           0.3327
      64            5.087           0.9167           0.1394
     512           27.192           0.9995           0.2701

Without the scaling, attention stops attending

A dot product of two random vectors of width d has a standard deviation proportional to the square root of d. At width 512 the scores are large enough that the softmax saturates: one position takes almost all the weight and the rest get nothing, which also means the gradient reaching them is nearly zero. Dividing by the square root of d holds the scores at a sensible size no matter how wide the model gets, and that single division is what the word scaled in "scaled dot product attention" refers to.

It is permutation equivariant, which is a problem

import torch

def attention(Q, K, V):
    return (Q @ K.transpose(-2, -1) / K.shape[-1] ** 0.5).softmax(-1) @ V

torch.manual_seed(0)
x = torch.randn(4, 6)
Wq, Wk, Wv = (torch.randn(6, 6) * 0.4 for _ in range(3))

a = attention(x @ Wq, x @ Wk, x @ Wv)
order = [2, 0, 3, 1]
b = attention(x[order] @ Wq, x[order] @ Wk, x[order] @ Wv)

print('shuffle the positions and the outputs are the same rows,')
print('in the shuffled order:', bool(torch.allclose(a[order], b, atol=1e-6)))
print('\nso attention alone cannot tell 12345 from 54321.')
print('day 5 puts the position back.')

shuffle the positions and the outputs are the same rows,
in the shuffled order: True

so attention alone cannot tell 12345 from 54321.
day 5 puts the position back.

Day 2 takeaway

Attention scores every position against a query with a dot product, divides by the square root of the width to stop the softmax saturating, and returns a weighted average of the values. It treats the sequence as a set, so position has to be supplied separately.

Week 08 · Day 3 of 7

Multiple Heads

Several relationships at once, for the same price

By Pawan Rama Mali546 words

One set of weights can only express one kind of relationship at a time. Multi-head attention runs several in parallel on different projections and concatenates the results.

Several heads at once

import torch
from torch import nn

class MultiHeadAttention(nn.Module):
    def __init__(self, width, heads):
        super().__init__()
        assert width % heads == 0
        self.heads = heads
        self.head_dim = width // heads
        self.qkv = nn.Linear(width, width * 3)
        self.out = nn.Linear(width, width)

    def forward(self, x, mask=None):
        B, T, C = x.shape
        q, k, v = self.qkv(x).split(C, dim=2)
        # split the width into heads and move heads next to the batch
        def heads(t):
            return t.view(B, T, self.heads, self.head_dim).transpose(1, 2)
        q, k, v = heads(q), heads(k), heads(v)

        scores = q @ k.transpose(-2, -1) / self.head_dim ** 0.5
        if mask is not None:
            scores = scores.masked_fill(mask == 0, float('-inf'))
        weights = scores.softmax(-1)
        out = weights @ v                       # (B, heads, T, head_dim)
        out = out.transpose(1, 2).reshape(B, T, C)
        return self.out(out), weights

torch.manual_seed(0)
mha = MultiHeadAttention(width=32, heads=4)
x = torch.randn(2, 6, 32)
out, w = mha(x)
print('input   ', tuple(x.shape))
print('output  ', tuple(out.shape), ' same shape as the input')
print('weights ', tuple(w.shape), ' = (batch, heads, query, key)')
print('parameters %d' % sum(p.numel() for p in mha.parameters()))

input    (2, 6, 32)
output   (2, 6, 32)  same shape as the input
weights  (2, 4, 6, 6)  = (batch, heads, query, key)
parameters 4224

The heads are free

Splitting a width of 32 into four heads of 8 costs exactly the same parameters and the same arithmetic as one head of 32. You are not buying capacity, you are buying the ability to attend to several things at once: one head can track which positions share a digit while another tracks position, and averaging them into a single set of weights would have forced a compromise.

The heads really do differ

import torch
from torch import nn

torch.manual_seed(0)
width, heads, T = 32, 4, 6
qkv = nn.Linear(width, width * 3)
x = torch.randn(1, T, width)
B, _, C = x.shape
q, k, v = qkv(x).split(C, dim=2)
hd = width // heads
shape = lambda t: t.view(B, T, heads, hd).transpose(1, 2)
q, k = shape(q), shape(k)
w = (q @ k.transpose(-2, -1) / hd ** 0.5).softmax(-1)

for h in range(heads):
    row = w[0, h, 0]
    print('head %d, what position 0 looks at: %s'
          % (h, ' '.join('%.2f' % val for val in row.tolist())))
print('\nuntrained, so the pattern is arbitrary. The point is that the')
print('four rows differ, which one head could not produce.')

head 0, what position 0 looks at: 0.08 0.09 0.25 0.20 0.24 0.14
head 1, what position 0 looks at: 0.21 0.16 0.16 0.14 0.15 0.18
head 2, what position 0 looks at: 0.15 0.24 0.08 0.15 0.16 0.22
head 3, what position 0 looks at: 0.22 0.18 0.17 0.14 0.16 0.13

untrained, so the pattern is arbitrary. The point is that the
four rows differ, which one head could not produce.

Day 3 takeaway

Multi-head attention splits the width into groups and runs attention independently in each, at no extra cost in parameters or computation. It buys the ability to attend to several different relationships at once rather than averaging them into one.

Week 08 · Day 4 of 7

Masks

Self against cross attention, the causal mask, and the padding mask

By Pawan Rama Mali790 words

Where the queries come from decides what kind of attention you have, and which positions you allow to be seen decides whether you can use it for generation.

Self-attention and cross-attention

	Queries from	Keys and values from	Used in
Self-attention	The sequence itself	The same sequence	Every transformer block
Cross-attention	The output sequence	The input sequence	Translation, captioning, any encoder to decoder model
Causal self-attention	The sequence itself	Earlier positions only	Language models

import torch
from torch import nn

torch.manual_seed(0)

def attend(q, k, v, mask=None):
    scores = q @ k.transpose(-2, -1) / k.shape[-1] ** 0.5
    if mask is not None:
        scores = scores.masked_fill(mask == 0, float('-inf'))
    return scores.softmax(-1) @ v

source = torch.randn(1, 7, 16)      # e.g. a sentence in French
target = torch.randn(1, 4, 16)      # what we have produced in English

print('self-attention on the target  ',
      tuple(attend(target, target, target).shape))
print('cross-attention to the source ',
      tuple(attend(target, source, source).shape))
print('\none output per query position either way. Cross-attention lets')
print('each output position read the whole input.')

self-attention on the target   (1, 4, 16)
cross-attention to the source  (1, 4, 16)

one output per query position either way. Cross-attention lets
each output position read the whole input.

The causal mask

import torch

T = 6
mask = torch.tril(torch.ones(T, T))
print('lower triangular mask:')
print(mask.int())

torch.manual_seed(0)
scores = torch.randn(T, T)
masked = scores.masked_fill(mask == 0, float('-inf'))
weights = masked.softmax(-1)
print('\nattention weights after masking:')
for i, row in enumerate(weights):
    print('  pos %d  %s' % (i, ' '.join('%.2f' % v for v in row.tolist())))
print('\nposition 0 sees only itself; position 5 sees everything.')
print('no position can see its own future, which is what makes')
print('next-token training possible in one parallel pass.')

lower triangular mask:
tensor([[1, 0, 0, 0, 0, 0],
        [1, 1, 0, 0, 0, 0],
        [1, 1, 1, 0, 0, 0],
        [1, 1, 1, 1, 0, 0],
        [1, 1, 1, 1, 1, 0],
        [1, 1, 1, 1, 1, 1]], dtype=torch.int32)

attention weights after masking:
  pos 0  1.00 0.00 0.00 0.00 0.00 0.00
  pos 1  0.86 0.14 0.00 0.00 0.00 0.00
  pos 2  0.15 0.45 0.40 0.00 0.00 0.00
  pos 3  0.27 0.33 0.23 0.17 0.00 0.00
  pos 4  0.13 0.08 0.29 0.22 0.28 0.00
  pos 5  0.04 0.03 0.54 0.18 0.05 0.16

position 0 sees only itself; position 5 sees everything.
no position can see its own future, which is what makes
next-token training possible in one parallel pass.

This is why language models train so efficiently

Without the mask you would have to run the model once per position to predict each next token without cheating. With it, a single forward pass over a sequence of length 1,024 produces 1,024 predictions, each conditioned only on what came before it. Every position is a training example and they are all computed at once. That efficiency, not the architecture itself, is most of why large language models are possible.

The padding mask, which is a different thing

import torch

# Two sequences of different lengths, padded with zeros.
tokens = torch.tensor([[3, 7, 2, 0, 0],
                       [4, 1, 8, 6, 5]])
pad_mask = (tokens != 0).unsqueeze(1)          # (batch, 1, positions)
print('padding mask, per sequence:')
print(pad_mask.squeeze(1).int())

torch.manual_seed(0)
scores = torch.randn(2, 5, 5)
masked = scores.masked_fill(pad_mask == 0, float('-inf'))
w = masked.softmax(-1)
print('\nsequence 0, weights from position 0:')
print(' ', ' '.join('%.3f' % v for v in w[0, 0].tolist()))
print('the two padded positions get exactly zero weight.')
print('\nwithout it, padding would contribute to every average, and')
print('a short sequence in a batch of long ones would be mostly padding.')

padding mask, per sequence:
tensor([[1, 1, 1, 0, 0],
        [1, 1, 1, 1, 1]], dtype=torch.int32)

sequence 0, weights from position 0:
  0.229 0.223 0.549 0.000 0.000
the two padded positions get exactly zero weight.

without it, padding would contribute to every average, and
a short sequence in a batch of long ones would be mostly padding.

Two masks, combined with a logical and

A causal language model on padded batches needs both: the causal mask stops a position seeing the future, and the padding mask stops any position seeing padding. They are combined by multiplying or by a logical and before the masked_fill. Getting one and not the other produces a model that trains and is quietly wrong, which is the recurring theme of this course.

Day 4 takeaway

Self-attention queries the sequence with itself; cross-attention queries one sequence with another. A lower triangular mask makes attention causal, which is what lets a language model produce a prediction at every position in one pass. Padding needs its own mask, and you usually need both at once.

Week 08 · Day 5 of 7

Positional Information

Putting back the order that attention discards

By Pawan Rama Mali802 words

Day 2 showed attention treating the sequence as a set. Since order usually matters, the position has to be added to the input, and there are three ways of doing it.

Learned position embeddings

import torch
from torch import nn

torch.manual_seed(0)
max_len, width = 8, 16
tokens = nn.Embedding(10, width)
positions = nn.Embedding(max_len, width)

x = torch.tensor([[3, 7, 2, 5]])
pos = torch.arange(x.shape[1])
combined = tokens(x) + positions(pos)

print('token embeddings   ', tuple(tokens(x).shape))
print('position embeddings', tuple(positions(pos).shape))
print('added together     ', tuple(combined.shape))
print('\nthe same token at two positions is now two different vectors:')
same = torch.tensor([[5, 5]])
out = tokens(same) + positions(torch.arange(2))
print('  difference %.4f' % (out[0, 0] - out[0, 1]).abs().mean().item())

token embeddings    (1, 4, 16)
position embeddings (4, 16)
added together      (1, 4, 16)

the same token at two positions is now two different vectors:
  difference 0.7602

Sinusoidal encodings

import torch
import math

def sinusoidal(max_len, width):
    pos = torch.arange(max_len).unsqueeze(1).float()
    i = torch.arange(0, width, 2).float()
    rate = torch.exp(-math.log(10000.0) * i / width)
    pe = torch.zeros(max_len, width)
    pe[:, 0::2] = torch.sin(pos * rate)
    pe[:, 1::2] = torch.cos(pos * rate)
    return pe

pe = sinusoidal(6, 8)
print('positional encodings, 6 positions of width 8:')
for i, row in enumerate(pe):
    print('  %d  %s' % (i, ' '.join('%+.2f' % v for v in row.tolist())))

print('\ndot product between position 0 and each other position:')
print(' ', ' '.join('%.2f' % v for v in (pe @ pe[0]).tolist()))
print('it falls away smoothly, so nearby positions are similar.')

positional encodings, 6 positions of width 8:
+0.00 +1.00 +0.00 +1.00 +0.00 +1.00 +0.00 +1.00
+0.84 +0.54 +0.10 +1.00 +0.01 +1.00 +0.00 +1.00
+0.91 -0.42 +0.20 +0.98 +0.02 +1.00 +0.00 +1.00
+0.14 -0.99 +0.30 +0.96 +0.03 +1.00 +0.00 +1.00
-0.76 -0.65 +0.39 +0.92 +0.04 +1.00 +0.00 +1.00
-0.96 +0.28 +0.48 +0.88 +0.05 +1.00 +0.00 +1.00

dot product between position 0 and each other position:
00 3.54 2.56 1.96 2.27 3.16
it falls away smoothly, so nearby positions are similar.

Scheme	Extends past training length?	Used by
Learned embeddings	No, it has no row for position 5000	BERT, the original GPT
Sinusoidal	In principle yes, in practice poorly	The original transformer
Rotary (RoPE)	Much better	Llama, most current models
Relative position bias	Yes	T5

Whether it matters, measured

import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset

def make_sequences(n, length=8, seed=0):
    """Week 7's task: is the digit at position 0 repeated later?
    Balanced by construction, and unanswerable from a bag of counts."""
    g = torch.Generator().manual_seed(seed)
    first = torch.randint(1, 10, (n, 1), generator=g)
    rest = torch.randint(1, 9, (n, length - 1), generator=g)
    rest = rest + (rest >= first).long()
    y = torch.randint(0, 2, (n,), generator=g)
    where = torch.randint(0, length - 1, (n, 1), generator=g)
    idx = torch.arange(length - 1).unsqueeze(0)
    plant = (y.unsqueeze(1) == 1) & (idx == where)
    rest = torch.where(plant, first.expand_as(rest), rest)
    return torch.cat([first, rest], dim=1), y

X_tr, y_tr = make_sequences(8000, seed=0)
X_va, y_va = make_sequences(2000, seed=1)
train_loader = DataLoader(TensorDataset(X_tr, y_tr), batch_size=64,
                          shuffle=True)
val_loader = DataLoader(TensorDataset(X_va, y_va), batch_size=512)
def fit(model, epochs=10, lr=3e-3, loader=None, val=None):
    torch.manual_seed(0)
    opt = torch.optim.AdamW(model.parameters(), lr=lr)
    loss_fn = nn.CrossEntropyLoss()
    for _ in range(epochs):
        model.train()
        for xb, yb in (loader or train_loader):
            opt.zero_grad()
            loss_fn(model(xb), yb).backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            opt.step()
    model.eval()
    right = seen = 0
    with torch.no_grad():
        for xb, yb in (val or val_loader):
            right += (model(xb).argmax(1) == yb).sum().item()
            seen += yb.numel()
    return right / seen
class AttentionClassifier(nn.Module):
    def __init__(self, width=32, heads=4, use_positions=True):
        super().__init__()
        self.use_positions = use_positions
        self.tok = nn.Embedding(10, width)
        self.pos = nn.Embedding(16, width)
        self.attn = nn.MultiheadAttention(width, heads, batch_first=True)
        self.norm = nn.LayerNorm(width)
        self.head = nn.Linear(width, 2)

    def forward(self, x):
        h = self.tok(x)
        if self.use_positions:
            h = h + self.pos(torch.arange(x.shape[1]))
        a, _ = self.attn(h, h, h)
        h = self.norm(h + a)
        return self.head(h.mean(dim=1))

print('%-22s %12s' % ('', 'accuracy'))
for use in [False, True]:
    torch.manual_seed(0)
    print('%-22s %12.4f'
          % ('with positions' if use else 'no positions',
             fit(AttentionClassifier(use_positions=use))))

                           accuracy
no positions                 0.5855
with positions               0.9995

This task asks whether the digit at position 0 reappears, so a model with no notion of position cannot express the question. The gap between the two rows is the cost of treating a sequence as a set.

Day 5 takeaway

Attention needs position supplied separately. Learned embeddings are simplest and cannot extend beyond the length they were trained on; rotary encodings are what current models use because they extrapolate. On a task where position is the question, leaving them out costs you the task.

Week 08 · Day 6 of 7

The Library and the Cost

nn.MultiheadAttention, the inverted mask, and quadratic growth

By Pawan Rama Mali546 words

The library version, and the cost that decides how long a sequence you can afford.

nn.MultiheadAttention

import torch
from torch import nn

torch.manual_seed(0)
mha = nn.MultiheadAttention(embed_dim=32, num_heads=4, batch_first=True)
x = torch.randn(2, 6, 32)

out, weights = mha(x, x, x)                 # query, key, value
print('output ', tuple(out.shape))
print('weights', tuple(weights.shape), ' averaged over heads by default')

out, weights = mha(x, x, x, average_attn_weights=False)
print('per head', tuple(weights.shape))

causal = torch.triu(torch.ones(6, 6, dtype=torch.bool), diagonal=1)
out, w = mha(x, x, x, attn_mask=causal)
print('\nwith a causal mask, position 0 attends to:',
      ' '.join('%.2f' % v for v in w[0, 0].tolist()))

output  (2, 6, 32)
weights (2, 6, 6)  averaged over heads by default
per head (2, 4, 6, 6)

with a causal mask, position 0 attends to: 1.00 0.00 0.00 0.00 0.00 0.00

The mask convention is the opposite of what you expect

nn.MultiheadAttention takes attn_mask where True means block this position, which is the reverse of the 1 means keep convention used when you write it yourself, and the reverse of key_padding_mask in some older code. Read the docstring every time. An inverted mask produces a model that trains happily while attending to exactly the positions it should not.

The cost

print('%10s %16s %18s' % ('length', 'attention scores', 'recurrent steps'))
for T in [8, 128, 1024, 8192]:
    print('%10d %16d %18d' % (T, T * T, T))
print('\nattention computes a score for every pair of positions, so')
print('memory and time grow with the square of the length. Recurrence')
print('grows linearly, but cannot be parallelised along the sequence.')

    length attention scores    recurrent steps
         8               64                  8
       128            16384                128
      1024          1048576               1024
      8192         67108864               8192

attention computes a score for every pair of positions, so
memory and time grow with the square of the length. Recurrence
grows linearly, but cannot be parallelised along the sequence.

import time
import torch
from torch import nn

torch.manual_seed(0)
attn = nn.MultiheadAttention(64, 4, batch_first=True)
lstm = nn.LSTM(64, 64, batch_first=True)

print('%10s %14s %14s' % ('length', 'attention', 'lstm'))
for T in [32, 128, 512]:
    x = torch.randn(8, T, 64)
    t = time.time()
    for _ in range(10):
        attn(x, x, x, need_weights=False)
    a = (time.time() - t) / 10
    t = time.time()
    for _ in range(10):
        lstm(x)
    b = (time.time() - t) / 10
    print('%10d %13.1fms %13.1fms' % (T, a * 1000, b * 1000))

    length      attention           lstm
        32           0.5ms           3.2ms
       128           1.8ms           2.3ms
       512           8.5ms           6.8ms

Read that table carefully

Attention is doing quadratically more arithmetic as the sequence grows, and it is still competitive or faster, because all of it happens in one large matrix multiplication that a modern processor is built for. The LSTM does linearly less work in a Python-level loop over time steps that cannot be parallelised at all.

This is the practical answer to why transformers won. Not that they do less work, but that the work they do is the shape hardware is fast at. It also means the crossover point moves as sequences get very long, which is why efficient attention variants exist.

Day 6 takeaway

Use nn.MultiheadAttention and check the mask convention, because True means blocked. Attention costs the square of the sequence length in memory and time, and wins anyway at moderate lengths because it parallelises where recurrence cannot.

Week 08 · Day 7 of 7

The Transformer Block

Attention assembled, measured against week 7, and read honestly

By Pawan Rama Mali1612 words

A complete attention block, and the comparison against week 7's recurrent models on identical data.

The block

import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset

def make_sequences(n, length=8, seed=0):
    """Week 7's task: is the digit at position 0 repeated later?
    Balanced by construction, and unanswerable from a bag of counts."""
    g = torch.Generator().manual_seed(seed)
    first = torch.randint(1, 10, (n, 1), generator=g)
    rest = torch.randint(1, 9, (n, length - 1), generator=g)
    rest = rest + (rest >= first).long()
    y = torch.randint(0, 2, (n,), generator=g)
    where = torch.randint(0, length - 1, (n, 1), generator=g)
    idx = torch.arange(length - 1).unsqueeze(0)
    plant = (y.unsqueeze(1) == 1) & (idx == where)
    rest = torch.where(plant, first.expand_as(rest), rest)
    return torch.cat([first, rest], dim=1), y

X_tr, y_tr = make_sequences(8000, seed=0)
X_va, y_va = make_sequences(2000, seed=1)
train_loader = DataLoader(TensorDataset(X_tr, y_tr), batch_size=64,
                          shuffle=True)
val_loader = DataLoader(TensorDataset(X_va, y_va), batch_size=512)
def fit(model, epochs=10, lr=3e-3, loader=None, val=None):
    torch.manual_seed(0)
    opt = torch.optim.AdamW(model.parameters(), lr=lr)
    loss_fn = nn.CrossEntropyLoss()
    for _ in range(epochs):
        model.train()
        for xb, yb in (loader or train_loader):
            opt.zero_grad()
            loss_fn(model(xb), yb).backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            opt.step()
    model.eval()
    right = seen = 0
    with torch.no_grad():
        for xb, yb in (val or val_loader):
            right += (model(xb).argmax(1) == yb).sum().item()
            seen += yb.numel()
    return right / seen
class AttentionBlock(nn.Module):
    """Attention, then a small feed forward network, each wrapped in a
    residual connection with layer normalisation. This is the transformer
    block, and week 9 assembles a stack of them."""
    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                               # residual
        x = x + self.ff(self.norm2(x))          # residual
        return x

torch.manual_seed(0)
block = AttentionBlock(32, 4)
x = torch.randn(2, 8, 32)
print('in ', tuple(x.shape), ' out', tuple(block(x).shape))
print('parameters %d' % sum(p.numel() for p in block.parameters()))

in  (2, 8, 32)  out (2, 8, 32)
parameters 12704

Normalise before, not after

The original transformer applied layer normalisation after the residual addition. Every model since about 2019 does it before, exactly as week 6 argued for pre-activation residual blocks: it keeps the shortcut path completely clear, so the gradient reaches the earliest block undiminished. Post-normalisation transformers need a careful warmup to train at all, and pre-normalisation ones largely do not.

Against the recurrent models

import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset

def make_sequences(n, length=8, seed=0):
    """Week 7's task: is the digit at position 0 repeated later?
    Balanced by construction, and unanswerable from a bag of counts."""
    g = torch.Generator().manual_seed(seed)
    first = torch.randint(1, 10, (n, 1), generator=g)
    rest = torch.randint(1, 9, (n, length - 1), generator=g)
    rest = rest + (rest >= first).long()
    y = torch.randint(0, 2, (n,), generator=g)
    where = torch.randint(0, length - 1, (n, 1), generator=g)
    idx = torch.arange(length - 1).unsqueeze(0)
    plant = (y.unsqueeze(1) == 1) & (idx == where)
    rest = torch.where(plant, first.expand_as(rest), rest)
    return torch.cat([first, rest], dim=1), y

X_tr, y_tr = make_sequences(8000, seed=0)
X_va, y_va = make_sequences(2000, seed=1)
train_loader = DataLoader(TensorDataset(X_tr, y_tr), batch_size=64,
                          shuffle=True)
val_loader = DataLoader(TensorDataset(X_va, y_va), batch_size=512)
def fit(model, epochs=10, lr=3e-3, loader=None, val=None):
    torch.manual_seed(0)
    opt = torch.optim.AdamW(model.parameters(), lr=lr)
    loss_fn = nn.CrossEntropyLoss()
    for _ in range(epochs):
        model.train()
        for xb, yb in (loader or train_loader):
            opt.zero_grad()
            loss_fn(model(xb), yb).backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            opt.step()
    model.eval()
    right = seen = 0
    with torch.no_grad():
        for xb, yb in (val or val_loader):
            right += (model(xb).argmax(1) == yb).sum().item()
            seen += yb.numel()
    return right / seen
import time

class AttentionBlock(nn.Module):
    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):
        h = self.norm1(x)
        a, _ = self.attn(h, h, h, need_weights=False)
        x = x + a
        return x + self.ff(self.norm2(x))

class AttentionClassifier(nn.Module):
    def __init__(self, width=32, heads=4, blocks=2):
        super().__init__()
        self.tok = nn.Embedding(10, width)
        self.pos = nn.Embedding(16, width)
        self.blocks = nn.ModuleList([AttentionBlock(width, heads)
                                     for _ in range(blocks)])
        self.norm = nn.LayerNorm(width)
        self.head = nn.Linear(width, 2)

    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.head(self.norm(h).mean(dim=1))

class Recurrent(nn.Module):
    def __init__(self, cell):
        super().__init__()
        self.emb = nn.Embedding(10, 32)
        self.rnn = {'rnn': nn.RNN, 'gru': nn.GRU,
                    'lstm': nn.LSTM}[cell](32, 64, batch_first=True)
        self.head = nn.Linear(64, 2)

    def forward(self, x):
        out, _ = self.rnn(self.emb(x))
        return self.head(out[:, -1, :])

print('%-26s %10s %10s %9s' % ('', 'params', 'accuracy', 'time'))
for name, maker in [('plain RNN', lambda: Recurrent('rnn')),
                    ('LSTM', lambda: Recurrent('lstm')),
                    ('attention, 1 block',
                     lambda: AttentionClassifier(blocks=1)),
                    ('attention, 2 blocks',
                     lambda: AttentionClassifier(blocks=2))]:
    torch.manual_seed(0)
    model = maker()
    start = time.time()
    acc = fit(model)
    print('%-26s %10d %10.4f %8.0fs'
          % (name, sum(p.numel() for p in model.parameters()), acc,
             time.time() - start))

                               params   accuracy      time
plain RNN                        6722     0.6380        5s
LSTM                            25538     0.9860        6s
attention, 1 block              13666     1.0000       10s
attention, 2 blocks             26370     1.0000       12s

Reading what it attended to

import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset

def make_sequences(n, length=8, seed=0):
    """Week 7's task: is the digit at position 0 repeated later?
    Balanced by construction, and unanswerable from a bag of counts."""
    g = torch.Generator().manual_seed(seed)
    first = torch.randint(1, 10, (n, 1), generator=g)
    rest = torch.randint(1, 9, (n, length - 1), generator=g)
    rest = rest + (rest >= first).long()
    y = torch.randint(0, 2, (n,), generator=g)
    where = torch.randint(0, length - 1, (n, 1), generator=g)
    idx = torch.arange(length - 1).unsqueeze(0)
    plant = (y.unsqueeze(1) == 1) & (idx == where)
    rest = torch.where(plant, first.expand_as(rest), rest)
    return torch.cat([first, rest], dim=1), y

X_tr, y_tr = make_sequences(8000, seed=0)
X_va, y_va = make_sequences(2000, seed=1)
train_loader = DataLoader(TensorDataset(X_tr, y_tr), batch_size=64,
                          shuffle=True)
val_loader = DataLoader(TensorDataset(X_va, y_va), batch_size=512)
def fit(model, epochs=10, lr=3e-3, loader=None, val=None):
    torch.manual_seed(0)
    opt = torch.optim.AdamW(model.parameters(), lr=lr)
    loss_fn = nn.CrossEntropyLoss()
    for _ in range(epochs):
        model.train()
        for xb, yb in (loader or train_loader):
            opt.zero_grad()
            loss_fn(model(xb), yb).backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            opt.step()
    model.eval()
    right = seen = 0
    with torch.no_grad():
        for xb, yb in (val or val_loader):
            right += (model(xb).argmax(1) == yb).sum().item()
            seen += yb.numel()
    return right / seen
class Simple(nn.Module):
    def __init__(self, width=32, heads=1):
        super().__init__()
        self.tok = nn.Embedding(10, width)
        self.pos = nn.Embedding(16, width)
        self.attn = nn.MultiheadAttention(width, heads, batch_first=True)
        self.norm = nn.LayerNorm(width)
        self.head = nn.Linear(width, 2)

    def forward(self, x, want_weights=False):
        h = self.tok(x) + self.pos(torch.arange(x.shape[1]))
        a, w = self.attn(h, h, h)
        out = self.head(self.norm(h + a).mean(dim=1))
        return (out, w) if want_weights else out

torch.manual_seed(0)
model = Simple()
print('accuracy %.4f\n' % fit(model))

model.eval()
with torch.no_grad():
    _, w = model(X_va[:1], want_weights=True)
seq = X_va[0].tolist()
print('sequence %s, label %d' % (''.join(str(d) for d in seq), y_va[0]))
print('\nattention weights from each position:')
print('      ' + ' '.join('%5d' % d for d in seq))
for i, row in enumerate(w[0]):
    print('%2d %s  %s' % (i, str(seq[i]).rjust(2),
                          ' '.join('%5.2f' % v for v in row.tolist())))

accuracy 0.8725

sequence 57959964, label 1

attention weights from each position:
          5     7     9     5     9     9     6     4
 0  5   0.00  0.00  0.00  1.00  0.00  0.00  0.00  0.00
 1  7   0.00  0.00  0.10  0.28  0.19  0.35  0.02  0.06
 2  9   0.00  0.00  0.36  0.00  0.35  0.29  0.00  0.00
 3  5   0.93  0.00  0.00  0.07  0.00  0.00  0.00  0.00
 4  9   0.00  0.00  0.13  0.00  0.30  0.57  0.00  0.00
 5  9   0.00  0.00  0.43  0.00  0.33  0.24  0.00  0.00
 6  6   0.00  0.00  0.00  0.00  0.00  0.00  1.00  0.00
 7  4   0.00  0.00  0.00  0.00  0.00  0.00  0.00  1.00

Read the first row. Position 0 holds a 5, and it has put essentially all of its weight on the one later position that also holds a 5. The model has discovered the structure of the task and is using attention to do exactly the lookup the task requires, which is the clearest picture of attention working that this course can offer.

Do not over-read an attention map

It is tempting to treat these weights as an explanation. They are not: they show what the model looked at, not why it decided. A well-documented result is that attention weights can be changed substantially while leaving the prediction almost unaltered, which means they are not a faithful account of the reasoning. Read them as a debugging aid, in the same spirit as week 5's occlusion test, and not as evidence.

Day 7 takeaway

A transformer block is attention plus a feed forward network, each in a residual wrapper with layer normalisation applied first. It matches or beats the recurrent models on this task, computes every position in parallel, and its attention maps are a debugging tool rather than an explanation. Week 9 stacks these into a full transformer.