Attention Score Computation with Einsum

hardPROEinsum OperationsPyTorch Advanced

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

Input

None

Output

{'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]}

Reasoning

Inputs: three random tensors $Q$ , $K$ , $V$ , each $2 \times 4 \times 8$ : batch 2, 4 positions, an 8-dimensional vector per position ( $d_k = 8$ ).
Scores: entry $(b, i, j)$ is the dot product $\sum_{d=1}^{8} Q_{bid} K_{bjd}$ , divided by $\sqrt{8} \approx 2.8284$ . Summing away $d$ leaves $2 \times 4 \times 4$ : every query against every key. For batch 0, query 0 against key 0 the raw dot product is $3.4281$ , giving $3.4281 / 2.8284 \approx 1.2120$ ; that whole row is $[1.2120, 0.7364, -0.3090, -0.8766]$ .
Weights: exponentiate each score and divide by its row total. That row becomes $[0.5092, 0.3165, 0.1113, 0.0631]$ : query 0 attends mostly to key 0. Shape unchanged.
Output: position $i$ gets $\sum_j w_{bij} V_{bj\cdot}$ , a weighted average of the 4 value vectors, so the key index $j$ disappears and the shape is $2 \times 4 \times 8$ .
Notice that each of the 4 rows of batch 0 sums to exactly 1 by construction, which is why attn_row_sums is [1.0, 1.0, 1.0, 1.0] after rounding to 2 decimal places.

Constraints

Use einsum for Q@K^T and attn@V
Scale by sqrt(d_k)
Apply softmax for attention weights

Learn this in

Machine TranslationNLP Study Plan

More from PyTorch Advanced

Implement SGD with Momentumhard Implement Custom SGD Optimizerhard Weight Decay for Specific Layersmedium Freeze Backbone, Train Headmedium Per-Layer Learning Rateseasy

Concept primerPREVIEW

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:

PRO

Continue the full explanation

You're reading the free preview. Unlock the complete walkthrough, the code editor, test runner and reference solution with Premium.

PRO

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.

solution.py

Python 3.13

Test results