[論文レビュー] The Swiss Army Knife for Image-to-Image Translation: Multi-Task Diffusion Models

Julia Wolleb, Robin Sandkühler+2 more|arXiv (Cornell University)|Apr 6, 2022

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

ひとこと要約

本論文は、分類タスクに限らない回帰・セグメンテーションモデルからの外部勾配を用いて、ベースの拡散モデルの微調整なしに、ノイズ除去拡散サンプリングをガイドするマルチタスク拡散モデルを提案する。この手法により、顔の年齢変化や脳腫瘍の成長シミュレーションといった制御可能な画像対画像翻訳が可能となり、1つの事前学習済みDDPMを用いて顔の年齢変化およびMR画像における腫瘍生成の両タスクで最先端の結果を達成した。

ABSTRACT

Recently, diffusion models were applied to a wide range of image analysis tasks. We build on a method for image-to-image translation using denoising diffusion implicit models and include a regression problem and a segmentation problem for guiding the image generation to the desired output. The main advantage of our approach is that the guidance during the denoising process is done by an external gradient. Consequently, the diffusion model does not need to be retrained for the different tasks on the same dataset. We apply our method to simulate the aging process on facial photos using a regression task, as well as on a brain magnetic resonance (MR) imaging dataset for the simulation of brain tumor growth. Furthermore, we use a segmentation model to inpaint tumors at the desired location in healthy slices of brain MR images. We achieve convincing results for all problems.

研究の動機と目的

• 1つの事前学習済み拡散モデルを用いて、柔軟かつタスク固有の画像対画像翻訳を実現すること。
• 背景や構造的詳細を保持しつつ、タスク関連の特徴のみを変更する課題に対処すること。
• 分類タスクに限らない、回帰およびセグメンテーションタスクへの分類器ガイドランスを拡張すること。
• 健康なMRスキャンに脳腫瘍を合成するなど、医療画像分野におけるリアルな合成データを生成すること。
• 推論時における勾配注入により、同じ拡散モデルを多様なタスクに再利用可能であることを示すこと。

提案手法

• タスク固有の適応なしに、ソースデータセット上で学習された事前学習済みのノイズ除去拡散確率的モデル（DDPM）をベースジェネレータとして使用する。
• 同じデータセット上で独立して学習された外部タスク固有モデル（回帰またはセグメンテーション）が、ガイド信号を提供する。
• DDIMサンプリング中に、各ステップでノイズの入力に対するタスク固有モデル出力の勾配によって、ノイズ除去プロセスがガイドされる。
• ガイド強度はスケーリング係数 $ s_t $ で制御され、正または負の値に設定することで、ターゲット属性（例：年齢や腫瘍サイズ）を増加または減少させることができる。
• 回帰の場合、ターゲット値 $ i $ を用いて $ s_t = i - R(x_t, t) $ を計算する。ここで $ R $ はステップ $ t $ における回帰モデルの予測値である。
• セグメンテーションの場合、ラベルマスクが望ましい腫瘍の位置を定義し、セグメンテーションモデルの損失の勾配を用いて、サンプリング中に腫瘍に類似する特徴を画像に注入する。

実験結果

リサーチクエスチョン

• RQ1 1つの事前学習済み拡散モデルを、再訓練なしに複数の画像対画像翻訳タスクに使用可能か？
• RQ2 外部の回帰およびセグメンテーションモデルからの勾配ガイドランスが、年齢や腫瘍サイズといった画像属性のサンプリング中に効果的に制御可能か？
• RQ3 背景、衣類、髪型などのターゲット外の画像特徴が、必要な属性のみを変更する際に保持されるか？
• RQ4 本アプローチにより、健康なMRスライスに脳腫瘍を合成するなど、リアルな医療画像の合成が可能か？
• RQ5 忠実度および属性制御の観点から、本手法はタスク固有の拡散モデルと比較してどの程度の性能を示すか？

主な発見

• 本手法は顔の画像に対して年齢回帰を効果的に実行し、40歳の入力から希望の年齢（例：10歳、20歳、60歳、80歳）のリアルな画像を生成した。
• 脳MR画像において、回帰損失の勾配を調整することで、腫瘍の成長または縮小を正確にシミュレートした。腫瘍サイズはターゲット値 $ i $ を用いて制御可能であった。
• セグメンテーションガイドランスによる手法は、ユーザーが指定した位置に健康な脳MRスライスに腫瘍をインpaintingするのに成功し、差分マップにより領域特有の変化が明確に確認された。
• 同じベースのDDPMが、回帰、腫瘍成長シミュレーション、腫瘍インpaintingのすべてのタスクで再利用された。これにより、フレームワークの汎用性が示された。
• 結果から、ターゲット属性に関連する特徴のみが変更されており、背景、顔の特徴、衣類などが入力画像と一致したままであった。
• 顔の年齢変化および脳腫瘍シミュレーションの両タスクにおいて、本手法は視覚的に説得力があり、最小限のアーティファクトと高い構造的忠実度を達成した。