The attention mechanism, specifically the confusing bits about the Q/K/V layers, are all single linear layers with bias followed by a nonlinearity. Each of Q, K

The attention mechanism, specifically the confusing bits about the Q/K/V layers, are all single linear layers with bias followed by a nonlinearity. Each of Q, K, and V are just square matrices with edge length the same as the “dimension” of the network (which you usually see denoted as “d”). So if you have an input sentence of 10 words, it gets embedded into a 10 x d space by looking up each word’s d-dimensional representation (from a |V| x d matrix). Add in some static position-dependent embedding, which is usually a function of the time index, so that the model can use position information and call the result E_0. Multiply E_0 by your d x d matrix of keys K and set aside as K_0. Do the same for your d x d query matrix Q and call the result Q_0. (The result in both cases is shape 10 x d). Take (Q • K’) which results in a 10 x 10 matrix where entry i,j denotes the amount that vector i attends to vector j. Squash this with a softmax along the time dimension* and set it aside as “attentions”. Create V_0 by running your inputs E_0 through the d x d layer of values V. Finally, calculate the output of the attention layer as (attentions) * V_0. This is just a linear cimbination of the values matrix V_0, where the coefficients are the softmaxed attention values from earlier. This is exactly what’s going on in Eq 1 from the “Attention is all you need” paper. A transformer block is usually the above operation, and then it immediately runs the result through a high-dimensional MLP ending with a layer with dimension d. And it will usually use multi-head attention which splits the input d-dimensional representation into m splits along the d axis (so that you end up with m 10 x m/d vectors) and then does the above calculation on each of the m inputs using layer- and head-specific QKV matrices, and then recombines all the results along the d axis to recover the 10 x d representation. It can be confusing when you read code that does this because it’s all vectorized and the splitting and head-specific matrices are implicit. * after dividing by the square root of d. this reduces absolute value of the inputs to the softmax and mitigates an issue that softmax has with gradients on large inputs becoming 0 on limited precision floating points.