Katkat
Back

Implement Masked Multi-Head Self-Attention

AppleMachine Learning Engineer · Technical Screen
EasyMachine Learning EngineerTechnical ScreenMachine Learning

Last updated

This question tests the ability to implement the core mechanics of a Transformer's multi-head attention mechanism, specifically focusing on linear projections, attention weight scaling, and causal masking.

Problem Statement

Implement a MultiHeadAttention class in Python using a framework like PyTorch or NumPy. Your implementation must support scaled dot-product attention with causal masking to prevent tokens from attending to future positions.

Requirements

  1. Linear Projections: Given an input embedding tensor x of shape (batch_size, seq_len, d_model), create projections for Query, Key, and Value using linear layers to reach a dimension of d_model.
  2. Multi-Head Split: Split the projected tensors into n_heads. Each head will have a dimension of d_k = d_model // n_heads.
  3. Scaled Dot-Product Attention: Compute Attention(Q, K, V) = softmax((Q @ K.transpose) / sqrt(d_k) + mask) @ V.
  4. Causal Masking: Apply a lower-triangular mask (filled with -inf or a very large negative number) to the attention scores before the softmax to ensure position i only attends to positions <= i.
  5. Concatenation: Concatenate the outputs of all heads and project them back to d_model.

Signature

class MultiHeadAttention:
    def __init__(self, d_model, n_heads):
        # Initialize linear layers for Q, K, V and output projection
        pass

    def forward(self, x):
        # Return the attended output tensor of shape (batch_size, seq_len, d_model)
        pass

Example

Input x has shape (1, 3, 8) (batch size 1, sequence length 3, embedding size 8). With n_heads=2, each head processes a dimension of 4. The causal mask ensures the third token cannot attend to the first or second token's future information.

Constraints

  • d_model is divisible by n_heads.
  • The output shape must match the input shape.

Hints

Hint 1

Before computing the softmax, create a `(seq_len, seq_len)` matrix where the upper triangle is filled with negative infinity.

Hint 2

Ensure you reshape the tensors from `(batch, seq_len, d_model)` to `(batch, n_heads, seq_len, d_k)` to allow parallel computation across heads.

Hint 3

Remember to scale the dot product of Q and K by `1 / sqrt(d_k)` to maintain stable gradients during training.

Follow-up questions

  • How would you modify this implementation to support padding masks for variable-length sequences in a batch?
  • Why is the scaling factor `1 / sqrt(d_k)` critical when the dimensionality of the embeddings increases?
  • If this model were used for inference in a large language model, how would you optimize the Key-Value (KV) cache to avoid redundant computation?