Implement Multi-Head Attention
ByteDance● HardNEW GRADPHONE SCREENProblem statement
A Google AI/ML interview report asks the candidate to hand-code multi-head attention. For this exercise, implement a projection-free scaled dot-product attention forward pass. query, key and value are row-major integer matrices with the same feature width d; key and value have the same number of rows. heads divides d.
Split each row into contiguous head slices. For each query row and head, compute dot products against every key slice, divide by sqrt(d / heads), apply a numerically stable softmax across the keys, and use those probabilities to combine the matching value slices. Concatenate the head outputs in head order.
Return one comma-separated string per query row, with each coordinate formatted to exactly six decimal places. There are no learned projections, masks, biases, dropout or backward pass in this judged practice.
Function
multiHeadAttention(query: int[][], key: int[][], value: int[][], heads: int) → String[]Examples
Example 1
query = [[1]]key = [[0],[1]]value = [[10],[20]]heads = 1return = ["17.310586"]Softmax over scores 0 and 1 weights the second value more heavily.
Example 2
query = [[1,0]]key = [[1,0],[0,1]]value = [[4,8],[6,2]]heads = 2return = ["4.537883,5.000000"]Each one-dimensional head computes its own attention distribution.
Constraints
1 <= query rows, key rows, d <= 40.- Every row has width
d; key and value have the same number of rows. 1 <= heads <= d, andheadsdividesd.- Every entry is an integer in
[-100, 100].