[論文レビュー] 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 は、Anchor Condition with Uncertain Spatial Attention と Spectrum-Space Transformer (SS-Former) を用いた医用画像セグメンテーションのための Transformer ベースの拡散フレームワークを導入し、複数モダリティの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

研究の動機と目的

• UNet バックボーンよりもセグメンテーション品質を向上させるために、拡散ベースの医用画像セグメンテーションとトランスフォーマーの統合を動機づける。
• 拡散を安定化し意味的相互作用を高めるために、Anchor Condition と Semantic Condition の二重条件付け戦略を提案する。
• 周波数領域でノイズと意味表現を橋渡しする SS-Former の開発。
• 時間ステップ間で拡散ノイズを意味特徴と整合させる適応型 Neural Band-pass Filter (NBP-Filter) の導入。
• 多様なモダリティにわたる20の臓器セグメンテーションタスクでSOTA性能を示す。

提案手法

• 二つの条件付けストリーム: Anchor Condition は Uncertain Spatial Attention (U-SA) を介してセグメンテーションデコード特徴を拡散エンコーダへ注入し、拡散分散を低減する。
• Semantic Condition は Spectrum-Space Transformer (SS-Former) を用いて周波数空間で Neural Band-pass Filter (NBP-Filter) を組み込み、意味セグメンテーション情報を拡散埋め込みへ埋め込む。
• 拡散バックボーンは、生画像特徴に条件付けられたUNetベースの逆過程を伴う Denoising Diffusion Probabilistic Model (DPM) に従う。
• U-SA は、アンカー特徴をガウスカーネルで平滑化し、元の特徴と最大値を取り、空間注意に類似した1x1畳み込み変調を適用して融合する。
• SS-Former は条件と拡散特徴をフーリエ空間へ転送し、情報を交換するクロスアテンション様のモジュールを用い、拡散時間ステップに条件付けられた NBP-Filter を用いてスペクトルを整合させる。
• 訓練はノイズ予測損失に加え、スケジュール制御された条件付け監督を伴うアンカー損失（soft 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 トランスフォーマーベースの条件付けを拡散モデルと統合することで、UNetベースの拡散手法を超える医用画像セグメンテーションの改善が期待できるか？
• RQ2 Anchor Condition はトランスフォーマー・バックボーンを用いた場合に拡散分散を低減し、安定性を向上させるか？
• RQ3 SS-Former は周波数領域で拡散ノイズ埋め込みと意味条件付けを効果的に結合し、より良いセグメンテーションを実現できるか？
• RQ4 複数モダリティにまたがる拡散ベースのセグメンテーションにおける精度・多様性・収束性に対する U-SA と SS-Former の影響はどの程度か？

主な発見

• MedSegDiff-V2 は 5 モダリティにわたる 20 の臓器セグメンテーションタスクで最先端の性能を達成。
• Anchor Condition を用いた U-SA は vanilla diffusion の性能を大幅に改善し、より頑健な初期点を提供する。
• SS-Former を用いた Semantic Condition は、特に NBP-Filter と組み合わせると、ノイズと意味埋め込みを整合させることで顕著な改善をもたらす。
• エンセmbles の反復回数が減少して収束し、Gflops が低減する一方で Dice/IoU 指標が向上する。
• アブレーション研究は、Anchor Conditioning および SS-Former の両方がセグメンテーション品質の向上に有効であることを確認した。

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