Diffusion Posterior Sampling for Linear Inverse Problem Solving: A Filtering Perspective

[bd3dowling]

Dou, Zehao, and Yang Song. Diffusion Posterior Sampling for Linear Inverse Problem Solving: A Filtering Perspective. 2023. openreview.net, https://openreview.net/forum?id=tplXNcHZs1&referrer=%5Bthe%20profile%20of%20Yang%20Song%5D(%2Fprofile%3Fid%3D~Yang_Song1).

[bd3dowling]

Code: https://github.com/ZehaoDou-official/FPS-SMC-2023

[bd3dowling]

Paste from obsidian:

• Formulates DDPM (ancestral sampling) and DDIM under one framework (in terms of discrete sampling).
• They use DDIM parameterizing $\sigma{k}$ by: $$ \sigma{k} = c \cdot \sqrt{ \beta{k}. \cdot \frac{1-\bar{\alpha}{k-1}}{1 - \bar{\alpha}_{k}} } $$
• Discretize over time steps application of the linear noise operator to form measurement sequence so that: $$\begin{align} \mathbf{x}{k} &= a{k}\cdot \mathbf{x}{k-1} + b{k}\mathbf{z}{k} \ \mathbf{y}{k} &= a{k}\cdot \mathbf{y}{k-1} + b{k}\cdot\mathbf{Az}{k}\ : k \geq 1 \ \implies \mathbf{y}{k} &\sim \mathcal{N}(\mathbf{Ax}{k}, c{k}^2\sigma^2\mathbf{I}),\ \dots c{k}=a{1}\cdots a{k} \end{align}$$
• This sequence can be equivalently generated forwards (using above) or backwards (using DDIM).
• Use these "observations" as basis for filtering problem, computing posterior distribution: $$\begin{align} p{\theta}(\mathbf{x}{k-1} \mid \mathbf{x}{k}, \mathbf{y}{k-1}) &\propto p{\theta}(\mathbf{x}{k-1} \mid \mathbf{x}{k}) \cdot p{\theta}(\mathbf{y}{k-1} \mid \mathbf{x}{k-1}) \end{align}$$ where $$\begin{align} p{\theta}(\mathbf{x}{k-1} \mid x{k}) &= \mathcal{N}(u{k}\hat{\mathbf{x}}{0}(\mathbf{x}{k}) + v{k}\mathbf{s}{\theta}(\mathbf{x}{k}, t{k}), w_{k}^2\mathbf{I}) \end{align}$$ with $$\begin{align} \hat{\mathbf{x}}{0}(\mathbf{x}{k}) &:= \frac{\mathbf{x}{k} + d{k}^2\mathbf{s}{\theta}(\mathbf{x}{k}, t{k})}{c{k}} \end{align}$$ the conditional expectation of $\mathbf{x}{0} \mid \mathbf{x}{k}$ (computed via Tweedie's formula).
• FPS approximates $p{\theta}(\mathbf{x}{k+1} \mid y{k:N})$ by $p{\theta}(\mathbf{x}{k+1} \mid \mathbf{y}{k+1:N})$ so to skip the update step of the filtering process.
• Generalized to continuous-time Markov chains with following SDE: $$ d\mathbf{x}{t} = \left[ - \frac{\beta(t)}{2}\mathbf{x}{t} - \beta(t)\nabla{\mathbf{x}{t}}\log p{\theta, t}(\mathbf{x}{t} \mid \mathbf{y}{t}) \right] dt + \sqrt{ \beta(t) }d \overline{\mathbf{W}}{t} $$ Hence, FPS approximates $\nabla{\mathbf{x}{t}}\log p{\theta, t}(\mathbf{x}{t} \mid \mathbf{y})$ by $\nabla{\mathbf{x}{t}}\log p{\theta, t}(\mathbf{x}{t} \mid \mathbf{y}_{t})$.
• Improve approximation with SMC.

  \begin{algorithm}
  \caption{Framework of FPS-SMC}
  \begin{algorithmic}
  \Input $N, \mathbf{y}, M, \mathbf{s}_{\theta}(\cdot, \cdot)$
  \State Sample the $\{\mathbf{y}_k\}$ sequence from $q(\mathbf{y}_{1:N} \mid \mathbf{y}_0)$.
  \State Sample $M$ i.i.d. samples $\mathbf{x}_N^{(j)} \sim p_{\theta}(\mathbf{x}_N \mid \mathbf{y}_N)$ for $j \in [M]$.
  \For{$k \gets N$ \To $1$}
      \State Generate $M$ i.i.d. samples $\overline{\mathbf{x}}_{k-1}^{(j)} \sim p_\theta(\mathbf{x}_{k-1} \mid \mathbf{x}_k^{(j)}, \mathbf{y}_{k-1})$ for $j \in [M]$.
      Randomly pick $M$ samples with replacement from $\left\{\overline{\mathbf{x}}_{k-1}^{(j)}\right\}_{j \in [M]}$. Denote the $M$ new samples as $\left\{\mathbf{x}_{k-1}^{(j)}\right\}_{j \in [M]}$.
  \EndFor
  \Output $\mathbf{x}_0$ which is uniformly sampled from $\left\{\mathbf{x}_0^{(j)}\right\}_{j \in [M]}$.
  \end{algorithmic}
  \end{algorithm}