Multi-Head Attention With Causal Mask | Dataford Interview Questions

Write a PyTorch implementation of multi-head attention including a causal mask for autoregressive decoding.

For grading, represent tensors as nested lists with shape (batch, sequence_length, embedding_dimension). Return a nested list with shape (batch, query_length, embedding_dimension). Support rectangular key and query lengths, where cached keys precede the current query tokens.

Constraints

1 <= batch_size <= 8
1 <= query_length, key_length <= 256
1 <= num_heads <= embedding_dimension
embedding_dimension is divisible by num_heads
query, key, and value have matching batch dimensions
key and value have the same sequence length and embedding dimension
All tensor entries are finite numbers

Function Signature

def multi_head_attention(query, key, value, num_heads, causal):

Your question is Multi-Head Attention With Causal Mask. Start with the requirements in the Question tab.

Run and submit as often as you like. When you're ready, talk me through your approach or go straight to the code.

Problem

Write a PyTorch implementation of multi-head attention including a causal mask for autoregressive decoding.

Constraints

1 <= batch_size <= 8
1 <= query_length, key_length <= 256
1 <= num_heads <= embedding_dimension
embedding_dimension is divisible by num_heads
query, key, and value have matching batch dimensions
key and value have the same sequence length and embedding dimension
All tensor entries are finite numbers

Function Signature

def multi_head_attention(query, key, value, num_heads, causal):

Your question is Multi-Head Attention With Causal Mask. Start with the requirements in the Question tab.

Run and submit as often as you like. When you're ready, talk me through your approach or go straight to the code.