[논문 리뷰] Improving Diffusion Models for Inverse Problems using Manifold Constraints

Hyungjin Chung, Byeongsu Sim외 2명|arXiv (Cornell University)|2022. 06. 02.

Advanced Neuroimaging Techniques and Applications인용 수 101

한 줄 요약

논문은 역문제에 대해 무조건 확산 모델 해를 보완하는 매니폴드 제약 그래디언트(MCG) 항을 도입하여 데이터 매니폴드를 따라 반복 경로를 유지하고 재구성을 향상시키며 인페인팅, 색상화, 그리고 희소 뷰 CT와 같은 작업에서 개선을 보인다. 이론적 정당화와 광범위한 실험을 통해 기존 확산 기반 및 지도 학습 방법 대비 상당한 이득을 보여준다.

ABSTRACT

Recently, diffusion models have been used to solve various inverse problems in an unsupervised manner with appropriate modifications to the sampling process. However, the current solvers, which recursively apply a reverse diffusion step followed by a projection-based measurement consistency step, often produce suboptimal results. By studying the generative sampling path, here we show that current solvers throw the sample path off the data manifold, and hence the error accumulates. To address this, we propose an additional correction term inspired by the manifold constraint, which can be used synergistically with the previous solvers to make the iterations close to the manifold. The proposed manifold constraint is straightforward to implement within a few lines of code, yet boosts the performance by a surprisingly large margin. With extensive experiments, we show that our method is superior to the previous methods both theoretically and empirically, producing promising results in many applications such as image inpainting, colorization, and sparse-view computed tomography. Code available https://github.com/HJ-harry/MCG_diffusion

연구 동기 및 목표

• 일반 선형 역문제에 대해 작업 특화 학습 없이 무감독 확산 모델 솔버를 동기 부여한다.
• 확산 경로가 데이터 매니폴드 위에 남도록 매니폴드 제약 그래디언트 항을 도입한다.
• 데이터 일관성 그래디언트와 Tweedie 기반 디노이즈를 통합하여 안정성과 정확성을 향상시킨다.
• 그래디언트가 데이터 매니폴드의 접선 공간에 남는다는 이론적 정당화를 제공한다.
• 인페인팅, 색상화, CT 재구성에 걸쳐 실증적 이득을 보여준다.

제안 방법

• 무조건 스코어 함수를 역 확산 단계 내에서 사용하고 데이터 일관성 투영은 기존 연구와 같이 수행한다.
• 데이터 매니폴드를 통해 Q_i 와 J_{Q_i}로 측정 용어의 그래디언트를 투영하는 매니폴드 제약 그래디언트(MCG)를 도입한다.
• MCG 보정을 접선 공간 투영으로 도출한다: ∂/∂x_i ||W(y−Hx̂0)||^2 on the tangent space T_{x̂0}M.
• Tweedie의 디노이징 철학을 적용하여 단일 스텝 디노이징 보정을 정당화하고 이를 매니폴드 제약과 연결한다.
• MCG 스텝을 일반적인 데이터 일관성 투영과 결합하여 데이터 충실도를 유지하면서 매니폴드에 남아있도록 한다.
• 강한 매니폴드 가정하에서 이론적 결과(Theorem 1)와 기하학적 해석을 제공한다.

실험 결과

연구 질문

• RQ1 문제 특화 학습 없이도 매니폴드 인식 보정 항이 확산 기반 솔버를 역문제에 대해 개선할 수 있는가?
• RQ2 매니폴드 제약 그래디언트가 확산 경로를 데이터 매니폴드에 더 가깝게 유지하고 오차 누적을 줄이는가?
• RQ3 MCG가 표준 투영 기반 데이터 일관성 및 Tweedie 디노이징과 실제로 어떻게 상호 작용하는가?
• RQ4 인페인팅, 색상화, 희소 뷰 CT와 같은 다양한 역문제에서 어떤 성능 이득이 나타나는가?
• RQ5 방법이 서로 다른 측정 마스크와 NFEs(함수 평가 수)에 대해 견고한가?

주요 결과

• MCG가 반복 경로를 데이터 매니폴드 위에 유지하여 오차 누적을 줄인다.
• MCG를 투영과 결합하면 데이터 일관성과 재구성이 향상된 결과를 얻고 다수의 작업에서 기존 확산 기반 방법을 능가한다.
• MCG는 Inpainting(FFHQ/ImageNet), 색상화(FFHQ/LSUN-bedroom), CT 재구성(AAPM 데이터셋)에서 Score-SDE, DDRM, RePAINT 및 기타 기준선보다 우수하다.
• 배치들에서 MCG만 사용할 때는 LPIPS가 향상되지만 측정 일관성은 해로울 수 있으며, MCG와 투영을 조합하면 최상의 결과를 얻는다.
• 정성적 결과에서 MCG로 더 높은 충실도 재구성이 가능하며 경계 아티팩트가 감소한다.