【用于医学图像分割的轻量级边界感知Transformer】

代码：BATFormer

发表时间：2023

发表期刊：IEEE TMI

这篇论文介绍了一种新的Transformer架构，称为BATFormer（Boundary-Aware Lightweight Transformer），旨在改进医学图像分割的效率和效果。

目标

• 解决现有问题：传统的卷积神经网络（CNNs）由于感受野有限，在处理需要长距离依赖关系的任务时表现不佳。而现有的Transformer虽然能解决这一问题，但其计算复杂度高，并且在进行全局表示学习时采用了固定的窗口分区策略，这在医学图像分割任务中可能导致边界失真。

创新点

• 跨尺度全局Transformer (CGT) 模块：为了利用Transformer在建立长距离依赖关系上的优势，同时减少计算复杂度，作者提出了一个CGT模块。该模块可以联合使用多个小规模特征图来提取更丰富的全局特征，从而降低计算成本。
• 边界感知局部Transformer (BLT) 模块：考虑到形状建模在医学图像分割中的重要性，设计了BLT模块。与传统Transformer采用固定窗口分区不同的是，BLT采用基于熵的自适应窗口分区方案，以减少计算复杂度的同时保持物体的形状。

方法

A. 概述

BATFormer的整体架构如图2所示，包括一个轻量级U形主干网络和两个主要组件，旨在解决前述的问题：

1. 跨尺度全局Transformer (CGT)
2. 边界感知局部Transformer (BLT)

这两个组件分别针对医学图像分割中遇到的计算复杂度高和边界失真问题进行了优化。

1. 跨尺度全局Transformer (CGT)

• 计算成本与输入序列长度的关系：在大规模表示（如F1和F2）上构建全局Transformer的成本非常高，因为计算复杂度与输入序列长度的平方成正比。

• 深度语义特征的需求：全局依赖关系主要用于构建深层语义特征。因此，CGT仅接收多个小规模（即高层次）特征图（即F3、F4和F5），通过跨尺度自注意力机制，CGT能够在较低计算复杂度的情况下提取更丰富的语义特征，用于全局表示学习。这反过来补充了CNN模块，实现了联合的局部和全局特征提取与融合。

2. 边界感知局部Transformer (BLT)

• 灵活窗口分区方案：为了解决边界失真的问题，提出了基于熵计算的灵活窗口分区方案来定位与边界相关的窗口。通过这种方法，BLT可以过滤掉非边界窗口，更专注于探索长距离特征以进行边界检测，并且降低了计算复杂度。

B.轻量级U形主干网络

BATFormer的轻量级U形主干网络是整个架构的基础，它主要由四个下采样块、四个上采样块和五个跳跃连接组成。这种设计借鉴了经典的U-Net架构，但进行了优化以减轻计算负担，同时保持高效的特征提取能力。

主干网络结构

1. 输入处理：

   • 每个输入图像首先经过两次3x3卷积操作。

2. 下采样阶段（Encoder）：

   • 经过初始卷积后，图像被送入四个连续的下采样块中。
   • 每个下采样块包含一个2x2的最大池化层（步长为2），用于减少空间维度，以及两个3x3的卷积层来提取特征。
   • 下采样的通道数依次设置为$2C,4C,8C,8C$，其中$C=16$。

3. 跳跃连接：

   • 相同尺度的下采样块和上采样块通过跳跃连接相连，这有助于保留和传递局部信息，防止在下采样过程中丢失重要细节。

4. 上采样阶段（Decoder）：

   • 四个上采样块负责将特征图逐步恢复到原始尺寸。
   • 每个上采样块包括双线性插值（用于放大图像）、特征图拼接（与对应的下采样块通过跳跃连接传来的特征图结合），以及两个3x3的卷积层。
   • 上采样的通道数依次设置为$4C,2C,C,C$，同样$C=16$。

C. 跨尺度全局Transformer (CGT)

CGT旨在通过两个跨尺度注意力模块和一个前馈网络（FFN）来构建三个小尺度特征图之间的全局依赖关系。这三个特征图由主干网络生成，分别是 $F3\in {\mathbb{R}}^{4C\times \frac{H}{4}\times \frac{W}{4}}$、$F4\in {\mathbb{R}}^{8C\times \frac{H}{8}\times \frac{W}{8}}$ 和 $F5\in {\mathbb{R}}^{8C\times \frac{H}{16}\times \frac{W}{16}}$，其中 $C$ 是第一个分支头的通道数，$(H,W)$ 是输入图像的分辨率。

跨尺度注意力模块

在跨尺度注意力模块中：

• 查询生成：最高分辨率的特征图 $F3$ 被投影成两个查询 $Q3,4\in {\mathbb{R}}^{\frac{HW}{16}\times d}$和 $Q3,5\in {\mathbb{R}}^{\frac{HW}{16}\times d}$，其中 $d$ 是Transformer模块的维度。

• 键值对生成：较低分辨率的特征图 $F4$ 和 $F5$ 分别被投影成键值对 { K 4 , V 4 } ∈ R H W 64 × d \{K4, V4\} \in \mathbb{R}^{\frac{HW}{64} \times d} {K4,V4}∈R64HW​×d和 { K 5 , V 5 } ∈ R H W 256 × d \{K5, V5\} \in \mathbb{R}^{\frac{HW}{256} \times d} {K5,V5}∈R256HW​×d。

• 跨尺度注意力计算：跨尺度注意力公式如下：
  $F_i^{ca}(Q_{3,i},K_i,V_i)=\text{softmax}\left(\frac{Q_{3,i}{K}_i^T}{\sqrt{d}}\right)V_i$
  其中 $i$分别为4和5，对应于 $F4$ 和 $F5$。

与标准自注意力机制不同的是，这里的键 $K$ 和值 $V$ 来自其他两个较小尺度的特征图，而不是来自同一输入。这样做有两个主要优点：

1. 降低计算复杂度：由于 $K$和 $V$ 的序列长度较短，计算复杂度可以减少 $2^2$ 或 $2^4$ 倍。

2. 引入多样化特征：不同的 $K$ 和 $V$ 组对应于不同尺度的感受野和语义信息，从而增强了依赖关系的多样性。

特征融合与前馈网络（FFN）

在送入FFN之前，两组跨尺度注意力的结果会被结合并精炼：

$F^{ca}=(F_4^{ca,1}\oplus \cdots \oplus F_4^{ca,g}\oplus F_5^{ca,1}\cdots \oplus F_5^{ca,g})\cdot W^{ca}+F3,$

其中 $g$是CGT中预定义的Transformer头的数量，$W^{ca}\in {\mathbb{R}}^{2gd\times d}$是用于组合的可学习投影矩阵，$\oplus$ 表示拼接操作，$+$ 表示通过元素级加法实现的残差连接。

经过跨尺度注意力后的FFN进一步处理以获得CGT的最终输出：

$F^{CGT}=(\max (0,F^{ca}{W}_{c}1+b_{c}1)\cdot W_{c}2+b_{c}2)+F^{ca},$

其中 $W_{c}1\in {\mathbb{R}}^{d\times 4d}$和 $W_{c}2\in {\mathbb{R}}^{4d\times d}$是可学习的投影矩阵，$b_{c}1$ 和 $b_{c}2\in \mathbb{R}$ 是偏置项。

CGT的主要特点:

• 多尺度特征融合：通过跨尺度注意力机制，CGT能够有效地融合来自不同尺度特征图的信息，建立更丰富的全局依赖结构。

• 计算效率提升：通过使用较小尺度的特征图作为键值对，显著降低了计算复杂度，同时保持了有效的特征交互。

• 增强的全局表示能力：跨尺度注意力不仅减少了计算量，还增强了模型对全局上下文的理解，有助于提高分割任务中的表现。

D. 边界感知局部Transformer (BLT)

BLT的设计旨在克服现有Transformer中刚性窗口划分带来的边界信息丢失问题。它包含三个主要阶段：动态边界感知窗口生成、边界附近的自注意力机制以及特征蒸馏的前馈网络（FFN）。这些设计使得BLT能够在保持计算效率的同时，更好地捕捉和处理医学图像中的边界信息。

1. 动态边界感知窗口生成

• 初始窗口集合：首先，在特征图 $F2\in {\mathbb{R}}^{2C\times \frac{H}{2}\times \frac{W}{2}}$ 上均匀密集地平铺窗口，收集所有可能的窗口位置，形成初始窗口集合 ({w})。

• 熵计算：为了确定每个窗口是否位于边界附近，首先需要计算由CGT产生的概率图 $P^{CGT}\in {\mathbb{R}}^{c\times \frac{H}{4}\times \frac{W}{4}}$中每个位置 $(m,n)$ 的熵：
  $C_p(m,n)=-\frac{1}{{\log }_{2}c}{\sum }_{l=1}^c{P}^{CGT}(l,m,n){\log }_2{P}^{CGT}(l,m,n)$
  其中 $c$ 是类别数量。

• 窗口评分：然后，根据每个窗口内的熵值计算窗口得分 $C_w(x,y)$，其中 $(x,y)$ 是窗口的左上角坐标，$(h,w)$ 是窗口大小：
  $C_w(x,y)=\frac{1}{\lfloor \frac{h}{2}\rfloor \lfloor \frac{w}{2}\rfloor }{\sum }_{m=0}^{\lfloor \frac{h}{2}\rfloor }{\sum }_{n=0}^{\lfloor \frac{w}{2}\rfloor }{C}_p\left(\lfloor \frac{x}{2}\rfloor +m,\lfloor \frac{y}{2}\rfloor +n\right).$

• 非极大值抑制（NMS）：对初始窗口集合 {w}$ 根据上述评分进行非极大值抑制，去除冗余窗口，并通过RoIAlign与 $F2$对齐，最终得到边界附近的窗口集合 { w ∗ } \{w*\} {w∗} 和对应的特征集合 ${f*}。

2. 边界附近的自注意力机制

• 多头自注意力：在每个边界附近的窗口内执行多头自注意力操作，以建立长距离依赖关系：
  $F^{sa}(f_j^{\ast })=\text{softmax}\left(\frac{f_j^{\ast T}{E}_q{f}_j^{\ast }{E}_k}{\sqrt{d}}\right)f_j^{\ast }{E}_v,$
  其中 $E_q\in {\mathbb{R}}^{2C\times d}$, $E_k\in {\mathbb{R}}^{2C\times d}$, $E_v\in {\mathbb{R}}^{2C\times d}$ 是可学习的投影矩阵，$W^{sa}\in {\mathbb{R}}^{gd\times d}$ 用于组合不同头部的结果。

• 特征融合：将多头自注意力结果结合并精炼：
  $F^{sa}=(F^{sa,1}\oplus \cdots \oplus F^{sa,g})\cdot W^{sa}+F2,$

3. 特征蒸馏的前馈网络（FFN）

• 最终输出：经过FFN进一步处理，获得BLT的最终输出：
  $F^{BLT}=(\max (0,F^{sa}{W}_{s}1+b_{s}1)\cdot W_{s}2+b_{s}2)+F^{sa},$
  其中 $W_{s}1\in {\mathbb{R}}^{d\times 4d}$, $W_{s}2\in {\mathbb{R}}^{4d\times d}$ 是可学习的投影矩阵，$b_{s}1$ 和 $b_{s}2\in \mathbb{R}$ 是偏置项。

计算复杂度分析

• BLT的计算复杂度：只有边界附近的区域会通过上述边界附近的局部多头自注意力进行表示学习，因此其计算复杂度为：
  ${\Omega }^{BLT}=k(6dhwC+2d(hw{)}^2+d^{2}hw),$
  其中 $k$ 是最大窗口数。在实验中，采用 $h=\frac{H}{32}$, $w=\frac{W}{32}$，并将 $k$ 设置为 $\alpha \frac{HW}{hw}$，$\alpha \in (0,1)$。因此，公式8可以重写为：
  ${\Omega }^{BLT}\approx O\left(\alpha \frac{1}{16\times 32}d(HW{)}^2\right).$

• 与标准自注意力的比较：直接在 (F2) 上应用标准自注意力的计算复杂度为：
  ${\Omega }^{VIT}\approx O\left(\frac{1}{8}d(HW{)}^2\right).$
  相比之下，BLT的计算复杂度最多低64倍。

BLT的主要特点

• 边界感知窗口生成：通过基于熵的评分机制和非极大值抑制，BLT能够自适应地定位边界附近的窗口，有效减少了对边界信息的破坏。

• 高效的边界识别：只对边界附近的区域应用多头自注意力，显著降低了计算复杂度，同时确保了对边界的精确建模。

• 轻量级设计：通过精心设计的计算流程和优化的参数配置，BLT在保持高效计算的同时实现了高质量的边界检测。

E. 多尺度软监督

为了应对医学图像分割中多尺度特征带来的边界不确定性问题，引入了软监督机制。这一机制旨在通过在不同尺度上提供监督信号，缓解由于下采样或上采样导致的边界信息丢失。具体来说，每个真实标签掩码 $M\in {\mathbb{R}}^{H\times W}$ 会被调整到多个尺度以用于监督（例如，$G^{BLT}\in {\mathbb{R}}^{c\times \frac{H}{2}\times \frac{W}{2}}$ 用于训练BLT，以及 $G^{CGT}\in {\mathbb{R}}^{c\times \frac{H}{4}\times \frac{W}{4}}$ 用于训练CGT，同时保持原始边界的分布。

软监督的具体实现

给定类别 $l$ 和下采样尺度 $s$，调整后的真值图 $G_{l,s}$ 的构建方法如下：

$G_{l,s}(i,j)=\frac{1}{2^{2(s-1)}}{\sum }_{(m,n)\in O_{i,j}}|M(m,n)==l|,$

其中 $O_{i,j}$ 表示真值图 $M$ 中对应于 $G_{l,s}$ 中位置 $(i,j)$ 的下采样块，定义为 $M((i-1)2^{s-1}+1:i{2}^{s-1},(j-1)2^{s-1}+1:j{2}^{s-1})$。这样，$G_{l,s}$ 中位置 $(i,j)$ 的概率值由下采样块内类别 $l$ 出现的频率决定，并且 $G_{l,s}$ 中的边界分布会接近原始掩码 $M$ 的边界分布。

对于给定的下采样尺度 $s$，相应地构造一组调整后的真值图 G s = { G l , s ∣ l = 1 , 2 , . . . , c } G_s = \{G_{l,s}|l = 1, 2, ..., c\} Gs​={Gl,s​∣l=1,2,...,c}。然后，定义相同尺度下的预测图 P s = { P l , s ∣ l = 1 , 2 , . . . , c } P_s = \{P_{l,s}|l = 1, 2, ..., c\} Ps​={Pl,s​∣l=1,2,...,c} 与真值图 $G_s$ 之间的损失函数为：

$L_s={\sum }_{l=1}^c{\sum }_{i=1}^{\frac{H}{2^{s-1}}}{\sum }_{j=1}^{\frac{W}{2^{s-1}}}|G_{l,s}(i,j)-P_{l,s}(i,j)|.$

设置 $s=1,2,3$ 分别定义用于CNN、BLT和CGT的损失函数 KaTeX parse error: Can't use function '\)' in math mode at position 8: L^{CNN}\̲)̲，\(L^{BLT}，和 $L^{CGT}$。这些特征进一步融合以生成最终预测，并通过交叉熵损失和Dice损失进行惩罚，记作 $L^C$。

总体损失函数

多尺度软监督的整体损失函数 $L$ 定义为：

$L={\beta }_1{L}^{CNN}+{\beta }_2{L}^{BLT}+{\beta }_3{L}^{CGT}+{\beta }_4{L}^C,$

其中 ${\beta }_1,{\beta }_2,{\beta }_3,$ 和 ${\beta }_4$ 是平衡超参数，在实验中分别设置为 0.2, 0.1, 0.1, 和 0.6。

多尺度软监督的主要特点

• 缓解边界不确定性：通过在不同尺度上提供监督信号，确保模型在各个尺度上都能学习到准确的边界信息，减少了因下采样或上采样带来的边界信息丢失。

• 保持原始边界分布：调整后的真值图保留了原始掩码中的边界分布，使得模型能够更好地学习边界特征。

• 综合多种损失：结合了CNN、BLT、CGT模块各自的损失以及最终预测的交叉熵和Dice损失，使得模型能够在多个层次上进行优化，提高了整体分割性能。

IV. 评估

本节对BATFormer进行了广泛的对比实验，涵盖了最先进的基于CNN和基于Transformer（或混合）的方法，并在公开可用的数据集上进行了评估。

A. 数据集

评估使用了两个涵盖2D和3D医学图像数据的数据集：

1. ACDC：由150名患者的心脏MRI组成

2. ISIC 2018：包含2596张带有良好标注的皮肤病变图像

消融实验：

对比实验：

如果有想发表论文的小伙伴，没有创新点，代码能力差，论文有疑惑等问题，可以找我呀