Problem
A language model must not peek at future words. Causal attention masks the future: query row i may only attend to key rows 0 to i.
For each row i (queries, keys and values all have n rows):
- Score each allowed key j ≤ i: score = (queries[i] · keys[j]) / √d, where d is the number of columns of keys.
- Apply softmax over those allowed scores only (subtract the largest before taking e^x).
- Output row i = Σ weight_j × values[j] over j ≤ i.
Implement causal_attention(queries, keys, values).
Input
queries, keys: n rows of d numbers each.
values: n rows of v numbers each.
Output
Return the n × v output matrix.
Constraints
1 ≤ n ≤ 50, 1 ≤ d, v ≤ 32
-10 ≤ each value ≤ 10
Examples
Input:
keys = [
[1.0],
[2.0],
[3.0]
]
values = [
[3.0],
[6.0],
[9.0]
]
queries = [
[0.0],
[0.0],
[0.0]
]
Expected output:
expected = [
[3.0],
[4.5],
[6.0]
]Input:
std::vector<std::vector<double>> keys = {
{1.0},
{2.0},
{3.0}
};
std::vector<std::vector<double>> values = {
{3.0},
{6.0},
{9.0}
};
std::vector<std::vector<double>> queries = {
{0.0},
{0.0},
{0.0}
};
Expected output:
std::vector<std::vector<double>> expected = {
{3.0},
{4.5},
{6.0}
};Input:
double[][] keys = new double[][] {
new double[] {1.0},
new double[] {2.0},
new double[] {3.0}
};
double[][] values = new double[][] {
new double[] {3.0},
new double[] {6.0},
new double[] {9.0}
};
double[][] queries = new double[][] {
new double[] {0.0},
new double[] {0.0},
new double[] {0.0}
};
Expected output:
double[][] expected = new double[][] {
new double[] {3.0},
new double[] {4.5},
new double[] {6.0}
};Explanation
All scores are 0, so each row averages the values it may see: row 0 sees only 3, row 1 averages 3 and 6 (4.5), row 2 averages all three (6).
Input:
keys = [
[1.0, 0.0],
[1.0, 0.0]
]
values = [
[1.0, 2.0],
[3.0, 4.0]
]
queries = [
[1.0, 0.0],
[1.0, 0.0]
]
Expected output:
expected = [
[1.0, 2.0],
[2.0, 3.0]
]Input:
std::vector<std::vector<double>> keys = {
{1.0, 0.0},
{1.0, 0.0}
};
std::vector<std::vector<double>> values = {
{1.0, 2.0},
{3.0, 4.0}
};
std::vector<std::vector<double>> queries = {
{1.0, 0.0},
{1.0, 0.0}
};
Expected output:
std::vector<std::vector<double>> expected = {
{1.0, 2.0},
{2.0, 3.0}
};Input:
double[][] keys = new double[][] {
new double[] {1.0, 0.0},
new double[] {1.0, 0.0}
};
double[][] values = new double[][] {
new double[] {1.0, 2.0},
new double[] {3.0, 4.0}
};
double[][] queries = new double[][] {
new double[] {1.0, 0.0},
new double[] {1.0, 0.0}
};
Expected output:
double[][] expected = new double[][] {
new double[] {1.0, 2.0},
new double[] {2.0, 3.0}
};Explanation
Row 0 can only see itself, so it copies values[0] = [1, 2]. Row 1 gives both keys the same score, so it averages them: [2, 3].
How to approach it
Try on your own first, then open one hint at a time.
Hint 1
Each output row can only use keys and values from the current row and earlier rows.
Hint 2
Compute dot products, scale by the square root of the key dimension, then apply a stable softmax over only the allowed rows.
Hint 3
Remember to normalize separately for every query row, and handle the first row with just one allowed key.
Expected complexity
O(n^2 * (d + v)) time, O(n * v) extra space
References
This challenge is adapted for CodeUtility. Follow the statement and constraints above.