[论文解读] PyramidTNT: Improved Transformer-in-Transformer Baselines with Pyramid Architecture

Kai Han, Jianyuan Guo+2 more|arXiv (Cornell University)|Jan 4, 2022

Advanced Memory and Neural ComputingEngineering被引用 18

一句话总结

PyramidTNT 引入了分层金字塔结构和卷积主干，以改进 Transformer-in-Transformer（TNT）主干网络，在图像分类和目标检测任务中实现了最先进性能。该模型在仅使用 3.3B FLOPs 的情况下，于 ImageNet-1K 上实现了 82.0% 的 top-1 准确率，优于 Swin-T 和原始 TNT，同时保持更低的计算成本。

ABSTRACT

Transformer networks have achieved great progress for computer vision tasks. Transformer-in-Transformer (TNT) architecture utilizes inner transformer and outer transformer to extract both local and global representations. In this work, we present new TNT baselines by introducing two advanced designs: 1) pyramid architecture, and 2) convolutional stem. The new "PyramidTNT" significantly improves the original TNT by establishing hierarchical representations. PyramidTNT achieves better performances than the previous state-of-the-art vision transformers such as Swin Transformer. We hope this new baseline will be helpful to the further research and application of vision transformer. Code will be available at https://github.com/huawei-noah/CV-Backbones/tree/master/tnt_pytorch.

研究动机与目标

• 通过引入分层特征学习，提升 Transformer-in-Transformer（TNT）架构在视觉任务中的性能。
• 通过用卷积主干替代标准的分块主干，提升训练稳定性和表征学习能力。
• 建立一个新的、更强的视觉 Transformer 基线模型，其性能优于现有模型（如 Swin Transformer 和原始 TNT）。
• 通过逐级减少空间分辨率的金字塔结构，实现高效的多尺度特征提取。

提出的方法

• 引入包含四个阶段的金字塔结构，每个阶段将空间分辨率降低两倍，实现分层特征学习。
• 采用包含五个 3x3 卷积的卷积主干，以改善嵌入表示并提升训练稳定性。
• 在各阶段之间使用 Patch Merging 层进行下采样，保留特征层次结构并减少序列长度。
• 采用双分支 Transformer 设计：外层 Transformer 处理由全局表示，内层 Transformer 建模局部 patch 间关系。
• 引入可学习空间下采样注意力（LSRA），通过可学习的下采样比率压缩局部 token。
• 采用标准训练技术，包括权重衰减、AdamW 优化器以及多尺度数据增强，以提升模型鲁棒性。

实验结果

研究问题

• RQ1 金字塔结构能否提升 TNT 模型在视觉任务中的表征能力？
• RQ2 将标准分块主干替换为卷积主干，是否能提升训练稳定性和性能？
• RQ3 PyramidTNT 是否能在减少 FLOPs 的前提下，实现优于现有最先进视觉 Transformer 模型（如 Swin Transformer）的准确率？
• RQ4 PyramidTNT 中的分层特征学习对 COCO 上的目标检测性能有何影响？

主要发现

• PyramidTNT-S 在 ImageNet-1K 上实现了 82.0% 的 top-1 准确率，仅消耗 3.3B FLOPs，优于原始 TNT-S 和 Swin-T。
• 在 COCO 目标检测任务中，PyramidTNT-S 在 1x 训练计划下使用 RetinaNet 达到 42.0 mAP，较 Swin-T 提升 0.5 mAP。
• 在多尺度训练和 3x 训练计划下，PyramidTNT-S 在 Mask R-CNN 上达到 51.0 mAP，超过 Hire-MLP-S 0.9 mAP，且 FLOPs 更少。
• 该模型在所有主干尺寸（Ti、S、M、B）上均保持强劲性能，FLOPs 范围为 0.6B 至 16.0B。
• 金字塔结构与卷积主干的结合，显著提升了特征表示能力，尤其在检测任务中对大目标的表征效果更优。