$\textit{Jump Your Steps}$: Optimizing Sampling Schedule of Discrete Diffusion Models

Yong-Hyun Park,Chieh-Hsin Lai,Satoshi Hayakawa,Yuhta Takida,Yuki Mitsufuji
2024-10-10
Abstract:Diffusion models have seen notable success in continuous domains, leading to the development of discrete diffusion models (DDMs) for discrete variables. Despite recent advances, DDMs face the challenge of slow sampling speeds. While parallel sampling methods like $\tau$-leaping accelerate this process, they introduce $\textit{Compounding Decoding Error}$ (CDE), where discrepancies arise between the true distribution and the approximation from parallel token generation, leading to degraded sample quality. In this work, we present $\textit{Jump Your Steps}$ (JYS), a novel approach that optimizes the allocation of discrete sampling timesteps by minimizing CDE without extra computational cost. More precisely, we derive a practical upper bound on CDE and propose an efficient algorithm for searching for the optimal sampling schedule. Extensive experiments across image, music, and text generation show that JYS significantly improves sampling quality, establishing it as a versatile framework for enhancing DDM performance for fast sampling.
Machine Learning,Artificial Intelligence,Computation and Language,Computer Vision and Pattern Recognition
What problem does this paper attempt to address?
### What problem does this paper attempt to solve? This paper aims to solve the problem of **slow sampling speed** faced by Discrete Diffusion Models (DDMs) when generating discrete variables. Although discrete diffusion models perform well in handling data with discrete features such as text and music, their sampling process is very time - consuming. Existing acceleration methods such as τ - leaping can speed up the sampling speed, but will introduce **Compounding Decoding Error (CDE)**, which leads to a decline in sample quality. To solve this problem, the author proposes a new method named **Jump Your Steps (JYS)**, which minimizes CDE by optimizing the allocation of discrete sampling time steps without increasing additional computational costs. Specifically, they derive a practical upper bound of CDE and propose an efficient algorithm to search for the optimal sampling schedule. Experimental results show that JYS significantly improves the sampling quality in image, music and text generation tasks and becomes a general framework for enhancing the fast sampling performance of DDM. #### Key problem summary: 1. **Slow sampling speed**: The step - by - step refinement process of discrete diffusion models results in a slow sampling speed. 2. **Compounding Decoding Error (CDE)**: Acceleration methods such as τ - leaping introduce CDE, which reduces sample quality. 3. **Optimize sampling schedule**: Reduce CDE by optimizing the allocation of sampling time steps to improve sampling efficiency and quality. Through these improvements, JYS can significantly accelerate the sampling process of discrete diffusion models while maintaining high - quality samples.