用希尔伯特曲线优化扩散模型注意力，提升高分辨率图像生成速度
HilbertA: Hilbert Attention for Image Generation with Diffusion Models
- 通过希尔伯特曲线重排图像令牌，实现空间局部性与内存连续性的统一
- 在1024×1024下提速2.3倍，2048×2048下最高达4.17倍，图像质量不降
- 适合追求高分辨率图像生成效率的研究者与开发者
为扩散模型中的变压器设计稀疏注意力机制，需在二维空间局部性与GPU效率之间取得平衡，现有方法常因非合并内存访问导致效率下降。本文提出HilbertA，一种兼顾二维感知与GPU高效性的稀疏注意力机制。HilbertA沿希尔伯特曲线重排图像令牌，实现连续内存布局同时保留空间邻近性，并采用层间滑动调度策略，在无需重复或非合并内存访问的前提下实现长程信息传播。为进一步增强跨块通信与位置感知，HilbertA引入小型中心共享区域。该方法基于Triton实现，在Flux.1-dev上展现显著加速：生成1024×1024图像时速度提升2.3倍，生成2048×2048图像时最高达4.17倍，图像质量与基线相当或更优，验证了面向硬件的二维稀疏注意力在高分辨率图像生成中的可行性。
原文摘要 · Abstract (English)
Designing sparse attention for diffusion transformers requires reconciling two-dimensional spatial locality with GPU efficiency, a trade-off that current methods struggle to achieve. Existing approaches enforce two-dimensional spatial locality but often incur uncoalesced memory access. We present HilbertA, a 2D-aware and GPU-efficient sparse attention mechanism. HilbertA reorders image tokens along Hilbert curves to achieve a contiguous memory layout while preserving spatial neighborhoods, and employs a sliding schedule across layers to enable long-range information propagation without repeated or uncoalesced memory access. To further enhance cross-tile communication and positional awareness, HilbertA introduces a small central shared region. Implemented in Triton, HilbertA delivers comparable image quality with significant acceleration over prior methods on Flux.1-dev, demonstrating the feasibility of hardware-aligned two-dimensional sparse attention for high-resolution image generation. HilbertA delivers attention speedups of $2.3\times$ when generating $1024\times 1024$ images, and up to $4.17\times$ at $2048\times 2048$, while achieving image quality comparable to or surpassing baselines.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。