Transformers are a really fun topic in AI/DL. There are so many parts that come together to make this incredible technology.
Self-attention figures out how tokens in a sequence ‘work together’ or ‘communicate with each other’. During inference, a token ‘attends to’ previous tokens in a sequence. ‘Attending to tokens’ determines how much info each of the previous tokens and the current token itself contributes to a new representation of the current token.
Both training and inference attend to the current token and previous tokens. Training could look ahead to next tokens, since all tokens in a training sample are known in advance. However, training needs to learn the same task performed by inference, so attention always looks at the current token and previous tokens and never next tokens. A causal mask enforces this while training.
We’ll focus on inference for the rest of the discussion. When the current token is processed, there are three projections, Q K and V, which are Query, Key and Value. They are often described as follows:
Query is “What am I looking for?”
Key is “What information do I represent?”
Value is “Here’s the information I will provide.”
For current token’s vector t, the projections are calc’d like so:
Q = WQ × t
K = WK × t
V = WV × t
WQ, WK, and WV are weight matricies learned during training. Each token in a sequence computes a Q K and V. K and V are cached for reuse as subsequent tokens are processed.
The current token’s Q is compared to all K in the sequence of tokens to create scores.
Qn · K1, ..., Qn · Kn -> [score 1, ..., score n] = scores
Softmax is applied to the scores to derive attention weights.
softmax(scores) -> [weight 1, ..., weight n] = attentions weights
Attention weights are used to calc all V’s contribution to the current token’s new representation.
w1V1 + ... + wnVn -> [x1, ..., xm] = weighted sum vector
The weighted sum vector is the output representation of the currently processed token for a single head. The same steps are taken for every head in multi-headed attention.
A ‘head’ concentrates on token patterns and relationships. In language, words have all sorts of relationships, like subject -> verb, word proximities, syntactic patterns, and more. So it’s helpful to think of multi-headed attention as focusing on something analogous to those word patterns and relationships, though what each head has learned to do is not easily evident.
A transformer layer concats the output representations from all its heads and passes that through linear projection then a feed forward neural net. Transformer layers are stacked, each layer builds on previous layers, which helps build more complex token relationships and representations.
The stack’s last transformer layer outputs a final representation. The representation goes through different processing to to calc scores that are converted to probabilities for every token in the model’s vocabulary. There could be tens of thousands of tokens or more. One token is selected from the entire vocabulary and that becomes the next token generated in the sequence.