Self-attention mechanism in PyTorch - Cheat Sheet & Quick Revision

Recall & Review

beginner

What is the main purpose of the self-attention mechanism in neural networks?

Self-attention helps the model focus on different parts of the input sequence to understand relationships and context better, improving tasks like language understanding.

intermediate

In self-attention, what are the Query, Key, and Value vectors?

They are vectors derived from the input that help compute attention scores: Query asks what to focus on, Key represents the content to compare against, and Value holds the actual information to be combined.

intermediate

How is the attention score calculated in self-attention?

Attention scores are calculated by taking the dot product of the Query and Key vectors, then scaling and applying a softmax to get weights that sum to 1.

advanced

Why do we scale the dot product by the square root of the key dimension in self-attention?

Scaling prevents the dot product values from becoming too large, which helps keep the softmax function stable and gradients well-behaved during training.

beginner

What is the output of the self-attention mechanism?

The output is a weighted sum of the Value vectors, where weights come from the attention scores, representing a context-aware combination of input elements.

What does the Query vector represent in self-attention?

A The final prediction

B The part of input asking what to focus on

C The weights for output

D The actual information to be combined

The Query vector represents what the model is trying to focus on in the input.

Why is softmax applied to the dot product of Query and Key vectors?

A To normalize scores into probabilities

B To increase the dot product values

C To reduce the size of vectors

D To create new vectors

Softmax converts raw scores into probabilities that sum to 1, making them easier to interpret as attention weights.

What is the role of the Value vector in self-attention?

A It is used to calculate dot product

B It asks what to focus on

C It holds the actual information to be combined

D It normalizes the output

Value vectors contain the information that will be weighted and combined to form the output.

What problem does scaling the dot product by sqrt(d_k) solve?

A Speeds up training

B Increases the size of the output

C Reduces the number of parameters

D Prevents large values that make softmax unstable

Scaling keeps the dot product values in a range that prevents softmax from producing very small gradients.

Which of these best describes self-attention?

A A way to relate different parts of the same input sequence

B A method to increase dataset size

C A technique to reduce model size

D A way to generate random noise

Self-attention relates different parts of the input to capture context and dependencies.

Explain how the self-attention mechanism computes its output from input vectors.

Think about how the model decides what parts of the input to focus on.

Why is self-attention important in models like Transformers for language tasks?

Consider how words in a sentence relate to each other.

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