Hilbert-Guided Sparse Local Attention

TL;DR

• Local attention for images (Swin windows, slide attention, neighborhood attention) is sparse on paper. A block-sparse kernel such as FlexAttention only saves time on blocks of the attention matrix that are completely empty.
• Flattening an image in row-major order scatters each 2D window across many rows of the sequence, so most non-empty blocks are partial and need element-wise masking. FlexAttention running plain window attention is slower than a dense kernel in every setting the paper tests.
• The fix is a reordering, not a new kernel: flatten the image along a Hilbert curve, then build windows and neighborhoods on the 1D sequence. Nearby tokens stay nearby, the attention pattern collapses toward the diagonal, and far more blocks are empty.
• On an RTX 3080, Hilbert window attention runs 4.0× faster than dense window attention at 128×128 tokens, and Hilbert slide attention runs 18× faster than naive slide attention at 56×56. Swin-T and NAT-mini variants built on these patterns lose at most 0.2 points of ImageNet top-1 accuracy.

Local attention is sparse, but the kernel cannot see it

Global self-attention over N image tokens costs O(N^2). Local attention restricts each query to a neighborhood: a fixed window in Swin Transformer, a sliding window in Slide Transformer, a neighborhood that shifts inward at the image border in NAT. Most of the N \times N attention matrix is then masked out.

Block-sparse kernels turn that mask into saved work. FlexAttention tiles the matrix into blocks of b_q \times b_k entries and classifies each one (sparse attention patterns covers the general idea):

• Empty block: every entry is masked. The kernel skips it and never loads its keys and values.
• Full block: every entry is kept. The kernel runs a plain dense tile.
• Partial block: some entries are masked. The kernel computes the whole tile and applies the mask element by element.

Empty blocks are cheapest, then full blocks, then partial blocks. The paper's Figure 1 shows FlexAttention's speedup rising with the share of empty blocks, most steeply above about 80%.

The catch is the order of the tokens. Vision models flatten the feature map row by row, so a 2D window such as tokens (1, 2, 5, 6) in a 4×4 map sits in two separate stretches of the sequence. Its attention lands in several blocks, and each of those blocks is only partly used. At 128×128 tokens with 16×16 windows, 87.5% of the blocks are already empty, yet window attention on FlexAttention takes 5.68 ms against 2.74 ms for a dense per-window kernel. Every remaining block is partial, and masking them costs more than the skipped blocks save.

Reordering tokens along a Hilbert curve

A Hilbert curve visits every cell of a grid, moving one step at a time, and fills each aligned square completely before moving to the next. Consecutive positions on the curve are neighbors in 2D, and tokens close in 2D tend to stay close in the sequence. The paper flattens the feature map along this curve and then defines local attention directly on the 1D sequence:

• Hilbert Window Attention (HWA): a window is a run of W^2 consecutive tokens. When the map and window sides are powers of two, this run is exactly an aligned W \times W square.
• Hilbert Slide Attention (HSA): each token attends to the K^2 tokens around it in the sequence, a band along the diagonal.
• Hilbert Neighborhood Attention (HNA): the same band, shifted inward at the ends of the sequence, which is 1D neighborhood attention.

The paper's toy example (Figure 3) uses 16 tokens, 2×2 windows and 4-token blocks. In row-major order the first eight tokens form four half-used partial blocks. In Hilbert order the first window is tokens 1 to 4 and the second is tokens 5 to 8, so the same attention forms two full blocks and two empty ones. Slide attention keeps its empty-block count but halves its partial blocks.

The authors' code uses a generalized Hilbert curve that works for any height and width. When the map or window size is not a power of two, a Hilbert window is an irregular region rather than a square, but its tokens are still adjacent in the image.

Why empty blocks translate into time

The paper models the runtime of a block-sparse kernel as the work of all its thread blocks (CTAs) divided by how many the GPU can run at once:

Here \alpha is the fixed cost of one CTA (launch, loading its query block, setup), \beta is the cost of one non-empty block (loading keys and values, the QK^\top tile, score modification, the online softmax and the product with V), r_i is the number of non-empty blocks the i-th CTA processes, and P_{\text{eff}} is the effective parallelism across the streaming multiprocessors. Hilbert order lowers the r_i, and for the window pattern it also turns partial blocks into full ones.

Block size and window size interact. At 128×128 tokens with 16×16 windows, 128- and 256-token blocks leave 98.44% of blocks empty and run 4.0× faster than dense window attention. With 1024-token blocks each block spans several windows, sparsity falls to 93.75%, and the speedup disappears (2.78 ms against 2.74 ms).

Two backbones: HWT and HNT

Hilbert Window Transformer (HWT) follows Swin Transformer. The patch tokens are reordered along the Hilbert path, which depends only on the feature map size, so it is computed once and cached. Blocks come in pairs as in Swin: the first applies HWA, and the second applies Hilbert shifted window attention, which moves the windows forward along the 1D sequence by a fixed offset. Tokens from the start and the end of the sequence that land in the same shifted window are not neighbors in 2D, so their attention is masked out. Because Hilbert windows can be irregular, Swin's per-window relative position bias no longer fits, and HWT uses a single relative position bias over the whole feature map. Both attention types are expressed as a FlexAttention mask_mod and score_mod, with no hand-written kernel.

Hilbert Neighborhood Transformer (HNT) follows NAT. After the Hilbert reordering, 2D neighborhood attention becomes 1D neighborhood attention, which can run on NATTEN's na1d or on FlexAttention. The attention pattern is the same either way, so the backend changes speed but not accuracy.

Implementation

Both backbones use the same attention core: reorder image tokens, define local connections in the new sequence, compute attention, and restore the spatial layout. Follow that core below with seeded Q/K/V tensors, a dense CPU reference, and an optional compiled CUDA FlexAttention backend.

Each row pairs a reference to the original conference PDF with our tensor explanation and source extracted from the downloadable example. The three images are original figure or equation details, not our notation. This is an educational attention implementation, not a pretrained classifier or a reproduction of the paper’s benchmarks. The authors’ official repository provides the full HWT and HNT backbones and pretrained weights.

1. Build and cache the image-token order

Original paper

Enlarge figure

Math and tensor shapes

N=HW,\quad p_i=r_iW+c_i

Here H and W are the image-token grid dimensions. The permutation p maps each sequence index i to one row-major flat index, visiting all N positions exactly once.

The wrapper calls the complete generalized Gilbert recursion in the download. It supports rectangular and non-power-of-two grids and caches the immutable path by geometry. This does not guarantee that every consecutive pair is Manhattan-adjacent on every odd rectangle.

Python implementation

python

@lru_cache(maxsize=128, typed=True)
def hilbert_order(height, width):
    """Return cached sequence index -> row-major flat index for any rectangle."""
    _positive_integer(height, "height")
    _positive_integer(width, "width")
    path = []
    if width >= height:
        _generate2d(0, 0, width, 0, 0, height, path)
    else:
        _generate2d(0, 0, 0, height, width, 0, path)
    return tuple(row * width + col for col, row in path)

2. Gather Q, K, and V along the token axis

Original paper

Math and tensor shapes

Q'_i=Q_{p_i}

Apply the same gather to K and V: K'_i=K_{p_i} and V'_i=V_{p_i}.

All three tensors keep shape [B, heads, N, d]. index_select(-2, indices) gathers only the token dimension; each head’s channels stay together.

Reordering all tokens alone leaves global attention unchanged after restoration. Defining new local predicates in that order changes which image tokens can interact.

Python implementation

python

def reorder_tokens(tensor, height, width):
    """Move row-major Q/K/V tokens into Hilbert sequence order."""
    order = hilbert_order(height, width)
    _validate_token_tensor(tensor, len(order))
    indices = torch.tensor(order, device=tensor.device, dtype=torch.long)
    return tensor.index_select(-2, indices)

3. Define windows and exclude shifted wraparound

Original paper

Enlarge figure

Math and tensor shapes

g(i)=\lfloor i/W_s^2\rfloor

A_{ij}=\mathbf{1}[g(i)=g(j)]

For unshifted HWA, W_s is the window side, distinct from image width W. Each window contains up to W_s^2 tokens; the final short window is retained. Power-of-two aligned windows can match square 2D windows, while other choices give irregular regions.

With shift s, the code compares window IDs of (i + s) % N and (j + s) % N. It also requires equal wrap groups (i + s) // N and (j + s) // N, so rotating the sequence cannot join its opposite ends.

Python implementation

python

def window_predicate(n, window_side, shift=0):
    """Build HWA in original Hilbert indices, guarding rotated wraparound."""
    area = window_side**2

    def allowed(i, j):
        a, b = (i + shift) % n, (j + shift) % n
        same_window = (a // area) == (b // area)
        same_wrap = ((i + shift) // n) == ((j + shift) // n)
        return same_window & same_wrap

    return allowed

4. Keep the neighborhood size at sequence boundaries

Original paper

Math and tensor shapes

The sequence radius is R=(K^2-1)/2.

a_i=\operatorname{clamp}(i-R,0,N-K^2)

A^{\mathrm{HNA}}_{ij}=\mathbf{1}[a_i\leq j<a_i+K^2]

With odd kernel side K and K^2\leq N, HNA shifts the interval inward to keep exactly K^2 keys. HSA instead uses abs(i - j) <= R and loses out-of-range neighbors at either endpoint.

These sequence neighborhoods generally differ from the original 2D receptive fields. Dense/Flex agreement checks the same new mask, not equivalence to the original SA or NA2D pattern.

Python implementation

python

def neighborhood_predicate(n, kernel_side):
    """Build HNA with a fixed area neighborhood shifted inward at the ends."""
    area = kernel_side**2
    half = (area - 1) // 2

    def allowed(i, j):
        start = i - half
        if isinstance(start, torch.Tensor):
            start = start.clamp(0, n - area)
        else:
            start = min(max(start, 0), n - area)
        return (j >= start) & (j < start + area)

    return allowed

5. Count empty, partial, and full blocks

Original paper

Enlarge figure

Math and tensor shapes

f_{\mathrm{empty}}=\frac{\#\text{empty blocks}}{\lceil N/b\rceil^2}

For each b-by-b tile, evaluate only valid indices inside N: no allowed pairs means empty, all valid pairs allowed means full, otherwise partial. The empty-block fraction counts blocks, not masked elements.

This diagnostic evaluates one clipped boolean tile at a time without retaining an N-by-N mask. It does not run a sparse attention kernel. Use b = 4 for Figure 3’s toy geometry; the native FlexAttention backend uses 128-token blocks.

At 64-by-64 tokens with window side 8 and b = 128, the example counts 992 of 1024 blocks empty for HWA (96.875%), versus 896 of 1024 for row-major WSA (87.5%). These are geometry counts, not measured speedups.

Python implementation

python

def block_statistics(predicate, n, block_size):
    """Classify valid cells only; empty_fraction counts blocks, not elements."""
    _positive_integer(n, "n")
    _positive_integer(block_size, "block_size")
    counts = {"empty": 0, "partial": 0, "full": 0}
    for q in range(0, n, block_size):
        rows = torch.arange(q, min(q + block_size, n))[:, None]
        for k in range(0, n, block_size):
            cols = torch.arange(k, min(k + block_size, n))[None, :]
            valid = predicate(rows, cols)
            counts[_block_kind(valid)] += 1
    total = math.ceil(n / block_size) ** 2
    return {
        "counts": counts,
        "total": total,
        "empty_fraction": counts["empty"] / total,
    }

6. Compile sparse attention and check the same mask

Original paper

Math and tensor shapes

S_{ij}=\frac{Q'_i(K'_j)^\top}{\sqrt d}

O'=\operatorname{softmax}(S_{\mathrm{masked}})V'

Disallowed scores become negative infinity before the row-wise softmax. The dense reference materializes N-by-N scores to check correctness; FlexAttention uses the predicate’s native BlockMask to skip empty GPU tiles.

The runner constructs the mask and compiled function once, then reuses them. CUDA support or compilation failures surface directly; the Flex path has no silent dense fallback. The CLI compares its output with dense attention under the identical mask.

Python implementation

python

def build_flex_attention(predicate, n, device):
    """Construct one native block mask and one reusable compiled function."""
    from torch.nn.attention.flex_attention import create_block_mask
    from torch.nn.attention.flex_attention import flex_attention

    def mask_mod(batch, head, i, j):
        return (i < n) & (j < n) & predicate(i, j)

    block_mask = create_block_mask(
        mask_mod,
        B=None,
        H=None,
        Q_LEN=n,
        KV_LEN=n,
        device=str(device),
        BLOCK_SIZE=128,
    )
    compiled = torch.compile(flex_attention)

    def attend(q, k, v):
        return compiled(q, k, v, block_mask=block_mask)

    return attend

7. Restore the spatial layout of the output

Original paper

Math and tensor shapes

O_{\mathrm{row},p_i}=O'_i

argsort(indices) builds the inverse permutation p^{-1}. Gathering with it returns [B, heads, N, d] in row-major token order and preserves gradient flow.

Restoration puts each attended vector back at its image position. It does not undo the values changed by attention. The path is cached, while these small index tensors are rebuilt during each forward call.

Python implementation

python

def restore_order(tensor, height, width):
    """Invert the permutation on the token axis, including its gradients."""
    order = hilbert_order(height, width)
    _validate_token_tensor(tensor, len(order))
    indices = torch.tensor(order, device=tensor.device, dtype=torch.long)
    inverse = torch.argsort(indices)
    return tensor.index_select(-2, inverse)

Run the complete attention core

Download hilbert_attention.py and requirements-hilbert.txt into the same directory. The requirements pin PyTorch 2.12.0. Start with the seeded CPU reference run; the optional Flex path needs a supported CUDA/Triton environment.

bash

python -m venv .venv
source .venv/bin/activate
python -m pip install -r requirements-hilbert.txt

# Default: 8x8 tokens, window side 4, dense CPU
python hilbert_attention.py

# Figure 3-sized diagnostic: 4x4 tokens, 4-token blocks
python hilbert_attention.py --height 4 --width 4 --size 2 --block-size 4

# Native 128-token CUDA Flex blocks and same-mask dense check
python hilbert_attention.py --device cuda --backend flex --height 64 --width 64 --size 8 --block-size 128

# Shifted HWA, HSA, and rectangular HNA on CPU
python hilbert_attention.py --size 4 --shift 8
python hilbert_attention.py --pattern slide --size 3
python hilbert_attention.py --height 3 --width 5 --pattern neighborhood --size 3

The CLI prints configuration, output shape, clipped block counts, and, for Flex, the numerical difference from the dense reference. It does not download a model or dataset. Use python hilbert_attention.py --help for all options.

bash

python hilbert_attention.py --device cuda --backend flex --height 64 --width 64 --size 8 --block-size 128 --benchmark --warmup 5 --iterations 20

Timing measures warmed, synchronized forward calls including reorder, attention, and restore. It excludes mask construction and initial compilation. Results depend on the machine and configuration; they are not the paper’s 4× or 18× measurements. With the core and its correctness check in view, the following results return to the authors’ kernel and full-model experiments.

Results

The kernel benchmarks use an RTX 3080 with CUDA 12.6 and PyTorch 2.7.0, batch size 16, 2 heads of dimension 64 and 128-token blocks; the appendix repeats them on an A100. Times below are forward passes.

Window attention (Table 1). Dense WSA computes each window separately. WSA (Flex) and HWA (Flex) run on FlexAttention.

The FlexAttention variants also use far less memory than dense WSA, which materializes each window's attention matrix: 66 MB against 520 MB in the forward pass at 128×128. HWA also runs on xFormers' block-diagonal dense kernel, which plain window attention cannot use.

Slide and neighborhood attention (Table 3). SA gathers each neighborhood explicitly. Because HSA is a plain band, it also runs on FlashAttention-2's built-in sliding-window kernel (FA2).

The headline 18× is HSA (Flex) against naive SA at 56×56 (5.85 ms to 0.32 ms). Against the optimized NATTEN kernel the gain is smaller but still there: HNA runs 1.8× faster than NA2D at 56×56 and 2.1× faster at 96×96.

Full layers. Counting the reshape (window partition or Hilbert reorder) and the QKV projection, HWA (Flex) takes 1.63 ms at 128×128 with 16×16 windows, against 4.22 ms for WSA and 6.39 ms for WSA (Flex). At 56×56 with 7×7 kernels, NA2D and HNA on NATTEN tie at 0.45 ms; HNA pulls ahead only at larger inputs and kernels.

Models (Table 4 and Figure 7). HWT-T and HNT-mini were trained on 8 V100 GPUs with the Swin-T and NAT-mini recipes. The V100 does not support FlexAttention well, so HWT trained with a dense kernel; the authors checked that the dense and FlexAttention versions give identical outputs. The speedups are therefore inference measurements.

The throughput runs use 7×7 windows or kernels at 224×224 and 16×16 at 512×512; the paper does not name the GPU for this figure. At 256×256, HWT-T reaches 81.5% against 81.6% for Swin-T. On CIFAR-10 and CIFAR-100, both models stay within 0.2 points of Swin and NAT.

Critical analysis

Strengths:

• A reordering instead of a kernel. The whole method is a token permutation plus a mask definition. It runs on FlexAttention, FlashAttention-2, xFormers and NATTEN without any custom CUDA.
• It fixes the right bottleneck. The paper identifies partial blocks, not attention FLOPs, as what makes local attention slow on block-sparse kernels, and the measured speedups track its sparsity numbers.
• Near-free accuracy. Changing the window shape from squares to Hilbert regions costs at most 0.2 points on ImageNet at the Swin-T and NAT-mini scale.

Limitations:

• It needs enough sparsity to pay off. At 56×56 with 7×7 windows, HWA (Flex) takes 0.31 ms against 0.22 ms for dense WSA, 0.7× the speed. The gain appears only at higher resolutions and with a good block-window pairing.
• Speedups depend on hardware and batch size. The main results are on a consumer RTX 3080. At batch size 1 on that GPU the gain nearly vanishes, and it stops growing beyond a batch of about 32. At high resolution with large windows, PyTorch's SDPA running dense window attention on FlashAttention can match HWA (Flex).
• Irregular windows complicate the model. Swin's per-window position bias has to become a global one, and shifted Hilbert windows need extra masks for tokens that wrap around the sequence.
• Narrow evaluation. Accuracy is reported only for image classification at small model sizes. Detection, segmentation and generation, where high resolution matters most, are left to future work.

Related Reading

• Swin Transformer: the window attention and shifted-window backbone that HWT reorders
• FlashAttention: the tiled attention kernel that FlexAttention and FlashAttention-2 build on
• Vision Transformer (ViT): global attention over image patches, the quadratic cost local attention avoids
• Making Deep Learning Go Brrrr: why overhead and memory traffic, not FLOPs, often decide GPU runtime
• VIOLIN: the same generalized Hilbert curve, used as a spatial prior inside global attention instead of for speed