I vaguely understand the mechanism and what it’s trying to do around building a stronger contextual representation of the current input with respect to the other supplied inputs but I struggled with understanding how the underlying operations lead to this effect.
Let’s work through the idea.
Let’s say we have a sequence of length of vectors of dimension , each vector acting as a representation of some kind for our input.
Pack them into a matrix with dimension . Each row containing a vector. When I first read about this I was confused about why we need to pack them into a matrix. The answer is that if we need to pack in context about other vectors in the sequence, we need to able to reason about them in the same place.
Let’s project them into 3 matrices, , and ,
We get,
The letter names are arbitrary for now, it should ideally make sense after we work through what the mechanism is doing.
We have,
Let’s dive into what’s happening. Why ?
We have,
This gives us a matrix that shows us how aligned each and vector is to each other. One way to look at a dot product is as a scalar measure of how directionally aligned each vector is.
The softmax operation is then done row-wise. Each row now consists of numbers that sum up to 1, which gives us a weight. The most ‘aligned’ pair will be the closest to 1.
We take these vectors and apply them to our third linear projection ,
The resultant matrix consists of vectors that are a weighted sum of the vectors in , each weight between 0 to 1.
Now, that’s the base operation. How does this become a contextual representation of the sequence?
If you trace how the gradients flow for an objective. Let’s say that this was part of a language model and the goal was to predict the next token autoregressively.
The matrices don’t really do anything when initialized. But they start learning projections that suit their position in the computation.
and are responsible for finding the right coefficients for the weighted sum. I struggled a lot on what the role of seemed to be. Why not just use directly and you’d get a weighted sum of the tokens to pay attention to. is simply a linear transformation of the tokens so that the model gets to make a linear transformation best suited for the current layer. Functionally, due to the way matrix multiplication works, is nothing but , in other words, really is a learned linear projection that’s transforming the representation to best suit the task at hand.
and get a positive signal from the loss function when they correctly assign the contributions of each token in the final weighted sum. gets a positive signal from the loss function when it transforms the result into a representation more useful for subsequent computations.
Over the course of training, these matrices are jointly optimized. , learn contributions should be weighted. learns to represent the information after it’s combined.