English

Neural Continuous-Time Markov Chain: Discrete Diffusion via Decoupled Jump Timing and Direction

Machine Learning 2026-05-11 v2 Probability

Abstract

Discrete diffusion models based on continuous-time Markov chains (CTMCs) have shown strong performance on language and discrete data generation, yet existing approaches typically parameterize the reverse rate matrix monolithically -- through proxies such as concrete scores (SEDD) or clean-data predictions (MDLM, GIDD) -- rather than aligning the parameterization with the intrinsic CTMC decomposition into jump timing and jump direction. We propose \textbf{Neural CTMC}, which exploits the underlying Poisson structure of CTMC dynamics by separately parameterizing the reverse process through an \emph{exit rate} (when to jump) and a \emph{jump distribution} (where to jump) via two dedicated network heads. We show that the evidence lower bound (ELBO) reduces to a path-space KL divergence between the true and learned reverse processes that factorizes into a Poisson KL for timing and a categorical KL for direction, and admits a tractable, gradient-equivalent and consistent loss. Experimentally, scored by Gemma2-9B, our pure-uniform Neural CTMC achieves 16.3616.36 generative perplexity on TinyStories (vs.\ GIDD 37.6037.60 and MDLM 42.6642.66). On OpenWebText, it attains the best perplexity at the same training-token budget across 16--128 sampling steps among the methods we compare (e.g., at 128 steps: Neural CTMC 183.6183.6 vs.\ MDLM 210.5210.5 and GIDD 249.8249.8). To facilitate reproducibility, we release our pretrained weights at https://huggingface.co/Jiangxy1117/Neural-CTMC.

Keywords

Cite

@article{arxiv.2604.15694,
  title  = {Neural Continuous-Time Markov Chain: Discrete Diffusion via Decoupled Jump Timing and Direction},
  author = {Jingyuan Li and Xiaoyi Jiang and Fukang Wen and Wei Liu and Renqian Luo and Yi Zhu and Zuoqiang Shi and Pipi Hu},
  journal= {arXiv preprint arXiv:2604.15694},
  year   = {2026}
}
R2 v1 2026-07-01T12:13:48.916Z