[논문 리뷰] Solving Linear Inverse Problems Provably via Posterior Sampling with Latent Diffusion Models

Litu Rout, Negin Raoof외 4명|arXiv (Cornell University)|2023. 07. 02.

Advanced Neuroimaging Techniques and Applications인용 수 5

한 줄 요약

이 논문은 사전 훈련된 잠재 확산 모델을 미세조정 없이 사용하여 선형 역문제를 해결하는 데 있어 최초의 프레임워크인 후행 샘플링을 통한 잠재 확산 모델(PSLD)을 제안한다. VAE 버블넥을 통해 복원 정밀도를 유지하기 위해 DPS 알고리즘을 확장하여 추가적인 기울기 업데이트 단계를 도입함으로써, PSLD는 증명 가능 샘플 복원과 함께, 인painting, 노이즈 제거, 블러 제거, 초해상도, 스트립 제거 작업 전반에서 최신 기술 수준의 성능을 달성한다. 이는 정량적 및 정성적 지표에서 이전 방법들을 능가한다.

ABSTRACT

We present the first framework to solve linear inverse problems leveraging pre-trained latent diffusion models. Previously proposed algorithms (such as DPS and DDRM) only apply to pixel-space diffusion models. We theoretically analyze our algorithm showing provable sample recovery in a linear model setting. The algorithmic insight obtained from our analysis extends to more general settings often considered in practice. Experimentally, we outperform previously proposed posterior sampling algorithms in a wide variety of problems including random inpainting, block inpainting, denoising, deblurring, destriping, and super-resolution.

연구 동기 및 목표

• 기존의 역문제 해법이 Stable Diffusion와 같은 사전 훈련된 잠재 확산 모델을 활용할 수 없는 격차를 해결하기 위해.
• 작업 특화의 미세조정 없이도 잠재 확산 모델을 사용하여 선형 역문제에 대한 후행 샘플링을 가능하게 하는 방법을 개발하기 위해.
• 선형 모델 설정 하에서 이중 단계 확산 프로세스를 통해 후행 샘플링에 대한 이론적 분석을 수행하고 증명 가능한 샘플 복원을 확립하기 위해.
• 다양한 역문제, 예를 들어 인painting, 노이즈 제거, 블러 제거, 초해상도, 스트립 제거 등에서 최신 기술 수준의 성능을 경험적으로 입증하기 위해.

제안 방법

• VAE 버블넥을 통해 복원 정밀도를 유지하기 위해 DPS 알고리즘을 확장하여, 인코더-디코더 맵이 약간의 항등함수에 가까워지도록 유도하는 추가 기울기 업데이트 단계를 도입한다.
• 잠재 공간에서 복원 오차를 최소화하기 위해 VAE의 인코더와 디코더를 활용하는 글루잉 목적함수를 사용하여, 샘플된 잠재 변수가 정확하게 재구성될 수 있도록 보장한다.
• 사전 훈련된 잠재 확산 모델의 무조건적 스코어 함수와 선형 열화 연산자로부터 유도된 측정 일致성 항을 포함하는 수정된 역방향 SDE를 적용한다.
• 이중 단계 확산 프로세스를 활용한다: 먼저 수정된 역방향 SDE를 사용해 잠재 공간에서 샘플링을 수행하고, 이후에 픽셀 공간으로 디코딩한다. 글루잉 목적함수를 통해 정밀도를 보장한다.
• 모델의 사전 훈련된 잠재 공간 내에서만 작동하여, 미세조정 또는 재훈련이 필요 없도록 한다.
• 이론적 분석을 통해 선형 모델 설정 하에서 샘플 복원이 증명 가능하며, 스코어 기반 생성 모델의 비점근적 분석을 통해 수렴 보장을 도출한다.

실험 결과

연구 질문

• RQ1 사전 훈련된 잠재 확산 모델을 사용한 후행 샘플링이 생성 모델의 미세조정 없이 선형 역문제를 해결하는 데 사용될 수 있는가?
• RQ2 역문제의 후행 샘플링 과정에서 잠재 확산 모델의 복원 정밀도를 어떻게 유지할 수 있는가?
• RQ3 선형 역문제 설정 하에서 잠재 확산 모델을 사용한 후행 샘플링에 대해 어떤 이론적 보장을 설정할 수 있는가?
• RQ4 제안된 방법이 다양한 역문제에서 DPS 및 DDRM과 같은 기존 후행 샘플링 알고리즘을 능가할 수 있는가?
• RQ5 잠재 공간에 글루잉 목적함수를 포함함으로써 표준 후행 샘플링에 비해 복원 품질이 어떻게 향상되는가?

주요 결과

• FFHQ 256 검증 세트에서 무작위 인painting 작업에서 PSLD는 PSNR 30.31, SSIM 0.851을 기록하여 DPS(25.23 PSNR, 0.851 SSIM)를 크게 능가한다.
• 4× 초해상도 작업에서 PSLD는 PSNR 30.73, SSIM 0.867을 기록하여 DPS(25.67 PSNR, 0.852 SSIM)와 DDRM(25.36 PSNR, 0.835 SSIM)를 초월한다.
• 외부 분포의 ImageNet 샘플에 대한 가우시안 블러 제거 작업에서 PSLD는 PSNR 30.10, SSIM 0.843을 기록하여 DPS의 24.25 PSNR, 0.811 SSIM을 상회한다.
• 운동 블러 제거 작업에서 PSLD는 LPIPS를 DPS의 0.613에서 0.244로 감소시켜 더 뛰어난 정서적 품질을 보여준다.
• 스트립 제거 작업에서 PSLD는 수평 방향 LPIPS 0.244, 수직 방향 LPIPS 0.255를 기록하여 DPS의 0.613 및 0.597을 상회한다.
• 이론적 분석을 통해 이중 단계 확산 프로세스를 사용한 선형 모델 설정 하에서 증명 가능한 샘플 복원이 확인되었으며, 제한된 스코어 추정 오차 하에서 수렴 보장을 확립하였다.