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

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

前言

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

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

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

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

目前存在的問題

目前帶有 diffusion model 的架構(gòu)訓(xùn)練和推理耗時。

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

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

ImgX-DiffSeg 架構(gòu)

概述

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

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

請?zhí)砑訄D片描述

DDPM with Variance Schedule Resampling

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

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

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

Diffusion Model for Segmentation

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

在訓(xùn)練期間，現(xiàn)有方法通過插值噪聲和金標(biāo)準(zhǔn)對噪聲掩模進(jìn)行采樣，這會導(dǎo)致數(shù)據(jù)信息丟失。為了解決這個問題，模型在上一個時間步中的預(yù)測被回收以取代金標(biāo)準(zhǔn)?；厥盏脑肼曆诖a是使用下面方程計算的。

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

實驗

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

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

請?zhí)砑訄D片描述

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

請?zhí)砑訄D片描述

總結(jié)

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

參考

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

https://arxiv.org/abs/2303.06040

審核編輯 ：李倩