기술 요약: MMSE 를 넘어선 ProxiMAP 을 통한 PnP 복원 향상

문제 제기

Plug-and-Play(PnP) 방법은 근사 분할 알고리즘에서 처리가 불가능한 최대 사후 확률 (MAP) 제거기를 학습된 최소 평균 제곱 오차 (MMSE) 제거기로 대체함으로써 이미지 역문제 (예: 디블러링, 인페인팅) 해결의 표준이 되었습니다. 효과적이기는 하지만, 이 대체는 근본적인 불일치를 초래합니다. MMSE 추정기는 사후 평균을 근사하는데, 이는 '평균 회귀 (regression-to-the-mean)'로 인해 과도하게 평활화된 재구성을 경향하는 반면, PnP 의 변분 형식화는 이론적으로 MAP 추정기를 요구합니다.

최근 확산 모델의 점수 (scores) 를 활용하여 직접 MAP 를 목표로 함으로써 이 격차를 해소하려는 시도는 실제로 문제점이 드러났습니다. 본 논문은 학습된 점수를 사용할 때, MAP 를 목표로 하는 반복 과정이 현실적인 이미지로 수렴하지 않음을 보여줍니다. 대신, 이러한 과정은 실제 데이터 분포와 일치하지 않는 학습된 분포의 고밀도 영역으로 편향되어 '만화 같은' 아티팩트 (평평한 색상, 과장된 가장자리) 를 초래합니다. 논문은 이러한 실패가 MAP 추정 자체의 고유한 문제가 아니라, 학습된 모델에서의 점수 근사 오차의 결과라고 주장합니다.

방법론: ProxiMAP

저자들은 학습된 제거기의 신뢰성 영역 내에서 작동하도록 설계된 반복적 MAP 근사 알고리즘인 ProxiMAP을 소개합니다. 핵심 설계 원칙은 노이즈 매칭입니다. 각 반복 $k$에서 현재 반복체의 잔여 노이즈 수준은 제거기가 훈련되고 쿼리된 노이즈 수준 $\sigma_k$와 일치해야 합니다.

주요 메커니즘

1. 노이즈 스케줄 및 잔여 추적: 근사 목적 함수에 대한 표준 경사 하강법과 달리, ProxiMAP 은 노이즈 스케줄이 반복체의 상태와 분리되지 않고 반복체 $x_k$의 잔여 노이즈 $\sigma_k$를 명시적으로 추적합니다. 이 알고리즘은 제거기가 항상 훈련 분포와 일치하는 노이즈 통계를 가진 입력으로 쿼리되도록 제약을 부과합니다.
2. 암시적 조기 종료: 이 일치를 유지함으로써 알고리즘은 학습된 점수가 신뢰할 수 없는 영역 (실패 영역) 으로 반복체가 편향되는 것을 방지합니다. 이는 원칙적인 조기 종료 메커니즘으로 작용합니다. 반복체가 데이터 다양체 (manifold) 에 접근함에 따라 속도가 느려져 만화 효과를 유발하는 불일치 쿼리를 피합니다.
3. 알고리즘적 공식화: ProxiMAP 은 Pesme 등 (2025) 의 연구에 기반한 평활화된 경사 하강 재귀를 수정하여 $\sigma_k$(반복체의 노이즈 수준) 가 제거기의 조건부 노이즈와 같아지도록 합니다. 단계 크기 $\gamma_k$와 노이즈 스케줄은 이 제약을 만족하도록 유도되어, 제거기가 항상 '분포 내 (in-distribution)'에 머무르도록 보장합니다.
4. 하이브리드 변형 (Fast ProxiMAP): 모든 PnP 반복에서 MMSE 를 ProxiMAP 으로 완전히 대체하는 것이 계산적으로 비용이 많이 든다는 점을 인식하여, 저자들은 하이브리드 접근법을 제안합니다. 이 변형은 초기 반복 (반복체가 다양체에서 멀고 점수 정확도가 덜 중요한 영역) 에서는 표준 MMSE 제거기를 사용하고, 후기 반복 (반복체가 다양체에 가깝고 제거기가 가장 신뢰할 수 있는 영역) 에만 ProxiMAP 으로 전환합니다. 이는 계산 비용의 일부분으로 동등하거나 더 우수한 결과를 달성합니다.

주요 기여

1. 학습된 사전 지식을 통한 MAP 추정 진단: 정확한 점수를 가진 가우시안 혼합 모델 (GMM) 과 학습된 점수를 가진 확산 모델을 사용한 통제된 비교를 통해, 논문은 MAP 타겟팅에서의 '만화 효과'가 MAP 추정 과정 자체에서 비롯된 것이 아니라, 확산 훈련의 암묵적 편향이 고밀도 영역을 이동시킴으로써 발생함을 규명합니다.
2. ProxiMAP 알고리즘: 표준 PnP 프레임워크에서 MMSE 제거기를 교체할 수 있는 모듈식, 바로 적용 가능한 솔루션입니다. 이는 제거기를 분포 내로 유지하고 암시적 조기 종료를 제공하는 단일 원칙 (노이즈 매칭) 에서 유도됩니다.
3. 체계적 개선: 이 방법은 여러 PnP 알고리즘 (DPIR, DiffPIR, DAPS), 작업 (디블러링, 인페인팅, 초해상도, 위상 복원), 그리고 데이터셋 (FFHQ, ImageNet) 에 걸쳐 일관된 지각적 품질 (LPIPS) 향상을 보여주면서도 경쟁력 있는 왜곡 (PSNR) 을 유지합니다.
4. 효율적인 하이브리드 변형: 'Fast ProxiMAP' 변형은 전체 교체 ProxiMAP 의 성능을 맞추거나 능가하면서도 계산 비용을 크게 절감합니다 (예: 일부 구성에서 함수 평가 횟수를 160 회에서 27 회로 감소).

실험 결과

저자들은 다양한 역문제에 걸쳐 FFHQ 및 ImageNet 데이터셋에서 ProxiMAP 을 평가했습니다.

• 독립적 제거: ProxiMAP 은 지각 - 왜곡 트레이드오프 곡선에서 조건부 MMSE 및 기타 결정론적 확산 샘플러 (DDIM, Flow Matching, InDI) 보다 우수한 성능을 발휘하여, 동등한 PSNR 에 대해 더 낮은 LPIPS(선명한 이미지) 를 달성합니다.
• PnP 통합: DPIR, DiffPIR, DAPS 에 통합되었을 때, ProxiMAP 변형은 파레토 프론트를 최적 모서리 (높은 PSNR, 낮은 LPIPS) 쪽으로 일관되게 이동시킵니다.
• 하이브리드 효율성: 최종 PnP 반복에서만 ProxiMAP 으로 전환하는 Fast ProxiMAP 변형은 전체 교체 변형과 동등한 결과를 달성하지만 네트워크 평가 횟수는 극적으로 줄입니다. 예를 들어, DiffPIR 에서 함수 평가 횟수 (NFEs) 를 160 회에서 27 회로 줄이면서도 재구성 품질을 유지하거나 향상시킵니다.
• 실패 모드 해결: 전체 ProxiMAP 대체가 고주파 아티팩트 (예: FFHQ 의 초해상도) 를 초래한 경우, 하이브리드 변형은 초기 불안정 반복 동안 견고한 MMSE 추정기에 의존함으로써 이러한 문제를 성공적으로 해결했습니다.

중요성 및 주장

본 논문은 ProxiMAP 이 PnP 에서 사용되는 MMSE 제거기와 기초 변분 형식에서 요구되는 MAP 추정기 사이의 오랜 불일치에 대한 실용적인 해결책을 제공한다고 주장합니다. 학습된 분포의 잠재적으로 잘 정의되지 않은 MAP 로의 수렴을 강요하는 대신, ProxiMAP 은 최적화 궤적을 제거기의 훈련 분포 내에 유지함으로써 학습된 점수의 신뢰성을 활용합니다.

저자들은 그들의 기여가 기존 PnP 계열을 강화하는 모듈식 구성 요소이며, 기초 확산 모델의 재훈련이나 새로운 전문 역문제 솔버 개발을 요구하지 않는다고 강조합니다. 그들은 ProxiMAP 이 지각적 품질을 향상시키지만, 재구성은 여전히 사전 훈련된 확산 모델의 편향을 계승한다는 점을 지적합니다. 이는 시각적 타당성이 정확성의 유일한 결정 요인이 되어서는 안 되는 안전-중요 응용 분야에서 특히 중요한 주의사항입니다. 이 연구는 부정확한 점수로 정확한 MAP 수렴을 추구하는 것보다, 반복체의 노이즈 수준과 제거기의 훈련 영역 간의 관계를 관리하는 것이 더 효과적인 설계 레버라고 시사합니다.