Problem
Attention lets each position in a sequence look at every other position and mix their information. It is the core of the Transformer architecture used by large language models:
Attention(Q, K, V) = softmax(Q × Kᵀ / √d) × V
For each query row q:
- Score every key row k: score = (q · k) / √d, where d is the number of columns of K.
- Turn the scores into weights with softmax (subtract the largest score before taking e^x).
- The output row is the weighted sum of the value rows: Σ weight_j × values[j].
Implement attention(queries, keys, values).
Input
queries: n rows, each with d numbers.
keys: m rows, each with d numbers.
values: m rows (one per key), each with v numbers.
Output
Return the n × v output matrix.
Constraints
1 ≤ n, m ≤ 50, 1 ≤ d, v ≤ 32
-10 ≤ each value ≤ 10
Examples
Input:
keys = [
[1.0, 0.0],
[1.0, 0.0]
]
values = [
[2.0, 4.0],
[6.0, 8.0]
]
queries = [
[1.0, 0.0]
]
Expected output:
expected = [
[4.0, 6.0]
]Input:
std::vector<std::vector<double>> keys = {
{1.0, 0.0},
{1.0, 0.0}
};
std::vector<std::vector<double>> values = {
{2.0, 4.0},
{6.0, 8.0}
};
std::vector<std::vector<double>> queries = {
{1.0, 0.0}
};
Expected output:
std::vector<std::vector<double>> expected = {
{4.0, 6.0}
};Input:
double[][] keys = new double[][] {
new double[] {1.0, 0.0},
new double[] {1.0, 0.0}
};
double[][] values = new double[][] {
new double[] {2.0, 4.0},
new double[] {6.0, 8.0}
};
double[][] queries = new double[][] {
new double[] {1.0, 0.0}
};
Expected output:
double[][] expected = new double[][] {
new double[] {4.0, 6.0}
};Explanation
Both keys give the same score, so the softmax weights are 0.5 and 0.5. The output is the average of the value rows: [(2 + 6) / 2, (4 + 8) / 2] = [4, 6].
Input:
keys = [
[3.0, 1.0],
[-2.0, 5.0],
[7.0, 7.0]
]
values = [
[1.0],
[2.0],
[6.0]
]
queries = [
[0.0, 0.0]
]
Expected output:
expected = [
[3.0]
]Input:
std::vector<std::vector<double>> keys = {
{3.0, 1.0},
{-2.0, 5.0},
{7.0, 7.0}
};
std::vector<std::vector<double>> values = {
{1.0},
{2.0},
{6.0}
};
std::vector<std::vector<double>> queries = {
{0.0, 0.0}
};
Expected output:
std::vector<std::vector<double>> expected = {
{3.0}
};Input:
double[][] keys = new double[][] {
new double[] {3.0, 1.0},
new double[] {-2.0, 5.0},
new double[] {7.0, 7.0}
};
double[][] values = new double[][] {
new double[] {1.0},
new double[] {2.0},
new double[] {6.0}
};
double[][] queries = new double[][] {
new double[] {0.0, 0.0}
};
Expected output:
double[][] expected = new double[][] {
new double[] {3.0}
};Explanation
A zero query scores 0 against every key, so attention is spread evenly: (1 + 2 + 6) / 3 = 3.
How to approach it
Try on your own first, then open one hint at a time.
Hint 1
For each query, compare it with every key using a dot product scaled by the square root of the key dimension.
Hint 2
Convert the scores into weights with a numerically stable softmax, then blend the value rows using those weights.
Hint 3
Each query produces one output row, and the value vectors determine the output width rather than the queries or keys.
Expected complexity
O(n m (d + v)) time, O(m) extra space
References
This challenge is adapted for CodeUtility. Follow the statement and constraints above.