[논문 리뷰] Conffusion: Confidence Intervals for Diffusion Models

Eliahu Horwitz, Yedid Hoshen|arXiv (Cornell University)|2022. 11. 17.

Model Reduction and Neural Networks인용 수 6

한 줄 요약

Conffusion는 초상화 복원 및 메시지 복원과 같은 이미지 간 작업에서 확산 모델의 픽셀 수준 신뢰구간을 생성하는 방법을 제안하며, 통계적으로 보장되는 불확실성 추정을 가능하게 한다. 사전 훈련된 확산 모델을 분위수 회귀 손실로 미세조정함으로써, 샘플링 기반 기준보다 3개 정도의 차수 빠른 추론 속도를 달성하면서도 더 좁은 간격을 확보하며 이론적 커버리지 보장을 유지한다.

ABSTRACT

Diffusion models have become the go-to method for many generative tasks, particularly for image-to-image generation tasks such as super-resolution and inpainting. Current diffusion-based methods do not provide statistical guarantees regarding the generated results, often preventing their use in high-stakes situations. To bridge this gap, we construct a confidence interval around each generated pixel such that the true value of the pixel is guaranteed to fall within the interval with a probability set by the user. Since diffusion models parametrize the data distribution, a straightforward way of constructing such intervals is by drawing multiple samples and calculating their bounds. However, this method has several drawbacks: i) slow sampling speeds ii) suboptimal bounds iii) requires training a diffusion model per task. To mitigate these shortcomings we propose Conffusion, wherein we fine-tune a pre-trained diffusion model to predict interval bounds in a single forward pass. We show that Conffusion outperforms the baseline method while being three orders of magnitude faster.

연구 동기 및 목표

• 사용자가 정의한 확률로 진짜 픽셀 값이 간격 내에 포함될 수 있도록, 확산 모델 출력에 대해 통계적으로 엄밀한 신뢰구간을 제공하는 것.
• 확산 모델에서 샘플링 기반 신뢰구간 추정의 느린 추론 속도를 극복하는 것.
• 기존 분위수 회귀 방법과 비교해 더 좁은 간격을 확보하면서도 이론적 커버리지 보장을 유지하는 것.
• 단일 사전 훈련된 확산 모델을 사용해 다양한 작업과 데이터셋 간 전이 학습을 가능하게 하는 것.
• 노이즈 제거 모델을 확산 과정에서 분리하여 단일 통과 추론을 효율적으로 수행할 수 있도록 하는 것.

제안 방법

• 사전 훈련된 확산 모델(예: ADM)을 분위수 회귀 손실로 미세조정하여, 단일 전방전파에서 각 픽셀의 하한 및 상한을 예측하는 것.
• 노이즈 제거 모델을 확산 과정에서 분리함으로써, 전체 확산 과정을 재훈련하지 않고도 간격 추정에 재사용할 수 있도록 하는 것.
• 공차 예측 원칙을 활용하여, 진짜 픽셀 값이 예측 간격 내에 사용자가 지정한 확률로 포함되도록 보장하는 것.
• ImageNet에서 유도된 강력한 사전 훈련된 특징을 활용하여 초상화 복원 및 메시지 복원과 같은 하류 작업으로 제로샷 전이를 가능하게 하는 것.
• 추론의 효율화를 위해 암시적 추론을 적용: 이미지당 수천 번의 출력을 샘플링하는 대신, 단일 전방전파에서 간격을 직접 예측하는 것.
• 단일 스칼라 분위수 손실을 사용하여 다중 분위수(예: 0.05 및 0.95)를 동시에 최적화하여 간격의 하한 및 상한을 정의하는 것.

실험 결과

연구 질문

• RQ1 이미지 간 작업에서 확산 모델 출력에 대해 통계적으로 타당한 신뢰구간을 구성할 수 있는가?
• RQ2 샘플링 기반 기준 대비 더 빠른 추론 속도를 확보할 수 있는가?
• RQ3 사전 훈련된 확산 모델을 간격 예측을 위해 미세조정하면 기존 분위수 회귀 방법보다 더 좁고 신뢰도 있는 간격을 얻을 수 있는가?
• RQ4 효율적인 간격 추정을 위해 확산 사전 훈련이 필수적인가, 아니면 일반적인 특징 추출기로도 충분한가?
• RQ5 단일 사전 훈련된 모델이 초상화 복원 및 메시지 복원과 같은 다양한 작업에 일반화될 수 있는가?

주요 결과

• Conffusion는 샘플링 기반 기준 대비 3개 정도의 차수 빠른 추론 속도를 확보하여, 추론 시간을 수천 번의 전방전파에서 단일 전방전파로 줄였다.
• 기존 최고 수준의 방법인 ADM_{UQ}보다 더 좁은 간격을 생성하면서도 이론적 커버리지 보장을 유지한다.
• N-Conffusion는 간격 크기와 시각적 품질 면에서 샘플링 기반 및 분위수 회귀 기준 모두를 능가하며, 현실적이고 선명한 간격을 생성한다.
• 확산 사전 훈련은 필수적이다: 분석 결과, 비확산 모델(예: ImageNet 또는 Instagram 사전 훈련된 인코더를 갖춘 FPN)은 Conffusion보다 훨씬 넓은 간격을 생성한다.
• 모델은 다양한 작업에서 우수한 성능을 유지하며, ImageNet 사전 훈련된 모델에서 CelebA-HQ 및 기타 데이터셋으로의 일반화가 가능하며, 작업별 재훈련 없이도 효과적으로 작동한다.
• Conffusion는 중간 시간 단계에서 신뢰구간을 추정함으로써 확산 생성 과정에서의 조기 정지가 가능하게 하여, 최종 출력 품질의 빠른 대체 지표를 제공한다.