Academic Paper reading log

22-05

Inverse problems using diffusion models

• What are the definition of inverse problems? $y=Ax+n$, where $n$ is noise (random, unavoidable errors), $x$ is the true state that we want to know, and $A$ is a measurement function (corruption or transformation function).
• Examples:
  • Blurring: y = blurred image, x = clean image, A = nearby pixel averaging
  • Mask & Delete: y = masked image, x = clean image, A = cuts gray box
  • Pose estimation: y = three floating dots, x = 3D human pose, A = motion capture system.
• Usually, forward problems (solving for y from x) are solved by using plain physics. And for inverse problems, When $A^{-1}$ is computable, x can be directly solved. But, for most inverse problems, A is invertible, and there are many possible $x$ for the same y. Thus, the problem is said to be ill-posed.
• The classical line of attack solves for minimum distance between the solution ($y,Ax$) based on some prior $R(x)$: $\arg \underset{x}{\min }J(x)=\underset{\text{distance from }y}{\underbrace{d(y,A(x))}}+\underset{\text{regularizer}}{\underbrace{\lambda R(x)}}$
• Many works have explored different types of regularizers:
  • Total variation (TV) regularization: $\min \frac{1}{2}\Vert Ax-y{\Vert }_2^2+\lambda \Vert x{\Vert }_{\text{TV}}$, where $\Vert x{\Vert }_{\text{TV}}=\Vert \nabla x{\Vert }_1={\sum }_i{\sum }_j\sqrt{(x_{i+1,j}-x_{i,j}{)}^2+(x_{i,j+1}-x_{i,j}{)}^2}$ which is the L1 norm of the gradient. So, the regularizer restricts total variation in a single update, and selects $x$ that prevents a spike change.
  • Truncated SVD: We decompose A into SVD components: $A=U\Sigma V^{\top }$, and write the naive pseudo-inverse solution as $x_{\text{naive}}={\sum }_{i=1}^n\frac{u_i^{\top }y}{{\sigma }_i}{v}_i$. For small eigenvalues, in ill-posed inverse problems some eigenvalues rapidly decay to zero. Dividing a noisy measurement $y$ with a small ${\sigma }_i$, scales the noise massively. TSVD truncates problematic small singular values by choosing a truncation threshold $k<n$, keeping the first k singular values, and discarding the rest by setting them to zero. Thus, TSVD solution can be formulated as $x_{\text{TSVD}}={\sum }_{i=1}^k\frac{u_i^{\top }y}{{\sigma }_i}{v}_i$.
  • Bayesian inference: We treat $x$ as random variable, and apply Bayes’ law to compute the posterior: $p(x|y)=\frac{p(y|x)p(x)}{\int p(y|x)p(x)dx}$, where the likelihood is regularized by a prior. Taking the argmax of the posterior gives a suitable solution $x$. But finding suitable prior is not straightforward.
• This is where generative models simplified the scope of the problem. Generative models are excellent at learning unknown distributions from samples of the distribution. By using a diffusion model to learn the prior distribution from data, and even generate prior samples based on feedback.

24-05

Pseudo-inverse GDM

• Measurement of some signal: $\mathbf{y}=\mathbf{H}{\mathbf{x}}_0+\mathbf{z}$, where $\mathbf{H}\in {\mathbb{R}}^{n\times m}$ and $\mathbf{z}\sim \mathcal{N}(0,{\sigma }_y^2\mathbf{I})$.
• Writing score of posterior using bayes law: ${\nabla }_{{\mathbf{x}}_t}\log p({\mathbf{x}}_t|\mathbf{y})=\underset{\text{likelihood}}{\underbrace{{\nabla }_{{\mathbf{x}}_t}\log p(\mathbf{y}|{\mathbf{x}}_t)}}+\underset{\text{score}}{\underbrace{{\nabla }_{{\mathbf{x}}_t}\log p({\mathbf{x}}_t)}}$
• We know that ${\nabla }_{{\mathbf{x}}_t}\log p(\mathbf{y}|{\mathbf{x}}_t)$ is intractable to compute, so we marginalize $p_t(\mathbf{y}|{\mathbf{x}}_t)={\int }_{{\mathbf{x}}_0}p({\mathbf{x}}_0|{\mathbf{x}}_t)p(\mathbf{y}|{\mathbf{x}}_0)d{\mathbf{x}}_0$.
• Approximate $p_t({\mathbf{x}}_0|{\mathbf{x}}_t)\sim \mathcal{N}({\hat{\mathbf{x}}}_t,r_t^2\mathbf{I})$ with a gaussian where $r_t^2$ is a time-dependent variance that depends on the data,
  • the mean ${\hat{\mathbf{x}}}_t$ is obtained using Tweedie’s formula: ${\hat{\mathbf{x}}}_t=\mathbb{E}[{\mathbf{x}}_0|{\mathbf{x}}_t]={\mathbf{x}}_t+{\sigma }_t^2{\nabla }_{{\mathbf{x}}_t}\log p_t({\mathbf{x}}_t)\approx {\mathbf{x}}_t+{\sigma }_t^2{S}_{\theta }(\mathbf{x};{\sigma }_t)$.
• When $\mathbf{H}$ is linear, and measurement is obtained by adding gaussian white noise along with $p({\mathbf{x}}_0|{\mathbf{x}}_t),p(\mathbf{y}|{\mathbf{x}}_0)$ both being Gaussian r.v. implying the distribution of $\mathbf{y}$ conditioned of ${\mathbf{x}}_t$ is also Gaussian
  $p(\mathbf{y}|{\mathbf{x}}_t)\approx \mathcal{N}(\mathbf{H}{\hat{\mathbf{x}}}_t,r_t^2\mathbf{H}{\mathbf{H}}^{\top }+{\sigma }_{\mathbf{y}}^2\mathbf{I})$
• And we finally get following approximation to the score. This is interpreted as vector-jacobian product and can be obtained using backpropagation through the network.
  ${\nabla }_{{\mathbf{x}}_t}\log p_t(\mathbf{y}|{\mathbf{x}}_t)\approx {\left((\underset{\text{vector}}{\underbrace{\mathbf{y}-\mathbf{H}{\hat{\mathbf{x}}}_t{)}^{\top }(r_t^2\mathbf{H}{\mathbf{H}}^{\top }+{\sigma }_{\mathbf{y}}^2\mathbf{I}{)}^{-1}\mathbf{H}}}\underset{\text{Jacobian}}{\underbrace{\frac{\partial {\hat{\mathbf{x}}}_t}{\partial {\mathbf{x}}_t}}}\right)}^{\top }$
• For cases when ${\sigma }_{\mathbf{y}}=0$, above equation can be written as following, and indicates that likelihood guidance term is influenced by the pseudo-inverse of the measurement matrix.

$$\begin{array}{rlr}{\nabla }_{{\mathbf{x}}_t}\log p(\mathbf{y}|{\mathbf{x}}_t)& \approx {\left((\underset{\text{vector}}{\underbrace{\mathbf{y}-\mathbf{H}{\hat{\mathbf{x}}}_t{)}^{\top }(r_t^2\mathbf{H}{\mathbf{H}}^{\top }+\text{cancelto}0{\sigma }_{\mathbf{y}}^2\mathbf{I}{)}^{-1}\mathbf{H}}}\underset{\text{Jacobian}}{\underbrace{\frac{\partial {\hat{\mathbf{x}}}_t}{\partial {\mathbf{x}}_t}}}\right)}^{\top }& \\ & ={\left((\mathbf{y}-\mathbf{H}{\hat{\mathbf{x}}}_t{)}^{\top }\mathbf{H}\left(\frac{1}{r_t^2\mathbf{H}{\mathbf{H}}^{\top }}\right)J_{{\mathbf{x}}_t}(\hat{\mathbf{x}})\right)}^{\top }& \\ & ={\left(\frac{1}{r_t^2}({\mathbf{y}}^{\top }-{\hat{\mathbf{x}}}^{\top }{\mathbf{H}}^{\top })(\mathbf{H}{\mathbf{H}}^{\top }{)}^{-1}\mathbf{H}\cdot J_{{\mathbf{x}}_t}({\hat{\mathbf{x}}}_t)\right)}^{\top }& \\ & ={\left(\frac{1}{r_t^2}({\mathbf{y}}^{\top }({\mathbf{HH}}^{\top }{)}^{-1}\mathbf{H}-{\hat{\mathbf{x}}}_t{\mathbf{H}}^{\top }(\mathbf{H}{\mathbf{H}}^{\top }{)}^{-1}\mathbf{H})\cdot J_{{\mathbf{x}}_t}({\hat{\mathbf{x}}}_t)\right)}^{\top }& \\ & ={\left(\frac{1}{r_t^2}({\mathbf{H}}^{\top }(\mathbf{H}{\mathbf{H}}^{\top }{)}^{-1}\mathbf{y}{)}^{\top }-({\mathbf{H}}^{\top }(\mathbf{H}{\mathbf{H}}^{\top }{)}^{-1}\mathbf{H}{\hat{\mathbf{x}}}_t{)}^{\top }\cdot J_{{\mathbf{x}}_t}({\hat{\mathbf{x}}}_t)\right)}^{\top }& \\ & ={\left(\frac{1}{r_t^2}({\mathbf{H}}^{†}\mathbf{y}-{\mathbf{H}}^{†}\mathbf{H}{\mathbf{x}}_t{)}^{\top }\cdot J_{{\mathbf{x}}_t}({\hat{\mathbf{x}}}_t)\right)}^{\top }& ({\mathbf{H}}^{\top }(\mathbf{H}{\mathbf{H}}^{\top }{)}^{-1}\text{ is Moore-Penrose PI})\end{array}$$

• As notable from the equation derived above, score of the likelihood doesn’t require $H$ to be differentiable, because $\mathbf{H}$ and the pseudo-inverse ${\mathbf{H}}^{†}$ are used directly as static functions inside the “vector” of the vector-jacobian product. This means $\mathbf{H}$ doesn’t need to linear or even differentiable.
• For complex non-linear cases, measurement function $h(\mathbf{x}):{\mathbb{R}}^n\to {\mathbb{R}}^m$ need to only satisfy the pseudo-inverse property, namely we need to find another function $h^{†}:{\mathbb{R}}^m\to {\mathbb{R}}^n$ such that $h(h^{†}(h(\mathbf{x})))=h(\mathbf{x}),\forall \mathbf{x}\in {\mathbb{R}}^n$.
• Variance of approximation $r_t^2$. Assume $p_0({\mathbf{x}}_0)\sim \mathcal{N}(0,\mathbf{I})$, then applying Bayes law to $p_t({\mathbf{x}}_0|{\mathbf{x}}_t)\propto p_0({\mathbf{x}}_0)p({\mathbf{x}}_t|{\mathbf{x}}_0)$. We know that forward noising process adds Gaussian noise with variance ${\sigma }_t^2$, and $p({\mathbf{x}}_t|{\mathbf{x}}_0)\sim \mathcal{N}({\mathbf{x}}_0,{\sigma }_t^2\mathbf{I})$.
  • The posterior turns out to be $p_t({\mathbf{x}}_0|{\mathbf{x}}_t)\propto \mathcal{N}\left(\frac{{\mathbf{x}}_t}{{\sigma }_t^2+1},\frac{{\sigma }_t^2}{{\sigma }_t^2+1}\right)$. So, for $t=0,{\sigma }_t^2\to 0⟹r_0=0$ and $t=T,{\sigma }_t^2{\gg }_1⟹r_T\approx 1$.
  • To see how the algorithm behaves (listening to unconditional score vs likelihood as measurement guidance) over time with respect to variance, assume $H=I$, the weighting factor for the likelihood guidance term comes out to be $r_t^2(r_t^2+{\sigma }_{\mathbf{y}}^2{)}^{-1}$.
  • For $t\to T$, ${\sigma }_t$ is large, the image is composed of purely white noise, and the model doesn’t know if it’s generating image of a face or cat or dog, $r_t\approx 1$ meaning the guidance term dominates the unconditional path, and forcefully take the sample towards noisy measurement $\mathbf{y}$.
  • For $t\to 0$, $r_t\to {\sigma }_t$ then $r_t^2(r_t^2+{\sigma }_{\mathbf{y}}^2{)}^{-1}\approx {\sigma }_t^2{\sigma }_{\mathbf{y}}^{-2}\to 0$. So, the noisy measurement term provides little impact, and the algorithm follow the unconditional generation path.