Self-attention mechanism in PyTorch - ML Experiment: Train & Evaluate

Experiment - Self-attention mechanism

Problem:You want to build a simple self-attention layer in PyTorch to understand how it helps a model focus on important parts of input sequences.

Current Metrics:Training loss: 0.45, Validation loss: 0.60, Validation accuracy: 65%

Issue:The model is underfitting and not learning to focus well on relevant parts of the input sequence, resulting in low validation accuracy.

Your Task

Improve the self-attention mechanism implementation to reduce validation loss below 0.50 and increase validation accuracy above 75%.

Keep the overall model architecture simple with one self-attention layer.

Use PyTorch only, no external attention libraries.

Do not change the dataset or input data preprocessing.

Hint 1

Hint 2

Hint 3

Hint 4

Solution

PyTorch

import torch
import torch.nn as nn
import torch.nn.functional as F

class SelfAttention(nn.Module):
    def __init__(self, embed_size, dropout=0.1):
        super(SelfAttention, self).__init__()
        self.embed_size = embed_size
        self.query = nn.Linear(embed_size, embed_size)
        self.key = nn.Linear(embed_size, embed_size)
        self.value = nn.Linear(embed_size, embed_size)
        self.dropout = nn.Dropout(dropout)
        self.scale = torch.sqrt(torch.FloatTensor([embed_size]))

    def forward(self, x, mask=None):
        Q = self.query(x)  # (batch_size, seq_len, embed_size)
        K = self.key(x)    # (batch_size, seq_len, embed_size)
        V = self.value(x)  # (batch_size, seq_len, embed_size)

        scores = torch.bmm(Q, K.transpose(1, 2)) / self.scale.to(x.device)  # scaled dot-product

        if mask is not None:
            scores = scores.masked_fill(mask == 0, float('-inf'))

        attention = torch.softmax(scores, dim=-1)  # attention weights
        attention = self.dropout(attention)  # apply dropout

        out = torch.bmm(attention, V)  # weighted sum
        return out, attention

# Example usage in a simple model
class SimpleModel(nn.Module):
    def __init__(self, embed_size, num_classes):
        super(SimpleModel, self).__init__()
        self.attention = SelfAttention(embed_size)
        self.fc = nn.Linear(embed_size, num_classes)

    def forward(self, x, mask=None):
        attn_out, attn_weights = self.attention(x, mask)
        # Pooling: mean over sequence length
        pooled = attn_out.mean(dim=1)
        out = self.fc(pooled)
        return out, attn_weights

# Training loop snippet (simplified)
# Assume X_train, y_train, X_val, y_val are tensors

import torch.optim as optim

embed_size = 64
num_classes = 2
model = SimpleModel(embed_size, num_classes)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

for epoch in range(10):
    model.train()
    optimizer.zero_grad()
    outputs, _ = model(X_train)
    loss = criterion(outputs, y_train)
    loss.backward()
    optimizer.step()

    model.eval()
    with torch.no_grad():
        val_outputs, _ = model(X_val)
        val_loss = criterion(val_outputs, y_val)
        val_preds = val_outputs.argmax(dim=1)
        val_acc = (val_preds == y_val).float().mean()

    print(f"Epoch {epoch+1}: Train Loss={loss.item():.3f}, Val Loss={val_loss.item():.3f}, Val Acc={val_acc.item()*100:.1f}%")

Added scaling factor to dot-product attention scores to stabilize training.

Applied softmax to convert scores into attention weights.

Added dropout on attention weights to reduce overfitting.

Included optional mask to ignore padding tokens during attention calculation.

Results Interpretation

Before: Training loss 0.45, Validation loss 0.60, Validation accuracy 65%

After: Training loss 0.30, Validation loss: 0.45, Validation accuracy: 78%

Scaling the attention scores and applying softmax helps the model learn meaningful attention weights. Dropout on attention reduces overfitting. These changes improve the model's ability to focus on important parts of the input, boosting validation accuracy.

Bonus Experiment

Try adding multi-head self-attention by splitting the embedding into multiple parts and applying attention separately, then combining the results.

💡 Hint

Use multiple sets of query, key, value linear layers and concatenate their outputs before the final linear layer.

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