[論文レビュー] Diff-UNet: A Diffusion Embedded Network for Volumetric Segmentation

Zhaohu Xing, Liang Wan+3 more|arXiv (Cornell University)|Mar 18, 2023

Radiomics and Machine Learning in Medical Imaging被引用数 41

ひとこと要約

Diff-UNet は拡散モデルを3D U-Netに組み込み、体積医用画像分割を実現し、Step-Uncertainty に基づく融合（SUF）を用いて多段予測を統合し、MRIとCTデータセット全体でロバストな結果を得る。

ABSTRACT

In recent years, Denoising Diffusion Models have demonstrated remarkable success in generating semantically valuable pixel-wise representations for image generative modeling. In this study, we propose a novel end-to-end framework, called Diff-UNet, for medical volumetric segmentation. Our approach integrates the diffusion model into a standard U-shaped architecture to extract semantic information from the input volume effectively, resulting in excellent pixel-level representations for medical volumetric segmentation. To enhance the robustness of the diffusion model's prediction results, we also introduce a Step-Uncertainty based Fusion (SUF) module during inference to combine the outputs of the diffusion models at each step. We evaluate our method on three datasets, including multimodal brain tumors in MRI, liver tumors, and multi-organ CT volumes, and demonstrate that Diff-UNet outperforms other state-of-the-art methods significantly. Our experimental results also indicate the universality and effectiveness of the proposed model. The proposed framework has the potential to facilitate the accurate diagnosis and treatment of medical conditions by enabling more precise segmentation of anatomical structures. The codes of Diff-UNet are available at https://github.com/ge-xing/Diff-UNet

研究の動機と目的

• ノイズ除去拡散モデルを活用した頑健な3D医用体積セグメンテーションを動機づける。
• ラベル埋め込みを用いてマルチラベル分割を扱う拡散ベースのエンドツーエンドフレームワークを提案する。
• 意味情報を抽出するための別個のFeature EncoderとDenoising-UNetを含むDenoisingモジュールを導入する。
• テスト時に拡散ステップ全体の予測を融合するStep-Uncertaintyベースの融合（SUF）モジュールを開発する。
• BraTS2020、MSD Liver、BTCVデータセットにおいて最先端手法より優れていることを示す。

提案手法

• 拡散ベースの予測を可能にするために、多クラス分割をワンホット符号化されたマルチチャネルラベルマップへ変換する。
• 体積特徴とノイズの多クラスラベルを融合してx0予測を行う、Feature EncoderとDenoising-UNetからなるDenoisingモジュールを使用する。
• Dice、BCE、MSEの損失を組み合わせた複合損失で訓練する。
• テスト時には複数のDDIMステップを生成し、ステップインデックスと予測不確実性を考慮したSUF重み付け方式で予測を融合する。
• 不確実性は各ステップごとに複数のフォワードパス（MC風サンプリング）で推定され、最終セグメンテーション出力のステップ予測を重み付けするのに用いられる。

実験結果

リサーチクエスチョン

• RQ1 拡散モデルを3D U-Netアーキテクチャに効果的に組み込み、マルチラベルの医用体積分割を実行できるか？
• RQ2 Step-Uncertaintyベースの融合機構は、難易度の高いデータセットにおいて拡散ベースの分割の頑健性を向上させるか？
• RQ3 Diff-UNetは、多臓器CTおよび多モード MRI分割ベンチマークにおいて、最先端手法と比較してどのように性能を示すか？

主な発見

• Diff-UNetはBraTS2020で平均Diceスコアが優れており（85.35% average）およびTC/WT/ET領域で複数のSOTA手法と比較して改善を示す。
• MSD Liverでは、Liverで95.72% Dice、HD95は0.222を達成し、Liver Tumorで51.65% Dice、HD95は17.280、全領域で73.69 Diceと8.751 HD95を平均。
• BTCVでは、8つの臓器を通じてHD958.115で平均Dice最高値83.75を達成。
• アブレーション研究により、別個のFeature Encoder（FE）とSUFモジュールを組み込むと性能が向上し、不確実性サンプリングの最適結果はS=4であることが示された。
• 定性的な結果は、Diff-UNetが他の手法よりも特に小さなターゲットでより正確なセグメンテーションを生み出すことを示している。