[论文解读] MEW-UNet: Multi-axis representation learning in frequency domain for medical image segmentation

Jiacheng Ruan, Mingye Xie+3 more|arXiv (Cornell University)|Oct 25, 2022

Advanced Neural Network Applications被引用 15

一句话总结

MEW-UNet 提出了一种新颖的多轴外部权重模块（Multi-axis External Weights block），通过沿三个空间轴进行二维傅里叶变换来学习全局频域表征，同时结合深度可分离卷积以提取局部特征，从而在 U-Net 架构中替代自注意力机制。该模型在 Synapse 数据集上的 HD95 指标上相比 MT-UNet 提升了 10.15mm，实现了最先进性能。

ABSTRACT

Recently, Visual Transformer (ViT) has been widely used in various fields of computer vision due to applying self-attention mechanism in the spatial domain to modeling global knowledge. Especially in medical image segmentation (MIS), many works are devoted to combining ViT and CNN, and even some works directly utilize pure ViT-based models. However, recent works improved models in the aspect of spatial domain while ignoring the importance of frequency domain information. Therefore, we propose Multi-axis External Weights UNet (MEW-UNet) for MIS based on the U-shape architecture by replacing self-attention in ViT with our Multi-axis External Weights block. Specifically, our block performs a Fourier transform on the three axes of the input feature and assigns the external weight in the frequency domain, which is generated by our Weights Generator. Then, an inverse Fourier transform is performed to change the features back to the spatial domain. We evaluate our model on four datasets and achieve state-of-the-art performances. In particular, on the Synapse dataset, our method outperforms MT-UNet by 10.15mm in terms of HD95. Code is available at https://github.com/JCruan519/MEW-UNet.

研究动机与目标

• 为解决现有医学图像分割模型仅关注空间域表征而忽略频域信息的局限性。
• 通过多轴傅里叶变换引入频域特征，提升基于 U-Net 架构的全局上下文建模能力。
• 通过将频域注意力与深度可分离卷积结合于统一模块中，实现全局与局部特征学习的平衡。
• 开发一种频域感知的注意力机制，增强表征学习能力，同时避免依赖视觉 Transformer 中的标准自注意力机制。
• 通过整合频域不变性与多尺度上下文建模，实现在多个医学影像基准上的最先进分割性能。

提出的方法

• 多轴外部权重（MEW）模块沿通道维将输入特征图划分为四个分支。
• 其中三个分支分别沿不同轴（高度-宽度、通道-宽度、通道-高度）应用二维离散傅里叶变换（2D DFT），将特征投影至频域。
• 在频域中应用可学习的外部权重，以调制频域分量，增强全局上下文建模能力。
• 对加权后的频域特征应用逆向二维傅里叶变换，将其转换回空间域。
• 第四个分支应用深度可分离卷积，以保留局部空间细节。
• 将所有四个分支的输出进行拼接，并与原始输入残差连接，形成多轴外部权重模块（MEWB）。

实验结果

研究问题

• RQ1 多轴频域表征学习是否能超越仅基于空间域建模的医学图像分割性能？
• RQ2 在多个轴上使用二维傅里叶变换进行频域学习，如何影响模型捕捉全局上下文的能力？
• RQ3 将频域注意力与局部卷积结合于统一模块中，对分割任务的贡献是什么？
• RQ4 所提出的 MEW-UNet 是否在医学图像分割任务中优于现有的混合 CNN-ViT 及纯 ViT 模型？
• RQ5 在频域中引入可学习外部权重如何影响模型的泛化能力与性能表现？

主要发现

• MEW-UNet 在四个公开医学图像分割数据集（ISIC17、ISIC18、Synapse 和 ACDC）上均实现了最先进性能。
• 在 Synapse 数据集上，与 MT-UNet 相比，MEW-UNet 的 HD95 提升了 10.15mm，DSC 提升了 0.33%，表明其在边界定位方面具有显著优势。
• 在 ACDC 数据集上，MEW-UNet 达到了最高的 DSC 得分 91.00%，优于 MT-UNet（90.43%）和 Swin-Unet（88.07%）。
• 消融实验表明，沿全部三个轴添加频域操作可逐步提升 mIoU 与 DSC，完整 MEWB 在 ISIC18 上实现 81.90% 的 mIoU 与 90.05% 的 DSC。
• 在消融研究中，将 GroupNorm 替换为 BatchNorm 导致性能下降，表明 GroupNorm 在此架构中更为有效。
• 在 Synapse 数据集上的可视化结果表明，MEW-UNet 生成的分割掩码更清晰、更精确，尤其在器官边界区域表现更优。