📘

Scaled Dot-Product Attention

EasyAttention Mechanisms 📚 LLM 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:

Line 1: n d (sequence length, dimension)
Next n lines: Q matrix (space-separated floats)
Next n lines: K matrix
Next n lines: V matrix

Output: The attention output matrix, values rounded to 4 decimal places.

Example:

Input:

Output:

[0.7311 0.2689]
[0.2689 0.7311]

Reasoning:

The input provides the sequence length n = 2 and dimension d = 2, along with the Q, K, and V matrices.
We compute the dot product of Q and K^T, which is $\begin{bmatrix} 1 & 0 \\ 0 & 1 \end{bmatrix} \begin{bmatrix} 1 & 0 \\ 0 & 1 \end{bmatrix} = \begin{bmatrix} 1 & 0 \\ 0 & 1 \end{bmatrix}$ , then scale it by $\frac{1}{\sqrt{d}} = \frac{1}{\sqrt{2}}$ .
We apply the softmax function to the scaled dot product: $\text{softmax}\left(\frac{1}{\sqrt{2}}\begin{bmatrix} 1 & 0 \\ 0 & 1 \end{bmatrix}\right) = \begin{bmatrix} \frac{e^{\frac{1}{\sqrt{2}}}}{e^{\frac{1}{\sqrt{2}} } + e^{\frac{1}{\sqrt{2}}}} & \frac{e^0}{e^{\frac{1}{\sqrt{2}}} + e^{\frac{1}{\sqrt{2}}}} \\ \frac{e^0}{e^{\frac{1}{\sqrt{2}}} + e^{\frac{1}{\sqrt{2}}}} & \frac{e^{\frac{1}{\sqrt{2}}}}{e^{\frac{1}{\sqrt{2}}} + e^{\frac{1}{\sqrt{2}}}} \end{bmatrix} \approx \begin{bmatrix} 0.7311 & 0.2689 \\ 0.2689 & 0.7311 \end{bmatrix}$ .
The final output is obtained by multiplying this softmax result with the V matrix, which yields $\begin{bmatrix} 0.7311 & 0.2689 \ 0.2689 & 0.7311 \end{bmatrix} \begin{bmatrix} 1 & 0 \ 0 & 1 \end{bmatrix} = \begin{bmatrix} 0.

Constraints:

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

More from LLM 1: Foundations

Transformer Block Forward PassHard Feed-Forward NetworkEasy Layer NormalizationEasy KV-Cache AttentionHard Multi-Head AttentionMedium

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

0/0

Run code to see test results.