Revealing example of self-attention, the building block of transformer AI models

No one here actually checks anything before upvoting it. Having said that, this repo seems worthwhile and I appreciate the effort of OP. It's open source which empowers people to fix such problems themselves and make a PR, or at least an issue.

The discussions are sometimes more interesting than the submission itself.

I made the same mistake recently! It doesn’t help that there are so many poorly documented repos that make similar mistakes (and worse yet, half baked medium articles that blatantly give false information). Directly translating papers to PyTorch is hard too, I really wish there was a better API to express these models in that aligns better with how they’re actually documented in papers.

Now that you have updated the code I think you have just flipped the error, it's still there, just doing the opposite (still wrong) thing: for example, if "num_steps" is 768 and "num_inputs" is 1 then K, Q, V is 1x1, okay but the the attention matrix cannot also be 1x1, it needs to be 768x768; you can also do the opposite K, Q, V being 768x768 and the attention matrix 1x1, before the error was that you made both the K, Q, V and the attention matrix 768x768 by applying the matrices on the wrong dimension, now the error is that you have made the the K, Q, V and the attention matrix 1x1.

For a simple fix just make the attention matrix W [num_steps, num_steps] instead of [num_inputs, num_inputs] then "y_i = torch.matmul(x_v_i, w_i)" would be "y_i = torch.matmul(w_i, x_v_i)" and should be correct.

TBH I don't find this exposition enlightening. If you don't already know the pytorch API, the code will be pretty obscure. And if you do understand the API, it's still not very clarifying.

Instead, it's a lot easier IMO to understand self attention in just plain English. If you understand the words, and also understand the torch API (or tensorflow or numpy, they're all more or less fungible), then it should also be straightforward to implement it.

Here's how I explain it to my grad students:

You have a collection of N things, with some fixed number of features each (doesn't matter how many, but call it d_f). Call that collection {x_i}. You have two learnable matrices Q and K, which learn to project those items into some collection of "questions" and "answers", respectively. Doesn't really matter what those questions or answers are, the NN will figure out what they should be during training.

These projection matrices map x_i -> k_i, and x_i -> q_i. k and q are each d_k-dimensional vectors (d_k is a number you choose) of features, so Q and K are d_k-by-d_f matrices. To measure the "compatibility" between questions q and answers k, we take the dot product of them. Ones which are similar (large, positive dot products) means the NN likes the answer k for question q.

So, here's what you do. For each token i, you compare it's "question" against every other token j's answer. I.e. you compute dot(q_i, k_j), which itself can be considered an NxN matrix of scalar numbers, qk_{ij}. Each element of this matrix contains the answer to the question "how well did item j answer the question of item i?"

Applying softmax to this matrix just converts the dot product values into a score from 0-1 over all the items. I.e., you convert the question to "on a scale from 0-1, how well did item j answer the question of item i?". The only subtlety here is that all the scores over all items j have to sum to 1 for each question i.

Finally, we have a third projection matrix, V, which maps x_j -> v_j. The v_j are d_v-dimensional vectors (doesn't have to equal d_k, and again something you arbitrarily choose). These represent the values that the NN would like to pass on from item j to any other item which decides it "likes" the answer of item j.

So now, for each item i we have a score for "how well does item j answer my question i?". And we also have the list of values v_j that each item j would like to pass to item i. So, to compute the new information to add to item i, we take the weighted sum of the values v_j that item i liked, using these compatibility scores as a weight. In math, this reads: y_i = sum_over_j [ softmax(qk_ij) v_j ]

* Technically it's also empirically found that normalizing the argument of softmax by a factor 1/sqrt(d_k) helps. But it's not really germane to understanding and also not strictly necessary since a NN is of course free to simply learn to include an overall factor proportional to 1/sqrt(d_k) in the parameters of one of the projection matrices. Actually the use of softmax at all is also arbitrary, but encourages the attention to be sparse, which again is empirically known to be helpful.