Mid-level interview questionAsked at Meta · interview write-up
Implement self-attention, give its time and memory complexity, and describe the efficient variants.
Model answer
def self_attn(x, wq, wk, wv): # x: (T, d)
q, k, v = x @ wq, x @ wk, x @ wv
s = q @ k.T / q.shape[-1] ** 0.5 # (T, T)
s = s - s.max(-1, keepdims=True)
p = np.exp(s)
return (p / p.sum(-1, keepdims=True)) @ v
Time is O(T² d) and the score matrix takes O(T²) memory, against O(T d²) for projections, so long sequences are attention-dominated. FlashAttention stays exact but tiles the work so the T-by-T matrix never reaches HBM. Sliding-window and sparse patterns make it linear in T; linear and low-rank attention approximate; MQA and GQA shrink the KV cache.
Practise the rest
20 Attention questions, 492 in all, each with a model answer. Free.
More from Meta 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.