ImgX-DiffSeg：基于DDPMs的3D醫(yī)學圖像分割

ImgX-DiffSeg：基于 DDPMs 的 3D 醫(yī)學圖像分割

前言

本篇文章繼續(xù)解讀醫(yī)學圖像 diffusion 系列，之前我們分別介紹過在自監(jiān)督和有監(jiān)督分割中的 diffusion 應用。鏈接：

ICLR 2023：基于 diffusion adversarial representation learning 的血管分割

MedSegDiff：基于 Diffusion Probabilistic Model 的醫(yī)學圖像分割

而這次的《Importance of Aligning Training Strategy with Evaluation for Diffusion Models in 3D Multiclass Segmentation》這篇文章并不是一種新的 diffusion 應用，而是對訓練和推理策略進行優(yōu)化，并適應 3D 的醫(yī)學圖像分割任務，參考鏈接在文末。

目前存在的問題

目前帶有 diffusion model 的架構訓練和推理耗時。

在一些分割任務中，并不確定 diffusion model 預測噪聲推斷分割圖和直接預測分割圖哪個效果更好。

模型過度依賴先前時間步中的信息。

ImgX-DiffSeg 架構

概述

DDPM 是一種生成模型，可用于圖像去噪和分割。工作原理是模擬干凈圖像的概率分布，然后在圖像中添加噪點以生成噪聲版本。相反的，模型嘗試通過移除添加的噪點來對圖像進行降噪。在圖像分割的情況下，模型會生成分割掩碼，可以根據(jù)輸入圖像的特征將圖像分成不同的區(qū)域，更細節(jié)的內容推薦閱讀前置文章（強烈建議）。

對于 ImgX-DiffSeg，整體流程如下圖所示。首先，該架構預測的是分割掩碼而不是采樣噪聲，并直接通過 Dice Loss 進行優(yōu)化。這意味著 ImgX-DiffSeg 可以直接預測圖像的分割圖，而不是生成噪點并用它來推斷分割。其次，回收上一個時間步中預測的掩碼，生成（noise-corrupted mask）噪音損壞的掩碼。這有助于減少信息泄露，當模型過度依賴先前時間步中的信息時，就會發(fā)生這種情況。最后，將訓練的擴散過程減少到五個步驟，與推理過程相同。擴散過程是一種平滑圖像中噪點的方法，減少步驟數(shù)有助于提高效率。

請?zhí)砑訄D片描述

DDPM with Variance Schedule Resampling

訓練過程的公式和 DDPM 基本是保持一致的，下面的公式分別表示反向過程中預測噪聲和原圖：

推理過程中，DDPM 中的生成過程從正常噪聲開始，由變量 xT 表示。該初始噪聲是從平均值為 0 且方差為 1 的正態(tài)分布中采樣的。在生成過程的每個步驟中，使用預測的平均值 μ 對變量 xtk-1 進行采樣。下標 k-1 表示上一個時間步。這意味著每步 x 的值取決于上一步中 x 的值以及分布的預測平均值。

上面這些流程均和 DDPM 相似，我們就不展開說明了。重點關注 Variance Schedule Resampling 的過程，也就是如何實現(xiàn)訓練的擴散過程減少到五個步驟的。可以理解為對方差值子序列進行采樣的過程。給出了方差表 {βt} Tt=1，子序列 {βk} Kk=1 可以用 {tk} Kk=1 進行采樣。簡而言之，在訓練或推理過程中，給出方差值序列，并對這些值的子序列進行采樣。子序列中的值是根據(jù)先前的值和重新計算的值計算的。目標是通過在訓練或推理期間調整方差值來優(yōu)化模型的性能。如果是在圖像去噪任務中進行方差的重采樣，一定會影響結果，但在分割任務中經(jīng)過驗證是有效的。

Diffusion Model for Segmentation

上一部分是對 DDPM 的方差重采樣，不涉及到圖像分割過程。對于分割任務上的優(yōu)化，ImgX-DiffSeg 可以依據(jù)時間步，使用預測噪聲和采樣噪聲之間的 L2 損失進行訓練。此外，ImgX-DiffSeg 計算預測掩碼和金標準之間定義的特定分割損失，例如 Dice Loss 或 CE Loss。

在訓練期間，現(xiàn)有方法通過插值噪聲和金標準對噪聲掩模進行采樣，這會導致數(shù)據(jù)信息丟失。為了解決這個問題，模型在上一個時間步中的預測被回收以取代金標準。回收的噪聲掩碼是使用下面方程計算的。

其中， x0theta 是使用金標準計算的上一個時間步的預測分割掩碼，xt 和 xt+1 是兩個獨立的采樣噪聲。梯度停止應用于 xt+1 以防止通過回收的噪聲掩碼進行反向傳播。αt 是超參數(shù)。第一個方程使用先前的預測和當前噪聲計算 xt，而第二個方程使用金標準和下一個噪聲計算 xt+1。

實驗

實驗分別基于 MRI 和 CT 圖像數(shù)據(jù)集， 值得注意的是，ImgX-DiffSeg 在 3D 上的表現(xiàn)是好于 2D 數(shù)據(jù)集的，如下表所示。

下圖是非擴散分割模型和擴散概率模型之間的可視化比較，其中 t 表示時間步，一直反向擴散到第一個時間步的效果最好。

請?zhí)砑訄D片描述

下表是四種消融實驗，分別是預測噪聲推斷分割圖和直接預測分割圖對比；損失函數(shù)對比；是否回收上一個時間步中預測的掩碼對比；訓練過程的時間步數(shù)量對比。5 個 steps 的效果優(yōu)于 1000 個 steps，說明 Variance Schedule Resampling 是有效果的。

請?zhí)砑訄D片描述

總結

ImgX-DiffSeg 是第一個用于 3D 圖像多類分割的 DDPM 模型，與現(xiàn)有的基于擴散的方法相比，該模型顯著提高了性能，但也沒有優(yōu)于普通的非擴散分割模型，還值得進一步改進。這篇文章的代碼目前已開源，我試了一下，訓練的收斂速度真的快，可以作為一個不錯的 benchmark。

參考

https://github.com/mathpluscode/ImgX-DiffSeg

https://arxiv.org/abs/2303.06040

審核編輯 ：李倩