跳到主要内容
返回时间线
arXiv来源发表:

医学影像迁移不再靠布局先验:全局注意力挪到特征轴,同样预训练换来约三倍增益

导语

把视觉主干的全局注意力从空间轴换到特征轴,在23个医学分类、分割与低标注任务上,同样的ImageNet预训练带来最高约三倍的迁移增益,乳腺超声分割Dice从53.1升到77.7,参数量约四分之一。

Source-provided article image: Learning Features, Not Layouts: Feature-Axis Attention as a Transfer Prior for Medical Imaging
Figure 1 · arXiv

正文

医学影像主干里,全局注意力放在哪条轴上会改变预训练知识的可迁移程度。把全局混合从空间token改到通道token后,同一份ImageNet预训练在下游医学任务上带来的增益明显变大,而从头训练的性能几乎不变。 此前的迁移主干沿用标准视觉Transformer,全局建模由空间token自注意力主导,通道交互交给基本静态的前馈网络;通道感知的混合模型也只是在空间混合主干上外挂通道注意力,两条轴仍然耦合。 在参数量匹配的对照中,特征轴主干把预训练转成乳腺超声分割Dice增益+15.1,空间轴ViT-S只有+4.3,约三倍差距;把空间基线放大到ViT-B(90.6M)或换成预训练ResNet50,增益仍只有+5.7和+6.0。

做法是把空间提取和全局混合拆给两个互不重叠的算子。深度卷积只做逐通道的局部空间提取,不混合通道;全局交互则沿特征轴做单头自注意力,每个通道token是一整张响应图,因此注意力重加权的是整幅图像的通道响应,而不是空间位置。 通道感知的既有设计保留逐点投影,会在特征注意力之前先把通道线性混合,空间混合也贯穿整个网络,因此无法把预训练效果归因到某一条轴。 在23个医学数据集与任务上按统一协议评测,分割用同一个U-Net解码器;11个留出分割任务中特征轴在10个上取得最好Dice,均值84.3,对比ViT-UNet的79.8和DINOv2-UNet的77.4,参数量23.1M对90.6M。

接下来

报告覆盖的是23个医学数据集与任务上的分类、分割和低标注线性探测,分割统一挂同一个U-Net解码器,分类端到端微调,线性探测冻结主干只训线性头。作者提出的后续方向是更强的流程(nnU-Net)、其他架构家族和自监督预训练。

作者把中心核对齐读作迁移差距的佐证性相关,而不是特征复用的直接证明。低标注线性探测的结果依数据集而异,特征轴在PathMNIST、OCTMNIST、PneumoniaMNIST和BloodMNIST上领先,在少数数据集上DINOv2与ViT-S/16领先。从头训练时轴交换只差约1 Dice,落在种子方差内,说明该轴本身并不让网络更强。

来源