Reparameterized Diffusion Loss

Reparameterization of Target

The original diffusion update was given by

$$z_t=\sqrt{{\alpha }_t}\cdot x+\sqrt{1-{\alpha }_t}\cdot ϵ$$

Recall that in the original diffusion loss function was:

$$\begin{array}{r}L[{\varphi }_{1\dots T}]=\sum _{i=I}^I\left(\underset{\text{reconstruction term}}{\underbrace{-\log [{\text{Norm}}_{x_i}[f_1[z_{i_1},{\varphi }_1],{\sigma }_1^{2}I]]}}+\sum _{t=2}^T\frac{1}{2{\sigma }_t^2}{\left|\left|\underset{\text{target, mean of }q(z_{t-1}|z_t,x)}{\underbrace{\frac{1-{\alpha }_{t-1}}{1-{\alpha }_t}\sqrt{1-{\beta }_t}{z}_{it}+\frac{\sqrt{{\alpha }_{t-1}{\beta }_t}}{1-{\alpha }_t}{x}_i}}-\underset{\text{predicted }{z}_{t-1}}{\underbrace{f_t[z_{it},{\varphi }_t]}}\right|\right|}^2\right)\end{array}$$

where the original form of the KL divergence as:

$$D_{KL}[q(z_{t-1}|z_t,x)\mid \mid Pr(z_{t-1}|z_t,{\varphi }_t)]=\frac{1}{2{\sigma }_t^2}{\left|\left|\frac{(1-{\alpha }_{t-1})}{1-{\alpha }_t}\sqrt{1-{\beta }_t}{z}_t+\frac{\sqrt{{\alpha }_{t-1}}{\beta }_t}{1-{\alpha }_t}x-f_t[z_t,{\varphi }_t]\right|\right|}^2+C$$

The data term $x$ in the KL divergence can be expressed as the diffused image minus the noise that was added to it:

$$x=\frac{1}{\sqrt{{\alpha }_t}}\cdot z_t-\frac{\sqrt{1-{\alpha }_t}}{\sqrt{{\alpha }_t}}\cdot ϵ$$

Substituting this into the target terms from the original loss function gives:

$$\begin{array}{rl}& \frac{(1-{\alpha }_{t-1})}{1-{\alpha }_t}\sqrt{1-{\beta }_t}{z}_t+\frac{\sqrt{{\alpha }_{t-1}}{\beta }_t}{1-{\alpha }_t}x\\ & =\frac{(1-{\alpha }_{t-1})}{1-{\alpha }_t}\sqrt{1-{\beta }_t}{z}_t+\frac{\sqrt{{\alpha }_{t-1}}{\beta }_t}{1-{\alpha }_t}\left(\frac{1}{\sqrt{{\alpha }_t}}{z}_t-\frac{\sqrt{1-{\alpha }_t}}{\sqrt{{\alpha }_t}}ϵ\right)\\ & =\frac{(1-{\alpha }_{t-1})}{1-{\alpha }_t}\sqrt{1-{\beta }_t}{z}_t+\frac{{\beta }_t}{1-{\alpha }_t}\left(\frac{1}{\sqrt{1-{\beta }_t}}{z}_t-\frac{\sqrt{1-{\alpha }_t}}{\sqrt{1-{\beta }_t}}ϵ\right)\end{array}$$

where we have used the fact that $\sqrt{{\alpha }_t}/\sqrt{{\alpha }_{t-1}}=\sqrt{1-{\beta }_t}$ between the second and third lines.

Simplifying further, we get:

$$\begin{array}{rl}\frac{1-{\alpha }_{t-1}}{1-{\alpha }_t}\sqrt{1-{\beta }_t}{z}_t+\frac{\sqrt{{\alpha }_{t-1}}{\beta }_t}{1-{\alpha }_t}x& =\left(\frac{(1-{\alpha }_{t-1})\sqrt{1-{\beta }_t}}{1-{\alpha }_t}+\frac{{\beta }_t}{(1-{\alpha }_t)\sqrt{1-{\beta }_t}}\right)z_t-\frac{{\beta }_t}{\sqrt{1-{\alpha }_t}\sqrt{1-{\beta }_t}}ϵ\\ & =\left(\frac{(1-{\alpha }_{t-1})(1-{\beta }_t)}{(1-{\alpha }_t)\sqrt{1-{\beta }_t}}+\frac{{\beta }_t}{(1-{\alpha }_t)\sqrt{1-{\beta }_t}}\right)z_t-\frac{{\beta }_t}{\sqrt{1-{\alpha }_t}\sqrt{1-{\beta }_t}}ϵ\\ & =\frac{(1-{\alpha }_{t-1})(1-{\beta }_t)+{\beta }_t}{(1-{\alpha }_t)\sqrt{1-{\beta }_t}}{z}_t-\frac{{\beta }_t}{\sqrt{1-{\alpha }_t}\sqrt{1-{\beta }_t}}ϵ\\ & =\frac{1-{\alpha }_t}{(1-{\alpha }_t)\sqrt{1-{\beta }_t}}{z}_t-\frac{{\beta }_t}{\sqrt{1-{\alpha }_t}\sqrt{1-{\beta }_t}}ϵ\\ & =\frac{1}{\sqrt{1-{\beta }_t}}{z}_t-\frac{{\beta }_t}{\sqrt{1-{\alpha }_t}\sqrt{1-{\beta }_t}}ϵ.\end{array}$$

where we multiplied the numerator and denominator of the first term by $\sqrt{1-{\beta }_t}$ between lines 2 and 3, multiplied out the terms, and then simplified the numerator in the first term between lines 3 and 4.

Substituting this back into the loss function, we have:

$$\begin{array}{rl}L[{\varphi }_{1\dots T}]& =\sum _{i=1}^I\left(-\log [{\text{Norm}}_{x_i}[f_1[z_{i1},{\varphi }_1],{\sigma }_1^{2}I]]+\sum _{t=2}^T\frac{1}{2{\sigma }_t^2}{\left|\left|\left(\frac{1}{\sqrt{1-{\beta }_t}}{z}_{it}-\frac{{\beta }_t}{\sqrt{1-{\alpha }_t}\sqrt{1-{\beta }_t}}{ϵ}_{it}-f_t[z_{it},{\varphi }_t{]}^2\right)\right|\right|}^2\right)\end{array}$$

Reparameterization of network

Now we replace the model ${\hat{z}}_{t-1}=f_t[z_t,{\varphi }_t]$ with a new model $\widehat{ϵ}=g_t[z_t,{\varphi }_t]$, which predicts the noise $ϵ$ that was mixed with $x$ to create $z_t$:

$$f_t[z_t,{\varphi }_t]=\frac{1}{\sqrt{1-{\beta }_t}}{z}_t-\frac{{\beta }_t}{\sqrt{1-{\alpha }_t}\sqrt{1-{\beta }_t}}{g}_t[z_t,{\varphi }_t]$$

Substituting the new model our loss function produces the criterion:

$$\begin{array}{rl}L[{\varphi }_{1\dots T}]& =\sum _{i=1}^I-\log [{\text{Norm}}_{x_i}[f_1[z_{i1},{\varphi }_1],{\sigma }_1^{2}I]]+\sum _{t=2}^T\frac{{\beta }_t^2}{(1-{\alpha }_t)(1-{\beta }_t)2{\sigma }_t^2}{\left|\left|g_t[z_{it},{\varphi }_t]-{ϵ}_{it}\right|\right|}^2\end{array}$$

The log normal can be written as a least squares loss plus a constant $C_i$:

$$\begin{array}{rl}L[{\varphi }_{1\dots T}]& =\sum _{i=1}^I\frac{1}{2{\sigma }_1^2}{\left|\left|x_i-f_1[z_{i1},{\varphi }_1]\right|\right|}^2+\sum _{t=2}^T\frac{{\beta }_t^2}{(1-{\alpha }_t)(1-{\beta }_t)2{\sigma }_t^2}{\left|\left|g_t[z_{it},{\varphi }_t]-{ϵ}_{it}\right|\right|}^2\end{array}$$

Substituting in the definitions of $x$ and $f_1[z_1,{\varphi }_1]$ from above, the first term simplifies to

$$\frac{1}{2{\sigma }_1^2}\left|\right|x_i-f_1[z_{i1},{\varphi }_1]|{|}^2=\frac{1}{2{\sigma }_1^2}{\left|\left|\frac{{\beta }_1}{\sqrt{1-{\alpha }_1}\sqrt{1-{\beta }_1}}{g}_1[z_{i1},{\varphi }_1]-\frac{{\beta }_1}{\sqrt{1-{\alpha }_1}\sqrt{1-{\beta }_1}}{ϵ}_{i1}\right|\right|}^2$$

Adding this back to the final loss function yields:

$$L[{\varphi }_{1\dots T}]=\sum _{i=1}^I\sum _{t=1}^T\frac{{\beta }_t^2}{(1-{\alpha }_t)(1-{\beta }_t)2{\sigma }_t^2}\left|\right|g_t[z_{it},{\varphi }_t]-{ϵ}_{it}|{|}^2$$

where we have disregarded the additive constants $C_i$.

In practice, the scaling factors (which might be different at each time step) are ignored, giving an even simpler formulation:

$$\begin{array}{rl}L[{\varphi }_{1\dots T}]& =\sum _{i=1}^I\sum _{t=1}^T\left|\right|g_t[z_{it},{\varphi }_t]-{ϵ}_{it}|{|}^2\\ & =\sum _{i=1}^I\sum _{t=1}^T\left|\right|g_t[\sqrt{{\alpha }_t}\cdot x_i+\sqrt{1-{\alpha }_t}\cdot {ϵ}_{it},{\varphi }_t]-{ϵ}_{it}|{|}^2\end{array}$$

where we have rewritten $z_t$ using the diffusion kernel in the second line.