Post

Attention: The Mathematics Behind Understanding Context

Demystifying attention through Queries, Keys, Values, and a step-by-step numerical example of scaled dot-product attention.

1,209 words7 min read
On this page

When ChatGPT first came out, I remember being amazed by it. But ironically, I also remember laughing whenever it failed to figure out what a pronoun referred to, or couldn't correctly establish the relationship between two subjects in a sentence.

It was a strange feeling. For some reason, a part of me wanted it to fail. Perhaps because the alternative felt almost magical: a machine actually understanding the meaning of what I was saying.

Fast forward to my Computer Science days, and it turns out that much of this magic is just mathematics. Well, mathematics is a kind of magic in itself, isn't it?

Enter Attention

One of the most important breakthroughs in helping models understand relationships within sequences was the concept of Attention.

The core idea is surprisingly simple.

When a model tries to understand a word in a sequence, attention helps it identify other words that might be relevant to understanding the current word.

Consider this sentence:

The animal in the jungle went out on a hunt because it was hungry.

When the model encounters the word "it", how does it figure out what "it" refers to?

Is it the animal? The jungle? Or perhaps something else?

We humans can establish this relationship almost effortlessly. But how does a machine do it?

This is where attention comes in. It allows the model to learn these relationships mathematically.

Attention as a Search System

A useful way to think about attention is as a search system.

It has three main components:

  • Query (QQ): What am I looking for? In our example, we are trying to understand the word "it".
  • Key (KK): What information can I match my query against? Think of words such as "animal", "jungle", and "hunt".
  • Value (VV): What information should I retrieve once I know which words are relevant?

Each token gets its own Query, Key, and Value, all represented as vectors. These vectors are produced through learned transformations of the token's representation.

The Query is compared with the Keys to calculate how relevant each token is. The resulting scores determine how much attention each token receives.

Finally, these attention weights are used to combine the Value vectors into a single output.

Ideally, when processing "it", the model would assign more attention to "animal" than to "jungle", since "animal" is the more relevant reference in our sentence.

But how are these attention scores calculated?

Scaled Dot-Product Attention

The mathematical expression for scaled dot-product attention is:

Attention⁡(Q,K,V)=softmax⁡(QKTdk)V\boxed{\operatorname{Attention}(Q,K,V)=\operatorname{softmax}\left(\frac{QK^{\mathsf T}}{\sqrt{d_k}}\right)V}

It looks intimidating at first, but let's break it down.

Step 1: Calculate the dot product

The dot product between the Query and a Key, Q⋅KiQ\cdot K_i, gives us a compatibility score.

In simple terms, it tells us how well the Query matches that Key.

A higher score generally means a stronger match.

Step 2: Scale the scores

We divide each score by dk\sqrt{d_k}, where dkd_k is the dimensionality of the Key vectors.

Why?

Because as the dimensionality increases, dot products can become very large. Scaling helps prevent excessively large scores from pushing the softmax into regions where its gradients become extremely small.

Step 3: Apply softmax

Next, we apply the softmax function to the scaled scores.

Softmax converts these scores into attention weights that:

  • Are between 0 and 1.
  • Add up to 1.
  • Give relatively greater importance to larger scores.

These weights tell the model how much information to gather from each corresponding Value vector.

A Quick Example: Understanding Softmax

Let's simplify our sentence and consider three words.

Suppose the Query for "it" produces the following scaled scores after being compared with their Keys:

WordScaled score
Animal3
Jungle1
Hunt0

The softmax function is:

softmax⁡(z)i=ezi∑jezj\operatorname{softmax}(z)_i=\frac{e^{z_i}}{\sum_j e^{z_j}}

First, exponentiate each score (rounded to two decimal places):

e3≈20.09,e1≈2.72,e0=1e^3\approx20.09,\quad e^1\approx2.72,\quad e^0=1

Now calculate the sum of these rounded values:

20.09+2.72+1=23.8120.09+2.72+1=23.81

Finally, divide each exponential by the total.

WordCalculationAttention weight (rounded)
Animal20.09/23.8120.09 / 23.810.844
Jungle2.72/23.812.72 / 23.810.114
Hunt1/23.811 / 23.810.042

Notice that:

0.844+0.114+0.042=10.844+0.114+0.042=1

So, in this simplified example, "animal" receives approximately 84.4% of the attention.

Why softmax instead of simply dividing by the sum?

Because softmax has three useful properties:

  • Always positive: Even negative scores produce positive attention weights.
  • Normalized: All weights add up to 1, making them easy to interpret as relative importance.
  • Emphasizes larger scores: Exponentiation increases the relative importance of higher scores.

At this point, we know which words deserve attention.

But we still haven't gathered any actual information from them.

And that's where Values come into the picture.

Bringing Values into the Picture

After applying softmax, we have our rounded attention weights:

α=[0.844, 0.114, 0.042]\alpha=[0.844,\ 0.114,\ 0.042]

Notice something interesting.

These numbers only describe how much attention each word receives.

They don't directly contain the information represented by those words. They simply tell us how strongly to weigh each corresponding Value vector.

Now, let's bring in the Values.

For simplicity, suppose our three words have the following two-dimensional Value vectors:

V1=[2,4](Animal)V2=[6,2](Jungle)V3=[1,8](Hunt)\begin{aligned} V_1&=[2,4] &&\text{(Animal)}\\ V_2&=[6,2] &&\text{(Jungle)}\\ V_3&=[1,8] &&\text{(Hunt)} \end{aligned}

These are arbitrary vectors chosen to make the calculations easier. In a real model, the Value vectors contain learned features rather than human-readable meanings.

The final attention output is calculated as a weighted sum of the Value vectors:

Output=α1V1+α2V2+α3V3\boxed{\text{Output}=\alpha_1V_1+\alpha_2V_2+\alpha_3V_3}

Substituting our rounded weights:

Output≈0.844[2,4]+0.114[6,2]+0.042[1,8]\text{Output}\approx0.844[2,4]+0.114[6,2]+0.042[1,8]

Multiplying each vector by its attention weight:

=[1.688,3.376]+[0.684,0.228]+[0.042,0.336]=[1.688,3.376]+[0.684,0.228]+[0.042,0.336]

Adding them together:

Output≈[2.414, 3.94]\boxed{\text{Output}\approx[2.414,\ 3.94]}

This result uses the rounded attention weights above. Keeping full precision throughout gives approximately [2.41477, 3.93965][2.41477,\ 3.93965].

And there we have it!

This is the final attention output vector for our Query, "it", in this simplified example.

Notice that the Value vector for "animal" contributes the most to our output because it received the highest attention weight.

The resulting vector carries information gathered from the other tokens, weighted according to their relevance to our Query.

It doesn't explicitly say "it = animal". Rather, it provides a contextual representation that the model can use to understand relationships like this.

And that's the beauty of attention.

The Query determines what we're looking for, the Keys help us decide where to look, and the Values provide the information we retrieve.

A Final Thought

What once felt like magic to me turned out to be a series of dot products, scaling operations, and weighted sums.

Of course, attention alone doesn't explain everything about how LLMs work, but it's a fascinating piece of the puzzle.

And perhaps that's what makes it even more magical: that something as abstract as mathematics can help a machine make sense of relationships in language that we humans barely even think about.

Turns out, understanding attention was all I needed. ;)

Share

← Back to the blog