Search papers, labs, and topics across Lattice.
This paper introduces a parallel-in-time sampling method for discrete diffusion models by modifying the $\tau$-leaping algorithm within a Continuous-Time Markov Chain framework. The authors achieve significant acceleration in the sampling process, reducing time complexity from ${\mathcal{O}}(d \log S)$ to ${\mathcal{O}}(\log (d\log S)\cdot \log d)$ while maintaining quality. Empirical results demonstrate up to a 9x speedup in runtime for synthetic distributions and 1.45-1.86x speedups in image and text tasks, highlighting the method's effectiveness across various applications.
Achieving up to 9x runtime speedup in discrete diffusion models could revolutionize their application in real-time tasks like molecular structure prediction and language generation.
Discrete diffusion models are widely used for learning and generating discrete distributions. As the generation process is inherently sequential, the acceleration of sampling is of significant importance. In this work, we parallelize the mainstream $\tau$-leaping algorithm for absorbing discrete diffusion in a Continuous-Time Markov Chain (CTMC) framework. By leveraging the continuous-time stochastic integral form of the $\tau$-leaping algorithm and the Picard iteration method, we achieve parallel-in-time sampling acceleration and provide a proof of exponential-factorial convergence for our algorithm. We improve the overall time complexity of $\tau$-leaping under absorbing settings from ${\mathcal{O}}(d \log S)$ to ${\mathcal{O}}(\log (d\log S)\cdot \log d)$ with respect to NFE. Empirically, our method shows consistent acceleration across synthetic and real-data settings. The new sampler achieves at most $7$--$9\times$ runtime speedup for synthetic distribution, and maintains the same quality with $50\%$ fewer NFE and $1.45$--$1.86\times$ runtime speedups in image/text tasks on a single GPU. Our research expands the potential of discrete diffusion models for efficient parallel inference, with broader implications for applications such as molecular structure and language generation.
