越轻越好:通过自适应剪枝重新审视医学图像分割中的Transformer
计算机视觉与模式识别
2022-07-13 v2 人工智能
摘要
视觉 Transformer 因其在各类计算机视觉任务上的卓越性能,近来在医学图像分析领域掀起新浪潮。然而,近期基于混合/Transformer 的方法主要关注 Transformer 在捕获长程依赖上的优势,却忽视了其令人望而生畏的计算复杂度、高昂训练成本与冗余依赖等问题。本文中,我们提出将自适应剪枝应用于医学图像分割的 Transformer,并给出轻量且高效的混合网络 APFormer。据我们所知,这是首个针对医学图像分析任务的 Transformer 剪枝工作。APFormer 的关键特性主要为:自监督自注意力(SSA)以改善依赖建立的收敛性,高斯先验相对位置嵌入(GRPE)以促进位置信息学习,以及自适应剪枝以消除冗余计算与感知信息。具体而言,SSA 与 GRPE 分别将良好收敛的依赖分布与高斯热图分布作为自注意力与位置嵌入的先验知识,以简化 Transformer 训练并为后续剪枝操作奠定基础。随后,通过调节门控参数,执行查询级与依赖级的自适应 Transformer 剪枝,从而实现复杂度降低与性能提升。在两个广泛使用数据集上的大量实验表明,APFormer 以远少参数与更低 GFLOPs 取得了优于最先进方法的突出分割性能。更重要的是,我们通过消融实验证明,自适应剪枝可作为即插即用模块提升其他基于混合/Transformer 方法的性能。代码见 https://github.com/xianlin7/APFormer。
引用
@article{arxiv.2206.14413,
title = {The Lighter The Better: Rethinking Transformers in Medical Image Segmentation Through Adaptive Pruning},
author = {Xian Lin and Li Yu and Kwang-Ting Cheng and Zengqiang Yan},
journal= {arXiv preprint arXiv:2206.14413},
year = {2022}
}