Sliding window to local attention

The fixed-width window you slide over an array to dodge $O(n^2)$ is the same window that turns full attention into cheap, linear local attention - a band hugging the diagonal.

You already slid this exact window - over an array, in Two pointers and sliding window. There the trick was: don't re-scan the whole array from every start; keep a fixed window of width $w$, move it one step, and patch the summary for the single element that entered and the one that left. Quadratic work collapses to linear. Local attention is that same window applied to the attention grid. Full attention lets every token look at all $n$ tokens - an $n\times n$ grid, $O(n^2)$. Local attention restricts each token to a fixed window of $w$ neighbours, so only a thin band around the diagonal gets computed: $O(n\cdot w)$. Same window, same escape from the quadratic wall - just a grid instead of an array.

You already know this window

The array version and the attention version are the same object seen twice. Read the correspondence term for term:

The band around the diagonal

Recall the shape of full self-attention: query $i$ dots against every key $j$, filling an $n\times n$ score grid. Row $i$ has $n$ entries, and there are $n$ rows, so the layer computes $n^2$ scores - that is the $O(n^2)$ “attention wall” from Big-O to the O(n squared) attention wall. Local attention keeps only the entries where the two positions sit within radius $r$ of each other, so every off-band cell is simply never computed:

For a fixed window $w$, that is linear in $n$: double the sequence and the cost merely doubles, instead of quadrupling. This is exactly the recipe in the Longformer - a sliding window plus a handful of designated global tokens - and in Mistral 7B, whose causal sliding window lets token $t$ attend only to the previous $w$ tokens. The masked score becomes:

Softmax then normalises over the surviving in-window keys, so each row still sums to $1$ - it is an honest attention distribution, just a local one. Slide the window and watch the op-count below.

Two things jump out as you drag. First, the lit-cell count tracks $n\cdot w$ and grows linearly with $r$, while full attention stays pinned at $n^2=100$ - the savings are real. Second, push $r$ up toward $4$ ($w=9$) and the band swallows the grid: “skipped” collapses toward $0$. The window only pays off while $w\ll n$. So how does a model built from myopic layers ever see the whole sentence?

Depth buys the global view

Here is the second half of the trick, and it is pure stacking. Layer 1 lets token $q$ see radius $r$. Layer 2 reads layer 1's outputs - but each of those already summarised its own radius $r$ - so through them, token $q$ now indirectly senses radius $2r$. Every layer you add hops the information one more window outward:

Slide the $L$ control in the demo: the grid band never changes - each layer stays exactly as cheap, still $O(n\cdot w)$ - yet the amber receptive field on the token strip fans out with depth. This is the identical mechanism by which stacked $3\times 3$ convolutions grow a CNN's receptive field. A deep tower of cheap local layers routes context across the whole sequence, one window-hop per layer.

import numpy as np

def local_attention(Q, K, V, r):
    n, d = Q.shape
    out = np.zeros((n, d))
    scores_computed = 0
    for i in range(n):                        # each query token
        lo, hi = max(0, i - r), min(n, i + r + 1)  # its window [i-r, i+r]
        s = Q[i] @ K[lo:hi].T / np.sqrt(d)      # only w = hi-lo scores, not n
        a = softmax(s)                          # normalise over the window
        out[i] = a @ V[lo:hi]                    # blend just those values
        scores_computed += hi - lo               # grand total ~ n*w, never n*n
    return out, scores_computed
# n=1024, r=127 (w=255): ~261k scores vs 1.05M for full attention - 4x fewer.

Check yourself