Skip to the solution
MediumAI Engineering15 min read

How does the transformer architecture and attention work?

99 views

The problem

Question presented to candidate: "Walk me through the transformer. What happens to a sentence from the moment it is tokenized until the model outputs the next token, and what exactly is attention computing?"

What a strong answer should cover:

  • The pipeline: token embeddings plus position information, a stack of identical blocks (multi-head self-attention, then a feed-forward network, each wrapped in a residual connection and normalisation), then a projection to logits over the vocabulary
  • Attention itself: each token produces a query, a key and a value; scores are query-key dot products scaled by the square root of the key size, softmax turns each row into weights, and the output is the weighted sum of values: softmax(QKแต€/โˆšd_k)V
  • Why it works: every token can pull information from any other token in one step, and all positions are computed in parallel, unlike an RNN
  • The causal mask in decoder-only (GPT-style) models, and the three families: encoder-only, encoder-decoder, decoder-only
  • Multi-head attention, and why position must be injected (attention alone ignores order) with sinusoidal, learned or rotary (RoPE) encodings
  • The cost: attention scores grow with the square of sequence length, which is what makes long context expensive and motivates the KV cache and efficient attention kernels

Code / implementation expected: Optional โ€” computing attention for a few tokens by hand (or in a few lines of code) is the best proof of understanding.

  • transformer
  • attention
  • fundamentals

Try it yourself

Sketch your approach before you read the solution. That is what the interview tests.

Stuck? Get a hint.

A short nudge toward the approach. It will not give the answer away.

Solution

The solution is hidden so you can try first

15 min read ยท 12 sections

Read the code

Scaled dot-product attention by hand: Q, K, V, causal mask, positions and the n-squared cost
// Scaled dot-product attention, by hand, on a 4-token sentence.
// One head, d_model = 4, d_k = 2. Numbers are small so you can follow them.
const words = ["the", "cat", "sat", "down"];
const emb = {
  the: [1, 0, 0, 0], cat: [0, 1, 1, 0], sat: [0, 0, 1, 1], down: [1, 0, 0, 1],
};
// Learned projection matrices (4 x 2). In a real model these come from training.
const Wq = [[1, 0], [0, 1], [1, 1], [0, 0]];
const Wk = [[0, 1], [1, 0], [1, 1], [0, 1]];
const Wv = [[1, 0], [0, 2], [2, 0], [0, 1]];

const matmul = (A, B) => A.map((row) => B[0].map((_, j) => row.reduce((s, a, k) => s + a * B[k][j], 0)));
const T = (A) => A[0].map((_, j) => A.map((row) => row[j]));
const f3 = (x) => (Object.is(x, -0) ? 0 : x).toFixed(3);
const show = (name, M, labels) => {
  console.log(name);
  M.forEach((row, i) => console.log("  " + (labels ? labels[i].padEnd(5) : "") + row.map((x) => (x === -Infinity ? "-inf" : f3(x)).padStart(7)).join("")));
};

function attention(X, causal, labels, verbose) {
  const Q = matmul(X, Wq), K = matmul(X, Wk), V = matmul(X, Wv);
  const dk = Q[0].length;
  const scores = matmul(Q, T(K)).map((row) => row.map((s) => s / Math.sqrt(dk)));
  if (causal) scores.forEach((row, i) => row.forEach((_, j) => { if (j > i) row[j] = -Infinity; }));
  const weights = scores.map((row) => {
    const m = Math.max(...row);
    const e = row.map((s) => Math.exp(s - m));
    const z = e.reduce((a, b) => a + b, 0);
    return e.map((x) => x / z);
  });
  const out = matmul(weights, V);
  if (verbose) {
    show("Q = X * Wq", Q, labels);
    show("K = X * Wk", K, labels);
    show("scores = Q * K^T / sqrt(" + dk + ")" + (causal ? " with causal mask" : ""), scores, labels);
    show("weights = softmax(scores) per row (each row sums to 1)", weights, labels);
    show("output = weights * V", out, labels);
  }
  return out;
}

const X = words.map((w) => emb[w]);
attention(X, true, words, true);

// Without position information, attention cannot tell word order apart.
const a = ["the", "cat", "sat"], b = ["sat", "cat", "the"];
const outA = attention(a.map((w) => emb[w]), false, a, false);
const outB = attention(b.map((w) => emb[w]), false, b, false);
console.log("no positions:   cat in [" + a.join(" ") + "] -> " + outA[1].map(f3).join(", ") + " | in [" + b.join(" ") + "] -> " + outB[1].map(f3).join(", "));

// Add a simple position vector to each embedding and the outputs differ.
const pos = (i) => [Math.sin(i), Math.cos(i), Math.sin(i / 10), Math.cos(i / 10)];
const addPos = (seq) => seq.map((w, i) => emb[w].map((x, d) => x + pos(i)[d]));
const pA = attention(addPos(a), false, a, false), pB = attention(addPos(b), false, b, false);
console.log("with positions: cat in [" + a.join(" ") + "] -> " + pA[1].map(f3).join(", ") + " | in [" + b.join(" ") + "] -> " + pB[1].map(f3).join(", "));

// Cost: every token scores every other token, so the score matrix is n x n.
for (const n of [1024, 8192, 131072]) {
  console.log("n=" + String(n).padStart(6) + " tokens -> " + (n * n).toLocaleString("en-US") + " scores per head per layer");
}

Discussion

Other approaches, follow-ups and how it went in real interviews.

Sign in to join the discussion.

No comments yet. Share how you would answer it.

Question 58 of 83 in AI engineering

Keep your momentum going.

Mark this question solved to track your progress, then take the next one while it is fresh.

Up nextMedium

What is the ReAct pattern and why did it become the default for agents?

Continue