Self-attention mechanism in PyTorch - Practice Problems & Coding Challenges

Challenge - 5 Problems

❓ Predict Output

intermediate

Output of scaled dot-product attention calculation

Given the following PyTorch code snippet implementing scaled dot-product attention, what is the output tensor?

PyTorch

import torch
import torch.nn.functional as F

query = torch.tensor([[1., 0., 0.]])  # shape (1, 3)
key = torch.tensor([[1., 0., 0.], [0., 1., 0.]])  # shape (2, 3)
value = torch.tensor([[1., 2.], [3., 4.]])  # shape (2, 2)

# Compute attention scores
scores = torch.matmul(query, key.T) / (3 ** 0.5)
weights = F.softmax(scores, dim=-1)
output = torch.matmul(weights, value)
print(output)

A [[1.9999, 2.9999]]

B [[3.0, 4.0]]

C [[1.718, 2.718]]

D [[1.5, 2.5]]

💡 Hint

Recall that softmax normalizes scores to probabilities, and output is weighted sum of values.

The query matches perfectly with the first key vector, so the attention weight for the first key is near 1, and near 0 for the second. The output is close to the first value vector [1., 2.]. After scaling and softmax, the output is approximately [[1.718, 2.718]].

❓ Model Choice

intermediate

Choosing the correct self-attention output shape

In a self-attention layer, if the input tensor has shape (batch_size=4, seq_len=10, embedding_dim=64), what will be the shape of the output tensor after applying multi-head self-attention with 8 heads and the same embedding dimension?

A (4, 10, 512)

B (4, 10, 64)

C (4, 8, 10, 8)

D (4, 64, 10)

💡 Hint

Multi-head attention splits embedding_dim into heads but concatenates back to original embedding_dim.

Multi-head attention splits the embedding dimension into 8 heads of size 8 each (64/8=8), processes attention in parallel, then concatenates back to shape (batch_size, seq_len, embedding_dim), which is (4, 10, 64).

❓ Hyperparameter

advanced

Effect of increasing number of attention heads

What is the most likely effect of increasing the number of attention heads in a multi-head self-attention model while keeping the total embedding dimension fixed?

A The model's total capacity increases linearly with the number of heads.

B The model becomes unable to learn because the embedding dimension is fixed.

C The embedding dimension per head increases, improving feature representation.

D Each head attends to a smaller subspace, potentially capturing more diverse features.

💡 Hint

Think about how embedding dimension is split among heads.

Increasing heads splits the fixed embedding dimension into smaller parts per head, allowing each head to focus on different aspects of the input, improving diversity of attention.

🔧 Debug

advanced

Identifying error in self-attention implementation

Consider this PyTorch code snippet for self-attention. What error will it raise when executed?

PyTorch

import torch
import torch.nn.functional as F

query = torch.randn(2, 5, 16)
key = torch.randn(2, 5, 16)
value = torch.randn(2, 5, 16)

scores = torch.matmul(query, key) / (16 ** 0.5)
weights = F.softmax(scores, dim=-1)
output = torch.matmul(weights, value)
print(output.shape)

A RuntimeError due to shape mismatch in torch.matmul(query, key)

B RuntimeError due to shape mismatch in torch.matmul(weights, value)

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

D SyntaxError due to missing colon

💡 Hint

Check the dimensions of query and key for matrix multiplication.

query has shape (2,5,16) and key has shape (2,5,16). torch.matmul expects the last two dims to be compatible. Here, 16 != 5, so it raises RuntimeError.

🧠 Conceptual

expert

Why use scaled dot-product in self-attention?

Why do self-attention mechanisms scale the dot product of query and key vectors by the square root of their dimension?

A To prevent the dot products from growing too large, which can cause softmax gradients to vanish.

B To increase the magnitude of dot products, making attention weights sharper.

C To normalize the query and key vectors to unit length before dot product.

D To reduce computational cost by scaling down the vectors.

💡 Hint

Think about how large dot products affect softmax.

Without scaling, large dot products can push softmax into regions with very small gradients, making training harder. Scaling by sqrt(dim) keeps values in a range that stabilizes gradients.

Practice

1. What is the main purpose of the self-attention mechanism in models like Transformers?

easy

A. To convert input data into images

B. To increase the size of the input data

C. To reduce the number of layers in the model

D. To help the model focus on important parts of the input data

Solution

1. Step 1: Understand self-attention's role

   Self-attention helps the model decide which parts of the input are important by comparing each part to others.

2. Step 2: Identify the main goal

   The main goal is to focus on important parts, not to change input size or model layers.

3. Final Answer:

   To help the model focus on important parts of the input data -> Option D

4. Quick Check:

   Self-attention = focus on important parts [OK]

Hint: Self-attention means focusing on important input parts [OK]

Common Mistakes:

• Thinking self-attention changes input size
• Confusing self-attention with model depth
• Assuming it converts data formats

2. Which of the following is the correct PyTorch code snippet to compute scaled dot-product attention scores given query Q and key K tensors?

easy

A. scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(Q.size(-1))

B. scores = torch.matmul(K, Q.transpose(-2, -1)) * math.sqrt(Q.size(-1))

C. scores = torch.matmul(Q, K) / Q.size(0)

D. scores = torch.matmul(Q.transpose(-2, -1), K) / math.sqrt(K.size(-1))

Solution

1. Step 1: Recall scaled dot-product attention formula

   Attention scores = (Q x K^T) / sqrt(d_k), where d_k is the dimension of keys.

2. Step 2: Match formula with code

   scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(Q.size(-1)) correctly transposes K, multiplies Q and K^T, and divides by sqrt of last dimension size.

3. Final Answer:

   scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(Q.size(-1)) -> Option A

4. Quick Check:

   Scaled dot-product = QxK^T / sqrt(d_k) [OK]

Hint: Transpose keys and divide by sqrt of key dimension [OK]

Common Mistakes:

• Multiplying K by Q instead of Q by K^T
• Forgetting to transpose keys
• Using multiplication instead of division by sqrt

3. Given the following PyTorch code for self-attention, what is the shape of the output tensor attn_output?

import torch
Q = torch.randn(2, 4, 8)  # batch=2, seq_len=4, dim=8
K = torch.randn(2, 4, 8)
V = torch.randn(2, 4, 8)

scores = torch.matmul(Q, K.transpose(-2, -1)) / (8 ** 0.5)
weights = torch.nn.functional.softmax(scores, dim=-1)
attn_output = torch.matmul(weights, V)

medium

A. (4, 2, 8)

B. (2, 4, 8)

C. (2, 8, 4)

D. (4, 8, 2)

Solution

1. Step 1: Analyze input tensor shapes

   Q, K, V have shape (batch=2, seq_len=4, dim=8).

2. Step 2: Understand matrix multiplications

   Scores shape: (2,4,4), weights shape: (2,4,4), multiplying weights (2,4,4) by V (2,4,8) results in (2,4,8).

3. Final Answer:

   (2, 4, 8) -> Option B

4. Quick Check:

   Output shape = (batch, seq_len, dim) = (2,4,8) [OK]

Hint: Output shape matches Q and V batch and seq_len dims [OK]

Common Mistakes:

• Confusing batch and sequence dimensions
• Mixing up transpose effects on shape
• Assuming output shape changes dimension size

4. Identify the error in this PyTorch self-attention code snippet:

import torch
Q = torch.randn(1, 3, 5)
K = torch.randn(1, 3, 5)
V = torch.randn(1, 3, 5)

scores = torch.matmul(Q, K) / (5 ** 0.5)
weights = torch.nn.functional.softmax(scores, dim=-1)
output = torch.matmul(weights, V)

medium

A. Q and V should have the same shape

B. Softmax should be applied on dim=0 instead of dim=-1

C. K should be transposed before multiplication with Q

D. Division by sqrt(5) is incorrect

Solution

1. Step 1: Check matrix multiplication for scores

   Q shape is (1,3,5), K shape is (1,3,5). To multiply Q by K^T, K must be transposed on last two dims.

2. Step 2: Identify missing transpose

   Code multiplies Q by K directly, causing shape mismatch or wrong attention scores.

3. Final Answer:

   K should be transposed before multiplication with Q -> Option C

4. Quick Check:

   QxK^T needed, missing transpose causes error [OK]

Hint: Always transpose keys before multiplying with queries [OK]

Common Mistakes:

• Skipping transpose of K
• Applying softmax on wrong dimension
• Confusing shapes of Q and V

5. You want to implement multi-head self-attention in PyTorch with 4 heads. If your input tensor has shape (batch=2, seq_len=6, dim=32), what should be the dimension size of each head's query, key, and value vectors?

hard

A. 8

B. 4

C. 16

D. 32

Solution

1. Step 1: Understand multi-head attention dimension split

   Total dimension 32 is split evenly across 4 heads, so each head gets 32 / 4 = 8 dimensions.

2. Step 2: Confirm dimension per head

   Each head's query, key, and value vectors have dimension 8 to keep total dimension consistent.

3. Final Answer:

   8 -> Option A

4. Quick Check:

   32 total dim / 4 heads = 8 per head [OK]

Hint: Divide total dim by number of heads for each head size [OK]

Common Mistakes:

• Using total dimension for each head
• Dividing by wrong number of heads
• Confusing batch size with dimension