Sticky Jump Diffusions: A Unifying View of Masked, Continuous, and Hybrid Diffusion¶
会议: NeurIPS 2026
arXiv: 2607.10951
代码: https://github.com/PascalJD/sticky-jump-diffusions
领域: 其他
关键词: 粘性跳跃扩散, 掩码扩散, 混合扩散, 时间反转, 去噪风险率匹配
一句话总结¶
本文提出粘性跳跃扩散(SJD),将离散 token 嵌入视为连续空间中的粘性锚点,在统一的连续时间跳跃扩散马尔可夫过程中实现了掩码扩散、连续扩散与混合扩散的理论统一,并通过去噪风险率匹配与单网络无仿真交叉熵训练,解锁了利用跨位置混合核注入先验的全新设计空间。
研究背景与动机¶
离散序列(如文本、代码、符号逻辑)的扩散生成模型长期被两大割裂的技术范式所主导。一类是以 MDM、MDLM 为代表的掩码离散扩散,其直接在 token 空间定义转移矩阵并向吸收态掩码演化,能够天然输出合法 token,但各位置被强制解耦映射到单一吸收态,彻底抹平了连续嵌入空间中度量相近候选词的平滑几何梯度信息;另一类是以 Diffusion-LM 为代表的连续嵌入扩散,直接在实数空间运行连续 SDE 并借助评分匹配进行学习,但生成轨迹终点漂浮在离散流形之外,必须依赖外部解码器、后处理量化或启发式投影强行映射回词表,破坏了动力学内部的一致性。近期涌现的混合扩散模型(如 CADD、CANDI、CCDD)尝试通过拼接掩码离散链与连续潜变量来吸取二者之长,但其采样器中的离散提交决策、提交速率、连续漂移项以及潜变量更新规则完全由人工拼凑组合,缺乏第一性原理的系统理论推导。
这一分裂局面的核心矛盾在于:离散 token 的原子分布(点测度)与连续空间的平滑扩散流(勒贝格绝对连续测度)在动力学表述上长期缺乏严密的数学统一工具。连续扩散无法解释为何必须在边界处进行离散化投影,掩码扩散无法解释为何掩码状态对源 token 具有完全的信息虚无性,而混合扩散则无法从统一的前向扰动流中直接内生推导出精确的时间反转采样规则。
针对上述困境,本文借鉴经典一维粘性布朗运动理论并将其拓展至高维多锚点流形,提出了粘性跳跃扩散(Sticky Jump Diffusions, SJD)。核心 idea:将离散 token 嵌入建模为连续空间中的粘性锚点,前向过程按危害率脱落进入环境空间扩散,逆向过程由通量平衡唯一确定为评分 SDE 与粘性跳入核的精确耦合,从而使掩码、连续与混合扩散全部分化为 SJD 的退化极限,并利用去噪危害率匹配实现单分类器无仿真训练。
方法详解¶
整体框架¶
粘性跳跃扩散在状态空间 \(\mathsf{X} = \mathbb{R}^d\) 上定义,其包含由离散词表嵌入组成的有限锚点集 \(\mathcal{A} = \{a_k\}_{k=1}^K \subset \mathbb{R}^d\),连续区域记为 \(\mathsf{X}_{\mathcal{A}} := \mathbb{R}^d \setminus \mathcal{A}\)。在前向时间流中,粒子初始依数据分布位于某一锚点 \(X_0 = a \in \mathcal{A}\);在滞留期间,该锚点以危害率(hazard rate)\(\vec{\lambda}_t\) 释放质量,并依脱粘核(unsticking kernel)\(R_t(\cdot \mid a)\) 跃入连续区域 \(\mathsf{X}_{\mathcal{A}}\);一旦脱离锚点,粒子便遵循由漂移项 \(\vec{b}_t\) 与扩散系数 \(g_t\) 控制的标准 Itô SDE 在连续环境中自由扩散,且前向过程不再跳回锚点。
在反向生成时间流中,依据测度通量平衡原理,时间反转必然对应一个精确耦合的随机动力系统:粒子在连续区域内沿包含分数函数修正的逆向漂移项连续演化,同时以确定的逆向危害率发生向锚点的粘性跳跃(sticky jump);一旦跳入某一锚点,便在该位置被永久冻结(absorbing)。整个生成过程的提交速率、提交目标词以及环境漂移项全部由前向动力学解析导出。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["先验噪声采样<br/>y ~ p_T 在连续空间初始化"] --> B["环境连续演化<br/>评分驱动逆向 SDE 漂移与扩散"]
B --> C["单网络去噪后验估计<br/>统一分类器 P_theta 预测 token 概率"]
C --> D["去噪危害率匹配 (DHM)<br/>解析重加权计算粘性跳入率与锚点分配"]
D --> E["通量平衡粘性跳跃决策<br/>依据 hazard 判定是否跃入并吸收至离散锚点"]
E -->|未提交| B
E -->|提交冻结| F["最终离散 Token 序列输出"]
关键设计¶
1. 基于通量平衡的时间反转与极限理论统一:推导包含原子跳跃的逆向生成动力学 针对现有方法人工设定反向采样步骤的理论断层,本文基于马尔可夫跳跃扩散的狄利克雷形式与分部积分恒等式,推导了 SJD 的时间反转定理。在连续区域 \(\mathsf{X}_{\mathcal{A}}\) 上,逆向漂移项为 \(\reflectbox{\)\vec{\reflectbox{\(b\)}}\(}_t(x) = -\vec{b}_t(x) + g_t^2 \nabla \log p_t(x)\);而逆向粘入核 \(\reflectbox{\)\vec{\reflectbox{\(J\)}}\(}_t(y, \mathrm{d}x)\) 由前向边缘分布与前向跳跃核之间的通量平衡方程唯一确定: $\(p_t(\mathrm{d}y) \reflectbox{\)\vec{\reflectbox{\(J\)}}\(}_t(y, \mathrm{d}x) = p_t(\mathrm{d}x) \vec{J}_t(x, \mathrm{d}y)\)$ 对于 \(y \in \mathsf{X}_{\mathcal{A}}\) 跃迁至锚点 \(x = a \in \mathcal{A}\),逆向单锚点强度严格等于 \(\reflectbox{\)\vec{\reflectbox{\(J\)}}\(}_t(y, \{a\}) = \frac{p_t(\{a\})}{p_t(y)} \vec{\lambda}_t r_t(y \mid a)\)。该结论完美统一了三大族:当脱粘核退化为单一狄拉克点 \(r_a \equiv \delta_m\) 且 \(g_t \equiv 0\) 时,严格退化为掩码扩散(MDM),解释了为何掩码 token 抹平了距离几何;当粒子在 \(t=0^+\) 瞬间脱粘(\(S_a(0^+)=0, \vec{\lambda}_t \equiv 0\))时,前向边际完全丧失原子测度,逆向通量为 0,因而连续扩散必须依靠外部终端量化投影;而在方差保持(VP)设定下且脱粘核为独立高斯时,精确还原了 CADD 混合扩散。
2. 去噪危害率匹配(DHM)与单分类器即插即用:无须额外头部的解析重加权 针对边缘概率密度 \(p_t(y)\) 包含脱粘时间隐变量混合导致逆向危害率解析不可求的难题,本文提出了去噪危害率匹配(DHM)。DHM 证明了难以计算的逆向全危害率 \(\reflectbox{\)\vec{\reflectbox{\(\lambda_t^{\star}\)}}\(}(y) = \sum_a \Lambda_t^{\star}(a \mid y)\) 正好等于一个仅依赖前向已知量的目标量 \(\widehat{\lambda}_t(y, X_0)\) 在给定 \(X_t=y\) 下的条件期望: $\(\widehat{\lambda}_t(y, X_0) := \frac{\vec{\lambda}_t(X_0) S_{X_0}(t) r_t(y \mid X_0)}{p_t^{\mathrm{ac}}(y \mid X_0)}\)$ 由于目标 token \(X_0\) 离散,该条件期望可展开为离散求和 \(\sum_a \widehat{\lambda}_t(y, a) p_t(a \mid y)\)。因此,只需训练一个以常规交叉熵为损失的去噪分类器 \(P_\theta(a \mid y, t)\) 逼近真实后验 \(p_t(a \mid y)\),便可直接获得 Rao-Blackwell 化的即插即用估计量 \(\Lambda_{\mathrm{plug},\theta}(a \mid y, t) = \widehat{\lambda}_t(y, a) P_\theta(a \mid y, t)\)。采样时锚点的归一化选择概率直接化简为对分类器 logits 叠加解析偏置 \(\ell_\theta(a, y, t) + \log \widehat{\lambda}_t(y, a)\) 后的 softmax,同时连续评分项亦可由该分类器后验期望解析推导,完全免除了独立的评分网络或 hazard 回归网络。
3. 跨位置混合脱粘核(Cross-Position Blending):将结构先验注入扰动几何 突破现有混合扩散仅在孤立局部加噪的限制,SJD 构造了非局部脱粘核。对于长度为 \(L\) 的序列,引入预定义的混合矩阵 \(W \in \mathbb{R}^{L \times L}\),将位置 \(i\) 的扰动中心由单纯的自身嵌入 \(E(a^{(i)})\) 替换为邻近位置的加权混合均值 \(\mu_i(X_0) = (W E(X_0))_i\)。在图像生成中,\(W\) 可实现为基于欧氏像素距离的高斯值域模糊核 \(B^{(\sigma_B)}\);在文本中,\(W\) 为一维高斯局部卷积;在数独等硬约束逻辑推理中,\(W\) 严格编码同行、同列及同宫格的 20 个约束邻居。这一设计将物理拓扑与组合约束直接转化为前向破坏的归纳偏置,逆向生成时通过评分残差修正引导模型优先满足结构依赖。
损失函数 / 训练策略¶
SJD 在训练阶段享有无仿真(simulation-free)特性。在线性高斯 VP-SDE 下,脱粘时刻 \(\tau \in (0, t)\) 可通过生存函数 \(S(t) = \exp(-\int_0^t \vec{\lambda}(s) \mathrm{d}s)\) 的反函数解析逆变换直接采样;脱粘核与扩散半群的卷积解析为单一各向同性高斯分布 \((r_\tau * K_{\tau \to t})(\cdot \mid X_0, i) = \mathcal{N}(\cdot; \alpha(t) \mu_i(X_0), v_t(\tau) I)\)。训练时单步仅需一次矩阵乘法计算混合均值并抽取单次高斯噪声,模型参数仅通过标准交叉熵损失进行端到端优化: $\(\mathcal{L}(\theta) = \mathbb{E}_{X_0 \sim p_0, t \sim \mathrm{Unif}(0, T), X_t \sim p_t^{\mathrm{ac}}(\cdot \mid X_0)} \left[ -\log P_\theta(X_0 \mid X_t, t) \right]\)$ 逆向采样采用一阶预测器:每一反向步长 \(\Delta \tau\) 内,单次前向推理同时得到分类后验与诱导评分,连续坐标执行 Euler-Maruyama 连续漂移更新,同时各未冻结位置以概率 \(1 - \exp(-\lambda_{\mathrm{plug},\theta} \Delta \tau)\) 独立判定是否粘入锚点,若粘入则依重加权概率采样锚点并永久冻结。
实验关键数据¶
主实验¶
论文在 CIFAR-10 无条件图像生成、Text8 字符级文本生成以及数独(Sudoku)逻辑求解三大代表性基准上进行了系统评测。在 CIFAR-10 上(256 NFE),SJD 超越了各类掩码、连续及基线混合扩散;在数独任务上,SJD 展现出极高解题准确率且彻底消除了模式崩溃。
| 数据集 / 任务 | 评估指标 | SJD (本文) | 对应 Baseline / 之前 SOTA | 表现对比与提升 |
|---|---|---|---|---|
| CIFAR-10 (32x32) | FID (↓) | 14.57 (\(\sigma_B=2.5\)) | 15.88 (CADD, 恒等混合) / 14.66 (DDPM) | 相比 CADD 降低 1.31;优于连续与掩码扩散 |
| Text8 (128 NFE) | 有效词比例 (长≥5, ↑) | 0.212 | 0.117 (连续反向 SDE) / 0.092 (连续 PF-ODE) | 相比连续扩散 SDE 提升 1.81 倍;32 NFE 超过其 128 NFE |
| Text8 (128 NFE) | 有效词比例 (长≥6, ↑) | 0.130 | 0.055 (连续反向 SDE) / 0.039 (连续 PF-ODE) | 相比连续扩散 SDE 提升 2.36 倍,相比 PF-ODE 提升 3.33 倍 |
| Sudoku 逻辑求解 | 全盘完全匹配率 (%, ↑) | 95.65 | 47.12 (CADD) / 89.49 (MDM) / 87.18 (AR-Solver) | 相比 CADD 提升 48.53 个百分点;超越自回归与掩码模型 |
| Sudoku 逻辑求解 | 起飞训练步数 (Takeoff, ↓) | 50k | 203k (CADD) | 收敛起飞速度加快约 4 倍 |
消融实验与分析¶
针对脱粘核混合带宽 \(\sigma\)、中间粘性提交相比终端量化解码的优势、以及可学习危害率机制的消融分析如下:
| 配置 / 实验变体 | 核心指标表现 | 说明 |
|---|---|---|
| CIFAR-10: CADD 恒等核 (\(\sigma_B = 0\)) | FID = 15.88 | 缺乏空间邻域混合,生成质量最差 |
| CIFAR-10: SJD 弱模糊 (\(\sigma_B = 1.5\)) | FID ≈ 15.05 | 引入轻微邻域耦合,性能显著改善 |
| CIFAR-10: SJD 强模糊 (\(\sigma_B = 2.5\)) | FID = 14.57 | 充分利用图像局部平滑性,达到最优 |
| Text8: 连续 Reverse SDE (终端最近邻投影) | 32 NFE: 0.081; 128 NFE: 0.117 | 依赖终端投影导致语义漂移,长词生成质量差 |
| Text8: SJD 粘性跳跃 (中间动态提交) | 32 NFE: 0.180; 128 NFE: 0.212 | 动态吸收提交机制使得 32 NFE 效果超越连续模型 128 NFE |
| Sudoku: CADD 恒等核 (\(W=I, \eta=1\)) | 准确率 47.12%(66 种子双峰分布崩溃) | 缺乏图结构先验,部分种子在训练中完全无法收敛 |
| Sudoku: SJD 约束图核 (\(\sigma_W = 1.5\)) | 准确率 95.65%(零崩溃) | 将行/列/宫格硬约束转化为扰动均值,收敛稳健 |
| Text8: 联合学习危害率 (ELBO Bounded) | 有效词比例暴跌至 0.066(展开比 4x) | 优化梯度诱导频率倾斜,破坏反向采样分配 |
| Text8: 冻结均匀危害率 (固定 \(w \equiv 1\), CE) | 有效词比例保持在 0.273 ~ 0.279 | 证实固定均匀危害率是最具鲁棒性的实用默认配置 |
关键发现¶
- 混合矩阵 \(W\) 所引入的非局部脱粘核是性能增益的核心驱动力。在图像上利用空间距离高斯模糊,在数独上利用邻域约束图编码,均大幅超越了将各位置独立加噪的经典混合扩散架构。
- 动态粘性跳跃相比连续扩散的“终端量化投影”具有压倒性优势。连续扩散的终点往往偏离真实离散流形,导致最近邻量化错误频发;而 SJD 在逆向轨迹中渐进吸收,有效避免了维度灾难带来的投影失真。
- 联合端到端学习危害率(hazard schedule)在理论与实验上均证明不可行。由于在最优点处 ELBO 边界对于危害率是平坦的(退化为数据熵),可学习参数极易走向极端的频率倾斜,进而严重损害反向分配;而固定的均匀脱粘率则能保持高度稳健的优异生成表现。
亮点与洞察¶
- 第一性原理统一扩散流派:将掩码扩散的“信息虚无”与连续扩散的“终端外力投影”精妙地解释为脱粘核退化和边缘原子测度缺失的极端特例,填补了混合扩散中反向离散提交启发式规则的理论空缺。
- 无须双头的解析即插即用机制:提出去噪危害率匹配(DHM),仅凭单网络标准交叉熵训练分类器,即可通过数学恒等式解析导出逆向跳跃速率与环境评分场,无须额外的评分网络或复杂的强化学习。
- 结构约束几何化迁移范式:将离散组合逻辑(如数独图约束、分子图拓扑或语法依赖树)直接编码到加噪过程的均值混合矩阵 \(W\) 中,为图生成、定理证明与离散推理扩散模型提供了极具启发性的先验注入通用范式。
局限与展望¶
- 实验规模尚未覆盖到现代超大规模语言模型(LLM)的数十亿参数预训练,且词表规模在数万级别时高斯混合求和与核矩阵相乘的显存与计算开销仍需优化。
- 脱粘核的混合矩阵 \(W\) 当前依赖人工设计的先验拓扑结构(如高斯模糊、行列邻域),缺乏针对通用无结构长文本自动学习自适应稀疏混合拓扑的机制。
- 进一步探索更丰富的跳跃扩散过程族(例如允许前向过程重新粘回锚点、或允许锚点之间直接跳跃的双向动力学系统)将是未来理论拓展的重要方向。
相关工作与启发¶
- vs CADD / CANDI (混合扩散):现有混合扩散将离散掩码链与连续潜变量机械拼接,反向提交规则纯属人工设计;本文将二者统一为单一跳跃扩散过程,证明了 CADD 仅是其 \((W, \eta) = (I, 1)\) 下的一阶离散化特例,并提供了超越恒等加噪的通用结构核设计空间。
- vs MDM / MDLM / MD4 (掩码扩散):掩码扩散对应脱粘核完全退化为狄拉克掩码且无环境扩散的极限;SJD 保留了连续几何梯度,在保持离散输出特性的同时彻底消除了掩码状态的信息孤立问题。
- vs Continuous Diffusion / Diffusion-LM:连续扩散在反向终点缺乏原子分布,必须强行依赖最近邻或外部解码器量化;SJD 将离散词表直接作为逆向吸引子,在动力学内部实现了严格的内生离散化。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ 建立了粘性跳跃扩散与去噪危害率匹配理论,完美统一连续、掩码与混合扩散。
- 实验充分度: ⭐⭐⭐⭐☆ 覆盖图像、文本与数独求解三类不同拓扑任务,消融实验与理论推导严密自洽。
- 写作质量: ⭐⭐⭐⭐⭐ 数学证明严谨,动机阐述层层递进,概念抽象与物理直觉结合极佳。
- 价值: ⭐⭐⭐⭐⭐ 为离散扩散模型奠定了扎实的随机分析理论基石,具有广阔的扩展与借鉴前景。