Puzzle Solution Revealed - Transformer: Need for Position Embedding

▶ Watch the 7-minute walkthrough

import torch

torch.set_printoptions(precision=2, sci_mode=False, linewidth=160)

In this notebook we are extending the previous notebook, attention_ann.ipynb, which introduced (plausibly) the world's tiniest hand-constructed transformer model. Please read the previous notebook for the context. This is the solution to the puzzle that was introduced as part of the previous notebook.

We add one word 'disobeys' which modifies the action-attribute words (swap-speech, keep-speech, swap-flight, keep-flight). In the previous notebook the order of the words did not matter. In the current notebook, the order does matter. The word disobeys modifies only the word right after it, so word order matters.

The solution adds an extra layer, and a residual connection that carries the input unmodified to the next block, so each next layer gets both the unmodified input and the modified input.

The residual stream is made 22 bits, where it carries the original inputs and each layer's/head's findings. The most interesting addition is the position information. Each token carries the position where it appears.

This example is constructed to illustrate the ideas. In the real implementation, we do not hand-construct in this manner. Vector embedding and position embedding are often also learned using the training loop and gradient descent. But the hand-construction allows us to see why those blocks and connections exist and how they are helpful to have.

idx = {"fly": 0, "speak": 1, "swap_fly": 2, "swap_speak": 3,
       "object": 4, "action_fly": 5, "action_speak": 6, "question": 7, "disobey": 8}


position_start, L = 9, 6             # position one-hot occupies 9 .. 14
previous_word_is_disobey = 15
object_attribute_fly, object_attribute_speak = 16, 17
is_swap_attr_fly, is_attr_fly_disobeyed = 18, 19
is_swap_attr_speak, is_attr_speak_disobeyed = 20, 21
num_bits = 22

slot_names = {v: k for k, v in idx.items()}
slot_names.update({position_start + p: f"pos{p}" for p in range(L)})
slot_names.update({previous_word_is_disobey: "previous_word_is_disobey",
                   object_attribute_fly: "object_attribute_fly",
                   object_attribute_speak: "object_attribute_speak",
                   is_swap_attr_fly: "is_swap_attr_fly",
                   is_attr_fly_disobeyed: "is_attr_fly_disobeyed",
                   is_swap_attr_speak: "is_swap_attr_speak",
                   is_attr_speak_disobeyed: "is_attr_speak_disobeyed"})

for i in range(num_bits):
    print(f"{i:2} {slot_names[i]}")
0 fly
 1 speak
 2 swap_fly
 3 swap_speak
 4 object
 5 action_fly
 6 action_speak
 7 question
 8 disobey
 9 pos0
10 pos1
11 pos2
12 pos3
13 pos4
14 pos5
15 previous_word_is_disobey
16 object_attribute_fly
17 object_attribute_speak
18 is_swap_attr_fly
19 is_attr_fly_disobeyed
20 is_swap_attr_speak
21 is_attr_speak_disobeyed
token_to_vector = {
    "Rock":             [0,  0,  0,  0,   1,  0,  0,  0,  0],
    "Human":            [0,  1,  0,  0,   1,  0,  0,  0,  0],
    "Crow":             [1,  0,  0,  0,   1,  0,  0,  0,  0],
    "Flying superhero": [1,  1,  0,  0,   1,  0,  0,  0,  0],
    "swap-flight":      [0,  0,  1,  0,   0,  1,  0,  0,  0],
    "swap-speech":      [0,  0,  0,  1,   0,  0,  1,  0,  0],
    "keep-flight":      [0,  0,  0,  0,   0,  1,  0,  0,  0],
    "keep-speech":      [0,  0,  0,  0,   0,  0,  1,  0,  0],
    "disobeys":         [0,  0,  0,  0,   0,  0,  0,  0,  1],
    "he-is?":           [0,  0,  0,  0,   0,  0,  0,  1,  0],
}

for tok, bits in token_to_vector.items():
    print(f"{tok:18} {torch.tensor(bits)}")
Rock               tensor([0, 0, 0, 0, 1, 0, 0, 0, 0])
Human              tensor([0, 1, 0, 0, 1, 0, 0, 0, 0])
Crow               tensor([1, 0, 0, 0, 1, 0, 0, 0, 0])
Flying superhero   tensor([1, 1, 0, 0, 1, 0, 0, 0, 0])
swap-flight        tensor([0, 0, 1, 0, 0, 1, 0, 0, 0])
swap-speech        tensor([0, 0, 0, 1, 0, 0, 1, 0, 0])
keep-flight        tensor([0, 0, 0, 0, 0, 1, 0, 0, 0])
keep-speech        tensor([0, 0, 0, 0, 0, 0, 1, 0, 0])
disobeys           tensor([0, 0, 0, 0, 0, 0, 0, 0, 1])
he-is?             tensor([0, 0, 0, 0, 0, 0, 0, 1, 0])
python
def embed(sentence):
    X = torch.zeros(len(sentence), num_bits)
    for current_token_position, current_token in enumerate(sentence):
        X[current_token_position, :9] = torch.tensor(token_to_vector[current_token], dtype=torch.float32)
        X[current_token_position, position_start + current_token_position] = 1
    return X

sentence = ["Crow", "disobeys", "keep-flight", "swap-speech", "he-is?"]
X = embed(sentence)

for tok, row in zip(sentence, X):
    print(f"{tok:14} {row[:9]}  {row[position_start:position_start + L]}  {row[previous_word_is_disobey:]}")
Crow           tensor([1., 0., 0., 0., 1., 0., 0., 0., 0.])  tensor([1., 0., 0., 0., 0., 0.])  tensor([0., 0., 0., 0., 0., 0., 0.])
disobeys       tensor([0., 0., 0., 0., 0., 0., 0., 0., 1.])  tensor([0., 1., 0., 0., 0., 0.])  tensor([0., 0., 0., 0., 0., 0., 0.])
keep-flight    tensor([0., 0., 0., 0., 0., 1., 0., 0., 0.])  tensor([0., 0., 1., 0., 0., 0.])  tensor([0., 0., 0., 0., 0., 0., 0.])
swap-speech    tensor([0., 0., 0., 1., 0., 0., 1., 0., 0.])  tensor([0., 0., 0., 1., 0., 0.])  tensor([0., 0., 0., 0., 0., 0., 0.])
he-is?         tensor([0., 0., 0., 0., 0., 0., 0., 1., 0.])  tensor([0., 0., 0., 0., 1., 0.])  tensor([0., 0., 0., 0., 0., 0., 0.])
python
def softmax(z, dim=-1):
    z = z - z.max(dim=dim, keepdim=True).values
    e = torch.exp(z)
    return e / e.sum(dim=dim, keepdim=True)

def head(X, Wq, Wk, Wv):
    Q, K, V = X @ Wq, X @ Wk, X @ Wv
    A = softmax(Q @ K.T / Wq.shape[1] ** 0.5)
    return A @ V, A
Wq_O = torch.zeros(num_bits, 2); Wq_O[idx["question"]] = torch.tensor([8.0, 0.0])
Wk_O = torch.zeros(num_bits, 2); Wk_O[idx["object"]]   = torch.tensor([1.0, 0.0])
Wv_O = torch.zeros(num_bits, 2); Wv_O[idx["fly"]] = torch.tensor([1.0, 0.0]); Wv_O[idx["speak"]] = torch.tensor([0.0, 1.0])
Wo_O = torch.zeros(2, num_bits); Wo_O[0, object_attribute_fly] = 1; Wo_O[1, object_attribute_speak] = 1

print("Wq_O nonzero rows:", (Wq_O != 0).any(1).nonzero().flatten())
print("Wk_O nonzero rows:", (Wk_O != 0).any(1).nonzero().flatten())
print("Wv_O nonzero rows:", (Wv_O != 0).any(1).nonzero().flatten())
print("Wo_O nonzero cols:", (Wo_O != 0).any(0).nonzero().flatten())
Wq_O nonzero rows: tensor([7])
Wk_O nonzero rows: tensor([4])
Wv_O nonzero rows: tensor([0, 1])
Wo_O nonzero cols: tensor([16, 17])
S = 24.0

Wq_P = torch.zeros(num_bits, L)
for p in range(1, L):
    Wq_P[position_start + p, p - 1] = S

Wk_P = torch.zeros(num_bits, L)
for p in range(L):
    Wk_P[position_start + p, p] = 1

Wv_P = torch.zeros(num_bits, 1); Wv_P[idx["disobey"], 0] = 1
Wo_P = torch.zeros(1, num_bits); Wo_P[0, previous_word_is_disobey] = 1

M = Wq_P @ Wk_P.T
print("M[9:15, 9:15] =")
print(M[position_start:position_start + L, position_start:position_start + L])
print("symmetric:", torch.allclose(M, M.T))
M[9:15, 9:15] =
tensor([[ 0.,  0.,  0.,  0.,  0.,  0.],
        [24.,  0.,  0.,  0.,  0.,  0.],
        [ 0., 24.,  0.,  0.,  0.,  0.],
        [ 0.,  0., 24.,  0.,  0.,  0.],
        [ 0.,  0.,  0., 24.,  0.,  0.],
        [ 0.,  0.,  0.,  0., 24.,  0.]])
symmetric: False
python
def layer1(X):
    oO, AO = head(X, Wq_O, Wk_O, Wv_O)
    oP, AP = head(X, Wq_P, Wk_P, Wv_P)
    X1 = X + oO @ Wo_O + oP @ Wo_P
    return X1, AO, AP

X1, AO, AP = layer1(X)

print("A_P")
print(AP)
print()
for tok, row in zip(sentence, X1):
    print(f"{tok:14} previous_word_is_disobey={row[previous_word_is_disobey]:.3f}   "
          f"object_attribute_fly={row[object_attribute_fly]:.3f}  "
          f"object_attribute_speak={row[object_attribute_speak]:.3f}")
A_P
tensor([[0.20, 0.20, 0.20, 0.20, 0.20],
        [1.00, 0.00, 0.00, 0.00, 0.00],
        [0.00, 1.00, 0.00, 0.00, 0.00],
        [0.00, 0.00, 1.00, 0.00, 0.00],
        [0.00, 0.00, 0.00, 1.00, 0.00]])

Crow           previous_word_is_disobey=0.200   object_attribute_fly=0.200  object_attribute_speak=0.000
disobeys       previous_word_is_disobey=0.000   object_attribute_fly=0.200  object_attribute_speak=0.000
keep-flight    previous_word_is_disobey=1.000   object_attribute_fly=0.200  object_attribute_speak=0.000
swap-speech    previous_word_is_disobey=0.000   object_attribute_fly=0.200  object_attribute_speak=0.000
he-is?         previous_word_is_disobey=0.000   object_attribute_fly=0.986  object_attribute_speak=0.000

Layer 2 has two heads: get flight attribute and get speech attribute.

Usually layers have a similar topology, so two heads are used in both layers.

Could this work be done with a single head, like the action head in the previous notebook? Not with this residual layout. In the previous notebook, one head attended to both action words, and that worked because each action word's own swap bit (swap_fly or swap_speak) says which attribute it swaps. Here, the disobey information sits in one shared slot, previous_word_is_disobey, on both action words. A single head attending to both action words would add the two disobey signals into the same number, and they could no longer be told apart. For example, Crow disobeys keep-flight swap-speech he-is? and Crow keep-flight disobeys swap-speech he-is? would give exactly the same head output, but the answers are Human and Crow. So we use one head for the flight word and one head for the speech word.

Wq_fly_head = torch.zeros(num_bits, 1); Wq_fly_head[idx["question"], 0] = 8
Wk_fly_head = torch.zeros(num_bits, 1); Wk_fly_head[idx["action_fly"], 0] = 1
Wv_fly_head = torch.zeros(num_bits, 2); Wv_fly_head[idx["swap_fly"]] = torch.tensor([1.0, 0.0]); Wv_fly_head[previous_word_is_disobey] = torch.tensor([0.0, 1.0])

Wq_speak_head = torch.zeros(num_bits, 1); Wq_speak_head[idx["question"], 0] = 8
Wk_speak_head = torch.zeros(num_bits, 1); Wk_speak_head[idx["action_speak"], 0] = 1
Wv_speak_head = torch.zeros(num_bits, 2); Wv_speak_head[idx["swap_speak"]] = torch.tensor([1.0, 0.0]); Wv_speak_head[previous_word_is_disobey] = torch.tensor([0.0, 1.0])

Wo_2 = torch.zeros(4, num_bits)
Wo_2[0, is_swap_attr_fly]        = 1
Wo_2[1, is_attr_fly_disobeyed]   = 1
Wo_2[2, is_swap_attr_speak]      = 1
Wo_2[3, is_attr_speak_disobeyed] = 1

print("Wv_fly_head nonzero rows:  ", (Wv_fly_head != 0).any(1).nonzero().flatten())
print("Wv_speak_head nonzero rows:", (Wv_speak_head != 0).any(1).nonzero().flatten())
print("Wo_2 nonzero cols:         ", (Wo_2 != 0).any(0).nonzero().flatten())
Wv_fly_head nonzero rows:   tensor([ 2, 15])
Wv_speak_head nonzero rows: tensor([ 3, 15])
Wo_2 nonzero cols:          tensor([18, 19, 20, 21])
python
def layer2_attn(X1):
    oF, AF = head(X1, Wq_fly_head, Wk_fly_head, Wv_fly_head)
    oS, AS = head(X1, Wq_speak_head, Wk_speak_head, Wv_speak_head)
    X2 = X1 + torch.cat([oF, oS], dim=1) @ Wo_2
    return X2, AF, AS

X2, AF, AS = layer2_attn(X1)
q = sentence.index("he-is?")

print("A_F[q]", AF[q])
print("A_S[q]", AS[q])
print()
print("is_swap_attr_fly       ", X2[q, is_swap_attr_fly])
print("is_attr_fly_disobeyed  ", X2[q, is_attr_fly_disobeyed])
print("is_swap_attr_speak     ", X2[q, is_swap_attr_speak])
print("is_attr_speak_disobeyed", X2[q, is_attr_speak_disobeyed])
A_F[q] tensor([0.00, 0.00, 1.00, 0.00, 0.00])
A_S[q] tensor([0.00, 0.00, 0.00, 1.00, 0.00])

is_swap_attr_fly        tensor(0.)
is_attr_fly_disobeyed   tensor(1.00)
is_swap_attr_speak      tensor(1.00)
is_attr_speak_disobeyed tensor(0.00)

The FFN is shown for completeness, also as a hand-constructed implementation. But for understanding the ideas of attention, position embedding, residuals and the need for layers, this part can be skipped.

Summary: bias values are used carefully to allow distinguishing 0, 1, 2, 3. This provides different ReLU activation levels corresponding to the number of flips. This FFN is a parity finder, while the FFN used in the previous notebook without the disobeys word was an XOR gate. Repeating the note that, in the real implementation, all the weights and biases are learned using gradient descent and the backpropagation algorithm.

W1 = torch.zeros(num_bits, 8)
b1 = torch.zeros(8)

for r in (object_attribute_fly, is_swap_attr_fly, is_attr_fly_disobeyed):
    W1[r, 0:4] = 1
for r in (object_attribute_speak, is_swap_attr_speak, is_attr_speak_disobeyed):
    W1[r, 4:8] = 1

b1[0:4] = torch.tensor([0.0, -1.0, -2.0, -3.0])
b1[4:8] = torch.tensor([0.0, -1.0, -2.0, -3.0])

W2 = torch.zeros(8, 2)
W2[0:4, 0] = torch.tensor([1.0, -2.0, 2.0, -2.0])
W2[4:8, 1] = torch.tensor([1.0, -2.0, 2.0, -2.0])

print("W1.T\n", W1.T)
print("\nb1", b1)
print("\nW2.T\n", W2.T)
print("\npi:", [float(torch.relu(torch.tensor([s, s - 1.0, s - 2.0, s - 3.0])) @ W2[0:4, 0]) for s in range(4)])
W1.T
 tensor([[0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 1., 1., 0., 0.],
        [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 1., 1., 0., 0.],
        [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 1., 1., 0., 0.],
        [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 1., 1., 0., 0.],
        [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 0., 1., 1.],
        [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 0., 1., 1.],
        [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 0., 1., 1.],
        [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 0., 1., 1.]])

b1 tensor([ 0., -1., -2., -3.,  0., -1., -2., -3.])

W2.T
 tensor([[ 1., -2.,  2., -2.,  0.,  0.,  0.,  0.],
        [ 0.,  0.,  0.,  0.,  1., -2.,  2., -2.]])

pi: [0.0, 1.0, 0.0, 1.0]
python
def forward(sentence):
    X = embed(sentence)
    X1, AO, AP = layer1(X)
    X2, AF, AS = layer2_attn(X1)
    H = torch.relu(X2 @ W1 + b1)
    Y = H @ W2
    return dict(X=X, X1=X1, X2=X2, Y=Y, AO=AO, AP=AP, AF=AF, AS=AS)

word_to_attributes = {
    "Rock":              (0, 0, 0),
    "Human":             (0, 0, 1),
    "Car":               (0, 1, 0),
    "Talking Tow Truck": (0, 1, 1),
    "Crow":              (1, 0, 0),
    "Flying Superhero":  (1, 0, 1),
    "Plane":             (1, 1, 0),
    "Talking Planes":    (1, 1, 1),
}
attributes_to_word = {b: w for w, b in word_to_attributes.items()}

def readout(sentence):
    r = forward(sentence)
    q = sentence.index("he-is?")
    fly, speak = (int(v.round()) for v in r["Y"][q])
    return attributes_to_word[(fly, 0, speak)], r["Y"][q]
sentences = [
    ["Human", "disobeys", "keep-flight", "disobeys", "swap-speech", "he-is?"],
    ["Crow", "disobeys", "keep-flight", "swap-speech", "he-is?"],
    ["Crow", "keep-flight", "disobeys", "swap-speech", "he-is?"],
    ["Crow", "disobeys", "keep-flight", "disobeys", "swap-speech", "he-is?"],
]

for s in sentences:
    word, y = readout(s)
    print(f"{' '.join(s):58} {y}  ->  {word}")
Human disobeys keep-flight disobeys swap-speech he-is?     tensor([1.00, 0.98])  ->  Flying Superhero
Crow disobeys keep-flight swap-speech he-is?               tensor([0.02, 1.00])  ->  Human
Crow keep-flight disobeys swap-speech he-is?               tensor([0.99, 0.00])  ->  Crow
Crow disobeys keep-flight disobeys swap-speech he-is?      tensor([0.02, 0.00])  ->  Rock
a = ["Crow", "disobeys", "keep-flight", "swap-speech", "he-is?"]
b = ["Crow", "keep-flight", "disobeys", "swap-speech", "he-is?"]

print(sorted(a) == sorted(b))
print(readout(a)[0])
print(readout(b)[0])
True
Human
Crow
for s in sentences:
    r = forward(s)
    print(" ".join(s))
    print("  A_P")
    for tok, row in zip(s, r["AP"]):
        print(f"    {tok:14} {row}")
    print("  previous_word_is_disobey", r["X1"][:, previous_word_is_disobey])
    q = s.index("he-is?")
    print("  gathered  ", r["X2"][q, [is_swap_attr_fly, is_attr_fly_disobeyed,
                                      is_swap_attr_speak, is_attr_speak_disobeyed]])
    print("  Y         ", r["Y"][q])
    print()
Human disobeys keep-flight disobeys swap-speech he-is?
  A_P
    Human          tensor([0.17, 0.17, 0.17, 0.17, 0.17, 0.17])
    disobeys       tensor([1.00, 0.00, 0.00, 0.00, 0.00, 0.00])
    keep-flight    tensor([0.00, 1.00, 0.00, 0.00, 0.00, 0.00])
    disobeys       tensor([0.00, 0.00, 1.00, 0.00, 0.00, 0.00])
    swap-speech    tensor([0.00, 0.00, 0.00, 1.00, 0.00, 0.00])
    he-is?         tensor([0.00, 0.00, 0.00, 0.00, 1.00, 0.00])
  previous_word_is_disobey tensor([0.33, 0.00, 1.00, 0.00, 1.00, 0.00])
  gathered   tensor([0.00, 1.00, 1.00, 1.00])
  Y          tensor([1.00, 0.98])

Crow disobeys keep-flight swap-speech he-is?
  A_P
    Crow           tensor([0.20, 0.20, 0.20, 0.20, 0.20])
    disobeys       tensor([1.00, 0.00, 0.00, 0.00, 0.00])
    keep-flight    tensor([0.00, 1.00, 0.00, 0.00, 0.00])
    swap-speech    tensor([0.00, 0.00, 1.00, 0.00, 0.00])
    he-is?         tensor([0.00, 0.00, 0.00, 1.00, 0.00])
  previous_word_is_disobey tensor([0.20, 0.00, 1.00, 0.00, 0.00])
  gathered   tensor([0.00, 1.00, 1.00, 0.00])
  Y          tensor([0.02, 1.00])

Crow keep-flight disobeys swap-speech he-is?
  A_P
    Crow           tensor([0.20, 0.20, 0.20, 0.20, 0.20])
    keep-flight    tensor([1.00, 0.00, 0.00, 0.00, 0.00])
    disobeys       tensor([0.00, 1.00, 0.00, 0.00, 0.00])
    swap-speech    tensor([0.00, 0.00, 1.00, 0.00, 0.00])
    he-is?         tensor([0.00, 0.00, 0.00, 1.00, 0.00])
  previous_word_is_disobey tensor([0.20, 0.00, 0.00, 1.00, 0.00])
  gathered   tensor([0.00, 0.00, 1.00, 1.00])
  Y          tensor([0.99, 0.00])

Crow disobeys keep-flight disobeys swap-speech he-is?
  A_P
    Crow           tensor([0.17, 0.17, 0.17, 0.17, 0.17, 0.17])
    disobeys       tensor([1.00, 0.00, 0.00, 0.00, 0.00, 0.00])
    keep-flight    tensor([0.00, 1.00, 0.00, 0.00, 0.00, 0.00])
    disobeys       tensor([0.00, 0.00, 1.00, 0.00, 0.00, 0.00])
    swap-speech    tensor([0.00, 0.00, 0.00, 1.00, 0.00, 0.00])
    he-is?         tensor([0.00, 0.00, 0.00, 0.00, 1.00, 0.00])
  previous_word_is_disobey tensor([0.33, 0.00, 1.00, 0.00, 1.00, 0.00])
  gathered   tensor([0.00, 1.00, 1.00, 1.00])
  Y          tensor([0.02, 0.00])

This notebook is part of the support files for the TechAarvam workshop

Build Your Own Model.

© TechAarvam. You are free to use, copy, modify, share and build on this

material, including for commercial purposes, provided you credit TechAarvam and link back to