提出新型位置偏置机制,让超分Transformer用上高效注意力,大幅提速且提升效果。
Rank-Factorized Implicit Neural Bias: Scaling Super-Resolution Transformer with FlashAttention
- 用低秩神经表示替代传统位置偏置,支持FlashAttention加速
- 在96×96窗口下训练,峰值信噪比达35.63dB,远超旧方法
- 适合追求高分辨率重建与推理速度的视觉算法研究者
当前超分辨率(SR)方法多采用Transformer以捕捉长程依赖,但主流模型依赖相对位置偏置(RPB),无法使用FlashAttention等硬件友好注意力核,导致训练和推理开销巨大,限制了模型规模扩展。本文提出秩分解隐式神经偏置(RIB),通过低秩隐式神经表示近似位置偏置,并以通道方式与像素内容拼接,将注意力分数中的逐元素加法转为点积运算,从而兼容FlashAttention。进一步引入卷积局部注意力与循环窗口策略,充分释放长程交互优势。在扩大至96×96注意力窗口的同时,联合提升训练图像块尺寸与数据集规模,最大化Transformer在超分任务中的潜力。实验显示,本方法在Urban100×2上达到35.63dB PSNR,训练与推理时间分别减少2.1倍和2.9倍,显著优于基于RPB的PFT模型。
原文摘要 · Abstract (English)
Recent Super-Resolution~(SR) methods mainly adopt Transformers for their strong long-range modeling capability and exceptional representational capacity. However, most SR Transformers rely heavily on relative positional bias~(RPB), which prevents them from leveraging hardware-efficient attention kernels such as FlashAttention. This limitation imposes a prohibitive computational burden during both training and inference, severely restricting attempts to scale SR Transformers by enlarging the training patch size or the self-attention window. Consequently, unlike other domains that actively exploit the inherent scalability of Transformers, SR Transformers remain heavily focused on effectively utilizing limited receptive fields. In this paper, we propose Rank-factorized Implicit Neural Bias~(RIB), an alternative to RPB that enables FlashAttention in SR Transformers. Specifically, RIB approximates positional bias using low-rank implicit neural representations and concatenates them with pixel content tokens in a channel-wise manner, turning the element-wise bias addition in attention score computation into a dot-product operation. Further, we introduce a convolutional local attention and a cyclic window strategy to fully leverage the advantages of long-range interactions enabled by RIB and FlashAttention. We enlarge the window size up to \textbf{96$\times$96} while jointly scaling the training patch size and the dataset size, maximizing the benefits of Transformers in the SR task. As a result, our network achieves \textbf{35.63\,dB PSNR} on Urban100$\times$2, while reducing training and inference time by \textbf{2.1$\times$} and \textbf{2.9$\times$}, respectively, compared to the RPB-based SR Transformer~(PFT).
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。