[論文レビュー] Memory-Efficient 3D Denoising Diffusion Models for Medical Image Processing

Florentin Bieder, Julia Wolleb+3 more|arXiv (Cornell University)|Mar 27, 2023

Generative Adversarial Networks and Image Synthesis被引用数 8

ひとこと要約

本論文は PatchDDM を提案する。メモリ効率の良いパッチベースの拡散モデルで、パッチで訓練しつつ全解像度ボリュームでの推論が可能、BraTS2020 のセグメンテーションへ適用。

ABSTRACT

Denoising diffusion models have recently achieved state-of-the-art performance in many image-generation tasks. They do, however, require a large amount of computational resources. This limits their application to medical tasks, where we often deal with large 3D volumes, like high-resolution three-dimensional data. In this work, we present a number of different ways to reduce the resource consumption for 3D diffusion models and apply them to a dataset of 3D images. The main contribution of this paper is the memory-efficient patch-based diffusion model extit{PatchDDM}, which can be applied to the total volume during inference while the training is performed only on patches. While the proposed diffusion model can be applied to any image generation tasks, we evaluate the method on the tumor segmentation task of the BraTS2020 dataset and demonstrate that we can generate meaningful three-dimensional segmentations.

研究の動機と目的

• 大規模な3D体積に起因する医用画像処理におけるメモリ効率の良い3D拡散モデルの必要性を動機づける。
• パッチでの訓練と全ボリューム推論を実現する PatchDDM を提案する。
• 性能を維持しつつメモリ使用を削減するように拡散モデルのアーキテクチャを適応・単純化する。
• BraTS2020 のセグメンテーションで全解像度および半解像度のベースラインと PatchDDM を比較評価する。
• 精度と計算量のバランスをとるためのアンサンブルと加速サンプリングを検討する。

提案手法

• 3Dデータに適応した標準的な U-Net バックボーンを用いた denoising diffusion implicit models (DDIM) を採用。
• 3D 入力のメモリ削減のためグローバル自己注意を除去。
• メモリを保持しつつ訓練を安定化させるためスキップ接続を平均化に置換し、総メモリ量を抑えつつより広いネットワークを実現。
• PatchDDM を導入：パッチを座標エンコードして訓練し、全ボリューム推論を可能にすることで1回辺のメモリ使用を低減。
• 入力とサンプリングされたパッチへ格子状のデカルト座標を付加して座標エンコーディングを提供。
• 補助体積による条件付けとアンサンブルによる精度向上を伴う拡散ベースのセグメンテーション枠組みを用いてセグメンテーション性能を評価。

実験結果

リサーチクエスチョン

• RQ1 座標エンコード済みパッチでのメモリ効率の訓練が大規模3D医用体積の拡散モデルベースのセグメンテーションに有効となるか。
• RQ2 PatchDDM は BraTS2020 において全解像度およびダウンサンプリングベースのベースラインと比較してセグメンテーション性能でどうなるか。
• RQ3 アンサンブルと加速サンプリングが 3D 拡散ベースのセグメンテーションにおける Dice および HD95 指標に与える影響は何か。
• RQ4 PatchDDM のリソース要件（メモリと時間）は3D医用画像処理のベースライン手法と比べてどうか。

主な発見

Method	Training Memory (GB)	Inference Memory (GB)	Training Time (s)	Inference Time (s)
FullRes	78.5	25.7	2.12	1.01
HalfRes	10.5	4.90	0.351	0.124
PatchDDM	10.6	24.0	0.340	1.02

• PatchDDM はパッチで訓練し全ボリューム推論で3D体積を扱えるため、一般的なGPUで訓練可能。
• BraTS2020 で PatchDDM は単一評価の Dice スコアが HalfRes (0.86) より高く、FullRes (0.82) より高い 0.88 を達成するが、nnU-Net が依然として高い 0.96。
• アンサンブルは PatchDDM および FullRes の性能を向上させ、比較的小さなアンサンブルサイズで最先端のアンサンブルに近づく。
• PatchDDM の訓練メモリは FullRes (~78.5 GB) に比べ顕著に小さく、約 10.6 GB、HalfRes と同程度である一方、推論時メモリは FullRes に近い (~24.0 GB)。
• 加速サンプリング（ステップ数を減らすこと）と控えめなアンサンブルを組み合わせると、性能がほぼトップクラスに近い Dice を、計算コストを大幅に削減して達成できる。例として 20 ステップでほぼ上位の性能に到達。