Scaled Dot-Product Attention

easyFREEAttention MechanismsLLM 1: Foundations

Implement scaled dot-product attention.

Given Query (Q), Key (K), and Value (V) matrices, compute: $\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V$

where d_k is the dimension of the keys (number of columns in K).

Input:

Q, K, V: NumPy float arrays, each of shape (n, d) (sequence length, dimension)

Output: The attention output matrix of shape (n, d). The tests print each row with values rounded to 4 decimal places.

Example

Input

Output

[0.6698 0.3302]
[0.3302 0.6698]

Reasoning

The sample header says $n = 2$ tokens with $d = 2$ dimensions, and $Q$ , $K$ and $V$ are all the $2 \times 2$ identity: token 0 is $[1, 0]$ , token 1 is $[0, 1]$ .
Entry $(i, j)$ of $QK^T$ is the dot product of query $i$ with key $j$ : each token matches itself with score $1$ and the other with $0$ . Dividing by $\sqrt{d_k} = \sqrt{2}$ turns the $1$ s into $0.7071$ .
Softmax runs along each row. Row 0 is $[0.7071, 0]$ , so the weights are $\frac{e^{0.7071}}{e^{0.7071} + e^{0}} = \frac{2.0281}{3.0281} = 0.6698$ and $\frac{1}{3.0281} = 0.3302$ . Row 1 is the mirror image, $[0.3302, 0.6698]$ .
Each output row is the weighted average of the rows of $V$ . Since $V$ is the identity, the weights come through unchanged, giving [0.6698 0.3302] and [0.3302 0.6698], an $n \times d$ result.
Notice that each row sums to $1$ . Without the $\sqrt{d_k}$ scaling the weights would be $[0.7311, 0.2689]$ ; scaling keeps them softer. If two keys were identical, the query would split its weight evenly between them.

Constraints

1 <= n <= 10, 1 <= d <= 10
Use numpy for matrix operations
softmax is applied row-wise
Round to 4 decimal places

Learn this in

Transformer ArchitectureLLM Study Plan CNNs & Sequence ModelsML Study Plan Machine TranslationNLP Study Plan Attention MechanismsNLP Study Plan

More from LLM 1: Foundations

Transformer Block Forward Passhard Feed-Forward Networkeasy Layer Normalizationeasy KV-Cache Attentionhard Multi-Head Attentionmedium

Concept primer

Background Knowledge

The Attention Mechanism is a key component in many deep learning models, particularly in Natural Language Processing (NLP) and Computer Vision. It allows the model to focus on specific parts of the input data that are relevant to the task at hand. The Scaled Dot-Product Attention is a specific type of attention mechanism that is widely used in models like Transformers. It computes the attention weights by taking the dot product of the Query (Q) and Key (K) matrices, scaling the result by the square root of the dimension of the keys ( $\sqrt{d_k}$ ), and then applying the softmax function.

The softmax function is a common activation function used in many neural network models. It takes a vector of real numbers as input and outputs a vector of values in the range (0, 1) that add up to 1. This is useful for modeling probabilities or attention weights. In the context of scaled dot-product attention, the softmax function is used to normalize the attention weights, ensuring that they add up to 1 and can be interpreted as probabilities.

The dimension of the keys ( $d_k$ ) plays a crucial role in scaled dot-product attention. It determines the scaling factor used to compute the attention weights. The scaling factor ( $\frac{1}{\sqrt{d_k}}$ ) helps to prevent the dot product from growing too large, which can lead to extremely small gradients during backpropagation. This, in turn, can make it difficult to train the model.

Algorithm/Approach

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

Compute the dot product of the Query (Q) and Key (K) matrices
Scale the result by the square root of the dimension of the keys ( $\sqrt{d_k}$ )
Apply the softmax function to the scaled dot product
Compute the attention output by multiplying the softmax output with the Value (V) matrix

This approach can be implemented using basic linear algebra operations and the softmax function.

Step-by-Step Strategy

To implement the solution, follow these steps:

Read the input data, including the sequence length (n), dimension (d), and the Query (Q), Key (K), and Value (V) matrices.
Compute the dot product of the Query (Q) and Key (K) matrices: QK^T.
Compute the scaling factor: $\frac{1}{\sqrt{d_k}}$ .
Scale the dot product: $\frac{QK^T}{\sqrt{d_k}}$ .
Apply the softmax function to the scaled dot product.
Compute the attention output by multiplying the softmax output with the Value (V) matrix.
Round the attention output values to 4 decimal places.

Common Pitfalls

When implementing the solution, watch out for the following:

Make sure to transpose the Key (K) matrix before computing the dot product.
Use the correct dimension ( $d_k$ ) to compute the scaling factor.
Apply the softmax function correctly, ensuring that the output values add up to 1.
Use a stable implementation of the softmax function to avoid numerical issues.

Time & Space Complexity

The time complexity of the solution is O(n^2 * d), where n is the sequence length and d is the dimension. This is because we need to compute the dot product of the Query (Q) and Key (K) matrices, which has a time complexity of O(n^2 * d). The space complexity is O(n^2 + n * d), as we need to store the input matrices, the dot product, and the attention output.

solution.py

Python 3.13

Test results