Cross-Attention Block

Cross-attention lets one sequence read another. Here the decoder sequence provides the queries, while the encoder sequence provides the keys and values. Every decoder position forms its own weighted mixture of encoder values. The two sequences can have different lengths.

This exercise adds an output projection, two residual connections, two layer normalizations and a small feed-forward network. The result stays aligned with decoder positions, even when the encoder sequence is longer.

Inputs and output

Implement cross_attention_block(decoder_x, encoder_out, w_q, w_k, w_v, w_o, w_mlp1, w_mlp2, num_heads).

• decoder_x: floating-point array of shape (N, T_dec, d_model).
• encoder_out: floating-point array of shape (N, T_enc, d_model).
• All six weight matrices: shape (d_model, d_model).
• num_heads: a positive integer that divides d_model evenly.

All dimensions are positive. Return shape (N, T_dec, d_model). There is no causal or padding mask: every decoder query reads every encoder position in its own batch item. There are no biases, dropout or learned normalization parameters. The feed-forward hidden width is d_model in this exercise.

Attention stage

1. Compute Q = decoder_x @ w_q, K = encoder_out @ w_k and V = encoder_out @ w_v. Use weights in the supplied orientation.
2. Set d_head = d_model // num_heads. Split the projected feature axis into (num_heads, d_head) and swap the head and sequence axes. Q becomes (N, num_heads, T_dec, d_head); K and V become (N, num_heads, T_enc, d_head).
3. Multiply Q by K with only K's last two axes swapped, then divide by sqrt(d_head). Scores have shape (N, num_heads, T_dec, T_enc). Apply a stable softmax over the final axis, which contains encoder keys.
4. Multiply the weights by V. Swap head and sequence axes back, then reshape to (N, T_dec, d_model), keeping heads in their original order.
5. Project the joined heads and add the decoder residual: x = layer_norm(decoder_x + joined @ w_o).

K and V come from the encoder, but the residual comes from the decoder. Using the encoder as the residual would also fail when the lengths differ.

Feed-forward stage

Apply the network independently to each decoder position:

hidden = x @ w_mlp1
ffn_out = gelu(hidden) @ w_mlp2
output = layer_norm(x + ffn_out)

Use this elementwise tanh approximation of GELU, where ** denotes a power:

gelu(t) = 0.5 * t * (1 + tanh(sqrt(2 / pi) * (t + 0.044715 * t**3)))

Do not substitute the exact normal-CDF version of GELU; it gives slightly different values.

Normalization and shape checks

Both calls to layer_norm subtract the mean over the final feature axis and divide by sqrt(population_variance + 1e-5). Population variance is the mean squared deviation, with divisor d_model. Keep batch and sequence positions separate, and add no learned scale or offset.

Both normalizations follow their residual additions. Even if both branches output zeros, the specified result is layer_norm(layer_norm(decoder_x)). Do not skip the second normalization: epsilon means it is not exactly redundant.

For T_dec = 2 and T_enc = 3, each head has a 2-by-3 score matrix and produces two output vectors. The cross-attention panel below shows this rectangular pattern with different example lengths. Filled cells show allowed connections, not their attention weights.