Attention Score Computation with Einsum
Problem Statement
Compute scaled dot-product attention scores using einsum, as used in Transformers.
Background
Attention computes softmax(Q @ K^T / sqrt(d_k)) @ V. Einsum handles the batched transpose elegantly by specifying which indices to sum over and which to keep.
Your Task
The starter code creates Q, K, V tensors (batch=2, seq=4, d_k=8). Use einsum to compute scaled dot-product attention: compute Q·K^T scores (scaled by sqrt(d_k)), apply softmax, then multiply by V.
Output Format
Returns a dictionary with "score_shape", "attn_shape", "output_shape", and "attn_row_sums" (should all be 1.0).
Example:
None
{'score_shape': [2, 4, 4], 'attn_shape': [2, 4, 4], 'output_shape': [2, 4, 8], 'attn_row_sums': [1.0, 1.0, 1.0, 1.0]}- We start by creating tensors Q, K, V of shape (2, 4, 8) and computing attention scores using
torch.einsum('bqd,bkd->bqk', Q, K) / sqrt(8), resulting in a tensor of shape (2, 4, 4). - The scores are then passed through a softmax function along the last dimension, yielding attention weights of the same shape (2, 4, 4).
- We compute the output by multiplying the attention weights with V using
torch.einsum('bqk,bkd->bqd', attn_weights, V), resulting in a tensor of shape (2, 4, 8). - Finally, we calculate the sum of each row of the attention weights for the first batch, which should all be approximately 1.0 due to the properties of the softmax function, and return the required information in a dictionary.
Constraints:
- Use einsum for Q@K^T and attn@V
- Scale by sqrt(d_k)
- Apply softmax for attention weights
Background Knowledge
The problem revolves around the concept of attention mechanisms in deep learning, particularly in the context of Transformers. Attention mechanisms allow models to focus on specific parts of the input data when generating outputs. In this case, we're dealing with scaled dot-product attention, which computes attention scores by taking the dot product of query (Q) and key (K) vectors, scaling them by the square root of the dimensionality (d_k), and then applying a softmax function to obtain attention weights.
The einsum operation is a key component in this problem. Einsum is a powerful tensor operation that can perform various linear algebra operations, such as matrix multiplication, transpose, and contraction, in a single step. It's particularly useful for batched operations, where we need to perform the same operation on multiple batches of data. In this case, we'll use einsum to compute the attention scores and the final output.
The softmax function is also crucial in this problem. Softmax is a activation function that takes a vector of real numbers and returns a vector of values in the range (0, 1) that add up to 1. In the context of attention mechanisms, softmax is used to normalize the attention scores, ensuring that the attention weights sum up to 1. This allows the model to focus on specific parts of the input data when generating outputs.
Algorithm/Approach
The general approach to solving this problem involves the following steps:
- Initialize the input tensors (Q, K, V) with the given shape
- Compute the attention scores using einsum and scale them by the square root of the dimensionality (d_k)
- Apply the softmax function to the attention scores to obtain attention weights
- Compute the final output by multiplying the attention weights with the value tensor (V) using einsum
The key insight here is to use einsum to perform the batched matrix multiplications and transposes, which can be tricky to implement using traditional tensor operations.
Step-by-Step Strategy
To implement the solution, follow these steps:
Continue the full explanation
You're reading the free preview. Unlock the complete walkthrough, the code editor, test runner and reference solution with Premium.
Editor locked
The code editor is locked for Pro problems. It is only available for free problems. Please upgrade to gain access to the code editor for all problems.