Attention Mechanisms

Why attention exists (the bottleneck story)

In sequence-to-sequence modeling, we often want an output sequence to depend on different parts of the input at different times. Translation, summarization, speech recognition, program synthesis—these tasks all have alignment structure:

• •When translating a noun, you want to focus on the source noun phrase.
• •When translating a verb, you might need to look at a different region.
• •When resolving a pronoun (“it”, “they”), you need to find an antecedent that could be far away.

Older encoder–decoder RNNs forced the entire input sequence into one fixed-size vector (or a narrow channel through the final hidden state). This creates an information bottleneck: long sequences degrade because the decoder can’t selectively retrieve what it needs.

Attention removes the bottleneck by turning “memory” into a set of vectors (one per input element) and letting the model compute a weighted combination of those vectors each time it needs context.

The three roles: Query, Key, Value

Attention is easiest to understand by analogy to retrieval:

• •A query (Q) encodes what you’re currently looking for.
• •A set of keys (K) encodes what each memory item is “about.”
• •A set of values (V) encodes the information you will actually retrieve.

The algorithm:

1. 1)Score how similar each key is to the query (relevance).
2. 2)Convert scores into a probability distribution (weights).
3. 3)Use weights to compute a weighted sum of values.

This is not just a metaphor; it’s literally what the math does.

A minimal single-query definition

Suppose we have one query vector q ∈ ℝᵈ, and n keys/values {(kᵢ, vᵢ)} for i=1..n.

1) Similarity scoring (dot-product attention):

2) Score-to-weight via softmax:

3) Aggregate values:

Here o is the attention output (sometimes called the “context vector”).

Why Q, K, V are usually learned projections

In neural networks, the input tokens already have embeddings xᵢ. We project them into Q/K/V spaces with learned affine transformations:

• •q = W_Q x
• •k = W_K x
• •v = W_V x

This matters because:

• •the model can learn what “matching” should mean (via W_Q, W_K), and
• •the model can learn what to retrieve (via W_V) even if that differs from what helps matching.

Self-attention vs cross-attention (source origin distinction)

This node emphasizes a crucial distinction:

• •Self-attention: Q, K, V come from the same sequence (same source). Each token can attend to other tokens in its own sequence.
• •Cross-attention: Q comes from one source (e.g., decoder states), while K and V come from another (e.g., encoder outputs). The decoder queries the encoder memory.

In translation terms:

• •self-attention inside the encoder mixes information among source tokens,
• •self-attention inside the decoder mixes information among generated tokens (with causal masking),
• •cross-attention lets the decoder pull in source information.

Preview: the axis and mask pitfalls

Attention is easy to write but easy to implement incorrectly.

• •Softmax axis mistake: weights must sum to 1 over keys for each query. If you normalize over the wrong axis, the layer still outputs something, gradients still flow, but the meaning of “attention” is broken.
• •Mask leakage: masks must align with the score matrix shape. Broadcasting mistakes can allow forbidden positions (future tokens or padding) to influence outputs.

We’ll keep returning to shapes and axes so you can debug these confidently.

Why scoring is the heart of attention

If attention is “weighted focus,” then the score function decides what counts as relevant. The score is computed between a query and each key.

In practice, the most common scoring rule is dot-product similarity because it is fast on GPUs and works well with learned projections.

From one query to many queries: matrix form

Assume we have:

• •Q ∈ ℝ^(m×dₖ): m queries, each dₖ-dimensional
• •K ∈ ℝ^(n×dₖ): n keys, each dₖ-dimensional
• •V ∈ ℝ^(n×dᵥ): n values, each dᵥ-dimensional

Compute all pairwise query–key scores:

Shapes:

• •Q is (m×dₖ)
• •Kᵀ is (dₖ×n)
• •S is (m×n)

Interpretation:

• •S_{ij} is the score between query i and key j.

This “score matrix” is the object you will mask, normalize, and use to weight V.

Why the scaling factor 1/√dₖ exists

In Transformers, the standard formula is scaled dot-product attention:

Motivation: dot products grow in magnitude with dimension.

A rough variance argument:

• •Suppose components of q and k are independent with mean 0 and variance 1.
• •Then qᵀk = ∑_{t=1}^{dₖ} q_t k_t.
• •Each term q_t k_t has mean 0, and (under independence) variance ≈ 1.
• •The sum has variance ≈ dₖ.

So typical score magnitudes scale like √dₖ. Large magnitudes push softmax into saturation:

• •one position gets probability ≈ 1
• •others ≈ 0
• •gradients become small for most keys

Dividing by √dₖ keeps the score distribution more stable as dₖ changes.

Alternative similarity scoring functions

Dot product is not the only option. Historically, early attention used additive scoring.

Because you already know cosine similarity: note that dot-product attention can learn to behave like cosine similarity if the model learns to normalize representations (or learns norm control via layer norm / projection matrices). But in standard Transformers, the scaling is the main explicit normalization.

Shape discipline (the first line of defense against bugs)

When implementing scoring, always write down:

• •batch size: B
• •number of queries: m (often sequence length L_q)
• •number of keys: n (often L_k)
• •feature dimension: dₖ

In a batched setting:

• •Q: (B×m×dₖ)
• •K: (B×n×dₖ)
• •scores S = QKᵀ: (B×m×n)

A common silent bug: transposing the wrong axes so you compute (B×dₖ×dₖ) or normalize over the wrong dimension.

Causal structure and “who can look at whom”

The score matrix S encodes potential connections:

• •In encoder self-attention, typically every token can attend to every token (except padding).
• •In decoder self-attention, we must enforce causality: token t cannot attend to tokens > t.

This is done not at the Q/K/V level but by masking the score matrix before softmax.

We’ll treat masking carefully in the next mechanic because it interacts directly with the probability distribution.

Why we need a distribution, not raw scores

Raw scores S_{ij} are unbounded real numbers. To create a “focus,” we need nonnegative weights that sum to 1 across keys for each query.

Softmax does exactly this, turning each query’s score row into a categorical distribution over keys.

The core formula (matrix form)

Given scores:

Compute attention weights:

Important: softmax is applied row-wise over the key dimension (size n). That means:

Then aggregate values:

Shapes:

• •A is (m×n)
• •V is (n×dᵥ)
• •O is (m×dᵥ)

Interpretation:

• •output for each query i is a convex combination of the value vectors.

Masking: forbidding attention to certain positions

Masking modifies S before softmax so forbidden positions get probability ≈ 0.

Two common masks:

1) Padding mask (ignore pad tokens)

• •If some keys correspond to padding positions, we must prevent attending to them.

2) Causal mask (prevent “future” access in autoregressive decoding)

• •For query position i, keys j > i should be forbidden.

Mechanically, we add a large negative number to masked scores:

Where M_{ij} = 0 if allowed, and M_{ij} = -\infty (or a large negative constant like -10^9) if disallowed.

Then:

Because exp(-∞) → 0, masked entries get weight 0.

The surprising failure mode: mask leakage via broadcasting

Masks are often stored with shape (B×1×1×n) or (B×1×m×n) depending on implementation (especially with multi-head attention).

A common bug pattern:

• •scores S has shape (B×h×m×n)
• •mask has shape (B×m×n)
• •you add them and broadcasting might align incorrectly

Result: you mask the wrong dimension or the wrong positions. The model may still train but exhibits “cheating” (decoder sees future) or ignores padding improperly.

Practical discipline:

• •assert shapes explicitly
• •test that masked attention weights at forbidden positions are ~0
• •test with tiny sequences where you can print A

Softmax axis mistake (the other silent bug)

Given S of shape (B×m×n):

• •correct: softmax over n (keys) so each row sums to 1
• •wrong: softmax over m (queries) or over d (features)

If you accidentally normalize over queries, you enforce that each key distributes probability over queries, which is not the retrieval interpretation.

A quick invariant check:

• •For correct attention: for each query i, ∑_{j=1}^n A_{ij} = 1.

Numerical stability: subtract max

Softmax can overflow if scores are large. The standard stable computation:

For each row i:

This doesn’t change results because subtracting a constant from all logits preserves softmax.

Temperature and sharpness

Sometimes you’ll see a temperature τ:

• •τ < 1 makes distributions sharper (more peaky)
• •τ > 1 makes them softer (more uniform)

The Transformer’s √dₖ scaling can be interpreted as a kind of dimension-dependent temperature.

Weighted sum as linear algebra (and why it’s differentiable)

Once you have A, the output is:

This is a linear combination of V with coefficients from A.

Differentiability:

• •Softmax is differentiable
• •Matrix multiply is differentiable
• •So the whole attention block is end-to-end trainable

To see the dependency explicitly for a single query i:

If A_{ij} increases, oᵢ moves toward vⱼ.

A useful mental model: attention is “content-addressable memory”

Keys provide an address space, queries pick addresses, values store content. The softmax makes it a soft (continuous) lookup rather than a hard index.

This is why attention can represent alignment: it’s literally learning a soft alignment matrix A.

At this point, you have the atomic concepts:

• •similarity scoring (Q vs K)
• •score → weights via softmax
• •weighted sum of V

Next we connect that to the self vs cross distinction in full architectural context.

Why the “origin of Q,K,V” matters

Attention is a general operator: it maps (Q,K,V) to O. The difference between self- and cross-attention is simply where these tensors come from.

This origin choice encodes a modeling decision:

• •Do we want interactions within a sequence? (self)
• •Or do we want one sequence to retrieve from another? (cross)

Self-attention: mixing information inside one sequence

Let X ∈ ℝ^(L×d_model) be a sequence of L token embeddings (after adding positional information).

We compute:

Where:

• •W_Q ∈ ℝ^(d_model×dₖ)
• •W_K ∈ ℝ^(d_model×dₖ)
• •W_V ∈ ℝ^(d_model×dᵥ)

Then:

Mask M depends on the setting:

• •encoder: padding mask only
• •decoder: padding mask + causal mask

Interpretation: each token representation is updated by blending information from other tokens.

A key property: self-attention can connect tokens at arbitrary distance in one step (unlike RNNs where information must travel sequentially).

Cross-attention: querying one sequence with another

In an encoder–decoder setup:

• •Encoder outputs memory H ∈ ℝ^(L_src×d_model)
• •Decoder has states/representations Y ∈ ℝ^(L_tgt×d_model)

Cross-attention uses:

So each decoder position forms a query based on what it has generated so far, and retrieves relevant source information.

Shape intuition:

• •Q: (L_tgt×dₖ)
• •K,V: (L_src×dₖ) and (L_src×dᵥ)
• •scores: (L_tgt×L_src)

This matrix is literally an alignment between target positions and source positions.

Where multi-head attention fits (high-level, because Transformers unlock next)

This node unlocks Transformers, where attention is typically multi-head.

Multi-head attention repeats the attention computation h times with different learned projections:

For head r:

• •Q_r = XW_Q^(r)
• •K_r = XW_K^(r)
• •V_r = XW_V^(r)

Each head produces O_r, then we concatenate and project:

Why multiple heads helps:

• •different heads can specialize (syntax, coreference, local context, long-range dependency)
• •it increases representational capacity without exploding per-head dimension

But the atomic mechanism remains exactly what you learned: score, softmax, weighted sum.

Deriving the batched, multi-head shapes (to prevent axis errors)

Let:

• •B = batch size
• •L_q = query length
• •L_k = key length
• •h = number of heads
• •dₖ = per-head key dimension
• •dᵥ = per-head value dimension

Typical shapes:

• •Q: (B×h×L_q×dₖ)
• •K: (B×h×L_k×dₖ)
• •V: (B×h×L_k×dᵥ)

Scores:

Softmax over L_k:

• •for each (B, head, query position), weights over keys sum to 1

Output:

• •O: (B×h×L_q×dᵥ)

This is where the earlier failure modes live:

• •wrong transpose can swap L_q and dₖ
• •softmax axis must be L_k
• •mask must broadcast to (B×h×L_q×L_k)

A concrete debugging checklist (practical connection)

When attention behaves oddly, check invariants:

1) Row sum invariant (per query):

• •For unmasked keys: ∑ weights = 1

2) Mask invariant:

• •masked positions have weights ≈ 0

3) Causality invariant (decoder self-attention):

• •for any query position t, weights on keys > t are ≈ 0

4) Sanity input test:

• •if all keys are identical, attention weights should become uniform (modulo mask)
• •if one key is much closer to the query, attention should be peaked there

How attention connects to the next node (Transformers)

Transformers stack attention layers with:

• •residual connections
• •layer normalization
• •position-wise feed-forward networks
• •positional encodings (to inject order)

But none of those change what attention is. They make it trainable, stable, and expressive at scale.

If you can derive the score matrix shape and explain why softmax is row-wise, you’re ready to understand multi-head attention, positional encoding, and full Transformer blocks.