Batched Attention Scores with einsum

The problem

Inside a transformer, every position of the text looks at the other positions and decides how much each one matters. To do that, each position has two vectors of $D$ numbers: a query (what this position is looking for) and a key (what this position offers). The score between a query and a key is their dot product: large when the two vectors point the same way.

The real arrays have four axes, because many scores are computed at once:

This exercise is the shape work you will reuse when you build attention, without the attention itself.

Part 1: scaled dot products

For every text $b$ and head $h$, the score between query position $t$ and key position $s$ is:

where:

• $q_{b,h,t,d}$ is number $d$ of the query vector at position $t$ (text $b$, head $h$);
• $k_{b,h,s,d}$ is number $d$ of the key vector at position $s$;
• the sum multiplies the two vectors number by number and adds up: their dot product;
• dividing by $\sqrt{D}$ keeps the scores around the same size whatever the head size (more on this in the Transformers section).

Example. One text and one head ($B = H = 1$), $D = 4$, two query positions and one key position.

• Queries: $q_0 = (1, 0, 0, 0)$, $q_1 = (0, 2, 0, 0)$. Key: $k_0 = (1, 1, 0, 0)$.
• Dot products: $q_0 \cdot k_0 = 1 \cdot 1 = 1$ and $q_1 \cdot k_0 = 2 \cdot 1 = 2$.
• Divided by $\sqrt{4} = 2$: the scores are $0.5$ and $1.0$.

The result has shape $(1, 1, 2, 1)$: one score for each of the 2 query positions and the 1 key position.

np.einsum writes this in one line. You give each axis a letter, and letters that appear in the inputs but not in the output are summed over. Here the output keeps $b, h, t, s$, and $d$ is summed.

Part 2: merging the heads

After attention, each head has produced a vector of $D$ numbers for every position. The next layer wants, for every position, one long vector with all the heads side by side: head 0's $D$ numbers, then head 1's, and so on. So shape $(B, H, T, D)$ must become $(B, T, H \cdot D)$.

Example. One text, $H = 2$ heads, $T = 2$ positions, $D = 2$:

• head 0 gives [0, 1] at position 0 and [2, 3] at position 1;
• head 1 gives [4, 5] at position 0 and [6, 7] at position 1.

The right result puts each position's heads together: position 0 is [0, 1, 4, 5] and position 1 is [2, 3, 6, 7].

A plain x.reshape(1, 2, 4) has the right shape but returns [0, 1, 2, 3] and [4, 5, 6, 7]: it glues together the two positions of head 0, then the two positions of head 1. That is the wrong grouping.

reshape reads the numbers in memory order, last axis fastest, and never moves them. To group by position, first swap the head and position axes with transpose, so the order is $(B, T, H, D)$, and then reshape.

What to implement

• attention_scores(q, k): q has shape $(B, H, T, D)$ and k has shape $(B, H, S, D)$. Return the $(B, H, T, S)$ array of scaled dot products.
• merge_heads(x): x has shape $(B, H, T, D)$. Return shape $(B, T, H \cdot D)$, where out[b, t] is x[b, 0, t], then x[b, 1, t], and so on, side by side.

Requirements

• No Python loops over batch or heads.
• merge_heads must keep each position's heads together, as in the example. The tests catch a plain reshape.

Hints

• For attention_scores, the einsum string has the form "....,....->....": four letters for q, four for k, four for the output. The two inputs share the letters of $b$, $h$ and $d$.
• For merge_heads, transpose takes the new order of the axes: on a table, m.transpose(1, 0) swaps rows and columns.