Importance of Aligning Training Strategy with Evaluation for Diffusion Models in 3D Multiclass Segmentation

要約

最近、ノイズ除去拡散確率モデル (DDPM) が、画像に条件付けられたセグメンテーション マスクを生成することによって画像セグメンテーションに適用されましたが、アプリケーションは主に 2D ネットワークに限定され、3D フォーミュレーションからの潜在的な利点を活用することはできませんでした。
この作業では、初めて DDPM が 3D マルチクラス イメージ セグメンテーションに使用されます。
私たちは、トレーニング戦略を評価方法と一致させ、効率を向上させることに焦点を当てた 3 つの主要な貢献を行っています。
まず、モデルはサンプリングされたノイズの代わりにセグメンテーション マスクを予測し、Dice 損失によって直接最適化されます。
次に、前の時間ステップで予測されたマスクを再利用して、ノイズで破損したマスクを生成し、情報漏えいを減らします。
最後に、トレーニング中の拡散プロセスは、評価と同じ 5 つのステップに削減されました。
2 つの大規模なマルチクラス データ セット (前立腺 MR と腹部 CT) に関する研究を通じて、既存の DDPM と比較して大幅に改善されたパフォーマンスを実証し、同じ計算予算内で、U-net に基づく非拡散セグメンテーション モデルと競合するパフォーマンスに達しました。
JAX ベースの拡散フレームワークは、https://github.com/mathpluscode/ImgX-DiffSeg でリリースされています。

要約(オリジナル)

Recently, denoising diffusion probabilistic models (DDPM) have been applied to image segmentation by generating segmentation masks conditioned on images, while the applications were mainly limited to 2D networks without exploiting potential benefits from the 3D formulation. In this work, for the first time, DDPMs are used for 3D multiclass image segmentation. We make three key contributions that all focus on aligning the training strategy with the evaluation methodology, and improving efficiency. Firstly, the model predicts segmentation masks instead of sampled noise and is optimised directly via Dice loss. Secondly, the predicted mask in the previous time step is recycled to generate noise-corrupted masks to reduce information leakage. Finally, the diffusion process during training was reduced to five steps, the same as the evaluation. Through studies on two large multiclass data sets (prostate MR and abdominal CT), we demonstrated significantly improved performance compared to existing DDPMs, and reached competitive performance with non-diffusion segmentation models, based on U-net, within the same compute budget. The JAX-based diffusion framework has been released on https://github.com/mathpluscode/ImgX-DiffSeg.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google