본문 바로가기
Paper Review/Inverse Problem

[논문 리뷰] PSLD: Solving Linear Inverse Problems Provably via Posterior Sampling with Latent Diffusion Models

by kongshin 2025. 4. 3.

[논문]


Summary

  • 기존의 DPS (Diffusion Posterior Sampling) - Pixel space에서 동작
    • LDM (Latent Diffusion Models)로 확장하면 잘 작동 X

 

  • Latent 공간에서 inverse problems을 이론적으로 정확한 posterior sampling 보장 입증
  • PSLD - Posterior Sampling with Latent Diffusion
  • Pre-trained LDMs 사용 가능 (Stable Diffusion 등)
    • inverse problem에 대해 훈련 없이 바로 적용 가능

Method

[Diffusion process in latent space]

  • $dz=f(z,t)dt + g(t)dw$
  • timesteps t에 따라 latent z에 noise 추가하는 Forward SDE

 

[DPS in latent space]

  • $p(y|z_t) \approx p(y|x_0=D(E[z_0|z_t])$
  • Encoder $E$: $\mathbb{R}^d$→$\mathbb{R}^k$ → many-to-one mapping
    • 다른 measurement $y$를 따르는 $x_0$이 동일한 $z_0$으로 encoding 될 수 있음
  • 하나의 $z_0$가 여러 $x_0$ (즉, 여러 $y$)에 대응됨 ⇒ $y$ gradient가 분산됨
  • 즉, $\nabla_{z_t}logp(y|z_t)$의 방향이 애매해짐 ⇒ 비일관성, 불안정성

 

[PSLD]

  • Vanilla DPS - latent space에서 gradient 불안정하게 작용
  • 2개의 추가 loss term 도입하여 안정화 ⇒ Encoder & Decoder 수정 X
    1. Goodness Term
    2. Gluing Term

 

  • Goodness Modified Latent DPS (GML-DPS)

    1. 기존 DPS 방식의 measurement 일치 조건
    2. Goodness term
      •   목표: latent $z_0$가 Decoder-Encoder 조합의 fixed point가 되도록 penalty 부여
      •   $z_0\approx E(D(z_0))$를 만족해야함
      •   latent → image → latent로 돌아왔을 때 거의 변하지 않은 latent가 좋은 latent
        •   ⇒ 해당 $z_0$은 실제 iamge manifold에 가깝다고 볼 수 있음

 

  • 해당 Method를 저자가 refer
  • $z_0$가 image manifold 위에 있어도 $D(z_0)$가 boundary에서 실제 $x_0^*$와 자연스럽게 이어지지 않을 수 있음 ⇒ PSLD 제안
    • $z_0$ - Unconditional sampling
    • $x^*_0$ - 관측된 부분 ground truth
    • $D(z_0)$ + $x^*_0$ → 최종 inverse problems의 output

 

  • Posterior Sampling with Latent Diffusion (PSLD)
    • if $A$: masked operation $\in \mathbb{R}^{m\times d}$
      • $A^TAx^*_0$: 관측된 부분만 그대로 유지 & 나머지는 0이 됨
      • $A^T$: 원래 dimension에 맞게 변경하는 효과
    • Encoder 안의 term의 의미
      • 생성 image에 관측된 실제 ground truth를 gluing(붙임)
    • latent $z_0$가 gluing image의 latent 표현과 가깝도록 penalty term 추가

⇒ 관측된 영역과 생성된 영역이 경계에서 부드럽게 연결되도록 latent space에서 유도하는 방법

⇒ discontinuity 방지

 

Experiments

  • DPS vs PSLD