[논문 리뷰] Masked Diffusion as Self-supervised Representation Learner

Zixuan Pan, Jianxu Chen외 1명|arXiv (Cornell University)|2023. 08. 10.

Generative Adversarial Networks and Image Synthesis인용 수 6

한 줄 요약

이 논문은 확산 모델에서 가우시안 노이즈를 마스킹으로 대체함으로써, 의미 분할을 위한 자기지도 학습 표현 학습기인 마스킹 확산 모델(Masked Diffusion Model, 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.

연구 동기 및 목표

• 확산 모델의 생성 능력과 표현 학습 잠재력 간의 분리(디커플링)를 시도하여, 노이즈 제거가 강력한 특징을 얻는 데 필수적이라는 가정에 도전한다.
• 가벼운 모델을 사용한 소수의 샘플 설정에서, 복원 사전 훈련과 후행 밀도 예측 작업 간의 불일치 문제를 해결한다.
• 노이즈를 마스킹으로 대체하는 새로운 사전 훈련 프레임워크인 마스킹 확산 모델(Masked Diffusion Model, MDM)을 제안한다.
• MSE 손실을 SSIM 손실으로 교체하여 사전 훈련 목표와 분할 성능 간의 일치를 향상시킨다.
• 의료 영상 및 자연 영상 분할 벤치마크에서 DDPM과 MAE를 포함한 다양한 방법들과의 실험적 비교를 통해 MDM의 우수성을 검증한다.

제안 방법

• 시간에 따라 이미지 패치를 확률적으로 마스킹하는 동적 마스킹 메커니즘을 사용하여, 기존의 가우시안 노이즈를 대체한다.
• T개의 시간 단계를 가진 확산 과정을 사용하여, 마스킹된 입력에서 원본 이미지를 재구성하는 시간 인식 U-Net을 훈련한다.
• 사전 훈련과 후행 분할 작업 간의 분포 불일치를 줄이기 위해 평균 제곱 오차(MSE) 대신 구조적 유사도(SSIM) 손실을 사용한다.
• 각 시간 단계에서 원본 이미지(또는 마스킹된 영역)를 재구성하도록 모델을 훈련하며, 의미 일관성을 고려한 재구성 목표를 최적화한다.
• DDPM 변종에서 노이즈 예측 전략을 도입하여, 손상 과정에서 추가된 노이즈를 예측하도록 모델을 학습함으로써 표현 품질을 향상시킨다.
• mIoU를 주요 지표로 사용하여 후행 의미 분할 작업에서의 미세 조정을 통해 학습된 표현을 평가한다.

실험 결과

연구 질문

• RQ1 확산 모델의 생성 능력은 표현 학습 성능와 분리될 수 있는가?
• RQ2 확산 모델에서 가우시안 노이즈를 마스킹으로 대체해도 여전히 분할을 위한 강력한 의미 표현을 얻을 수 있는가?
• RQ3 SSIM 손실이 MSE 손실보다 사전 훈련 목표와 후행 밀도 예측 작업 간의 일치에 더 효과적인가?
• RQ4 MDM은 DDPM과 MAE에 비해 소수의 샘플 설정에서 어떻게 성능을 내는가?
• RQ5 재구성 목표로 노이즈를 예측하는 것과 원본 이미지를 예측하는 것 중 어느 것이 후행 분할 정확도에 더 큰 영향을 미치는가?

주요 결과

• MDM는 소수의 샘플 설정에서 GlaS(84.51 ± 1.15)와 FFHQ-34(59.18 ± 0.11)에서 최신 기술 수준의 mIoU 성능을 달성하며, DDPM과 MAE를 모두 능가한다.
• SSIM 손실을 사용할 경우, MDM는 10%의 레이블 데이터에서 GlaS에서 84.51 ± 1.15 mIoU를 기록하며, DDPM(82.32 ± 0.77)과 MAE(79.47 ± 1.08)보다 뚜렷한 향상을 보였다.
• SSIM 손실은 MDM와 DDPM 모두에서 성능 향상을 가져왔지만, MAE에서는 그렇지 않았으며, 이는 SSIM 손실의 효과가 손상 메커니즘의 맥락에 따라 달라진다는 것을 시사한다.
• MSE 손실을 사용한 MDM(82.70 ± 0.79)는 이미 MAE(79.47 ± 1.08)를 능가하며, SSIM 손실을 사용할 경우 더 높은 성능(84.51 ± 1.15)을 기록했다.
• 단일 시간 단계 형태로도 열악한 형태로 변형된 MDM과 DDPM는 MAE를 능가하며, 이는 확산 과정이 단순한 오토인코딩을 넘어서 표현 학습을 향상시킨다는 것을 시사한다.
• 노이즈(εt)를 예측하도록 훈련한 모델이 원본 이미지(x0)를 예측하는 것보다 더 높은 후행 정확도를 달성했으며, 이는 이전의 확산 모델 문헌과 일치한다.