Maths for ML interview question · Capstone: embeddings and attention by hand · senior level
Compute scaled dot-product attention by hand for two tokens with d_k = 2.
Free to read in full. This is 1 of the 120 questions in the Maths for ML and Deep Learning interview pack, with 147 diagrams. Get the pack, ₹999 · $20
Why it is asked
Interviewers at LLM teams ask this to see that you know every step of attention, not just the formula.
The answer
Attention gives each token a new vector that is a weighted average of the value vectors, with weights from how well its query matches each key: Attention(Q, K, V) = softmax(QKᵀ / √d_k) V.
Symbols: Q (queries), K (keys) and V (values) hold one row per token. A query is what a token is looking for, a key is what each token offers, a value is the information it passes on. d_k is the length of each query and key; softmax turns a row of scores into positive weights that add up to 1: softmax(zᵢ) = exp(zᵢ) / Σⱼ exp(zⱼ).
Inputs (2 tokens, d_k = 2): Q = [[1, 0], [1, 1]], K = [[1, 0], [0, 1]], V = [[1, 2], [3, 4]].
1. Scores S = QKᵀ. Token 1: (1, 0)·(1, 0) = 1 and (1, 0)·(0, 1) = 0. Token 2: 1 and 1. So S = [[1, 0], [1, 1]]. 2. Divide by √2 ≈ 1.4142: [[0.7071, 0], [0.7071, 0.7071]]. 3. Softmax per row. Row 1: exp(0.7071) ≈ 2.0281 and exp(0) = 1, so the weights are 2.0281 / 3.0281 ≈ 0.6698 and 0.3302. Row 2: two equal scores give 0.5 and 0.5. 4. Output = weights × V. Token 1: 0.6698·(1, 2) + 0.3302·(3, 4) ≈ (1.6605, 2.6605). Token 2: 0.5·(1, 2) + 0.5·(3, 4) = (2, 3). (Outputs use the unrounded weights; all values rounded to 4 decimals.)
In a decoder such as GPT, a causal mask (which stops a token from looking at later tokens) sets the score for token 1 looking at token 2 to minus infinity before the softmax, so token 1 gets weights (1, 0) and output (1, 2).
What the interviewer listens for
- softmax(QKᵀ/√d_k)V, row by row.
- Worked: weights (0.6698, 0.3302) and (0.5, 0.5).
- Outputs (1.6605, 2.6605) and (2, 3).
- Causal mask uses minus infinity before softmax.
Follow-ups the interviewer may ask
- What changes in cross-attention? Hint: Q comes from one sequence, K and V from another, such as the encoder output.
- Why softmax rather than plain normalization? Hint: it handles negative scores and gives smooth gradients.
- How does the KV cache use this? Hint: during generation, earlier K and V rows are stored and only the new query row is computed.
Common mistakes
- Applying softmax over the whole matrix instead of each row.
- Forgetting the √d_k scaling or dividing by d_k.
Written by Roni Das, who builds distributed systems and AI systems for a living.
Learn it properly
This is taught step by step in the course lesson Attention: How Each Word Looks Back at the Words Before It.
One question is not enough for the interview.
This is 1 of the 120 questions in the Maths for ML and Deep Learning pack
Each question has why it is asked, a full answer, what the interviewer listens for, the follow-ups, the common mistakes and a diagram (147 diagrams in all). A PDF you keep, junior to staff level.
India: UPI, cards or netbanking through Razorpay. Elsewhere: PayPal. Download at once from My packs.
More free questions from this pack
- Compute the cosine similarity of a = (3, 4) and b = (4, 3) by hand. What does the result mean? (junior)
- A batch of 32 MNIST images goes into a dense layer with 128 units. What are the shapes, and what is the rule for when two matrices can be multiplied? (mid)
- A disease affects 1% of people. A test catches 90% of sick people (sensitivity) and wrongly flags 9% of healthy people (false positive rate). You test positive. What is the chance you are sick? (mid)
- Why do we not use mean squared error for classification with a sigmoid output? Compare the gradients numerically. (staff)
Free questions from the other packs
- What do temperature and top-p control, and what values would you use for data extraction versus brainstorming? (AI Engineering)
- What is hybrid search, and how do you merge keyword and vector results with Reciprocal Rank Fusion? (AI Engineering)
- Users ask "list all X" questions and RAG returns only five items. Why, and what do you do? (AI Engineering)
- How do you check that an LLM judge agrees with human reviewers? (AI Engineering)
- Cut the cost of an LLM feature serving 10 million requests a day by 60% without hurting quality. Walk through your plan. (AI Engineering)
- Why are LLMs bad at counting letters or reversing a word? (GenAI and LLM)
- Write out scaled dot-product attention with the shape of every tensor. (GenAI and LLM)
- How does speculative sampling keep exactly the target model's distribution? (GenAI and LLM)
- What is GRPO, and how do verifiable rewards train reasoning models? (GenAI and LLM)
- Design a system that answers questions from 300-page contracts accurately. (GenAI and LLM)
- What happens, step by step, when an LLM server handles one request? (Inference Engineering)
- Derive the KV cache size formula and work it out for Llama 3 70B. (Inference Engineering)
- What is PagedAttention? (Inference Engineering)
- When does speculative decoding stop helping or make things worse? (Inference Engineering)
- Design a service to serve Llama 3 70B at 10 QPS with latency targets. (Inference Engineering)
- A customer says: "We want AI for our support team." What do you do in the first meeting? (Forward Deployed Engineer)
- The customer's admin offers you a global admin account "to save time". What do you ask for instead, and why? (Forward Deployed Engineer)
- Your connector writes tickets into the customer's system. A request timed out, the connector retried, and now there are duplicate tickets. How do you make the writes safe to retry? (Forward Deployed Engineer)
- The customer wants RAG over 10 million documents, all behind their firewall. Size it for me on the whiteboard. (Forward Deployed Engineer)
- One deployment of your AI platform must serve 50 business units of the same enterprise. How do you design multi-tenancy? (Forward Deployed Engineer)