Katkat
Back

Implement multi-head attention and LLM sampling

Scale AIMachine Learning Engineer · Onsite
EasyMachine Learning EngineerOnsiteCoding & Algorithms

Last updated

This problem tests your ability to implement the core mechanics of a Transformer's attention mechanism and the logic required to sample the next token from a model's output logits.

Multi-Head Attention and LLM Sampling

You are building a simplified inference engine for a Transformer model. You need to implement two core components: scaled dot-product attention and a sampling function for autoregressive generation.

Part 1: Scaled Dot-Product Attention

Implement a function that calculates attention weights and the resulting context vector. Given query Q, key K, and value V tensors (all 2D matrices of shape [seq_len, d_k]), calculate the output.

output = softmax((Q @ K.T) / sqrt(d_k)) @ V

Ensure numerical stability by subtracting the maximum value from the logits before exponentiation.

Part 2: Sampling

Implement a function that selects a token index given raw logits. The function must support:

  1. Temperature: Scale the logits before softmax (division by T).
  2. Top-k: Only keep the top k tokens, setting the rest to negative infinity.

Constraints

  • d_k is the dimension of the keys.
  • Use NumPy for matrix operations.
  • Handle cases where k is greater than the vocabulary size.

Example 1: Attention

# Q, K, V are [2, 3]
# Output should be [2, 3]

Example 2: Sampling

logits = [0.1, 0.5, 0.8, 0.2]
temperature = 0.5
k = 2
# Output should be one of the top-k indices based on the modified distribution

Hints

Hint 1

For the attention mechanism, remember that subtracting the row-wise maximum from logits before computing the softmax does not change the result but prevents overflow.

Hint 2

When implementing top-k sampling, you can use `np.argpartition` to efficiently find the indices of the largest elements without fully sorting the array.

Hint 3

Temperature scaling is applied by dividing logits by `T` before the softmax; as `T` approaches 0, the distribution becomes a one-hot vector at the maximum logit.

Follow-up questions

  • How would you modify the attention mechanism to support causal masking, ensuring the model cannot attend to future tokens?
  • Explain the trade-offs between Top-k and Top-p (nucleus) sampling regarding diversity and quality of generated text.
  • What are the memory complexity implications of calculating the full attention matrix in long-context scenarios, and what techniques (e.g., FlashAttention) help mitigate this?