Multi-Head Attention and Transformers

The attention layer carried a single query and reduced the whole sequence to a single weight distribution. The task was suited to this: the label was read from a single step and the order in which the steps arrived changed nothing. In real sequences two things change at once. A decision often comes out of the combination of information read from several places in the sequence, and where that information stands matters as much as the information itself.

This lesson measures two additions. The first is the multiplication of the query: the same inner dimension is divided among several heads, each head produces its own weight distribution, and the outputs are concatenated. The second is the positional encoding: the place of the step in the sequence is added next to the input as a number. These two, together with a per-step projection and skip connections, are the load-bearing parts of the transformer architecture — a structure without recurrence that looks at all the steps at once. Here only the two load-bearing parts are written and measured.

• AR55. The sequence setup of the Architectures topic was built in architectures/06 and is fictional: 720 sequences, length 12, seed 20260218. The same sequences are regenerated with the same seed.
• AR56. The question of this lesson is different and depends on order: the label is whether the average of the seventh through eleventh steps exceeds the average of the first through fifth steps by more than 0.30. The sixth step is left out of both regions.
• AR57. The baseline is again the most frequent class. Because the label is read from the difference of two regions, the general level of the sequence is not enough to pass the baseline.
• AR58. The positional encoding is a sine channel and a cosine channel: $\sin(\pi t / 11)$ and $\cos(\pi t / 11)$ for step $t$. A single frequency was chosen to keep the scale small; real setups use several frequencies.
• AR59. The comparison is under an equal parameter budget. The total inner dimension is fixed at 4 and the size per head shrinks as the number of heads grows. The parameter count is the same in all three setups.
• AR60. The layer is written inside the lesson and is a model. No library is called.
• AR61. Training is descent with momentum: 150 epochs, learning rate 0.15, momentum 0.8. Momentum was built in the neural-network-fundamentals topic of this course and is not repeated here.
• AR62. Each setup is printed with two runs; the runs differ only in the initialization seed.
• AR63. The order test is this: the readings of a sequence are shuffled among themselves, the position channels are left in place, and how much the output probability moved is measured.

The Same Budget, a Different Number of Heads

The sequences of the previous lesson are used as they are; the question changes. Because the new label is read from the difference between the averages of two regions, it is hard to produce with a single weighted average: a sequence-wide average cannot separate two regions.

# common.py — MODEL. The FICTIONAL sequence set built in `architectures/06` is regenerated with the same seed.
import math

SEED, M32, STEPS = 20260218, 0xFFFFFFFF, 12


def generator(t):
    x = ((t ^ (t >> 16)) * 2246822507) & M32
    x = ((x ^ (x >> 13)) * 3266489909) & M32
    s = [(x ^ (x >> 16)) & M32]

    def next_value():
        s[0] = (s[0] * 1664525 + 1013904223) & M32
        return s[0] / 4294967296
    return next_value


def build_sequences(n, seed):
    D = []
    for i in range(n):
        r = generator(seed + i)
        r()  # region draw: not used in this lesson, consumed so the sequence stays the same
        reading = [0.0] + [round((r() + r() + r() - 1.5) * 1.9, 3) for _ in range(STEPS - 1)]
        r()
        D.append({"reading": reading})
    return D


DATA = build_sequences(720, SEED)
FIRST, LAST = range(1, 6), range(7, 12)
for x in DATA:
    a = sum(x["reading"][t] for t in FIRST) / 5
    b = sum(x["reading"][t] for t in LAST) / 5
    x["y"] = 1 if b - a > 0.30 else 0
TRAIN, TEST = DATA[:200], DATA[540:]
rate = sum(x["y"] for x in TEST) / len(TEST)
BASELINE = max(rate, 1 - rate)
print(f"sequences {len(DATA)}, length {STEPS}, training {len(TRAIN)}, held-out {len(TEST)}")
print(f"rise label rate train {sum(x['y'] for x in TRAIN) / len(TRAIN):.4f}, "
      f"held-out {rate:.4f}")
print(f"baseline (most frequent class) {BASELINE:.4f}")
print(f"example sequence {DATA[0]['reading']}")

sequences 720, length 12, training 200, held-out 180
rise label rate train 0.3150, held-out 0.3389
baseline (most frequent class) 0.6611
example sequence [0.0, 0.893, -0.222, -0.542, -0.404, -1.114, 0.234, -1.03, -1.493, 0.534, -1.807, 0.05]

In multi-head attention every head has its own key, value and query projections. Because the total inner dimension is held fixed, when the number of heads doubles the size per head is halved; the parameter count does not change. The meaning of the comparison lies here — if there is a gain, it comes not from capacity but from the way the same capacity is divided.

DTOP, DIN = 4, 3


def inputs(x, encode=True):
    K = [[math.sin(math.pi * t / (STEPS - 1)), math.cos(math.pi * t / (STEPS - 1))]
         for t in range(STEPS)]
    return [[x["reading"][t]] + (K[t] if encode else [0.0, 0.0]) for t in range(STEPS)]


def init_params(seed, H):
    d = DTOP // H
    r = generator(seed)
    g = lambda: (r() + r() + r() - 1.5) * 0.9
    return {"H": H, "d": d, "scale": 1.0 / math.sqrt(d),
            "Wk": [[[g() for _ in range(DIN)] for _ in range(d)] for _ in range(H)],
            "Wv": [[[g() for _ in range(DIN)] for _ in range(d)] for _ in range(H)],
            "q": [[g() for _ in range(d)] for _ in range(H)],
            "w": [g() for _ in range(DTOP)], "b": [0.0]}


def weights(P, X, h):
    # the query and the key projection fold into a single direction: u = q^T Wk
    d, wk, q = P["d"], P["Wk"][h], P["q"][h]
    u = [sum(q[i] * wk[i][j] for i in range(d)) for j in range(DIN)]
    s = [(u[0] * x[0] + u[1] * x[1] + u[2] * x[2]) * P["scale"] for x in X]
    mx = max(s)
    e = [math.exp(v - mx) for v in s]
    z = sum(e)
    return [v / z for v in e]


def head(P, X, h):
    a = weights(P, X, h)
    Xa = [sum(a[t] * X[t][j] for t in range(STEPS)) for j in range(DIN)]
    return a, Xa, [sum(P["Wv"][h][i][j] * Xa[j] for j in range(DIN))
                   for i in range(P["d"])]


def forward(P, X):
    d, o = P["d"], P["b"][0]
    for h in range(P["H"]):
        c = head(P, X, h)[2]
        o += sum(P["w"][h * d + i] * c[i] for i in range(d))
    return o


def train_model(seed, H, TR, epochs=150, lr=0.15, mom=0.8):
    P = init_params(seed, H)
    d, scale, n, vel = P["d"], P["scale"], len(TR), None
    for _ in range(epochs):
        gWk = [[[0.0] * DIN for _ in range(d)] for _ in range(H)]
        gWv = [[[0.0] * DIN for _ in range(d)] for _ in range(H)]
        gq = [[0.0] * d for _ in range(H)]
        gw, gb = [0.0] * DTOP, [0.0]
        for X, y in TR:
            B = [head(P, X, h) for h in range(H)]
            o = P["b"][0] + sum(P["w"][h * d + i] * B[h][2][i]
                                for h in range(H) for i in range(d))
            df = 1.0 / (1.0 + math.exp(-max(-30.0, min(30.0, o)))) - y
            gb[0] += df
            for h in range(H):
                a, Xa, c = B[h]
                dc = [df * P["w"][h * d + i] for i in range(d)]
                mv = [sum(dc[i] * P["Wv"][h][i][j] for i in range(d)) for j in range(DIN)]
                da = [mv[0] * x[0] + mv[1] * x[1] + mv[2] * x[2] for x in X]
                mean = sum(a[t] * da[t] for t in range(STEPS))
                S = [0.0, 0.0, 0.0]
                for t in range(STEPS):
                    ds, x = a[t] * (da[t] - mean) * scale, X[t]
                    S[0] += ds * x[0]
                    S[1] += ds * x[1]
                    S[2] += ds * x[2]
                for i in range(d):
                    gw[h * d + i] += df * c[i]
                    gq[h][i] += sum(P["Wk"][h][i][j] * S[j] for j in range(DIN))
                    for j in range(DIN):
                        gWk[h][i][j] += P["q"][h][i] * S[j]
                        gWv[h][i][j] += dc[i] * Xa[j]
        params = [P["Wk"][h][i] for h in range(H) for i in range(d)] \
            + [P["Wv"][h][i] for h in range(H) for i in range(d)] \
            + P["q"] + [P["w"], P["b"]]
        grads = [gWk[h][i] for h in range(H) for i in range(d)] \
            + [gWv[h][i] for h in range(H) for i in range(d)] + gq + [gw, gb]
        if vel is None:
            vel = [[0.0] * len(u) for u in params]
        for u, v, g in zip(params, grads, vel):   # gradient descent with momentum
            for i in range(len(u)):
                g[i] = mom * g[i] - lr * v[i] / n
                u[i] += g[i]
    return P


def accuracy(P, G):
    return sum((forward(P, X) > 0) == (y == 1) for X, y in G) / len(G)


ALL_IN = [(inputs(x), x["y"]) for x in DATA]
TR_IN, TE_IN = ALL_IN[:200], ALL_IN[540:]
NETS = {}
print(f"total inner dimension {DTOP}, epochs 150, training examples {len(TR_IN)}, "
      f"baseline {BASELINE:.4f}")
print(f"{'heads':>6}{'head size':>12}{'parameters':>11}{'run 1':>10}{'run 2':>10}"
      f"{'mean':>10}{'run difference':>16}")
for H in (1, 2, 4):
    NETS[H] = [train_model(SEED + k, H, TR_IN) for k in (1, 2)]
    a, b = (accuracy(P, TE_IN) for P in NETS[H])
    print(f"{H:>6}{DTOP // H:>12}{2 * DTOP * DIN + 2 * DTOP + 1:>11}{a:>10.4f}"
          f"{b:>10.4f}{(a + b) / 2:>10.4f}{abs(a - b):>16.4f}")

total inner dimension 4, epochs 150, training examples 200, baseline 0.6611
 heads   head size parameters     run 1     run 2      mean  run difference
     1           4         33    0.8889    0.8556    0.8722          0.0333
     2           2         33    0.8778    0.8833    0.8806          0.0056
     4           1         33    0.8722    0.9222    0.8972          0.0500

The means move in one direction: 0.8722, 0.8806, 0.8972. Raising the number of heads from one to four adds 0.0250 to the held-out figure and carries it 0.2361 above the baseline. The column on the right puts the gain in context: the difference between the two runs of the four-head setup is 0.0500, which is twice what the number of heads adds. The difference that dividing the same capacity among heads brings is smaller than the difference that rerunning that division itself brings. At this scale the number of heads is a hyperparameter, and if it is chosen with a single run it is chosen wrongly.

What the Heads Look At

Why the number of heads adds anything can be asked by looking at where the weights fall. The task wants the average of two regions; if the heads split into these two regions, the structure has found the form of the task. The measure is two numbers: the weight center of each head and the separation between heads — half the difference between two distributions, from zero to one.

def center(P, G, h):
    return sum(sum(t * weights(P, X, h)[t] for t in range(STEPS)) for X, _ in G) / len(G)


def separation(P, G):
    H, s = P["H"], 0.0
    if H < 2:
        return 0.0
    for X, _ in G:
        A = [weights(P, X, h) for h in range(H)]
        for i in range(H):
            for j in range(i + 1, H):
                s += sum(abs(A[i][t] - A[j][t]) for t in range(STEPS)) / 2
    return s / (len(G) * H * (H - 1) / 2)


print(f"the two regions the task wants: steps {min(FIRST)}-{max(FIRST)} and {min(LAST)}-{max(LAST)}")
print(f"{'':18}{'weight centers':>34}{'separation':>12}{'held-out':>10}")
for H in (1, 2, 4):
    for k in (0, 1):
        P = NETS[H][k]
        m = ", ".join(f"{center(P, TE_IN, h):.2f}" for h in range(H))
        name = f"{H} head" + ("s" if H > 1 else "") + f", run {k + 1}"
        print(f"{name:<18}{m:>34}{separation(P, TE_IN):>12.4f}"
              f"{accuracy(P, TE_IN):>10.4f}")

the two regions the task wants: steps 1-5 and 7-11
                                      weight centers  separation  held-out
1 head, run 1                                   5.51      0.0000    0.8889
1 head, run 2                                   6.02      0.0000    0.8556
2 heads, run 1                            5.57, 6.30      0.4096    0.8778
2 heads, run 2                            3.30, 8.55      0.7264    0.8833
4 heads, run 1                5.84, 5.48, 5.79, 5.61      0.4782    0.8722
4 heads, run 2                3.24, 5.81, 8.37, 4.59      0.4291    0.9222

The second run of the two-head setup makes exactly the expected distinction: one head stands at 3.30, the other at 8.55, and these two points fall in the middle of the two regions the task wants. The separation is 0.7264, that is, the two distributions largely look at different steps.

The first run of the same setup does not do this. The heads are at 5.57 and 6.30, both standing in the middle of the sequence, and the separation stays at 0.4096. Yet the held-out figure is 0.8778 — only 0.0055 below the run that separates. The same observation is even clearer in the four-head setup: while the four heads of the first run are squeezed between 5.48 and 5.84, its accuracy is 0.8722. The structure “finding the form of the task” comes not from the architecture but from the run, and a run that has not found it can also produce a close number.

Positional Encoding Is a Prerequisite, Not an Improvement

Attention looks at all the steps at once and carries no order link between steps. If there is no number about its place in the input of a step, then for the layer the sequence is a multiset: the same readings arriving in another order give the same output. This is not a loss of performance but an absence of information, and it can be tested directly.

def shuffle_readings(X, seed):
    r, s = generator(seed), list(range(STEPS))
    for i in range(STEPS - 1, 0, -1):
        j = int(r() * (i + 1))
        s[i], s[j] = s[j], s[i]
    return [[X[s[t]][0], X[t][1], X[t][2]] for t in range(STEPS)]


def probability(P, X):
    return 1.0 / (1.0 + math.exp(-max(-30.0, min(30.0, forward(P, X)))))


def order_test(P, G):
    e, mx = 0.0, 0.0
    for i, (X, _y) in enumerate(G):
        d = abs(probability(P, shuffle_readings(X, SEED + 4000 + i)) - probability(P, X))
        e, mx = e + d, max(mx, d)
    return e / len(G), mx


NO_POS = [(inputs(x, False), x["y"]) for x in DATA]
NO_ENC = [train_model(SEED + k, 2, NO_POS[:200]) for k in (1, 2)]
print(f"baseline (most frequent class) {BASELINE:.4f}")
print(f"{'':28}{'run 1':>9}{'run 2':>9}{'mean |change|':>15}{'largest':>10}")
for name, M, G in (("positional encoding present", NETS[2], TE_IN),
                   ("positional encoding absent", NO_ENC, NO_POS[540:])):
    e, mx = order_test(M[1], G)
    print(f"{name:<28}{accuracy(M[0], G):>9.4f}{accuracy(M[1], G):>9.4f}"
          f"{e:>15.4f}{mx:>10.4f}")

baseline (most frequent class) 0.6611
                                run 1    run 2  mean |change|   largest
positional encoding present    0.8778   0.8833         0.3173    0.9753
positional encoding absent     0.6611   0.6611         0.0000    0.0000

The bottom row is sharp in two places at once. The setup without the encoding gives exactly 0.6611 in both runs — the baseline itself. There is no difference even in the decimal places, because the network cannot learn anything other than giving every sequence the same class. The order test says why: when the readings are shuffled, the output probability moves 0.0000 on average and 0.0000 at most. This is not rounding; because the layer sees the input as an unordered set, the output is the same bit for bit.

In the setup with the encoding, the same shuffling moves the output 0.3173 on average and 0.9753 at most. The difference between them is not that the encoding “helps”; it is that the order enters the reach of the network only through the encoding. In a recurrent network the order comes from the structure itself and does not need to be encoded separately; because attention does not give this assurance, writing the order into the input is not a choice but a necessity.

Summary

• When the same inner dimension is divided among different numbers of heads, the parameter count stays fixed at 33; the held-out mean is 0.8722, 0.8806 and 0.8972, and what the number of heads adds is 0.0250.
• The difference between the two runs of the four-head setup is 0.0500, twice what the number of heads adds; if the number of heads is chosen with a single run, it is chosen wrongly.
• The heads separate into the two regions the task wants only in some runs: the second run of the two-head setup gives a separation of 0.7264 with centers 3.30 and 8.55, its first run stays at 0.4096 with 5.57 and 6.30, and its accuracy is only 0.0055 lower.
• Without the positional encoding both runs fall exactly to the baseline, 0.6611; when the readings are shuffled the output moves 0.0000, because the layer sees the sequence as an unordered set.
• In the setup with the positional encoding, the same shuffling moves the output 0.3173 on average and 0.9753 at most.

Next Step

In these two lessons every step of the input was a ready-made number: the reading a floating-point value, the marker a zero or a one, the position a sine. Most sequences do not arrive like this. A counter type, a tariff tier, a region code or an event label is a discrete symbol and has no natural numeric value. The encoding lessons of the Data Preparation and Feature Engineering course had turned these symbols into one-hot columns; as the number of categories grows, this grows the input dimension in the same proportion and carries no closeness between the symbols. The next lesson builds the embedding layer that maps a discrete input to a learned vector, counts the parameter count against one-hot encoding, and compares this learned representation, on the same measure, with the non-learned projection of the Unsupervised Learning course.