前言

在前面的一篇文章ICLR 2023：基于 diffusion adversarial representation learning 的血管分割中，我們已經(jīng)介紹過(guò)了 diffusion model 在醫(yī)學(xué)圖像分割上的一個(gè)應(yīng)用，推薦對(duì) diffusion model 不了解的同學(xué)優(yōu)先閱讀，其中講了一些基本概念。上一篇文章是將 diffusion 應(yīng)用到自監(jiān)督學(xué)習(xí)中，而 MedSegDiff 是一個(gè)有監(jiān)督的框架，現(xiàn)在已更新到 V2 版本， V2 版本區(qū)別于 V1 使用了 Transformer，且適用于多分類(lèi)。MedSegDiff-V1 已被接收在 MIDL 2023。

MedSegDiff

MedSegDiff 在原版 DPM 的基礎(chǔ)上引入了動(dòng)態(tài)條件編碼，增強(qiáng) DPM 在醫(yī)學(xué)圖像分割上的分步注意力能力。特征頻率解析器（FF-Parser）可以消除分割過(guò)程中損壞的給定掩碼中的高頻噪聲。DPM 是一種生成模型，由兩個(gè)階段組成，正向擴(kuò)散階段和反向擴(kuò)散階段。在正向過(guò)程中，通過(guò)一系列步驟 T，將高斯噪聲逐漸添加到分割標(biāo)簽 x0 中。在反向過(guò)程中，訓(xùn)練神經(jīng)網(wǎng)絡(luò)通過(guò)反向噪聲過(guò)程來(lái)恢復(fù)原始數(shù)據(jù)：

其中 theta 是反向過(guò)程參數(shù)。從高斯噪聲開(kāi)始，p(xT) 表示原始圖像，反向過(guò)程將潛在變量分布 p(xT) 轉(zhuǎn)換為數(shù)據(jù)分布 p(x0)。反向過(guò)程逐步恢復(fù)噪聲圖像，以獲得最終的清晰分割。該模型使用 U-Net 作為學(xué)習(xí)網(wǎng)絡(luò)，步長(zhǎng)估計(jì)函數(shù)由原始圖像先驗(yàn)條件確定：

其中 EI 是條件特征嵌入，即原始圖像嵌入，Ex 是當(dāng)前步驟的分割映射特征嵌入。這兩個(gè)組件被添加并發(fā)送到 U-Net 的解碼器進(jìn)行重建。步長(zhǎng)索引 t 與新增的嵌入和解碼器功能集成在一起，使用共享的 look-up table 進(jìn)行嵌入，這在 DDPM 的論文中有介紹?？偠灾?，MedSegDiff 模型基于 DPM，使用 U-Net 進(jìn)行學(xué)習(xí)。步長(zhǎng)估計(jì)函數(shù)由原始圖像先驗(yàn)得到，步長(zhǎng)索引與新增的嵌入和解碼器功能集成在一起。使 MedSegDiff 在三項(xiàng)具有不同圖像模式的醫(yī)學(xué)分割任務(wù)中表現(xiàn)不錯(cuò)。先看下 MedSegDiff 整體流程圖：

下面我們將逐一介紹動(dòng)態(tài)條件編碼和 FF-Parser。

動(dòng)態(tài)條件編碼

對(duì)于 MRI 或超聲之類(lèi)的低對(duì)比度圖像，很難將感興趣的對(duì)象與背景分開(kāi)。所以使用動(dòng)態(tài)條件編碼方法來(lái)解決這個(gè)問(wèn)題。可以注意到，原始圖像包含準(zhǔn)確的目標(biāo)分割信息，不過(guò)很難與背景區(qū)分開(kāi)，而當(dāng)前步驟的 grand truth 包含增強(qiáng)的目標(biāo)區(qū)域，但不準(zhǔn)確。

為了整合這兩個(gè)信息來(lái)源，使用類(lèi)似注意力的機(jī)制將條件特征圖的每個(gè)尺度與當(dāng)前步驟的編碼特征融合。這種融合是首先對(duì)兩個(gè)特征圖分別應(yīng)用層歸一化，然后將它們相乘以獲得 affine map 來(lái)實(shí)現(xiàn)的，再將 affine map 與條件編碼特征相乘以增強(qiáng)注意力區(qū)域。如 MedSegDiff 流程圖所示，此操作應(yīng)用于中間兩個(gè)階段，其中每個(gè)階段都是在 Resnet34 之后實(shí)現(xiàn)的卷積階段。但是，集成當(dāng)前條件編碼功能可能會(huì)產(chǎn)生額外的高頻噪聲。為了解決這個(gè)問(wèn)題，使用 FF-Parser 來(lái)限制特征中的高頻分量。

FF-Parser

FF-Parser 的流程如下圖所示：

FF-Parser 是頻率濾波器的可學(xué)習(xí)版本，它可以全局調(diào)整特定頻率的分量，以限制高頻分量進(jìn)行自適應(yīng)集成。首先使用二維 FFT（快速傅立葉變換）沿空間維度對(duì)解碼器特征圖 m 進(jìn)行變換，生成頻譜 M。然后，將參數(shù)化的注意力地圖 A 與 M 相乘以調(diào)整頻譜，得出 M'。最后，使用逆向 FFT 將 M' 反向回空間域，以獲得修改后的特征圖 m'。使用 FF-Parser 可以學(xué)習(xí)適用于傅里葉空間特征的權(quán)重圖，該權(quán)重圖可用于全局調(diào)整特定頻率的分量。這種技術(shù)不同于空間注意力，后者調(diào)整特定空間位置的組成部分。

實(shí)驗(yàn)

下圖分別是腦部 MRI、眼底視盤(pán)和甲狀腺結(jié)節(jié)的超聲圖像分割結(jié)果的可視化，可以看出 MedSegDiff 在簡(jiǎn)單解刨結(jié)構(gòu)的二分類(lèi)上效果還是不錯(cuò)的。

對(duì)比其他 SOTA 方法的結(jié)果如下表：

MedSegDiff-V2

MedsegDiff-v2 通過(guò)將 Transformer 機(jī)制整合到原始的U-Net骨干中，增強(qiáng)了基于擴(kuò)散的 MedSegDiff-v1。具體來(lái)說(shuō)，MedsegDiff-v2 引入了新的 Spectrum-Space Transformer（SS-former），對(duì)噪聲和語(yǔ)義特征之間的相互作用進(jìn)行建模。驗(yàn)證了 Medsegdiff-v2 對(duì)具有不同模態(tài)圖像的五個(gè)分割數(shù)據(jù)集的十八個(gè)器官的有效性。

概述

如下圖所示，MedsegDiff-v2 結(jié)合了錨點(diǎn)條件和語(yǔ)義條件兩種不同的條件方式，以提高擴(kuò)散模型的性能。錨點(diǎn)條件將錨分割特征（條件模型的解碼分割特征）集成到擴(kuò)散模型的編碼特征中。即允許使用粗略但靜態(tài)的參照來(lái)初始化擴(kuò)散模型，有助于減少擴(kuò)散方差。

然后將語(yǔ)義條件強(qiáng)加于擴(kuò)散模型的 embedding，理解為將條件模型的語(yǔ)義 embedding 集成到擴(kuò)散模型的 embedding 中。這種條件集成由 SS-former 實(shí)現(xiàn)，它彌合了噪聲和語(yǔ)義嵌入之間的鴻溝，并利用 Transformer 的全局和動(dòng)態(tài)特性抽象出更強(qiáng)的特征表達(dá)形式。

Medsegiff-v2 是使用 DPM 的標(biāo)準(zhǔn)噪聲預(yù)測(cè)損失 Lnoise 和錨損失 Lanchor 進(jìn)行訓(xùn)練的。Lanchor 是 Dice loss 和 CE loss 的組合。總損失函數(shù)表示為：

其中 t ≡ 0 (mod α) 通過(guò)超參數(shù) α 控制監(jiān)督條件模型的時(shí)間，β 是另一個(gè)用于加權(quán)交叉熵?fù)p失的經(jīng)驗(yàn)超參數(shù)?？偠灾琈edsegdiff-v2 顯著提高了 MedsegDiff 的性能。該方法采用了新的基于 Transformer 的條件 U-Net 框架和兩種不同的條件方式，以提高擴(kuò)散模型的性能。

Anchor Condition with Gaussian Spatial Attention

與卷積層相比，Transformer 具有更強(qiáng)的表示性，但對(duì)輸入方差更敏感。為了克服這種負(fù)面影響，所以使用了錨條件運(yùn)算，如上面概述中的介紹，該運(yùn)算將條件模型的解碼分割特征（錨點(diǎn)）集成到擴(kuò)散模型的編碼器特征中。此外，還使用了高斯空間注意力來(lái)表示條件模型中給定分割特征的不確定性（概率）。

在第一個(gè)公式中，表示在錨點(diǎn)特征上應(yīng)用高斯核以進(jìn)行平滑激活，因?yàn)殄^點(diǎn)可能不完全準(zhǔn)確，且高斯核的均值和方差是可以學(xué)習(xí)的。選擇平滑 ground truth 和原始圖之間的最大值以保留最相關(guān)的信息，從而生成平滑的錨特征。在第二個(gè)公式中，將平滑錨點(diǎn)特征集成到擴(kuò)散模型中以獲得增強(qiáng)特征。首先應(yīng)用 1x1 卷積將錨特征中的通道數(shù)減少到 1（經(jīng)常作用于解碼器的最后一層）。最后，在錨點(diǎn)特征上使用 sigmoid 激活函數(shù)，將其添加到擴(kuò)散模型的每個(gè)通道中，類(lèi)似于空間注意力的實(shí)現(xiàn)。

Semantic Condition with SS-Former

關(guān)于 SS-Former 的作用，我們只做簡(jiǎn)單的總結(jié)。對(duì)比 MedSegDiff-v1，是一種將條件模型分割 embedding 集成到擴(kuò)散模型 embedding 中的新架構(gòu)，其使用頻譜空間注意力機(jī)制來(lái)解決擴(kuò)散和分割 embedding 之間的域差距。此外，注意力機(jī)制在傅里葉空間中合并語(yǔ)義和噪聲信息，和 MedSegDiff-v1 是類(lèi)似的。

實(shí)驗(yàn)

Medsegiff-v2 是可以進(jìn)行多分類(lèi)任務(wù)的，只是目前 Github 還沒(méi)有維護(hù)多分類(lèi)的代碼，不過(guò) Medsegiff-v2 的實(shí)現(xiàn)代碼已經(jīng)提交了。下圖為在腹部 CT 圖像中的多器官分割上的表現(xiàn)。

對(duì)比其他 SOTA 方法的結(jié)果如下表：

總結(jié)

關(guān)于 diffusion model 在醫(yī)學(xué)圖像分割上的應(yīng)用，其是否能適應(yīng)多分類(lèi)且復(fù)雜的解刨結(jié)構(gòu)還需要進(jìn)一步探索，訓(xùn)練和推理時(shí)的效率也低于常規(guī)的有監(jiān)督神經(jīng)網(wǎng)絡(luò)，可以進(jìn)一步優(yōu)化。

審核編輯：劉清