# Denoising Diffusion Models on Model-Based Latent Space
## 3 Proposed Method
### Latent-Space Diffusion Models
#### 3.2.2 Categorical Posterior Distribution  
Despite the promising results in training the DDPM in latent space, thanks to the approximation introduced above by optimizing for the training objective of Equation (4), the sample quality of the decoded images As an additional enhancement, we revise the formulation of the inverse process to better align with the inherent characteristics of VQ-encoded images. Building upon the training objective presented in Equation (5), learning the reverse process within the latent space by targeting the uncorrupted sample \(z\in\mathbb{R}^{\frac{n}{\tau}\times\frac{1}{\tau}\times\mathcal{I}}\) implies predicting the index \(i\) corresponding to the suitable codebook entry \(c_{i}\in\mathcal{C}_{j}\) for each codebook \(j\in(1,\ldots,\bar{c})\) at each spatial location. Instead of directly regressing the value of \(c_{i}\), we propose an alternative strategy, wherein the inverse process targets a categorical distribution for each codebook entry across all utilized codebooks. Assuming that all codebooks contain the same number of entries, the output for the reverse process can be parametrized as \(\bar{z}\in\mathbb{R}^{\frac{n}{\tau}\times\frac{1}{\tau}\times\mathcal{I}| \mathcal{C}|}\) so that the recovered sample \(\bar{z}\) could be obtained by sampling the corresponding categorical distribution at each spatial location. The forward  
Figure 4: Effect of temperature smoothing of categorical distribution. diffusion process instead is kept unchanged, with \(z_{t}\) being defined by introducing noise to \(z\) in the index space. Since the sampling operation is nondifferentiable, in the training phase, we leverage the Gumbel-Softmax trick [30] to approximate the sampling of the categorical distribution; this is achieved by perturbing the logits \(\hat{z}^{j}\) for each category with Gumbel-distributed noise \(G_{i}\sim Gumbel(0,1)\), and then approximating the one-hot vector representing the categorical choice by applying the Softmax function to the perturbed logits. Mathematically, this can be expressed as  
\[\overline{z}^{j}=softmax\left(\frac{\hat{z}^{j}+G}{\tau}\right)\qquad j=1, \ldots,\tilde{c} \tag{20}\]  
\[G\sim Gumbel(0,\mathrm{Id})\]  
The hyperparameter \(\tau\) is a temperature parameter which increases the entropy of the distribution; hence, for \(\tau\to 0\), it approaches the \(argmax\) function, and for \(\tau\to\infty\), it approaches the uniform distribution, as depicted in Figure 4. Notably, this approximation allows us to explicitly sample from the codebook distribution instead of the index distribution. Consequently, we can decode the predicted codewords in a fully differentiable manner, as demonstrated by the equation  
\[\tilde{c}=[\overline{z}^{1}\mathcal{C}_{1},\cdots,\overline{z}^{\tilde{c}} \mathcal{C}_{\tilde{z}}] \tag{21}\]  
and the decoded image \(\tilde{x}\) can be recovered from \(\tilde{c}\) depending on the utilized VQ specific (the complete procedure is schematized in Algorithm 1). By leveraging this formulation, we could compute the training objective in image space as in Equation (5), while still performing the generation and the gradient computation in the latent space. It is worth noticing that, in this way, the latent codes for \(z\) are never elicited in the training objective.