Improving Diffusion Posterior Samplers with Lagged Temporal Corrections for Image Restoration

기술 요약: 지연된 시간 보정 (LAMP) 을 통한 확산 후사 분포 샘플러 개선

문제 제기

확산 기반 후사 분포 샘플링 (PS) 은 이미지 역문제 (예: 디블러링, 초해상도, 인페인팅) 해결을 위한 선도적인 프레임워크로 부상했습니다. 이러한 방법들은 학습된 확산 사전 지식과 측정 제약을 결합하여 $y = Kx_0 + e$인 조건에서 후사 분포 $p(x_0|y)$를 샘플링합니다. DiffPIR 및 DDRM 과 같은 표준 PS 공식은 확산 모델의 디노이징 단계와 데이터 일관성 보정 $D_t$를 사용하여 현재 추정치 $x_t$를 반복적으로 업데이트합니다.

그러나 저자들은 이러한 표준 접근법에서 근본적인 구조적 한계를 식별했습니다: 즉, 모든 역방향 반복 단계마다 재계산되는 순간적, 단계별 데이터 일관성 추정치($D_t$) 에 의존한다는 점입니다. 이러한 설계는 역방향 역학에서 시간적 변동성과 불안정성을 초래합니다. 동역학 시스템 관점에서 표준 PS 업데이트는 확산 역학의 1 차 이산화와 잔차 보정 항으로 구성됩니다. 이 잔차 항은 시간에 따라 변하는 강제 신호로 작용하여 궤적을 교란시키므로, 무조건부 생성에 사용되는 고차 솔버 (예: DPM-Solver) 가 가정하는 매끄럽고 잘 정의된 연속 시간 역학에서 벗어나게 만듭니다. 결과적으로 표준 PS 방법의 시간적 기억 부재는 일관성 없는 업데이트와 비최적의 안정성, 특히 잘 정의되지 않은 (ill-posed) 영역에서 이를 초래합니다.

방법론: LAMP

본 논문은 지연된 시간 보정을 도입하여 이러한 불안정성을 해결하는 모듈형 플러그인 방식인 **LAMP(LAgged Multistep Posterior)**를 제안합니다.

1. 동역학적 재해석: 저자들은 표준 PS 업데이트를 강제 확산 역학의 1 차 이산화로 재해석합니다. 업데이트 $x^{PS}_{t-\Delta t}$는 측정 인지형 목표 $D_t$에 의해 주도되는 1 차 지수 단계와 불일치 $(D_t - \hat{x}_{0|t})$에 비례하는 잔차 보정 항으로 분해될 수 있음을 보여줍니다.
2. 2 차 이산화: 안정성을 개선하기 위해 LAMP 는 기본 역학의 2 차 이산화를 적용합니다. 매끄러운 ODE 를 가정하는 표준 고차 솔버와 달리, LAMP 는 역문제에 내재된 잔차 의존적 강제 항을 명시적으로 고려합니다.
3. 지연 보정 메커니즘: 핵심 혁신은 연속 추정치의 변동을 기반으로 한 시간 보정의 도입입니다. LAMP 업데이트는 2 차 이산화와 잔차 보정을 결합합니다:
   $$x^{LAMP}_{t-\Delta t} = x^{PS}_{t-\Delta t} + \alpha_{t-\Delta t} A_1(h) \gamma \frac{D_t - D_{t+\Delta t}}{h_{prev}}$$
   여기서 항 $(D_t - D_{t+\Delta t})/h_{prev}$는 이산 시간 미분으로 작용합니다. 음수 계수 $\gamma$를 선택함으로써, 이 방법은 외삽하는 대신 현재 추정치를 이전 추정치 쪽으로 약간 끌어당기는 지연 보정을 적용합니다. 이는 궤적을 매끄럽게 하고 순간적 보정으로 인한 변동성을 상쇄합니다.
4. 구조 보존: 중요하게도, 저자들은 (Proposition 1) LAMP 가 표준 후사 분포 샘플링 구조를 보존함을 증명합니다. 업데이트는 $x_{t-\Delta t} = \alpha_{t-\Delta t} \tilde{D}_t + \sigma_{t-\Delta t} \epsilon_\Theta(x_t, t)$라는 정형화된 형태로 다시 쓸 수 있으며, 여기서 $\tilde{D}_t$는 현재 및 이전 데이터 일관성 추정치의 볼록 결합으로 정의된 시간 필터링 추정치입니다: $\tilde{D}_t = (1-\beta_t)D_t + \beta_t D_{t+\Delta t}$. 이는 LAMP 가 디노이저나 학습 절차를 수정하지 않고도 기존 PS 백본 (예: DiffPIR, DDRM) 과 호환되는 모듈형 확장임을 보장합니다.

이론적 분석

본 논문은 지연 보정의 편차 - 분산 트레이드오프를 특징짓는 1 단계 위험 분석(Proposition 2) 을 제공합니다. 분석 결과, 시간적 평활화로 인한 추정 오차 분산의 감소가 지연으로 인한 편차 (후사 목표의 드리프트로 인한) 보다 클 때 LAMP 가 역방향 전이를 개선함을 보여줍니다. 저자들은 이 조건이 일반적으로 확산 과정의 후기 단계에서 충족됨을 입증하여, 짧은 웜업 단계 후 LAMP 를 활성화하는 하이브리드 전략을 동기화합니다.

주요 기여

• 구조적 불일치 식별: 저자들은 역문제 후사 분포 샘플링에 존재하는 잔차 의존적 강제 항을 포착하지 못해 시간적 불안정성을 초래하는 확산 솔버에 대한 표준 ODE 기반 해석을 강조합니다.
• LAMP 프레임워크: 후사 분포 샘플링 업데이트에 지연된 시간 보정을 통합하는 2 차 보정 방식인 LAMP 를 제안합니다.
• 모듈성 및 호환성: LAMP 는 재학습이나 전방 모델 변경 없이 기존 PS 방법 (예: DiffPIR, DDRM) 위에 플러그인으로 구현될 수 있는 유효한 후사 분포 샘플러임이 증명되었습니다.
• 이론적 정당성: 국소 오차 분석은 지연 보정이 위험을 감소시키는 조건, 즉 인접한 후사 목표가 추정 노이즈에 비해 매끄럽게 진화할 때를 확립합니다.

실험 결과

저자들은 FFHQ, ImageNet, CelebA 세 가지 데이터셋에 걸쳐 표준 선형 역문제 (가우시안 디블러링, 모션 디블러링, 초해상도) 에서 LAMP 를 평가했습니다. LAMP 강화 버전 (DiffPIR-LAMP, DDRM-LAMP) 을 베이스라인 및 단순 고차 변형 (1M, 2M) 과 비교했습니다.

• 성능: LAMP 는 강력한 베이스라인보다 일관되게 재구성 품질을 향상시킵니다. 잡음 환경 ($\sigma_y = 0.05$) 에서 DDRM-LAMP는 모든 작업에서 PSNR, SSIM, LPIPS 에서 일관된 개선을 보이며 전체적으로 가장 좋은 성능을 달성했습니다. DiffPIR-LAMP 또한 특히 디블러링 및 초해상도에서 개선을 보였습니다.
• 효율성: 이러한 개선은 디노이징 평가 횟수 (NFE) 를 증가시키지 않고 달성되었으며, 기본 방법의 계산 비용을 유지했습니다 (예: DDRM 의 경우 20 NFE, DiffPIR 의 경우 100 NFE).
• 단순 고차 방법과의 비교: 특정 잔차 보정 없이 단순한 1 차 및 2 차 ODE 이산화는 일관되게 성능을 저하시켜, 제안된 지연된 잔차 보정이 개선의 핵심 요소임을 확인시켜 주었습니다.
• 정성적 결과: 시각적 비교 결과, LAMP 는 베이스라인에 비해 더 매끄러운 질감, 더 선명한 경계, 그리고 세부 사항 (예: CelebA 의 헤어 질감, Iguana 이미지의 비늘) 의 더 나은 보존을 보여주며, 입자감과 블러링 아티팩트를 감소시켰습니다.

중요성 및 주장

본 논문은 측정 인지형 추정치의 시간 필터링이 확산 기반 역 솔버를 안정화시키는 중요하지만 이전에는 충분히 탐구되지 않은 메커니즘이라고 주장합니다. 후사 분포 샘플링을 잔차 의존적 강제 항을 가진 동역학 시스템으로 재정의함으로써, 저자들은 간단한 지연 보정이 재구성 품질과 안정성을 크게 향상시킬 수 있음을 입증했습니다.

이 접근법의 중요성은 모듈성에 있습니다: LAMP 는 새로운 사전 지식, 재학습, 또는 복잡한 조건부 전략을 요구하지 않습니다. 대신 기존 최첨단 후사 분포 샘플러에 대한 범용적이고 저비용의 개선을 제공합니다. 저자들은 LAMP 가 백본 방법의 한계 (예: 근본적인 모델 불일치를 보상할 수 없음) 를 계승하며, 그 효과는 기본 샘플러의 보정 시퀀스의 시간적 행동에 의존한다고 겸손하게 언급합니다. 그러나 결과는 안정적인 베이스라인의 경우 LAMP 가 미세 구조에 특히 데이터 충실도에서 일관되고 지각적으로 의미 있는 개선을 제공함을 시사합니다.