Focus Unfolded: A Deep Dive into Attention Mechanisms in AI

As a person with Attention Deficit Hyperactivity Disorder (ADHD), reading a long paragraph can be quite challenging. I might start off strong, but as the paragraph goes on, I find it increasingly difficult to maintain my focus and understand the overall meaning. My mind tends to wander, and I often have to go back and re-read sentences multiple times to grasp the information fully. This can make reading comprehension tasks frustrating and time-consuming, same thing can be found in AI models.

In this article we are going to explain how we can make models pay attention.

Why Attention ?

Let’s imagine that I want to translate the following sentence “The match ignited the firewood, filling the room with warmth.” from English to Arabic, but I can only remember one word at a time what is the expected output?

Sure, it will be “المبارة أشعلت الحطب، يملئ الغرفة بالدفئ” which is a completely wrong translation as I thought of “match” like a “football match” because I didn’t pay attention to the other parts of the word like “Ignited” which will change my translation to be “"أشغل عود الثقاب الحطب، مالئاً الغرفة بالدفئ” .

One can say that we can solve this problem using Long Short Term Memory (LSTM) as it contains long memory for very far away words and short memory for recent ones, Okay lets break it down and see how Long memory and short memory are calculated.

First we have an input I which represents a sequence of T time-steps, C as Long Memory and H as Short memory, and we have 3 gates which is forget gate that determine how much to remember from the long memory, input gate which updates the long term memory with respect to current input, and finally the output gate which updates the short term memory.

Forget Gate:
1. as mentioned before forget gate uses the input I at time-step t, and the short term memory H from time-step t-1 to determine how much to forget from the long term memory as follows:
  $RememberRate = Sigmoid([I_{t} * W_{i}] + [H_{t-1} * W_{h}])$
  $C_{t} = C_{t-1} * RememberRate$
  - Note that bias term was ignored for simplification and W is different for each gate.
Input Gate:
1. Input gate is simply meant to update the long term memory using the previous short term memory and current input as follows:
  $PotentialRate = Sigmoid([I_{t} * W_{t}] + [H_{t-1} * W_{h}])$
  $PotentialLongMemory = Tanh([I_{t} * W_{i}] + [H_{t-1} * W_{h}])$
  $C_{t} = C_{t} + (PotentialRate * PotentialLongMemory)$
Output Gate:
1. Output gate is the last gate that meant to return the short term memory H as output after update it using the long term memory as follows:
  $PotentialRate = Sigmoid([I_{t} * W_{t}] + [H_{t-1} * W_{h}])$
  $H_{t} = Tanh(C_{t}) * PotentialRate$

As you can see we have some problems:

At any point in time if we forget too much of long term memory we can not restore it.
It is sequential which means that only the 2nd word knows the 1st but the 1st will never know that there is a 2nd word, which means that in our example “match” can’t be affected by “Ignited”

So, How to solve this problem ?!

Attention

What if we have a linear connection between every pair of words?! this means that “match” will have a direct connection with each word in the sentence? and by training the network will learn how much the word “match” is affected by each of the other words, but this will create a fixed connections which means that in our example the network will learn that the 2nd word “match” will be affected only by the 3rd word “ignited” even if we changed the example.

What if we have a different input “Both teams were ready to play the match” the location of “match” changed and the meaning too, can we have some sort of connection that changes when input itself changes?

This is what attention is all about, “allows a network to give different weights to different inputs, with weighting coefficient that themselves depend on the input values.”1

Suppose that we have an input X of length N and we want to transform it to another space Y with same length N but with richer information then, each value of Y equals:

y_{n} = \sum_{m=1}^{N}{a_{nm} * x_{m}}

which means that the single value in Y will depend on all values of X by factor of a associated to each value of y and x, but what if a takes too large or too small values resulting in more weight on only one input? we want a to be close to 0 if x have a little influence in y and to be close to 1 if it have a large influence which means:

0 \le a_{nm} \le 1

and to ensure that when we pay more attention to one input we pay less attention to another it must follow:

\sum_{m=1}^{N}{a_{nm}} = 1

Until now we discovered that we are going to make a connection between every pair of the input data, and that connection will depend on the input value, but how to calculate it ?! how to find this a matrix ?

In typical NLP situation each word will be converted to set of numerical features before being passed to the network, this process is called “Embedding”.

So one way to find the value of a is measuring the similarity between words in each pair then normalize the similarity matrix (using softmax) to be between 0 and 1.

One way to calculate the similarity between n-th input and m-th input is dot product

a_{nm} = \frac{e^{x_n^Tx_m}}{\sum_{m=1}^{N}e^{x_n^Tx_m}} = softmax(x_{n}^Tx_{m})

to simplify things lets use the matrix notation:

Y = Softmax(XX^T) X

This process is called “self-attention” or “dot-product self attention” because we are using the same sequence to determine the value of attention.

Until now I tried to explain self attention without the need to mention the information retrieval field, but to explain more we will need to be aware of some terminology of the field.

Figure 3 from “Focus Unfolded: A Deep Dive into Attention Mechanisms in AI”

Consider that you own a simple library shop online. you have many books and you want to recommend books for your clients based on the books they are reading, you started by clustering the books into categories where each book can have number of attributes (features) based on its cluster

for example:

Book 1:
- [language: en, publish date: 2023, auther: Auther1, genre: technology]
Book 2:
- [language: fr, publish date: 1998, auther: Auther2, genre: romance]

In this scenario Book 1 and Book 2 are called “Values”, the set of attributes for each book is called “Key”

Let’s say that one of the users read the following book:

Book 3:
- [language: en, publish date: 2019, auther: Auther3, genre: technology]

In this case the set of attributes of the book read by the user is called “Query”.

Simply, your system should recommend a Value (V) to the user that reduces the distance between the Query (Q) and the Key (K), this situation is called “Hard Attention” what we used earlier is called “Soft Attention”

Soft Attention in which we use continuous variables to measure the degree of match between queries and keys and we then use these variables to weight the influence of the value vectors on the outputs.2

Using what we mentioned in information retrieval we can think of the n-th input x_n as value vector will be used to create the output, and also as the key which represents the set of attributes of n-th input.

Why this? because each input word was converted to a set of features that represents its attributes, so x is the input value and also the attributes.

while m-th input x_m is the query that we want to measure the distance between it and the key x_n As we did earlier when we calculated the dot product.

So we simply can change the notation only in this equation to be more clear

from:

Y = Softmax(XX^T) X

To:

Y = Softmax(QK^T) V

Where (util now) Q = X, K = X and V = X

But util now the transformation from x to y is fixed and it isn’t learning anything from the data (no trainable parameters), and all values of x_n plays an equal role in determining the attention coefficients.

Parametric Self Attention

What if we added trainable parameters U (a neural network layer weights) to X before calculating the attention coefficients?

\hat{X} = XU

Y = Softmax(\hat{X}\hat{X}^T) \hat{X}

This now has more flexibility but it has a property that the matrix is symmetric which means that the attention coefficients are mutual, where we need our attention mechanism to support significant asymmetry.

For example, we might expect that ‘chisel’ should be strongly associated with ‘tool’ since every chisel is a tool, whereas ‘tool’ should only be weakly associated with ‘chisel’ because there are many other kinds of tools besides chisel.3

What if we can use different weights (parameters) for each values of Q, K, and V, we can create much more flexible model by allowing each one of them to have an independent parameters.

Q = XW^{(Q)}

K = XW^{(K)}

V = XW^{(V)}

Y = Softmax(QK^T) V

You can think of it like, we have an input X will pass it to neuron with weights Wq to get Q and will pass X to another neuron with weights Wk to get K, and will pass X to neuron with weights Wv to get V.

Figure 4 from “Focus Unfolded: A Deep Dive into Attention Mechanisms in AI”

After reaching this point we have only one modification to our attention mechanism, which is re-scaling, because the gradients of softmax function can be exponentially small for very large input values, to prevent this we can re-scale QK^t before applying the softmax function.

But how to find a suitable scaling factor ?!

Note that if all the elements of the query and key vectors are all independent random numbers with zero mean and unit variance, then the variance of the dot product would be Dk (the size of the key vector). 4

As an explanation, when the elements of these vectors are independent random numbers with zero mean and unit variance, it means that each element in the vectors is randomly distributed around zero with a spread of one, and they are statistically uncorrelated with each other.

In this scenario, when computing the dot product between the query and key vectors, each element in the query vector is multiplied by the corresponding element in the key vector, and then all these products are summed up. Since the elements are independent, the sum of their products behaves according to the properties of the variance.

In particular, if there are Dk elements in each of the query and key vectors (where Dk is the dimensionality of the key vectors), then according to the properties of the variance, the variance of the dot product would be Dk. This is because the variance of the sum of independent random variables is equal to the sum of their variances.

So if variance is Dk then standard deviation is the square root of Dk, and our attention takes the form:

Y = Softmax(\frac{QK^T}{\sqrt{D_k}}) V

Now we successfully finalized our first attention head that will capture data dependent patterns of input vectors.

But what if we had multiple patterns ?!

Multi-Head Attention

Here’s comes multi-head attention into play, it’s an identical copies of the attention head mentioned before stacked together (each one with different weights) and then transformed linearly with a wight for each attention head

So if

H_h = Attention(Q_h,K_h,V_h)

And we have H attention heads, then

Y(X) = concat[H_1, ..., H_H]W^{(o)}

Up to this point, we described attention in deep learning briefly, which is the key concept behind the Transformers which is leading the AI scene currently.

Finally, In my opinion many of deep learning algorithms have some sort of ADD (No hyperactivity) they are not that good in seeing the full picture, understanding the full text, and in this article, we try to address this problem through the attention mechanism.

If you still can't understand it fully, remember ..

“ATTENTION IS ALL YOU NEED.”.

Resources:

Deep Learning Foundations and concepts
StatQuest

Deep Learning Foundations and concepts, page 357

Deep Learning Foundations and concepts, pages 362-363

Deep Learning Foundations and concepts, page 364

Deep Learning Foundations and concepts, page 366