Scaled Dot-Product Attention
Learn this problem
AmazonINTERNPHONE SCREENProblem statement
Implement unmasked single-head scaled dot-product attention for three nonempty matrices:
querieshas shapen by d.keyshas shapem by d.valueshas shapem by v.
For each query row, compute one score per key as the dot product divided by sqrt(d). Apply maximum-subtracted softmax across that score row, then return the weighted sum of the value rows. The result has shape n by v.
There is no mask, batch dimension, learned projection, bias, or dropout. Implement the matrix loops, stable softmax, and weighted sum directly without a machine-learning or numerical-array library.
Function
scaledDotProductAttention(queries: double[][], keys: double[][], values: double[][]) → double[][]Examples
Example 1
queries = [[1.0,0.0]]keys = [[1.0,0.0],[0.0,1.0]]values = [[10.0,0.0],[0.0,20.0]]return = [[6.697615493266569,6.604769013466862]]The scaled scores are approximately [0.7071067812, 0]. Their softmax weights are approximately [0.6697615493, 0.3302384507], which weight the two value rows.
Example 2
queries = [[0.0,0.0]]keys = [[1.0,0.0],[-1.0,0.0]]values = [[2.0,4.0],[6.0,8.0]]return = [[4.0,6.0]]The zero query has equal score for both keys, so softmax gives each value row weight 0.5.
Example 3
queries = [[1.0,2.0],[-3.0,4.0]]keys = [[5.0,6.0]]values = [[7.0,-2.0,1.0]]return = [[7.0,-2.0,1.0],[7.0,-2.0,1.0]]With only one key, every query gives its value row softmax weight 1.
Constraints
1 <= n, m <= 201 <= d, v <= 10- Every matrix is rectangular and follows the stated compatible shapes.
-20 <= matrix entry <= 20- Every matrix entry is finite: no value is NaN or infinity.
- The result is compared entry by entry with absolute tolerance
1e-9.