NNX Sliding-Window Attention

Attend to the current position and recent history

This exercise adds a local, causal mask to multi-head self-attention. For window size W, query position i can attend to keys from max(0, i - W + 1) through i. W includes the current position, so there are at most W - 1 earlier positions in the window.

The first few positions have shorter windows because no earlier input exists. W=1 allows only self-attention at the same position; W at least as large as the sequence length gives ordinary causal attention.

Build the band by index difference

Scores have shape (H, T, T): head, query position, key position. For each pair, let diff = i - j. Keep it exactly when (diff >= 0) & (diff < W). The first condition excludes the future; the second excludes history outside the window. The resulting (T, T) Boolean mask broadcasts across heads.

For T=5 and W=3, allowed connections are:

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

Using diff <= W would include one extra key. Using abs(diff) < W would allow future keys too. Sliding windows can be symmetric in other applications; this exercise specifically asks for the backward-looking, causal version.

Mask before softmax

Replace scores outside the band with -jnp.inf, then apply softmax along the key axis. Positive W keeps the diagonal, ensuring a finite allowed score in each row under the input contract. Masked weights are zero, and gradients remain well-defined for finite allowed scores.

A finite sentinel such as -1e9 is not a reliable exclusion rule: if allowed logits are even smaller, masked positions can receive weight. Negative infinity avoids that problem. An entirely masked row would still need explicit handling, but valid positive windows do not create one.

A complete example with moving averages

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

class SlidingWindowMHA(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, window):
        T, _ = 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 = jnp.matmul(q, k.transpose(0, 2, 1)) / jnp.sqrt(Dh)
        i = jnp.arange(T)[:, None]
        j = jnp.arange(T)[None, :]
        diff = i - j
        in_band = (diff >= 0) & (diff < window)
        scores = jnp.where(in_band, scores, -jnp.inf)
        weights = jax.nn.softmax(scores, axis=-1)
        per_head = jnp.matmul(weights, v)
        concat = per_head.transpose(1, 0, 2).reshape(T, H * Dh)
        return self.out_proj(concat)

model = SlidingWindowMHA(2, 2, nnx.Rngs(0))
model.q_proj.kernel[...] = jnp.zeros((2, 2))
model.k_proj.kernel[...] = jnp.zeros((2, 2))
model.v_proj.kernel[...] = jnp.eye(2)
model.out_proj.kernel[...] = jnp.eye(2)
x = jnp.array([[2., 4.], [6., 8.], [10., 12.], [14., 16.], [18., 20.]])
print(model(x, window=3))
# [[2., 4.], [4., 6.], [6., 8.], [10., 12.], [14., 16.]]
changed = x.at[0].set(jnp.array([100., 200.]))
assert jnp.allclose(model(changed, 3)[3:], model(x, 3)[3:])

Q and K are zero in this example, so the allowed keys have uniform weights. With identity V and output kernels and zero default biases, each output is the mean of the input rows inside its window. Position 3 averages rows 1, 2 and 3; it no longer uses row 0. Restore the default random initialization when implementing the graded entrypoint.

A local mask is not a sparse implementation

The code here forms the entire (H, T, T) score and weight arrays and uses dense matrix multiplications. Its attention arithmetic remains O(T * T * D) and its attention matrices use O(H * T * T) storage, where D is model width. Masking those arrays does not avoid computing the excluded dot products. Halving T roughly quarters this quadratic part.

An implementation that computes only the allowed windows can reduce attention arithmetic to O(T * min(W, T) * D). That requires different kernels or data access, not just a Boolean mask. Projection costs remain separate. This exercise teaches the mask semantics, not that optimization, and it does not implement a rolling KV cache.

Across L layers of this local causal attention, information can in principle travel back up to L * (W - 1) positions, giving at most 1 + L * (W - 1) positions in a token's receptive field, clipped at the sequence start. This assumes no other operation mixes sequence positions; it describes possible paths, not guaranteed learned influence.

Choose Sliding window in the visualization to inspect a backward band. It shows allowed connections rather than numerical attention weights; its displayed window setting is separate from the function's input W.