Score and diffusion methods for inverse data

Let $d\ge 1$, let $\sigma >0$, and let $\mu$ be a probability measure on the Borel subsets of ${\mathbb{R}}^d$. Write ${\phi }_{\sigma }(v)=c_{\sigma ,d}\exp (-\mathrm{\parallel }v{\mathrm{\parallel }}^2\mathrm{/}(2{\sigma }^2))$ for the normalized Gaussian density with respect to Lebesgue measure $dx$, define $P_{\sigma }(du,dx)=\mu (du){\phi }_{\sigma }(x-u)dx$, and set $p_{\sigma }(x)=\int {\phi }_{\sigma }(x-u)\mu (du)$. Define $Q_{\sigma }(x,du)={\phi }_{\sigma }(x-u)\mu (du)\mathrm{/}{p}_{\sigma }(x)$ and $m_{\sigma }(x)=\int u{Q}_{\sigma }(x,du)$; under $P_{\sigma }$, write $X_0(u,x)=u$ and $X_{\sigma }(u,x)=x$ for the coordinate random vectors, and set $Z=(X_{\sigma }-X_0)\mathrm{/}\sigma$, so $Z\sim N(0,I_d)$ is independent of $X_0$. For Borel maps $h,s:{\mathbb{R}}^d\to {\mathbb{R}}^d$ that are square-integrable with respect to $p_{\sigma }(x)dx$, define $R(h)=\int \mathrm{\parallel }h(x)-u{\mathrm{\parallel }}^2{P}_{\sigma }(du,dx)$ and $J(s)=\int \mathrm{\parallel }s(x)-(u-x)\mathrm{/}{\sigma }^2{\mathrm{\parallel }}^2{P}_{\sigma }(du,dx)$. Then $p_{\sigma }$ is positive and continuously differentiable, $Q_{\sigma }$ is the conditional law of $X_0$ given $X_{\sigma }=x$, and $\mathrm{\nabla }\log p_{\sigma }(x)=(m_{\sigma }(x)-x)\mathrm{/}{\sigma }^2$ for every $x\in {\mathbb{R}}^d$. If $\int \mathrm{\parallel }u{\mathrm{\parallel }}^2\mu (du)<\mathrm{\infty }$, $m_{\sigma }$ uniquely minimizes $R$ up to $p_{\sigma }$-almost-everywhere equality. The unique minimizer of $J$ is $\mathrm{\nabla }\log p_{\sigma }$, also up to $p_{\sigma }$-almost-everywhere equality. Under the same Gaussian noising, suppose an observation variable $Y$ is conditionally independent of $Z$ given $X_0$, and, for the fixed observation $y$, ${\mathrm{\ell }}_y:{\mathbb{R}}^d\to (0,\mathrm{\infty })$ is the bounded Borel conditional density of $Y$ at $y$ given $X_0=u$, relative to a fixed sigma-finite observation measure, and assume $0<Z_y=\int {\mathrm{\ell }}_y(u)\mu (du)<\mathrm{\infty }$. Let ${\mu }_y(du)={\mathrm{\ell }}_y(u)\mu (du)\mathrm{/}{Z}_y$ and $p_{\sigma }(x\mid y)=\int {\phi }_{\sigma }(x-u){\mu }_y(du)$. Then, for $r_{\sigma }(y\mid x)=\int {\mathrm{\ell }}_y(u)Q_{\sigma }(x,du)$, the conditional score satisfies ${\mathrm{\nabla }}_x\log p_{\sigma }(x\mid y)={\mathrm{\nabla }}_x\log p_{\sigma }(x)+{\mathrm{\nabla }}_x\log r_{\sigma }(y\mid x)$. No absolute-continuity or bounded-support assumption on $\mu$ is required; $\sigma >0$ supplies the full-support smoothing density.

Daras et al., §2.2 and Appendix A.1–A.2

Song et al., §§3.3 and 5

Prince, Chapter 18, “Diffusion models”

Let $p_{\sigma }$, $Q_{\sigma }$, $m_{\sigma }$, and $r_{\sigma }(y\mid x)$ be formed under the hypotheses of the conditional-score and Tweedie identities for Gaussian smoothing. The diffusion prior score is $s_{\sigma }(x)={\mathrm{\nabla }}_x\log p_{\sigma }(x)$, the exact score of the smoothed prior density. The measurement score is $g_{\sigma }(x,y)={\mathrm{\nabla }}_x\log r_{\sigma }(y\mid x)$, where $r_{\sigma }(y\mid x)=\int {\mathrm{\ell }}_y(u)Q_{\sigma }(x,du)$ is the likelihood density averaged over the conditional law of the clean state. For general inverse problems, evaluating this clean-state likelihood integral is generally computationally intractable. The conditional score is $c_{\sigma }(x,y)={\mathrm{\nabla }}_x\log p_{\sigma }(x\mid y)$; it equals $s_{\sigma }(x)+g_{\sigma }(x,y)$, and the clean-state conditional mean is $m_{\sigma }(x)=x+{\sigma }^2{s}_{\sigma }(x)$. A learned field ${\widehat{s}}_{\theta }$ is an estimate of $s_{\sigma }$ and is exact only when it agrees with $s_{\sigma }$ almost everywhere under $p_{\sigma }(x)dx$. Replacing $r_{\sigma }(y\mid x)$ by the likelihood evaluated at $m_{\sigma }(x)$ is a plug-in likelihood approximation, not part of the exact measurement-score definition.

Conditional-score and Tweedie identities for Gaussian smoothing

Daras et al., §§2.2 and 2.4.3–2.4.4

Song et al., §5

Chung et al., §3.1

Let $N\ge 1$ be finite, let $0=t_0<t_1<\cdots <t_N$, let $d\ge 1$ and let the state space be ${\mathbb{R}}^d$, and let $(\mathsf{Y},\mathcal{Y})$ be a measurable observation space. Fix a measurable forward operator $\mathcal{A}:{\mathbb{R}}^d\to \mathsf{Y}$, an observation $y\in \mathsf{Y}$, and a nonnegative measurable likelihood function ${\mathrm{\ell }}_y:{\mathbb{R}}^d\to [0,\mathrm{\infty })$ obtained by fixing $y$ in a likelihood density relative to a fixed sigma-finite measure on $\mathsf{Y}$. A measurement-guided reverse-kernel reconstruction procedure takes a finite Borel score field ${\widehat{s}}_{\theta }:{\mathbb{R}}^d\times \{t_0,\dots ,t_N\}\to {\mathbb{R}}^d$; a terminal reference probability ${\nu }_N$ on the Borel sets of ${\mathbb{R}}^d$ and a sampler that returns an exact draw or failure within finite initialization budget $b_{\mathrm{i}\mathrm{n}\mathrm{i}\mathrm{t}}$; and, for each $k$, a measurable correction space $({\mathsf{C}}_k,{\mathcal{C}}_k)$, its failure-augmented measurable space ${\mathsf{C}}_k^{†}={\mathsf{C}}_k\bigsqcup \{\mathrm{\perp }\}$ with $\sigma$-algebra ${\mathcal{C}}_k^{†}=\{D,D\cup \{\mathrm{\perp }\}:D\in {\mathcal{C}}_k\}$, a measurable correction map $C_{k,{\widehat{s}}_{\theta },\mathcal{A},{\mathrm{\ell }}_y}:{\mathbb{R}}^d\times \mathsf{Y}\to {\mathsf{C}}_k^{†}$, a probability kernel $Q_{\theta ,k}(\cdot \mid x,y,a)$ from ${\mathbb{R}}^d\times \mathsf{Y}\times {\mathsf{C}}_k$ with its product sigma-algebra to ${\mathbb{R}}^d$ with its Borel sigma-algebra, and a sampler that returns an exact draw from that kernel or failure within finite per-step budget $b_k$. The correction map may use the supplied score, forward operator and likelihood and, when supplied for the exact-score option, the prior and scale map. It may also take, for each step, a declared Borel map ${\mathrm{\Pi }}_{k-1}:{\mathbb{R}}^d\to {\mathbb{R}}^d$ and a Borel admissible set $A_{k-1}\subseteq {\mathbb{R}}^d$. Each per-step budget includes correction evaluation, kernel sampling and any declared projection. First draw $x_N$ from ${\nu }_N$ within $b_{\mathrm{i}\mathrm{n}\mathrm{i}\mathrm{t}}$; if initialization fails, return the terminal-initialization failure status. For $k=N,N-1,\dots ,1$, compute $a_k=C_{k,{\widehat{s}}_{\theta },\mathcal{A},{\mathrm{\ell }}_y}(x_k,y)$; if $a_k=\mathrm{\perp }$, return the correction failure status at $k$; otherwise draw $x_{k-1}$ from $Q_{\theta ,k}(\cdot \mid x_k,y,a_k)$ within $b_k$; if drawing fails, return the sampling failure status at $k$; if ${\mathrm{\Pi }}_{k-1}$ is supplied, apply it to $x_{k-1}$ and return the admissibility failure status at $k-1$ if evaluation fails or its output is outside $A_{k-1}$. If $k>1$, continue at $k-1$ with the resulting state $x_{k-1}$; if $k=1$, return $x_0$ with completed status. A correction may use a declared surrogate without additional likelihood or noising assumptions. If it uses the exact measurement score at step $k$, the caller must also declare a Borel probability measure $\mu$ on ${\mathbb{R}}^d$ and a map $\sigma :\{t_1,\dots ,t_N\}\to (0,\mathrm{\infty })$ associating each reverse-schedule time with its Gaussian-noising scale. For every such step, the forward noising convention at $t_k$ must be $X_{\sigma (t_k)}=X_0+\sigma (t_k)Z$, where $X_0\sim \mu$ and $Z\sim N(0,I_d)$ is independent of $X_0$, and the observation variable $Y$ with conditional density ${\mathrm{\ell }}_y$ must be conditionally independent of $Z$ given $X_0$, and ${\mathrm{\ell }}_y$ must be bounded, strictly positive and Borel measurable with $0<Z_y=\int {\mathrm{\ell }}_y(u)\mu (du)<\mathrm{\infty }$. The exact correction then uses $g_{\sigma (t_k)}(x_k,y)={\mathrm{\nabla }}_x\log \int {\mathrm{\ell }}_y(u)Q_{\sigma (t_k)}(x_k,du)$, with $Q_{\sigma (t_k)}$ formed from this prior and Gaussian noising as in the conditional-score and Tweedie identities. The procedure returns an algorithmic sample; it makes no posterior-exactness claim unless the terminal law and all transition and projection conditions are separately certified.

diagram failed to render:
\begin{tikzpicture}[>=stealth]
  \node[draw,rounded corners,align=center] (init) at (0,0) {draw $x_N\sim\nu_N$ within $b_{\mathrm{init}}$};
  \node[draw,rounded corners,align=center] (corr) at (3.5,0) {$a_k=C_k(x_k,y)$; check $a_k\ne\bot$};
  \node[draw,rounded corners,align=center] (kernel) at (7.5,0) {draw $x_{k-1}\sim Q_{\theta,k}(\cdot\mid x_k,y,a_k)$};
  \node[draw,rounded corners,align=center] (proj) at (11.5,0) {optional $\Pi_{k-1}$; check $A_{k-1}$};
  \node[draw,diamond,aspect=2,align=center] (next) at (15.3,0) {$k=1?$};
  \node[draw,rounded corners,align=center] (out) at (15.3,-1.8) {return $x_0$};
  \node[draw,rounded corners,align=center] (stop) at (7.5,-2.2) {first failure status};
  \draw[->] (init)--(corr);
  \draw[->] (corr)--(kernel);
  \draw[->] (kernel)--(proj);
  \draw[->] (proj)--(next);
  \draw[->] (next.south)--node[right]{$k=1$}(out.north);
  \draw[->] (next.north) -- ++(0,1.5) -| node[pos=0.25,above,align=center]{$k>1$: continue at $k-1$\\with the just-produced state}(corr.north);
  \draw[->] (init.south)--(stop.north);
  \draw[->] (corr.south)--(stop.north);
  \draw[->] (kernel.south)--(stop.north);
  \draw[->] (proj.south)--(stop.north);
\end{tikzpicture}

The upper return arrow repeats the reverse step with the newly produced state and decremented index. The downward arrows record immediate stopping at a failed initialization, correction, kernel draw, or declared projection. The diagram is an operational flow description and asserts no distributional identity.

Diffusion prior and measurement scores

Measure kernel and probability kernel

Daras et al., §§2.4 and 3

Song et al., §5 and predictor-corrector algorithms

Chung et al., §3.1 and Algorithms 1–2

Prince, Chapter 18, “Diffusion models”

Let $X_0$ take values $-1$ and $1$ with equal probability, let $X_{\sigma }=X_0+\sigma Z$ with $\sigma >0$ and standard Gaussian noise independent of $X_0$, and let $Y=X_0+\tau W$ for $W\sim N(0,1)$ independent of $(X_0,Z)$, with $\tau >0$. At the observation $y=0$, the exact measurement density $r_{\sigma }(0\mid x)=\mathbb{E}[{\phi }_{\tau }(-X_0)\mid X_{\sigma }=x]$ is constant in $x$, so its score is zero. The exact posterior mean is $m_{\sigma }(x)=\mathbb{E}[X_0\mid X_{\sigma }=x]=\tanh (x\mathrm{/}{\sigma }^2)$, where $\tanh (t)=(e^{2t}-1)\mathrm{/}(e^{2t}+1)$. The plug-in density ${\phi }_{\tau }(-m_{\sigma }(x))$ has score $-m_{\sigma }(x)m_{\sigma }^{\mathrm{\prime }}(x)\mathrm{/}{\tau }^2$, which is nonzero for $x\mathrm{\ne }0$. Thus evaluating the measurement likelihood at the exact posterior mean does not in general recover the exact measurement score; this example does not rule out equality in special models.

Conditional-score and Tweedie identities for Gaussian smoothing

Diffusion prior and measurement scores

Daras et al., §2.4.4, Eq. 2.20, and §3.1.4, Eq. 3.8

Chung et al., §3.1, Eqs. 7–15 and Theorem 1