Discrete Inversion: A Controllable Latent Space for Multinomial Diffusion and Masked Generative Models

W1: Cache burden and space efficiency.

A1: Thank you for raising this important point. We acknowledge that the caching requirement could be a potential drawback if cache efficiency is a priority, as it's an inherent limitation of DDPM inversion-based methods. One possible way to mitigate this is by exploring models like Glauber Generative Models $1$ which turns multiclass into binary classification and thus significantly reduces logit space size. We consider this an exciting avenue for future work.

W2: Theoretical analysis.

A2: We completely agree with the reviewer on the importance of providing a theoretical foundation. We would like to clarify that analyzing the mutual information in the context of inversion is challenging, as it involves the model's forward call, which is difficult to analyze directly. To address this, we presented a simple yet prototypical example of DDPM where the posterior mean can be computed in closed form. This enabled us to explicitly compute the mutual information, providing valuable insights into the behavior of DDPM Inversion type of approaches.
That said, we fully acknowledge that a more rigorous theoretical analysis of mutual information and the convergence properties of both continuous and discrete inversion algorithms is necessary. We consider this an important direction for future research and appreciate the reviewer’s suggestion.

W3: Missing reference.

A3: Thank you for pointing this out, we have referenced the two papers in the revised version.

W4: Gumbel-Softmax trick.

A4: Yes, that is correct. We have clarified in the revised manuscript.

W5: Minus sign in Eq 7.

A5: If we rearrange Eq 7 to $y_{t-1} = \hat{y}\_{t-1} + z_t$, we see that $z_t$ is the Gumbel noise or residual that is needed to sample indices $x_{t-1}$ from categorical distribution defined by logits $\hat{y}\_{t-1}$.

W6: Dependency on random seed.
A6: We thank the reviewer for raising this interesting perspective. While we respectfully do not see this as a drastic issue, we acknowledge that it is an inherent property of stochastic methods such as non-ODE-based inversion. The algorithm is designed or intended to keep the same seed. To further address this point, we have added additional results in Figure 16 that explore the behavior of our method when the seed is changed under three settings:

1. Row 1: Different noise maps are used during inversion and inference (editing).
2. Row 2: A different noise map is used at every renoising step in the sampling process, but the inversion steps only use the same noise map.
3. Row 3: A different noise map is used at every renoising step in both inversion and sampling steps. In this case, while the same noise map is typically used for all renoising steps in Paella by default, we ensure that the noise maps in corresponding steps during inversion and inference are the same.

Empirical visual results show that our Discrete Inversion is somewhat robust to the variations and randomness in the seed.

W7: Details on $\lambda\_1$ and $\lambda\_2$.

A7: $\lambda_1$ controls the amount of information injected from the original image, and $\lambda_2$ controls the amount of random noise we want to introduce in inference. Detailed ablations are provided in supplementary section C & D.

W8: Remark 3.1.
A8: The analysis in Remark 3.1 is to provide an explicit analysis of the encoded information in the DDPM Inversion type of latents, which is not provided in the original DDPM Inversion paper.

W9: Clarification of Table 1.
A9: The "Inf" value in Table 1 indicates perfect reconstruction, which is a notable property of our method.
Regarding the difference between the second and third rows in Table 1:

• The second row measures the distance between the reconstructed image and the original image.
• The third row measures the distance between the reconstructed image and the VAE-reconstructed image.

To further clarify, we have added additional visualizations in Figure 14, showing examples of reconstruction with different noise injection mechanisms. Please note that when using a max function to merge the recorded noise and sampled Gumbel noise, perfect reconstruction is not guaranteed.

We appreciate the reviewer’s prompt response.

To store the $z_t$’s for Paella for an image of 256×256 (token map 64×64), each $z_t$ is stored as torch.float16, with the shape of $[1, 8192, 64, 64]$, so storing each $z_t$ needs 67.11MB in GPU memory. We need to store a total of 31 maps, which is about 2080MB (~2GB). However, this tensor does not require gradients and can be loaded on demand. Two possible approaches to improve efficiency:

1. Compression: Since the logit map does not require very high precision, we can use further quantization to reduce cache size or employ sparse representations.
2. Alternative Models: Using models like Glauber Generative Models, which convert multiclass classification into binary classification, could achieve up to an 8,000× memory saving.

Overall, the entire process fits within 16 GB of VRAM, making it feasible to run on consumer-grade GPUs.

W1: Related work discussion.
A1: We thank the reviewer for the valuable suggestion, and we have updated this in the revised submission. We agree that exploring the integration of SEGA-like approaches in discrete settings could be an interesting direction for future work.

W2 & Q2: Operator definition.
A2: We thank the reviewer for pointing this out, and we have addressed this issue in the revised submission.

W3 & Q4: Comparison with ControlNet.
A3: We thank the reviewer for this valuable suggestion and have added comparisons with ControlNet in Table 7 and Figure 11 in the revised submission. In the original manuscript, we did not include ControlNet for the following reasons:

1. Base model differences: ControlNet is built on a continuous diffusion base model (e.g., Stable Diffusion), while our method focuses on discrete diffusion models (e.g., Paella). These fundamental differences make direct comparisons less straightforward.
2. Limitations of ControlNet:
   While ControlNet offers significant flexibility and controllability for introducing additional conditioning inputs (e.g., skeletons, edges, or depth), it has certain limitations. For instance:
   • Users often need to manually specify the subject or object to edit, and the success depends on the availability of corresponding detection models (e.g., there might not be a detection model for "dog skeletons").
   • It requires dataset collection and pretraining if no pre-trained ControlNet exists for a specific use case.
   • It cannot guarantee perfect reconstruction. For example, while it may accurately control the pose of a dog, the generated dog might not resemble the original or specific target (e.g., a cat transformed into a similar-looking dog).
3. Scope and intent of our paper:
   Our primary objective is to introduce a complementary perspective for discrete diffusion models, expanding the existing toolbox. Importantly, our method and ControlNet are not mutually exclusive but rather complementary. If a ControlNet were developed for discrete models like Paella or VQ-Diffusion, our approach could be seamlessly integrated and used alongside it.

We believe that these clarifications, along with the newly added comparisons, provide a more comprehensive understanding of our method in relation to ControlNet.

W1: Highlight key differences from DDPM Inversion.
A1: While extending DDPM Inversion to multinomial diffusion is more straightforward, our main contribution is discrete inversion for masked generative models. The key differences between our approach and DDPM inversion are as follows:

1. Inference process differences. In DDPM Inversion, the inference process involves "fitting" one inference step by recording the Gaussian noise required to sample a given $x\_{t-1}$ from the Gaussian posterior distribution given $x\_t$. This is very different from masked generative models, where $x\_{t-1}$ is NOT sampled from a posterior given $x\_t$. instead, $x\_t$ is obtained from sampled $\\hat{x}\_{0|t}$ by renoising. Thus our discrete inversion is to find the Gumbel noise needed in order to sample the logit of $x\_0$ given $x\_t$. So the method is NOT straightforwardly adaptable from continuous to discrete.
2. Noise type difference. Our method operates in a discrete domain using Gumbel distributions, as opposed to the continuous domain with Gaussian distributions used in DDPM Inversion. This necessitates addressing challenges to discrete generative processes, such as determining appropriate noise injection mechanisms and studying a suitable noise injection function.
3. Another challenge lies in handling the complexities of operating in the logit space. Converting indices to logits accurately is itself a non-trivial problem that we address in our work. For Paella model, we use the model conditioned on $t=0$ as the conversion function, i.e., $\\mathcal{D}(\\cdot, t=0)$. We empirically verified the accuracy of this approach in index space and the mean squared error in logit space. For VQ-Diffusion, we use log onehot as Multinomial Diffusion paper implements everything in logit space and this function is compatible with their implementation.
4. We also conducted ablation studies on several key decision choices in our framework. These include the mask schedule, which differs from DDPM inversion. Unlike DDPM inversion, where an independent q-sample is always used (corresponding roughly to random masking). Visual comparison results are shown in revision Figure 15. We also tested various masking schedules over $t$, with results shown in revision Figure 13 and Table 8. Additionally, we investigated the impact of initial noise and the renoising process on the overall performance (see revision Figure 16). These experiments provide valuable insights into the roles of these components. For instance, in contrast to MaskGIT, which observed that a concave schedule function (e.g., cosine) performs best, we discovered that a convex schedule yields better results in our framework.
5. Noise injection. We introduce a mechanism to control the amount of information injected in editing inference, i.e., down weighting the recorded residual (or latent) and add random Gumbel noise. We also discussed and ablated various injection functions (linear, max, variance preserving). Without this mechanism, simply adding back the recorded residual with a scale of 1 often results in the model failing to edit and instead producing a reconstruction.
6. We highlight our key contributions:
   1. We introduce the first inversion algorithm for discrete diffusion models including both masked generative models and multinomial diffusion.
   2. We experimented on both image and text modalities. For text, we showed an interesting application to turn a BERT model into a text editing model.
   3. We introduce a new dataset for text editing evaluation.

A1:

Thank you for your feedback.

It is indeed fundamental to think of gumbel noise for discrete model in replace of Gaussian noise for continuous model. While our algorithm 2 can be considered as the extension of DDPM-inversion for discrete space, our algorithm 1 is totally different from DDPM-inversion. As shown in the line 7 of algorithm 1, our gumbel noise $z_t$ is obtained via the difference between ground truth $y_0$ and predicted $\hat{y}\_{0\|t}$ , which is different with the line 7 of algorithm 2, the gumbel noise is obtained from the difference between $y\_{t-1}$ and $\hat{y}\_{t-1|t}$, which logit at time $t-1$ is obtained from posterior sampling. Below here, we will reason the design of 2 proposed inverse algorithms 1 and 2 which is heavily based on the training procedure of generative model.

1: Paella, MaskGit and Bert training algorithm

Given data $x\_{0}$, we first create $x\_{t}$, then we train model $f\_{\theta}(x\_{t})$ to predict given $x\_{0}$, using following objective: $CE(f\_\theta(x\_{t}), x\_0)$.

2: VQ-Diffusion training algorithm

Given data $x\_{0}$, we first perturb to obtain $x\_{t}$, then we train model $g\_\phi(x\_{t})$ to predict $x\_{0}$ but using different objective:

$KL(q(x\_{t-1}|x\_t, x\_0))||q(x\_{t-1}|x\_t, g\_\phi(x_t))$

We can see that VQ-diffusion is based on the posterior $q(x_{t-1}|x_t, x_0)$ to train the model while the MaskGit, Bert and Paella does not. This indicate that the $x_0$ predicted from $g_\phi$ is put under an additional posterior constraint while $x_0$ predicted from $f_\theta$ does not. Therefore, if we use the Algorithm 2 (DDPM-inversion) into such models like MaskGit, Bert and Paella, it will limit the ability of editing controllability since the $f_\theta$ prediction is not under posterior constraint. Whereas with VQ-diffusion, if using algorithm 1, we cannot edit since the noise inversion in algorithm 1 does not take into account the posterior distribution while the pretrained model $g_\phi$ is trained under that constraint.

In summary, we present two editing algorithm. While the first one is used for mask generative modelling such as Bert, Maskgit and Paella to maximize the editing controllability of inverse noise, the second algorithm is the first editing algorithm extending DDPM-inversion for discrete diffusion. Our paper is not only limited for VQ-diffusion but also notably can be used for mask generative model.

These discussion will be added to the paper to avoid the confusion about novelty for reader.

We hope that the above response could address your W1 and please let us know if you have any other concern. We are happy to provide more clarification and positively engage into discussion.

Dear reviewer Vpe1, thank you for your detailed and thoughtful review of our paper. We just want to let you know that we have recently provided the rebuttal answer for weakness 1 (the novelty weakness) and weakness 3 (mutual analysis). We would like to know if our rebuttal is reasonable and well address your concern. If yes, could you please consider to revise your score. Thank you very much for your time and consideration!

If you have any other concerns, we would like to hear and provide clarification.