[논문 리뷰] Solving Inverse Problems with Latent Diffusion Models via Hard Data Consistency

Bowen Song, Soo Min Kwon외 4명|arXiv (Cornell University)|2023. 07. 16.

Generative Adversarial Networks and Image Synthesis인용 수 13

한 줄 요약

ReSample은 일반적인 선형 및 비선형 역문제를 해결하기 위해 하드 데이터 일관성을 갖춘 잠재 확산 모델을 활용하여 최첨단 결과를 달성하면서 메모리 사용량을 줄입니다.

ABSTRACT

Diffusion models have recently emerged as powerful generative priors for solving inverse problems. However, training diffusion models in the pixel space are both data-intensive and computationally demanding, which restricts their applicability as priors for high-dimensional real-world data such as medical images. Latent diffusion models, which operate in a much lower-dimensional space, offer a solution to these challenges. However, incorporating latent diffusion models to solve inverse problems remains a challenging problem due to the nonlinearity of the encoder and decoder. To address these issues, we propose extit{ReSample}, an algorithm that can solve general inverse problems with pre-trained latent diffusion models. Our algorithm incorporates data consistency by solving an optimization problem during the reverse sampling process, a concept that we term as hard data consistency. Upon solving this optimization problem, we propose a novel resampling scheme to map the measurement-consistent sample back onto the noisy data manifold and theoretically demonstrate its benefits. Lastly, we apply our algorithm to solve a wide range of linear and nonlinear inverse problems in both natural and medical images, demonstrating that our approach outperforms existing state-of-the-art approaches, including those based on pixel-space diffusion models.

연구 동기 및 목표

• 고차원 역문제에 대한 효율적 프라이어로 잠재 확산 모델(LDM)을 활용하도록 동기를 부여한다.
• LDM에서 역 샘플링 과정에서 측정 일관성을 강제로 보장하기 위해 하드 데이터 일관성을 도입한다.
• 측정 일관성을 갖는 잠재 샘플을 데이터 매니폴드로 다시 매핑하는 재샘플링 스킴을 개발한다.
• 향상된 성능과 메모리 효율성에 대한 이론적 보장과 실증적 근거를 제시한다.
• 자연 및 의료 영상의 다양한 선형 및 비선형 역문제에 대한 적용 가능성을 입증한다.

제안 방법

• 잠재 벡터 z를 가진 인코더 E와 디코더 D를 통해 잠재 공간에서 확산 모델을 작동시킨다.
• ||y - A(D(z))||^2를 최소화하는 z에 대한 제약 최적화를 해결하여 무조건적 역 업데이트를 하드 데이터 일관성으로 대체한다.
• 잠재 최적화를 지상 진짜 잠재에 근접하도록 Tweedie 기반 추정 z0_hat으로 초기화한다.
• 최적화된 잠재를 노이즈가 있는 데이터 매니폴드에 다시 매핑하기 위한 확률적 재샘플링을 도입하여 계속되는 역 확산을 위한 사후-유사 샘플을 형성한다.
• 정확도와 계산량의 균형을 맞추기 위해 선택된 역 시간 스텝에서만 데이터 일관성을 적용하는 건너뛰기(스킵) 단계 기법을 사용한다.
• 재샘플링 잠재의 분산 감소와 무편향성을 보이는 이론적 결과를 제공하고, 하이퍼파라미터 sigma^2를 통한 분산 제어에 대해 논의한다.

실험 결과

연구 질문

• RQ1 잠재 확산 모델(LDM)이 일반적인(선형 및 비선형) 역문제의 프라이어로 효과적으로 작동할 수 있는가?
• RQ2 잠재 공간 최적화를 통한 하드 데이터 일관성이 소프트 경사 기반 업데이트와 비교하여 측정 신뢰도와 이미지 품질을 향상시키는가?
• RQ3 제안된 확률적 재샘플링이 분산을 감소시키고 잠재 공간의 역확산 동안 데이터 일관성을 더 잘 보존하는가?
• RQ4 자연 및 의료 영상 작업 전반에서 ReSample과 픽셀 공간 확산 접근 방식의 메모리 및 계산 비용 차원은 어떤가?

주요 결과

• ReSample은 자연 및 의료 영상의 여러 선형 및 비선형 역문제에서 최첨단 성능을 달성한다.
• 하드 데이터 일관성은 잠재 공간 최적화와 확률적 재샘플링과 결합되어 경쟁 방법들보다 더 높은 PSNR/SSIM과 더 낮은 LPIPS를 산출한다.
• 확률적 재샘플링은 분산을 줄이고 노이즈에 강건한 재구성을 생성하며, 확산 시간 t가 0에 가까워질수록 측정 일관성을 유지한다.
• 이 방법은 잠재 공간에서 작동함으로써 메모리 효율성을 제공하고, 픽셀 공간 기준선에 비해 상당한 메모리 사용 감소를 보인다.
• 가지치기(아블레이션) 결과는 재샘플링이 확률적 인코딩보다 이점을 가지며 CT 재구성 작업에서 재샘플링 빈도의 가치가 있음을 보여준다.