Fast Solvers for Discrete Diffusion Models: Theory and Applications of High-Order Algorithms

Abstract: Discrete diffusion models have emerged as a powerful generative modeling framework for discrete data with successful applications spanning from text generation to image synthesis. However, their deployment faces challenges due to the high dimensionality of the state space, necessitating the development of efficient inference algorithms. Current inference approaches mainly fall into two categories: exact simulation and approximate methods such as $\tau$-leaping. While exact methods suffer from unpredictable inference time and redundant function evaluations, $\tau$-leaping is limited by its first-order accuracy. In this work, we advance the latter category by tailoring the first extension of high-order numerical inference schemes to discrete diffusion models, enabling larger step sizes while reducing error. We rigorously analyze the proposed schemes and establish the second-order accuracy of the $\theta$-trapezoidal method in KL divergence. Empirical evaluations on GPT-2 level text and ImageNet-level image generation tasks demonstrate that our method achieves superior sample quality compared to existing approaches under equivalent computational constraints.

• The paper presents two new high-order schemes, θ-RK-2 and θ-trapezoidal, that enhance inference accuracy in discrete diffusion models.
• The methodology leverages τ-leaping and weighted estimators to achieve second-order convergence, reducing discretization error from κT to κ²T.
• Empirical results demonstrate improved performance in text and image generation with lower perplexity and better FID scores than first-order methods.

Fast Solvers for Discrete Diffusion Models: Theory and Applications of High-Order Algorithms

Introduction to Discrete Diffusion Models

The paper "Fast Solvers for Discrete Diffusion Models: Theory and Applications of High-Order Algorithms" presents novel high-order numerical solvers specifically designed for discrete diffusion models. Discrete diffusion models are crucial for generative modeling of discrete data types such as text and images. These models operate within a discrete state space, unlike continuous diffusion models. A major challenge with discrete diffusion models lies in efficiently performing inference, which often suffers from either high computational cost or first-order accuracy limitations.

Development of High-Order Schemes

Two primary algorithms introduced are the $\theta$-Runge-Kutta-2 ($\theta$-RK-2) and $\theta$-trapezoidal methods. These high-order schemes are adapted from stochastic differential equations and chemical reaction simulations to the discrete diffusion domain:

1. $\theta$-Runge-Kutta-2 Method: This two-stage algorithm leverages $\tau$-leaping to an intermediate step, then improves upon it using a weighted estimator for an entire timestep. The algorithm's formulation mirrors that of second-order Runge-Kutta methods used in ordinary differential equations.
2. $\theta$-Trapezoidal Method: This method features a two-phase approach within each timestep, separating it into sub-intervals with varied intensities for balanced computation. The trapezoidal method is resilient and achieves unconditional second-order accuracy.

Theoretical Analysis

The analysis establishes that the $\theta$-trapezoidal method guarantees second-order convergence, providing improvements over traditional first-order approaches:

• Assumptions: Several assumptions underpin the theoretical framework, such as the exponential convergence of the forward process and the bounded regularity of the intensity functions.
• Convergence Guarantees: The paper rigorously demonstrates that the $\theta$-trapezoidal method reduces discretization error significantly, from $(\kappa T)$ to $(\kappa^2 T)$, compared to first-order methods.

Empirical Validation

The implementations were validated across various generative tasks, using both toy models and large-scale benchmarks such as text and image generation:

• 15-Dimensional Toy Model: Demonstrates the second-order accuracy in practice, with a clear reduction in error as step sizes increase.

  Figure 1: An illustrative application of the uniformization algorithm to discrete diffusion models for text generation. The x-axis denotes the time of the backward process, and the y-axis denotes the frequency of jumps reflected by NFE. Perplexity convergence occurs well before the NFE experiences unbounded growth.

• Text Generation: Using RADD for masked discrete diffusion models, the $\theta$-trapezoidal method consistently outperformed baseline methods in terms of generative perplexity under equivalent computational constraints.
• Image Generation: When generating high-resolution images in the ImageNet dataset, the proposed method achieved better FID scores, indicating higher sample quality compared to existing techniques.

Discussions and Future Directions

The paper opens pathways for improvements in discrete diffusion model inference, specifically through high-order numerical algorithms:

• Algorithmic Robustness: The $\theta$-trapezoidal method's robustness across varying parameters showcases its versatility in generative tasks.
• Potential Extensions: Future work includes exploring adaptive step sizes and further integrating these techniques with advanced sampling methodologies to enhance performance across diverse discrete data domains.

Conclusion

The work presents significant advances in inference algorithms for discrete diffusion models, providing both theoretical insights and practical implementations that offer improved accuracy and efficiency. The high-order schemes developed here serve as a benchmark for future investigations into discrete generative models and their applications in complex scientific and industrial domains.