[논문]
- Solving Linear Inverse Problems Provably via Posterior Sampling with Latent Diffusion Models
- NeurIPS 2023
- Citations:73
- [https://arxiv.org/abs/2307.00619]
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
- Goodness Term
- Gluing Term
- Goodness Modified Latent DPS (GML-DPS)
- 기존 DPS 방식의 measurement 일치 조건
- 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