[논문 리뷰] Fast and Stable Diffusion Inverse Solver with History Gradient Update

Linchao He, Hongyu Yan외 9명|arXiv (Cornell University)|2023. 07. 22.

MRI in cancer diagnosisMedicine인용 수 3

한 줄 요약

이 논문은 사전 훈련된 잠재 확산 모델(LDM)을 데이터 사전지식으로 통합하여 반복적 CT 재구성에 사용하는 제로샷 방법인 잠재 확산 반복 재구성(LDIR)을 제안한다. 데이터 일관성 항의 기울기 지도를 통해 쌍화된 데이터나 역투영 행렬 근사가 필요 없도록 하며, 희박한 시야 및 제한된 각도 CT 재구성에서 최신 기술을 초월하는 성능을 달성한다. 이는 고해상도(512×512) 재구성과 함께 더 빠른 추론 및 낮은 메모리 사용을 가능하게 한다.

ABSTRACT

Diffusion models have recently been recognised as efficient inverse problem solvers due to their ability to produce high-quality reconstruction results without relying on pairwise data training. Existing diffusion-based solvers utilize Gradient Descent strategy to get a optimal sample solution. However, these solvers only calculate the current gradient and have not utilized any history information of sampling process, thus resulting in unstable optimization progresses and suboptimal solutions. To address this issue, we propose to utilize the history information of the diffusion-based inverse solvers. In this paper, we first prove that, in previous work, using the gradient descent method to optimize the data fidelity term is convergent. Building on this, we introduce the incorporation of historical gradients into this optimization process, termed History Gradient Update (HGU). We also provide theoretical evidence that HGU ensures the convergence of the entire algorithm. It's worth noting that HGU is applicable to both pixel-based and latent-based diffusion model solvers. Experimental results demonstrate that, compared to previous sampling algorithms, sampling algorithms with HGU achieves state-of-the-art results in medical image reconstruction, surpassing even supervised learning methods. Additionally, it achieves competitive results on natural images.

연구 동기 및 목표

• 쌍화된 훈련 데이터나 역투영 행렬 근사가 필요 없이 희박한 데이터에서의 CT 재구성 문제의 불안정성을 해결하기 위해.
• 사전 훈련된 잠재 확산 모델을 비지도 데이터 사전지식으로 사용하여 제로샷 고해상도(512×512) 이미지 재구성을 가능하게 하기 위해.
• 데이터 일관성 항의 기울기 정보로 잠재 확산 샘플링 과정을 지도함으로써 재구성 품질과 효율성을 향상시키기 위해.
• 픽셀 공간 반복 방법에 비해 계산 비용과 메모리 소비를 감소시키면서도 성능을 유지하거나 초월하기 위해.
• 의료 영상 및 자연 이미지 복원(예: 인painting, 슈퍼해상도 등)을 포함한 다양한 작업에 대한 일반화 능력을 입증하기 위해.

제안 방법

• LDIR는 사전 훈련된 무조건적 잠재 확산 모델을 데이터 사전지식으로 사용하여 반복 재구성 문제를 최적화 문제로 공식화한다. 기존의 수작업 또는 학습된 사전지식을 대체한다.
• 잠재 확산 모델의 역방향 샘플링 과정은 데이터 일관성 항의 기울기로 지도되며, 이는 역투영 행렬이 필요 없이 희박한 측정치와 일치하는 재구성을 생성할 수 있도록 한다.
• 역방향 샘플링 과정 중에 이전 기울기 정보를 통합함으로써 샘플 일관성을 향상시키기 위해 역사 기울기 업데이트 메커니즘을 도입한다.
• 모든 계산이 잠재 공간에서 이루어지므로 계산 복잡도가 감소하고 고해상도(512×512) 재구성에 있어 효율성이 향상된다.
• 프레임워크는 한 번만 훈련되며, 희박한 시야 및 제한된 각도 스캔과 같은 다양한 CT 재구성 작업에 재훈련 없이 또는 아키텍처 변경 없이 적용 가능하다.
• 투영 기반 최적화를 직접 피하기 위해 잠재 공간에서 스코어 기반 샘플링을 사용함으로써 다양한 검출기 기하학에 적응 가능하다.

실험 결과

연구 질문

• RQ1 쌍화된 데이터나 역투영 행렬 근사 없이도 사전 훈련된 잠재 확산 모델을 반복 재구성에서 효과적으로 데이터 사전지식으로 사용할 수 있는가?
• RQ2 데이터 일관성 항의 기울기 지도가 제로샷 환경에서 재구성 품질을 어떻게 향상시키는가?
• RQ3 역사 기울기 업데이트 메커니즘이 잠재 확산 기반 재구성에서 샘플 일관성과 수렴 안정성을 향상시키는가?
• RQ4 잠재 공간 기반 반복 재구성은 픽셀 공간 대비 속도, 메모리 효율성, 재구성 품질 측면에서 뛰어나지 않는가?
• RQ5 단일 통합 모델이 의료 CT 및 자연 이미지 복원과 같은 다양한 희박 측정 작업에 일반화 가능한가?

주요 결과

• LDIR는 희박한 시야 CT 재구성에서 PSNR 39.01, SSIM 0.9552를 달성하여 다음으로 우수한 방법(DPS)의 PSNR 28.55, SSIM 0.8140를 크게 능가한다.
• 제한된 각도 CT 재구성에서 LDIR는 PSNR 39.77, SSIM 0.9612를 기록하여 DPS(PSNR 28.97, SSIM 0.8242)와 DIR(PSNR 32.82, SSIM 0.8898)를 모두 능가한다.
• LDIR는 초당 36.67회의 반복을 수행하고 메모리 소비는 단 4268 MB로, 픽셀 기반 방법(DIR: 20.75 it/s, 6338 MB)에 비해 속도와 메모리 효율성에서 뛰어나다.
• 자연 이미지 작업에서는 DPS(최신 기술)와 경쟁 가능한 성능을 기록했으며, 특히 99% 랜덤 인페인팅과 최소한의 잡음으로 인한 슈퍼해상도에서 뛰어난 성능을 보였다.
• 절단 실험 결과, 역사 기울기 지도를 사용한 잠재 공간 재구성은 픽셀 공간 방법(예: DPS의 PSNR 28.55)에 비해 뛰어난 성능(PSNR 39.01)을 기록했으며, 속도와 메모리 효율성 측면에서 상당한 향상이 있었다.
• 이 방법은 다양한 작업에 대해 잘 일반화되며, 의료 영상 및 자연 이미지 복원 모두에서 뛰어난 성능을 보였고, 매우 희박한 측정치 조건에서도 강력한 성능을 유지했다.