The Attention Mechanism — Explained Simply

What the Attention Mechanism Actually Is

You Are Reading a Message From a Friend

"The bank by the river was flooded yesterday but the bank refused to give me a loan because of it."

When you read the second "bank" your brain automatically connected it to the financial institution meaning — not the river meaning. It looked at surrounding words like "refused" and "loan" and assigned them high importance. It ignored "river" and "flooded" as irrelevant to this second bank.

That selective focusing — deciding which parts of the input deserve attention right now — is the attention mechanism.

The Problem It Solves

Before attention, models read text left to right one word at a time with no ability to look back. By the end of a long sentence the beginning had faded. Words far apart could not be connected.

Attention allows every word to directly look at every other word simultaneously — no matter how far apart they are.

The Library Analogy

You are a researcher with a specific question — your query. You walk through the library comparing your question to every book spine — each is a key. The most relevant books get pulled and their content — their values — combine to form your answer.

Query:   "What does bank mean in this sentence?"
Keys:    ["river", "flooded", "refused", "loan", "yesterday"]
Most relevant: "refused" and "loan" ← high attention scores
Least relevant: "river" and "flooded" ← low attention scores
Answer:  Financial institution meaning

Queries, Keys, and Values

Every word becomes three things simultaneously:

"bank" Query:  "What context should I use to understand myself?"
"bank" Key:    "I am potentially a financial or geographical feature"
"bank" Value:  "The actual representation of bank to share with others"

"bank" compares its Query against every other word's Key:
  bank → river:   low score   → pay little attention
  bank → loan:    high score  → pay lots of attention
  bank → refused: medium score → pay some attention

The Attention Score

Dot product of Query and Key — large when they are relevant to each other.

After Softmax the weights sum to 1:

Attention weights for "bank":
  bank → the:       0.02
  bank → bank:      0.05
  bank → refused:   0.18
  bank → loan:      0.68  ← pays most attention here
  bank → yesterday: 0.07
                   ──────
                   sum = 1.0

The Final Step — Weighted Combination

New "bank" = 0.02 × Value("the")
           + 0.05 × Value("bank")
           + 0.18 × Value("refused")
           + 0.68 × Value("loan")   ← dominates
           + 0.07 × Value("yesterday")

Result: "bank" now carries information from "loan" and "refused"
        It knows it means a financial institution in this context

In Code

import torch
import torch.nn.functional as F

def attention(Q, K, V):
    d_k = Q.shape[-1]

    # Step 1 — Every query looks at every key
    scores = Q @ K.transpose(-2, -1)        # (seq_len, seq_len)

    # Step 2 — Scale to prevent large values
    scores = scores / (d_k ** 0.5)

    # Step 3 — Softmax → weights that sum to 1
    weights = F.softmax(scores, dim=-1)

    # Step 4 — Weighted combination of values
    output = weights @ V                    # (seq_len, d_v)

    return output, weights

seq_len, d_k = 5, 64
Q = torch.randn(seq_len, d_k)
K = torch.randn(seq_len, d_k)
V = torch.randn(seq_len, d_k)

output, weights = attention(Q, K, V)
print(output.shape)   # (5, 64) — 5 contextualised word representations
print(weights.shape)  # (5, 5)  — attention matrix

Multi-Head Attention — Multiple Perspectives

One attention head captures one type of relationship. Language has many:

Head 1: learns grammatical relationships (subject → verb)
Head 2: learns coreference (pronoun → its noun)
Head 3: learns positional relationships
...
Head 8: learns semantic similarity

attention_layer = torch.nn.MultiheadAttention(
    embed_dim   = 256,
    num_heads   = 8,      # 8 parallel attention heads
    dropout     = 0.1,
    batch_first = True
)

GPT-3 uses 96 attention heads. Each specialises in different linguistic patterns.

Self-Attention vs Cross-Attention

Self-attention:
  Words in a sentence attend to other words in the SAME sentence
  Used in: encoders, understanding the input
  "The bank refused the loan" → every word looks at every other

Cross-attention:
  Words in one sequence attend to words in a DIFFERENT sequence
  Used in: decoders, translation, connecting input to output
  Output word "le" (French) looks at English input words to find its translation

Why Attention Changed Everything

LSTM reading "The bank by the river... the bank refused to give a loan":
  By the second "bank" it has almost forgotten "river" context
  Too much has passed — limited memory

Attention reading the same sentence:
  Second "bank" directly looks at ALL words simultaneously
  Immediately finds "loan" is most relevant
  Distance does not matter at all

This is why "Attention Is All You Need" (2017) is one of the most important papers ever. It showed you could build powerful language models using attention alone — no recurrence needed. That architecture became the transformer.

The Real Words Mapped to the Story

The One Thing to Remember

Attention allows every word to look at every other word simultaneously and decide how much to focus on each based on relevance. The same word gets a different representation depending on its context. "Bank" next to "loan" becomes a financial institution. "Bank" next to "river" becomes a geographical feature. This context-sensitivity is what makes transformers so powerful.