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

前言

本篇文章繼續(xù)解讀醫(yī)學(xué)圖像 diffusion 系列，之前我們分別介紹過(guò)在自監(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)用，而是對(duì)訓(xùn)練和推理策略進(jìn)行優(yōu)化，并適應(yīng) 3D 的醫(yī)學(xué)圖像分割任務(wù)，參考鏈接在文末。

目前存在的問(wèn)題

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

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

模型過(guò)度依賴先前時(shí)間步中的信息。

ImgX-DiffSeg 架構(gòu)

概述

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

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

請(qǐng)?zhí)砑訄D片描述

DDPM with Variance Schedule Resampling

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

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

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

Diffusion Model for Segmentation

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

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

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

實(shí)驗(yàn)

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

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

請(qǐng)?zhí)砑訄D片描述

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

請(qǐng)?zhí)砑訄D片描述

總結(jié)

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

參考

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

https://arxiv.org/abs/2303.06040

審核編輯 ：李倩