Skip to content

← Flax step 30 of 100

Medium Primitives

Block-Diagonal Attention Mask

Jump to the task

Keep attention inside each block

Divide the input sequence into consecutive blocks of B positions. Each query can attend to every key in its own block and no keys in other blocks. The last block can be shorter than B.

This pattern is bidirectional within a block: a query can read later positions in the same block. It is not a causal language-model mask. This exercise initializes one attention layer and applies the pattern; it does not implement a sparse attention kernel.

Compare block IDs

Position i belongs to block i // B. The mask permits a query/key pair exactly when their block IDs match:

mask[i, j] = (i // B) == (j // B)

For T = 6 and B = 2, the allowed entries are:

[[1, 1, 0, 0, 0, 0],
 [1, 1, 0, 0, 0, 0],
 [0, 0, 1, 1, 0, 0],
 [0, 0, 1, 1, 0, 0],
 [0, 0, 0, 0, 1, 1],
 [0, 0, 0, 0, 1, 1]]

Rows are queries and columns are keys. Query zero can read the later key one, but cannot read key two. Softmax normalizes each query's weights over its allowed keys. The weights sum to one per row; raw attention scores do not.

Initialize and check a block-isolated layer

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

x = jnp.arange(30, dtype=jnp.float32).reshape(6, 5) / 30.0
T, B = x.shape[0], 2
positions = jnp.arange(T)
block_ids = positions // B
mask = block_ids[:, None] == block_ids[None, :]
attn = nn.MultiHeadDotProductAttention(num_heads=2, qkv_features=8)
variables = attn.init(jax.random.PRNGKey(0), x, mask=mask)
out = attn.apply(variables, x, mask=mask)

changed = x.at[2:].add(1.0)
changed_out = attn.apply(variables, changed, mask=mask)
assert out.shape == (6, 5)
assert jnp.allclose(out[:2], changed_out[:2])
assert jnp.array_equal(mask[0], jnp.array([1, 1, 0, 0, 0, 0], dtype=bool))

The first block's outputs stay unchanged when the other blocks' inputs change, with fixed parameters. A boolean (T, T) mask broadcasts over heads. An equivalent zero/one float mask also works; casting to float32 is optional and does not change broadcasting.

Boundary cases

  • B = 1: each query reads only itself. The output is its projected value followed by the final output projection, including their biases.
  • B >= T: all positions share one block, giving full unmasked attention.
  • T % B != 0: the remaining positions form a smaller final block. Integer division handles this without padding.

B must be positive. Every query belongs to a nonempty block, so its own key is always allowed.

Scope and cost

This dense Linen call still computes a score tensor of shape (H, T, T) before masking. Blocking entries does not reduce that tensor's size or make the call use sparse computation.

Computing attention separately within blocks could reduce query/key interactions. For block lengths b_1, ..., b_n, the count per head is b_1^2 + ... + b_n^2. If T is divisible by B, that count is T * B. These counts describe allowed pairs, not a measured speedup of this code.

Fixed blocks also remain isolated across stacked layers if every cross-position operation uses the same partition. Connecting blocks requires another path, such as shifted partitions or a different mask. For packed documents of unequal lengths, use document IDs instead of position // B. Autoregressive training additionally requires causal restrictions within each document.

Problem

Implement mha_masked_block(seed, x, num_heads, qkv_features, block_size):

  1. Convert seed, num_heads, qkv_features and block_size to Python integers.
  2. Construct the (T, T) mask by comparing block IDs.
  3. Build nn.MultiHeadDotProductAttention(num_heads=H, qkv_features=D) with the remaining options at their defaults.
  4. Initialize with jax.random.PRNGKey(seed), x and mask=mask.
  5. Apply with the returned variables, x and the same mask.
  6. Flatten the (T, D_in) result with .reshape(-1) and return it.

x is a finite float32 array of shape (T, D_in). All dimensions, H and B are positive, and D is a positive multiple of H. D_in need not equal D or be divisible by H. The output has length T * D_in, with each position's features together. Keep default biases, initializers, zero dropout and disabled decoding. Add no position encodings, residual connections or normalization. Under jax.jit, make all arguments except x static to match the integer conversions.

Loading visualization…