[论文解读] SegDiff: Image Segmentation with Diffusion Probabilistic Models

Tomer Amit, Tal Shaharbany+2 more|arXiv (Cornell University)|Dec 1, 2021

Advanced Neural Network Applications参考文献 47被引用 13

一句话总结

本文提出 SegDiff，一种新型的扩散概率模型，用于端到端图像分割，无需预训练主干网络。它通过将输入图像的编码特征与当前分割预测的特征相加来条件化去噪过程，迭代地优化输出；该方法通过融合多组随机生成结果并进行模型平均，在 Cityscapes、Vaihiken 和 MoNuSeg 基准测试中实现了最先进性能。

ABSTRACT

Diffusion Probabilistic Methods are employed for state-of-the-art image generation. In this work, we present a method for extending such models for performing image segmentation. The method learns end-to-end, without relying on a pre-trained backbone. The information in the input image and in the current estimation of the segmentation map is merged by summing the output of two encoders. Additional encoding layers and a decoder are then used to iteratively refine the segmentation map, using a diffusion model. Since the diffusion model is probabilistic, it is applied multiple times, and the results are merged into a final segmentation map. The new method produces state-of-the-art results on the Cityscapes validation set, the Vaihingen building segmentation benchmark, and the MoNuSeg dataset.

研究动机与目标

• 将扩散概率模型从此前主要用于图像生成的任务，扩展至图像分割任务，其中真实标签是唯一且确定的。
• 通过从零开始端到端训练整个模型，消除对预训练主干网络的依赖。
• 开发一种新颖的条件化机制，通过特征相加而非拼接的方式，融合输入图像与当前分割估计的特征。
• 通过生成多组随机输出并取平均，提升模型校准性与准确性。
• 在多样化的分割基准上展示最先进性能，尤其在小样本数据集上表现突出。

提出的方法

• 模型采用 U-Net 架构作为去噪网络，并配备两个独立编码器：一个用于输入图像 I，另一个用于当前分割估计 xₜ。
• 两个编码器的特征在 U-Net 编码器路径的早期阶段即被相加，实现图像与分割信息的联合建模。
• 去噪网络利用相加后的特征预测下一状态 xₜ₋₁，通过 T 个扩散步骤迭代优化分割图。
• 模型采用变分推理框架进行训练，以最小化反向过程损失，从而优化去噪目标。
• 对每个输入执行多组独立生成，最终通过平均结果获得分割输出，以提升鲁棒性与准确性。
• 采用类似残差的条件化机制，即将输入图像编码后加到当前预测表示上，而非进行拼接。

实验结果

研究问题

• RQ1 扩散概率模型能否有效应用于图像分割任务（该任务具有唯一真实标签），而不仅限于生成任务（允许多个合理输出）？
• RQ2 在避免使用预训练主干网络的前提下，如何有效对输入图像进行条件化处理，以应用于扩散模型的分割任务？
• RQ3 与单一确定性预测相比，生成多组随机样本并取平均是否能提升分割准确率与模型校准性？
• RQ4 在分割扩散设置中，特征融合机制的选择（如相加与拼接）对性能有何影响？
• RQ5 在该新框架中，推理速度（扩散步数）与分割准确率之间的权衡关系如何？

主要发现

• SegDiff 在 Cityscapes 验证集上实现了最先进性能，mIoU 相较于先前方法有显著提升，尤其在 'bus' 等挑战性类别上表现突出。
• 在 Vaihingen 建筑物分割基准上，该方法显著优于现有方法，尤其当采用特征相加而非拼接作为条件化方式时优势更明显。
• 在 MoNuSeg 数据集上，模型在细胞核分割任务中也取得了优异结果，证明了其在多样化医学影像领域的泛化能力。
• 将扩散步数从 100 减少到 25 仅导致性能下降 1–2 mIoU，实现了高达四倍的推理速度提升，且准确率损失极小。
• RRDB 模块数量对性能影响有限，最优配置下最佳与最差设置之间的 mIoU 差异小于 1。
• 对九组生成样本进行平均能显著提升准确率与模型校准性，证实了在确定性任务中使用随机采样的有效性。