What does it mean order of input sequence does not matter for transformer self-attention head?

Question

The need for positional encoding in transformer models is justified by permutation invariance of self-attention heads, because, without it, transformer wouldn't have any mechanism to take into account the order of the words.

Suppose we trained a simple trigram transformer model without positional encoding to predict C as the next token, if the input is AB. Because the self-attention head output has T (time) dimension, we actually predicted BC, in other words, head.forward('AB')='BC'. We use only the last token C as the prediction of our model.

Now, because of the head permutation invariance, head.forward('BA')='CB'. Thus, the next token prediction becomes B (last token in output of 'forward`). So why is it said that transformer models without positional encoding are position agnostic. In the above example, we permuted input tokens and obtained a different prediction (B instead of C)

UPDATE: made a colab notebook illustrating the concept: https://colab.research.google.com/drive/1ItIQTCg3sVGRUIGbrh1pxTNlfLU5C7kq?usp=sharing

Answer 1

When the order of input sequence is changed, the positions of self attention values change (rows interchanged), but the values remain same (like you mentioned head.forward('AB')='BC' and head.forward('BA')='CB'. However, in sequential networks, in most cases the position also adds meaning and is crucial for capturing the relationships between different words. For example, consider these sentences, "A hit B" and "B hit A". In these sentences, the relation between the words is different, in the first one, A is hitting while in the second one A is being hit, but the self attention values will be same corresponding to each word. After the attention layer, there is no mechanism by which model can incorporate the sequence order or position value in calculating the output. The model will output the same values for both the sentences. That's why positional encoding is needed.
Order of input sequence does not matter for self attention heads -> This simply means that the self-attention value corresponding to each input word remains the same if its position is changed when positional encoding is not used.

Answer 2

+1 for a great reproducible code sample.

The answer to your main question can be seen by looking at the linear algebra, of what self-attention does. Go back to what matrix multiplication is doing, and you can see it is just summing a bunch of multiplies. The order of the tokens has to be lost.

So, safe in that knowledge, I then spent some time trying to work out why your script still works. I noticed the following.

With torch.manual_seed(42) I get as you reported:

A hit B predicting !
B hit A predicting ?

But with torch.manual_seed(77) I get:

A hit B predicting B
B hit A predicting hit

And with torch.manual_seed(88) I get:

A hit B predicting hit
B hit A predicting ?

To be honest I've not convinced myself, as quite a few other random seeds give !/? still. I also seem to get ? for the second one more often than ! for the first. I guess the proper thing to do would be to run the experiment for a large number of different seeds and report how random the results are.

If I reduce head_size to 2, so it is smaller than T, rather than larger, I get the desired result for a seed of 88, but every other seed I've tried fails.

I find that curious, rather than an explanation. If we believe that positional embeddings are not needed, a dim-2 transformer still ought to have enough capacity to memorize ("over-fit") two sentences of length 3 with a vocab size of 5? And if we believe positional embeddings are needed, the head size shouldn't matter?

Answer 3

Your code sample works because you are using two different tokens to do the prediction of "!" and "?". The model is trained to transform the token "A" to "?" and the token B to "!".

Now, suppose your dataset consists of "A B hit ?" and "B A hit !". Without positional encoding, the model won't be able to know when to transform the token "hit" to "!" and when to "?".