How should I practice Attention With Absolute Index Bias?

Start by reading the full statement, solve the visible examples, then compare the pattern against related NVIDIA phone screen problems in FastPrep.

Is Attention With Absolute Index Bias associated with NVIDIA?

FastPrep tags this as a NVIDIA phone screen coding practice problem when the problem catalog has that company and stage context.

Attention With Absolute Index Bias

NVIDIA● MediumFULLTIMEPHONE SCREEN

Learn

Problem statement

Implement the reported attention formula softmax(Q * K + B) * V, where B[i][j] = abs(i-j). Return the attention output matrix.

For this exercise, assume Q has shape n by d, K has shape d by m, and V has shape m by p. Indices start at zero. Apply softmax across the m keys independently for each query row: softmax(s)[j] = exp(s[j]) / sum(exp(s)).

Use the positive absolute-distance bias exactly as written. For this exercise, assume there is no causal mask or additional scale factor. Compute in double precision; relative-or-absolute error of 10^-6 is accepted. Subtracting the maximum score before exponentiation is mathematically equivalent and avoids overflow.

Function

biasedAttention(q: double[][], k: double[][], v: double[][]) → double[][]

Examples

Example 1

q = [[0]]k = [[0,0]]v = [[0],[10]]return = [[7.310585786300049]]

Scores are 0 and 1, so the second value receives e/(1+e) of the weight.

Example 2

q = [[1,0],[0,1]]k = [[2,0],[0,2]]v = [[10,0],[0,20]]return = [[7.310585786300049,5.3788284273999025],[2.6894142136999513,14.621171572600097]]

Row softmax is computed independently; output has two rows and two value coordinates.

Example 3

q = [[10,-10]]k = [[10],[-10]]v = [[3,-4]]return = [[3.0,-4.0]]

A single key has probability one even with a large dot product; its value vector is returned.

Constraints

For this exercise, assume 1 <= n,m,d,p <= 64.
Every matrix has the specified rectangular shape, with finite entries in [-20,20].