Transformer Attention

Scaled dot-product attention

Prerequisites

graph TD
  LinAlg[Vectors / dot product] --> Attn[Scaled dot-product attention]
  Prob[Softmax / probability] --> Attn
  IT[Information theory] -.-> Attn

When to use

• Long-range dependencies in language (pronouns, coreference, translation alignment).
• Parallel sequence modeling instead of sequential RNN state bottlenecks.
• Context mixing where each position reweights all keys via softmax attention.

When not to use

• Ultra-long contexts without sparse/linear attention approximations — full $O(n^2)$ attention is costly.
• Tasks with fixed local windows only — simpler convolutions or sliding windows may suffice.
• When interpretability of routing matters more than raw quality — attention maps are hints, not proofs.

Lab

On the interactive attention lab, tune the 2-D query sliders, read the softmax heatmap over four toy keys, compare attention α to a uniform baseline, try “it” → animal or street focus presets, predict the top key, then reveal which word wins.

Simple Fundamental Explanation

Imagine you are reading a long legal contract.

• The old way (RNN/LSTM): You read it word by word, left to right. By the time you get to page 10, you’ve mostly forgotten the exact phrasing on page 1. You only have a vague “hidden state” summarizing what you read.
• The Transformer way (Attention): You read the whole document at once. Every time you look at a word, you draw a mental string connecting it to every other relevant word on the page, regardless of how far apart they are. If you see the word “bank”, you draw strings to “river” and “mud”, and ignore “money”.

The Attention Mechanism allows an AI model to probabilistically weigh the importance of every single word in a sentence relative to every other word simultaneously. This allows it to understand deep, complex context that previous AI architectures fundamentally missed.

Deep Dive: How It Works Under the Hood

The Transformer architecture (introduced in the famous paper “Attention Is All You Need”) discarded sequential processing in favor of parallelized matrix math. The core engine is “Scaled Dot-Product Attention”.

Every word (token) in a sentence is represented by a vector (a list of numbers). To compute attention, the model splits this vector into three new vectors:

1. Query ($Q$): “What am I looking for?”
2. Key ($K$): “What do I contain?”
3. Value ($V$): “If you select me, what information do I provide?”

The Matching Process

For every word in the sentence, the model takes its Query vector and compares it against the Key vectors of all other words using a dot product. If the Query of word A aligns well with the Key of word B, the dot product is a large positive number.

These scores are passed through a Softmax function, converting them into a probability distribution that sums to 1.0. Finally, the model multiplies these probabilities by the Value ($V$) vectors. The word’s new, context-aware meaning is a probabilistically weighted sum of the Values of all the words it paid attention to.

Visual Diagram: Self-Attention

Sentence: "The animal didn't cross the street because it was too tired."
What does "it" refer to? The animal or the street?

Focus Word: "it" (Query)

Dot Product (Query of "it" * Keys of all words):
Key("The")    -> Score: 2
Key("animal") -> Score: 95   <-- Massive alignment!
Key("street") -> Score: 10
Key("tired")  -> Score: 80

Softmax (Probabilities):
"animal" -> 0.85
"tired"  -> 0.12
"street" -> 0.01
...

New Contextual Vector for "it":
V_new = (0.85 * Value("animal")) + (0.12 * Value("tired")) + ...
The model now mathematically understands that "it" means "the tired animal".

Practical Example: Machine Translation

Before Transformers, Google Translate used Recurrent Neural Networks (Seq2Seq). When translating a 50-word German sentence to English, the network had to compress the entire German sentence into a single, fixed-size vector bottleneck before decoding it into English. Long sentences suffered severe quality drops.

With Transformers, when generating the 10th English word, the attention mechanism looks across the entire 50-word German input and assigns probabilities to which German words are most relevant to the current English word being translated. There is no bottleneck.

Furthermore, because Attention is just giant matrix multiplications, it can be computed for all words simultaneously on GPUs. This allowed models to scale from translating single sentences to digesting entire 100-page books (like ChatGPT’s 128k context window).

The Mathematics: Equations and In-Depth Analysis

1. The Attention Equation

Let $Q$ be the matrix of queries, $K$ be the matrix of keys, and $V$ be the matrix of values. The Scaled Dot-Product Attention is defined as:

Explanation of terms:

• $Q{K}^T$: The dot product of every query with every key. This produces an $N\times N$ matrix of raw compatibility scores (where $N$ is the sequence length).
• $\sqrt{d_k}$: The scaling factor. $d_k$ is the dimension of the key vectors. Without this, the dot products for large vectors grow massive, pushing the Softmax function into regions with extremely small gradients, which stalls the neural network training.
• $\text{Softmax}(\dots )$: Converts the raw scores into a valid probability distribution where each row sums to 1.
• $\times V$: The probabilistically weighted sum of the values.

2. Multi-Head Attention

Instead of performing a single attention function, Transformers use Multi-Head Attention. The $Q,K,V$ vectors are linearly projected into $h$ different smaller subspaces (heads). Attention is computed for each head in parallel, and the results are concatenated.

Why? It allows the model to probabilistically attend to different types of relationships simultaneously. Head 1 might learn to attend to grammar (verbs pointing to nouns). Head 2 might learn to attend to historical facts. Head 3 might attend to emotional sentiment.

3. The Scaling Bottleneck (The $O(N^2)$ Problem)

Look at the matrix $Q{K}^T$. It requires multiplying every word by every other word. If the sequence length is $N$, the compute and memory required scale quadratically: $O(N^2)$.

• 1,000 words = 1 million calculations.
• 100,000 words = 10 billion calculations.

This quadratic complexity is why LLMs historically struggled with massive context windows, and has sparked intense research into “Sparse Attention” and linear approximations to bypass the $O(N^2)$ math.

Lab