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

2024-10-09

arXiv.org

AI Summary

Discrete diffusion models are inherently non-differentiable, posing a fundamental challenge for solving continuous-space image linear inverse problems. This work pioneers the integration of discrete diffusion models as priors into variational inversion frameworks. We propose a gradient-guided variational approximation coupled with Gumbel-Softmax continuous relaxation to enable end-to-end differentiable optimization. To mitigate absorption-state trapping inherent in standard discrete noise schedules, we introduce a star-shaped noise scheduling scheme. Furthermore, we develop a fully differentiable posterior approximation framework grounded in categorical distributions. Extensive experiments on denoising, super-resolution, and compressive sensing demonstrate that our method achieves performance on par with state-of-the-art continuous diffusion models. These results validate the effectiveness, generalizability, and computational feasibility of discrete diffusion priors for inverse problems in imaging.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Search and Optimization: Mixed Discrete/Continuous Search
• Constraint Satisfaction and Optimization: Mixed Discrete/Continuous Optimization

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• User Modeling, Personalization and Recommendation: User privacy protection in personalized systems

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.

Problem

• Solving linear inverse problems using discrete diffusion models
• Overcoming non-differentiability limitations through variational approximation
• Reducing GPU memory consumption compared to continuous methods

Innovation

• Discrete diffusion models with variational approximation
• Continuous relaxation for non-differentiable problems
• Star-shaped noise process reducing memory consumption