Multi-Head Attention

Example

Two heads over the same 3 tokens. Head 1 compares the first feature, head 2 the second, so they attend differently.

import math
def softmax(r):
    e = [math.exp(x - max(r)) for x in r]
    return [x / sum(e) for x in e]

tokens = ["order", "late", "it"]
X = [[1.0, 0.2], [0.1, 1.0], [0.9, 0.8]]

def head(feature):  # each head here projects onto one feature (a stand-in for learned Wq/Wk)
    q = [x[feature] for x in X]
    weights = [softmax([q[i] * q[j] * 3 for j in range(3)]) for i in range(3)]
    values = [x[feature] for x in X]
    return weights, [sum(w * v for w, v in zip(row, values)) for row in weights]

w1, o1 = head(0)
w2, o2 = head(1)
print("head 1 weights for 'it':", [round(w, 3) for w in w1[2]])
print("head 2 weights for 'it':", [round(w, 3) for w in w2[2]])
concat = [[round(a, 3), round(b, 3)] for a, b in zip(o1, o2)]
for t, c in zip(tokens, concat):
    print(f"{t:6} concat output {c}")

Output

head 1 weights for 'it': [0.54, 0.048, 0.412]
head 2 weights for 'it': [0.083, 0.566, 0.351]
order  concat output [0.926, 0.732]
late   concat output [0.713, 0.889]
it     concat output [0.916, 0.863]