Self-attention and multi-head attention in NLP - Practice Problems & Coding Challenges

Challenge - 5 Problems

🧠 Conceptual

intermediate

What does self-attention compute in a transformer?

In a transformer model, self-attention helps the model focus on different parts of the input sequence. What exactly does self-attention compute?

A It applies a convolution operation over the input sequence to extract local features.

B It computes a weighted sum of the input elements where weights depend on the similarity between elements.

C It computes the average of all input tokens without considering their relationships.

D It sorts the input tokens based on their frequency in the sequence.

💡 Hint

Think about how the model decides which words to pay attention to when processing a sentence.

Self-attention calculates how much each word in a sentence should attend to every other word. It does this by computing similarity scores and then using these scores as weights to combine the input elements.

❓ Predict Output

intermediate

Output shape of multi-head attention layer

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

NLP

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)

output, _ = mha(query, key, value)
print(output.shape)

A (5, 2, 16)

B (2, 16, 5)

C (5, 16, 2)

D (2, 5, 16)

💡 Hint

Check the expected input and output shapes for PyTorch's MultiheadAttention module.

PyTorch's MultiheadAttention expects inputs with shape (sequence_length, batch_size, embedding_dim) and outputs the same shape.

❓ Hyperparameter

advanced

Effect of increasing number of heads in multi-head attention

What is the main effect of increasing the number of heads in a multi-head attention mechanism while keeping the total embedding dimension fixed?

A The attention weights become uniform across all tokens.

B The total embedding dimension increases, making the model slower but more accurate.

C The model ignores positional information and treats all tokens equally.

D Each head has a smaller dimension, allowing the model to focus on different representation subspaces.

💡 Hint

Think about how splitting embedding dimension among heads affects representation.

Increasing the number of heads splits the embedding dimension into smaller parts per head. This allows each head to learn different aspects or subspaces of the input features.

❓ Metrics

advanced

Interpreting attention weights in self-attention

In a trained transformer model, you extract the attention weights from a self-attention layer. What does a high attention weight between two tokens indicate?

A The tokens appear consecutively in the input sequence.

B The two tokens are identical in the input sequence.

C The model considers the two tokens highly related or important to each other for the current task.

D The tokens have the same part of speech tag.

💡 Hint

Attention weights show how much one token focuses on another during processing.

High attention weights mean the model is focusing more on the related token when processing the current token, indicating importance or relevance.

🔧 Debug

expert

Identifying error in custom multi-head attention implementation

Consider this simplified custom multi-head attention code snippet. What error will it raise when run?

NLP

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

class SimpleMultiHeadAttention(nn.Module):
    def __init__(self, embed_dim, num_heads):
        super().__init__()
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads
        self.q_linear = nn.Linear(embed_dim, embed_dim)
        self.k_linear = nn.Linear(embed_dim, embed_dim)
        self.v_linear = nn.Linear(embed_dim, embed_dim)
        self.out_linear = nn.Linear(embed_dim, embed_dim)

    def forward(self, x):
        batch_size, seq_len, embed_dim = x.size()
        Q = self.q_linear(x)
        K = self.k_linear(x)
        V = self.v_linear(x)

        # reshape for multi-head
        Q = Q.view(batch_size, seq_len, self.num_heads, self.head_dim)
        K = K.view(batch_size, seq_len, self.num_heads, self.head_dim)
        V = V.view(batch_size, seq_len, self.num_heads, self.head_dim)

        # transpose to get dimensions batch_size, num_heads, seq_len, head_dim
        Q = Q.transpose(1, 2)
        K = K.transpose(1, 2)
        V = V.transpose(1, 2)

        scores = torch.matmul(Q, K.transpose(-2, -1)) / (self.head_dim ** 0.5)
        attn = F.softmax(scores, dim=-1)
        out = torch.matmul(attn, V)

        # concatenate heads
        out = out.transpose(1, 2).contiguous().view(batch_size, seq_len, embed_dim)
        out = self.out_linear(out)
        return out

x = torch.rand(2, 5, 16)
model = SimpleMultiHeadAttention(embed_dim=16, num_heads=4)
output = model(x)

A No error, code runs successfully and output shape is (2, 5, 16)

B AttributeError because 'SimpleMultiHeadAttention' object has no attribute 'out_linear'

C RuntimeError due to shape mismatch in torch.matmul(Q, K.transpose(-2, -1))

D RuntimeError due to incorrect batch size in input tensor

💡 Hint

Check the tensor shapes carefully during reshaping and matrix multiplication.

The code correctly reshapes and transposes the tensors for multi-head attention. The matrix multiplications have matching dimensions, and the final output shape matches the input embedding dimension. Therefore, no error occurs.