Scaled Dot-Product Attention
Implement Scaled Dot-Product Attention from the Transformer architecture.
Given query matrix Q, key matrix K, and value matrix V:
Attention(Q,K,V)=softmax(dk​​QKT​)V
where dk​ is the dimension of the keys (number of columns in K).
Steps:
- Compute QKT (matrix multiplication)
- Scale by dk​​1​
- Apply softmax row-wise (each row sums to 1)
- Multiply by V
Return the attention output matrix, rounded to 4 decimal places.
Example:
Q = [[1, 0]] K = [[1, 0], [0, 1]] V = [[1, 2], [3, 4]]
[[1.6605, 2.6605]]
- We start by computing the matrix product QKT, which is [[1,0]]â‹…[[1,0],[0,1]]T=[[1,0]]â‹…[[1,0],[0,1]]=[[1,0]].
- Then, we scale this result by dk​​1​=2​1​, yielding [[2​1​,0]].
- Next, we apply the softmax function row-wise: since [[2​1​,0]] is a single row, this results in [[1+e−2​1​1​,1+e01​]]≈[[0.7311,0.2689]] after normalization, but because we apply softmax to [2​1​,0] we actually calculate softmax([[2​1​,0]])=[[e2​1​+e0e2​1​​,e2​1​+e0e0​]]=[[e2​1​+1e2​1​​,e2​1​+11​]], which then gets multiplied by V.
- Finally, multiplying this result by V yields $[[\frac{e^{\frac{1}{\sqrt{2}}}}{e^{\frac{1}{\sqrt{2}}} + 1}, \frac{1}{e^{\frac{1}{\sqrt{2}}} + 1}]] \cdot [[1, 2], [3, 4]] = [[\frac{e^{\frac{1}{\sqrt{2}}}}{e^{\frac{1}{\sqrt{2}}} + 1} \cdot 1 + \frac
Constraints:
- Q: 2D list (n x d_k)
- K: 2D list (m x d_k)
- V: 2D list (m x d_v)
- Return 2D list (n x d_v) rounded to 4 decimal places
- Use numerical stability trick for softmax
Background Knowledge
The Scaled Dot-Product Attention mechanism is a key component of the Transformer architecture, introduced in the paper "Attention Is All You Need" by Vaswani et al. in 2017. This mechanism allows the model to attend to different parts of the input sequence simultaneously and weigh their importance. The attention mechanism is based on the idea of dot-product attention, which computes the attention weights by taking the dot product of the query and key vectors.
The Transformer architecture relies heavily on self-attention mechanisms, which allow the model to attend to different parts of the input sequence and capture long-range dependencies. The Scaled Dot-Product Attention mechanism is a variant of the self-attention mechanism that scales the dot product of the query and key vectors by the square root of the dimensionality of the key vectors. This scaling helps to prevent the dot product from growing too large, which can lead to extremely small gradients during backpropagation.
The softmax function is used to normalize the attention weights, ensuring that they sum to 1 for each row. This is important because it allows the model to interpret the attention weights as probabilities. The matrix multiplication operation is used to compute the dot product of the query and key matrices, and the value matrix is used to compute the final attention output.
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.