[论文解读] DFormer: Diffusion-guided Transformer for Universal Image Segmentation

Hefeng Wang, Jiale Cao+4 more|arXiv (Cornell University)|Jun 6, 2023

Medical Image Segmentation Techniques被引用 7

一句话总结

DFormer 通过将任务建模为去噪过程，提出了一种基于扩散模型的 Transformer 架构用于通用图像分割：其通过深度特征和掩码注意力机制，训练 Transformer 从噪声掩码中重建干净掩码。该方法实现了最先进性能，在 COCO 全景分割任务上优于先前基于扩散的方法 3.6%，在 ADE20K 语义分割任务上优于 2.2%。

ABSTRACT

This paper introduces an approach, named DFormer, for universal image segmentation. The proposed DFormer views universal image segmentation task as a denoising process using a diffusion model. DFormer first adds various levels of Gaussian noise to ground-truth masks, and then learns a model to predict denoising masks from corrupted masks. Specifically, we take deep pixel-level features along with the noisy masks as inputs to generate mask features and attention masks, employing diffusion-based decoder to perform mask prediction gradually. At inference, our DFormer directly predicts the masks and corresponding categories from a set of randomly-generated masks. Extensive experiments reveal the merits of our proposed contributions on different image segmentation tasks: panoptic segmentation, instance segmentation, and semantic segmentation. Our DFormer outperforms the recent diffusion-based panoptic segmentation method Pix2Seq-D with a gain of 3.6% on MS COCO val2017 set. Further, DFormer achieves promising semantic segmentation performance outperforming the recent diffusion-based method by 2.2% on ADE20K val set. Our source code and models will be publicly on https://github.com/cp3wan/DFormer

研究动机与目标

• 通过统一架构解决在语义、实例和全景分割等多任务上泛化分割模型的挑战。
• 克服现有基于扩散的分割方法性能低于非扩散最先进方法的局限性。
• 探索在统一分割框架中使用扩散模型作为掩码预测生成先验的可行性。
• 提升对小物体和遮挡实例的性能，这些在当前分割模型中仍具挑战性。

提出的方法

• 通过向真实标签掩码添加多层级高斯噪声来训练模型，将分割任务转化为去噪问题。
• 使用深层像素级特征和噪声掩码作为输入，通过基于扩散的解码器生成掩码特征和注意力掩码。
• 在 Transformer 解码器中使用掩码自注意力机制，逐步从噪声输入预测干净掩码。
• 在推理阶段，生成一组随机噪声掩码，并使用相同的基于扩散的解码器预测最终掩码及其对应类别。
• 对噪声掩码应用二值阈值处理以生成注意力掩码，使模型在预测过程中能够聚焦于相关区域。
• 采用可学习的掩码编码策略并结合尺度控制（例如二值化或打乱编码），以提升掩码表征学习能力。

实验结果

研究问题

• RQ1 基于扩散的生成过程能否有效适配到单一 Transformer 框架中，以统一多种图像分割任务？
• RQ2 与直接预测掩码相比，使用高斯噪声对掩码进行去噪范式在分割精度和泛化能力方面表现如何？
• RQ3 在不同分割任务中实现高性能的最优解码器层数和推理时间步数是多少？
• RQ4 不同主干网络（如 ResNet-50 与 Swin-T）以及不同掩码编码策略对模型在全景、实例和语义分割任务上的性能有何影响？

主要发现

• 使用 ResNet-50 主干网络时，DFormer 在 MS COCO val2017 上取得 51.1% 的全景质量（PQ）分数，优于 Pix2Seq-𝒟 3.6 个百分点。
• 在 ADE20K val 数据集上，使用 Swin-T 主干网络时，DFormer 取得 48.3% 的平均交并比（mIoU），优于近期基于扩散的方法 DDP 2.2 个百分点。
• 采用 9 层解码器时，在实例分割任务上取得最佳性能（AP = 42.6），表明更深的解码器可提升掩码预测质量。
• 使用单步前向过程进行推理即可达到近似最优性能，表明模型即使在极少去噪步骤下也能良好泛化。
• 与随机打乱相比，采用尺度为 0.1 的二值掩码编码策略表现更优，尤其在低尺度下优势更明显。
• DFormer 在视频实例分割任务中也表现出良好泛化能力，在 YouTube-VIS 2019 val 数据集上取得 46.5% 的 AP，证明其在静态图像之外的可迁移性。