Puzzle Solution Revealed - Transformer: Need for Position Embedding

A developer has published a solution to a puzzle from a workshop on building a minimal hand-constructed transformer, demonstrating why position embeddings and residual connections are necessary. The extended model adds a 'disobeys' token that modifies only the immediately following word, forcing word order to matter, and introduces a 22-bit residual stream carrying position one-hot encodings plus each layer's findings. The author notes that while real implementations learn vector and position embeddings via gradient descent, the hand-constructed example clarifies why those blocks and connections exist.

▶ Watch the 7-minute walkthrough https://youtu.be/Hav2zlqxzUI python import torch torch.set printoptions precision=2, sci mode=False, linewidth=160 In this notebook we are extending the previous notebook, attention ann.ipynb https://colab.research.google.com/github/techaarvam/byom workshop/blob/main/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 | slot | |---|---| | 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 | The new list of token bits, with the additional word 'disobeys' idx = {"fly": 0, "speak": 1, "swap fly": 2, "swap speak": 3, "object": 4, "action fly": 5, "action speak": 6, "question": 7, "disobey": 8} constants used for bit-slicing and locating the portion of the residual we need L is the maximum sentence length, i.e. the number of position slots. It is used to name the position slots below, and to build the "Disobey Position Finder" head Wq P, Wk P in layer 1. 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 = { fly spk swF swS obj actF actS q dis "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 https://www.techaarvam.com/workshops/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