拋棄編碼器-解碼器架構，用擴散模型做邊緣偵測效果更好，國防科大提出DiffusionEdge

目前的深度邊緣偵測網路通常採用編碼器-解碼器架構，其中包含上下採樣模組，以更好地提取多層次的特性。然而，這種結構限制了網路輸出準確且細緻的邊緣檢測結果。

針對這個問題，一篇 AAAI 2024 上的論文給出了新的解決方案。

• #論文主題：DiffusionEdge: Diffusion Probabilistic Model for Crisp Edge Detection
• #作者：葉雲帆（國防科技大學），徐凱（國防科技大學），黃雨行（國防科技大學），易任嬌（國防科技大學），蔡志平（國防科技大學）
• 論文連結: https://arxiv.org/abs/2401.02032
• 開源程式碼: https://github.com/GuHuangAI/DiffusionEdge

國防科技大學iGRAPE Lab提出了一種新的方法，用於二維邊緣偵測任務。此方法利用擴散機率模型，在學習迭代的去噪過程中產生邊緣結果圖。為了減少計算資源的消耗，該方法採用隱空間訓練網絡，並引入了不確定性蒸餾模組以優化性能。同時，此方法也採用解耦架構加速去噪過程，並引入自適應傅立葉濾波器來調整特徵。透過這些設計，該方法能夠以有限的資源進行穩定訓練，並以更少的增強策略預測出清晰準確的邊緣圖。實驗結果表明，該方法在四個公共基準資料集上的準確度和精度方面均顯著優於其他方法。

圖1 基於擴散機率模型的邊緣偵測過程與優勢範例

本文的創新點包括：

提出了針對邊緣偵測任務的擴散模型DiffusionEdge，該模型無需後處理即可預測出更細、更準確的邊緣圖。

為了解決擴散模型應用時的困難點，我們設計了多種技術來確保方法在隱空間中穩定學習。同時，我們也保留了像素級的不確定性先驗知識，並自適應地過濾傅立葉空間中的隱特徵。

3、在四個邊緣偵測公共基準資料集上進行的大量對比實驗展示了 DiffusionEdge 在準確度和細度方面均具有卓越的效能優勢。

相關工作

基於深度學習的方法通常採用包含上下取樣的編解碼結構整合多層特徵[1-2]，或是整合多個標註的不確定性資訊以提升邊緣檢測的準確度[3]。然而，自然受限於這樣的結構，其產生的邊緣結果圖對於下游任務來說太過粗厚而嚴重依賴後處理的問題仍然亟待解決。儘管許多工作已經在損失函數[4-5] 和標籤修正策略[6] 方面做出了探索以使網路能輸出更細的邊緣，但本文認為該領域仍然需要一個可以不借助任何額外模組，就能直接滿足準確度和細度的邊緣偵測器，而無需任何後處理步驟。

擴散模型是一類基於馬可夫鏈的生成模型，透過學習去雜訊過程逐漸恢復目標資料樣本。擴散模型在電腦視覺、自然語言處理和音訊生成等領域都表現出了卓越的性能。不僅如此，透過將影像或其他模態的輸入作為額外條件時，在感知任務中也表現出了巨大的潛力，例如影像分割[7]、目標偵測[8] 和姿態估計[9] 等。

方法描述

本文所提出的 DiffusionEdge 方法總體架構如圖 2 所示。受以往工作的啟發，該方法在隱空間中訓練具有解耦結構的擴散模型，並將圖像作為額外的條件線索輸入。此方法引入了自適應傅立葉濾波器來進行頻率解析，且為了保留來自多個標註者的像素級不確定性資訊並減少對計算資源的要求，還以蒸餾的方式直接使用交叉熵損失優化隱空間。

#

圖2 DiffusionEdge 的整體結構示意圖

針對目前的擴散模型受到取樣步數太多，推理時間太長等問題的困擾，本方法受DDM [10] 的啟發，同樣使用解耦的擴散模型架構來加速取樣推理過程。其中，解耦的前向擴散過程由顯式的轉移機率和標準Wiener 過程的組合來控制：

其中和分別代表初始邊緣和雜訊邊緣，指的是反向邊緣梯度的明確轉換函數。與DDM 類似，本文方法預設使用常數函數，則其對應的逆過程可以被表示為：

其中 。為了訓練解耦的擴散模型，此方法需要同時監督資料和雜訊分量，因此，訓練目標可以參數化為：

其中是去噪網路中的參數。由於擴散模型如果要在原始影像空間中訓練的話會佔用太多的計算成本，因此參考 [11] 的思路，本文提出的方法將訓練過程轉移到了具有 4 倍下採樣空間大小的隱空間中。

如圖2 所示，該方法首先訓練了一對自編碼器和解碼器的網絡，該編碼器將邊緣標註壓縮為一個隱變量，而解碼器則用於從這個隱變數中恢復出原來的邊緣標註。如此一來，在訓練基於U-Net 結構的去噪網絡階段，該方法便固定這一對自編碼和解碼器網絡的權重，並在隱空間中訓練去噪過程，這樣可以大幅減少網絡對計算資源的消耗，同時維持不錯的效能。

為了提升網路最後的效能，本文提出的方法在解耦操作中引入了一個可以自適應地過濾掉不同頻率特徵的模組。如圖2 左下角所示，該方法將自適應的快速傅立葉變換濾波器（Adaptive FFT-filter）整合到了去噪Unet 網路中，位置在解耦操作之前，以在頻域中自適應過濾並分離出邊緣圖和雜訊分量。具體來說，給定編碼器特徵  ，該方法首先沿著空間維度執行二維的傅立葉轉換（FFT），並將變換後的特徵表示為  。緊接著，為了訓練這個自適應頻譜濾波模組，建構了一個可學習的權重圖  並將其W乘以Fc。頻譜濾波器可以全域調整特定頻率，並且學習到的權重可以適應不同資料集目標分佈的不同頻率情況。透過自適應地濾除無用成分，本方法透過快速傅立葉逆變換（IFFT）操作將特徵從頻域映射回空間域。最後，透過額外引入了來自的殘差連接，避免完全過濾掉了所有有用的信息。上述過程可以由以下公式來描述：

#其中是輸出特徵，o表示哈達瑪積（Hadamard Product）。

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

#

其中,其中是決定真值標註中不確定的邊緣像素的閾值，如果像素值大於0 而小於此閾值，則此類模糊的，置信度不夠高的像素樣本將在後續的最佳化過程中被忽略（損失函數為0）。 和分別表示真值標註邊緣圖中邊緣和非邊緣像素的數量。是用來平衡和的權重（設為 1.1）。因此，每個邊緣圖的最終損失函數計算方式為。

在最佳化過程中忽略模糊的低置信度的像素可以避免網路混亂，使訓練過程更穩定的收斂，並提高模型的表現。然而，將二元交叉熵損失直接應用於在數值和空間上均未對齊的隱空間中是幾乎不可能的。尤其是，不確定性感知的交叉熵損失中利用閾值（一般從0 到1）來判斷像素是否為邊緣，這是從圖像空間上定義的，而隱變量遵循正態分佈並且具有完全不同的範圍和實際意義。此外，像素級的不確定性很難與不同大小的編碼和下採樣的隱特徵保持一致，二者是很難直接相容的。因此，將交叉熵損失直接應用於最佳化隱變量不可避免地會導致不正確的不確定性感知。

另一方面，人們可以選擇將隱變量解碼回圖像級別，從而可以直接使用不確定性感知交叉熵損失來直接監督預測出的邊緣結果圖。不幸的是，這種實現讓反向傳播的參數梯度通過了冗餘的自編碼器網絡，使得梯度很難得到有效的傳遞。此外，自編碼器網路中額外的梯度運算會帶來龐大的 GPU 記憶體消耗成本，這違背了本方法設計實用的邊緣偵測器的初衷，很難推廣到實際應用中。因此，本方法提出了不確定性蒸餾損失，可以直接優化隱空間上的梯度，具體來說，設重構出的隱變量為，自編碼器網絡的解碼器是D，解碼出的邊緣結果是eD，本方法考慮基於鍊式法則直接計算不確定性感知的二元交叉熵損失的梯度，具體計算方式為：

為了消除自編碼器網路的負面影響，本方法直接跳過了自編碼器傳遞梯度並將梯度的計算方式修改調整為：

這樣的實作方式大大降低了計算成本，並允許利用不確定性感知的損失函數直接在隱變數上最佳化。如此一來，再結合一個隨步數t自適應變化的時變損失權重，本方法最終訓練最佳化目標可以被表示為：

實驗結果

本方法在四個領域內被廣泛使用的邊緣檢測公共標準資料集上進行了實驗：BSDS、NYUDv2、Multicue 和BIPED。由於邊緣偵測資料標註比較難，標註資料量都比較少，以往的方法通常會使用各種策略來增強資料集。例如說，BSDS 中的影像透過水平翻轉 (2×)、縮放 (3×) 和旋轉 (16×) 進行增強，能夠產生比原始版本擴大了 96 倍的訓練集。以往的方法在其他資料集上使用的通用增強策略總結在了表格 1 中，其中 F 代表水平翻轉，S 代表縮放，R 代表旋轉，C 代表裁剪，G 代表伽馬校正。不同的是，本方法只需要使用隨機裁切的 320320 的圖像塊來訓練所有資料。在 BSDS 資料集中，本方法則僅採用隨機的翻轉和縮放，其定量對比結果展示在了表 2 中。在 NYUDv2、Mu lticue 和 BIPED 資料集中，此方法僅需採用隨機翻轉訓練。在使用較少增強策略的情況下，本方法在各個資料集，各個指標的表現都優於先前的方法。透過觀察圖3-5 的預測結果可以看出，DiffusionEdge 能夠學習並預測出跟gt 分佈幾乎一樣的邊緣檢測結果圖來，預測結果準確而清晰的優勢對於有精細化需求的下游任務來說非常重要，也展現了其能直接應用於後續任務的巨大潛力。

表1 過去方法在四個邊緣偵測資料集上所使用的增強策略

表2 不同方法在BSDS 資料集上的定量比較

#圖3 不同方法在BSDS 資料集上的定性對比

圖4 不同方法在NYUDv2 資料集上的定性對比

圖5 不同方法在BIPED 資料集上的定性對比

以上是拋棄編碼器-解碼器架構，用擴散模型做邊緣偵測效果更好，國防科大提出DiffusionEdge的詳細內容。更多資訊請關注PHP中文網其他相關文章！