[论文解读] Learning Multi-axis Representation in Frequency Domain for Medical Image Segmentation

Jiacheng Ruan, Jingsheng Gao+2 more|arXiv (Cornell University)|Dec 28, 2023

Advanced Neural Network Applications被引用 4

一句话总结

本文提出 MEW-UNet，一种基于 U-Net 的架构，通过新型多轴外部权重模块学习多轴频域表征，以增强医学图像分割。通过在三个空间与通道轴上独立应用二维离散傅里叶变换（2D DFT），并在频域中利用可学习权重调制特征，该模型捕获了更丰富的全局与局部上下文信息，在四个医学影像数据集上实现了最先进性能，ISIC18 数据集上的平均 Dice 分数达到 91.00%。

ABSTRACT

Recently, Visual Transformer (ViT) has been extensively used in medical image segmentation (MIS) due to applying self-attention mechanism in the spatial domain to modeling global knowledge. However, many studies have focused on improving models in the spatial domain while neglecting the importance of frequency domain information. Therefore, we propose Multi-axis External Weights UNet (MEW-UNet) 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 features and assigns the external weight in the frequency domain, which is generated by our External 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, including Synapse, ACDC, ISIC17 and ISIC18 datasets, and our approach demonstrates competitive performance, owing to its effective utilization of frequency domain information.

研究动机与目标

• 解决现有基于视觉 Transformer 的模型在医学图像分割中仅关注空间域建模、忽略频域信息的局限性。
• 通过在多个轴上利用频域信号，改进医学图像分割中的全局上下文建模。
• 开发一种可学习机制，有效融合频域表征与局部空间特征，以增强特征学习。
• 证明多轴频域分析相较于单轴分析能提供更具判别性的信号模式，尤其适用于细微病理结构。

提出的方法

• 多轴外部权重模块（MEWB）在三个轴上独立应用二维离散傅里叶变换（2D DFT）：高-宽、通道-宽、通道-高，以提取多轴频域表征。
• 通过外部权重生成器生成可学习的外部权重，并将其应用于频域特征，以调制全局上下文信息。
• 第四分支应用深度可分离（DW）卷积，以保留局部空间细节，确保局部与全局表征的互补学习。
• MEWB 取代 U-Net 编码器-解码器架构中视觉 Transformer 的自注意力模块，构成 MEW-UNet 模型。
• 加权调制后的频域特征图通过逆 DFT 转换回空间域，支持端到端训练。
• 模型在四个公开医学图像分割数据集上使用 Dice 分数和 mIoU 等标准指标进行端到端训练。

实验结果

研究问题

• RQ1 与单轴或仅空间域方法相比，多轴频域表征学习是否能提升医学图像分割模型的性能？
• RQ2 在频域中集成可学习外部权重如何影响模型捕捉长距离依赖与上下文信息的能力？
• RQ3 在分割任务中，频域特征与局部空间特征的相对贡献如何？二者如何相互作用？
• RQ4 所提出的 MEW-UNet 是否在多种数据集上均优于现有的最先进模型？

主要发现

• MEW-UNet 在 ISIC18 数据集上达到 91.00% 的平均 Dice 分数，优于先前最先进模型如 MT-UNet（90.43%）和 Swin-Unet（88.07%）。
• 消融实验表明，若移除 DW 卷积分支，mIoU 下降 0.61%，证实了局部特征学习的重要性。
• 若在外部权重生成器中不使用内部表征模块（IRB）而仅使用随机权重，mIoU 下降 1.31%，证明了可学习权重生成机制的必要性。
• 完整 MEWB（包含全部三个频域轴与 DW 分支）在 ISIC18 上达到 81.90% 的 mIoU，显著优于单轴或无 DW 的基线模型。
• 可视化结果表明，MEW-UNet 在小病灶或不规则病灶等挑战性病例中，边界更精确、更完整，优于 MT-UNet 和 TransFuse。
• 频域信号强度分析表明，多轴 DFT 在不同组织类型间揭示了非重叠的信号模式，验证了该设计在提升判别性方面的合理性。