by imtringued | 2025-05-12 09:32:28+00:00
I personally don't really see the point in giving meaning to the Q, K, V parts. It doesn't actually matter what Q, K, V do, it's the training algorithms' job to assign it a role automatically. It makes more sense to think of it as modeling capacity or representational space.

One of the biggest things people don't understand about machine learning is that there is a lot of information in the model that is only relevant to the training phase. It's similar to having test points for your probes on a production PCB or trace/debug logging that is disabled in production. This means that you could come up with an explanation that makes sense at training time, but is actually completely irrelevant at inference time.

For example, what you really want from attention is the pairing of all vectors in Q with all vectors in K. Why? Not necessarily because you need it for inference. It's so that you don't have to know or predict where the gradient will propagate in advance when designing your architecture. There are a lot of sparse attention variants that only really apply to inference time. They show you that transformers are doing a lot of redundant and unnecessary work, but you can only really know that after you're done training.

There is a pattern in LSTMs and Mamba models that is called gating [0], which in my opinion is a huge misnomer. Gating implies that the gate selectively stops the flow of information. What they really mean by it is element-wise multiplication.

If you look at this from the concept of model capacity, then what additional concepts and things does this multiplication allow you to represent? LSTMs are really good at modelling dynamical systems. Why is that the case? It's actually quite simple. Given a discretized linear dynamical system x_next = Ax + Bu, you run into a problem: You can only model linear systems.

So, assuming we only had matrix multiplications with activation functions, we would be limited to modeling linear dynamical systems. The key problem is that the matrices in the neural network can model A and B of the dynamical system, but they cannot have a time varying A and B. You could add additional parameters to x that contain the time varying parameters of A, but you will not be able to use these parameters to model non-linearity effectively.

Your linearization of a function f(x) might be written as g(x) = f(x_0) + f'(x_0)x = a_0 + a_1 * x. The problem is very apparent, while you can add an additional parameter to modify a_0, you can never modify a_1, since a_1 is baked into the matrix and multiplied with your x. What you want is a function like this h(x) = f(x_0) + f'(x_0)x = (a_0+m_0) + (a_1+m_1) * x, where m is a modifier value in x. In other words, we need the model to represent the multiplication m_1 * x and it turns out that this is exactly what the gating mechanism in LSTM models does.

This might look like a small victory, but it is actually enough of a victory to essentially model most non-linear behavior. You can now model the derivative in the hidden state, which also means you can model the derivative of the derivative, or the derivative of the derivative of the derivative. Of course it's still going to be an approximation, but a pretty good one if you ask me.

[0] https://en.wikipedia.org/wiki/Gating_mechanism