Hilbert-Guided Sparse Local Attention

AI summary

Overview

• Research area: Efficient computer vision — specifically sparse/local self-attention kernels for Vision Transformers.
• Technical level: Intermediate. Familiarity with Transformer attention, 2D window/sliding-window attention (Swin, NAT), and block-sparse attention kernels (FlexAttention) is helpful.
• Scope: A method that reorders image tokens along a Hilbert curve so that 2D local attention becomes contiguous in 1D, raising the empty-block ratio of block-sparse kernels and thereby accelerating window, slide, and neighborhood attention with minimal accuracy loss.

What This Paper Is About

Global self-attention scales quadratically with sequence length, so it is impractical for high-resolution images; local attention limits each token to a neighborhood, but its 2D patterns (square windows in row-major order) do not map cleanly onto 1D token sequences, leaving many "partial blocks" that still must be computed. The authors reorder image tokens along a Hilbert curve before building windows and neighborhoods, so tokens in a local window become contiguous in the 1D sequence and far more attention blocks are entirely empty. The goal is to convert 2D local attention into a highly block-sparse, kernel-friendly form that runs faster without changing the underlying model design or requiring hand-written kernels.

Key Contributions

1. Hilbert-guided window/neighborhood construction: A new way to form local attention regions by reordering tokens along a Hilbert curve and then creating windows and neighborhoods on the reordered 1D sequence, yielding three attention variants — Hilbert Window Attention (HWA), Hilbert Slide Attention (HSA), and Hilbert Neighborhood Attention (HNA), plus Hilbert Shifted Window Attention (HSWA).
2. Block-sparse efficiency gains: The patterns raise the empty-block ratio and reduce partial blocks, and are implemented through FlexAttention's programmable mask_mod and score_mod interfaces without hand-written kernels. Reported gains include roughly 4× for window attention and roughly 18× for slide attention.
3. Two end-to-end models: Hilbert Window Transformer (HWT) and Hilbert Neighborhood Transformer (HNT), instantiated from HWA and HNA respectively, achieving end-to-end speedups with minimal accuracy loss. HWT additionally shifts windows along the 1D sequence for cross-window interaction and replaces Swin-style window-relative position bias with a global relative position bias.
4. Systematic kernel-level evaluation: Comparisons across input sizes, window/kernel sizes, block sizes, and multiple backends (dense kernels, xFormers, FlashAttentionv2, NATTEN, FlexAttention), plus a released implementation at https://github.com/Yunge6666/Hilbert-Local-Attention.

Main Findings

• Sparsity drives speed: The paper's motivating analysis shows that for fixed sequence lengths of 1024, 2048, and 4096, the speedup of attention computation is positively correlated with block sparsity (the percentage of empty blocks), with the effect especially pronounced above 80% sparsity and larger for longer sequences.
• Hilbert reordering raises empty-block ratios: On a 16-token example with a 2×2 window and block size 4, HWA increases the empty-block ratio from 0% to 50% compared with row-major window attention. For slide attention in the same example, the empty-block ratio is unchanged but the partial-block ratio drops from 100% to 50%.
• Window attention speedups up to 4.0×: In Table 1 (batch size 16, head_num 2, head_dim 64, block size 128), HWA (Flex) reaches 0.12 ms (2.3×) at input 64×64 with an 8×8 window and 98.61% empty-block sparsity at input 96×96 with an 8×8 window versus 91.67% for WSA (Flex). The largest reported gain is 0.68 ms (4.0×) over dense WSA at input 128×128 with a 16×16 window, at 98.44% sparsity versus 87.50% for WSA (Flex).
• Block-sparse alone is not automatically faster: WSA (Flex) runs slower than dense WSA in every setting of Table 1 because it contains many partial blocks requiring element-wise masking. HWA (Flex) is faster than WSA (Flex) for both inference and training.
• Block and window sizes matter: In Table 2 (input 128, window 16), HWA (Flex) keeps 98.44% sparsity and a 4.0× speedup at block sizes 128 and 256, but drops to 1.43 ms (1.9×, 96.88% sparsity) at block size 512 and 2.78 ms (1.0×, 93.75%) at block size 1024, while WSA (Flex) stays at 5.68 ms and 87.50% for all block sizes.
• Slide/neighborhood attention gains: At input 56×56 with kernel 7×7 (Table 3), HNA (NAT) is 0.12 ms forward, reported as 48.8× faster than SA and 1.8× faster than NA2D. At input 96×96, SA runs out of memory, and HSA (Flex) and HNA (NAT) are about 1.8–2.1× faster than NA2D (NAT). HSA (Flex) shows 95.87% sparsity versus 87.89% for SA (Flex) at 96×96.
• HSA benefits from banded kernels: Because FlashAttentionv2 natively supports sliding-window (banded) patterns, HSA can use it directly (HSA (FA2) 0.25 ms forward at 56×56 versus 5.85 ms for unoptimized SA), whereas WSA and HWA cannot be accelerated by FlashAttentionv2 without manually modifying the kernel.
• Memory behavior: WSA (Flex) and HWA (Flex) use the same memory for the same input, dominated by QKV tensor size and weakly affected by block and window size; for input 96×96 their memory stays unchanged across window sizes, whereas dense WSA memory grows with input and window size because it materializes intermediate attention results.
• Input-processing overhead included: With window partitioning, QKV projection, and relative position bias included for a 128×128 input and 16×16 window (Figure 6a), HWA (Flex) takes 1.63 ms, which is 2.6× faster than WSA's 4.22 ms and 3.9× faster than WSA (Flex).
• Accuracy is largely preserved: On ImageNet-1K (Table 4), Swin-T scores 81.2 and HWT-T 81.0 at 224×224; NAT-mini scores 81.8 and HNT-mini 81.6. At 256×256, Swin-T scores 81.6 and HWT-T 81.5 (NAT-mini and HNT-mini results are not reported for 256×256). The reported drop is at most 0.2%. On CIFAR10/CIFAR100 with three-layer, depth [2,6,4], head [3,6,12] configurations trained for 300 epochs, SWIN scores 92.3%/76.6%, HWT 92.2%/76.4%, NAT 94.2%/79.9%, and HNT 94.1%/79.9%.
• A low-resolution limitation: In Appendix Table 6, at 56×56 with a 7×7 window, HWA (Flex) fails to outperform dense WSA (0.31 ms, labeled 0.7×, versus WSA's 0.22 ms) because sparsity is too low at that configuration to yield meaningful acceleration.

Methodology in Plain English

Attention over N tokens can be viewed as an N×N matrix where rows are queries and columns are keys. Block-sparse kernels cut this matrix into fixed-size tiles and label each tile full, partial, or empty; empty tiles are skipped entirely, partial tiles need expensive element-wise masking. The authors' key observation is that ordinary square windows in row-major order scatter a spatially local neighborhood across the 1D token order, producing many partial tiles. Their fix is to trace a Hilbert space-filling curve over the image and reorder the tokens along that path. Because the Hilbert curve preserves 2D locality, a window or neighborhood drawn on the reordered 1D sequence still corresponds to spatially adjacent image content, but now the attention interactions land on contiguous 1D blocks — so many tiles become wholly empty and can be skipped. The Hilbert path for a given feature-map size (H,W) is computed once and cached for reuse. The patterns are expressed through FlexAttention's mask_mod and score_mod hooks, so no custom CUDA kernels are needed. The researchers then benchmark HWA against WSA and HSA/HNA against SA/NA across input sizes, window and kernel sizes, and block sizes, using dense kernels, xFormers, FlashAttentionv2, NATTEN, and FlexAttention, with timings averaged over 10 runs of 100 iterations (25% warm-up, 75% measurement) on an RTX 3080 with CUDA 12.6 and PyTorch 2.7.0, plus an A100 reference in the appendix. Finally they train full models: HWT mirrors Swin Transformer's architecture and HNT mirrors NAT's, trained on 8 V100 GPUs with settings matching Swin-Tiny and NAT-mini.

Why This Matters

The paper shows that the practical cost of local attention is limited as much by kernel-level block structure as by theoretical complexity — a shift in framing that connects algorithm design with hardware-aware implementation. It offers a general recipe for making any 2D local attention pattern more block-sparse without redesigning the model or writing kernels, and it does so on top of an existing programmable framework.

Real-world applications implied by the work:

• High-resolution image recognition and classification, where local attention is standard and sequence length dominates cost (the paper evaluates on ImageNet-1K, CIFAR10, and CIFAR100).
• Object detection and semantic segmentation, since HWT and HNT retain standard hierarchical backbones commonly used in these pipelines; the paper notes these downstream tasks are left to future work.
• Image and video generation pipelines, which the authors also name as a setting where the same hierarchical backbones appear.
• Deployment on varied hardware and block-sparse backends, since the approach is described as extensible across backends and model backbones.

Industry relevance: The efficiency gains are large at high resolution (up to 4.0× for window attention, with HWA (xFormers) reaching 0.44 ms versus 12.17 ms for WSA (xFormers) at 96×96 with a 16×16 window), and reduced memory (for example, 37 MB versus 288 MB at 96×96 with a 16×16 window) directly translates into lower inference cost for vision workloads. The integration path is also low-friction: it works through FlexAttention's Python-level interfaces rather than requiring new hand-tuned kernels.

Future Directions

• Downstream task evaluation: The authors state that a comprehensive investigation into object detection, semantic segmentation, and image/video generation using the same Hilbert local attention modules will be conducted in future work.
• Automatic block-window configuration: Results show that speedup depends heavily on the combination of block size and window size (Table 2 shows a 4.0× gain collapsing to 1.0× as block size grows from 128/256 to 1024), so a principled way to choose these settings is an open problem.
• Extending to more backends and hardware: The paper positions "Hilbert-guided local attention + block-sparse kernel" as general and extensible across block-sparse backends and model backbones, raising the question of which other kernels and accelerator generations it can exploit — particularly given that V100 does not support FlexAttention well and HWT had to be trained with a dense kernel.
• Overcoming the low-sparsity regime: HWA (Flex) fails to beat dense WSA at 56×56 with a 7×7 window, indicating that beating a full-block kernel requires crossing a sparsity threshold; how to guarantee that threshold at small resolutions or small windows remains unresolved.

Target Audience

Researchers and engineers working on efficient vision backbones, attention kernel optimization, or block-sparse attention implementations. It is most useful to readers who already understand window/sliding-window attention and block-sparse kernels and want a concrete, low-effort way to make 2D local attention faster, and to practitioners deciding whether to adopt Hilbert-based token ordering in a Vision Transformer pipeline. Readers looking for a purely theoretical treatment of attention complexity, or for results on non-vision modalities, will find less here.

Authors’ abstract

The quadratic compute and memory costs of global self-attention severely limit its use in high-resolution images. Local attention reduces complexity by restricting attention to neighborhoods. Block-sparse kernels can further improve the efficiency of local attention, but conventional local attention patterns often fail to deliver significant speedups because tokens within a window are not contiguous in the 1D sequence. This work proposes a novel method for constructing windows and neighborhoods based on the Hilbert curve. Image tokens are first reordered along a Hilbert curve, and windows and neighborhoods are then formed on the reordered 1D sequence. From a block-sparse perspective, this strategy significantly increases block sparsity and can be combined with existing block-sparse kernels to improve the efficiency of 2D local attention. Experiments show that the proposed Hilbert Window Attention and Hilbert Slide Attention can accelerate window attention and slide attention by about $4\times$ and $18\times$, respectively. To assess practicality, the strategy is instantiated as the Hilbert Window Transformer and the Hilbert Neighborhood Transformer, both of which achieve end-to-end speedups with minimal accuracy loss. Overall, combining Hilbert-guided local attention with block-sparse kernels offers a general and practical approach to enhancing the efficiency of 2D local attention for images.