Cross-Attention & Encoder Masking

The problem, from first principles

Self-attention lets a sequence look at itself. Cross-attention lets one sequence look at a different one: queries come from the decoder, while keys and values come from the encoder's output. This is how a translation model consults the source sentence, how a captioning model consults image features and how a diffusion model consults a text prompt. The decoder and encoder usually have different lengths, and the encoder side is padded in a batch, so padded encoder positions must be masked out before the softmax or the decoder would attend to nothing.

From theory to code

Implement cross_attention.

Constraints

• x_dec is (Tq, d) and enc is (Tk, d). W_q, W_k and W_v are (d, d_h) projections: queries come from x_dec @ W_q and keys and values from enc @ W_k and enc @ W_v.
• Scores are Q @ K.T / sqrt(d_h), shape (Tq, Tk). enc_mask is a length-Tk array of 1 (real) and 0 (padding): scores at padding positions are set to -inf before the softmax.
• Return (output, weights) with output of shape (Tq, d_h) and weights of shape (Tq, Tk).
• At least one encoder position is real.

Hints