Intution behind Masked and Cross Attention

Need for masked self-attention

Before talking about masked self-attention, need to discuss autoregressive models concept. The transformer decoder is autoregressive at inference time and non-autoregressive at training time. Autoregressive models are a class of of models which generate data points in a sequence by conditioning each new point on the previously generated points. Example is the simple encoder decoder based NMT model.

So why is it producing a word at a time ? Like why not output all the words at a time (parallely), so we can't do that because the next word depends on previous words and thus we need autoregressive models as if input data was sequential then output data must also be sequential. Thus transformers are also autoregressive.

But then why do transformers (specifically its decoder block) behave differently during training and during inference ? Answer behind this masked self-attention.

Let's assume decoder to be auto-regressive model for the sake of this discussion. Say we have to build a MT (machine translation) system for english to hindi translation.

Inference will be done as in the figure below considering auto-regressive model, everything will be sequential.

Now, lets also see how training is carried out considering it to be auto-regressive

Now if this assumption of ours is correct then the problem will be that the training procedure will be slow.
How is it slow ? So here due to it being sequential, every operation in decoder executes 4 times for 3 word sentence which is fine, but now let's say we have a paragraph of 300 words then every operation executes 301 times which makes training very slow. Now for some dataset we can have multiple such paragraphs which will make overall training extremely slow.

Now at the inference time it needs to be auto-regressive / sequential, because output of one decoder state will be input to other decoder state.

But for training time it need not be auto-regressive / sequential beacuse we are using "teacher forcing" i.e we are not relying on the output of previous decoder state. So why not carry out all these steps in parallel with each other as there is no need for sequence / time-steps. With proper optimizations we actual transformer does it using parallel execution.

Problem with Parallelization

Now there's a problem with parallelization that we need to address, if calculate the contextual embedding of "आप" we use "कैसे" and "हैं", so we are in a way looking at future tokens to calculate this embedding, which is kind of cheating or formally known as "data leakage". During real-world application the model has only generated "आप" and the words "कैसे" and "हैं" do not exist yet and is mathematically impossible for the model to use them. This is an issue we encounter during inference. In machine learning, we can't carry out training in one way and inference in other way, and in training it will perform well and during inference performance will be worse due to this data leakage.

So what kind of problem are we facing ?

If we use auto-regressive training => no data leakage but slow training

If we use non auto-regressive training => training is faster but data leakage is there

Masked self-attention

Masked self-attention is just attention with extra steps, where we have an additional masking step. We don't want to include future token's contribution while calculating the contextual embedding so what if we mask them / cover them, mathematically make them zero. In order to implement this we add a masked matrix where weights corresponding to future tokens is "-infinity" and after addition and softmax the weights regading the future tokens becomes zero thus achieving what we wanted.

Cross Attention

So what really is cross attention ? It is a mechanism used in transformers, particularly in sequential tasks like machine translation / summarization. Allows model to focus on different parts of an input sequence when generating an output sequence. In short we are trying to find relation between two different sequences.

Self-Attention vs Cross-Attention

Conceptually they are very similar but they do have some differences :

1. The Input Difference

• Self Attention: Receives only one single sequence as input. For example, if you are processing the English sentence "We are friends", the module only looks at the embeddings for those three words.

• Cross Attention: Receives two different sequences as input. In a translation task, it receives both the input sequence from the encoder (e.g., English: "We are friends") AND the output sequence you are currently generating in the decoder (e.g., Hindi: "Hum dost hain").

2. The Processing Difference (Q, K, V Vectors)

This is the most critical mechanical difference in how the math works inside the block.

• Self Attention: The Query (Q), Key (K), and Value (V) vectors are all generated from the exact same sequence. Every word acts as a query looking for matches, and also acts as a key/value for other words to look at.

• Cross Attention: The vectors are split between the two sequences:

  • The Query (Q) vectors are generated ONLY from the Output Sequence (e.g., the Hindi words in the decoder).

  • The Key (K) and Value (V) vectors are generated ONLY from the Input Sequence (e.g., the English words from the encoder).

  • Analogy: The Hindi words are "asking the questions" (Queries), and the English words hold the "labels and answers" (Keys and Values).

  (Dot product happens between query (from output sequence) and key (from input sequence) )

3. The Output Difference

• Self Attention: Outputs a contextual embedding for every word that shows how it relates to other words in its own sentence. (e.g., How much does "friends" relate to "we"?).

• Cross Attention: Outputs contextual embeddings strictly for the words in the output sequence. These embeddings represent how much each generated word relates to the words in the original input sentence. (e.g., How much does the Hindi word "dost" relate to the English word "friends"?).