G2D2: Gradient-guided Discrete Diffusion for image inverse problem solving

Authors: Naoki Murata, Chieh-Hsin Lai, Yuhta Takida, Toshimitsu Uesaka, Bac Nguyen, Stefano Ermon, Yuki Mitsufuji
Primary area: generative models
Keywords: diffusion model, discrete diffusion model, inverse problems, categorical data

Abstract:
Recent literature has effectively utilized diffusion models trained on continuous variables as priors for solving inverse problems. Notably, discrete diffusion models with discrete latent codes have shown strong performance, particularly in modalities suited for discrete compressed representations, such as image and motion generation. However, their discrete and non-differentiable nature has limited their application to inverse problems formulated in continuous spaces. This paper presents a novel method for addressing linear inverse problems by leveraging image-generation models based on discrete diffusion as priors. We overcome these limitations by approximating the true posterior distribution with a variational distribution constructed from categorical distributions and continuous relaxation techniques. Furthermore, we employ a star-shaped noise process to mitigate the drawbacks of traditional discrete diffusion models with absorbing states, demonstrating that our method performs comparably to continuous diffusion techniques. To the best of our knowledge, this is the first approach to use discrete diffusion model-based priors for solving image inverse problems.

Decision: reject
Reviewer scores: 5, 6, 6, 5

Reviews and discussion:
<thread>
creation date:
1729012649938

thread_type:
official_review

<thread_object>
creator: reviewer_gc4q

<content>
summary:
This paper solves inverse problems in image processing by leveraging discrete diffusion models. Unlike prior approaches using continuous diffusion models, G2D2 introduces a star-shaped noise process and variational relaxation to enable the use of discrete diffusion models in continuous image reconstruction tasks, such as super-resolution and Gaussian deblurring.

strengths:
1. G2D2 is original in its use of discrete diffusion models for solving inverse problems, which were previously limited to continuous domains. The introduction of the star-shaped noise process and continuous relaxation techniques shows a creative solution to a challenging problem of applying discrete models to continuous tasks.

weaknesses:
1. While G2D2 introduces innovative techniques, the computational complexity of the variational optimization and continuous relaxation might be higher than other methods. The paper could benefit from more detailed discussions on the computational trade-offs and potential optimizations to make the approach more efficient.


2. The paper focuses primarily on image processing tasks such as super-resolution and deblurring, with limited exploration of other possible applications. Extending the method to more diverse and complex inverse problems, such as in other imaging modalities or higher-dimensional data (e.g., 3D medical imaging), would improve the robustness and generalizability of G2D2.

3. Although G2D2 is compared with several continuous diffusion models, a more thorough comparison with the latest emerging techniques, such as latent diffusion models or score-based generative models, would strengthen the evaluation. 

4. The paper emphasizes that the star-shaped noise process can correct early-stage errors, but it does not explore the limitations of this mechanism in depth. There could be scenarios where certain types of errors are less tractable, and a discussion on the failure modes of G2D2


5. More related work should be cited such as "Diffusion modeling with domain-conditioned prior guidance for accelerated mri and qmri reconstruction".

questions:
1. How does the star-shaped noise process enable the "re-masking" operation, and why is this beneficial for solving inverse problems?

2. What are the advantages of using a star-shaped noise process over a standard Markov process for sampling in discrete diffusion models?

3. What does the term "conditionally independent" imply about the relationship between the noisy discrete latents z1, z2,...,zT given z0 in the star-shaped noise process?

4. What does Property 2 suggest about the relationship between the marginal distributions of the joint distribution q_sampling and the star-shaped noise process graphical model?

5. How does the mean-field structure of both q(zt - 1 | z0) and p_\alpha (z0|zt,y) impact the marginalization process in the G2D2 method?

6. In the algorithm, What challenges might arise if the optimization process does not incorporate the values from the previous time step during initialization?
</content>
</thread_object>

</thread>
<thread>
creation date:
1729217400150

thread_type:
official_review

<thread_object>
creator: reviewer_zmjs

<content>
summary:
The authors propose G2D2, a diffusion model-based inverse problem solver (DIS) that uses a discrete diffusion model. To the best of my knowledge, this is the first work to demonstrate that this is possible. Since using the usual mask-observing Markov process of discrete diffusion models makes it hard to correct for the errors arising in the earlier stages of the sampling, the authors propose to use a star-shaped diffusion, where similar to DDIM, $z_t$s are conditionally independent given $z_0$. During inference, the parameters $\alpha$ of the variational reverse categorical distribution are optimized by balancing the prior and the likelihood, which is grounded by sound theory. Experiments are conducted on a standard FFHQ/ImageNet settings.

strengths:
1. To the best of my knowledge, this is the first work to target using discrete diffusion models for solving inverse problems. G2D2 will open up new opportunities for testing out different discrete diffusion priors.

2. The theory is sound, and the resulting algorithm is straightforward to understand and implement. This resembles how most of the current DIS is implemented in practice, where the predicted $x_0$ is used to compute the likelihood loss, and the sampling to $x_{t-1}$ is conducted with a DDIM sampling step.

weaknesses:
1. The results are weak. This is somewhat understandable given that the pre-trained diffusion prior is suboptimal, and latent diffusion models (whether continuous or discrete) tend to be inferior to pixel-based methods due to the existence of decoders.

2. The presentation could be improved. Using a star-shaped noise process is one of the crucial contributions of the work, but this is first introduced in 3.1. A brief review of this before the main section would be beneficial for better understanding. 

3. Many readers interested in the work will already be familiar with the family of continuous diffusion-based methods, and in many points, G2D2 resembles them. It would be beneficial for the authors to draw links to the continuous counterpart, especially in the construction of the sampling.

4. Motion inverse problem solving is demonstrated without any comparison.

5. It is discussed that G2D2 can be used with MaskGIT, which is one of the main strength of the work. However, MaskGIT is only used for motion inverse problem solving, which only consists of a small proportion of the experiments.

questions:
Did the authors try using MaskGIT prior to image inverse problems? It is a better prior than VQ-diffusion, and it would be surprising if the results did not improve by simply switching the pre-trained model.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730319585329

thread_type:
official_review

<thread_object>
creator: reviewer_eysn

<content>
summary:
Authors propose Gradient-guided Discrete Diffusion (G2D2), a novel method for solving linear inverse problems using discrete diffusion as priors. Limitations of discrete diffusion priors are overcame by approximating posterior with variational distribution constructed from categorical distributions and continuous relaxation techniques. Authors demonstrate their method on super-resolution and Gaussian deblurring tasks on ImageNet and FFHQ datasets. G2D2 is compared against popular baselines such as DPS, DDRM, PSLD and ReSample and found to be competitive against continuous diffusion models.

strengths:
* The paper is written very well.
* Using discrete diffusion priors for solving inverse problems is a novel direction that is not well explored.
* Usage of star-shaped noise process is motivated well with big gain in downstream performance compared to the standard Markov noise process.

weaknesses:
* While I appreciate the novelty of using discrete diffusion models in the context of inverse problems, the advantages of using them against continuous counterparts are not motivated well.
* See the questions below.

questions:
* Line 436: "For the image inverse problem experiments, we used text prompts for the VQ-Diffusion model...", do the authors use similar text conditioning for competing methods? If not would it give unfair advantage to G2D2?
* In the limitations section, it is mentioned that "G2D2 does not significantly surpass its continuous counterparts in terms of computational speed or performance". Is there a clear advantage of using discrete diffusions? 
* In the appendix, authors provide the hyperparameters used for competing methods (DPS, DDRM, etc.). Are those values taken directly from the corresponding papers or a hyperparemeter search was conducted to find them? 
	* If taken directly, I would recommend tuning them separately since in my experience these values are not robust against small changes in the problem setup (some papers add noise to the image in the range [-1,1] some in [0,1] etc.).
	* If the latter one, it would be good to describe which range was searched over how many validation samples.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730708079575

thread_type:
official_review

<thread_object>
creator: reviewer_twfi

<content>
summary:
The authors propose a method for solving inverse problems using discrete diffusion models with gradient guidance. They utilize the Gumbel trick to relax the categorical distribution, allowing for gradient computation. Additionally, they introduce star-shaped diffusion models to address the limitations of conventional discrete diffusion models. This approach increases the probability of transitioning to the absorbing state, enabling correction of erroneous codes. In experiments, the proposed method outperforms previous continuous pixel/latent domain diffusion model-based inverse solvers in super-resolution and deblurring tasks. They also apply it to a path-following task using a generative masked motion model.

strengths:
The authors propose novel methods to address challenges arising when adapting previous approaches to the new task of discrete inverse problems.

weaknesses:
The motivation for using a discrete diffusion model in the image domain is lacking. Given the availability of options to reduce computational burden, such as low-precision floating points or lighter models, the necessity of a discrete representation is not clearly justified. It would be more convincing if the proposed method demonstrated its advantages in areas where discrete representations are essential, such as language or molecular modeling. To strengthen the paper, consider the following suggestions:
- Compare the discrete approach directly with continuous models that use low-precision or lightweight architectures.
- Include experiments in domains where discrete representations are inherently suitable, such as text generation or molecular modeling.
- Discuss any potential advantages of discrete representations in image domains that may not be immediately clear from the current results.

questions:
- How does the discrete star-shaped diffusion process compare with its continuous counterparts in terms of similarities and differences? Could you discuss the specific challenges the authors faced when adapting star-shaped diffusion to discrete spaces and the strategies used to address them? Additionally, please elaborate on any unexpected advantages or limitations of the discrete version compared to the continuous model.
- Could discrete star-shaped diffusion models be integrated with diffusion model-based solvers like DDRM or DDNM?
- The proposed method outperforms even continuous-domain algorithms in image tasks. Were the experiments conducted under computational constraints? If so, please specify these restrictions and provide evaluations under unrestricted conditions as well.

Possible Errors
- Proof of Lemma B.3: The proof appears incorrect. In Lines 975 to 980, the claim that the reverse transition inverts the forward process and leads to equality with $\delta_{z_0, z_0{\prime}}$ should be revised. While the lemma’s conclusion is unaffected, replace exact equalities with equality in distribution where appropriate.
- Line 265: $z_t \to z_{t-1}$.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732366558127

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Message to the Reviewers

comment:
Dear Reviewers,

We would like to express our sincere gratitude for your thorough review of our paper. We have carefully considered each of your valuable comments and questions.

## Revisions

Based on the feedback we received, we have made the following revisions to our paper. Please find below a summary of the modifications:

1. In the sections explaining our theoretical framework, we corrected typographical errors and rewrote the content to make it clearer and more easily understandable.
2. The proof of Lemma B.3 in the Appendix has been corrected.
3. We have added several references to practical applications in the Related Works section of the Appendix.
4. As additional experiments in the Appendix C.7, we have included the results of our investigation into the impact of text conditioning on performance.
5. We have added an analysis of failure modes to Appendix C.8.
6. We updated the notation in the graphical model figure.

## Response to Key Feedback

We received feedback from our reviewers, primarily pointing out that "The motivation for using discrete diffusion models in the image domain is lacking." We consider this observation to be extremely important for exploring this topic. We would like to address this point as follows:

1. It is acknowledged that this research is the first paper to use discrete diffusion models for solving image inverse problems. However, it may be difficult to demonstrate significant advantages of using discrete diffusion models based solely on the current results.

2. Nevertheless, it is noteworthy that we were able to show competitive results with less GPU memory usage and faster inference compared to similar continuous space methods like PSLD and Resample, in terms of applying diffusion models on the learned latent space. Additionally, since performance is also influenced by the quality of pre-trained models, we may expect further performance improvements if better models than VQ-Diffusion become available in the future.

3. Furthermore, when considering the potential benefits of discrete diffusion models in inverse problem solving, recent studies such as Transfusion [1] and Show-o [2] might expand the possibilities in this field. These models enable the generation of multimodal data (text and images) using a common discrete token space and Transformer model, which we believe could be compatible with our approach.

## References

[1] Zhou et al., "Transfusion: Predict the next token and diffuse images with one multi-modal model", arXiv preprint, 2024, https://arxiv.org/abs/2408.11039

[2] Xie et al., "Show-o: One single transformer to unify multimodal understanding and generation", arXiv preprint, 2024, https://arxiv.org/abs/2408.12528
</content>
</thread_object>

</thread>


Metareview:
metareview:
The main claim of the work is to be the first approach to use discrete diffusion model-based priors for solving image inverse problems. 

While this is an interesting direction, and the paper is solid (and minor mistakes were fixed in the reviewer discussion phase), the motivation of using discrete diffusion models in the image domain still remains unclear. Empirical results presented in the paper are lagging behind the state-of-the-art. If other domains like molecular generation or motion modeling are potential candidates for this method, the work will benefit from a major revision with another round of reviews to make this point more appreciable. 

I recommend to reject the paper at this point and encourage the authors to take the reviewers feedback into account to make the motivation of the paper more appealing.