Deep Learning for Generative AI — Part 4: Math Behind Attention Mechanism

Part 3 showed the idea: the decoder looks back at the encoder states and focuses on the most relevant ones. This part shows the math that makes it work — one small step at a time.

Where this fits: this is the math of decoder-to-encoder attention. The same ingredients (dot products, softmax, Query/Key/Value) appear again in Part 9 for self-attention inside a Transformer.

1. Attention in one sentence

Attention lets the decoder ask every encoder state, “How useful are you to what I need right now?” and then take a weighted average of the answers.

The weights are not fixed. They are recomputed at every decoder step.

2. The three roles and how they work together

Think of a library search:

The decoder writes a query. It compares the query to every key. The better the match, the more of that value it keeps.

Step 1: the encoder turns each source word into a memory

Before the decoder can use attention, the encoder must read the source sentence.

Take the source sentence:

I    love   cats

Each word is first turned into an embedding vector:

"I"    →  e1
"love" →  e2
"cats" →  e3

The encoder then processes these vectors one by one. In a simple recurrent encoder, this looks like:

hidden_0 = zero vector
hidden_1 = RNN(e1, hidden_0) = h1   (memory of "I")
hidden_2 = RNN(e2, hidden_1) = h2   (memory of "I love")
hidden_3 = RNN(e3, hidden_2) = h3   (memory of "I love cats")

After reading the whole sentence, the encoder has produced three hidden states:

h1 = memory of "I"
h2 = memory of "I love"
h3 = memory of "I love cats"

These hidden states are the Keys and Values for attention. Each one is a compact summary of what the encoder has read up to that point.

Step 2: the decoder asks a question

Now the decoder starts generating the translation. Imagine we are translating:

Source: I    love   cats
Target: J'aime chats

The query is not a sentence. It is a vector that comes from the decoder’s current hidden state. To make this concrete, we will use the same 3-dimensional hidden states from the encoder:

h1 = [1, 0, 0]  → "I"
h2 = [0, 1, 0]  → "love"
h3 = [0, 0, 1]  → "cats"

First target word: J’

At the first step, suppose the decoder’s hidden state happens to be:

Q1 = [2, 0, 1]

This vector is the query. It represents what the decoder is currently looking for. In this toy example, the high numbers in positions 1 and 3 mean the decoder is interested in “I” and “cats”.

Compare the query with each key using the dot product:

Q1 · h1 = 2×1 + 0×0 + 1×0 = 2
Q1 · h2 = 2×0 + 0×1 + 1×0 = 0
Q1 · h3 = 2×0 + 0×0 + 1×1 = 1

scores = [2, 0, 1]

After scaling by √3 and applying softmax, the weights become:

weights = [0.533, 0.168, 0.299]

The context vector is mostly built from "I":

context = 0.533 × [1, 0, 0]   (I)
        + 0.168 × [0, 1, 0]   (love)
        + 0.299 × [0, 0, 1]   (cats)

The decoder uses this context to predict the first target word: J'.

Second target word: aime

After producing J', the decoder updates its hidden state. Its new query is now different because it has already generated one word and is looking for the next part of the sentence:

Q2 = [0, 2, 1]

Compare again:

Q2 · h1 = 0×1 + 2×0 + 1×0 = 0
Q2 · h2 = 0×0 + 2×1 + 1×0 = 2
Q2 · h3 = 0×0 + 2×0 + 1×1 = 1

scores = [0, 2, 1]

After softmax, the weights shift toward "love":

weights = [0.168, 0.533, 0.299]

Now the context vector focuses on "love", and the decoder predicts aime.

The point is simple: the query is a vector, and when that vector changes, the attention weights change with it. The decoder does not think in English sentences — it just has a new hidden-state vector at every step.

One decoder vector per step, not one per source word. The decoder is also a recurrent network, so it has one hidden-state vector at each step — just like the encoder had one hidden-state vector after each source word. That single decoder vector becomes the query for attention. It is compared against all the encoder states at once, and the result is one context vector for that step.

The attention call

So the decoder is calling a function like this at every step:

context_i = attention(
    query  = decoder_state_i,
    keys   = [encoder_state_1, encoder_state_2, encoder_state_3],
    values = [encoder_state_1, encoder_state_2, encoder_state_3]
)

The query changes at every step, so the answer changes too.

Why are keys and values the same list? In this encoder-decoder attention setup, both keys and values come from the same source — the encoder hidden states. The key is used to decide which source state is relevant, and the value is the actual content that gets blended. Because the encoder states are good for both jobs, we use them twice. In self-attention (inside a Transformer) the Q/K/V are usually produced by three different learned matrices, but the encoder-decoder attention you see in translation typically reuses the encoder states for keys and values.

Where the vectors come from

• Keys and Values come from the encoder hidden states.
• Query comes from the decoder hidden state at the current step.

In a Transformer, Q/K/V are produced by learned weight matrices (W_Q, W_K, W_V), but the idea is the same: keys and values describe the source, and the query describes what the decoder currently needs. We cover that in Part 9.

3. The three-step recipe

1. Compare the query with every key → raw scores
2. Turn the scores into percentages (weights) with softmax
3. Blend the values using those weights → context vector

That is the whole attention mechanism. Everything else is just details.

4. A concrete example with small numbers

We will use 3-dimensional vectors and three source words:

h1 = [1, 0, 0]  → encoder state for "I"
h2 = [0, 1, 0]  → encoder state for "love"
h3 = [0, 0, 1]  → encoder state for "cats"

The keys and values are the same as the encoder states:

K1 = V1 = h1 = [1, 0, 0]
K2 = V2 = h2 = [0, 1, 0]
K3 = V3 = h3 = [0, 0, 1]

Suppose the decoder query is:

Q = [2, 0, 1]

This query is like the decoder saying: “I am interested in the first and third source words right now.”

Step 1: raw scores — ask each memory how relevant it is

The score is the dot product:

Q · K1 = 2×1 + 0×0 + 1×0 = 2
Q · K2 = 2×0 + 0×1 + 1×0 = 0
Q · K3 = 2×0 + 0×0 + 1×1 = 1

scores = [2, 0, 1]

The query matches "I" best, then "cats", then "love".

Step 2: scale — keep the numbers friendly

The vectors have dimension d = 3, so we divide by √d ≈ 1.732:

scaled scores = [2/1.732, 0/1.732, 1/1.732]
              = [1.155, 0.000, 0.577]

Scaling keeps the numbers from growing too large when we use long vectors.

Step 3: softmax — turn scores into weights

Softmax converts any list of numbers into percentages that add up to 1:

e^1.155 ≈ 3.174
 e^0    = 1.000
e^0.577 ≈ 1.781

sum = 3.174 + 1.000 + 1.781 = 5.955

weights = [3.174/5.955, 1.000/5.955, 1.781/5.955]
        = [0.533, 0.168, 0.299]

Check: 0.533 + 0.168 + 0.299 = 1.000.

Step 4: weighted sum — build the final answer

c = 0.533 × V1 + 0.168 × V2 + 0.299 × V3
  = 0.533 × [1, 0, 0] + 0.168 × [0, 1, 0] + 0.299 × [0, 0, 1]
  = [0.533, 0.168, 0.299]

The final context vector is [0.533, 0.168, 0.299]. It is a blend of all three source words, with most weight on "I".

What happens for the next word?

After producing one word, the decoder updates its hidden state and creates a new query. Then it runs the same four steps again. Because the query is different, the weights are different, and the decoder gets a fresh context vector for the next word.

5. Why do we need softmax?

Without softmax, the weights could be negative or not add up to 1. Softmax guarantees two useful things:

• every weight is between 0 and 1,
• all weights add up to exactly 1.

That makes the context vector a true blend of the values, not a random sum.

6. Why do we scale by √d?

When vectors are long (for example, 512 dimensions), dot products can become huge. Huge numbers make softmax push one weight close to 1 and all others close to 0. The gradients then become very small and learning slows down.

Dividing by √d keeps the scores in a comfortable range.

score(Q, K) = (Q · K) / √d

That is all the scaling does.

7. The full attention formula

Putting the three steps into one line gives the famous scaled dot-product attention:

Attention(Q, K, V) = softmax( (Q · K^T) / √d ) · V

Read it from left to right:

1. Q · K^T — compare the query with every key.
2. / √d — scale the scores down.
3. softmax(...) — turn scores into weights.
4. · V — blend the values with those weights.

Why K^T?

Q is one vector with shape (1, d) — one query of size d. K is a matrix with shape (n, d) — n keys, each of size d.

To compare the query with every key, we need the dot product Q · K1, Q · K2, …, Q · Kn. Mathematically that is the same as multiplying Q by the transpose of K:

Q shape   = (1, d)
K shape   = (n, d)
K^T shape = (d, n)

Q · K^T   = (1, d) · (d, n) = (1, n)

The result is a single row of n scores — one score for each key. That is why the formula writes Q · K^T instead of Q · K.

The output is the context vector c.

What is matrix multiplication actually doing here? A kid-friendly view

If we strip away the big words, matrix multiplication is just a way to do many multiplications and additions at once.

Imagine you have three toy boxes:

• Box 1: 2 red blocks, 1 blue block
• Box 2: 0 red blocks, 3 blue blocks
• Box 3: 1 red block, 1 blue block

Your wish list is 2 red blocks and 1 blue block. Which box matches best?

You multiply matching colors and add them up:

Box 1: 2×2 + 1×1 = 5
Box 2: 2×0 + 1×3 = 3
Box 3: 2×1 + 1×1 = 3

Box 1 gets the biggest number, so it is the best match. That is exactly what a dot product does.

In attention, Q · K^T does the same thing for every key at the same time:

Q = [2, 1]       (my wish list)

K1 = [2, 1]      (Box 1)
K2 = [0, 3]      (Box 2)
K3 = [1, 1]      (Box 3)

scores = [5, 3, 3]   ← tells us how well each box matches

Softmax turns those scores into percentages — for example, mostly Box 1 and a little bit of the others.

Then · V mixes the boxes using those percentages. If Box 1, Box 2, and Box 3 are also your values, the final answer is mostly Box 1 with small amounts of the other boxes mixed in.

So matrix multiplication is just a fast way of saying:

1. Check every key.
2. Score every match.
3. Mix the values using the scores.

That is all the scary-looking formula is doing.

8. What changes in the decoder?

In Part 2 the decoder updated its state using a fixed context vector:

s_i = f(s_{i-1}, y_{i-1}, c)

With attention, the only change is that c becomes c_i — a fresh context vector at every step:

s_i = f(s_{i-1}, y_{i-1}, c_i)
                          ↑
                 recomputed at every decoder step

The previous decoder state s_{i-1} and the previous output word y_{i-1} stay the same. The context is rebuilt each time because the query changes.

One decoder hidden state at a time

The decoder does not hold many hidden states at once. Like any recurrent network, it has one hidden-state vector at each timestep:

step 1:  s0  →  produces word 1
step 2:  s1  →  produces word 2
step 3:  s2  →  produces word 3

At step i only s_i exists. It is used as the query, attention returns c_i, and then s_i plus c_i plus the previous output y_{i-1} are used to compute the next hidden state s_{i+1}.

So the decoder has one hidden state at a time, but that hidden state is different at every step.

A concrete walkthrough

Translating English → French:

Source: The   cat    sat
Target: Le    chat   s'est   assis

Here is what happens inside the decoder:

Each row is the same four-step recipe: query → scores → softmax → weighted sum. Only the query changes, so the context vector changes with it.

The full picture:

Source words → Encoder → hidden states h1, h2, h3 (keys + values)
                                          ↓
Decoder step 1:  Q1 = s0  →  attention  →  c1  →  "Le"
Decoder step 2:  Q2 = s1  →  attention  →  c2  →  "chat"
Decoder step 3:  Q3 = s2  →  attention  →  c3  →  "s'est"
Decoder step 4:  Q4 = s3  →  attention  →  c4  →  "assis"

How does the decoder know how many words to produce?

It does not know in advance. The decoder generates one word at a time and stops when it predicts a special end-of-sentence token, usually written <eos>.

During training, the target sentence is already known:

Target: Le    chat   s'est   assis   <eos>

The decoder is fed the true previous word at every step (this is called teacher forcing). It produces one output for each position in the target, including <eos>. The loss is computed at every position, so the model learns to predict the right word and, eventually, when to emit <eos>.

At inference time, when there is no target sentence available, the process looks like this:

Input:  <sos>
Output: "Le"

Input:  "Le"
Output: "chat"

Input:  "chat"
Output: "s'est"

Input:  "s'est"
Output: "assis"

Input:  "assis"
Output: "<eos>"   ← stop here

The previous output becomes the next input. The decoder keeps going until it either produces <eos> or reaches a maximum length limit.

So the decoder decides the sentence length itself. It does not get a fixed number of words from the encoder.

What if the word order changes?

English → German:

English: I     have    seen    him
German:  Ich   habe    ihn     gesehen

The English verb phrase “have seen” becomes “habe … gesehen”, with the second half jumping to the end. A fixed word-to-word mapping would fail. Attention learns to look at “seen” when generating “gesehen”, even though the positions differ.

9. Where do the scores come from?

The scores are not hand-coded. They come from a small learnable function:

score(decoder_state, encoder_state) = a(s_{i-1}, h_j)

The function a has weights that are updated during training. If the model focuses on the wrong word and produces a bad translation, the loss pushes the weights to focus better next time. Over millions of sentences, the model learns where to look.

The simplest version uses a dot product. A richer version, called Bahdanau (additive) attention, uses a tiny neural network instead. In modern Transformers, the dot-product version is used because it is fast and works well.

10. A tiny PyTorch check

Here is the whole recipe in code:

import torch
import torch.nn.functional as F
import math

def attention(query, key, value):
    d = query.size(-1)
    scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d)
    weights = F.softmax(scores, dim=-1)
    context = torch.matmul(weights, value)
    return context, weights

# one query, three keys/values of dimension 3
Q = torch.tensor([[2.0, 0.0, 1.0]])
K = torch.tensor([[1.0, 0.0, 0.0],
                  [0.0, 1.0, 0.0],
                  [0.0, 0.0, 1.0]])
V = K.clone()

context, weights = attention(Q, K, V)
print("weights:", weights)   # tensor([[0.533, 0.168, 0.299]])
print("context:", context)   # tensor([[0.533, 0.168, 0.299]])

The numbers match the hand calculation from section 4.

11. Summary

The full formula, worth remembering, is:

Attention(Q, K, V) = softmax( (Q · K^T) / √d ) · V

Next: Part 5: Drawbacks of Attention-Based Encoder-Decoder Architecture