[논문 리뷰] Plug-and-Play image restoration with Stochastic deNOising REgularization

Marien Renaud, Jean Prost외 2명|arXiv (Cornell University)|2024. 02. 01.

Photoacoustic and Ultrasonic Imaging인용 수 1

한 줄 요약

이 논문은 SNORE(Soft deNOising REgularization)를 제안하며, 이는 각 반복 단계에서 현재 이미지 추정치의 노이즈가 있는 복제본에 사전 훈련된 디노이징 모델을 적용함으로써, 디노이징 모델이 훈련된 노이즈 분포 내에서 작동하도록 보장하는 새로운 플러그 앤 플레이(Plug-and-Play) 이미지 복원 프레임워크이다. 이 방법은 이론적 보장 하에 수렴을 보장하며, 디블러링과 인painting 과제에서 최신 기술 수준의 성능을 달성하여, 표준 PnP 및 디퓨전 기반 방법보다 정량적·정성적 기준에서 모두 뛰어난 성능을 보인다.

ABSTRACT

Plug-and-Play (PnP) algorithms are a class of iterative algorithms that address image inverse problems by combining a physical model and a deep neural network for regularization. Even if they produce impressive image restoration results, these algorithms rely on a non-standard use of a denoiser on images that are less and less noisy along the iterations, which contrasts with recent algorithms based on Diffusion Models (DM), where the denoiser is applied only on re-noised images. We propose a new PnP framework, called Stochastic deNOising REgularization (SNORE), which applies the denoiser only on images with noise of the adequate level. It is based on an explicit stochastic regularization, which leads to a stochastic gradient descent algorithm to solve ill-posed inverse problems. A convergence analysis of this algorithm and its annealing extension is provided. Experimentally, we prove that SNORE is competitive with respect to state-of-the-art methods on deblurring and inpainting tasks, both quantitatively and qualitatively.

연구 동기 및 목표

• 표준 플러그 앤 플레이(PnP) 방법의 한계를 해결하기 위해, 디노이징 모델이 점점 더 깨끗한 이미지에 적용되면서 발생하는 분포 이탈(distribution shift) 문제를 해결한다.
• 각 반복 단계에서 노이즈를 주입함으로써 디노이징 모델의 효과성을 유지하면서도, 이론적으로 탄탄한 PnP 알고리즘을 개발한다.
• 디퓨전 모델의 고품질 결과와 PnP 방법의 이론적 수렴 보장을 연결한다.
• 입력 노이즈 수준을 디노이징 모델의 훈련 노이즈 수준과 일치시킴으로써, 디노이징 모델이 훈련 도메인 내에서 작동하도록 보장한다.
• 불안정한 이미지 역문제에서 강건성과 수렴성을 향상시키는 확률적 정규화 프레임워크를 제공한다.

제안 방법

• SNORE는 이미지의 노이즈가 있는 복제본에 대한 부드러운 로그-프리오르의 기대값으로 정의된 확률적 정규화 항을 도입하며, 이는 표준 음의 로그-프리오르를 대체한다.
• 각 반복 단계에서 알고리즘은 ˜xk = xk + σϵ 형태의 노이즈가 있는 복제본을 생성하며, 여기서 ϵ ∼ N(0, Id)이다. 이 노이즈가 있는 이미지에 대해 디노이징 모델 Dσ를 적용한다.
• 이 방법은 변분 목표를 최소화하기 위해 확률적 경사 하강법(Algorithm 2)을 사용하며, 정규화 단계는 정규화의 프록시멀 연산자로 대체된다.
• 냉각 기반 변형(Algorithm 3)이 제안되며, 이는 냉각된 중요도 샘플링에 영감을 받아 수렴성과 안정성을 향상시킨다.
• 이론적 분석을 통해 유한성 가정 하에 거의 확실히 수렴함을 증명하며, 정확하지 않은 디노이징 모델에 대해서도 오차 한계를 제시한다.
• 이 프레임워크는 프록시멀 기반 및 경사 기반 PnP 방법 모두와 호환되어 기존 솔버와의 탄력적인 통합이 가능하다.

실험 결과

연구 질문

• RQ1 PnP 프레임워크에 매 반복 단계에서 노이즈를 주입함으로써, 디노이징 모델 적용의 일관성이 훈련된 분포와 일치하는가?
• RQ2 디노이징 모델의 훈련 분포와 일치하는 노이즈 수준을 가진 이미지에 디노이징 모델을 적용할 경우, 수렴성과 복원 품질이 향상되는가?
• RQ3 확률적 정규화 프레임워크는 디퓨전 모델 수준의 성능을 유지하면서도 이론적 수렴 보장을 달성할 수 있는가?
• RQ4 SNORE 반복 과정에서 이미지 추정치의 노이즈 수준은 어떻게 변화하며, 디노이징 모델의 노이즈 수준 근처에서 안정화되는가?
• RQ5 정확하지 않거나 MMSE가 아닌 디노이징 모델을 사용할 경우, SNORE 알고리즘의 수렴성과 최종 이미지 품질에 어떤 영향을 미치는가?

주요 결과

• SNORE는 디블러링 및 인painting 과제에서 표준 PnP 및 최근의 디퓨전 기반 방법을 모두 능가하는 최신 기술 수준의 성능을 달성하며, 정량적 지표와 시각적 품질 양면에서 뛰어난 성능을 보였다.
• SNORE의 반복 과정에서의 노이즈 수준은 디노이징 모델의 노이즈 수준 σ로 수렴함을 확인하였으며, 이는 디노이징 모델이 의도한 분포 도메인 내에서 작동하고 있음을 뒷받침한다.
• 반면, RED는 수렴 시에도 입력 노이즈 수준과 디노이징 모델의 노이즈 수준 사이에 지속적인 이격이 존재함을 보여주며, 분포 이탈이 발생하고 있음을 시사한다.
• SNORE의 냉각 버전은 점진적으로 노이즈 수준을 감소시킴으로써 초기 반복 단계에서 수렴성과 안정성을 향상시킨다.
• 이론적 분석을 통해 유한성 가정 하에 확률적 경사 하강법이 거의 확실히 수렴함을 확인하였으며, 정확하지 않은 디노이징 모델에 대해서도 오차 한계를 제시하였다.
• SNORE는 이미지 추정치의 잔여 노이즈에 대해 강건성을 보이며, 노이즈 추정치가 자연 노이즈 수준으로 수렴하는 것으로 나타났다.