NNX MHA From Scratch

Several attention heads, one output

Multi-head self-attention learns several ways to compare positions in the same sequence. Each head computes its own attention weights and combines its value vectors. Concatenate those results, then apply a learned output projection to mix features from the heads.

NNX provides a built-in nnx.MultiHeadAttention. Here you build a smaller version from four nnx.Linear layers to see the projections and axis changes explicitly. It has no mask, dropout, position encoding, residual connection, or normalization. Every position can attend to every position.

Project first, then split

Let D = d_model, H = num_heads, and Dh = D // H, with positive H dividing D. The input x has shape (T, D).

Create four independent nnx.Linear(D, D, rngs=rngs) layers, in this order: q_proj, k_proj, v_proj, out_proj. Use their default initializers and biases, passing the same RNG container through all four constructors. Initialization order matters to the seeded values.

Q, K and V each start as a projection of the full input feature vector. Only then split each projected vector into H contiguous chunks of width Dh. Heads do not simply receive disjoint slices of the original input. Each projection kernel can combine all D input features.

Keep the axes in order

After projection, reshape (T, D) to (T, H, Dh), then transpose to (H, T, Dh). For each head:

• Scores are Q @ K.transpose(0, 2, 1) / sqrt(Dh), shaped (H, T, T).
• Softmax over the last axis gives weights across keys for each query.
• Multiplying by V gives per-head outputs shaped (H, T, Dh).

The scale uses the width of one head, not the full model width. Transpose outputs back to (T, H, Dh) before reshaping to (T, D). A direct reshape from (H, T, Dh) would mix sequence positions with head chunks. Finally apply out_proj, including its bias.

A complete example with known weights

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

class MHA(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):
        T, D = x.shape
        H, Dh = self.num_heads, self.head_dim
        q = self.q_proj(x).reshape(T, H, Dh).transpose(1, 0, 2)
        k = self.k_proj(x).reshape(T, H, Dh).transpose(1, 0, 2)
        v = self.v_proj(x).reshape(T, H, Dh).transpose(1, 0, 2)
        scores = (q @ k.transpose(0, 2, 1)) / jnp.sqrt(Dh)
        weights = jax.nn.softmax(scores, axis=-1)
        per_head = weights @ v
        combined = per_head.transpose(1, 0, 2).reshape(T, D)
        return self.out_proj(combined)

model = MHA(4, 2, nnx.Rngs(0))
model.q_proj.kernel[...] = jnp.zeros((4, 4))
model.k_proj.kernel[...] = jnp.zeros((4, 4))
model.v_proj.kernel[...] = jnp.eye(4)
model.out_proj.kernel[...] = jnp.flip(jnp.eye(4), axis=1)
model.out_proj.bias[...] = jnp.array([1., 0., -1., 2.])
x = jnp.array([[1., 2., 3., 4.], [5., 6., 7., 8.]], dtype=jnp.float32)
print(model(x))  # [[7., 5., 3., 5.], [7., 5., 3., 5.]]

The default projection biases start at zero. Setting Q and K kernels to zero makes every score zero, so both heads average the two value rows. The concatenated average is [3, 4, 5, 6]. The chosen output kernel reverses its feature order, then the output bias gives [7, 5, 3, 5]. These weight assignments are only for the example; keep the initialized weights when solving the exercise.

The four biased projections contain 4 * (D * D + D) scalar parameters. For fixed D, changing the head count changes the attention calculation, not this parameter count. Heads need not learn different patterns merely because they occupy different chunks.

Explore the split

The visualization partitions an already-projected feature vector. It shows head widths, not attention probabilities or learned relationships.