Mid-level interview questionAsked at Mistral AI, ByteDance, Meta, Apple, Uber · interview write-up
Implement batched multi-head self-attention with a causal mask.
Model answer
Project once, split into heads, attend in each head, merge.
def mhsa(x, w_qkv, w_o, h): # x: (B, T, d)
B, T, d = x.shape
q, k, v = (x @ w_qkv).split(d, -1)
heads = lambda t: t.view(B, T, h, -1).transpose(1, 2)
q, k, v = heads(q), heads(k), heads(v) # (B, h, T, d/h)
s = q @ k.transpose(-2, -1) / (d // h) ** 0.5
future = torch.ones(T, T, dtype=torch.bool,
device=x.device).triu(1)
s = s.masked_fill(future, float("-inf"))
y = s.softmax(-1) @ v
return y.transpose(1, 2).reshape(B, T, d) @ w_o
Interviewers watch for the mask before softmax, scaling by the head size, not d, and the transpose before the final reshape.
Practise the rest
20 Attention questions, 492 in all, each with a model answer. Free.
More from Mistral AI interview questions, ByteDance interview questions, Meta interview questions, Apple interview questions, Uber 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?
- A pre-norm transformer block trains badly. What do you check around the norms and residuals?