Think While You Generate: Discrete Diffusion with Planned Denoising (2410.06264v2)

Abstract: Discrete diffusion has achieved state-of-the-art performance, outperforming or approaching autoregressive models on standard benchmarks. In this work, we introduce Discrete Diffusion with Planned Denoising (DDPD), a novel framework that separates the generation process into two models: a planner and a denoiser. At inference time, the planner selects which positions to denoise next by identifying the most corrupted positions in need of denoising, including both initially corrupted and those requiring additional refinement. This plan-and-denoise approach enables more efficient reconstruction during generation by iteratively identifying and denoising corruptions in the optimal order. DDPD outperforms traditional denoiser-only mask diffusion methods, achieving superior results on LLMing benchmarks such as text8, OpenWebText, and token-based image generation on ImageNet $256 \times 256$. Notably, in LLMing, DDPD significantly reduces the performance gap between diffusion-based and autoregressive methods in terms of generative perplexity. Code is available at https://github.com/liusulin/DDPD.

• The paper demonstrates that decoupling planning and denoising tasks significantly improves the accuracy and efficiency of discrete generative models.
• It introduces a planner that strategically identifies corrupted tokens and uses an adaptive Gillespie algorithm for optimized sampling.
• Experiments on text and image benchmarks reveal enhanced performance and a better quality-diversity trade-off compared to traditional autoregressive methods.

Discrete Diffusion with Planned Denoising: An Analytical Overview

The paper presents a novel framework, Discrete Diffusion with Planned Denoising (DDPD), aimed at enhancing the generative modeling of discrete data. The research introduces a paradigm shift by delineating the generative process into two distinct components: planning and denoising, unlike traditional approaches which primarily focus on denoising. This separation allows for more adaptive and efficient reconstruction during the generation process.

Key Contributions and Methodology

The authors propose a model architecture where the planner identifies the most corrupted elements that need denoising. This structured approach enables more strategic corrections, optimizing the sequence in which errors are addressed. The framework comprises:

1. Separation into Planner and Denoiser: By decoupling the tasks, the system can explicitly identify and rectify corrupted tokens. The planner estimates the likelihood of corruption, while the denoiser predicts the correct values based on existing noisy sequences.
2. Adaptive Sampling Scheme: Utilizing the Gillespie algorithm enables adaptive time-stepping based on the noise level detected by the planner. This contrasts with fixed-step tau-leaping methods, allowing efficient allocation of computational resources and reducing unnecessary denoising steps.
3. Training Efficiency: The research provides simplified, yet effective, cross-entropy-based training objectives for both components—grounded in maximizing the Evidence Lower Bound (ELBO).

The experiments on text and image benchmarks (such as text8, OpenWebText, and ImageNet data) demonstrate improved quality-diversity trade-offs compared to existing methods. Importantly, DDPD closes the performance gap with autoregressive models, particularly in language tasks measured by generative perplexity.

Experimental Insights

• LLMing: On benchmarks like OpenWebText, the proposed method exhibits significant improvement in generative perplexity against state-of-the-art models including the autoregressive baselines. This indicates that the planner’s strategic ordering of denoising tasks enhances the model’s accuracy in generating text.
• Image Generation: The method outperforms mask-based diffusion in token-based generation on ImageNet, capitalizing on the planner's ability to guide sample corrections flexibly and ensuring that steps are not wasted.

Implications and Future Directions

The paper lays a strong foundation for rethinking the traditional paradigms in discrete diffusion models. It suggests that decoupling complex tasks into simpler sub-tasks—achieving what can be compared to a division of labor—results in more robust learning and generation processes. This could potentially extend beyond language and image applications to other domains where discrete data generation is critical.

Future Work: The approach holds promise for further refinement and scalability in generative AI. Improving planner accuracy and exploring joint tasks where planning and denoising interact could lead to even higher quality outcomes. An exploration into reducing planner computation costs, or explicitly modeling interactions between denoised tokens, might push the capabilities of such frameworks further.

In summary, the proposed DDPD framework enriches the toolkit available for discrete generative modeling, offering a fresh lens through which the efficiency and quality of data generation can be enhanced. The separation of planning from denoising fosters more informed and precise generative processes, setting a new benchmark for researchers in this domain.