日期:2024.1.29 - 2024.2.4
本周工作:
1. 阅读论文
本周主要对这篇文献进行了阅读:
《用可学习的跳跃连接缩小U-Net中的语义差距:以医学图像分割为例》
背景
医学图像分割和随后对目标对象的定量评估为疾病诊断和治疗规划提供了有价值的信息。最近的语义分割方法通常依赖于类UNet的编码器-解码器架构,其中编码器产生高级语义特征,解码器逐渐对这些隐藏特征进行上采样,以产生具有每像素概率的分割图。大多数最先进的医学图像分割方法都采用编码器-解码器架构。然而,这种U形框架在通过简单的跳跃连接捕获非局部多尺度信息方面仍然存在局限性。为了解决这个问题,首先探讨了U-Net中跳跃连接在多个分割任务上的潜在弱点,得出了两条发现,然后基于他们的发现提出了一个新的分割框架,命名为UDTransNet,以解决U-Net中的三个语义缺口。
主要内容
-
跳跃连接的分析:研究发现,并非所有的跳跃连接对分割性能都有正面贡献。在不同的数据集上,某些跳跃连接可能对性能产生负面影响,这意味着需要更加定制化的跳跃连接实现方法。
-
UDTransNet框架:为了解决这些问题,作者们提出了一种新的框架,命名为UDTransNet。这个框架包含两个关键模块:双注意力变换器(Dual Attention Transformer, DAT)和解码器引导的校准注意力(Decoder-guided Recalibration Attention, DRA)模块。这些模块主要关注于捕获编码器特征的通道和空间关系,以及融合编解码器阶段语义不兼容的特征。
-
双注意力变换器(DAT):DAT旨在通过利用变换器模型全局上下文的能力,融合编码器多个阶段的特征。它包括多尺度特征嵌入、通道融合注意力(Channel-wise Fusion Attention, CFA)、空间选择注意力(Spatial-wise Selection Attention, SSA)和多层感知器(MLPs),目的是学习跨尺度的有意义的表示。
-
解码器引导的校准注意力(DRA):这个模块旨在有效地连接DAT令牌和解码器特征,消除不一致性,有效地解决编码器和解码器之间的语义差距。
-
全面的评估:论文提供了全面的实验结果,显示其提出的UDTransNet模型在各种公共数据集上与其他最先进的分割方法相比,获得了更高的评估分数和更精细的分割结果,同时参数相对较少。
方法
UDTransNet整体框架
多尺度特征嵌入( Multi-scale Feature Embedding)
为了使多个阶段的补丁能够集中在原始图像 X ∈ {\mathbf{X} \in} X∈ R H × W {\mathbb{R}^{H \times W}} RH×W中的相应区域,
-
首先通过将不同尺度的特征 E i ∈ R H 2 i − 1 × W 2 i − 1 × C i , ( i = 1 , 2 , 3 , 4 ) {\mathbf{E}_{i} \in \mathbb{R}^{\frac{H}{2^{i-1}} \times \frac{W}{2^{i-1}} \times C_{i}},(i=1,2,3,4}) Ei∈R2i−1H×2i−1W×Ci,(i=1,2,3,4)重塑为补丁大小为 P , P 2 , P 4 , P 8 {P, \frac{P}{2}, \frac{P}{4}, \frac{P}{8}} P,2P,4P,8P的标记序列来进行标记化,通道尺寸全部更改为C=128。
-
然后,我们将四个尺度的令牌序列 T i ∈ R C × d , d = H W P 2 T_i∈R^{C×d},d=\frac{HW}{P^2} Ti∈RC×d,d=P2HW
通过通道轴作为key和value T Σ = c o n c a t c ( [ T 1 , T 2 , T 3 , T 4 ] ) , T Σ ∈ R C σ × d T_{\varSigma}=concat_c\left( \left[ T_1,T_2,T_3,T_4 \right] \right) ,T_{\varSigma}∈R^{C_{\sigma}×d} TΣ=concatc([T1,T2,T3,T4]),TΣ∈RCσ×d
其中 C σ = 4 C C_{\sigma}=4C Cσ=4C
通道融合注意力(CFA)
不同的通道包含特定尺度的语义特征,因此自适应地融合这些特征有利于复杂的医学图像分割任务。为了融合来自不同阶段的多尺度信息,我们提出了CFA子模块,这是一种沿着特征图的通道轴的通道注意操作,它允许编码器学习通道之间的关系并捕获全局语义依赖关系。该过程如下图所示:
从数学上讲,我们在线性层之后获得一个查询、值和键,如公式1所示:
Q = T Σ W Q , K = T Σ W V , V = T Σ W V ( 1 ) Q=T_{\varSigma}W_Q,K=T_{\varSigma}W_V,V=T_{\varSigma}W_V\ \ \ \ \ \ \ \ \left( 1 \right) Q=TΣWQ,K=TΣWV,V=TΣWV (1)
其中,WQ,WK,WV∈RCσ×d是变换权重,d是序列长度(片数),Cσ是全尺度的信道维数(在我们的方法中,Cσ=4C)。当Q,K,V∈RCσ×d时,产生信道相似矩阵Mc∈RC∑×Cσ,并且值V由Mc
M c V T = σ [ ψ ( Q T K ) ] V T ( 2 ) M^cV^T=\sigma \left[ \psi \left( Q^TK \right) \right] V^T\ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \left( 2 \right) McVT=σ[ψ(QTK)]VT (2)
其中ψ(·)和σ(·)表示实例归一化和softmax函数。与原始的自注意不同,CFA沿着通道轴而不是patchaxis操作自注意,然后我们利用实例归一化来归一化相似矩阵,以平滑地传播梯度。
空间选择注意力(SSA)
CFA的输出被拆分为T Plot 1、T Plot 2、T Plot3、T Plot4,并沿着补片轴重新连接T Plot∑=凹点([Plot T1,Të2,T Plot 3,T vx 4]),在我们的方法中,将Të∑∈RC×dσ用作SSA的密钥和值,dσ=4d。在CFA之后,来自所有尺度的令牌序列都包含全局上下文信息和一致的语义信息(见图2)。我们进一步提出了SSA子模块来学习跨多个尺度的斑块之间的空间相关性。给定来自所有编码阶段的标记,我们将来自特定规模的标记视为查询,以从空间角度突出重要的空间区域。它允许我们的模型通过利用多尺度的空间相关性来捕捉每个尺度和全尺度信息之间的空间相关性。
更具体地说,SSA具有五个输入,包括作为查询的四个令牌序列Tëi和作为密钥和值的级联令牌序列Plot T∑:
Q
i
=
T
i
∧
W
Q
,
K
i
=
T
i
∧
W
K
,
V
i
=
T
i
∧
W
V
(
3
)
Q_i=\overset{\land}{T_i}W_Q,K_i=\overset{\land}{T_i}W_K,V_i=\overset{\land}{T_i}W_V\ \ \ \ \ \ \ \ \left( 3 \right)
Qi=Ti∧WQ,Ki=Ti∧WK,Vi=Ti∧WV (3)
其中WQi∈RC×d,WK,WV∈RC x dσ是不同输入的权重。三个特征表示被划分为NH个不同的子空间,例如,Qi=[Q1 i,Q2 i,…,QNH i]。在Qh i∈RC NH×d和Kh,Vh∈RC NH3×dσ的情况下,计算空间相似矩阵Msi,并通过交叉注意(CA)机制由Msi加权Vh中的每个标记:
Θ
i
h
=
M
i
s
V
h
=
σ
[
ψ
(
Q
i
h
K
h
T
)
]
V
h
∈
R
C
N
h
×
d
(
4
)
\varTheta _{i}^{h}=M_{i}^{s}V^h=\sigma \left[ \psi \left( Q_{i}^{h}K^{hT} \right) \right] V^h\ ∈\ R^{\frac{C}{N_h}×d}\ \left( 4 \right)
Θih=MisVh=σ[ψ(QihKhT)]Vh ∈ RNhC×d (4)
然后将子空间的输出级联以产生最终输出
O
i
S
S
A
=
c
o
n
c
a
t
(
[
Θ
i
1
,
Θ
i
2
,
.
.
.
,
Θ
i
N
H
]
)
,
∈
R
C
×
d
O_{i}^{SSA}\ =\ concat\left( \left[ \varTheta _{i}^{1},\varTheta _{i}^{2},...,\varTheta _{i}^{N_H} \right] \right) ,∈\ R^{C×d}
OiSSA = concat([Θi1,Θi2,...,ΘiNH]),∈ RC×d
第i个尺度的SSA过程如图6所示。
多层感知机(MLPs)
在SSA之后,开发了具有残差结构的MLP,以对信道依赖性进行编码,从而从SSA中细化特征。每个MLP的输出如下所示:
O i = O i S S A + M L P ( Q i + Q i S S A ) ( 5 ) O_i=O_{i}^{SSA}+MLP\left( Q_i+Q_{i}^{SSA} \right) \ \ \ \ \ \ \ \ \ \ \ \ \ \left( 5 \right) Oi=OiSSA+MLP(Qi+QiSSA) (5)
为了简单起见,方程中省略了层归一化。将等式(1)-(6)中的操作重复NL次,以构建NL层变换器。在我们的实现中,NH和NL都设置为4。第NL层的输出O1、O2、O3和O4然后被传送到下面的DRA模块中。
解码器引导的Transformer重新校准注意力模块(DRA)
解码器部分旨在根据从编码器级学习的特征来生成分割图。然而,在以往的研究中,解码器的重要性一直被忽视。如何融合通过跳过连接从编码器传输的特征对于增强解码器特征至关重要。如图6所示,2,DAT-Oi的输出和上采样特征Di之间存在语义间隙,这不可避免地阻碍了DAT和U-Net解码器的特征的直接级联。因此,为了缓解在融合这两个不兼容的特征集时的差异,我们提出了DRA模块来利用解码器特征来指导DAT获得的增强特征的信道信息过滤,并消除不一致性。因此,DRA模块可以被认为是自动融合不兼容特征的特征校准器。它可以很容易地应用于任何存在语义差距的场景。
在数学上,我们将第i级DAT输出Oi∈RC×d和第i级解码器特征图Di∈RC x H×W作为输入。然后,Di进一步转换为令牌序列TDi作为查询,Oi作为密钥和值:
Q i = T D i W Q i , K = O i W K , V = O i W V ( 6 ) Q_i=T_{D_i}W_{Q_i},K=O_iW_K,V=O_iW_V\ \ \ \ \ \ \ \ \ \ \left( 6 \right) Qi=TDiWQi,K=OiWK,V=OiWV (6)
其中WQi∈RC×d,WK,WV∈RC x dσ是不同输入的权重。因此,在Qi∈RC×d的指导下,通过交叉注意重新校准K,V∈RC?d:其中M~i是按信道的相似性矩阵。
M i V = σ [ ψ ( Q i T K ) ] V T ( 7 ) \overset{~}{M_i}V=\sigma \left[ \psi \left( Q_{i}^{T}K \right) \right] V^T\ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \left( 7 \right) Mi V=σ[ψ(QiTK)]VT (7)
最后,DRA的输出,Plot O1,O Plot 2,通过上采样操作和随后的卷积层来重构,并与解码器特征D1、D2、…连接,Di。DRA的这种设计可以通过适当地融合编码器特征来在解码过程中保留纹理细节。然后使用上采样层和两个卷积层将顶层特征D1转换为K类预测掩码P∈RH×W×K。
KaTeX parse error: Expected '}', got 'EOF' at end of input: …t) \right] V^T
总之,我们的UDTransNet以三种不同的方式使用基于矩阵乘法的注意力:
- CFA:查询、键和值来自编码器中四个尺度的级联特征。
- SSA:查询来自编码器中的每个刻度特征,键和值来自跨刻度的CFA输出。
- DRA:查询来自解码器特征,而密钥和值都来自SSA子模块获得的多尺度融合特征。
实验
我们选择了五个公共数据集来评估我们提出的模型,包括两个小规模数据集(GlaS和MoNuSeg)和三个大规模数据集(Synapse、ISIC2018和ACDC)。
数据集:Synapse
数据集GlaS、MoNuSeg、COVID-19
数据集ISIC2018
数据集ACDC
消融研究
所提出模块的消融
所提出模块顺序的消融
对不同注意力模式的消融
结论
医学图像的精确和自动分割是临床诊断和分析的关键步骤。然而,关于将注意力机制纳入跳跃连接的研究相对较少。在这项工作中,我们通过探索编码器-解码器级之间连接的信道和空间关注机制,提出了一种基于变换器的分割框架(UDTransNet),用于精确的医学图像分割。通过结合多尺度双注意力转换器(DAT)和解码器引导的重新校准注意力(DRA)的优势,该方法在多个医学图像分割任务上取得了最先进的结果,包括腹部器官分割、皮肤损伤分割和腺核分割。通过彻底的比较和深入的分析,我们验证了UDTransNet中每个组件的有效性,它成功地缩小了语义差距,并融合了具有不同注意力机制的多尺度特征。我们的结果为注意力机制在分割任务中增强跳跃连接的重要性提供了新的线索。我们还获得了一些关于通道和空间注意的可解释性的有趣发现,这是未来研究的重要研究课题。