NNX Cross-Attention

Query one sequence using another

Cross-attention combines information from a source sequence for each position in a query sequence. Query vectors come from x_q; keys and values come from x_kv. Each query gets its own weighted average of the source values, so the output has one row per query.

For example, an encoder-decoder model can let a decoder position use representations from the encoded input. This exercise covers the attention layer itself, without a full encoder, decoder, mask or generation loop.

Keep the two lengths separate

Let query length be T_q, source length be T_k, model width be D, and head count be H. This exercise gives both inputs width D, with H dividing D and head width Dh = D // H.

• Project x_q through q_proj, then reshape to (H, T_q, Dh).
• Project x_kv through k_proj and v_proj, then reshape each to (H, T_k, Dh).
• Compute scaled query-key scores shaped (H, T_q, T_k).
• Softmax over the last axis normalizes across source keys for each query.
• Multiply by values to obtain (H, T_q, Dh), combine heads into (T_q, D), then apply out_proj.

To obtain the head-first shape, first reshape a projection to (sequence_length, H, Dh), then transpose axes (1, 0, 2). Reverse that transpose before concatenating heads. Scale scores by sqrt(Dh).

K and V rows must correspond to the same source positions. They use separate learned projections of x_kv here. More general cross-attention can accept different input feature widths or separately supplied key and value arrays, provided the projected shapes are compatible. Equal input widths are this exercise's contract, not a universal restriction.

A complete example with unequal lengths

import jax
import jax.numpy as jnp
from flax import nnx

class CrossMHA(nnx.Module):
    def __init__(self, d_model, num_heads, rngs):
        assert num_heads > 0 and d_model % num_heads == 0
        self.num_heads = num_heads
        self.head_dim = d_model // num_heads
        self.q_proj = nnx.Linear(d_model, d_model, rngs=rngs)
        self.k_proj = nnx.Linear(d_model, d_model, rngs=rngs)
        self.v_proj = nnx.Linear(d_model, d_model, rngs=rngs)
        self.out_proj = nnx.Linear(d_model, d_model, rngs=rngs)

    def __call__(self, x_q, x_kv):
        Tq, _ = x_q.shape
        Tk, _ = x_kv.shape
        H, Dh = self.num_heads, self.head_dim
        q = self.q_proj(x_q).reshape(Tq, H, Dh).transpose(1, 0, 2)
        k = self.k_proj(x_kv).reshape(Tk, H, Dh).transpose(1, 0, 2)
        v = self.v_proj(x_kv).reshape(Tk, H, Dh).transpose(1, 0, 2)
        scores = jnp.matmul(q, k.transpose(0, 2, 1)) / jnp.sqrt(Dh)
        weights = jax.nn.softmax(scores, axis=-1)
        per_head = jnp.matmul(weights, v)
        concat = per_head.transpose(1, 0, 2).reshape(Tq, H * Dh)
        return self.out_proj(concat)

model = CrossMHA(2, 2, nnx.Rngs(0))
for name in ("q_proj", "k_proj", "v_proj", "out_proj"):
    getattr(model, name).kernel[...] = jnp.eye(2)
x_q = jnp.array([[jnp.log(2.), jnp.log(3.)], [0., 0.]], dtype=jnp.float32)
x_kv = jnp.array([[1., 0.], [0., 1.], [0., 0.]], dtype=jnp.float32)
print(model(x_q, x_kv))  # approximately [[0.5, 0.6], [0.3333333, 0.3333333]]

The two query rows attend to three source rows. With the identity projections and zero default biases, each head has width one. For the first query, head 0 has relative weights [2, 1, 1] and source values [1, 0, 0], giving 2/4 = 0.5. Head 1 has relative weights [1, 3, 1] and values [0, 1, 0], giving 3/5 = 0.6.

The zero query gives uniform weights in both heads, producing [1/3, 1/3]. The identity output projection leaves those results unchanged. Set weights this way only for the example; the graded function uses random initialization.

Masks depend on the task

This exercise is unmasked: every query can attend to all source rows. Cross-attention is not inherently mask-free. Padding masks exclude invalid source positions, and streaming or aligned tasks may restrict which source positions are available to a query. Such restrictions need the task's alignment or availability rule, not an assumed triangle based only on row numbers from two different sequences.

In an offline encoder-decoder setting, the complete source is often available even while the decoder's own self-attention is causal. That is one common arrangement, not a requirement for every cross-attention layer.

The visualization compares self, causal and cross-attention connections. Use its cross-attention panel to follow the two inputs; it illustrates connections, not this example's computed weights.