Cross-Attention with Flax MHA

Read from a different sequence

Cross-attention takes queries from one input and keys and values from another. For example, a decoder can use its hidden states as the query source and an encoder's outputs as the key/value source. Each query produces a weighted combination of information from the source sequence.

Self-attention uses one input as the source of all three projections. Cross-attention separates those sources. Their sequence lengths can be equal or different; unequal lengths are not what defines cross-attention.

This exercise initializes one Linen attention module and performs a forward pass. It does not train a model, run an encoder or generate tokens.

Separate projections allow different input widths

Call the module through init and apply with the query input first and the key/value input second. Linen projects them separately:

Here H = num_heads, D = qkv_features and d_head = D // H. The query projection reads D_q features; the separate key and value projections each read D_kv. The two input feature widths need not match. Neither input width needs to be divisible by H, but the total projected width D does.

Each head divides query/key dot products by sqrt(d_head) and applies softmax over key positions. After the heads combine projected values, the final projection returns to the query input width D_q by default. The output width therefore does not come from the key/value input width.

A complete example with unequal lengths and widths

import jax
import jax.numpy as jnp
import flax.linen as nn

x_q = jnp.arange(12, dtype=jnp.float32).reshape(2, 6) / 12.0
x_kv = jnp.arange(15, dtype=jnp.float32).reshape(3, 5) / 15.0
cross = nn.MultiHeadDotProductAttention(num_heads=2, qkv_features=8)
variables = cross.init(jax.random.PRNGKey(0), x_q, x_kv)
out = cross.apply(variables, x_q, x_kv)
params = variables["params"]

assert params["query"]["kernel"].shape == (6, 2, 4)
assert params["key"]["kernel"].shape == (5, 2, 4)
assert params["value"]["kernel"].shape == (5, 2, 4)
assert params["out"]["kernel"].shape == (2, 4, 6)
assert out.shape == (2, 6)

Each of the two query positions can read all three key/value positions. Keys and values share an input here, but their projection weights are separate, so the projected arrays generally differ.

Once initialized, the parameter shapes are tied to D_q and D_kv. Reusing those variables with different feature widths is not the same as creating a module for those new widths. Sequence lengths can change without changing the projection shapes.

Keep input order and masking explicit

Swapping x_q and x_kv changes which input requests information and which supplies it. With newly initialized parameters, it also changes the output length and default output width to match the new query input.

This exercise passes no mask. Every query can attend to every key/value position, and dropout is zero. A causal triangle over these two different sequences is not automatically appropriate. For example, an offline encoder-decoder model may let each decoder position read the entire encoded source. Padding or streaming restrictions require their own masks.

Use the default biases, initializers and disabled decoding. Do not add position encodings, residual connections or normalization. The parameters are newly initialized, not pretrained; apply does not train them.