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'

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.