U-Net结构改进：双交叉注意力模块（DCA）

前言

本文由迈阿密大学的研究团队于2023年3月30日发表，提出了双交叉注意力模块（Dual Cross-Attention , DCA），其目标是在轻微的参数和复杂性增加的情况下改进U-Net及其变体，能够简单而有效地增强u-net结构中的跳跃连接（skip-connection）。

DCA通过按顺序捕获多尺度编码器特征之间的通道和空间依赖关系来解决编码器特征和解码器特征之间的语义差距。

1. 首先，通道交叉注意（CCA）通过利用多尺度编码器特征的跨通道token的交叉注意提取全局通道依赖关系。
2. 然后，空间交叉注意（SCA）模块进行交叉注意操作，来捕获跨空间令牌的空间依赖性。
3. 最后，将这些细粒度的编码器特征上采样并连接到相应的解码器部分，形成skip-connection方案。

传统的U-Net改进方法，例如残差和循环连接等存在以下缺陷：

1. 卷积的局部性无法捕获不同特征之间的长距离依赖关系。
2. skip-connection在简单地连接编码器和解码器特征时引起的语义差距。

受到顺序双重注意力和通道交叉注意力的启发，提出了双交叉注意力模块（DCA），有效提取多尺度编码器特征之间的通道和空间依赖，以解决语义差距问题。

1. 模型的特点

模型大致示意如下。DCA模块的结构不受编码器stage数量的影响，给定n+1个多尺度编码器stage，DCA将前n个stage的特征层作为输入，产生增强表示，并将它们连接到相应的n个解码器stage。

如下图所示，DCA可以分为两个主要阶段，三个步骤：

1. 第一阶段由多尺度patch embedding模块组成，以获得编码器Token。
2. 第二阶段，在这些编码器token上使用通道交叉注意（CCA）和空间交叉注意（SCA）模块来实现DCA，以捕获长距离依赖关系。
3. 最后，使用层归一化和GeLU对这些token进行序列化和上采样，将它们连接到解码器对应部分。

2. 基于多尺度编码器的Patch Embedding

首先从n个多尺度编码器stage中提取Patch。

给定n个不同尺度的编码器stage，$E_i\in {\mathbb{R}}^{C_i\times \frac{H}{2^{i-1}}\times \frac{W}{2^{i-1}}}$，并且块大小$P_i^S=\frac{P^S}{2^{i-1}}$，其中$i=1,2,\dots ,n$。使用大小和步长为$P_i^S$的平均池化来提取patch，并在展平的2维patch上使用$1\times 1$深度可分离卷积来进行映射。

$$T_i=DConv1D_{E_i}(Reshape(AvgPool2D_{E_i}(E_i)))$$

其中，$T_i\in {\mathbb{R}}^{p\times C_i},(i=1,2,\dots ,n)$表示第i个编码器stage展平后的patch。注意，$P$代表patch的数量，对于每个$T_i$都是相同的，所以可以利用这些token之间的交叉注意。

3. 通道交叉注意力（CCA）

如下图所示，使用CCA对每个token $T_i$进行处理。

首先对每个$T_i$进行层归一化（LN），然后沿通道维度对$T_i,(i=1,2,\dots ,n)$进行拼接，得到$T_c$，来产生Key和Value，同时使用$T_i$作为Query。将深度可分离卷积引用到自注意力中，以便捕获局部信息并降低计算复杂性。

$$Q_i=DConv1D_{Q_i}(T_i)andK=DConv1D_K(T_c)andV=DConv1D_V(T_c)$$

其中，$Q_i\in {\mathbb{R}}^{P\times C_i}$，$K\in {\mathbb{R}}^{P\times C_c}$，$V\in {\mathbb{R}}^{P\times C_c}$，分别为映射的queries，keys，values。从而CCA表示如下：

$$CCA(Q_i,K,V)=Softmax(\frac{Q_i^{T}K}{\sqrt{C_c}})V^T$$

其中$\frac{1}{\sqrt{C_c}}$是比例因子。交叉注意的输出是values的加权和，权重由queries和key之间的相似性决定。最后使用深度可分离卷积对交叉注意的输出进行处理，并将其输入SCA模块。

4. 空间交叉注意力（SCA）

SCA模块如下图所示。

给定CCA模块处理后的输出$\overline{T_i}\in {\mathbb{R}}^{P\times C_i},(i=1,2,\dots ,n)$，沿通道维度进行层归一化和拼接。与CCA模块不同，利用拼接后的token $\overline{T_c}$作为queries和key，而将每个$\overline{T_i}$作为value。对queries，keys，values上使用$1\times 1$深度可分离卷积进行投影。

$$Q_i=DConv1D_Q(\overline{T_c})andK=DConv1D_K(\overline{T_c})andV=DConv1D_{V_i}(\overline{T_i})$$

然后SCA可表示为：

$$SCA(Q,K,V_i)=Softmax(\frac{Q{K}^T}{\sqrt{d_k}})V_i$$

其中，$\frac{1}{\sqrt{d_k}}$是比例因子。对于多头的情况，$d_k=\frac{C_c}{h_c}$，其中$h_c$是head的数目。然后使用深度可分离卷积对SCA的输出进行处理得到最终DCA的输出。

然后对DCA的输出进行层归一化和GeLU处理。最后，DCA的n个输出通过上采样连接到它们对应的解码器部分。

注意交叉注意和自注意之间的主要区别在于，交叉注意通过将多尺度编码器特征融合在一起，而不是单独利用每个stage来创建注意力图，这也允许交叉注意捕获编码器不同stage之间的长距离依赖关系。