Attention Mechanisms

Attention Explained

The attention mechanism lets models focus on the most relevant parts of an input sequence, which is especially important in sequence-to-sequence tasks. It works by comparing a query to a set of keys and assigning weights to the corresponding values. The output is a weighted sum of values, where higher weights go to elements most relevant to the query. A softmax turns these weights into a probability distribution, emphasizing the most important elements. In self-attention, where queries, keys, and values all come from the same sequence, the model can capture long-range dependencies and train in parallel.

Step 1 - Linear Transformations: seq of embeddings passed through transformer’s input layers (encoder and decoder) and each embedding experiences three separate linear transformations (by multiplying it by 3 weight matrices) resulting in three vectors — query, key, and value. The proper weights are learned through training.

Each weight matrix is initialized randomly => resultant vectors each learn some different information about the embedding. This is important when calculating the attention score since we don’t want to just derive the dot product of the vector with itself.

Once we have these 3 vectors for each embedding in seq, we can calculate attention score which measures the strength of relationship between a word in the seq with all other words:

1. Take the query vector for a word and calculate it’s dot product with the transpose of the key vector of each word in the sequence — including itself. This is the attention score / weight.

2. Then divide each of the results by square root of the dimension of key vector. This is the scaled attention score.

3. Pass them through a softmax function, so that values are contained between 0 and 1.

4. Take each value vectors and calculate the dot product with output of softmax f(x).

5. Add all weighted value vectors together.

Notice in the figure below that we are doing matrix operations on seq_length x embedding_size matrices. This shows a toy example of two words with an embedding size of 3.

	SELF-ATTENTION: Encoder - each token attends to all tokens in seq. Decoder - each token attends only to tokens on the LEFT - future tokens are masked (to avoid cheating during next token prediction). CROSS ATTENTION: Encoder - keys and values Decoder - queries

The dot product results in a seq_length x seq_length matrix. Think correlation matrix where the relationships between any two members can be found by their intersections. In this case, the members are words and their intersections are attention scores.

Multiplying value matrix by attention matrix results once again in a seq_length x embedding_size matrix. This matrix holds the contextual information for each embedding.

Self-Attention Encoder-Decoder Cross Attention

Another way to visualize self-attention:

SELF-ATTENTION: Encoder - each token attends to all tokens in seq. Decoder - each token attends only to tokens on the LEFT - future tokens are masked (to avoid cheating during next token prediction). CROSS ATTENTION: Encoder - keys and values Decoder - queries