Multi-head attention in PyTorch - Model Metrics & Evaluation

Metrics & Evaluation - Multi-head attention

Which metric matters for Multi-head attention and WHY

Multi-head attention is often used in tasks like language translation or text understanding. The key metrics to check are accuracy for classification tasks, BLEU score for translation quality, and loss to see how well the model learns. Accuracy tells us how many predictions are correct. Loss shows how far off predictions are during training. For sequence tasks, BLEU score measures how close the output is to the correct sentence. These metrics help us know if the attention mechanism is helping the model focus on important parts of the input.

Confusion matrix example for classification with Multi-head attention

      | Predicted Positive | Predicted Negative |
      |--------------------|--------------------|
      | True Positive (TP) = 85  | False Negative (FN) = 15 |
      | False Positive (FP) = 10 | True Negative (TN) = 90  |

      Total samples = 85 + 15 + 10 + 90 = 200

      Precision = TP / (TP + FP) = 85 / (85 + 10) = 0.8947
      Recall = TP / (TP + FN) = 85 / (85 + 15) = 0.85
      F1 Score = 2 * (Precision * Recall) / (Precision + Recall) = 2 * (0.8947 * 0.85) / (0.8947 + 0.85) ≈ 0.871
    

This matrix helps us understand how well the multi-head attention model classifies inputs by showing correct and incorrect predictions.

Precision vs Recall tradeoff with Multi-head attention

Imagine a spam email detector using multi-head attention to focus on important words. If we want high precision, the model marks emails as spam only when very sure, so fewer good emails are wrongly marked. But it might miss some spam emails (lower recall).

If we want high recall, the model catches almost all spam emails but might wrongly mark some good emails as spam (lower precision).

Choosing precision or recall depends on what is worse: missing spam or wrongly blocking good emails. Multi-head attention helps by focusing on key parts of the email to improve both metrics.

Good vs Bad metric values for Multi-head attention use cases

• Good: Accuracy above 85%, Precision and Recall both above 80%, and steady loss decrease during training.
• Bad: Accuracy below 60%, Precision or Recall below 50%, or loss that does not improve or fluctuates wildly.

Good metrics mean the multi-head attention is helping the model understand important input parts. Bad metrics suggest the model is not learning well or focusing incorrectly.

Common pitfalls in metrics for Multi-head attention

• Accuracy paradox: High accuracy can be misleading if data is imbalanced (e.g., many more negatives than positives).
• Data leakage: If test data leaks into training, metrics look better but model fails in real use.
• Overfitting: Low training loss but high test loss means model memorizes training data but does not generalize.
• Ignoring sequence metrics: For tasks like translation, only accuracy is not enough; use BLEU or similar.

Self-check question

Your multi-head attention model has 98% accuracy but only 12% recall on fraud detection. Is it good for production? Why or why not?

Answer: No, it is not good. The model misses 88% of fraud cases (low recall), which is dangerous. High accuracy is misleading because most transactions are not fraud. For fraud detection, high recall is critical to catch as many frauds as possible.

Key Result

Precision, recall, and loss are key metrics to evaluate multi-head attention models, ensuring they focus well and generalize.

Practice

1. What is the main purpose of multi-head attention in PyTorch's MultiheadAttention layer?

easy

A. To reduce the size of the input data

B. To perform simple addition of input vectors

C. To increase the number of layers in the model

D. To allow the model to focus on different parts of the input simultaneously

Solution

1. Step 1: Understand the concept of multi-head attention

   Multi-head attention splits the input into multiple parts (heads) to focus on different information simultaneously.

2. Step 2: Relate to PyTorch's MultiheadAttention layer

   This layer implements this idea to improve model understanding by attending to various input features at once.

3. Final Answer:

   To allow the model to focus on different parts of the input simultaneously -> Option D

4. Quick Check:

   Multi-head attention = focus on multiple input parts [OK]

Hint: Multi-head means multiple focus points at once [OK]

Common Mistakes:

• Thinking it reduces input size
• Confusing it with adding layers
• Assuming it just adds vectors

2. Which of the following is the correct way to create a multi-head attention layer with 8 heads and embedding dimension 64 in PyTorch?

easy

A. torch.nn.MultiheadAttention(8, 64)

B. torch.nn.MultiheadAttention(embed_dim=64, num_heads=8)

C. torch.nn.MultiheadAttention(num_heads=64, embed_dim=8)

D. torch.nn.MultiheadAttention(64)

Solution

1. Step 1: Check PyTorch MultiheadAttention parameter order

   The constructor requires embed_dim first, then num_heads as keyword arguments.

2. Step 2: Verify correct syntax

   torch.nn.MultiheadAttention(embed_dim=64, num_heads=8) correctly uses embed_dim=64 and num_heads=8. torch.nn.MultiheadAttention(num_heads=64, embed_dim=8) swaps parameters incorrectly, A and D use positional arguments wrongly.

3. Final Answer:

   torch.nn.MultiheadAttention(embed_dim=64, num_heads=8) -> Option B

4. Quick Check:

   embed_dim first, num_heads second [OK]

Hint: Remember: embed_dim first, num_heads second in constructor [OK]

Common Mistakes:

• Swapping embed_dim and num_heads
• Using positional args incorrectly
• Missing one of the required parameters

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

import torch
mha = torch.nn.MultiheadAttention(embed_dim=16, num_heads=4)
query = torch.rand(5, 3, 16)  # (sequence_length, batch_size, embed_dim)
key = torch.rand(5, 3, 16)
value = torch.rand(5, 3, 16)
attn_output, attn_weights = mha(query, key, value)

medium

A. (5, 3, 16)

B. (3, 5, 16)

C. (5, 16, 3)

D. (3, 16, 5)

Solution

1. Step 1: Understand input shapes for MultiheadAttention

   Inputs are (sequence_length, batch_size, embed_dim). Here, query shape is (5, 3, 16).

2. Step 2: Check output shape of attn_output

   Output shape matches query shape: (5, 3, 16).

3. Final Answer:

   (5, 3, 16) -> Option A

4. Quick Check:

   Output shape = query shape = (5, 3, 16) [OK]

Hint: Output shape matches query input shape [OK]

Common Mistakes:

• Mixing batch and sequence dimensions
• Confusing embed_dim position
• Assuming output shape changes with heads

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

import torch
mha = torch.nn.MultiheadAttention(embed_dim=32, num_heads=4)
query = torch.rand(10, 2, 32)
key = torch.rand(10, 2, 16)
value = torch.rand(10, 2, 32)
output, weights = mha(query, key, value)

medium

A. The embed_dim of key does not match embed_dim of query and value

B. The batch size is incorrect

C. The sequence length is too large

D. The number of heads is too high

Solution

1. Step 1: Check embed_dim consistency

   MultiheadAttention requires query, key, and value to have the same embedding dimension. Here, key has embed_dim 16, while query and value have 32.

2. Step 2: Identify the error cause

   This mismatch causes a runtime error because the attention mechanism cannot align different embedding sizes.

3. Final Answer:

   The embed_dim of key does not match embed_dim of query and value -> Option A

4. Quick Check:

   Embed dims must match for query, key, value [OK]

Hint: Check all embed_dims match before calling attention [OK]

Common Mistakes:

• Ignoring embed_dim mismatch
• Assuming batch size causes error
• Thinking sequence length is the problem

5. You want to apply multi-head attention to a batch of sentences with varying lengths. Which PyTorch feature helps you handle padding tokens correctly during attention computation?

hard

A. Increasing the number of heads to cover padding

B. Setting batch_first=True in MultiheadAttention

C. Using the key_padding_mask argument in MultiheadAttention

D. Manually removing padding tokens before attention

Solution

1. Step 1: Understand padding in variable-length sequences

   Sentences have different lengths, so padding tokens are added to make batch tensors uniform.

2. Step 2: Use key_padding_mask to ignore padding

   PyTorch's MultiheadAttention accepts key_padding_mask to mask out padded positions during attention, preventing them from affecting results.

3. Final Answer:

   Using the key_padding_mask argument in MultiheadAttention -> Option C

4. Quick Check:

   Mask padding tokens with key_padding_mask [OK]

Hint: Use key_padding_mask to ignore padded tokens [OK]

Common Mistakes:

• Removing padding manually (inefficient)
• Confusing batch_first with padding handling
• Increasing heads does not fix padding