Pronoun Looks Back
A translation model reads: "The parcel did not fit in the locker because it was too big." To translate it correctly, the model has to work out which earlier word it is about. In a transformer, each word sends out a query vector ("what am I looking for?") and every word offers a key vector ("what do I contain?").
The model scores each key against the query with a dot product, shrinks the scores by the square root of the vector size so they do not get too extreme, and turns them into weights with softmax. Those weights say how much it pays attention to each word. This is the core step inside every transformer layer.
Your task
Write attentionWeights(query, keys). Let d be the length of query. For each key, compute the score (query · keys[i]) / √d. Return the softmax of the scores: e^scoreᵢ / Σ e^scoreⱼ, one weight per key, in the same order.
Scores can be large, so your answer must not overflow. The weights must add up to 1.
attentionWeights(query: float[], keys: float[][]) → float[] Numbers are accepted within 1e-6.
Examples
query = [1, 0, 1, 0] keys = [[2, 0, 2, 0], [0, 1, 0, 1], [1, 1, 1, 1]]
[0.6652409557748218, 0.09003057317038046, 0.24472847105479764]
d = 4, so divide by √4 = 2. Dot products: 4, 0, 2. Scores: 2, 0, 1. Softmax: e² ≈ 7.389, e⁰ = 1, e¹ ≈ 2.718, total 11.107, so the weights are about 0.665, 0.090 and 0.245. The first key matches the query best and gets most of the attention.
query = [0, 0] keys = [[3, 1], [-2, 5]]
[0.5, 0.5]
A zero query has dot product 0 with every key. All scores are 0, so attention is spread evenly: 0.5 and 0.5.
query = [30, 40] keys = [[30, 40], [40, 30]]
[1, 1.9531818567219233e-31]
Dot products are 2500 and 2400, scores 2500 / √2 ≈ 1767.8 and 2400 / √2 ≈ 1697.1. e^1767.8 overflows, so shift by the largest score first. The gap is about 70.7, so the second weight is e^-70.7, tiny but not zero, and the first is almost exactly 1.
Constraints
- 1 ≤ keys.length ≤ 1,000
- 1 ≤ d = query.length = keys[i].length ≤ 256
- -1,000 ≤ query[j], keys[i][j] ≤ 1,000
- Answers within 0.000001 of the expected values are accepted.
Start with brute force
There is no faster way for one query: you must look at every number of every key once, which is O(n · d). Real models do this for every word at once as a matrix multiplication, so a sentence of n words costs O(n² · d). That quadratic cost is why long inputs are expensive.
The mistakes to avoid are leaving out the √d scaling, dividing by d instead of √d, and exponentiating raw scores, which overflows for large vectors.
Hints
0 of 4 shown- Hint 1
In the first example, which key looks most like the query? Compute the three dot products before doing anything else. Does the biggest dot product match your guess?
- Hint 2
The dot product is large when two vectors point the same way and are long. Without the
√ddivision, longer vectors (biggerd) give bigger scores, and softmax would put nearly all the weight on one key. - Hint 3
Once you have the scores, this is plain softmax. What do you subtract from every score so that no
e^xcan overflow? - Hint 4
For each key: dot product with the query, divide by
√d. Find the maximum score. Computee^(score − max)for each, sum them, and divide each by the sum.
Walkthrough and solution
Unlocks when you pass every test. Stuck for a while? Looking is fine too, as long as you then close it and write the code yourself.
How to solve it
Attention weights are two familiar pieces glued together: a dot product to score how well each key matches the query, and a softmax to turn scores into weights that add up to 1. The √d division sits between them.
Trace the first example, query [1, 0, 1, 0], so d = 4 and √d = 2:
- Key 0
[2, 0, 2, 0]: dot 2 + 2 = 4, score 2. - Key 1
[0, 1, 0, 1]: dot 0, score 0. - Key 2
[1, 1, 1, 1]: dot 1 + 1 = 2, score 1. - Max score 2. Shifted: 0, -2, -1. Exponentials: 1, 0.135, 0.368. Total 1.503.
- Weights: 0.665, 0.090, 0.245.
Why √d? If each entry is roughly random, a dot product of d entries spreads out by about √d. Dividing by it keeps scores in a similar range whatever the vector size, so softmax does not collapse onto a single key and training stays steady.
function attentionWeights(query, keys) {
const scale = Math.sqrt(query.length);
const scores = keys.map((key) => {
let dot = 0;
for (let i = 0; i < query.length; i++) dot += query[i] * key[i];
return dot / scale;
});
const top = Math.max(...scores); // stable softmax: shift by the max
const weights = scores.map((s) => Math.exp(s - top));
let total = 0;
for (const w of weights) total += w;
return weights.map((w) => w / total);
}import math
def attention_weights(query, keys):
scale = math.sqrt(len(query))
scores = [sum(q * k for q, k in zip(query, key)) / scale for key in keys]
top = max(scores) # stable softmax: shift by the max
weights = [math.exp(s - top) for s in scores]
total = sum(weights)
return [w / total for w in weights]Run your code to see results here. Run checks the examples. Submit checks every test, including hidden edge cases.