bd3dowling
commented
1 month ago Paste from obsidian:
- Nice introduction to inverse problems using unconditional models...
- "Approximating the score is equivalent to learning a denoising neural network $\hat{x}{\theta}(x^t, t)$ to approximate $\mathbb{E}{q}[x^0 \mid x^t]$. The reason is that by Tweedie's formula: $$ \nabla{x^t} \log q(x^t) = \frac{\mathbb{E}{q}[x^0 \mid x^t] - x^t}{t\sigma^2} $$ and one can set $s{\theta}(x^t, t) := \frac{\hat{x}{\theta}(x^t, t) - x^t}{t\sigma^2}$."
- 3.1 -> Why need twisting.
- Twisting: reduce number of required particles for good approximation accuracy by reducing discrepancy between successive intermediate target distributions by bringing them closer to the final target using twisting functions which define new twisted proposals and weighting functions.
- Key idea is to approximate the optimal twisting function $p{\theta}(y \mid x^t)$ by $p{y \mid x^0}(y \mid \hat{x}_{\theta}(x^t))$ (likelihood evaluated at the denoising estimate of $x^0$ at step $t$ from the diffusion model).
- Then yield sequence of twisted proposals to approximate the optimal ones: $$\begin{align} \tilde{p}{\theta}(x^t \mid x^{t+1}, y) &:= \mathcal{N}(x^t; x^{t+1} + \sigma^2s{\theta}(x^{t+1}, y); \tilde{\sigma}^2) \end{align}$$ where $$\begin{align} s{\theta}(x^{t+1}, y) := s{\theta}(x^{t+1}) + \nabla{x^{t+1}}\log \tilde{p}{\theta}(y \mid x^{t+1}) \approx \nabla{x^{t+1}} \log p{\theta}(x^{t+1} \mid y) \end{align}$$ with $\tilde{\sigma}^2$ the variance of the proposal which can be taken as $\sigma^2$ to match the variance of the unconditional model for simplicity...
- Define twisted weighting functions accordingly using $\tilde{p}_{\theta}$.
\begin{algorithm} \caption{Twisted Diffusion Sampler (TDS)} \begin{algorithmic} \State Initial proposal and weight \For{$k = 1$ to $K$} \State Sample $x^T_k \sim p(x^T)$ \State Compute initial weight $w_k \gets \tilde{p}_T^k = \tilde{p}_\theta(y \mid x^T_k)$ \EndFor \For{$t = T-1$ to $0$} \State Resample $\{(x^{t+1}_k, \tilde{p}^{t+1}_k)\}^K_{k=1}$ from Multinomial based on $\{(x^{t+1}_k, \tilde{p}^{t+1}_k)\}^K_{k=1}$ and $\{w_k\}^K_{k=1}$ \For{$k = 1$ to $K$} \State Calculate conditional score $\tilde{s}_k = (\hat{x}_\theta(x^{t+1}_k) - x^{t+1}_k) / t\sigma^2 + \nabla_{x^{t+1}_k} \log \tilde{p}^{t+1}_k$ \State Propose $x^t_k \sim \tilde{p}_\theta(\cdot \mid x^{t+1}_k, y) = \mathcal{N}(x^{t+1}_k + \sigma^2 \tilde{s}_k, \tilde{\sigma}^2)$ \State Twising function $\tilde{p}_k^t \gets \tilde{p}_\theta(y \mid x^t)$ \State Update weight $w_k \gets p_\theta(x^t_k \mid x^{t+1}_k) \tilde{p}^t_k / (\tilde{p}_\theta(x^t_k \mid x^{t+1}_k, y) \tilde{p}^{t+1}_k)$ \EndFor \EndFor \end{algorithmic} \end{algorithm}
Wu, Luhuan, et al. Practical and Asymptotically Exact Conditional Sampling in Diffusion Models. arXiv:2306.17775, arXiv, 30 June 2023. arXiv.org, http://arxiv.org/abs/2306.17775.