[论文解读] MedSegDiff-V2: Diffusion based Medical Image Segmentation with Transformer

Junde Wu, Ji, Wei+4 more|arXiv (Cornell University)|Jan 19, 2023

Radiomics and Machine Learning in Medical Imaging被引用 14

一句话总结

MedSegDiff-V2 引入基于 Transformer 的扩散框架用于医学图像分割，采用带不确定空间注意的锚点条件 Anchor Condition with Uncertain Spatial Attention 和 Spectrum-Space Transformer (SS-Former) 进行语义条件，实现跨多模态的20个分割任务的最先进结果。

ABSTRACT

The Diffusion Probabilistic Model (DPM) has recently gained popularity in the field of computer vision, thanks to its image generation applications, such as Imagen, Latent Diffusion Models, and Stable Diffusion, which have demonstrated impressive capabilities and sparked much discussion within the community. Recent investigations have further unveiled the utility of DPM in the domain of medical image analysis, as underscored by the commendable performance exhibited by the medical image segmentation model across various tasks. Although these models were originally underpinned by a UNet architecture, there exists a potential avenue for enhancing their performance through the integration of vision transformer mechanisms. However, we discovered that simply combining these two models resulted in subpar performance. To effectively integrate these two cutting-edge techniques for the Medical image segmentation, we propose a novel Transformer-based Diffusion framework, called MedSegDiff-V2. We verify its effectiveness on 20 medical image segmentation tasks with different image modalities. Through comprehensive evaluation, our approach demonstrates superiority over prior state-of-the-art (SOTA) methodologies. Code is released at https://github.com/KidsWithTokens/MedSegDiff

研究动机与目标

• 动机：推动将 Transformer 与基于扩散的医学影像分割相结合，以提升分割质量，相较于 UNet 骨干网。
• 提出两重条件化策略——Anchor Condition 与 Semantic Condition，以稳定扩散过程并增强语义交互。
• 开发 SS-Former，将噪声和语义嵌入在频域中连接起来。
• 引入自适应神经带通滤波器（NBP-Filter），在不同时间步对齐扩散噪声与语义特征。
• 在跨多模态的20个器官分割任务中展示前沿性能。

提出的方法

• 两条条件化流：Anchor Condition 通过不确定空间注意（U-SA）将解码后的分割特征注入扩散编码器，以降低扩散方差。
• Semantic Condition 使用 Spectrum-Space Transformer (SS-Former) 在频域中将语义分割信息嵌入扩散嵌入，并搭配 Neural Band-pass Filter (NBP-Filter)。
• 扩散骨干遵循去噪扩散概率模型（DPM），其基于 UNet 的反向过程以原始图像特征为条件。
• U-SA 通过高斯核平滑锚点特征、与原始特征取最大，以及应用类似空间注意的 1x1 卷积调制，来融合特征。
• SS-Former 将条件与扩散特征转移到傅里叶空间，使用类似跨注意力的模块进行信息交换，并在扩散时间步条件下使用 NBP-Filter 以对齐频谱。
• 训练使用噪声预测损失，加上带有排程条件监督的锚点损失（软 Dice + 交叉熵）。

Figure 1: An illustration of MedSegDiff-V2, which starts from (a) an overview of the pipeline, and continues with zoomed-in diagrams of individual Models, including (b) SS-Former, and (c) NBP-Filter.

实验结果

研究问题

• RQ1 将基于 Transformer 的条件化与扩散模型结合是否能在医学影像分割中超越基于 UNet 的扩散方法？
• RQ2 Anchor Condition 在使用 Transformer 验背时是否能降低扩散方差并提高稳定性？
• RQ3 SS-Former 能否在频域中有效地将扩散噪声嵌入与语义条件耦合，以实现更好的分割？
• RQ4 在多模态下，U-SA 与 SS-Former 对扩散式分割的准确性、多样性和收敛性的影响如何？

主要发现

• MedSegDiff-V2 在跨 5 种模态的 20 个器官分割任务上达到最先进的性能。
• 带有 U-SA 的 Anchor Condition 显著提升了基础扩散性能，并提供更稳健的起点。
• Semantic Condition 配合 SS-Former，尤其是结合 NBP-Filter 时，通过对齐噪声与语义嵌入实现显著改进。
• 该模型需要更少的集成迭代即可收敛，且在提高 Dice/IoU 指标的同时展现更高的效率（更低的 Gflops）。
• 消融研究证实 Anchor Conditioning 与 SS-Former 在提升分割质量方面的有效性。

Figure 2: The visual comparison with SOTA segmentation models on BTCV.