A Recycling Training Strategy for Medical Image Segmentation with Diffusion Denoising Models

要約

ノイズ除去拡散モデルは、画像に条件付けされたセグメント化されたマスクを生成することにより、画像のセグメンテーションに応用できることがわかりました。
既存の研究は主に、モデル アーキテクチャの調整や、テスト時のサンプリング戦略などの推論の改善に焦点を当てています。
この研究では、トレーニング戦略の改善に焦点を当て、新しいリサイクル方法を提案します。
各トレーニング ステップ中に、画像とランダム ノイズを考慮してセグメンテーション マスクが最初に予測されます。
この予測マスクは従来のグラウンド トゥルース マスクに代わるもので、トレーニング中のノイズ除去タスクに使用されます。
このアプローチは、ノイズの多いサンプルを生成するためのグラウンド トゥルース マスクへの依存を排除​​することで、トレーニング戦略を推論と調整するものと解釈できます。
私たちが提案した方法は、筋肉超音波、腹部 CT、前立腺 MR、脳 MR といった複数の医用画像データセットにわたって、標準的な拡散トレーニング、セルフコンディショニング、および既存のリサイクル戦略を大幅に上回ります。
これは、広く採用されている 2 つのサンプリング戦略、つまりノイズ除去拡散確率モデルとノイズ除去拡散暗黙モデルに当てはまります。
重要なのは、既存の拡散モデルは推論中にパフォーマンスの低下または不安定を示すことが多いのに対し、私たちの新しいリサイクルは一貫してパフォーマンスを向上または維持します。
同じネットワーク アーキテクチャとコンピューティング予算との公正な比較の下で、提案されたリサイクル ベースの拡散モデルが、非拡散ベースの教師ありトレーニングと同等のパフォーマンスを達成したことを示します。
提案された拡散モデルと非拡散モデルをアンサンブルすることにより、すべてのアプリケーションにわたって非拡散モデルの大幅な改善が観察され、この新しいトレーニング方法の価値が実証されました。
このペーパーでは、https://github.com/mathpluscode/ImgX-DiffSeg でリリースされている、完全に再現可能な JAX ベースの実装を使用して、これらの定量的結果を要約し、その値について説明します。

要約(オリジナル)

Denoising diffusion models have found applications in image segmentation by generating segmented masks conditioned on images. Existing studies predominantly focus on adjusting model architecture or improving inference, such as test-time sampling strategies. In this work, we focus on improving the training strategy and propose a novel recycling method. During each training step, a segmentation mask is first predicted given an image and a random noise. This predicted mask, which replaces the conventional ground truth mask, is used for denoising task during training. This approach can be interpreted as aligning the training strategy with inference by eliminating the dependence on ground truth masks for generating noisy samples. Our proposed method significantly outperforms standard diffusion training, self-conditioning, and existing recycling strategies across multiple medical imaging data sets: muscle ultrasound, abdominal CT, prostate MR, and brain MR. This holds for two widely adopted sampling strategies: denoising diffusion probabilistic model and denoising diffusion implicit model. Importantly, existing diffusion models often display a declining or unstable performance during inference, whereas our novel recycling consistently enhances or maintains performance. We show that, under a fair comparison with the same network architectures and computing budget, the proposed recycling-based diffusion models achieved on-par performance with non-diffusion-based supervised training. By ensembling the proposed diffusion and the non-diffusion models, significant improvements to the non-diffusion models have been observed across all applications, demonstrating the value of this novel training method. This paper summarizes these quantitative results and discusses their values, with a fully reproducible JAX-based implementation, released at https://github.com/mathpluscode/ImgX-DiffSeg.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google