Split heads, merge heads, multi-head attention

Multi-head attention runs several small attentions in parallel. Each head sees one slice of the features. The heads are then joined again and mixed by an output projection. The hard part is the reshape: a wrong axis order gives the right shape but the wrong numbers.

Task

Write three functions.

1. split_heads(x, n_heads) turns x of shape (B, T, D) into an array of shape (B, H, T, d_head), with H = n_heads and d_head = D / H. Head h holds the columns h * d_head to (h + 1) * d_head - 1 of x, in the same order. So split_heads(x, H)[b, h, t, j] equals x[b, t, h * d_head + j].
2. merge_heads(x) is the inverse. It turns (B, H, T, d_head) into (B, T, H * d_head), so that merge_heads(split_heads(x, H)) equals x.
3. multi_head_attention(x, w_q, w_k, w_v, w_o, n_heads) returns the output and the attention weights of multi-head self-attention.

The steps of multi_head_attention are:

1. Q = x W_q, K = x W_k, V = x W_v. There are no bias terms.
2. Split Q, K and V into heads.
3. In each head, compute \mathrm{softmax}(Q_h K_h^T / \sqrt{d_{head}})\, V_h. The softmax is over the key axis. There is no mask.
4. Merge the head outputs and multiply by W_o.

Input and output

• split_heads returns a float array of shape (B, H, T, D / H).
• merge_heads takes an array of shape (B, H, T, d_head) and returns a float array of shape (B, T, H * d_head).
• multi_head_attention returns a tuple (output, weights). output has shape (B, T, D). weights has shape (B, H, T, T), and weights[b, h, i, j] is the weight that query i gives to key j in head h.

Constraints

• 1 <= B, T <= 16 and 1 <= H <= D <= 64.
• When D is not divisible by n_heads, split_heads raises ValueError.
• Activations can be large. The output must stay finite.
• Do not change the input arrays.
• Use numpy only.

Examples

x = np.arange(8.0).reshape(1, 2, 4)
split_heads(x, 2)

[[[[0. 1.]
   [4. 5.]]

  [[2. 3.]
   [6. 7.]]]]

merge_heads of that result gives x again.

x = np.array([[[1.0, 0.0, 0.0, 1.0], [0.0, 2.0, 1.0, 0.0], [1.0, 1.0, 1.0, 1.0]]])
eye = np.eye(4)
output, weights = multi_head_attention(x, eye, eye, eye, eye, 2)

output
[[[0.80222419 0.79666372 0.59888791 0.80222419]
  [0.23208206 1.72252957 0.80222419 0.59888791]
  [0.59888791 1.20333628 0.75174492 0.75174492]]]
weights
[[[[0.40111209 0.19777581 0.40111209]
   [0.04538836 0.76791794 0.1866937 ]
   [0.19777581 0.40111209 0.40111209]]

  [[0.40111209 0.19777581 0.40111209]
   [0.19777581 0.40111209 0.40111209]
   [0.24825508 0.24825508 0.50348984]]]]