基于Discrete Diffusion的模型不可知分割細(xì)化

這次這篇文章介紹一篇很有意思的工作：SegRefiner，來(lái)自 NeurIPS, 2023，目前代碼已開(kāi)源。

SegRefiner 提出一種新的的任務(wù)解釋?zhuān)瑢⒎指罴?xì)化視為一個(gè)數(shù)據(jù)生成過(guò)程。因此，細(xì)化可以通過(guò)一系列去噪擴(kuò)散步驟來(lái)實(shí)現(xiàn)，其中 Coarse Mask 是 Ground Truth 的噪聲版本。此外，為了處理二值掩模，進(jìn)一步設(shè)計(jì)了一種新穎的離散擴(kuò)散過(guò)程，在該過(guò)程中，每個(gè)像素執(zhí)行單向隨機(jī)狀態(tài)轉(zhuǎn)換。所提出的過(guò)程可以在訓(xùn)練期間逐漸將 Ground Truth 轉(zhuǎn)換為 Coarse Mask，并在推理期間使用 Coarse Mask 作為采樣起點(diǎn)。換句話(huà)說(shuō)，SegRefiner 將掩模細(xì)化任務(wù)表述為一個(gè)條件生成問(wèn)題，其中輸入圖像（指原圖）作為條件，用于迭代更新/細(xì)化 Coarse Mask 中的錯(cuò)誤預(yù)測(cè)。

SegRefiner 是模型不可知的，因此適用于不同的分割模型和任務(wù)。SegRefiner 驗(yàn)證的任務(wù)包括：語(yǔ)義分割，實(shí)例分割和二分圖像分割。

相關(guān)工作

分割細(xì)化的目標(biāo)是提高現(xiàn)有分割模型中掩模的質(zhì)量。一些工作專(zhuān)注于增強(qiáng)特定的分割模型。還有一些模型不可知方法的細(xì)化方法，這些策略強(qiáng)調(diào)使用多種形式的輸入，包括整個(gè)圖像、邊界補(bǔ)丁和邊緣條帶等。盡管這些技術(shù)可以細(xì)化來(lái)自不同模型的粗糙掩模，但它們的適用性仍然局限于特定的分割任務(wù)。

擴(kuò)散模型在檢測(cè)和分割任務(wù)中的應(yīng)用也成為越來(lái)越多研究的焦點(diǎn)，這些研究主要遵循 DDPM 的高斯擴(kuò)散過(guò)程，并利用額外的圖像編碼器來(lái)提取圖像特征作為生成掩模的條件。SegRefiner 是第一個(gè)將擴(kuò)散模型應(yīng)用于圖像分割細(xì)化任務(wù)的工作，它還在基于擴(kuò)散的分割任務(wù)中首次放棄連續(xù)的高斯假設(shè)，轉(zhuǎn)而采用新設(shè)計(jì)的離散擴(kuò)散過(guò)程。

Forward diffusion process

在介紹前向過(guò)程之前，先看一下整體框架的實(shí)現(xiàn)：

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

右側(cè)是提出的 transition sample 模塊，它基于輸入狀態(tài)轉(zhuǎn)換概率從當(dāng)前掩模中隨機(jī)采樣像素，并改變它們的值以匹配目標(biāo)掩模中的值（也就是完成我們?cè)谝灾刑岬降摹懊總€(gè)像素執(zhí)行單向隨機(jī)狀態(tài)轉(zhuǎn)換”）?？催@張框架圖的時(shí)候，注意區(qū)分不同 pipeline 的顏色區(qū)分，這里吐槽一下，我第一看 transition sample 沒(méi)有清楚的看到三種輸入。

在訓(xùn)練期間，轉(zhuǎn)換樣本模塊將 Ground Truth 轉(zhuǎn)換為 Coarse Mask，因此 Coarse Mask 是目標(biāo)掩模。在推理過(guò)程中，目標(biāo)掩模指的是預(yù)測(cè)的細(xì)致掩模，這個(gè)模塊根據(jù)預(yù)測(cè)的細(xì)致掩模和轉(zhuǎn)換概率在每個(gè)時(shí)間步中更新 Coarse Mask 中的值。

接下來(lái)仔細(xì)介紹下前向過(guò)程：

這種方法涉及將一個(gè)精細(xì)的 mask 逐漸轉(zhuǎn)變?yōu)橐粋€(gè)更粗糙的 mask，包含以下關(guān)鍵要素：

起始狀態(tài) (mo)：與 ground truth 的精細(xì) mask 相對(duì)應(yīng)。

目標(biāo)狀態(tài) (mr)：一個(gè)粗糙的 mask。

中間狀態(tài) (mt)：隨著時(shí)間步 t 的增加，逐漸從 mo 向 mr 演變的中間狀態(tài)。

轉(zhuǎn)移采樣模塊

引入了“轉(zhuǎn)移采樣”模塊，該模塊負(fù)責(zé)根據(jù)當(dāng)前 mask mt、粗 mask mr 以及狀態(tài)轉(zhuǎn)移概率來(lái)進(jìn)行狀態(tài)的轉(zhuǎn)移。狀態(tài)的轉(zhuǎn)移是單向的，保證了最終會(huì)收斂到粗糙 mask mr。

重參數(shù)技巧

此外，SegRefiner 還提到了重參數(shù)技巧（reparameterization trick），通過(guò)引入二元隨機(jī)變量 x 來(lái)描述這一過(guò)程。該過(guò)程允許直接獲取任何中間時(shí)間步的 mask mt，而無(wú)需逐步采樣。具體表述為：

定義 x 表示為一個(gè) one-hot 向量，表示中間掩模 mt 中像素 (i, j) 的狀態(tài)。

設(shè)置 xi = [1, 0] 和 xi = [0, 1] 分別表示精細(xì)狀態(tài)和粗糙狀態(tài)。

因此，前向過(guò)程可以表示為公式：

where

為超參數(shù)，而 對(duì)應(yīng)了上述的狀態(tài)轉(zhuǎn)移概率。狀態(tài)轉(zhuǎn)移矩陣 的邊緣分布可以表示為：

整體上還是遵循 DDPM，但是引入 二元隨機(jī)變量 x 。

Reverse diffusion process

將粗糙 mask mr 逐漸修正為精細(xì) mask mo。

由于此時(shí)精細(xì) mask mo 和狀態(tài)轉(zhuǎn)移概率未知，對(duì)照著 SegRefiner 框架圖來(lái)看，訓(xùn)練一個(gè)神經(jīng)網(wǎng)絡(luò) 來(lái)預(yù)測(cè)精細(xì) ，表示為：

其中 I 是相應(yīng)的圖像。

和 分別表示預(yù)測(cè)的精細(xì) mask 和其置信度分?jǐn)?shù)。相應(yīng)的， 可視作 中每個(gè)像素處于“精細(xì)狀態(tài)”的概率。

反向狀態(tài)轉(zhuǎn)移概率

根據(jù)前向過(guò)程的設(shè)定和貝葉斯定理，延續(xù) DDPM 的方法，我們可以由前向過(guò)程的后驗(yàn)概率和預(yù)測(cè)的 得到反向過(guò)程的概率分布，表示為：

where

其中 為反向過(guò)程的狀態(tài)轉(zhuǎn)移概率。

迭代修正過(guò)程

給定粗糙 mask mr 以及相應(yīng)的圖像 I，首先將所有像素初始化為粗糙狀態(tài) xi = [1, 0]。通過(guò)不斷迭代地狀態(tài)轉(zhuǎn)移，逐漸修正 m_T 中的預(yù)測(cè)值。

推理過(guò)程

給定一個(gè)粗糙的掩模和其對(duì)應(yīng)的圖像，我們首先初始化所有像素為粗糙狀態(tài)。我們?cè)谝韵虏襟E之間迭代：

前向過(guò)程：以獲取 和 。

計(jì)算反向狀態(tài)轉(zhuǎn)移矩陣： 并得到 。

計(jì)算精細(xì)化的掩模：基于 ， 和 計(jì)算精細(xì)化的掩模 。

這個(gè)過(guò)程（1）-（3）迭代進(jìn)行，直到獲得精細(xì)的掩模。

實(shí)驗(yàn)

分別訓(xùn)練了 LR-SegRefiner 和 HR-SegRefiner，數(shù)據(jù)集和具體的 settings 在上 。

定性對(duì)比其他方法，覺(jué)得對(duì)比 U-Net 和 ISNet 的效果的確很明顯。

總結(jié)

SegRefiner 是首個(gè)基于擴(kuò)散的圖像分割細(xì)化方法，采用了離散擴(kuò)散過(guò)程。SegRefiner 執(zhí)行模型不可知的分割細(xì)化，并在各種分割任務(wù)的細(xì)化中取得了強(qiáng)有力的實(shí)證結(jié)果。雖然它在準(zhǔn)確度上取得了顯著的提升，但其局限性在于擴(kuò)散過(guò)程由于多步迭代策略而導(dǎo)致推理速度變慢。