Attention Mechanism

Evolution Toward the Attention Mechanism: A Paradigm Shift in Sequence Modeling

Evolution Timeline of Neural Language Model

The evolution of conversational AI has progressed through several key stages. ELIZA (1966) was the first widely recognized chatbot, relying on rule-based pattern matching without learning or memory.

Recurrent Neural Networks (RNNs), introduced in the 1980s, enabled sequential processing and basic memory through hidden states, allowing models to capture contextual information from previous inputs.

Long Short-Term Memory (LSTM) networks, proposed in 1997, improved upon RNNs by introducing gated memory mechanisms that mitigated the vanishing gradient problem and enhanced long-range dependency modeling.

In 2014, the Attention Mechanism revolutionized sequence modeling by allowing models to selectively focus on the most relevant parts of the input, overcoming the context bottleneck of RNN-based architectures. This breakthrough paved the way for the Transformer architecture (2017), which replaced recurrence with self-attention, enabling parallel computation and superior context modeling.

Building upon Transformers, modern Large Language Models (LLMs) such as the GPT series have emerged, leveraging massive datasets and billions of parameters to achieve state-of-the-art performance in natural language understanding and generation.

Evolution

RNN Architecture

RNN

The figure illustrates the Encoder–Decoder architecture based on Recurrent Neural Networks (RNNs) for machine translation.

The encoder processes the input sequence word by word (“I”, “will”, “eat”), generating hidden states (h_1), (h_2), and (h_3), where each hidden state captures information from the current word and all preceding words.

The final hidden state (h_3) serves as a context vector, which is intended to summarize the meaning of the entire input sentence.

This context vector is then passed to the decoder, which generates the translated output sequence (“Je”, “vais”, “manger”) one token at a time through decoder states (s_1), (s_2), and (s_3).

However, because the entire input sequence must be compressed into a single context vector, the model suffers from a context bottleneck, making it difficult to retain information from long sentences.

This limitation motivated the development of the attention mechanism, which allows the decoder to access all encoder hidden states instead of relying solely on a single context vector.

Attention Mechanism

The first paper that introduce attention mechanism is Bahdanau Attention Mechanism. Link: https://arxiv.org/abs/1409.0473

They introduce idea of how much attention to give to the hidden state. They also proved that we can do sequence2sequence translation much better if we use the attention mechanism.

We can see in x-axis there is English and y-axis there is translated language. We can see attention is not always diagonal, model was able to detect below diagonal and understand economique comes before europeenne in translated language. European in x-axis is at position 5 and y-axis is 7. Translation is not always word by word.

Bahdanau Attention Mechanism

Bahdanau Attention Mechanism= RNN + Attention Mechanism

So, researcher proposed transformer architecture.

Original Transformer paper has encoder and decoder but diagram has only decoder only.

Self Attention

“Mechanism which allows each position in the input sequence to attend to all positions in same sequence.”

In RNN+ Attention , here the attention is between sequences. While decoding, we see how much attention we give to English sequences whereas self attention is predicting next token we don’t have different languages we just have bunch of data.

Here, in self attention we take same sequences or sentences. Let’s look at next in below diagram. How much attention should i give to all token in sentences.

If we take one word and see how important neighboring words. We need information about context of sentence i.e how different word relate with each other.

In figure below, look at word speak need to have maximum attention to Pune or India that information gonna help me to predict which language I speak. That’s why self attention mechanism become very important.

self attention

Token embedding vector contains no information about neighboring vector where as context vector contains information about neightboring vector.

Dot product with query and all token gives respective alpha.

Scale the dot product with respective token vector and do summation which is now called context vector.

Why we take dot product?

If two vectors are similar, their dot product will be higher. That’s what done in self attention mechanism.

What’s wrong with this approach?

We can see both dog and ball have same attention score. Here, second it should have more attention to ball than dog. So, if we take simple dot product there is no provision for me to encode the information the ball should be given more priority than dog.

We need to encode a mechanism so that we can capture these complexities and we don’t know what the mechanism would be then we use a trick which researchers have used for long period of time now. If you don’t know what the underlying relationship between things is just to replace it with neural network or buch of trainable matrices and let backpropagation figure it out. That’s exactly what happend in field of attention also.

Researcher invented two matrixes: Query Matrix and Key Matrix what this means is that instead of just looking at input embedding vectors what if I multiply every input embedding with a matrix. Here Query “It ” is multiply with Query Matrix(WQ) which is high dimensional matrix.

Dog and Ball are keys. So, they are multipled with keys Matrix repectively.

The advantage here is that dot product can’t get contextual relationship we are hopping that WQ and WK(intialized randomly and changed by backpropagation)

sdf

Now we are moving from input embedding space to different space which we get after multiplying with query and key matrix. Now, I calculate attention score between these vectors not original vectors.

Now, let’s explore in detail how we convert input embedding vectors into context vectors

Here, below we can see input embedding vectors for each token is 8 dimensional vector. As for now these embedding vectors doesn’t carry any information about how much information to give to other token. So, we need to convert input embedding vectors to context vectors.

We have three trainable matrix Query Weight Matrix: Wq(8,4) , Key Weight Matrix: Wk(8,4) and Value Key Matrix: Wv(8,4). These matrix came in picture because we wanna transform input embedding to different space so that our expressibility increases and we can capture underlying complexities which can’t be done through simple dot product.

All of these trainable weight matrix are initialized randomly. We backpropagate LLM architecture and these values get upadated.

No of column or d_out can be different or same with respect to input dimension and for gpt2, gpt3 etc usually the same. i.e d_out=8 but here for simplicity we have taken d_out=4.

When Input(5,8) multiplied with trainable matrix(8,4) output will be (5,4) called query/key/value vectors respectively. Now, we have moved from 8 dimension to 4 dimension space. After this point we don’t look back to input embedding vectors but look to Query,Key and Value vectors.

We already observed before when we take directly dot product with input embedding vectors it doesn’t work so, we project in different space. This trick we do in deep learning in many different fields for eg: if linear classifer not working with data you agument it with features and you project it into higher dimension space. If handwritten features for computer vision is not working you use a convolution neural network which discovers features own its on higher dimension space. This is what we are doing here.

In step 1, we obtain Q, K and V matrix.

Step: 1

Step 2: This step is all about finding attention scores. To make dot product possible we transpose key so, that query vector is as it is and multiply with query. After multiplication we get attention scores. Here, we find how one query relates to other tokens.

Let’s say I wanna find attention scores between next and other token.

Now, this entire second row represents attention scores between next and all the other keys i.e (The next day is bright).

Similary fifth row represents attention scores between bright and all other key.

Evey row of attention matrix represent attention scores between query of that row and all other key.

There is one problem with these attention scores I want to make interpreatable statement when I am looking at query next such as give 10% attention to “the”, 20% attention to “next”, 20% attention to day, 30% attention to “is” and 20% attention to “bright”. So, I can say maximum attention is given to next and is. I want all this percentage to sum to 100 and probability sum to 1. So, that I can look and piechart and based on this piechart I can see how much attention needs to be paid to each token.

But, when you see the attention scores for next, this values doesn’t sum up to 1 so, I can’t make any statement in term of percentage attention to other token as above. The values in rows doesn’t sum to 1.

Step 3:

So, we do softmax to this second row. Now, if you see these five values and add them you get sum=1 and they will lie between 0 and 1.

Softmax also give lot of weightage to high values and less to low values which makes classification easy.

The main problem is that Softmax also give lot of weightage to high values and doesn’t pay lot of attention to the low values. This is big problem.

Attention scores converted by softmax

This attention weight pay lot attention to one key and don’t pay attention to all other keys at all. That’s why there is need of scaling to perform before we apply softmax. So, we want to make sure before softmax we wanna divide these values by some value then only apply softmax.

Reason 1: For stability in learning

The softmax function is sensitive to the magnitudes of its inputs. When the inputs are large, the differences between the exponential values of each input become much more pronounced. This causes the softmax output to become “peaky,” where the highest value receives almost all the probability mass, and the rest receive very little.

In attention mechanisms, particularly in transformers, if the dot products between query and key vectors become too large (like multiplying by 8 in this example), the attention scores can become very large. This results in a very sharp softmax distribution, making the model overly confident in one particular “key.” Such sharp distributions can make learning unstable.

Scaling is done before softmax as shown below: dkey=4 which is dimension of key.

Why scaled by key dimension?

Because of concept of varience.

The figure shows 100 query vectors and 100 transposed key vectors being multiplied to produce a matrix of attention scores ((QK^T)).

Since the query and key vectors are randomly initialized, each attention score is obtained by combining many random values.

When the key dimension becomes larger, more random values contribute to each dot product, causing the variance of the attention scores to increase.

As a result, some scores become extremely large while others become very small, leading to unstable and highly peaked softmax outputs.

To prevent these large fluctuations and keep the attention scores within a stable range, the scores are scaled by sqrt{dk}. This normalization keeps the variance close to 1, avoids excessive oscillations, and ensures stable and effective learning during Transformer training.

Imagine you’re rolling dice. Consider two cases:

Case 1: Rolling one standard die (1–6)

The average (mean) is 3.5. The variance is relatively small (≈ 2.9). You have predictable outcomes.

Case 2: Rolling and summing 100 dice

The mean is 100 × 3.5 = 350. The variance significantly grows (100 × 2.9 = 290). Now, outcomes fluctuate widely (e.g., you might get sums like 320, 350, or 380). The distribution spreads out drastically. Outcomes become unpredictable.

Dot Product without normalization

Think of dimensions as “dice.” Increasing the number of dimensions is like rolling more dice and summing results. Each dimension (dice) contributes some variance. As dimensions grow, variance accumulates.

Result: Dot products (before softmax) become either extremely large or small, making attention weights unstable and erratic.

Dot Product with normalization (dividing by sqrt(d))

This effectively scales down the variance, ensuring the summed results remain stable. It’s like taking the average roll per dice rather than summing them up, stabilizing your expected outcomes.

Result: Attention weights become more stable, predictable, and informative, enabling the model to learn effectively.

import numpy as np

# Function to compute variance before and after scaling
def compute_variance(dim, num_trials=1000):
    dot_products = []
    scaled_dot_products = []

    # Generate multiple random vectors and compute dot products
    for _ in range(num_trials):
        q = np.random.randn(dim)
        k = np.random.randn(dim)

        # Compute dot product
        dot_product = np.dot(q, k)
        dot_products.append(dot_product)

        # Scale the dot product by sqrt(dim)
        scaled_dot_product = dot_product / np.sqrt(dim)
        scaled_dot_products.append(scaled_dot_product)

    # Calculate variance of the dot products
    variance_before_scaling = np.var(dot_products)
    variance_after_scaling = np.var(scaled_dot_products)

    return variance_before_scaling, variance_after_scaling

# For dimension 5
variance_before_5, variance_after_5 = compute_variance(5)
print(f"Variance before scaling (dim=5): {variance_before_5}")
print(f"Variance after scaling (dim=5): {variance_after_5}")

# For dimension 100
variance_before_100, variance_after_100 = compute_variance(100)
print(f"Variance before scaling (dim=100): {variance_before_100}")
print(f"Variance after scaling (dim=100): {variance_after_100}")

Output:

Variance before scaling (dim=5): 5.119910845892176
Variance after scaling (dim=5): 1.0239821691784352

Variance before scaling (dim=100): 91.07977142757677
Variance after scaling (dim=100): 0.9107977142757677

The output demonstrates that as the vector dimension increases, the variance of the query–key dot products grows significantly, causing the attention scores to become more spread out and unstable.

For a small dimension, the variance remains relatively low, whereas for a larger dimension it increases substantially.

After applying the scaling factor, the variance is reduced to approximately the same stable range in both cases.

This shows that scaling prevents attention scores from becoming excessively large, resulting in smoother softmax distributions and more stable learning during Transformer training.

Step 4: Computing context vector from Attention Weight

What does context vector mean?

Value vector is input embedding vector but which is transformed to different or higher dimension space.

We scale first vector with 0.1 , second with 0.5, third with 0.2 , forth with 0.1 and fifth with 0.1.

The visual representation of context vector is given below.

Blue vectors are input embedding vectors, Scaled vector is shown by green.

When green vectors are added together gives context vector for next.

In fig below we can see how context vector differs from input embedding vector. Input embedding vector contain meaning of that token where as context vector of next contains information of all these attention weights also. All these weighted sum give context vector. This is how attention scores are added together to get the context vector itself.

Similary, we can get context vectors for bright also by multiply 0.1 with first row , 0.05 with 2nd row , 0.1 with 3rd row , 0.25 with 4th row and 0.5 with fifth row and sum them all which is last low of context vector.

Step 5: Input embedding matrix is converted to context vector using self attention block which include all above step.

class SelfAttention_v2(nn.Module):

    def __init__(self, d_in, d_out, qkv_bias=False):
        super().__init__()

        self.W_query = nn.Linear(d_in, d_out, bias=qkv_bias)
        self.W_key   = nn.Linear(d_in, d_out, bias=qkv_bias)
        self.W_value = nn.Linear(d_in, d_out, bias=qkv_bias)

    def forward(self, x):
        keys    = self.W_key(x)
        queries = self.W_query(x)
        values  = self.W_value(x)

        attn_scores = queries @ keys.T
        attn_weights = torch.softmax(
            attn_scores / keys.shape[-1]**0.5,
            dim=-1
        )

        context_vec = attn_weights @ values
        return context_vec

Example:

torch.manual_seed(789)

inputs = torch.tensor(
    [[0.43, 0.15, 0.89],  # Your    (x^1)
     [0.55, 0.87, 0.66],  # journey (x^2)
     [0.57, 0.85, 0.64],  # starts  (x^3)
     [0.22, 0.58, 0.33],  # with    (x^4)
     [0.77, 0.25, 0.10],  # one     (x^5)
     [0.05, 0.80, 0.55]]  # step    (x^6)
)

d_in = 3
d_out = 2

sa_v2 = SelfAttention_v2(d_in, d_out)
print(sa_v2(inputs))

Output:

tensor([[-0.0739,  0.0713],
        [-0.0748,  0.0703],
        [-0.0749,  0.0702],
        [-0.0760,  0.0685],
        [-0.0763,  0.0679],
        [-0.0754,  0.0693]],
       grad_fn=<MmBackward0>)

Casual Attention

I divide entire dataset into chunks of 4 token and based on my batch size I have matix to be created i.e Input Batch 1 and Output Batch 1. Each matrix has 8 row since batch size is 8 and 4 column since context size is 4.

If context size is 4 then there are 4 Input and Target prediction task written below.

For any output input is token before them. In casual attention, we can’t look to the future token. For every output, input is only token which come before output token.

• Causal attention, also known as masked attention, is a special form of self attention.
• It restricts the model to only consider previous and current inputs in a sequence, when processing any given token.

• This is in contrast to the self attention mechanism, which allows access to the entire input sequence at once.
• To achieve this in GPT like LLMs, for each token processed, we mask out the future tokens, which come after the current token in the input text.

All value in diagonal set to 0 than how you calculate attention weights. Ensure each row sum to 1. So, divide in each row by each row total sum.

In, this step we do normalization by softmax and do again normalization each row sum.

Can we do smarter way of normalization?

Yes, try to make diagnoal top part zero in attention scores instead of attention weight.

Multi-Head Attention

Problem of self Attention

This diagram shows a limitation of self-attention: sentence ambiguity. The phrase “woman with a brush” can refer either to the artist using a brush to paint a woman or to a woman who is holding a brush in the portrait. Self-attention may struggle to choose the correct meaning without additional context.

This image shows how self-attention in a Transformer can focus on different words depending on the meaning being learned.

In the sentence “The artist painted the portrait of a woman with a brush”, the phrase “with a brush” is ambiguous.

In Figure 1, the attention pattern links “brush” strongly to “painted”, meaning the model interprets brush as the tool used for painting.

In Figure 2, the attention shifts so that “brush” is more connected to “woman”, suggesting a different interpretation where the phrase may describe the woman.

The red dots indicate strong attention weights between words, showing which tokens the model considers most relevant when understanding the sentence.

This demonstrates how attention helps Transformers capture different semantic relationships and resolve ambiguities in natural language.

Self attention can only capture a single perspective in a given input sequence. It cannot capture multiple perspectives

Implementing 2 head step by step

a. Input Embedding Matrix

a. Input Embedding Matrix

b) Start with a single Wq, Wk, Wv

b) Start with a single Wq, Wk, Wv

c) Split Wq, Wk, Wv into multiple heads (i.e) 2

c) Split Wq, Wk, Wv into multiple heads (i.e) 2

d) Multiple copies of Q, K, V matrices are created.

d) Multiple copies of Q, K, V matrices are created.

e) Computing Attention scores for each head.

e) Computing Attention scores for each head.

f) Computing Attention weight for each head.

f) Computing Attention weight for each head.

g) Compute Context vector matrix for each head.

g) Compute Context vector matrix for each head.

h) Merge Context matrix for both head.

h) Merge Context matrix for both head.

Visualization of Attentions Heads

In the attention visualization, Layer refers to the specific Transformer block being examined within the pre-trained language model (e.g., BERT). Since BERT contains multiple Transformer layers, selecting a layer allows us to inspect the attention patterns learned at that depth. Head refers to the specific attention head within that layer. Each head learns a different relationship or perspective among the tokens. By choosing a particular layer and head, we can visualize which words a token attends to and understand the linguistic patterns captured by that attention head.