Mid-level interview questionAsked at TikTok · interview write-up
Write pseudocode for single-head scaled dot-product self-attention with masks and tensor shapes, and for the cross-entropy loss.
Model answer
State shapes as you go: x is T by d, scores T by T. Mask before the softmax, and normalise stably.
def attention(x, Wq, Wk, Wv, mask): # x: (T, d)
q, k, v = x @ Wq, x @ Wk, x @ Wv # (T, d_k), (T, d_v)
s = q @ k.T / np.sqrt(k.shape[-1]) # (T, T)
s = np.where(mask, s, -1e9) # mask: True = keep
s = s - s.max(-1, keepdims=True)
p = np.exp(s)
p /= p.sum(-1, keepdims=True)
return p @ v # (T, d_v)
def cross_entropy(logits, y, ignore=-100): # (N, V), (N,)
keep = y != ignore
z, t = logits[keep], y[keep]
m = z.max(-1, keepdims=True)
lse = m[:, 0] + np.log(np.exp(z - m).sum(-1))
return (lse - z[np.arange(len(t)), t]).mean()
Scroll sideways to see the whole diagram
Rows are the token looking (its query), columns the token looked at (its key). Each cell is a score: the query's dot product with the key, divided by √d.
Text version
- Rows are the token looking (its query), columns the token looked at (its key). Each cell is a score: the query's dot product with the key, divided by √d.
- The causal mask: before the softmax, every score where the key comes after the query is set to minus infinity.
- Softmax turns those cells into exactly 0, and the remaining weights in each row still add up to 1.
- Token 1 sees only itself, token 2 sees tokens 1 and 2, and so on. Adding tokens later never changes an earlier row, which is why the KV cache is exact.
The causal mask as taught in lesson 3, on the attention lab's three tokens “The river bank”.
Cross-entropy is log-sum-exp minus the target logit, averaged over non-ignored tokens. Interviewers check the scale by the key width, masking scores rather than zeroing probabilities, and that padding labels are skipped.
Practise the rest
20 Attention questions, 492 in all, each with a model answer. Free.
More from TikTok interview questions.
More Attention interview questions
- Walk me through scaled dot-product attention. What are queries, keys and values, and why divide by the square root of d?
- What's inside one transformer block, and where do most of its parameters sit?
- What does the causal mask do, and what does it imply for a decoder model reading 'The ladder did not reach the window because it was too short'?
- A teammate shows you a GPT-2 head that puts most of a pronoun's attention on the right noun and concludes that the head resolves pronouns. How do you respond?
- Product wants to double the context you support, from 1,000 to 2,000 tokens. From the attention side, what gets more expensive, and what families of fixes would you weigh?
- Implement batched multi-head self-attention with a causal mask.