FastPrep
Problem · Matrix

Scaled Dot-Product Attention

Learn this problem
● MediumMeshy logoMeshyFULLTIMEOA

Problem statement

You are given three real-valued matrices:

  • query, with shape m x d
  • key, with shape n x d
  • value, with shape n x v

Compute single-head scaled dot-product attention without masking or learned projections:

softmax(query * transpose(key) / sqrt(d)) * value.

Apply softmax independently to every row. Return the resulting m x v matrix. Use a numerically stable softmax by subtracting the largest score in a row before exponentiation.

Function

scaledDotProductAttention(query: double[][], key: double[][], value: double[][]) → double[][]

Examples

Example 1

query = [[0]]key = [[2],[3]]value = [[4,5],[8,9]]return = [[6,7]]

Both scores are zero, so the attention weights are one half and the output is the elementwise average of the two value rows.

Example 2

query = [[1]]key = [[0],[1.0986122886681098]]value = [[0],[4]]return = [[3]]

The scores are 0 and ln(3), producing weights one quarter and three quarters. The weighted value is 3.

Example 3

query = [[1,0],[0,1]]key = [[1,0],[0,1]]value = [[10,0],[0,10]]return = [[6.697615493266569,3.302384506733431],[3.302384506733431,6.697615493266569]]

Each query gives the matching key a score of 1 / sqrt(2) and the other key a score of zero, so the two rows use reversed attention weights.

Constraints

  • 1 <= m, n, d, v <= 100
  • query.length == m and every query row has length d.
  • key.length == value.length == n, every key row has length d, and every value row has length v.
  • Every matrix entry is finite and lies between -100 and 100, inclusive.
  • Answers are accepted with relative-or-absolute tolerance 1e-6.
drafts saved locally
public double[][] scaledDotProductAttention(double[][] query, double[][] key, double[][] value) {
    // Write your code here.
}
query[[0]]
key[[2],[3]]
value[[4,5],[8,9]]
expected[[6,7]]
checking account