点击下方卡片,关注「集智书童」公众号
本文主要解决了什么问题
-
- 当前基于注意力机制的特征融合方法计算复杂度高,实现困难,限制了其在资源受限环境中的应用。
-
- 现有的轻量化注意力机制(如Sparse或低秩近似)虽然降低了理论成本,但往往破坏了硬件友好的数据流。
-
- 在检测和分类任务中,如何以极小的计算开销提升模型性能,同时保持空间细节和上下文完整性。
本文的核心创新是什么
金字塔SparseTransformer(PST)模块 :提出了一种轻量级、即插即用的特征融合模块,结合从粗到细的token选择策略和共享注意力参数,显著降低计算复杂度,同时保留空间细节。
两阶段注意力机制 :
- • 粗粒度注意力 :通过高层特征图作为
对,关注低层特征图中的重要区域,将计算复杂度从
降低到
。
- • 细粒度注意力 :选择性处理由top-
粗略Q-K相似性识别的最信息量大的区域,复杂度为
。
参数共享与训练-推理灵活性 :仅通过粗分支进行训练,推理时可灵活激活细粒度注意力,无需重新训练,提供可适应的精度-效率权衡。
高效轻量化设计 :整个架构仅相当于一个
卷积的参数量,适合实际应用。
结果相较于以前的方法有哪些提升
目标检测任务(COCO数据集) :
- • 在YOLOv11-N/S/M中分别带来0.9%、0.5%和0.4%的mAP提升,且对延迟影响极小。
- • ResNet-18+PST-N实现了46.3%的mAP,超越R-50-AFPN(38.4%)和R-50-A2-FPN-Lite(39.8%)。
- • ResNet-101+PST-S达到50.9%的mAP,显著优于R-101-AFPN(40.2%)和R-101-A2 FPN(42.8%)。
图像分类任务(ImageNet数据集) :
- • 将ResNet-18/50/101的top-1准确率分别提升6.5%、6.1%和1.0%。
- • YOLOv11-cls变体的top-1精度提升约1%。
实时性与效率 :
- • YOLOv11-PST-N实现40.3%的mAP,优于Baseline模型(39.4%),延迟仅为1.24 ms。
- • PST增强的YOLOv11-S和YOLOv11-M分别达到47.4%和52.1%的mAP,且具有相当或更低的FLOPs。
局限性总结
Top-k选择的敏感性 :较大的top-
值会引入额外计算开销而未显著提升性能,需谨慎选择适中的值。 2. 2. 堆叠设计的冗余性 :增加PSA堆叠阶段会导致延迟显著增加,且未能进一步提升准确性。
自门控机制的影响 :尽管理论上可以增强特征融合能力,但实验表明禁用该机制可能获得更高的性能。
线性注意力的适用性 :尝试用线性注意力替代标准多头注意力机制并未带来预期效果,反而降低了性能。
导读
特征融合对于高性能视觉模型至关重要,但通常会导致过高的复杂度。然而,当前主流的基于注意力机制的融合方法往往涉及显著的计算复杂度和实现挑战,限制了它们在资源受限环境中的效率。
为解决这些问题,作者提出了金字塔SparseTransformer(PST),这是一种轻量级、即插即用的模块,它集成了由粗到细的token选择和共享注意力参数,以减少计算量同时保留空间细节。PST仅使用粗注意力机制即可进行训练,并在推理阶段无缝激活以进一步提升精度,无需重新训练。当将其添加到最先进的实时检测模型(如YOLOv11-N/S/M)中时,在MS COCO数据集上,PST带来了0.9%、0.5%和0.4%的mAP提升,且对延迟影响极小。
同样地,将PST嵌入ResNet-18/50/101作为 Backbone 网络,分别将ImageNet top-1准确率提升了6.5%、6.1%和1.0%。这些结果表明,PST作为一种简单且对硬件友好的增强方法,在检测和分类任务中均表现出良好的效果。作者的代码可在以下链接获取。
1 引言
特征融合,即跨层或跨分支组合特征,是现代视觉架构的基本组成部分。基于Transformer的[41]和基于CNN的[7]模型均表明,内容自适应注意力机制通过将图像特征处理为动态token序列[13, 28, 18],显著优于固定的卷积方案。在大语言模型/视觉模型时代,将图像分割为patch或潜在向量,使得自然语言处理技术能够直接应用于视觉任务[31, 23, 19],同时也促进了图像编辑[32, 17]和多模态理解的精确视觉语言对齐。因此,基于Transformer的主干网络和 Head 结构现已成为识别、检测和分割领域最先进模型的统治者[41, 13, 28]。
尽管现有的基于注意力机制的融合方案(例如FPN、
1 FPN、Deformable DETR)具有强大的表征能力,但它们仍然存在高FLOPs、不规则内存访问和工程复杂度等问题。Sparse或低秩近似虽然降低了理论成本,但往往破坏了硬件友好的数据流 [5, 49]。即便是复杂的库(FlashAttention、SageAttention)也必须融合 Kernel 并分块I/O来恢复性能,但它们仍然在密集的token网格上运行。
这促使作者提出一个问题:作者能否设计一个即插即用的Transformer模块,用于视觉特征融合,该模块满足以下要求 (i) 轻量级但具有最先进的性能,(ii) 大幅降低复杂度,(iii) 保持空间保真度,以及 (iv) 训练-推理灵活性?
为应对这些挑战,作者提出了金字塔SparseTransformer(PST),这是一种分层融合架构,能够在保持计算效率的同时实现高效的跨尺度特征集成。作者的关键技术贡献包括:
- • 作者提出了一种新的分层注意力机制,通过两阶段过程结合粗粒度的全局上下文和细粒度的局部细节:
- •
-
- 一层跨层粗注意力阶段,利用高层特征图作为
对,以关注更精细的 Query ,将计算复杂度从
降低到
- •
-
- 一种Sparse精细注意力阶段,选择性地处理由top-
粗略Q-K相似性识别的最信息量大的区域,其中每个选定区域对应于原始特征图中的
块,实现
复杂度。
-
• 作者开发了一种高效的参数共享方案,该方案仅通过粗分支进行训练,同时在推理过程中允许灵活激活细粒度注意力,从而提供可适应的精度-效率权衡。
-
• 作者设计整个架构以极轻量级为目标,其参数量仅相当于一个
卷积,使其在实际应用中具有极高的实用性。
大量实验表明,PST可作为有效的即插即用组件,无缝替换检测Head中的FPN,或集成到包括ResNet [15] 和YOLO [14, 20] 在内的各种 Backbone 架构中。PST在不同模型规模上始终提升性能,在COCO [26] 上至少获得0.4%的mAP提升,并在ImageNet [9] 上提高top-1准确率,同时引入的计算开销极小。
2 相关工作
2.1 基于视觉 Transformer 的多尺度特征融合
Transformer架构通过建模全局依赖和长距离交互,彻底改变了计算机视觉领域。视觉Transformer(ViT)[12]将图像视为一系列 Patch 的序列,在图像分类方面表现出色,但其固定的 Patch 尺寸限制了对于检测和分割任务至关重要的多尺度表示能力。为解决这一问题,Swin Transformer [28]和Pyramid Vision Transformer(PVT)[46, 47]等分层设计采用了多阶段架构和空间缩减机制。此外,T2T-ViT [50]和CSWin Transformer [10]等创新通过专门化的注意力模式增强了局部上下文建模和多尺度特征聚合能力。
与此同时,特征融合架构如FPN[25]对于整合不同深度的语义仍然至关重要。例如
-FPN [16]、AC-FPN [1]、G-FPN[38]和BiFPN [37]等变体通过引入注意力机制或增强连接性来改进原始设计,从而实现跨尺度特征的精细化。这些设计在效率和表示丰富性之间取得了平衡,成为许多检测流程的核心支撑。
在YOLO系列中,为了增强多尺度特征融合,已集成了目标注意力模块。YOLOv11[20]引入了跨阶段部分与空间注意力(C2PSA)模块,该模块在拼接特征图上应用空间注意力,以强调显著区域(尤其是小或被遮挡的物体),而不会显著增加计算成本。YOLOv12[40]利用区域注意力机制将特征图划分为多个区域,以保持较大的有效感受野,同时集成FlashAttention以优化内存访问并维持实时吞吐量。这些进展表明,经过精心设计的注意力可以在最小的速度损失下带来显著的精度提升,突显了对更轻量级、即插即用模块(如作者的PST)的需求。
2.2 视觉任务中的高效和动态注意力机制
标准自注意力机制的高计算成本促使了高效注意力变体的设计。Sparse注意力模型,如MSVi-Longformer,通过限制注意力模式来降低复杂度。此外,Sparse注意力模块使用Sparse注意力聚合上下文信息,以关注语义丰富的区域。Deformable DETR 从Sparse参考点采样以加速密集预测。线性注意力近似 [5, 35] 以及GPU优化的 Kernel 如FlashAttention 进一步提升了高分辨率输入的可扩展性。
动态token选择为效率提升提供了另一条途径。TokenLearner 和 DynamicViT 等方法基于学习到的token重要性进行剪枝,从而优化了计算性能的权衡。然而,许多此类策略需要额外的监督信号或复杂的架构变更。与此同时,在从粗到细的层次化注意力机制中,动态token选择策略通常构建在细粒度部分。QuadTree Attention 首先构建四叉树token金字塔,并在每一层基于粗粒度注意力分数选择top-
区域以指导细粒度注意力。基于这一思路,CF-ViT 采用两阶段推理:首先对少量粗粒度块进行快速预测,然后仅对最信息量大的块进行重新分区和重新注意力计算。
相比之下,PST模块采用top-k token选择策略,以无训练开销的方式动态识别信息区域。它具有粗粒度和细粒度注意力分支之间的参数共享特性,使得模型仅需激活粗粒度注意力即可进行训练。在推理过程中,可以激活细粒度注意力以实现无缝的精度提升,而无需修改已训练的参数。这种设计提供了一个独特的权衡:训练过程中的简洁性和推理过程中的精度提升。结合其模块化和高效性,PST为现有的基于注意力的特征融合机制提供了一种实用且强大的替代方案。
3 方法论
3.1 模块架构
在深度学习领域已得到充分验证的是,High-Level特征表示编码了Low-Level细节的紧凑摘要,反映了卷积网络和基于注意力网络的分层抽象过程[51, 2]。基于这一见解,可以利用High-Level激活来识别和选择最显著的Low-Level符号,从而提供有针对性的局部特征增强。为了保持上下文完整性和信息完整性,这些经过优化的局部特征随后通过全局特征进行补充,确保既有细粒度细节又具备整体场景理解能力。
基于这一共识,作者提出了金字塔SparseTransformer(PST)。图2展示了PST的架构,其与标准自注意力机制不同,其中
源自同一特征图[41, 13],而是采用交叉注意力机制:高层特征作为
,用于 Query
从低层特征图中提取,利用层次化图像语义指导细粒度融合[46]。通过允许语义丰富的高层token引导注意力聚焦于低层细节,PST增强了上下文对齐能力和判别力,同时减少了token间的交互。
作者进一步采用卷积位置编码(CPE)模块代替传统的可学习或正弦嵌入,该模块借鉴自CvT [48]。具体而言,作者在注意力操作后对每个特征图应用一个
的深度卷积[6, 40]。
粗粒度到细粒度特征选择:为了控制计算量,PSA采用粗粒度到细粒度的范式选择特征:粗粒度注意力在降采样特征上计算
,然后对每个 Query 的相似度分数进行平均,并在更精细的图中选择前
个键值对token。通过仅关注这
个细粒度token,PST将复杂度从
降低到
。作者在粗粒度和细粒度阶段共享所有注意力参数,因此训练可以省略细粒度分支,推理时可以启用它而无需重新训练。
对于不同的模型尺寸(N、S、M),通道维度
按照1:2:4的比例进行缩放(上限为2048),头数设置为
。除非另有说明,作者设置
并仅考虑相似度得分高于
的情况。
具体而言,忽略批次大小和头数,给定两个相邻的特征图
和
,作者计算:
粗略注意力输出是:
粗
该方法将token交互从
降低到
,因为
。然后,作者通过对注意力矩阵的每一行进行平均来为每个键推导出全局相似度分数:
选择使
最大的前-
个索引
(默认
,阈值
)。这些索引映射到更细的网格,从而在
中产生
个细粒度token,从中作者收集
然后,精细化的注意力为:
参数共享:最后,在两个阶段之间共享参数,作者融合输出:
通过结合跨层粗注意力机制与Sparse细粒度注意力机制,PST实现了总体复杂度为
,同时保留了其空间归纳偏置和简洁性。
即插即用卷积模块。整个架构仅依赖于标准卷积和注意力机制,通过在现代加速器上无缝集成高度优化的计算库,能够实现极致的计算效率。
整个PST模块的可学习组件由10个
卷积层组成,每个卷积层与BatchNorm配对。参数总数由输入特征通道数
、高层特征通道数
以及token嵌入维度
决定,具体如下:
显然,PST是一种轻量级设计,其参数规模与单个
卷积近似相当。
3.2 PST-DET结构
PST-DET结构如图3a所示,通过将PST模块作为插件集成来改进传统FPN。在该架构中,FPN内的原始卷积层被PST模块替换,这些模块处理来自金字塔层P3、P4和P5的拼接特征。这些模块通过上采样和下采样操作适应输入,确保与不同特征分辨率兼容。经过优化的特征随后被输入检测Head,实现跨多尺度的鲁棒目标检测。这种即插即用设计通过利用PST的交叉注意力能力,无缝增强了FPN框架。
3.3 PST-CLS结构
PST-CLS结构如图3b所示,采用PST融合P4层和P5层的High-Level特征,作为 Backbone 网络后置的插件。通过PST模块整合这些语义丰富的特征,该架构捕获了全面的上下文信息,并将其传递给分类头。这种直接的集成在不需对现有 Backbone 网络进行大量修改的情况下提高了分类精度,展示了PST作为特征融合增强器的多功能性和有效性。
4 实验
4.1 设置
为评估所提出的PST在特征融合方面的有效性,作者在目标检测和图像分类任务上进行了实验。实验配置如下所述。所有模型均在
NVIDIA RTX 4090 GPU上进行训练和测试。
检测任务:作者在MS COCO 2017数据集上验证了PST。检测实验分为两部分。第一部分使用ResNet-18、ResNet-50和ResNet-101作为 Backbone 网络,集成了PST-DET结构,并将其性能与其他具有相同 Backbone 网络的基于注意力机制的特征金字塔网络(FPN)方法进行了比较。第二部分专注于实时检测,采用YOLOv11变体(Nano、Small、Medium),将 Head 替换为PST架构,并与最先进的实时检测框架进行了基准测试。所有模型均使用SGD优化器进行训练,初始学习率为0.01,与YOLOv11保持一致。
分类任务:在分类任务中,作者使用ResNet-18、ResNet50和ResNet-101作为 Backbone 网络评估PST-CLS架构。作者还探索了YOLOv11-cls变体(Nano和Small)的轻量级配置,将P4和P5层特征融合结构替换为PST-CLS。使用SGD优化器,动量为0.9,学习率为0.1,批处理大小为256,在ImageNet数据集上进行200个epoch的训练。性能指标使用top-1和top-5准确率进行评估。
4.2 检测实验
对于基于ResNet的检测器,如表1所示,PST-DET在低FLOPs和高FLOPs尺度上始终优于基于注意力的FPN方法。使用ResNet-18+PST-N,实现了46.3%的mAP,尽管 Backbone 网络较轻,但仍超越了R-50-AFPN(38.4%)和R-50-A2-FPN-Lite(39.8%),AP75达到49.6%(对比41.9%和43.4%)。在更高的FLOPs下,ResNet-101+PST·S实现了50.9%的mAP,显著超过R-101-AFPN(40.2%)和R-101-A2 FPN(42.8%),AP50达到68.1%。这些结果表明PST具有更优越的多尺度特征融合能力,突显了其相对于仅使用更重 Backbone 网络的有效性。图1a提供了更直观的视觉对比。
为了实时检测,作者在COCO 2017验证集上对比了PST增强的YOLOv11模型与 Baseline 模型以及其他最先进的检测器,如表2所示。YOLOv11-PST
实现了40.3%的mAP,优于YOLOv11-N(39.4%),其使用FlashAttention的延迟为1.24 ms。YOLOv11-PST-S和YOLOv11-PST-M分别达到了47.4%和52.1%的mAP,超越了YOLOv11-S(46.9%)和YOLOv11-M(51.5%),且具有相当或更低的FLOPs。这些改进验证了PST在实时环境下提升准确性和效率的能力。图1b提供了更直观的视觉对比。
4.3 分类实验
作者还使用ImageNet验证了PST作为分类插件(PST-CLS)的效果,采用ResNet-18/50/101和轻量级YOLOv11-cls Backbone 网络(表3)。在所有架构中,集成PST均能在相应 Baseline 模型上带来稳定的精度提升——例如,ResNet-18精度提升约6%,ResNet-50和ResNet-101提升1-2%,YOLOv11-cls变体在top-1精度上提升约1%。这些结果表明,PST能够以极小的开销增强重载和轻量级 Backbone 网络,在无需专门 Backbone 网络设计的情况下,实现了精度与效率的良好权衡。
4.4 消融实验
为评估PST中关键组件的贡献,作者在COCO 2017数据集上开展消融实验,聚焦于使用YOLOv11-PST-N模型进行的检测任务。实验评估了top-k选择、PSA堆叠设计、自门控机制、线性注意力变体、参数共享以及注意力核类型的影响。结果汇总于表4,其中加粗条目表示性能最优的配置。
Top-
选择:一个有趣的现象是,将top-
替换为Gumbel-Softmax软top-
[33]并在训练期间启用精细注意力导致了显著的训练困难,并且精度下降超过5%。因此,表4 a展示了仅在推理阶段启用精细注意力的结果,它表明top-
值为8实现了最高的
为40.3%,延迟为1.24 ms。将top-
增加到16和32分别将
降低到39.5%和38.9%,同时显著增加延迟到1.50 ms和1.94 ms。这表明选择适中的top-
token数量(例如8)优化了准确性和效率,而较大的值则引入了计算开销而没有性能提升。
PSA堆叠设计:PSA堆叠结构的工作方式是,PSA模块的 Query
从前一个PSA层接收输入,而 Key
和 Value
在所有PSA模块中保持一致。最后,所有PSA模块的输出被连接起来以进行后续处理。表4b显示,单个PSA阶段获得了最佳的
为
,并具有最低的延迟
。堆叠额外的阶段(2、4、6)将
降低到
和
,同时显著增加了延迟至
、
和
。这表明单个PSA阶段足以实现有效的特征融合,而额外的阶段引入冗余和计算负担,却未提升准确性。
自门控机制:受NSA [49]的启发,自门控机制设计了一种使用门控张量融合粗粒度和细粒度注意力输出的方法。给定形状相同的粗粒度输出
和细粒度输出
,作者将它们连接起来,并应用一个一维卷积,然后是一个sigmoid层
。
因此,最终输出为:
表4c表明禁用自门控机制可以提高Nano(N)和Small(S)模型的
,分别达到40.3%和47.4%,而启用门控机制时分别为40.0%和46.2%。这表明自门控机制可能引入不必要的复杂性,略微降低性能。
线性注意力变体:为了探索更轻量级的模型,作者研究了用线性注意力[21]替换标准的多头注意力机制(MHA)(注意该操作使得精细的注意力计算变得不可行)。作者比较了原始线性注意力、将激活函数替换为ReLU的线性注意力以及将top-
设置为0的标准MHA。表4d显示,标准线性注意力变体实现了最佳的
为
,而Linear-ReLU变体降至
,Linear-EL
变体未能收敛(NaN)。这突出了标准注意力机制在PST中的稳定性和有效性。
参数共享:在PSA的设计中,粗粒度注意力和细粒度注意力共享同一组卷积来计算 Key和Value 。在此,作者将此与它们分别计算的情况进行比较。表
表明,粗粒度和细粒度注意力阶段之间的参数共享不会影响
(共享和不共享设置均为-40.3%)或延迟(1.24 ms)。然而,参数共享减少了模型的参数数量和训练复杂度,使其成为在不牺牲性能的情况下的一种有利的设
注意力核:表4f比较了不同的注意力核,SageAttention(SA)在Nano模型上实现了最低的延迟1.21毫秒,在Small模型上实现了2.33毫秒,优于Naive PyTorch(1.94毫秒和4.62毫秒·)和FlashAttention(FA)(1.24毫秒和2.50毫秒)。这突显了SageAttention在减少计算开销的同时保持精度的效率。
5 结论
作者介绍了金字塔SparseTransformer(PST),这是一种通过交叉注意力引导机制和动态token选择来增强多尺度特征融合的新型架构。PST利用从粗到细的注意力策略和参数共享,高效地整合多尺度特征,优化了计算效率和性能。
通过用交叉注意力替代传统自注意力机制,并结合
卷积,PST有效地捕获了跨特征层级的层次依赖关系。当与各种 Backbone 网络集成并应用于目标检测(PST-DET)和图像分类(PST-CLS)框架时,PST始终优于 Baseline 方法,在COCO和ImageNet数据集上实现了准确性的显著提升。这些结果验证了PST在提供鲁棒特征融合方面的能力,使其成为适用于广泛视觉任务的通用且高效的解决方案。
参考
[1]. Pyramid Sparse Transformer: Efficient Multi-Scale Feature Fusion with Dynamic Token Selection
扫码加入👉「集智书童」交流群
(备注:方向+学校/公司+昵称)