Attention Weights

mediumPROSequence ProcessingNLP 2: Applications

Compute scaled dot-product attention weights given Query, Key, and Value matrices.

The attention mechanism:

Compute scores: $\text{scores} = Q \cdot K^T$
Scale: $\text{scaled} = \frac{\text{scores}}{\sqrt{d_k}}$ where $d_k$ is the key dimension
Apply softmax row-wise to get attention weights
Compute output: $\text{output} = \text{weights} \cdot V$

Softmax: $\text{softmax}(x_i) = \frac{e^{x_i}}{\sum_j e^{x_j}}$

Input format:

Line 1: seq_len d_k d_v (space-separated ints)
Next seq_len lines: Q matrix (d_k columns)
Next seq_len lines: K matrix (d_k columns)
Next seq_len lines: V matrix (d_v columns)

Output:

Line 1: Attention weights matrix (rounded to 4 decimals)
Line 2: Output matrix (rounded to 4 decimals)

Example

Input

2 2 2
1.0 0.0
0.0 1.0
1.0 0.0
0.0 1.0
0.5 0.5
0.3 0.7

Output

[[0.5987, 0.4013], [0.4013, 0.5987]]
[[0.4205, 0.5795], [0.3795, 0.6205]]

Reasoning

Step 1: Scores = Q @ K^T Q @ K^T = [[1,0],[0,1]] @ [[1,0],[0,1]]^T = [[1,0],[0,1]]

Step 2: Scale by sqrt(d_k) = sqrt(2) = 1.4142 Scaled = [[0.7071, 0.0], [0.0, 0.7071]]

Step 3: Softmax row-wise Row 0: softmax([0.7071, 0.0]) = [e^0.7071, e^0] / sum = [2.028, 1.0] / 3.028 = [0.5987, 0.4013] Row 1: softmax([0.0, 0.7071]) = [0.4013, 0.5987]

Step 4: Output = weights @ V [[0.5987, 0.4013], [0.4013, 0.5987]] @ [[0.5, 0.5], [0.3, 0.7]] = [[0.59870.5+0.40130.3, 0.59870.5+0.40130.7], ...] = [[0.4205, 0.5795], [0.3795, 0.6205]]

Constraints

Use numpy for matrix operations
Softmax is applied row-wise
Round all output values to 4 decimal places

Learn this in

Machine TranslationNLP Study Plan Attention MechanismsNLP Study Plan

More from NLP 2: Applications

BIO Entity Span Extractoreasy Simple Entity Taggereasy Sentiment Word Countereasy Keyword Classifiereasy One-Hot Sequence Encodereasy

Concept primerPREVIEW

Background Knowledge

The problem revolves around the attention mechanism, a key concept in Natural Language Processing (NLP) and Deep Learning. The attention mechanism allows a model to focus on specific parts of the input data when generating output. This is particularly useful in sequence-to-sequence models, such as machine translation, where the model needs to weigh the importance of different input elements when generating each output element. The attention mechanism is composed of three main components: Query (Q), Key (K), and Value (V) matrices.

In the context of this problem, we are dealing with scaled dot-product attention, a specific type of attention mechanism. This mechanism computes attention weights by taking the dot product of the Query and Key matrices, scaling the result, and then applying a softmax function. The softmax function is a common activation function in Deep Learning, used to normalize the output of a layer to ensure it forms a valid probability distribution. The key dimension ( $d_k$ ) and value dimension ( $d_v$ ) are crucial in determining the shape of the input matrices and the output.

Understanding the matrix operations involved is essential to solving this problem. The problem requires performing matrix multiplication, transposition, and element-wise operations. Additionally, the problem involves applying the softmax function row-wise to the scaled scores, which requires a good grasp of linear algebra and vectorized operations.

Algorithm/Approach

The general approach to solving this problem involves the following steps:

Compute the dot product of the Query and Key matrices
Scale the result by the square root of the key dimension
Apply the softmax function row-wise to obtain the attention weights
Compute the output by multiplying the attention weights with the Value matrix

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