How does attention work in a transformer?

TL;DR

• Attention lets each token look at other tokens using queries $Q$, keys $K$, and values $V$ to compute a weighted sum of information.
• We score similarity with dot products $q_i \cdot k_j$, scale by $\sqrt{d_k}$, apply softmax, and multiply by $V$ to get outputs.
• Scaling by $\sqrt{d_k}$ keeps the softmax from becoming too peaky as $d_k$ grows; this stabilizes gradients.
• Multi head attention runs several of these in parallel with different linear projections, letting the model focus on different relationships.

Key tradeoffs: simpler attention is flexible and parallelizable but costs $O(n^2)$ memory/time in sequence length and can struggle with very long contexts.

In this question, we will learn how the attention mechanism in a transformer turns token embeddings into context-aware representations. We will walk from a simple intuition to the formal formula, then run a small worked example on the sentence "The cat sat on the mat".

We will cover the following:

• The intuition (an analogy)
• How it actually works
• Why we scale by $\sqrt{d_k}$
• Multi head attention
• Tradeoffs and failure modes
• Questions the interviewer might ask
• What the interviewer is really testing

Direct answer: Attention lets each token form a query $q_i$ that compares to keys $k_j$ of all tokens, producing softmax weights that mix the values $v_j$ into a context vector. The core computation is a scaled dot-product followed by softmax and a weighted sum. This is computed efficiently with matrix operations on $Q$, $K$, and $V$.

The intuition (an analogy)

Think of a classroom where each student holds a sticky note with facts. When one student wants to answer a question, they look at other students' notes. Their question is a query $q_i$. Each other student presents a key $k_j$ that expresses what their note contains. The similarity between the question and a note determines how much the asking student reads that note. That reading is the weight, and the actual content copied is the value $v_j$. The final answer is a weighted combination of those notes.

This lets every token selectively copy information from any other token, whether that information is nearby or far away in the sequence.

How it actually works

The formal, compact formula for scaled dot-product attention is:

$\text{Attention}(Q,K,V) = \text{softmax}\!\left(\frac{QK^T}{\sqrt{d_k}}\right)V$

Here $Q$, $K$, and $V$ are matrices of queries, keys, and values; $d_k$ is the dimensionality of the keys. The inner product $QK^T$ gives all pairwise scores.

Worked example on "The cat sat on the mat". For clarity we use $d_k=2$ and set embeddings so that $Q=K=V$ (a common simplification in a single self-attention layer).

Token order: The | cat | sat | on | the | mat We will compute attention for the query of the token "sat", which we label $q_{\text{sat}}$.

Token vectors (rows are components):

So $q_{\text{sat}} = [0, 1]$. The dot-product scores $s_j = q_{\text{sat}}\cdot k_j$ are:

We scale by $\sqrt{d_k}$ with $d_k=2$, so $\sqrt{d_k} \approx 1.414$. The scaled scores are:

Apply softmax to the scaled scores to get weights $a_j$:

Exp of scaled scores (approx):

Softmax weights (normalized):

Finally, the attention output for the query $q_{\text{sat}}$ is the weighted sum of the values $v_j$:

$\text{output}_{\text{sat}} = \sum_j a_j v_j \approx [0.643,\;0.562].$

This output is a context-aware vector for "sat" that mixes nearby lexical items like "cat" and "mat" according to the learned similarities.

Why we scale by $\sqrt{d_k}$

When $d_k$ grows, typical dot products $q_i \cdot k_j$ have variance proportional to $d_k$. Without scaling, as $d_k$ increases the raw scores become larger in magnitude and softmax turns into a near one-hot distribution, which hurts gradient flow. Dividing by $\sqrt{d_k}$ normalizes the scale so the softmax operates in a stable range.

A small table showing how $e^x$ grows and why large $x$ dominate softmax:

Because $e^x$ increases rapidly, a modest increase in raw dot-product magnitude can make one token dominate the weights. The $\sqrt{d_k}$ factor keeps the $x$ values small enough that multiple tokens contribute.

Multi head attention

Multi head attention runs several parallel attention computations with different learned linear projections of $Q$, $K$, and $V$. Each head has its own subspace (often smaller, e.g., $d_k = d_{\text{model}}/h$). The separate heads can capture different kinds of relationships, such as syntactic connections in one head and semantic co-reference in another. The per-head outputs are concatenated and projected back to the model dimension.

Tradeoffs and failure modes

Questions the interviewer might ask

Some follow-up questions you might get:

• Why use dot product and not cosine similarity? Dot product is efficient to compute in matrix form and, with learned projections, the model can scale and shift vectors to represent useful comparisons; cosine would require extra normalization steps that complicate gradients.
• What is the role of positional encoding? Attention is permutation invariant, so positional encodings add order information so the model can distinguish "cat sat" from "sat cat".
• How does attention handle variable-length sequences? Attention naturally handles variable length because $QK^T$ just computes pairwise scores over whatever tokens are present; padding and masks are used to ignore positions.
• Why have multiple heads instead of one wide head? Multiple heads let the model attend to different subspaces and relationships in parallel; splitting often helps learning more diverse patterns than a single head of equal total size.
• How do masks work in attention? A mask adds large negative values (e.g., $-10^9$) to $QK^T$ before softmax to zero out unwanted positions, used for padding or causal/auto-regressive attention.
• What happens if $\sqrt{d_k}$ is omitted? Softmax can become extremely peaky for large $d_k$, causing vanishing gradients and learning difficulties.

Some things to note:

• Attention parameters are learned via the linear projections that produce $Q$, $K$, and $V$.
• In practice we use residual connections and layer norm around attention for stable training.

What the interviewer is really testing

They want to see that you understand both the math and the intuition: how queries, keys, and values interact, why scaling matters, and what practical consequences attention has for computation and model behavior. Showing that you can walk through a small numeric example and mention implementation details like masking and complexity signals depth of understanding.