[论文解读] Masked Diffusion as Self-supervised Representation Learner

Zixuan Pan, Jianxu Chen+1 more|arXiv (Cornell University)|Aug 10, 2023

Generative Adversarial Networks and Image Synthesis被引用 6

一句话总结

本文提出掩码扩散模型（MDM），一种用于语义分割的自监督表征学习方法，通过在扩散模型中用掩码替代高斯噪声，并采用结构相似性（SSIM）损失来对齐预训练与下游密集预测任务。MDM在少样本场景下表现尤为出色，实现了最先进性能，表明生成能力并非强表征学习的必要条件。

ABSTRACT

Denoising diffusion probabilistic models have recently demonstrated state-of-the-art generative performance and have been used as strong pixel-level representation learners. This paper decomposes the interrelation between the generative capability and representation learning ability inherent in diffusion models. We present the masked diffusion model (MDM), a scalable self-supervised representation learner for semantic segmentation, substituting the conventional additive Gaussian noise of traditional diffusion with a masking mechanism. Our proposed approach convincingly surpasses prior benchmarks, demonstrating remarkable advancements in both medical and natural image semantic segmentation tasks, particularly in few-shot scenarios.

研究动机与目标

• 解耦扩散模型的生成能力与其表征学习潜力，挑战去噪对强特征学习至关重要的假设。
• 解决预训练重建任务与下游密集预测任务之间的不匹配，特别是在轻量级模型的少样本设置下。
• 提出一种新型预训练范式——掩码扩散模型（MDM），用掩码替代噪声，实现可扩展的自监督学习。
• 通过将MSE损失替换为SSIM损失，改善预训练目标与分割性能之间的对齐。
• 通过实证验证MDM在医学图像和自然图像分割基准上优于DDPM和MAE。

提出的方法

• 在扩散模型中用基于时间步的动态掩码机制替代传统的加性高斯噪声，该机制基于时间步概率性地掩码图像块。
• 训练一个时间感知U-Net，从掩码输入中重建原始图像，采用具有T个时间步的扩散过程。
• 使用结构相似性（SSIM）损失替代均方误差（MSE）损失，以最小化预训练与下游分割任务之间的分布差异。
• 在每个时间步训练模型以预测原始图像（或掩码区域），并优化重建目标以实现语义一致性。
• 在DDPM变体中采用噪声预测策略，即模型学习预测污染过程中添加的噪声，从而提升表征质量。
• 通过在下游语义分割任务上微调并以mIoU为主要指标，评估学习到的表征。

实验结果

研究问题

• RQ1 能否将扩散模型的生成能力与其表征学习性能解耦？
• RQ2 在扩散模型中用掩码替代高斯噪声是否仍能生成用于分割的强语义表征？
• RQ3 与MSE损失相比，SSIM损失在对齐预训练目标与下游密集预测任务方面是否更有效？
• RQ4 MDM在少样本分割场景下的表现与DDPM和MAE相比如何？
• RQ5 重建目标的选择（噪声 vs. 图像）是否显著影响下游分割准确率？

主要发现

• 在少样本设置下，MDM在GlaS（84.51 ± 1.15）和FFHQ-34（59.18 ± 0.11）上达到最先进mIoU性能，优于DDPM和MAE。
• 采用SSIM损失时，MDM在GlaS上10%标注数据下的mIoU达到84.51 ± 1.15，显著优于DDPM（82.32 ± 0.77）和MAE（79.47 ± 1.08）。
• SSIM损失在MDM和DDPM上均提升性能，但在MAE上未见提升，表明其有效性取决于污染机制的上下文依赖性。
• 即使使用MSE损失，MDM（82.70 ± 0.79）也优于MAE（79.47 ± 1.08），而采用SSIM损失后性能进一步提升至84.51 ± 1.15。
• 即使在退化形式（如单时间步）下，MDM和DDPM仍优于MAE，表明扩散过程在表征学习方面的作用超越了简单的自编码。
• 训练为预测噪声（εt）的模型比预测干净图像（x0）的模型获得更高的下游准确率，与先前扩散模型研究结果一致。