Skip to content

← Flax step 28 of 100

Hard Primitives

Sliding-Window Self-Attention

Jump to the task

Attend to a recent prefix

Causal sliding-window attention lets each query read itself and a limited number of earlier positions. Here the window size W includes the current position, so at most W - 1 earlier keys are visible.

This exercise builds the window mask and passes it to Linen's multi-head attention module. It initializes fresh weights and returns one forward pass. It does not implement a complete language model or a sparse kernel.

State the boundary precisely

Query i may attend to key j when:

0 <= i - j < W

Equivalently, allowed key indices run from max(0, i - W + 1) through i, inclusive. The lower bound on distance blocks the future; the upper bound limits how far into the past the query can look.

For T = 6 and W = 3, the mask is:

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

Query five can read keys three, four and five. Using i - j <= W would allow one extra earlier key. Dropping i - j >= 0 would allow every future key, not just nearby ones. A symmetric local window needs a bound on the absolute distance instead.

W = 1 permits only self-attention at each position. If W >= T, the mask becomes the full causal lower triangle. Keep W positive so every query has at least its diagonal entry available.

Build and apply the mask

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, W = x.shape[0], 3
i = jnp.arange(T)[:, None]
j = jnp.arange(T)[None, :]
mask = (i >= j) & (i - j < W)
assert jnp.array_equal(mask[-1], jnp.array([0, 0, 0, 1, 1, 1], dtype=bool))

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)

# The last query cannot read row zero through this single layer.
changed = x.at[0].add(1.0)
changed_out = attn.apply(variables, changed, mask=mask)
assert out.shape == (6, 5)
assert jnp.allclose(out[-1], changed_out[-1])

A (T, T) mask broadcasts over the head axis. A (1, T, T) mask also works for this unbatched attention. Linen accepts boolean masks and equivalent zero/one float masks; a float cast is optional.

A dense mask does not make the calculation sparse

This implementation still constructs attention scores for every query/key pair before masking them. Its score tensor has shape (H, T, T). It therefore retains quadratic score storage and attention work in T for fixed head dimensions.

An implementation that computes only the allowed windows can reduce the number of query/key interactions to order T * W. Passing a band mask to this ordinary dense attention call does not provide that optimization.

What changes when layers are stacked?

A single layer can read up to W - 1 positions backward. With L such layers, information can travel up to L * (W - 1) positions backward through intermediate representations. The potential input range at position i is therefore:

[max(0, i - L * (W - 1)), i]

This assumes all cross-position mixing follows the same window and other operations act separately at each position. It describes possible dependencies, not a guarantee that training learns useful long-range behavior. This exercise applies only one layer.

Problem

Implement mha_sliding_window(seed, x, num_heads, qkv_features, window):

  1. Convert seed, num_heads, qkv_features and window to Python integers.
  2. Build the (T, T) mask for 0 <= i - j < W.
  3. Construct nn.MultiHeadDotProductAttention(num_heads=H, qkv_features=D) with its other 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) output with .reshape(-1) and return it.

x is a finite float32 array of shape (T, D_in) with positive dimensions. H, D and W are positive and D is divisible by 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 encoding, residual connection or normalization. Under jax.jit, make all arguments except x static to match the integer conversions.

Choose Sliding window in the diagram to compare its allowed entries with the full causal pattern.

Loading visualization…