Multi-head attention in PyTorch - Cheat Sheet & Quick Revision

Recall & Review

beginner

What is the main idea behind multi-head attention?

Multi-head attention splits the attention mechanism into several smaller parts (heads) that run in parallel. Each head learns different relationships in the data, and their results are combined to capture richer information.

beginner

In multi-head attention, what are the Query, Key, and Value?

Query, Key, and Value are three sets of vectors derived from the input data. The Query is what we want to find information about, the Key helps match the Query, and the Value holds the actual information to be gathered.

intermediate

Why do we use multiple heads instead of one in attention?

Using multiple heads allows the model to focus on different parts or aspects of the input simultaneously. This helps the model understand complex patterns better than a single attention head.

intermediate

What is the shape of the output from a multi-head attention layer in PyTorch?

The output shape is usually (sequence_length, batch_size, embedding_dim), where embedding_dim is the original input dimension after combining all heads.

intermediate

How is the scaled dot-product attention computed inside each head?

Scaled dot-product attention is computed by taking the dot product of Query and Key, dividing by the square root of the key dimension to scale, applying softmax to get weights, and then multiplying by the Value vectors.

What does each head in multi-head attention learn?

A Different relationships or features from the input

B The same information repeatedly

C Only the first token of the input

D Random noise

Each head focuses on different parts or aspects of the input to capture diverse information.

What is the purpose of scaling the dot product in scaled dot-product attention?

A To normalize the output to zero mean

B To prevent large values that can make softmax gradients too small

C To increase the dot product values

D To reduce the number of heads

Scaling by the square root of the key dimension keeps the dot products in a range that avoids very small gradients after softmax.

In PyTorch, which class implements multi-head attention?

A torch.nn.AttentionHead

B torch.nn.AttentionLayer

C torch.nn.MultiHeadLayer

D torch.nn.MultiheadAttention

PyTorch provides torch.nn.MultiheadAttention for multi-head attention implementation.

What is combined after all attention heads compute their outputs?

A The outputs are concatenated and projected back to the original dimension

B Only the first head's output is used

C The outputs are averaged without projection

D The outputs are discarded

Outputs from all heads are concatenated and then passed through a linear layer to match the original embedding size.

Which of these is NOT a component of multi-head attention?

A Query

B Value

C Bias

D Key

Query, Key, and Value are core components; bias is not a fundamental part of multi-head attention.

Explain how multi-head attention works and why it is useful in simple terms.

Think about how looking at something from different angles helps understand it better.

Describe the role of Query, Key, and Value in the attention mechanism.

Imagine searching for a book in a library using a catalog.