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

会议: NeurIPS2026
arXiv: 2604.15694
代码: https://huggingface.co/Jiangxy1117/Neural-CTMC
领域: 其他
关键词: 离散扩散模型、连续时间马尔可夫链、速率矩阵分解、泊松KL散度、自纠错采样

一句话总结¶

本文提出 Neural Continuous-Time Markov Chain (Neural CTMC)，将离散扩散逆过程的速率矩阵内生地解耦为出射率（何时跳跃）与跳跃分布（跳向何处），在路径空间将变分下界严格分解为泊松 KL 与分类 KL，使纯均匀噪声前向过程在文本生成与自纠错任务中显著超越主流掩码扩散基线。

研究背景与动机¶

扩散模型在连续领域（如图像、音频生成）已取得巨大成功，其核心建立在连续分数函数和朗之万动力学之上。然而，将这一成功推广至语言、蛋白质和分子图等离散数据时面临本质阻碍，因为离散状态空间中无法直接定义空间导数与连续漂移。为此，离散扩散模型普遍采用有限状态空间上的连续时间马尔可夫链（Continuous-Time Markov Chain, CTMC）来形式化前向加噪与逆向生成。在现有主流方法中，逆向转移速率矩阵要么通过 concrete score 比率进行隐式参数化（如 SEDD），要么依托吸收态掩码过程退化为对干净数据的预测（如 MDLM 与 GIDD）。

然而，这种整体性参数化（monolithic parameterization）忽略了 CTMC 动力学一个极其根本的数学结构特性：根据经典随机过程理论与 Gillespie 随机模拟算法，任意非负转移速率矩阵都可以唯一且正交地分解为一个控制停留时间的长短与何时转移的出射率标量，以及一个控制离开当前状态后具体跳转至哪里的条件类别概率分布。现有方法强行用单一网络代理拟合整体速率，不仅破坏了这一自然的概率测度分解，更导致训练目标中的时序与类别更新紧密耦合，限制了模型学习更灵活跳跃动力学的能力。

更为严峻的矛盾体现在前向加噪过程的选择上。当前表现较优的模型高度依赖吸收态（掩码）前向过程，但这导致生成过程中一旦某个位置被“解掩码”，该位置的状态便在后续反向轨迹中被彻底冻结，丧失了自我纠错与状态回退的能力；相比之下，允许在所有离散状态间自由跳转的均匀噪声前向过程虽具备可逆性与可编辑性，但在既有参数化下生成性能长期被掩码模型大幅压制。核心 idea：本文遵循 CTMC 内生结构，将逆过程速率矩阵显式解耦为出射率与跳跃分布双预测头，并在路径空间导出天然解耦的泊松-分类变分目标，从而在纯均匀前向噪声下彻底激活离散扩散的自纠错与高质量生成潜力。

方法详解¶

整体框架¶

Neural CTMC 将离散空间 $\mathcal{S}=\{1,\ldots,S\}$ 上的逆过程建模为一个时间非齐次 CTMC。针对当前含噪状态 $x_t$ 与连续时间 $t$，网络不直接拟合整体转移速率，而是通过两个专门的解耦预测头分别输出两类物理量：出射率标量 $\lambda_t^\theta(x_t)$ 用于控制泊松跳转强度（何时发生状态改变），跳跃分布向量 $r_t^\theta(\cdot \mid x_t)$ 用于给出在发生跳跃时各非自身状态的类别概率（跳向哪个新 token）。

在训练阶段，模型针对样本在连续时间上的路径测度优化证据下界（ELBO），该目标在数学上被严格证明等价于真实逆向过程与神经网络逆向过程在路径测度上的 KL 散度，并进一步完美分裂为衡量跳跃时刻的泊松 KL 项与衡量跳跃方向的分类交叉熵项。在推理采样阶段，该解耦形式可直接适配类似化学反应模拟的 $\tau$-Leaping 并行跳跃算法或欧拉离散化更新，并在全轨迹保留每个 token 的双向转移能力，支持即插即用的轻量级自纠错。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入数据：干净文本序列 x0<br/>前向加噪得到时间 t 处的离散序列 xt"] --> B["Transformer 主干特征提取器<br/>DiT 编码器提取时序与上下文表征"]
    B --> C["双头解耦参数化架构<br/>Φ_θ(xt, t) = (λ_t^θ, r_t^θ)"]
    C --> D["路径空间变分目标分解<br/>泊松 KL（何时跳跃） + 加权分类交叉熵（跳向何处）"]
    D --> E["反向动力学采样与状态生成<br/>τ-Leaping 采样 / Euler 离散化转移"]
    E --> F["轻量后验自纠错迭代<br/>恢复无偏干净分布并重采样错误位置"]

关键设计¶

1. 逆过程速率矩阵的正交双头参数化：打破代理参数化的整体性耦合

传统离散扩散往往将逆过程矩阵 $\hat{R}_t(i,j) = R_t(j,i) \frac{q_t(j)}{q_t(i)}$ 当作黑盒，采用 concrete score $s_\theta(x_t, t)_j \approx \frac{q_t(j)}{q_t(x_t)}$ 或预测干净数据 $x_\theta(x_t, t)$ 进行整体拟合。而根据任意 CTMC 的物理分解，离开状态 $i$ 的总速率与具体跳向状态 $j$ 的概率在统计上完全独立： $$ R_t^\theta(i,j) = \lambda_t^\theta(i) \cdot r_t^\theta(j \mid i), \quad j \neq i $$ 其中出射率 $\lambda_t^\theta(i) = \sum_{j \neq i} R_t^\theta(i,j) \in \mathbb{R}_{>0}$，跳跃分布 $r_t^\theta(\cdot \mid i) \in \Delta^{S-1}$ 位于 $S-1$ 维单纯形上。网络直接分出出射率头（标量回归）与跳跃分布头（Softmax 输出），使神经网络能够将表征容量正交分配给“当前 token 此时是否不稳定、是否需要被修改”与“如果需要修改，上下文最支持替换成什么语义”，显著降低了学习复杂度。

2. 路径空间 ELBO 的泊松-分类因式分解：推导一致且梯度等价的变分损失

为训练解耦参数化网络，作者从连续时间跳跃过程的路径测度（Path Measure）比率入手。在跳跃时刻 $0 \le t_1 < \cdots < t_n < T$ 与转移路径上，逆向路径与前向重要性采样的对数 Radon-Nikodym 导数由先验、跳跃比率和存活概率积分差三部分组成。经过 Campbell-Mecke 公式求期望化简，反向测度之间的路径 KL 散度在数学上展现出了极具美感的分解形式： $$ \mathrm{KL}(\hat{\mathbb{Q}}|\mathbb{P}^{\theta}) = \int_{0}^{T}\sum_{i\in\mathcal{S}}q_{t}(i)\left[\mathrm{KL}^{\text{Poi}}!\bigl(\hat{\lambda}{t}(i)|\lambda^{\theta}}(i)\bigr) + \hat{\lambda{t}(i)\cdot\mathrm{KL}^{\text{Cat}}!\bigl(\hat{r}t $$ 其中泊松 KL 散度 }(\cdot\mid i)|r^{\theta}_{t}(\cdot\mid i)\bigr)\right]\mathrm{d$\mathrm{KL}^{\text{Poi}}(\lambda\|\lambda^\theta) = \lambda \log\frac{\lambda}{\lambda^\theta} - \lambda + \lambda^\theta$ 严格惩罚出射率的偏差，而类别 KL 散度 $\mathrm{KL}^{\text{Cat}}(\hat{r}\|r^\theta)$ 以前向出射率期望为权重惩罚跳转方向。由于真实边缘分布 $q_t(i)$ 难求，论文通过期望变换推导出条件变分代理损失 $\mathcal{L}_{\mathrm{KL}}(\theta)$ 与 $\mathcal{L}_{\mathrm{cond}}(\theta)$，并证明了在正则条件下 $\nabla_\theta \mathcal{L}_{\mathrm{KL}}(\theta) = \nabla_\theta \mathrm{KL}(\hat{\mathbb{Q}}\|\mathbb{P}^{\theta})$，且临界点完全一致。特别地，当约束为掩码前向过程时，该损失严格退化为 MDLM 的交叉熵损失，表明 MDLM 仅是出射率退化为固定常数、跳跃分布等价于 $x_\theta$ 时的特例。

3. 基于内生速率的 $\tau$-Leaping 采样与纯均匀噪声自纠错机制

得益于参数化与动力学模拟逻辑的高度同构，模型在反向生成时支持经典的 $\tau$-Leaping 采样：序列各位置 $\ell$ 在步长 $\tau$ 内独立抽取指数分布等待时间 $\Delta_\ell \sim \mathrm{Exp}(\lambda_{t,\ell}^\theta)$，仅当 $\Delta_\ell < \tau$ 时才按照跳跃分布 $r_{t,\ell}^\theta(\cdot \mid x_t^{(\ell)})$ 采样转移。相比于掩码扩散“一旦还原即锁定”的缺陷，均匀前向过程使得任意生成步的 token 都保留了被重新修改的概率。作者进一步提出了轻量级自纠错算法：从学得的模型输出中反解当前位置模型所隐含的无偏干净分布 $\hat{p}_{0,t,\ell}^\theta$，在局部识别出最大置信度不一致的 token，并按照退火温度 $\tau_{\mathrm{sc}}$ 进行重采样替换，以极低开销消除了生成文本中的局部语法和逻辑瑕疵。

损失函数 / 训练策略¶

在具体训练实现中，为避免在时间接近终点 $t \approx T$ 时出现数值抵消不稳定，作者采用对数值下溢更具鲁棒性的等价数值稳定损失形式： $$ \mathcal{L}{\mathrm{cond}}(\theta) = \mathbb{E}\right) \right] $$ 其中 } \left[ \sum_{j \neq i} \lambda_t^\theta(i) r_t^\theta(j \mid i) - \sum_{j \neq i} R_t(j,i) \frac{q_{t|0}(j \mid x_0)}{q_{t|0}(i \mid x_0)} \log \frac{\lambda_t^\theta(i) r_t^\theta(j \mid i)}{R_t(j,i)} + \sum_{j \neq i} R_t(j,i) K!\left(\frac{q_{t|0}(j \mid x_0)}{q_{t|0}(i \mid x_0)$K(a) = a(\log a - 1)$ 为常数核项。前向噪声采用标准线性调度 $\alpha_t = 1 - t, \beta_t = t$；训练时时间采样于 $\mathcal{U}(\epsilon, T-\epsilon)$。主干网络统一采用 12 层 DiT（隐藏层维度 768，12 个注意力头，参数量约 163M），使用 AdamW 优化器，在均匀前向噪声下进行端到端优化。

实验关键数据¶

主实验¶

论文在 TinyStories 和大规模网络语料 OpenWebText 上对 Neural CTMC 进行了系统评估。所有对比模型均采用完全相同的 12 层 DiT 主干、参数量（~163M）以及 GPT-2 分词器（词表大小 50,257）。生成样本均抽取 1024 个无条件样本，并使用外部高性能大语言模型 Gemma2-9B 计算生成困惑度（PPL $\downarrow$）。

在 TinyStories 任务上，50 步采样下 Neural CTMC 在 45 个 epoch 内取得 16.36 的生成困惑度，大幅超越掩码前向过程的代表模型 GIDD ($p_{\mathrm{unif}}=0.0$ 时 PPL 为 37.60) 与 MDLM (PPL 为 42.66)。

下表为 OpenWebText 上不同采样步数下的主实验生成困惑度对比（加粗表示同预算 262B 训练 token 下最优）：

前向加噪类型	方法	训练 Token 预算	最大序列长度	16 步 PPL (↓)	32 步 PPL (↓)	64 步 PPL (↓)	128 步 PPL (↓)
Mask (吸收态)	SEDD	682B	1024	825.5	337.9	186.5	127.2
Mask (吸收态)	MDLM	262B	1024	1432.8	553.7	301.6	210.5
Mask (吸收态)	GIDD ($p_u=0.0$)	262B	512	2773.1	993.7	529.6	414.3
Mixture (混合)	GIDD ($p_u=0.1$)	262B	512	702.0	398.9	270.8	249.8
Mixture (混合)	GIDD ($p_u=0.2$)	262B	512	770.4	430.1	344.3	293.0
Uniform (纯均匀)	SEDD (复现)	262B	512	963.9	353.2	247.7	204.1
Uniform (纯均匀)	Neural CTMC (本文)	262B	512	578.3	258.8	189.7	183.6

消融实验¶

论文通过在相同生成的 OpenWebText 原始样本上施加自纠错更新，探究不同温度与前向设置对文本多维度质量的提升（由 GPT-5.2 在 1-10 分维度进行裁判，表中百分比为相对原始样本的变化，* 代表超过 $2\sigma$ 显著提升）：

模型配置 / 自纠错变体	清晰度 (Clarity)	语法性 (Grammaticality)	事实性 (Factuality)	写作风格 (Writing Style)	创意性 (Creativity)
原始样本基线 (Euler 128步)	3.78	3.91	2.91	3.58	5.21
GIDD 自纠错 ($p_u=0.0, \tau_{\mathrm{sc}}=0.1$)	3.19 (-15.7%)*	3.01 (-23.2%)*	3.03 (+4.3%)	3.08 (-13.9%)*	4.57 (-12.3%)*
GIDD 自纠错 ($p_u=0.1, \tau_{\mathrm{sc}}=0.1$)	3.81 (+0.6%)	4.00 (+2.2%)	3.08 (+5.7%)	3.63 (+1.5%)	5.19 (-0.4%)
GIDD 自纠错 ($p_u=0.2, \tau_{\mathrm{sc}}=0.5$)	4.00 (+5.8%)*	4.25 (+8.6%)*	2.98 (+2.5%)	3.88 (+8.4%)*	5.54 (+6.2%)*
Neural CTMC 自纠错 ($\tau_{\mathrm{sc}}=0.1$)	3.81 (+0.6%)	4.40 (+12.4%)*	3.09 (+6.2%)	3.77 (+5.4%)	5.63 (+8.0%)*
Neural CTMC 自纠错 ($\tau_{\mathrm{sc}}=0.5$)	3.82 (+1.1%)	4.06 (+3.8%)	3.01 (+3.4%)	3.74 (+4.5%)	5.57 (+6.8%)*

关键发现¶

纯均匀扩散击败掩码扩散：在相同 262B token 预算下，纯均匀设置的 Neural CTMC 在 16 到 128 各步数区间全面击败吸收态基线 MDLM 与 GIDD。尤其在少步数（32 步）下，Neural CTMC 相比 MDLM 生成困惑度从 553.7 骤降至 258.8（提升超过 2.1 倍），证明了解耦结构能极大加速反向轨迹的状态坍缩与收敛。
与大规模算力基线对比：SEDD (mask) 虽在 128 步下取得 127.2，但其耗费了高达 682B 的训练 tokens（比本文多 2.6 倍）且上下文窗口为 1024；在 16 与 32 步低步数推理时，Neural CTMC 依然保持明显优势。
Scaling Law 规律：在参数量（80.1M, 169.7M, 424.5M）与数据量（100M, 300M, 1B tokens）的缩放研究中，验证集损失与计算 FLOPs 呈现稳健的幂律边界（拟合幂律指数 $\alpha \approx 0.16, R^2 \approx 0.93$），表明解耦 CTMC 在大模型规模下依然具有优异的可扩展性。
状态可逆性带来的自纠错红利：掩码扩散由于不可逆转，在强制自纠错下反而破坏生成连贯性（GIDD $p_u=0$ 各项指标暴跌）；而 Neural CTMC 在 $\tau_{\mathrm{sc}}=0.1$ 下不仅大幅提升语法正确率 (+12.4%)，更全面激活了语言模型的事实性与创造力。

亮点与洞察¶

理论与经典物理/随机模拟算法的完美统一：将 CTMC 的速率矩阵参数化直接对齐 Gillespie 算法中出射率与跳转方向的因式分解，既消除了代理变量拟合的间接偏差，又在路径空间推导出了严密的泊松-分类变分下界。
为纯均匀离散扩散正名：长期以来学术界普遍认为吸收态掩码过程是离散扩散的最佳选择，而本文证明均匀加噪表现不佳的主因是过去的参数化架构未能适配 CTMC 内在动力学，一旦解耦便能释放其无需单向冻结、可自我纠偏的独特优势。
极具潜力的后训练迁移空间：出射率（何时改）与跳跃分布（改成什么）天然解耦，可方便地引入强化学习、偏好对齐或条件引导（例如仅对某些敏感 token 调节出射率），为可控离散生成提供了优雅的控制接口。

局限与展望¶

模型规模与序列长度受限：目前验证的最大模型为 4.24 亿参数、序列长度为 512，尚未在十亿至百亿级参数大语言模型及长上下文任务上进行验证。
前向噪声形式探索：虽然理论框架自然支持混合前向过程，但本文主要实验聚焦于纯均匀噪声，尚未系统探索特定结构化转移矩阵与解耦反向头的协同效应。
缺乏强化学习后训练探索：当前工作主要立足于预训练生成与轻量无监督纠错，如何利用解耦动力学进行 RLHF、DPO 或推理强化后训练仍有待未来工作深入挖掘。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 首次从 CTMC 本质的跳跃时刻与跳跃方向分解出发重构离散扩散参数化与变分下界，理论干净深刻。
实验充分度: ⭐⭐⭐⭐☆ 覆盖语言模型 TinyStories、OpenWebText、MNIST 图像、Scaling Law 及自纠错评测，对比详实；若能扩展至 1B+ 参数更佳。
写作质量: ⭐⭐⭐⭐⭐ 数学推导严密自洽，定理与引理层层递进，与物理模拟算法对应清晰。
价值: ⭐⭐⭐⭐⭐ 打破了掩码离散扩散的长期垄断，为离散生成模型的理论与架构设计提供了全新的里程碑式范式。