TLDR: A new research paper introduces a novel method to optimize local attention in image processing by reordering image tokens along a Hilbert curve. This approach significantly increases ’empty blocks’ in block-sparse attention kernels, leading to substantial speedups (up to 4x for window attention and 18x for slide attention) and reduced memory usage with minimal accuracy loss. The Hilbert Window Transformer (HWT) and Hilbert Neighborhood Transformer (HNT) demonstrate the practical benefits of this Hilbert-guided strategy for efficient 2D local attention in high-resolution images.
The world of artificial intelligence, particularly in computer vision, has seen remarkable advancements thanks to models based on self-attention mechanisms, like the Transformer architecture. However, a significant hurdle remains: the computational and memory demands of global self-attention grow quadratically with the size of the input, severely limiting its application in high-resolution images.
To tackle this, researchers developed local attention, which restricts each token’s focus to its immediate neighborhood. While this significantly reduces complexity, conventional local attention patterns, often based on a simple row-by-row ordering of image pixels, don’t fully leverage the efficiency benefits of modern hardware. Specifically, when using ‘block-sparse’ kernels – a technique that skips computations for irrelevant parts of the attention matrix – traditional local attention often results in many ‘partial blocks’ rather than truly ’empty blocks,’ hindering speedups.
A new research paper, titled “HILBERT-GUIDED BLOCK-SPARSE LOCAL ATTENTION” by Yunge Li and Lanyu Xu from Oakland University, introduces an innovative solution to this problem. Their work proposes a novel method for constructing attention windows and neighborhoods by reordering image tokens along a Hilbert curve. The Hilbert curve is a continuous fractal space-filling curve known for its strong locality-preserving property, meaning that points close to each other in 2D space remain close to each other when mapped to a 1D sequence.
The Hilbert Curve Advantage
By reordering image tokens according to the Hilbert curve, the researchers found that attention windows and neighborhoods become much more contiguous in the 1D sequence. From a block-sparse perspective, this reordering dramatically increases the ratio of ’empty blocks’ – areas where no computation is needed – and reduces the number of ‘partial blocks,’ which still require some processing overhead. This increased sparsity directly translates to significant reductions in computational and memory costs.
The paper demonstrates that this Hilbert-guided strategy can be seamlessly integrated with existing block-sparse kernels, such as those provided by the FlexAttention framework. FlexAttention allows developers to define custom sparse attention patterns, and the Hilbert reordering makes these patterns much more efficient for 2D image data.
Also Read:
- Random Batch Attention: A New Approach to Efficient Graph Transformers
- Accelerating Image AI Training with CoMA: A New Approach to Masked Autoencoders and Dynamic Vision Transformers
Practical Applications and Performance
To assess the real-world practicality of their approach, the researchers instantiated their strategy into two end-to-end models: the Hilbert Window Transformer (HWT) and the Hilbert Neighborhood Transformer (HNT). These models adapt existing Transformer architectures (Swin Transformer and Neighborhood Attention Transformer, respectively) by incorporating the Hilbert-guided attention mechanisms.
Experiments showed impressive results. The proposed Hilbert Window Attention (HW A) and Hilbert Slide Attention (HSA) accelerated window attention and slide attention by approximately 4 times and 18 times, respectively, compared to their conventional counterparts. These speedups were observed with minimal loss in accuracy on the ImageNet-1K dataset, confirming the effectiveness and feasibility of the approach. Furthermore, the Hilbert-guided methods significantly reduced memory consumption during both training and inference.
The authors emphasize that the combination of Hilbert-guided local attention with block-sparse kernels offers a general and practical method for enhancing the efficiency of 2D local attention for images. This approach is flexible, allowing it to interface with various block-sparse backends and model architectures, paving the way for more efficient deep learning models in computer vision. You can find the full research paper here: HILBERT-GUIDED BLOCK-SPARSE LOCAL ATTENTION.
