Multi-head attention in PyTorch - Practice Problems & Coding Challenges

Challenge - 5 Problems

❓ Predict Output

intermediate

Output shape of multi-head attention

Given the following PyTorch code snippet using multi-head attention, what is the shape of the output tensor?

PyTorch

import torch
import torch.nn as nn

batch_size = 2
seq_len = 5
embed_dim = 16
num_heads = 4

mha = nn.MultiheadAttention(embed_dim=embed_dim, num_heads=num_heads)
query = torch.rand(seq_len, batch_size, embed_dim)
key = torch.rand(seq_len, batch_size, embed_dim)
value = torch.rand(seq_len, batch_size, embed_dim)
out, _ = mha(query, key, value)
print(out.shape)

A torch.Size([5, 2, 16])

B torch.Size([2, 5, 16])

C torch.Size([5, 16, 2])

D torch.Size([2, 16, 5])

💡 Hint

Remember that PyTorch's MultiheadAttention expects input shape (sequence_length, batch_size, embedding_dim) and outputs the same shape.

The MultiheadAttention module in PyTorch takes inputs of shape (seq_len, batch_size, embed_dim) and returns output of the same shape. So the output shape is (5, 2, 16).

❓ Model Choice

intermediate

Choosing number of heads in multi-head attention

You have an embedding dimension of 64 for your model. Which choice of number of heads is valid for PyTorch's MultiheadAttention module?

A 10

B 7

C 9

D 8

💡 Hint

The embedding dimension must be divisible by the number of heads.

The number of heads must divide the embedding dimension evenly. 64 divided by 8 is 8, which is an integer. 10, 7, and 9 do not divide 64 evenly.

❓ Hyperparameter

advanced

Effect of increasing number of heads on model capacity

What is the most accurate effect of increasing the number of heads in a multi-head attention layer while keeping the embedding dimension fixed?

A It increases the model's ability to focus on different parts of the input simultaneously without increasing the total embedding size.

B It increases the embedding dimension, making the model larger and slower.

C It decreases the model's ability to learn because each head gets fewer parameters.

D It reduces the number of parameters by sharing weights across heads.

💡 Hint

Think about how multi-head attention splits the embedding dimension into multiple heads.

Increasing heads splits the embedding dimension into more parts, allowing the model to attend to different information in parallel without increasing total embedding size.

🔧 Debug

advanced

Identifying error in multi-head attention input shapes

What error will this PyTorch code raise when running the multi-head attention layer?

PyTorch

import torch
import torch.nn as nn

mha = nn.MultiheadAttention(embed_dim=32, num_heads=4)
query = torch.rand(10, 3, 32)
key = torch.rand(10, 3, 16)
value = torch.rand(10, 3, 32)
out, _ = mha(query, key, value)

A TypeError: Expected 3D tensor for query

B RuntimeError: The key and query must have the same embedding dimension

C ValueError: Number of heads must divide embedding dimension

D No error, code runs successfully

💡 Hint

Check the embedding dimension of key and query tensors.

The key tensor has embedding dimension 16, while query has 32. They must match for multi-head attention to work.

🧠 Conceptual

expert

Why use multi-head attention instead of single-head attention?

Which of the following best explains the main advantage of multi-head attention over single-head attention?

A Multi-head attention reduces the computational cost compared to single-head attention.

B Multi-head attention always produces smaller output tensors than single-head attention.

C Multi-head attention allows the model to jointly attend to information from different representation subspaces at different positions.

D Multi-head attention eliminates the need for positional encoding in sequence models.

💡 Hint

Think about how splitting into multiple heads affects the model's focus.

Multi-head attention splits the embedding into multiple heads, each learning different aspects of the input, improving representation power.