Multi-Head Self-Attention with Flax

Compute several attention patterns

Multi-head attention uses several sets of query, key and value features. Each head computes its own attention weights and combines values with those weights. An output projection then combines the heads' results. Training can make heads useful in different ways, but it does not assign each head a guaranteed role such as syntax or meaning.

In this exercise, all three projections use the same input sequence x. That makes it self-attention. The module initializes new parameters from the supplied seed and performs one forward pass; it does not train them.

Project first, then separate heads

Write the input shape as (T, D_in). Let H = num_heads, D = qkv_features and d_head = D // H.

Linen learns a separate affine projection for queries, keys and values. Each projection reads all D_in input features and produces D features across the heads. It does not simply split the original input into chunks. D may differ from D_in.

Within each head, divide query/key dot products by sqrt(d_head) and apply softmax across keys. The final projection combines all head features. Its default output width is the input query width, D_in.

Choose the head dimensions

Both H and D must be positive, and D must be divisible by H. The input width does not need to be divisible by H: its projection creates the required head dimensions.

At fixed D = 8, two heads have width 4, four heads have width 2, and eight heads have width 1. Width 1 is valid, though it gives each head fewer features. More heads at the same total width does not guarantee a better model or increase the number of projection parameters here.

Initialize and apply the module

This complete example deliberately uses different input and projection widths, so their roles are visible:

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

x = jnp.arange(24, dtype=jnp.float32).reshape(4, 6) / 24.0
attn = nn.MultiHeadDotProductAttention(num_heads=2, qkv_features=8)
variables = attn.init(jax.random.PRNGKey(0), x)
out = attn.apply(variables, x)
params = variables["params"]

assert params["query"]["kernel"].shape == (6, 2, 4)
assert params["out"]["kernel"].shape == (2, 4, 6)
assert out.shape == (4, 6)
assert jnp.allclose(attn.apply(variables, x), out)

init creates the variables. apply reads the supplied variables and evaluates the module; it does not update the weights. Constructing the module alone does not initialize its parameters.

A two-dimensional input represents one unbatched sequence. The module also accepts leading batch dimensions, but this problem supplies only (T, D_in) inputs.

Keep the defaults used by the exercise

Use biases and Linen's default parameter initializers. Dropout is zero, decoding is disabled, and no mask is passed. Each query can attend to every position, including later ones. There are no position encodings, residual connections or normalization layers around this attention call.

Calling apply with only x makes it the source for all three separate projections. Those projections generally produce different arrays even though their input is shared. The initialized weights are not pretrained.