Local Epistemic Uncertainty Guided Active Sampling for Plug-and-play Diffusive Image Restoration

기술 요약: 플러그 앤 플레이 확산 기반 이미지 복원을 위한 국소적 인식론적 불확실성 유도 능동 샘플링

1. 문제 정의

확산 모델 기반 이미지 복원(DMIR) 방법론은 사전 학습된 무조건부 확산 모델을 범용적인 사전 지식(prior)으로 활용하여 역문제(예: 초해상도, 디블러링, 압축 센싱)를 해결함으로써 최첨단 성능을 달ر 성취했습니다. 그러나 기존의 DMIR 접근 방식은 다음과 같은 두 가지 주요 한계점을 가집니다.

• 공간적 경직성(Spatial Rigidity): 현재의 방법들은 일반적으로 전역적으로 고정된 데이터 제약 가중치를 사용합니다. 이는 실제 세계의 비균일한 열화 특성을 반영하지 못합니다. 결과적으로, 복원이 가능한 영역에서는 과도하게 강한 사전 지식이 적용되어 디테일이 뭉개지는 오버 스무딩(over-smoothing) 현상이 발생하며, 심하게 열화된 영역에서는 제약이 부족하여 구조적 왜곡과 아티팩트가 발생합니다.
• 시간적 중복성(Temporal Redundancy): 기존 방법들은 역방향 샘플링 과정에서 균일한 타임 스텝을 사용합니다. 이러한 정적 스케줄링은 확산 궤적의 동적인 진화를 무시합니다. 이는 생성의 안정적인 단계에서 계산적 중복을 초래하며, 적응형 조절 없이 단계를 줄일 경우 이산화 오류(discretization errors)를 유발하여 수렴 속도와 복원 품질을 저하시킵니다.

2. 방법론: LEADer

저자들은 LEADer(Local Epistemic Uncertainty Guided Active Sampling)를 제안합니다. 이는 **국소적 인식론적 불확실성(local epistemic uncertainty)**을 기반으로 확산 역과정을 동적으로 조절하는 플러그 앤 플레이 프레임워크입니다. 이 프레임워크는 두 가지 차원에서 작동합니다.

A. 국소적 인식론적 불확실성 정량화
고정된 휴리스틱에 의존하는 대신, LEADer는 각 공간 위치와 타임 스텝에서의 모델 불확실성을 정량화합니다.

• 이 방법은 사전 학습된 확산 모델의 스코어 함수(score function)로부터 **피셔 정보 행렬(Fisher Information Matrix, FIM)**을 도출합니다.
• 사후 분포(posterior distribution)의 2차 테일러 전개를 사용하여, 저자들은 FIM과 사후 공분산 행렬($\Sigma_t$) 사이의 관계를 분석적으로 연결합니다.
• $\Sigma_t$의 대각 성분은 국소적 인식론적 불확실성의 척도로 사용됩니다. 높은 불확실성은 모델이 확신이 없는 영역(흔히 열화 연산자의 영공간/null space에 해당하는 영역)을 나타내며, 낮은 불확실성은 생성된 구조에 대한 높은 확신을 나타냅니다.

B. 공간적 변조: 불확실성 보정 사전 지식 변조(UCPM)
공간적 불균일성을 해결하기 위해, LEADer는 물리적 사전 지식(예: 그래디언트 또는 구조적 제약)의 강도를 동적으로 조절하는 UCPM을 도입합니다.

• 이 문제는 $\Sigma_t$가 공간적으로 불균일한 프리컨디셔너(preconditioner) 역할을 하는 최대 사후 확률(MAP) 최적화 문제로 정식화됩니다.
• 메커니즘: 불확실성이 낮은 영역(모델 확신도가 높은 영역)에서는 확산 모델이 생성한 고주파 디테일을 보존하기 위해 사전 지식 변조를 약화시킵니다. 불확실성이 높은 영역에서는 구조적 편차를 억제하고 강력한 규제를 적용하기 위해 변조를 강화합니다.
• 데이터 일관성: 이 방법은 변조된 추정치를 열화 연산자의 영공간(null space)에 투영함으로써 엄격한 데이터 충실도를 보장하며, 이를 통해 사전 지식의 조정이 알려진 랭지 공간(range-space) 성분을 손상시키지 않도록 합니다.

C. 시간적 가속: 상태 인지 궤적 프루닝(SATP)
계산적 중복을 해결하기 위해, LEADer는 적응형 스텝 크기 조절을 가능하게 하는 SATP를 도입합니다.

• 메커니즘: 이 방법은 샘플링 궤적의 국소 절단 오차(local truncation error)를 불확실성 행렬의 트레이스(trace)인 $U_t = \text{Tr}(\Sigma_t)$와 연관시킵니다.
• 적응형 건너뛰기(Adaptive Skipping): 미리 정의된 "정보 손실 예산(information loss budget)"($B$)이 허용 가능한 국소 오차를 제한합니다. 현재 불확실성이 낮으면(안정적인 복원 단계임을 나타냄), 알고리즘은 더 큰 허용 가능한 스텝 크기($\Delta t_{Active}$)를 계산하여 중복된 반복 과정을 건너뜁니다. 반대로, 불확실성이 높은 영역에서는 정확도를 유지하기 위해 스텝 크기를 줄입니다.
• 이론적 보장: 저자들은 그뢴월 부등식(Grönwall's inequality)을 통해 이 전략이 비균일 샘플링 환경에서도 결정론적인 누적 오차 상한을 유지한다는 이론적 증명을 제공합니다.

3. 주요 기여

1. 새로운 관점: 본 논문은 국소적 인식론적 불확실성 관점을 도입하여 역방향 샘플링 과정을 가이드함으로써, 기존 DMIR 방법들의 시공간적 한계를 식별하고 해결합니다.
2. 능동 샘플링 프레임워크: 적응형 사전 지식 변조를 위한 UCPM과 상태 인지 궤적 프루닝을 위한 SATP를 통합한 LEADer를 제안합니다. 이 프레임워크는 네트워크 파라미터의 재학습 없이도 국소적 왜곡과 샘플링 중복을 선제적으로 해결합니다.
3. 플러그 앤 플레이 호환성: 이 방법은 추론 과정에서 사전 학습된 모델의 사후 통계량만을 활용하므로, 기본 네트워크 파라미터를 수정하지 않고도 다양한 기존 DMIR 베이스라인(예: DDNM, DDPG, ProjDiff, SITCOM)에 원활하게 통합될 수 있습니다.
4. 이론적 및 경험적 검증: 본 논문은 엄격한 데이터 일관성과 유계된 오차 누적에 대한 이론적 증명을 제공하며, 동시에 향상된 품질과 효율성을 입증하는 광범적인 실험을 수행했습니다.

4. 실험 결과

저자들은 두 가지 데이터셋(CelebA-HQ 1K 및 ImageNet 1K)에 대해 다섯 가지 이미지 복원 작업(4x 초해상도, 가우시안 디블러링, 모션 디블러링, 압축 센싱)에서 LEADer를 평가했습니다.

• 성능 향상: 다양한 베이스라인과 결합했을 때, LEADer는 일관되게 복원 품질을 개선했습니다. CelebA-HQ에서 평균 PSNR은 0.88%에서 2.38% 향상되었으며, LPIPS(지각적 오차)는 7.13%에서 12.78% 감소했습니다.
• 효율성: SATP를 통해 중복 단계를 건너뜀으로써, 동일한 초기 단계를 사용하는 베이스라인 대비 평균 **3.04%에서 8.42%**의 샘플링 시간을 단축했습니다.
• 리소스 오버헤드: 이 방법은 무시할 만한 수준의 메모리 오버헤드(GPU 메모리 사용량 증가폭 0.07% ~ 0.28%)만을 도입했습니다.
• 절제 연구(Ablation Studies): 실험을 통해 UCPM과 SATP가 상호 보완적임을 확인했습니다. UCPM은 샘플링 궤적을 안정화하여 SATP가 품질 저하 없이 안전하게 단계를 건너뛸 수 있도록 합니다.

5. 의의 및 주장

본 논문은 LEADer가 현재의 제로샷(zero-shot) DMIR 방법들이 가진 근본적인 한계, 즉 이미지 생성의 동적인 특성에 적응하지 못하는 정적이고 균일한 제약에 의존한다는 점을 해결한다고 주장합니다. 국소적 인식론적 불확실성을 활용함으로써 저자들은 다음을 달성할 수 있음을 보여줍니다:

• 지각-왜곡 트레이드오프(Perception-Distortion Trade-off)의 균형: 사전 지식의 강도를 동적으로 조절함으로써, 확신이 있는 영역에서는 미세한 디테일을 보존하고 불확실한 영역에서는 아티팩트를 억제할 수 있습니다.
• 수렴 가속화: 상태 인지 궤적 프루닝은 안정적인 수렴에 필요한 결정론적 오차 상한을 희생하지 않으면서 계산 비용을 줄입니다.
• 일반화: 플러그 앤 플레이 모듈로서, LEADer는 다양한 역 이미지 생성 작업에 걸쳐 다양한 최첨단 DMIR 방법의 성능을 향상시키며, 확산 기반 이미지 복원의 품질과 효율성을 모두 개선할 수 있는 통합된 솔루션을 제공합니다.