将离散扩散模型的跳变时机与方向解耦建模,提升生成质量。
Neural Continuous-Time Markov Chain: Discrete Diffusion via Decoupled Jump Timing and Direction

- 分离建模跳变时机和方向,利用泊松过程优化逆向扩散。
- 在TinyStories上生成困惑度16.36,优于现有方法。
- 适合关注离散数据生成与模型可解释性的研究者。
基于连续时间马尔可夫链(CTMC)的离散扩散模型在语言和离散数据生成中表现优异,但现有方法通常通过代理(如concrete score或干净数据预测)对反向速率矩阵进行整体参数化,而非与CTMC内在的跳变时机与方向分解相一致。本文提出神经CTMC,利用CTMC动态中的泊松结构,通过两个独立网络头分别参数化反向过程的退出率(何时跳变)与跳变分布(跳向何处)。我们证明证据下界(ELBO)可简化为路径空间KL散度,分解为泊松KL(时机)与类别KL(方向),并导出可计算、梯度等价且一致的损失函数。实验表明,在Gemma2-9B评分下,纯均匀初始化的Neural CTMC在TinyStories上达到16.36的生成困惑度(优于GIDD的37.60和MDLM的42.66);在OpenWebText上,相同训练词元预算下,16至128步采样中均取得最优结果(如128步时:Neural CTMC 183.6 vs. MDLM 210.5, GIDD 249.8)。为促进复现,代码与预训练权重已公开于Hugging Face。
原文摘要 · Abstract (English)
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.36$ generative perplexity on TinyStories (vs.\ GIDD $37.60$ and MDLM $42.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.6$ vs.\ MDLM $210.5$ and GIDD $249.8$). To facilitate reproducibility, we release our pretrained weights at https://huggingface.co/Jiangxy1117/Neural-CTMC.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。