DAE-Former:用于医学图像分割的双注意引导高效transformer
-
论文标题:DAE-Former: Dual Attention-guided Efficient Transformer for Medical Image Segmentation
论文地址:https://arxiv.org/pdf/2212.13504.pdf
论文代码:https://github.com/mindflow-institue/DAEFormer
摘要
- 背景介绍:最近在计算机视觉领域获得了关注,因为它们能够对远程依赖关系进行建模。
- 现存问题: 作为Transformer模型的核心部分的自注意机制通常会受到tokens数量的二次方计算复杂性的影响。许多体系结构试图通过将自注意机制限制在局部区域或通过重新设计tokenization过程来降低模型复杂性。
- 解决方法: 作者提出了DAE-Former,这是一种新的方法,试图通过有效地设计自注意机制来提供另一种视角。更具体地说,重新制定了自注意机制,在保持计算效率的同时,捕获整个特征维度上的空间和通道关系。此外,重新设计了跳跃连接路径,加入了交叉注意模块,以确保特征的可重用性,增强定位能力。
- 实验结果: DAE-Former在多器官心脏和皮肤病变分割数据集上的性能优于最先进的方法,而不需要预训练权重。
算法
DAE-former(图1)是一种无卷积的U型分层纯Transformer。在给定具有空间维度H×W和C通道的输入图像xH×W×C的情况下,DAE-Former利用patch嵌入模块从输入图像中获得大小为4×4的重叠patch tokens。
然后,tokenized的输入(xn×d)通过编码器模块,具有3个堆叠的编码块,每个编码块由两个连续的双transformer层和一个patch合并层组成。在patch合并过程中,通过合并2×2的patch tokens来降低空间维度,同时使通道维度加倍。这允许网络以分层方式获得多尺度表示。
在解码器中,tokens在每个块中被再次扩展2倍。然后使用SCCA将每个patch扩展层的输出通过跳跃连接与并行编码层的特征进行融合。所产生的特征被馈送到两个连续的dual Transformer layers。最后,线性投影层产生输出分割图。在接下来的几节中,将首先对高效转置注意力进行简要概述。然后,介绍高效dual attention和SCCA模块
Efficient Attention
标准的自注意机制存在二次计算复杂度(O(N2)),这限制了该体系结构在高分辨率图像上的适用性。
沈等人利用自注意产生冗余的上下文矩阵的事实来提出一种计算自注意过程的有效方法(2):
其中ρq和ρk是查询和键的归一化函数。沈等人研究表明,当使用ρq和ρk这两个归一化函数时,该模块产生等同的点积注意输出。因此,高效注意力首先将键和查询归一化,然后将键和值相乘,最后将得到的全局上下文向量乘以查询以产生新的表示。
与点积注意力不同,Efficient Attention不会首先计算像素点之间的成对相似性。取而代之的是,keys被表示为dk注意图kjT,其中j指的是输入特征中的位置j。这些全局注意图代表了整个输入特征的语义方面,而不是与输入位置的相似性。这种顺序的改变大大降低了注意机制的计算复杂性,同时保持了较高的表征能力。
作者使用Efficient Attention来捕捉输入特征图的空间重要性。
Transpose Attention
交叉方差注意力,也称为转置注意(Ali等人,2021),是一种通道注意机制。该策略仅使用转置注意力来实现对较大输入大小的处理。然而,作者重新表述了这个问题,并提出了一种转置注意机制,以有效地捕获整个通道维度。转置注意力如等式(3)所示:
关键字和查询被转置,因此,注意力权重基于交叉协方差矩阵。CT指的是转置注意的语境向量。在计算注意力权重之前,引入温度参数τ来抵消应用于查询和关键字的L2范数的缩放。这增加了训练过程中的稳定性,但减少了一定的自由度,从而降低了模块的表征能力。
Efficient Dual Attention
一篇关于注意机制的文献综述(Guo等人,2022年)表明,与单一注意相比,空间和通道注意的结合提高了模型捕捉更多上下文特征的能力。因此,作者构建了一个结合了转置(通道)注意和高效(空间)注意的dual Transformer块。有效双重注意区块的结构如图2所示。有效dual Transformer块(8)包括有效注意(4),然后是加法和范数(5),以及执行通道注意的转置注意力区块(6),然后是加法和范数(7)。
Skip Connection Cross Attention
SCCA模块如图1(b)所示。不是简单地连接编码层和解码层的功能,而是交叉关注它们,以更有效地保留底层特征。提出的模块可以有效地为每个解码器提供空间信息,使其在生成输出掩码时能够恢复细粒度的细节。
跳过连接交叉注意(SCCA)应用了efficient attention,但是,与对键、查询和值使用相同的输入特性不同,用于查询的输入是由跳过连接X2的编码层的输出,因此得名。用于关键字和值的输入是较低解码器层X1的输出。为了融合这两个特征,需要使用线性层(10)将X1缩放到与X2相同的嵌入维度。使用X2作为查询输入背后的动机是对高效关注块内的多层表示进行建模。
实验
数据集和评估指标
使用公开可用的Synapse数据集进行评估,该数据集构成了一个包含30个病例的3779幅腹部临床CT图像的多器官分割数据集。遵循(Chen等人,2021b)中提出的评估设置。
使用ISIC 2018(Codella等人,2019)数据集进一步评估了在皮肤病变分割方面的挑战。遵循(Azad等人,2022d)进行实验设置。
Quantitative and Qualitative Results