提出SAGA机制,让线性注意力更高效且表达力更强。
SAGA: Selective Adaptive Gating for Efficient and Expressive Linear Attention
- 用自适应门控选择性聚合关键信息,避免均匀压缩导致的特征退化。
- 在1280×1280分辨率下,吞吐提升1.76倍,显存降低2.69倍。
- 适用于需要高分辨率视觉建模的高效模型设计,如大图分类任务。
尽管Transformer在建模长程依赖方面表现优异,广泛应用于视觉任务,但基于softmax的注意力机制具有二次复杂度,处理高分辨率图像时成为主要瓶颈。线性注意力通过将注意力计算从$QK)V$重构为$Q(KV)$,将复杂度从$/mathcal{O}(N^2)$降至$/mathcal{O}(N)$,同时保持全局感受野。然而,现有方法对历史键值(KV)信息进行统一压缩,易造成特征冗余并破坏查询(Q)与键值间的方向对齐,导致KV特征图低秩,性能低于softmax注意力。为此,本文提出选择性自适应门控机制(SAGA),引入输入自适应的可学习门控,有选择地调节到KV特征图的信息聚合,增强语义多样性,缓解传统线性注意力固有的低秩限制。此外,提出一种高效的哈达玛积分解方法用于门控计算,不增加额外内存开销。实验表明,SAGA在$1280 \times 1280$分辨率下相较PVT-T实现1.76倍吞吐提升和2.69倍峰值显存降低,且在ImageNet上最高提升4.4%的准确率,验证了其在计算效率与模型有效性上的双重优势。
原文摘要 · Abstract (English)
While Transformer architecture excel at modeling long-range dependencies contributing to its widespread adoption in vision tasks the quadratic complexity of softmax-based attention mechanisms imposes a major bottleneck, particularly when processing high-resolution images. Linear attention presents a promising alternative by reformulating the attention computation from $(QK)V$ to $Q(KV)$, thereby reducing the complexity from $\mathcal{O}(N^2)$ to $\mathcal{O}(N)$ while preserving the global receptive field. However, most existing methods compress historical key-value (KV) information uniformly, which can lead to feature redundancy and the loss of directional alignment with the query (Q). This uniform compression results in low-rank $KV$ feature maps, contributing to a performance gap compared to softmax attention. To mitigate this limitation, we propose \textbf{S}elective \textbf{A}daptive \textbf{GA}ting for Efficient and Expressive Linear Attention (SAGA) , which introduces input-adaptive learnable gates to selectively modulate information aggregation into the $KV$ feature map. These gates enhance semantic diversity and alleviate the low-rank constraint inherent in conventional linear attention. Additionally, we propose an efficient Hadamard-product decomposition method for gate computation, which introduces no additional memory overhead. Experiments demonstrate that SAGA achieves a 1.76$\times$ improvement in throughput and a 2.69$\times$ reduction in peak GPU memory compared to PVT-T at a resolution of $1280 \times 1280$. Moreover, it improves top-1 accuracy by up to 4.4\% on the ImageNet dataset, demonstrating both computational efficiency and model effectiveness.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。