在本节中,作者首先描述了网络的总体架构。随后,作者介绍了窗口 Transformer (WT)模块,它包括窗口化多 Head 注意力(W-MSA)和局部信息感知增强(LiPe)。这些元素构成了编码器和解码器中的基础组件。最后,作者讨论了嵌套密集块,其特点在于嵌套密集的跳跃路径,这些路径促进了编码器特征图与解码器特征图之间的跳跃连接。
WiTUnet被提出为一种U形的网络架构,如图1 (a)所示。该网络由编码器、瓶颈、解码器和中间嵌套的密集块组成。特别是,在处理损坏的图像时,即表示为
y
∈
R
1
×
H
×
W
y \in \mathbb{R}^{1\times H \times W}
y∈R1×H×W的低剂量计算机断层扫描(LDCT)图像,WiTUnet首先使用一个输入嵌入层,一个步长和填充设置为1的
3
×
3
3\times 3
3×3卷积层,将原始数据转换为特征图
a
∈
R
C
×
H
×
W
a \in \mathbb{R}^{C \times H \times W}
a∈RC×H×W。遵循U形架构[14],这些特征图通过
D
D
D个编码器块进行处理,每个编码器块包含几个窗口Transformer (WT)块,其中包含一个窗口化多Head自注意力(W-MSA)块。W-MSA块通过其自注意力机制捕捉非局部信息,并通过非重叠窗口减少计算复杂性,结合局部信息感知增强(LiPe))以精确捕捉特征图中的局部细节和全局信息。编码器每一级的输出都通过一个步长为2的
4
×
4
4\times 4
4×4卷积层进行下采样,作为下一编码器层的输入,其中特征图的通道数翻倍,其尺寸减半。经过
k
k
k层编码后,特征映射表示为
a
k
0
∈
R
2
C
×
H
′
×
W
′
a_k^0 \in \mathbb{R}^{2C \times H' \times W'}
ak0∈R2C×H′×W′,其中
k
∈
[
0
,
D
)
k \in [0, D)
k∈[0,D)。
在最后一个编码器层和解码器之间,引入了一个由WT块组成的瓶颈层,输出特征图为
Z
D
0
∈
R
C
×
H
′
×
W
′
Z_{D}^{0} \in \mathbb{R}^{C \times H' \times W'}
ZD0∈RC×H′×W′。这个瓶颈层被设计为更有效地捕捉全局信息。例如,当编码器深度足够时,最深特征图的尺寸与窗口大小相匹配,使得W-MSA能够高效地捕捉全局信息,而LiPe模块则保持对局部细节的敏感性。因此,网络可以在不需要循环移位[9]的情况下感知全局信息。
在解码器部分,特征图
z
p
0
z_{p}^{0}
zp0经历D个Level的解码处理。解码器的结构与编码器相似,采用堆叠的WT块。特征图通过使用
2
×
2
2 \times 2
2×2转置卷积核进行上采样,该卷积核的步长为2,这会使得通道数减半,尺寸加倍。与传统的U形网络设计[14][22]一致,解码器接收来自相应编码器Level的跳跃路径连接,并将它们与来自前一个解码层的上采样输入相结合。在引入嵌套密集块之后,输入特征图的通道数增加,因此,在与编码器输出结合后,解码器的第k级输入为
a
k
,
v
,
i
n
0
∈
R
(
D
−
k
+
1
)
×
C
×
H
′
×
W
′
a_{k,v,in}^{0} \in \mathbb{R}^{(D-k+1) \times C \times H' \times W'}
ak,v,in0∈R(D−k+1)×C×H′×W′。为了与传统的U形网络解码器输出保持一致,解码器中的WT块调整输出特征图的通道数,得到
a
k
,
v
,
o
u
t
0
∈
R
2
C
×
H
′
×
W
′
a_{k,v,out}^{0} \in \mathbb{R}^{2C \times H' \times W'}
ak,v,out0∈R2C×H′×W′,其中
k
∈
[
0
,
D
)
k \in [0,D)
k∈[0,D)且
0
=
D
−
k
0 = D - k
0=D−k。这种命名约定旨在简化在后续讨论中引入嵌套密集块。经过D级解码处理后,最终的输出是
a
0
,
4
,
o
u
t
0
∈
R
C
×
H
×
W
a_{0,4,out}^{0} \in \mathbb{R}^{C \times H \times W}
a0,4,out0∈RC×H×W,然后使用一个填充且步长设置为1的
3
×
3
3 \times 3
3×3卷积输出投影层将其投影到
r
∈
R
1
×
H
×
W
r \in \mathbb{R}^{1 \times H \times W}
r∈R1×H×W,并将其加到原始LDCT图像
y
y
y上以产生最终的重建图像
g
+
r
g + r
g+r,其中
r
∈
R
1
×
H
×
W
r \in \mathbb{R}^{1 \times H \times W}
r∈R1×H×W。
WiTUnet架构专门为满足LDCT去噪的独特要求而设计。最初,嵌套的密集块的使用加强了不同网络层之间的信息流动,这有助于在保留关键图像细节的同时减少噪声。此外,W-MSA模块与LiPe模块的结合使得通过窗口多头自注意力机制有效地捕获全局信息,并通过局部信息增强组件细化局部细节。这种整合优化了全局与局部信息之间的协同作用,使WiTUnet在LDCT图像处理中能够有效地平衡计算效率与重建质量。总的来说,WiTUnet以其创新的U形网络结构,优化了信息流和特征整合,为LDCT去噪提供了一个有效的解决方案。它显著提高了图像质量,而不会牺牲细节丰富性,从而为临床诊断提供了更清晰、更可靠的图像。
Window Transformer block
将Transformer应用于低剂量计算机断层扫描(LDCT)图像去噪领域面临着多重挑战。首先,标准的Transformer需要对所有标记进行全局自注意力计算,这导致了高计算复杂性[4][20],特别是考虑到LDCT图像的高分辨率特性,这会增加特征图的维度并复杂化全局注意力计算。其次,尽管Transformer通过自注意力机制擅长捕捉长距离信息,但在图像去噪任务中保留局部信息至关重要,尤其是在对后续临床诊断至关重要的LDCT去噪中。考虑到Transformer可能不如卷积神经网络(CNN)那样有效地捕捉局部细节,融合CNN以获得更多局部细节变得尤为重要。
为了应对这些挑战,本研究采用了窗口Transformer (WT)块,如图1©所示。WT块利用W-MSA有效捕获长距离信息,并通过窗口化方法显著降低计算成本。此外,LiPe替换了传统的MLP层,以增强WT块捕获局部信息的能力。图1(b)展示了多个WT块的堆叠效果。在解码器中,由于WT块后通道数的变化,在WT块后进行通道投影,而在编码器和瓶颈层中不这样做。值得注意的是,来自第(l一1)块的输入,表示为
X
i
−
1
X_{i-1}
Xi−1,通过W-MSA和LiPe进行处理。WT块内的计算可以用以下数学表达式表示:
[X_{t}^{1} = \text{W-MSA}(\text{LN}(X_{i-1})) + X_{i-1} = \text{LiPe}(\text{LN}(X_{t-1})) + X_{t-1}]
在公式中,
X
t
X_{t}
Xt 和
X
X
X 分别代表W-MSA和LiPe的输出,而LN表示层归一化[1]。
基于非重叠窗口的多头自注意力(W-MSA)。在本文中,作者采用了非重叠的W-MSA机制,与视觉Transformer中使用的全局自注意力机制相比,大大降低了计算复杂性。考虑一个二维特征图
X
∈
R
C
×
H
×
W
X \in \mathbb{R}^{C \times H \times W}
X∈RC×H×W,其中C、H和W分别表示通道数、高度和宽度。作者将X划分为N=M个大小为M×M的非重叠窗口,每个窗口被展平并转换为
X
i
∈
R
M
′
×
C
X_{i} \in \mathbb{R}^{M' \times C}
Xi∈RM′×C。然后,每个窗口都通过W-MSA进行处理,如果它有d个头,那么每个头分配一个维度
d
=
C
/
d
d = C/d
d=C/d。每个窗口内的自注意力计算如下:
[X = [X_{1}, X_{2}, X_{3}, …, X_{N}]; \text{ where } N = \frac{H \times W}{M^{2}}]
[U = \text{Attention}(X_{i}W^{Q}, X_{i}W^{K}, X_{i}W^{V}), i = 1, 2, …, N = {Y_{1}', Y_{2}, Y_{3}, …, Y_{N}}]
[O = \text{concat}(Y_{1}, Y_{2}, Y_{3}, …, Y_{a})]
在这里,
W
Q
W^{Q}
WQ,
W
K
W^{K}
WK,
W
V
W^{V}
WV 分别属于第k个Head的Query、键和值的投影矩阵,且
W
R
W^{R}
WR
[W^{Q}, W^{K}, W^{V} \in \mathbb{R}^{C \times d_{s}}]。其中,
Y
i
Y_{i}
Yi 表示第k个Head从所有窗口获得的输出。通过连接所有Head的输出并应用线性投影,可以得到最终的输出。与之前的工作 [9][17]一致,作者将相对位置编码纳入到作者的注意力机制中。因此,注意力的计算在数学上定义为:
[ \text{Attention}(Q, K, V) = \text{SoftMax} \left( \frac{QK^T}{\sqrt{d_s}} + B \right) V ]
在公式中,B代表相对位置编码偏差,这是从可学习的参数
B
∈
R
(
2
M
−
1
)
×
(
2
M
−
1
)
B \in \mathbb{R}^{(2M-1) \times (2M-1)}
B∈R(2M−1)×(2M−1)获得的。W-MSA将全局自注意力计算复杂度从
O
(
H
2
W
2
C
)
O(H^2W^2C)
O(H2W2C)降低到了
O
(
M
2
H
W
C
)
O(M^2HWC)
O(M2HWC)。
局部图像视角增强(LiPe)。标准Transformer中的前馈网络(FFN),通常采用多层感知机(MLP),在捕捉局部信息方面能力有限。对于低剂量CT图像去噪,精确恢复图像细节至关重要,因为这些细节对准确疾病诊断至关重要。由于卷积神经网络(CNN)卷积核的性质,CNN对局部信息表现出更高的敏感性。为了弥补MLP的不足,作者采用了先前研究中的策略,用卷积块替换传统的FFN。如图2所示,作者首先使用线性投影来增加每个窗口特征图的通道数。随后,作者将窗口Reshape回原始特征图形式,并应用一个3×3的卷积核来捕捉局部细节。之后,特征图被窗口化和展平,然后通过另一个线性投影将通道数恢复到原来的维度。在每一层之间,作者使用高斯误差线性单元(GELU)作为激活函数。
嵌套密集块
图3(a)展示了WiTUnet结构的另一种视角,特别突出了嵌套密集块中的复杂嵌套密集跳跃路径连接。在图3中,
X
k
,
o
X_{k,o}
Xk,o 代表编码器的阶段,其中
k
∈
[
0
,
D
)
k \in [0, D)
k∈[0,D),而
X
k
,
v
X_{k,v}
Xk,v表示解码器的阶段,其中
k
∈
[
0
,
D
)
k \in[0,D)
k∈[0,D)且
v
=
D
−
k
v= D - k
v=D−k;
X
p
,
o
X_{p,o}
Xp,o表示瓶颈分。剩余的图解重点介绍了本文中提到的嵌套密集块,与传统的U-Net架构的主要区别在于重新设计的跳跃路径(由绿色和蓝色箭头指示)。这些重新设计的关系改变了编码器与解码器之间的交互方式。不同于传统U-Net中,编码器的特征图直接传递给解码器,在这里,它们首先通过一系列密集卷积块传输。这些块中的层数与从跳跃连接收到的特征图的总通道数相关。本质上,密集卷积块将编码器特征图的语义层次与解码器中将要处理的特征图拉近。作者假设,当接收到的编码器特征图在语义上与对应的解码器特征图相似时,优化器在解决优化问题时面临更简单的任务。正式地,跳跃路径可以表示如下。设
L
k
,
o
L_{k,o}
Lk,o为节点
X
k
X_{k}
Xk的输出。
k
k
k表示下采样的深度,
v
v
v表示在嵌套块中的横向位置。
需要指出的是,
X
0
,
0
=
X
o
,
o
(
T
i
n
)
X_{0,0} = X_{o,o}(T_{in})
X0,0=Xo,o(Tin)。其中符号
U
U
U表示上采样层,上采样操作如第3.1节所述。通常,
v
=
1
v=1
v=1的节点接收来自两个源的输入,
v
=
2
v=2
v=2的节点接收来自三个源的输入,以此类推,
v
=
i
v=i
v=i的节点接收
i
+
1
i+1
i+1个输入。为了更清晰地理解所涉及的计算,图3(b)直观地描述了每个节点的计算过程。