Attention mechanism: Why apply multiple different transformations to obtain query, key, value

Question

I have two questions about the structure of attention modules:

Since I work with imagery I will be talking about using convolutions on feature maps in order to obtain attention maps.

1. If we have a set of feature maps with dimensions [B, C, H, W] (batch, channel, height, width), why do we transform our feature maps before we calculate their affinity/correlation in attention mechanisms? What makes this better than simply taking the cosine distance between the feature vectors (e.g. resizing the maps to [B, C, HW] and [B, HW, C] and multiplying them together). Aren't the feature maps already in an appropriate feature/embedding space that we can just use them directly instead of transforming them first?

2. Most of the time, attention mechanisms will take as input some stack of feature maps (F), and will apply 3 transformations on them to essentially produce a "query", "key" and "value". The query and key will be multiplied together to get the affinity/correlation between a given feature vector and all other feature vectors. In computer vision these transformation will typically be performed by the different 1x1 convolutions. My question is, how come we use 3 different 1x1 convolutions? Wouldn't it make more sense to apply the same 1x1 convolution to the input F? My intuition tells me that since we want to transform/project the feature maps F into some embedding/feature space that it would make the most sense if the "query", "key" and "value" were all obtained by using the same transformation. To illustrate what I mean lets pretend we had a 1x1 feature map and we wanted to see how well the pixel correlates with itself. Obviously it should correlate 100% because it is the same pixel. But wouldn't applying two sets of 1x1 convs to the pixel lead to the chance that the pixel would undergo a different transformation and in the end would have a lower correlation than it should?

Answer 1

I assume you're talking about this design: (image source)

But wouldn't applying two sets of 1x1 convs to the pixel lead to the chance that the pixel would undergo a different transformation and in the end would have a lower correlation than it should?

Yes, that's the point. We are not trying to measure a pixel's correlation with itself. Rather we are trying to allow it to query different related data. We are giving it freedom to change both the data and the queries.

It is true that the space for the queries and the keys is the same - but we shouldn't use the same transformation for both, or else each instance of the attention layer is just trying to fetch its own value! Generally the purpose of an attention layer is to query different parts of the input.

The first half of your question was essentially "why do we have a convolution at all?" and I think this has the same answer: you'd just be able to detect similar pixels, you wouldn't be able to pay attention to noses whenever eyes are detected.

It is also true that you could probably skip the convolution on the h(x) input. It looks like this one is somewhat redundant because the convolutions on h(x) and v(x) apply in series - which makes it a two-layer convolution, not quite the same as a one-layer convolution, but perhaps only one layer is needed.

It is possible that if you removed the conv layer on either the keys or the queries (but not both) the model would learn to generate the keys directly as the features, but this would hinder it because it would be unable to output any data in the values and queries that wasn't part of the keys (or vice versa). Seems silly. Don't do that.