Reparameterization of loss function

Reparameterization of target

\[\begin{eqnarray}\label{eq:diffusion_diff_kernel2} \mathbf{z}_{t} = \sqrt{\alpha_{t}}\cdot \mathbf{x} + \sqrt{1-\alpha_{t}}\cdot\boldsymbol\epsilon. \end{eqnarray}\]

\[\begin{eqnarray}\label{eq:diffusion_x_reparam} \mathbf{x} = \frac{1}{\sqrt{\alpha_{t}}}\cdot \mathbf{z}_t - \frac{\sqrt{1-\alpha_{t}}}{\sqrt{\alpha_{t}}}\cdot\boldsymbol\epsilon. \end{eqnarray}\]

\[\begin{eqnarray} \frac{(1-\alpha_{t-1})}{1-\alpha_{t}}\sqrt{1-\beta_{t}}\mathbf{z}_{t}+\frac{\sqrt{\alpha_{t-1}}\beta_{t}}{1-\alpha_{t}}\mathbf{x} \\ &&\hspace{-3cm} =\frac{(1-\alpha_{t-1})}{1-\alpha_{t}}\sqrt{1-\beta_{t}}\mathbf{z}_{t}+\frac{\sqrt{\alpha_{t-1}}\beta_{t}}{1-\alpha_{t}}\left(\frac{1}{\sqrt{\alpha_{t}}} \mathbf{z}_t - \frac{\sqrt{1-\alpha_{t}}}{\sqrt{\alpha_{t}}}\boldsymbol\epsilon\right)\nonumber\\ &&\hspace{-3cm}= \frac{(1-\alpha_{t-1})}{1-\alpha_{t}}\sqrt{1-\beta_{t}}\mathbf{z}_{t}+\frac{\beta_{t}}{1-\alpha_{t}}\left(\frac{1}{\sqrt{1-\beta_{t}}} \mathbf{z}_t - \frac{\sqrt{1-\alpha_{t}}}{\sqrt{1-\beta_{t}}}\boldsymbol\epsilon\right),\nonumber \end{eqnarray}\]

\[\begin{eqnarray} \frac{(1-\alpha_{t-1})}{1-\alpha_{t}}\sqrt{1-\beta_{t}}\mathbf{z}_{t}+\frac{\sqrt{\alpha_{t-1}}\beta_{t}}{1-\alpha_{t}}\mathbf{x}\\ &&\hspace{-3cm}=\left(\frac{(1-\alpha_{t-1})\sqrt{1-\beta_{t}}}{1-\alpha_{t}}+\frac{\beta_{t}}{(1-\alpha_{t})\sqrt{1-\beta_{t}}}\right) \mathbf{z}_t - \frac{\beta_{t}}{\sqrt{1-\alpha_{t}}\sqrt{1-\beta_{t}}}\boldsymbol\epsilon\nonumber \\ &&\hspace{-3cm}=\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) \mathbf{z}_t - \frac{\beta_{t}}{\sqrt{1-\alpha_{t}}\sqrt{1-\beta_{t}}}\boldsymbol\epsilon\nonumber \\ &&\hspace{-3cm}=\frac{(1-\alpha_{t-1})(1-\beta_{t})+\beta_{t}}{(1-\alpha_{t})\sqrt{1-\beta_{t}}} \mathbf{z}_t - \frac{\beta_{t}}{\sqrt{1-\alpha_{t}}\sqrt{1-\beta_{t}}}\boldsymbol\epsilon\nonumber \\ &&\hspace{-3cm}=\frac{1-\alpha_{t}}{(1-\alpha_{t})\sqrt{1-\beta_{t}}} \mathbf{z}_t - \frac{\beta_{t}}{\sqrt{1-\alpha_{t}}\sqrt{1-\beta_{t}}}\boldsymbol\epsilon\nonumber \\ &&\hspace{-3cm}=\frac{1}{\sqrt{1-\beta_{t}}} \mathbf{z}_t - \frac{\beta_{t}}{\sqrt{1-\alpha_{t}}\sqrt{1-\beta_{t}}}\boldsymbol\epsilon,\nonumber \end{eqnarray}\]

\[\begin{eqnarray}\label{eq:diffusion_loss_function_2} L[\boldsymbol\phi_{1\ldots T}] &=& \sum_{i=1}^{I}\biggl(-\log\Bigl[\mbox{Norm}_{\mathbf{x}_i}\left[\mbox{\bf f}_{1}[\mathbf{z}_{i1},\boldsymbol\phi_{1}],\sigma_{1}^{2}\mathbf{I}\right]\Bigr]\\ &&\vspace{-2.0cm} + \sum_{t=2}^{T}\frac{1}{2\sigma_{t}^{2}}\left\lVert \left(\frac{1}{\sqrt{1-\beta_{t}}} \mathbf{z}_{it} - \frac{\beta_{t}}{\sqrt{1-\alpha_{t}}\sqrt{1-\beta_{t}}}\boldsymbol\epsilon_{it}\right)-\mbox{\bf f}_{t}[\mathbf{z}_{it},\boldsymbol\phi_{t}] \right\rVert^2\biggr).\nonumber \end{eqnarray}\]

Reparameterization of network

\[\begin{eqnarray}\label{eq:diffusion_f_reparam} \mbox{\bf f}_{t}[\mathbf{z}_{t},\boldsymbol\phi_{t}] = \frac{1}{\sqrt{1-\beta_{t}}} \mathbf{z}_t - \frac{\beta_{t}}{\sqrt{1-\alpha_{t}}\sqrt{1-\beta_{t}}}\mbox{\bf g}_t[\mathbf{z}_{t},\boldsymbol\phi_t]. \end{eqnarray}\]

\[\begin{eqnarray}\label{eq:diffusion_loss_function_3} L[\boldsymbol\phi_{1\ldots T}] =\\ &&\hspace{-1.5cm}\sum_{i=1}^{I}-\log\Bigl[\mbox{Norm}_{\mathbf{x}_i}\left[\mbox{\bf f}_{1}[\mathbf{z}_{i1},\boldsymbol\phi_{1}],\sigma_{1}^{2}\mathbf{I}\right]\Bigr] + \sum_{t=2}^{T}\frac{\beta_t^2}{(1-\alpha_{t})(1-\beta_{t})2\sigma_{t}^{2}}\Big\lVert\mbox{\bf g}_{t}[\mathbf{z}_{it},\boldsymbol\phi_{t}]-\boldsymbol\epsilon_{it}\Big\rVert^2.\nonumber \end{eqnarray}\]

\[\begin{eqnarray} L[\boldsymbol\phi_{1\ldots T}] = \sum_{i=1}^{I}\frac{1}{2\sigma_1^2}\Big\lVert\mathbf{x}_i-\mbox{\bf f}_{1}[\mathbf{z}_{i1},\boldsymbol\phi_{1}]\Bigr\rVert^2 + \sum_{t=2}^{T}\frac{\beta_t^2}{(1-\alpha_{t})(1-\beta_{t})2\sigma_{t}^{2}}\Big\lVert\mbox{\bf g}_{t}[\mathbf{z}_{it},\boldsymbol\phi_{t}]-\boldsymbol\epsilon_{it}\Big\rVert^2+C_i. \end{eqnarray}\]

\[\begin{eqnarray}\label{eq:diffusion_ls_term} \frac{1}{2\sigma_1^2}\Big\lVert\mathbf{x}_i-\mbox{\bf f}_{1}[\mathbf{z}_{i1},\boldsymbol\phi_{1}]\Bigr\rVert^2 = \frac{1}{2\sigma_1^2}\biggl\lVert \frac{\beta_{1}}{\sqrt{1-\alpha_{1}}\sqrt{1-\beta_{1}}}\mbox{\bf g}_1[\mathbf{z}_{i1},\boldsymbol\phi_1]- \frac{\beta_{1}}{\sqrt{1-\alpha_{1}}\sqrt{1-\beta_{1}}}\boldsymbol\epsilon_{i1}\bigg\rVert^2. \end{eqnarray}\]

\[\begin{eqnarray} L[\boldsymbol\phi_{1\ldots T}] \!\!&\!\!=\!\!&\!\! \sum_{i=1}^{I}\sum_{t=1}^{T}\frac{\beta_t^2}{(1-\alpha_{t})(1-\beta_{t})2\sigma_{t}^{2}}\Big\lVert\mbox{\bf g}_{t}[\mathbf{z}_{it},\boldsymbol\phi_{t}]-\boldsymbol\epsilon_{it}\Big\rVert^2, \end{eqnarray}\]

\[\begin{eqnarray}\label{eq:diffusion_loss_function_3b} L[\boldsymbol\phi_{1\ldots T}] &=& \sum_{i=1}^{I}\sum_{t=1}^{T}\Big\lVert\mbox{\bf g}_{t}[\mathbf{z}_{it},\boldsymbol\phi_{t}]-\boldsymbol\epsilon_{it}\Big\rVert^2\\ &=& \sum_{i=1}^{I}\sum_{t=1}^{T}\bigg\lVert\mbox{\bf g}_{t}\Bigl[\sqrt{\alpha_t}\cdot\mathbf{x}_i+\sqrt{1-\alpha_t}\cdot\boldsymbol\epsilon_{it},\boldsymbol\phi_{t}\Bigr]-\boldsymbol\epsilon_{it}\bigg\rVert^2\nonumber, \end{eqnarray}\]

Implementation

Application to images

U-Net as used in diffusion models for images. The network aims to predict the noise that was added to the image. It consists of an encoder which reduces the scale and increases the number of channels and a decoder which increases the scale and reduces the number of channels. The encoder representations are concatenated to their partner in the decoder. Connections between adjacent representations consist of residual blocks, and periodic global self-attention in which every spatial position interacts with every other spatial position. A single network is used for all time steps, by passing a sinusoidal time embedding (figure 12.5) through a shallow neural network and adding the result to the channels at every spatial position at every stage of the U-Net.

Improving generation speed

Conditional generation

\[\begin{eqnarray} \mathbf{z}_{t-1} = \hat{\mathbf{z}}_{t-1} + \sigma_{t}^2\frac{\partial \log\bigl[\Pr(c|\mathbf{z}_{t})\bigr]}{\partial \mathbf{z}_{t}} + \sigma_{t}\boldsymbol\epsilon. \end{eqnarray}\]

Equations

\[\begin{eqnarray} \mathbf{x}_{t} = \sqrt{1-\beta_{t}}\cdot\mathbf{x}_{t-1} + \sqrt{\beta_{t}}\cdot\boldsymbol\epsilon_{t}, \end{eqnarray}\]

\[\begin{eqnarray} z = a\cdot \epsilon_{1}+b\cdot\epsilon_{2}, \end{eqnarray}\]

\[\begin{eqnarray} \mathbb{E}[z] &=& 0\nonumber \\ \mbox{Var}[z] &=& a^2+b^2, \end{eqnarray}\]

\[\begin{eqnarray} \mathbf{z}_{3} &=&\sqrt{(1-\beta_{3})(1-\beta_{2})(1-\beta_{1}})\cdot\mathbf{x} + \sqrt{1-(1-\beta_{3})(1-\beta_{2})(1-\beta_{1})}\cdot\boldsymbol\epsilon', \end{eqnarray}\]

\[\begin{eqnarray} \mbox{Norm}_{\mathbf{v}}\left[\mathbf{A}\mathbf{w},\mathbf{B}\right] \propto \mbox{Norm}_{\mathbf{w}}\Bigl[(\mathbf{A}^T\mathbf{B}^{-1}\mathbf{A})^{-1}\mathbf{A}^T\mathbf{B}^{-1}\mathbf{v},(\mathbf{A}^T\mathbf{B}^{-1}\mathbf{A})^{-1}\Bigr]. \end{eqnarray}\]

\[\begin{eqnarray} \mbox{Norm}_{\mathbf{x}}[\mathbf{a},\mathbf{A}]\mbox{Norm}_{\mathbf{x}}[\mathbf{b},\mathbf{B}] &\!\! \propto\!\! & \mbox{Norm}_{\mathbf{x}}\Bigl[(\mathbf{A}^{-1}+\mathbf{B}^{-1})^{-1}(\mathbf{A}^{-1}\mathbf{a}+\mathbf{B}^{-1}\mathbf{b}),(\mathbf{A}^{-1}+\mathbf{B}^{-1})^{-1} \Bigr].\nonumber\\ \end{eqnarray}\]

\[\begin{eqnarray} D_{KL}\Bigl[\mbox{Norm}_{\mathbf{w}}[\mathbf{a},\mathbf{A}]\bigl|\bigr|\mbox{Norm}_{\mathbf{w}}[\mathbf{b},\mathbf{B}]\Bigr]\!\! &\!=\!&\!\! \frac{1}{2}\left(\mbox{tr}\left[\mathbf{B}^{-1}\mathbf{A}\right]-d+(\mathbf{a}-\mathbf{b})^{T}\mathbf{B}^{-1}(\mathbf{a}-\mathbf{b})+\log\left[\frac{ |\mathbf{B}|}{|\mathbf{A}|}\right]\right).\nonumber\\ \end{eqnarray}\]

\[\begin{eqnarray} \sqrt{\frac{\alpha_{t}}{\alpha_{t-1}}}=\sqrt{1-\beta_{t}}. \end{eqnarray}\]

\[\begin{eqnarray} \frac{(1-\alpha_{t-1})(1-\beta_{t})+\beta_{t}}{(1-\alpha_{t})\sqrt{1-\beta_{t}}} =\frac{1}{\sqrt{1-\beta_{t}}}. \end{eqnarray}\]