An Expectation-Maximization Algorithm for Training Clean Diffusion Models from Corrupted Observations

문제 정의
확산 모델(Diffusion Models, DMs)은 복잡한 이미지 사전 확률(prior)을 모델링함으로써 계산적 이미징 역문제(computational imaging inverse problems)를 해결하는 데 탁월한 능력을 입증해 왔습니다. 그러나 이들의 실질적인 적용은 대규모의 고품질 깨끗한 데이터셋에 대한 의존성으로 인해 심각하게 제약받고 있습니다. 많은 실제 시나리오, 특히 과학 및 생물 의학 이미징(예: X-선 CT, cryo-EM)에서는 이러한 깨끗한 데이터를 확보하는 것이 불가능하거나 비용이 지나치게 많이 듭니다. 반대로, 미분 가능한 순방향 모델(forward models)을 갖춘 오염된 관측치(노이즈, 블러 또는 마스킹 처리된 데이터)는 풍부합니다. 이는 "닭과 달걀"의 딜레마를 생성합니다: 정확한 확산 모델을 훈련하려면 깨끗한 이미지가 필요하지만, 오염된 관측치로부터 깨끗한 이미지를 재구성하려면 잘 훈련된 확산 모델이 필요합니다. AmbientDiffusion 또는 SURE-Score와 같은 오염된 데이터로부터 학습하는 기존 방법들은 특정 유형의 오염(예: 대규모 인페인팅 또는 상당한 노이즈)에서 어려움을 겪거나 다양한 역문제에 걸쳐 일반화하는 데 실패하는 경우가 많습니다.

방법론: EMDiffusion
저자들은 이미지 역문제를 공동으로 해결하고 오염된 관측치로부터 직접 확산 모델을 훈련하도록 설계된 새로운 기댓값 최대화(Expectation-Maximization, EM) 프레임워크인 EMDiffusion을 제안합니다. 이 방법은 두 단계의 반복적인 사이클을 통해 작동합니다:

1. 초기화(Initialization): 프로세스는 매우 적은 양의 깨끗한 이미지(최소 10~50장)로 훈련된 "모호한(vague)" 확산 사전 확률로부터 시작됩니다. 이 초기화는 모델이 오염된 입력값에 의해서만 정의되는 분포로 붕괴(collapse)되는 것을 방지하는 데 매우 중요합니다. 저자들은 이러한 초기 깨끗한 이미지들이 타겟 작업과 동일한 분포일 필요는 없으며, 분포 외(out-of-distribution) 데이터도 효과적인 초기화 도구로 사용될 수 있다고 언급합니다.

2. E-단계 (Expectation): 현재의 확산 모델을 사전 확률로 사용하여, 오염된 관측치로부터 깨끗한 이미지를 재구성하기 위해 **적응형 확산 사후 샘플링(Adaptive Diffusion Posterior Sampling)**을 수행합니다.

   • 표준적인 사후 샘플링은 초기 반복 단계에서 사전 확률이 약하기 때문에 종종 실패하며, 이는 샘로가 관측치를 무시하고 훈련 분포로 되돌아가는 "모드 붕괴(mode collapse)"를 초래합니다.
   • 이를 완화하기 위해, EMDiffusion은 역시간 확률 미분 방정식(reverse-time SDE) 내의 데이터 가능도(data likelihood) 항을 재조정하는 하이퍼파라미터 $\lambda$를 도입합니다.
   • $\lambda$는 각 E-단계에서 데이터 손실($\|y - A(\hat{x}_0)\|_2^2$)을 최소화하도록 적응적으로 조정되어, 불완전한 사전 확률과 관측된 데이터 사이의 균형을 맞춥니다. 이를 통해 초기 사전 확률이 열악할 때도 더 깨끗한 재구성 이미지를 생성할 수 있습니다.

3. M-단계 (Maximization): 알고리즘은 E-단계에서 생성된 재구성된 이미지(사후 샘플)를 사용하여 확산 모델의 가중치를 업데이트합니다.

   • 이 단계는 모델이 재구성된 "깨끗한" 데이터의 스코어 함수(score function)를 근사하도록 학습하는 표준적인 디노이징 스코어 매칭(denoising score matching)과 유사합니다.
   • 훈련을 가속화하고 노이즈가 섞인 초기 재구성 결과에 과적합되는 것을 방지하기 위해, 이 방법은 2단계 훈련 전략을 채택합니다: 초기에는 이전 반복 단계에서 상속된 가중치를 미세 조정(fine-tuning)하고, 재구성 품질이 개선되면(최적의 $\lambda$가 1 미만으로 떨어지는 것으로 표시됨), 모델을 재초기화하여 전체 사후 샘플에 대해 처음부터 다시 훈련합니다.

주요 기여

• 오염된 데이터를 위한 EM 프레임워크: 본 논문은 EM 접근 방식을 사용하여 임의로 오염된 관측치로부터 스코어 기반 확산 모델을 훈련하는 최초의 일반적인 프레임워크를 도입하여, 깨끗한 데이터셋에 대한 의존성을 효과적으로 타파했습니다.
• 적응형 사후 샘플링: 역 SDE를 위한 가능도 항에 대한 적응형 스케일링 인자($\lambda$)의 제안은 초기 EM 반복 단계의 고유한 모드 붕괴 문제를 해결하며, 약한 사전 확률에서도 안정적인 재구성을 가능하게 합니다.
• 반복적 정교화: 이 방법은 적은 양의 깨끗한 데이터가 암묵적인 정규화 도구로 작용하여, 모델이 재구성과 훈련의 교차 단계를 통해 실제 깨끗한 데이터 분포로 점진적으로 수렴할 수 있음을 보여줍니다.
• 일반화 가능성: 특정 오염(예: 마스킹)에 특화된 기존 연구들과 달리, EMDiffusion은 무작위 인페인팅, 디노이징, 디블러링에 적용 가능한 통합된 접근 방식을 제공합니다.

실험 결과
저자들은 세 가지 작업(무작위 인페인팅(60% 픽셀 마스킹), 가우시안 디노이징, 가우시안 디블러링)에 대해 CIFAR-10 및 CelebA 데이터셋에서 EMDiffusion을 검증했습니다.

• 인페인팅: EMDiffusion은 AmbientDiffusion 및 SURE-Score와 같은 베이스라인을 크게 능가하며, 전체 깨끗한 데이터셋으로 훈련된 사전 확률을 사용하는 Diffusion Posterior Sampling(DPS)과 대등한 재구성 품질(PSRE/LPIPS)을 달달성했습니다.
• 디노이징 및 디블러링: 이 방법은 SURE-Score 및 자기 지도 학습 베이스라인(Noise2Self)을 능가하며, 다른 방법들이 매끄럽게 뭉개버리는 미세한 디테일(예: 자동차 바퀴, 얼굴 특징)을 복원해 냅니다.
• 생성 품질: 학습된 확산 모델은 오염된 데이터로 훈련된 다른 방법들과 비교했을 때 최첨단(state-of-the-art) FID 점수를 달성했습니다. 특히, EM을 통해 50,000개의 오염된 이미지로부터 학습된 모델이 약 15,000개의 깨끗한 이미지로 훈련된 모델보다 성능이 우수하다는 점이 주목할 만합니다.
• 효율성: 이 방법은 표준 DDPM 아키텍처(35.7M 파라미터)를 사용하여 이러한 결과를 달성했으며, 이는 경쟁 베이스라인(예: AmbientDiffusion의 113M+ 파라미터)에서 사용되는 아키텍처보다 현저히 작습니다.

의의 및 주장
본 논문은 EMDiffusion이 오염된 측정치로부터 근본적인 깨끗한 데이터 분포를 학습하는 것이 정보 이론적으로 가능하다는 것을 입증한다고 주장합니다. 그 의의는 다음과 같습니다:

1. 데이터 부족 문제 해결: 깨끗한 데이터가 없는 영역을 위해, 풍부한 오염된 관측치를 활용하는 실질적인 솔루션을 제공합니다.
2. 강건한 초기화: 최소한의 깨끗한 이미지(심지어 분포 외 데이터라도)가 매니폴드를 초기화하는 데 충분하며, 이를 통해 모델이 편향된 오염된 입력 분포로 붕괴되는 것을 방지할 수 있음을 확립합니다.
3. 최첨단 성능: 대규모 깨끗한 데이터셋에 대한 사전 훈련된 모델 없이도 역 이미징 작업에서 새로운 최첨단 성능을 달성하여, 이론적 사전 확률과 실제 데이터 제약 사이의 간극을 메웁니다.

저자들은 이 방법이 순방향 모델이 알려져 있고 미분 가능하다는 가정에 의존한다는 점을 인정하며, 향후 연구가 알려지지 않은 순방향 모델이나 노이즈 통계량을 학습하는 방향으로 확장될 수 있음을 명시하며 범위를 겸손하게 설정하였습니다.