Pyramid Sparse Transformer: Enhancing Multi-Scale Feature Fusion with Dynamic Token Selection

TL;DR

The paper tackles the high computational cost of attention-based feature fusion in vision models by introducing the Pyramid Sparse Transformer (PST), a lightweight cross-layer fusion module that uses coarse-to-fine token selection with shared attention parameters. PST achieves this efficiency through a two-stage process: a coarse cross-layer attention that reduces complexity, followed by a sparse, fine-attention stage guided by top-$k$ coarse similarities, with all parameters shared between stages. The authors demonstrate PST as a plug-in for both detection (PST-DET) and classification (PST-CLS), yielding consistent accuracy gains on COCO and ImageNet with minimal latency impact and without retraining when enabling the fine-attention path at inference. The practical impact is a hardware-friendly, adaptable enhancement for real-time detection and large-scale classification that can augment backbones like ResNet and detectors like YOLOv11.

Abstract

Feature fusion is critical for high-performance vision models but often incurs prohibitive complexity. However, prevailing attention-based fusion methods often involve significant computational complexity and implementation challenges, limiting their efficiency in resource-constrained environments. To address these issues, we introduce the Pyramid Sparse Transformer (PST), a lightweight, plug-and-play module that integrates coarse-to-fine token selection and shared attention parameters to reduce computation while preserving spatial detail. PST can be trained using only coarse attention and seamlessly activated at inference for further accuracy gains without retraining. When added to state-of-the-art real-time detection models, such as YOLOv11-N/S/M, PST yields mAP improvements of 0.9%, 0.5%, and 0.4% on MS COCO with minimal latency impact. Likewise, embedding PST into ResNet-18/50/101 as backbones, boosts ImageNet top-1 accuracy by 6.5%, 1.7%, and 1.0%, respectively. These results demonstrate PST's effectiveness as a simple, hardware-friendly enhancement for both detection and classification tasks.

Figures (5)

• Figure 1: Comparison with others in FLOPs-accuracy (left) and Latency-accuracy (right) trade-offs.
• Figure 2: Overview of the Pyramid Sparse Transformer. Left: PST takes two adjacent feature maps as inputs and replaces standard attention with our Pyramid Sparse Attention (PSA) block. Drawing on the hardware friendliness and efficient design of EfficientFormer li2022efficientformer, all Linear and LayerNorm layers are substituted with 1$\times$1 Convs followed by BatchNorm. Additionally, the concluding feature concatenation leverages insights from well-established architectures such as CSPNet wang2020cspnet, ELAN wang2211designing, and C3K2 wang2024yolov9, renowned for their robust feature aggregation capabilities as demonstrated in prior empirical studies. Right: PSA first applies a cross-layer coarse attention, then selects the top-$k$ fine tokens to perform sparse refined attention and sums both outputs.
• Figure 3: Individual framework of downstream tasks for detection and classification with our PST.
• Figure 4: Feature heatmaps of YOLOv8-N, YOLOv11-N, and YOLOv11-PST-N Models on COCO val2017. Compared to SOTAs, our PST shows a purer and cleaner features. [Zoom in for a better view]
• Figure 5: Comparison of object detection results between YOLOv8-N, YOLOv11-N and YOLOv11-PST-N on COCO val2017. It can be observed that our PST based method can detect more correct instances with higher scores than other SOTAs. [Zoom in for a better view]