拋棄編碼器-解碼器架構(gòu)，用擴(kuò)散模型做邊緣檢測(cè)效果更好，國(guó)防科大提出DiffusionEdge

作者：機(jī)器之心 2024-02-07 12:33:00

國(guó)防科技大學(xué) iGRAPE Lab 提出了首個(gè)用于二維邊緣檢測(cè)任務(wù)的擴(kuò)散概率模型方法，通過(guò)學(xué)習(xí)迭代的去噪過(guò)程獲得邊緣結(jié)果圖

現(xiàn)有的深度邊緣檢測(cè)網(wǎng)絡(luò)通常基于包含了上下采樣模塊的編碼器 - 解碼器架構(gòu)，以更好的提取多層次的特征，但這也限制了網(wǎng)絡(luò)輸出既準(zhǔn)又細(xì)的邊緣檢測(cè)結(jié)果。

針對(duì)這一問(wèn)題，一篇 AAAI 2024 上的論文給出了新的解決方案。

論文題目：DiffusionEdge: Diffusion Probabilistic Model for Crisp Edge Detection
作者：葉云帆（國(guó)防科技大學(xué)），徐凱（國(guó)防科技大學(xué)），黃雨行（國(guó)防科技大學(xué)），易任嬌（國(guó)防科技大學(xué)），蔡志平（國(guó)防科技大學(xué)）
論文鏈接: https://arxiv.org/abs/2401.02032
開(kāi)源代碼: https://github.com/GuHuangAI/DiffusionEdge

國(guó)防科技大學(xué) iGRAPE Lab 提出了首個(gè)用于二維邊緣檢測(cè)任務(wù)的擴(kuò)散概率模型方法，通過(guò)學(xué)習(xí)迭代的去噪過(guò)程獲得邊緣結(jié)果圖，去噪過(guò)程如圖 1 所示。為了在保留最終性能的同時(shí)減少計(jì)算資源的消耗，該方法在隱空間中訓(xùn)練網(wǎng)絡(luò)，并引入了不確定性蒸餾模塊以更好的優(yōu)化。同時(shí)，本方法還采用解耦架構(gòu)來(lái)加速去噪過(guò)程，并提出了相應(yīng)的自適應(yīng)傅立葉濾波器來(lái)調(diào)整特征?；谶@些設(shè)計(jì)，本方法得以用有限的資源進(jìn)行穩(wěn)定的訓(xùn)練，用更少的增強(qiáng)策略來(lái)預(yù)測(cè)清晰準(zhǔn)確的邊緣圖。在四個(gè)公共基準(zhǔn)數(shù)據(jù)集上的大量實(shí)驗(yàn)表明，本文提出的方法在準(zhǔn)確度和粗細(xì)度上全面超越了其他方法。

圖 1 基于擴(kuò)散概率模型的邊緣檢測(cè)過(guò)程與優(yōu)勢(shì)示例

本文的創(chuàng)新點(diǎn)包括：

1、提出了領(lǐng)域內(nèi)第一個(gè)針對(duì)邊緣檢測(cè)任務(wù)的擴(kuò)散模型 DiffusionEdge，無(wú)需任何后處理即可預(yù)測(cè)出更細(xì)更準(zhǔn)確的邊緣圖。

2、針對(duì)擴(kuò)散模型應(yīng)用時(shí)的難點(diǎn)，設(shè)計(jì)了多種技術(shù)以確保方法能在隱空間中穩(wěn)定學(xué)習(xí)，同時(shí)保留像素級(jí)的不確定性先驗(yàn)知識(shí)并自適應(yīng)地過(guò)濾傅立葉空間中的隱特征。

3、在四個(gè)邊緣檢測(cè)公共基準(zhǔn)數(shù)據(jù)集上開(kāi)展的大量對(duì)比實(shí)驗(yàn)展示了 DiffusionEdge 在準(zhǔn)確度和細(xì)度方面均具有卓越的性能優(yōu)勢(shì)。

方法描述

本文所提出的 DiffusionEdge 方法總體框架如圖 2 所示。受以往工作的啟發(fā)，該方法在隱空間中訓(xùn)練具有解耦結(jié)構(gòu)的擴(kuò)散模型，并將圖像作為額外的條件線(xiàn)索輸入。該方法引入了自適應(yīng)傅里葉濾波器來(lái)進(jìn)行頻率解析，且為了保留來(lái)自多個(gè)標(biāo)注者的像素級(jí)不確定性信息并減少對(duì)計(jì)算資源的要求，還以蒸餾的方式直接使用交叉熵?fù)p失優(yōu)化隱空間。

圖 2 DiffusionEdge 的整體結(jié)構(gòu)示意

針對(duì)目前的擴(kuò)散模型受到采樣步數(shù)太多，推理時(shí)間太長(zhǎng)等問(wèn)題的困擾，本方法受 DDM [10] 的啟發(fā)，同樣使用解耦的擴(kuò)散模型架構(gòu)來(lái)加速采樣推理過(guò)程。其中，解耦的前向擴(kuò)散過(guò)程由顯式的轉(zhuǎn)移概率和標(biāo)準(zhǔn) Wiener 過(guò)程的組合來(lái)控制：

其中和分別代表初始邊緣和噪聲邊緣，指的是反向邊緣梯度的顯式轉(zhuǎn)換函數(shù)。與 DDM 類(lèi)似，本文方法默認(rèn)使用常量函數(shù)，則其相應(yīng)的逆過(guò)程可以被表示為：

其中。為了訓(xùn)練解耦的擴(kuò)散模型，該方法需要同時(shí)監(jiān)督數(shù)據(jù)和噪聲分量，因此，訓(xùn)練目標(biāo)可以被參數(shù)化為：

其中是去噪網(wǎng)絡(luò)中的參數(shù)。由于擴(kuò)散模型如果要在原始圖像空間中訓(xùn)練的話(huà)會(huì)占用太多的計(jì)算成本，因此參考 [11] 的思路，本文提出的方法將訓(xùn)練過(guò)程轉(zhuǎn)移到了具有 4 倍下采樣空間大小的隱空間中。

如圖 2 所示，該方法首先訓(xùn)練了一對(duì)自編碼器和解碼器的網(wǎng)絡(luò)，該編碼器將邊緣標(biāo)注壓縮為一個(gè)隱變量，而解碼器則用于從這個(gè)隱變量中恢復(fù)出原來(lái)的邊緣標(biāo)注。如此一來(lái)，在訓(xùn)練基于 U-Net 結(jié)構(gòu)的去噪網(wǎng)絡(luò)階段，該方法便固定這一對(duì)自編碼和解碼器網(wǎng)絡(luò)的權(quán)重，并在隱空間中訓(xùn)練去噪過(guò)程，這樣可以大幅減少網(wǎng)絡(luò)對(duì)計(jì)算資源的消耗，同時(shí)維持不錯(cuò)的性能。

為了提升網(wǎng)絡(luò)最后的性能，本文提出的方法在解耦操作中引入了一種可以自適應(yīng)地過(guò)濾掉不同頻率特征的模塊。如圖 2 左下角所示，該方法將自適應(yīng)的快速傅立葉變換濾波器（Adaptive FFT-filter）集成到了去噪 Unet 網(wǎng)絡(luò)中，位置在解耦操作之前，以在頻域中自適應(yīng)過(guò)濾并分離出邊緣圖和噪聲分量。具體來(lái)說(shuō)，給定編碼器特征，該方法首先沿空間維度執(zhí)行二維的傅里葉變換（FFT），并將變換后的特征表示為。緊接著，為了訓(xùn)練這個(gè)自適應(yīng)頻譜濾波模塊，構(gòu)造了一個(gè)可學(xué)習(xí)的權(quán)重圖并將其W乘以Fc。頻譜濾波器可以全局地調(diào)整特定頻率，并且學(xué)習(xí)到的權(quán)重可以適應(yīng)不同數(shù)據(jù)集目標(biāo)分布的不同頻率情況。通過(guò)自適應(yīng)地濾除無(wú)用成分，本方法通過(guò)快速傅里葉逆變換（IFFT）操作將特征從頻域映射回空間域。最后，通過(guò)額外引入了來(lái)自的殘差連接，避免完全過(guò)濾掉了所有有用的信息。上述過(guò)程可以由以下公式來(lái)描述：

其中是輸出特征，o表示哈達(dá)瑪積（Hadamard Product）。

由于邊緣和非邊緣像素的數(shù)量高度不平衡（大多數(shù)像素都是非邊緣的背景），參考以往的工作，我們同樣引入不確定性感知的損失函數(shù)進(jìn)行訓(xùn)練。具體來(lái)說(shuō)，將作為第i個(gè)像素的真值邊緣概率，對(duì)于第j個(gè)邊緣圖中的第i個(gè)像素，其值為，則不確定性感知 WCE 損失的計(jì)算方式如下：

其中,其中是決定真值標(biāo)注中不確定的邊緣像素的閾值，如果像素值大于 0 而小于此閾值，則此類(lèi)模糊的，置信度不夠高的像素樣本將在后續(xù)的優(yōu)化過(guò)程中被忽略（損失函數(shù)為 0）。和分別表示真值標(biāo)注邊緣圖中邊緣和非邊緣像素的數(shù)量。是用來(lái)平衡和的權(quán)重（設(shè)為 1.1）。因此，每個(gè)邊緣圖的最終損失函數(shù)計(jì)算方式為。

在優(yōu)化過(guò)程中忽略模糊的低置信度的像素可以避免網(wǎng)絡(luò)混亂，使訓(xùn)練過(guò)程更穩(wěn)定的收斂，并提高模型的表現(xiàn)。然而，將二元交叉熵?fù)p失直接應(yīng)用于在數(shù)值和空間上均未對(duì)齊的隱空間中是幾乎不可能的。尤其是，不確定性感知的交叉熵?fù)p失中利用閾值（一般從 0 到 1）來(lái)判斷像素是否為邊緣，這是從圖像空間上定義的，而隱變量遵循正態(tài)分布并且具有完全不同的范圍和實(shí)際意義。此外，像素級(jí)的不確定性很難與不同大小的編碼和下采樣的隱特征保持一致，二者是很難直接兼容的。因此，將交叉熵?fù)p失直接應(yīng)用于優(yōu)化隱變量不可避免地會(huì)導(dǎo)致不正確的不確定性感知。

另一方面，人們可以選擇將隱變量解碼回圖像級(jí)別，從而可以直接使用不確定性感知交叉熵?fù)p失來(lái)直接監(jiān)督預(yù)測(cè)出的邊緣結(jié)果圖。不幸的是，這種實(shí)現(xiàn)讓反向傳播的參數(shù)梯度通過(guò)了冗余的自編碼器網(wǎng)絡(luò)，使得梯度很難得到有效的傳遞。此外，自編碼器網(wǎng)絡(luò)中額外的梯度計(jì)算會(huì)帶來(lái)巨大的 GPU 內(nèi)存消耗成本，這違背了本方法設(shè)計(jì)一個(gè)實(shí)用的邊緣檢測(cè)器的初衷，很難推廣到實(shí)際應(yīng)用中。因此，本方法提出了不確定性蒸餾損失，可以直接優(yōu)化隱空間上的梯度，具體來(lái)說(shuō)，設(shè)重構(gòu)出的隱變量為，自編碼器網(wǎng)絡(luò)的解碼器是D，解碼出的邊緣結(jié)果是eD，本方法考慮基于鏈?zhǔn)椒▌t直接計(jì)算不確定性感知的二元交叉熵?fù)p失的梯度，具體計(jì)算方式為：

為了消除自編碼器網(wǎng)絡(luò)的負(fù)面影響，本方法直接跳過(guò)了自編碼器傳遞梯度并將梯度的計(jì)算方式修改調(diào)整為：

這樣的實(shí)現(xiàn)方式大大降低了計(jì)算成本，并允許利用不確定性感知的損失函數(shù)直接在隱變量上優(yōu)化。如此一來(lái)，再結(jié)合一個(gè)隨步數(shù)t自適應(yīng)變化的時(shí)變損失權(quán)重，本方法最終訓(xùn)練優(yōu)化目標(biāo)可以被表示為：

實(shí)驗(yàn)結(jié)果

本方法在四個(gè)領(lǐng)域內(nèi)被廣泛使用的邊緣檢測(cè)公共標(biāo)準(zhǔn)數(shù)據(jù)集上進(jìn)行了實(shí)驗(yàn)：BSDS、NYUDv2、Multicue 和 BIPED。由于邊緣檢測(cè)數(shù)據(jù)標(biāo)注比較難，標(biāo)注數(shù)據(jù)量都比較少，以往的方法通常會(huì)使用各種策略來(lái)增強(qiáng)數(shù)據(jù)集。例如說(shuō)，BSDS 中的圖像通過(guò)水平翻轉(zhuǎn) (2×)、縮放 (3×) 和旋轉(zhuǎn) (16×) 進(jìn)行增強(qiáng)，能夠生成比原始版本擴(kuò)大了 96 倍的訓(xùn)練集。以往的方法在其他數(shù)據(jù)集上使用的通用增強(qiáng)策略總結(jié)在了表格 1 中，其中 F 代表水平翻轉(zhuǎn)，S 代表縮放，R 代表旋轉(zhuǎn)，C 代表裁剪，G 代表伽馬校正。不同的是，本方法僅需要使用隨機(jī)裁剪的 320320 的圖像塊來(lái)訓(xùn)練所有數(shù)據(jù)。在 BSDS 數(shù)據(jù)集中，本方法則僅僅采用隨機(jī)的翻轉(zhuǎn)和縮放，其定量對(duì)比結(jié)果展示在了表 2 中。在 NYUDv2、Mu lticue 和 BIPED 數(shù)據(jù)集中，該方法僅需采用隨機(jī)翻轉(zhuǎn)訓(xùn)練。在使用更少增強(qiáng)策略的情況下，本方法在各個(gè)數(shù)據(jù)集，各個(gè)指標(biāo)上的表現(xiàn)都優(yōu)于之前的方法。通過(guò)觀察圖 3-5 的預(yù)測(cè)結(jié)果可以看出，DiffusionEdge 能夠?qū)W習(xí)并預(yù)測(cè)出跟 gt 分布幾乎一樣的邊緣檢測(cè)結(jié)果圖來(lái)，預(yù)測(cè)結(jié)果準(zhǔn)確而清晰的優(yōu)勢(shì)對(duì)于有精細(xì)化需求的下游任務(wù)來(lái)說(shuō)非常重要，也展示了其能直接應(yīng)用于后續(xù)任務(wù)的巨大潛力。