PIXELBANKv9.1.0
Menu

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

The attention mechanism:

  1. Compute scores: scores=Qâ‹…KT\text{scores} = Q \cdot K^T
  2. Scale: scaled=scoresdk\text{scaled} = \frac{\text{scores}}{\sqrt{d_k}} where dkd_k is the key dimension
  3. Apply softmax row-wise to get attention weights
  4. Compute output: output=weightsâ‹…V\text{output} = \text{weights} \cdot V

Softmax: softmax(xi)=exi∑jexj\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
🔒

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

Test Results

0/0
Run code to see test results.
Attention Weights - Medium | PixelBank