TLDR: HilbertA is a new sparse attention mechanism for diffusion models that generates high-resolution images faster and with comparable quality. It reorders image tokens using Hilbert curves to maintain spatial locality and ensure GPU-efficient memory access, then uses a tiling and sliding strategy for local attention and cross-tile information flow. This approach significantly speeds up image generation, especially at higher resolutions, by optimizing memory access patterns.
Generating high-quality images, especially at large resolutions, has become a hallmark of modern artificial intelligence. Diffusion models, like those using Transformer backbones, are at the forefront of this capability. However, a significant challenge arises with these models: the computational cost of self-attention, which grows quadratically with the number of image tokens. This means that as image resolution increases, the time it takes to generate an image can become prohibitively long.
To tackle this, researchers often turn to “sparse attention” mechanisms. These methods reduce the computational burden by limiting how many tokens interact with each other, rather than having every token interact with every other token. While sparse attention has been widely explored for text-based models, its application to two-dimensional images presents a unique dilemma. Image data has a crucial property called “2D locality,” meaning that nearby pixels or regions in an image are highly related. An effective sparse attention mechanism for images must preserve this 2D locality while also being efficient for Graphics Processing Units (GPUs), which prefer data arranged in a specific, contiguous way in memory.
Many existing sparse attention methods for images struggle with this balance. They might preserve 2D locality but lead to inefficient memory access on GPUs, or they might be GPU-efficient but break the natural spatial relationships in the image. This trade-off often undermines the expected speedups.
Introducing HilbertA: A New Approach to Sparse Attention
A team from New York University, including Shaoyi Zheng, Wenbo Lu, Yuxuan Xia, Haoming Liu, and Shengjie Wang, has introduced a novel solution called HilbertA (Hilbert Attention). This mechanism is designed to be both 2D-aware and GPU-efficient, overcoming the fundamental dilemma faced by previous methods. HilbertA achieves this through three key components:
First, Reordering: HilbertA reorders image tokens along a special path known as a Hilbert curve. Imagine taking a complex, winding path that touches every part of a square grid without crossing itself. This path ensures that tokens that are spatially close in the original image remain close to each other in the reordered, one-dimensional sequence. This contiguous memory layout is highly efficient for GPUs.
Second, Tiling: After reordering, the sequence of tokens is divided into smaller, non-overlapping “tiles.” Attention is then restricted to within these tiles. Because the Hilbert curve preserves spatial locality, these tiles still represent coherent neighborhoods in the original image. This tiling strategy significantly reduces computation and memory usage, as attention is only performed on smaller, local blocks.
Third, Sliding: To ensure that information can still flow across these local tiles and enable long-range communication, HilbertA employs a clever “sliding schedule” across different layers of the diffusion model. Instead of costly reallocations or complex cross-tile attention patterns, the attention window effectively slides along the Hilbert-ordered sequence by a fixed offset in each layer. This allows tokens to gather information from different adjacent tiles over time, expanding their their “receptive field” without sacrificing memory efficiency.
Additionally, HilbertA incorporates a small, fixed “central shared region.” This region acts as a global relay, allowing all tiles to exchange information and providing a consistent positional reference point for the model, which is especially useful when using Rotary Positional Embeddings (RoPE).
Performance and Impact
The researchers implemented HilbertA using Triton, a high-performance programming language for GPUs, and evaluated it on the Flux.1-dev model for image generation. The results are impressive. HilbertA delivered significant speedups, achieving up to 2.3 times faster attention when generating 1024×1024 images, and an even more remarkable 4.17 times faster at 2048×2048 resolution. These speedups translate to faster end-to-end image generation times, up to 1.10 times faster at 1024×1024 and 1.51 times faster at 2048×2048.
Crucially, HilbertA achieves these gains while maintaining image quality comparable to or even surpassing existing baselines. This demonstrates that memory efficiency, driven by hardware-aligned layouts, is as important as raw computational reduction (sparsity) for real-world performance.
The overhead introduced by HilbertA, primarily for reordering tokens, is minimal and occurs only twice per inference process, making it negligible over many denoising steps. While minor visual artifacts at tile boundaries were occasionally observed, the overall quality and efficiency trade-off position HilbertA favorably for practical high-resolution image generation.
Also Read:
- Novel Training Approaches for Diffusion Models Significantly Enhance Generative AI Efficiency
- ViTSP: A Hybrid AI Framework for Large-Scale Traveling Salesman Problems
Future Directions
The principles behind HilbertA could extend beyond 2D images. The researchers suggest that a three-dimensional Hilbert curve could be used for video generation, capturing both spatial and temporal dependencies. This would open new avenues for efficient video diffusion models.
HilbertA represents a significant step forward in designing sparse attention mechanisms for image generation. By intelligently reordering tokens along Hilbert curves, it reconciles the conflicting demands of 2D spatial locality and GPU memory efficiency, paving the way for faster and more accessible high-resolution image synthesis. You can read the full research paper here.
