A Sentence Reads Itself

Point attention at its own sequence and every token queries every token, itself included. What comes out is the attention map: one row per token doing the attending, one column per token being attended to, and every row summing to $1$ because focus is a fixed budget.

The sequence arrives as $n$ token vectors stacked as the rows of $X$, each of width $d_{model}$. Two learned matrices, $W_Q$ and $W_K$, both of shape $d_{model} \times d_k$, turn every token into a query and a key:

$Q = X W_Q \qquad K = X W_K \qquad \text{weights} = \mathrm{softmax}\!\left(\frac{Q K^{\top}}{\sqrt{d_k}}\right)$

with the softmax taken across each row.

Task: write self_attention_weights(tokens, w_query, w_key) returning the $n \times n$ weight matrix, every entry rounded to 4 decimal places.

• Entry [i][j] is how much of token i's attention lands on token j.
• Every row sums to $1$ (up to rounding). There is no masking here — every token may look anywhere, itself included.
• The scale is $\sqrt{d_k}$: the width of the projected queries and keys, which is not in general the width of the token vectors.
• A one-token sequence has nowhere else to look, so the map is [[1.0]].

$Q$ and $K$ come from the same $X$ but through two different matrices, and that is exactly what lets the map be lopsided: token 3 can attend heavily to token 1 while token 1 all but ignores token 3. One shared matrix would force the grid to be symmetric and throw that freedom away.