(cvpr2024) Modality-agnostic Domain Generalizable Medical Image Segmentation by Multi-Frequency in Multi-Scale Attention

发表于 更新于 分类于

论文链接：arxiv

摘要

提出了一种模型：MADGNet，它由多频度多尺度注意模块（MFMSA）与集合子解码模块（ESDM）构成。前者通过结合多频率和多尺度特征，完善了空间特征提取过程，尤其是在捕捉边界特征方面，从而为组织轮廓和解剖结构提供了信息线索；而后者则减少了深度监督下的多任务学习中的信息丢失问题，尤其是在从低分辨率大幅上采样时。

论文的理论：频率方差是实现模型的模式无关性和领域通用性的关键因素。而各种模型都没有考虑到其重要性；各种模型都没有考虑到深度监督下的多任务学习可能造成的潜在信息损失。

改进的过程

在绘制比例和频率分布图的时候发现，医学图像在两个维度上呈现出不同的分布。并且发现：多频率与多尺度之间有相互依存的关系的同时，多频率自身也有很大的差异（最直接的证据是：频率方差高于比例方差）。（独立性与互补性）。模型不但要考虑相互依存的关系，也要考虑每个维度所能提供的信息。

创新

MFMSA 模块

该模块采用多频率信道注意（MFCA）与二维离散余弦变换（2 D DCT），通过提取频率统计信息生成信道注意图。然后使用多尺度空间注意力（MSSA）来提取具有区分性的边界特征，并从每个尺度中将其汇总。

MFMSA 可以全面捕捉广泛信息的核心：通过 MFCA 使用顺序注意法来抑制噪声信道的影响，然后使用 MSSA 提取不同尺度边界的判别特征图。（使用多频分析提取不同频段的各种特征，使用多尺度分析，捕捉复杂的细节与更广泛的结构信息）

E-SDM 模块

作用：可以防止在深度监督的多任务学习过程中因大幅提高采样率而造成信息丢失的问题。

网络架构

特征提取

图像在传入时，会经由 ResNeSt 网络的 split-attention 残差块进行特征图提取。上图中 $Block_{i}$ 就是由其组成。当特征图从 $Block_{i}$ 输出后，为了让网络的复杂度降低，会经过一个 1*1 的卷积层进行降维处理。之后再与下一阶段的经过上采样，MFMSA 模块处理后的特征图进行拼接后再传入这一阶段的 MFMSA Block 进行处理。当为第 5 阶段（初始阶段），则无须拼接，可直接上采样到第 4 阶段使用，无需经过 MFMSA 模块处理。

MFMSA 处理

MFMSA 模块结构图

规模分解

该部分的作用是加快网络训练，并且通过不同尺度的特征图学习更加广泛的数据。

对于输入 $X_{i}$，为了提高计算效率，将会使用缩减比为 $\gamma$ 的倍率来减小输入的通道数，使用下采样来降低图像的分辨率。输入的 $X_{i}$ 会分为 3 个分支：$X_{i}^{j},j=1,2,3$，不同的 j 对应着不同卷积核大小的二维卷积。卷积核大小为 $2j+1$。当 j=1 时，不会进行下采样。下采样的参数为 $Down_{s} = 2^{s - 1}$，若输入的宽与高为 $W$, $H$，则通过下采样后的大小为 $\frac{W}{2^{s-1}}$ 与 $\frac{H}{2^{s-1}}$。同时，对于图片的宽，高，通道数，均有一个最小值，确保传入 MFCA 模块之前的特征图有足够的特征。

MFCA

该模块的作用是通过在频域中进行特征提取来增强图像的通道信息表示，从而减少噪声对医学图像的影响。可以生成通道注意力图。

对于 $X_{i}^{j}$ 都有一个对应的 MFCA模块进入处理。特征图传入 MFCA 后，会先经过二维 DCT 将图像表示为由在不同频率上振荡的余弦函数产生的基图像的加权和。之后将其结果使用全局平均池化，全局最大池化，全局最小池化进行处理，将数据压缩成 $Z_{avg}$，$Z_{max}$，$Z_{min}$。然后使用两个全连接层（$C_{s} * \frac{C_{s}}{r}$ 与 $\frac{C_{s}}{r}*C_{s}$ ，其中，$C_{s}$ 为输入的通道数，$r$ 为压缩比）将这些输出聚合起来（先分别经过全连接层，再相加得到结果）。得到结果后，使用 sigmoid 函数将权值变为 0 到 1 的区间，再用该值进行特征重标定。

MSSA

该模块的作用是生成对各种尺度和噪声具有鲁棒性的边界线索特征图。由于在转入 MSSA 模块之前就已经经过了 MFCA 的处理，因此通过双重处理过程可以使模型理解各种模式之间的微妙的解剖学差异，和嘈杂医学图像中不规则病变的复杂特征。

具体的处理过程如上图所示，其中要注意的是：

  • 前景注意力图 $F_{i}^{s}$ 为 $sigmoid(Conv2D_{1}(\hat{X_{i}^{s}}))$，而背景注意力图 $B_{i}^{s}$ 为 $1-F_{i}^{s}$
  • 对于上面两条分支，由于在传输 MFCA 时经过了下采样处理，因此，在 MFCA 中进行汇总时，需要通过上采样处理，使分辨率匹配
  • $\alpha_{s}^{i}$ 与 $\beta_{s}^{i}$ 为可学习的参数

E-SDM 模块

使用的原因：带有深度监督的多任务学习可以增强模型的表示能力并防止梯度消失的问题，而当低分辨率的特征图上采样到高分辨率时，会在预测细节边界和结构时造成信息损失。

创新：E-SDM 通过级联地补充信息损失，将子任务的预测结果整合到核心任务的预测中，从而改进核心任务的预测精度。

该模块的伪代码如下：

  • 返回的结果是多个预测的集合

其损失函数与并行方式是相同的，公式如下：

性能

经过对比其他的医学图像分割模型，该模型在各个领域上均获取到了第一或第二的成绩。在医学图像分析中，该模型也可以成功描绘给出的所有模式的详细边界。

对于 MFMSA 的有效性研究

通过使用单尺度单频率网络，单尺度多频率网络，多尺度单频率网络，证明了多尺度，多频率网络的有效性：多尺度，多频率可显著提高特征表示能力，并提取出增强的边界特征。

对于 E-SDM 的有效性研究

通过使用并行方式进行多任务学习和深度监督实验，可以发现性能下降的问题。通过研究前向和后向流对分割性能的影响，可以发现后向流对于保持整体输出质量至关重要，因为它弥补了上采样导致的信息丢失。相比之下，前向流仅识别需要强调的区域，而后向流则保留了关键信息。

总结

MADGNet 的有效性可以归纳为三个方面：

  1. 在医学图像分割设计中，多尺度和多频率特性被证明是至关重要的。
  2. MFCA 擅长从嘈杂的医学图像特征图中提取辨别特征，通过 MSSA 获取更复杂的边界线索。
  3. E-SDM 在通过任务集合进行深度监督的多任务学习时，能补偿因大幅提高采样率而造成的信息损失。