抛弃编码器-解码器架构，用扩散模型做边缘检测效果更好，国防科大提出DiffusionEdge-51CTO.COM

现有的深度边缘检测网络通常基于包含了上下采样模块的编码器 - 解码器架构，以更好的提取多层次的特征，但这也限制了网络输出既准又细的边缘检测结果。

针对这一问题，一篇 AAAI 2024 上的论文给出了新的解决方案。

论文题目：DiffusionEdge: Diffusion Probabilistic Model for Crisp Edge Detection
作者：叶云帆（国防科技大学），徐凯（国防科技大学），黄雨行（国防科技大学），易任娇（国防科技大学），蔡志平（国防科技大学）
论文链接: https://arxiv.org/abs/2401.02032
开源代码: https://github.com/GuHuangAI/DiffusionEdge

国防科技大学 iGRAPE Lab 提出了首个用于二维边缘检测任务的扩散概率模型方法，通过学习迭代的去噪过程获得边缘结果图，去噪过程如图 1 所示。为了在保留最终性能的同时减少计算资源的消耗，该方法在隐空间中训练网络，并引入了不确定性蒸馏模块以更好的优化。同时，本方法还采用解耦架构来加速去噪过程，并提出了相应的自适应傅立叶滤波器来调整特征。基于这些设计，本方法得以用有限的资源进行稳定的训练，用更少的增强策略来预测清晰准确的边缘图。在四个公共基准数据集上的大量实验表明，本文提出的方法在准确度和粗细度上全面超越了其他方法。

图 1 基于扩散概率模型的边缘检测过程与优势示例

本文的创新点包括：

1、提出了领域内第一个针对边缘检测任务的扩散模型 DiffusionEdge，无需任何后处理即可预测出更细更准确的边缘图。

2、针对扩散模型应用时的难点，设计了多种技术以确保方法能在隐空间中稳定学习，同时保留像素级的不确定性先验知识并自适应地过滤傅立叶空间中的隐特征。

3、在四个边缘检测公共基准数据集上开展的大量对比实验展示了 DiffusionEdge 在准确度和细度方面均具有卓越的性能优势。

方法描述

本文所提出的 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 分布几乎一样的边缘检测结果图来，预测结果准确而清晰的优势对于有精细化需求的下游任务来说非常重要，也展示了其能直接应用于后续任务的巨大潜力。