Stochastic Generative Plug-and-Play Priors

1. 연구 배경 및 문제 정의 (Problem)

• 배경: 이미지 역문제 (Image Inverse Problems, 예: 초해상도, 인페인팅, MRI 재구성 등) 를 해결하기 위해 Plug-and-Play (PnP) 방법이 널리 사용됩니다. PnP 는 최적화 알고리즘의 사전 (Prior) 단계에 학습된 디노이저 (Denoiser) 를 삽입하여, 명시적인 분석적 사전 (예: Total Variation) 대신 강력한 학습된 이미지 통계를 활용합니다.
• 문제점: 기존 PnP 방법은 심하게 ill-posed(잘못 정의된) 역문제에서 성능이 저하되는 한계가 있습니다.
  1. 분포 불일치 (Distribution Mismatch): PnP 의 중간 반복(iterates) 은 학습된 디노이저가 가정하는 단순한 가우시안 잡음이 아닌, 구조화된 아티팩트를 포함합니다. 이로 인해 디노이저가 훈련된 잡음 영역 밖에서 사용되어 성능이 떨어집니다.
  2. 제한된 잡음 범위: 기존 PnP 에 주로 사용되는 CNN 기반 디노이저 (예: DRUNet) 는 저잡음 (low-noise) 복원에 최적화되어 있어, 심하게 손상된 이미지 (대규모 결손 등) 를 처리하는 데 한계가 있습니다.
• 현재의 대안 및 한계: 최근 Score-based Diffusion Models (SBDM) 은 다양한 잡음 수준에서 훈련되어 강력한 생성 능력을 보입니다. 그러나 SBDM 을 PnP 프레임워크에 통합할 때, 기존에는 매 최적화 단계마다 역확산 (reverse diffusion) 샘플링을 수행해야 하는 복잡성이나, 단순히 디노이저만 교체하는 것만으로는 ill-posed 문제에서 충분한 성능 향상을 얻지 못했습니다.

2. 제안 방법론 (Methodology)

저자들은 PnP 와 SBDM 간의 이론적 연결고리를 확립하고, 이를 활용한 새로운 Stochastic Generative PnP (SGPnP) 프레임워크를 제안합니다.

A. PnP 를 위한 Score 기반 해석 (Score Adaptation for PnP)

• 이론적 연결: Tweedie's 공식을 활용하여, SBDM 이 학습한 Score Function (로그 확률 밀도의 기울기) 과 MMSE 디노이저 간의 수학적 관계를 규명했습니다.
• 매핑 전략: PnP 알고리즘이 요구하는 연속적인 잡음 수준 ($\sigma$) 과 SBDM 이 사용하는 이산적인 시간 단계 ($t$) 간의 불일치를 해결하기 위해, 훈련된 SBDM 을 PnP 디노이저로 직접 매핑하는 방법을 제시했습니다.
  • VE (Variance-Exploding) 모델: $D_\theta(x; t) = x + \sigma_t^2 s_\theta^{VE}(x, t)$
  • VP (Variance-Preserving) 모델: $D_\theta(x; t) = x + \frac{1-\bar{\alpha}_t}{\sqrt{\bar{\alpha}_t}} s_\theta^{VP}(\sqrt{\bar{\alpha}_t}x, t)$
• 이를 통해 역확산 샘플링 없이도 SBDM 을 PnP 의 사전 단계에서 직접 사용할 수 있게 되었습니다.

B. 확률적 생성 PnP (SGPnP) 프레임워크

• 핵심 아이디어: 디노이저 입력에 의도적으로 잡음을 주입 (Noise Injection) 하는 두 가지 목적을 가집니다.
  1. 정합성 (Alignment): PnP 의 중간 반복 결과물에 SBDM 이 훈련된 가우시안 잡음 분포를 재적합시켜, 디노이저가 가장 잘 작동하는 영역으로 유도합니다.
  2. 최적화 개선 (Optimization): 주입된 잡음이 최적화 과정에서 Strict Saddle Point (엄격한 안장점) 에서 벗어나도록 돕는 확률적 동역학을 생성합니다.
• 구현:
  • 데이터 일관성 (Data Consistency) 업데이트 후, 디노이저에 입력되기 전에 $x_{input} = z_k + \sigma_{inject} n$ 형태로 잡음을 추가합니다.
  • 잡음 수준 분리 (Decoupling): 주입된 잡음 수준 ($\sigma_{inject}$) 과 디노이저의 조건부 잡음 수준 ($\sigma_{cond}$) 을 서로 다르게 설정할 수 있습니다. 이는 데이터 일관성 단계에서 발생하는 잔류 아티팩트와 주입된 잡음을 모두 고려하여 디노이저가 더 넓은 범위의 손상을 처리하도록 합니다.

3. 주요 기여 (Key Contributions)

1. PnP 에 대한 Score 기반 해석: PnP 반복과 Score-based 디노이징 사이의 직접적인 수학적 연결을 확립하여, 사전 훈련된 SBDM 을 PnP 알고리즘 내에서 직접 사용할 수 있는 이론적 근거를 마련했습니다.
2. SGPnP 프레임워크 제안: 잡음 주입을 통해 ill-posed 문제에서 강건한 생성적 복원을 가능하게 하는 새로운 알고리즘을 제안했습니다. 이는 기존 PnP 가 실패하는 대규모 결손 (Large-mask inpainting) 등의 문제에서 특히 효과적입니다.
3. 이론적 보장 (Theoretical Guarantees):
   • 안장점 탈출 (Saddle-point Escape): 주입된 잡음이 목적 함수의 Strict Saddle Point 에서 벗어나게 하는 것을 수학적으로 증명했습니다 (비볼록 최적화 관점).
   • 수렴성 (Convergence): Annealed Noise Schedule (점점 감소하는 잡음 스케줄) 하에서, 알고리즘이 원래의 비부드러운 (un-smoothed) 목적 함수의 정상점 (Stationary Point) 으로 수렴함을 증명했습니다.

4. 실험 결과 (Results)

• 데이터셋: Flickr-Faces-HQ (FFHQ, 자연 이미지) 와 fastMRI (다중 코일 뇌 MRI).
• 비교 대상:
  • 기존 PnP (DRUNet 기반: DPIR, PnP-ADMM, PGM)
  • 기존 확률적 PnP (SNORE)
  • 확산 기반 솔버 (DiffPIR, DPS)
• 주요 성과:
  • 정량적 지표: 인페인팅, 초해상도, 디블러링, MRI 재구성 등 모든 작업에서 기존 PnP 방법론과 확률적 PnP (SNORE) 를 능가하는 PSNR, SSIM, LPIPS 점수를 기록했습니다.
  • 정성적 결과: 특히 대규모 박스 마스크 인페인팅 (Box Inpainting) 과 같이 정보가 크게 결손된 경우, 기존 PnP 는 불완전한 복원을 보여주거나 아티팩트가 발생했으나, SGPnP 는 의미 있는 시맨틱 콘텐츠를 생성하여 현실적인 이미지를 복원했습니다.
  • 확산 솔버 대비 성능: 역확산 샘플링을 수행하는 DiffPIR/DPS 와 비교해도 경쟁력 있는 성능을 보였으며, 최적화 기반 접근법의 효율성을 유지하면서 생성 모델의 성능을 달성했습니다.
  • Ablation Study:
    • Score 기반 사전의 효과: CNN 기반 디노이저를 SBDM 으로 교체하는 것만으로도 (SDPnP) 성능이 향상됨을 확인했습니다.
    • 잡음 주입의 효과: 동일한 SBDM 사전이라도 잡음 주입 (SGPnP) 을 적용하면 ill-posed 문제 해결 능력이 획기적으로 향상됨을 입증했습니다.
    • 잡음 수준 분리: 주입 잡음과 조건부 잡음 수준을 분리하는 것이 성능 향상에 기여함을 확인했습니다.

5. 의의 및 결론 (Significance)

• 최적화와 생성 모델의 융합: 본 논문은 최적화 기반 PnP 프레임워크와 최신 생성 모델 (Diffusion Models) 간의 간극을 메웠습니다. 역확산 샘플링을 반복하지 않고도 SBDM 을 효율적인 디노이저로 활용할 수 있음을 보였습니다.
• 심한 ill-posed 문제 해결: 기존 PnP 가 실패했던 극단적으로 정보가 부족한 상황에서도, 잡음 주입을 통한 확률적 탐색과 SBDM 의 강력한 생성 능력을 결합하여 안정적인 복원을 가능하게 했습니다.
• 이론적 토대: PnP 방법론에 대한 최초의 안장점 탈출 (Saddle-point escape) 이론적 분석을 제공하여, 왜 잡음 주입이 성능 향상에 기여하는지에 대한 깊은 이해를 제공했습니다.

결론적으로, SGPnP 는 이미지 역문제 해결을 위해 학습된 생성 모델의 표현력과 최적화 알고리즘의 구조적 강건성을 결합한 차세대 프레임워크로, 특히 정보가严重하게 결손된 복잡한 복원 작업에서 혁신적인 성능 향상을 이끌어냅니다.