Primal-Dual Guided Decoding for Constrained Discrete Diffusion¶

会议: NeurIPS 2026
arXiv: 2605.09749
OpenReview: https://openreview.net/forum?id=RdcJpbmm7F
领域: 其他
关键词: 离散扩散模型、约束生成、原始对偶优化、镜像下降、解码引导

一句话总结¶

针对离散掩码扩散模型缺乏全局属性约束机制且现有方法计算开销大或依赖额外分类器的问题，本文提出基于镜像下降的原始对偶引导解码算法（SPDD），在去噪推断时自适应调整 token logits，以极低开销实现多目标约束满足并保持极佳分布保真度。

研究背景与动机¶

离散扩散模型（如掩码扩散语言模型 MDLM、SEDD）近年来作为自回归模型的重要替代范式迅速兴起，其通过对完全掩码序列逐步解掩码并利用全双向上下文生成离散序列，已广泛应用于自然语言处理、分子 SMILES 生成、音乐歌单推荐以及蛋白质设计等多个领域。然而在实际落地场景中，生成输出往往需要严格满足某些全局属性约束，例如文本生成中指定主题词频、药物分子设计中满足目标分子量或特定官能团数量、歌单生成中维持目标子流派占比。这种对全局属性的强制要求与保持基础模型学到的数据分布之间，存在着根本性的张力。

现有离散扩散引导方法存在不可忽视的瓶颈。以分类器引导（D-CBG）或预测器引导为代表的方案，必须为每个属性单独训练额外的分类器网络，且在面对稀疏关键词约束时梯度平滑往往导致引导失效；而无需重新训练的离散约束扩散（CDD）等方法，则依赖于在每个去噪步内执行内层 Gumbel-softmax 优化循环，带来 5 至 20 倍的巨大推断延迟开销，难以支持多目标联合约束。

本文切入角度是将带约束离散扩散推断建模为带期望约束的 KL 散度正则化优化问题，并在去噪过程中利用随机原始对偶近似技术进行在线求解。核心 idea：将离散去噪步中的约束满足转化为 KL 正则化投影问题，通过镜像下降在线自适应更新拉格朗日乘子，推导出闭式加性 Logit 偏置，无需辅助模型与内层循环即可达成帕累托最优约束控制。

方法详解¶

整体框架¶

在离散掩码扩散（MDLM）的反向去噪过程中，序列从完全掩码状态 \(x_T = [\texttt{M}]^L\) 开始逐步向无掩码样本 \(x_0\) 演化。假设每个 token \(j\) 被放置在位置 \(\ell\) 时会贡献非负属性值 \(b_{\ell j} \ge 0\)，目标是使得序列期望总贡献满足 \(\sum_{\ell=1}^L \sum_{j=1}^V b_{\ell j} r(x_0^\ell = j) \ge R\)。为了在满足约束的同时最小化对原始模型分布 \(p_\theta(x_{0:T})\) 的干扰，算法通过一阶最优条件推导出针对掩码位置的加性对偶偏置 \(\lambda_{t+1} b_{\ell j}\)，并在每个去噪步根据当前累积的约束满足亏损动态更新对偶乘子 \(\lambda_t\)。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["完全掩码输入 x_T 与初始对偶乘子 λ_T"] --> B["阶段 1：Logit 加性对偶偏置注入<br/>基于当前 λ 调整未掩码 token 采样分布"]
    B --> C["阶段 2：离散类别分布采样<br/>更新部分掩码位置并固定解掩码 token"]
    C --> D["阶段 3：约束松弛量评估与镜像下降更新<br/>计算累积赤字并自适应调整 λ_t"]
    D -->|未达到 t=0| B
    D -->|完成 T 步去噪| E["满足全局约束的高保真离散序列 x_0"]

关键设计¶

1. KL 正则化投影与闭式 Logit 偏置：以能量场形式无缝引导反向去噪

为了避免训练辅助判别器或陷入复杂的内层数值优化，本文将目标分布 \(r(x_{0:T})\) 限制在与原扩散模型相同的马尔可夫分解族中，目标是最小化轨迹级 KL 散度 \(D_{\mathrm{KL}}(r(x_{0:T}) \,||\, p_\theta(x_{0:T}))\) 并满足全局线性约束。由于掩码一旦解开在反向过程中即被吸收固定（Absorbing State），每个去噪步观测到的模型对数概率可视为常数。引入拉格朗日乘子 \(\lambda_{t+1} \ge 0\) 后，拉格朗日函数关于采样分布的一阶导数可以直接解析求得，解掩码位置的采样概率呈现为对原始概率的指数重加权形式：

\[a_{\ell j}^t = \frac{p_\theta(x_t^\ell = j \mid x_{t+1}) e^{\lambda_{t+1} b_{\ell j}}}{\sum_{i=1}^V p_\theta(x_t^\ell = i \mid x_{t+1}) e^{\lambda_{t+1} b_{\ell i}}}\]

这一闭式解在对数概率（Logit）空间中等价于注入了一个加性偏置项 \(\lambda_{t+1} b_{\ell j}\)。在能量模型视角下，该偏置充当了施加在无偏扩散模型能量曲面上的外加场，使得有利于满足约束的 token 具有更低的“自由能”，从而以最低的分布失真引导采样轨迹。

2. 基于在线镜像下降的乘子自适应演变：跟踪累积赤字实现闭环控制

静态加性偏置无法根据生成过程中的实际进展动态调节力度，极易导致前期过度约束或后期达标后的过冲破坏连贯性。本文通过负熵诱导的 Bregman 散度对内层对偶问题执行镜像下降（Mirror Descent），推导出乘子闭式递归演化机制：

\[\lambda_t = \lambda_T \exp\left\{ -\eta \sum_{\tau=t}^{T-1} \sum_{\ell=1}^L \sum_{j=1}^V \left(b_{\ell j} - \frac{R}{L}\right) a_{\ell j}^\tau \right\}\]

当实际采样出的 token 累积贡献低于按位置折算的期望目标 \(R/L\) 时，乘子 \(\lambda_t\) 迅速攀升，加大对高贡献 token 的偏置力度；一旦序列中生成的属性累积达标，乘子便迅速衰减回低位，将后续生成的主导权完全交还给原始模型的无偏分布，确保句末或后续结构的自洽性。

3. 多松弛量更新机制与乐观预测（OMD）：针对非均匀约束的时间动态补偿

将全局目标 \(R\) 均匀按步分摊到 \(R/L\)（累积松弛量，Accumulated Slack）在属性均匀分布时表现出极强的稳定性与极低的 KL 散度，但在属性高度稀疏时可能反应迟缓。为此，本文扩展了瞬时松弛量（Instantaneous Slack）以及基于乐观镜像下降（Optimistic Mirror Descent, OMD）的预测松弛量机制。OMD 在每一步估计未来未掩码位置的预期属性短缺量 \(M_\tau = R - \sum_{\ell \in \text{masked}} \sum_j p_\theta(x_\tau^\ell = j \mid x_{\tau+1}) b_{\ell j}\)：

\[\lambda_t = \lambda_T \exp\left\{ -\eta \sum_{\tau=t}^{T-1} \left[ \left(\sum_{\ell=1}^L \sum_{j=1}^V b_{\ell j} a_{\ell j}^\tau - R\right) - M_\tau + M_{\tau+1} \right] \right\}\]

该机制在去噪初期退化为标准镜像下降以防过度激进，在去噪中后期则利用模型已趋稳定的预测前瞻性修正乘子，消除了过冲现象并提升了生成连贯性。

4. 理论约束违反上界与对偶遗憾界分析：明确保真度与满足率的权衡

在假设扩散模型 logits 满足弱时间一致性（漂移均值受限、跨位置协方差受限的高斯过程）的前提下，本文给出了严谨的理论保障：算法的预期约束违反上界为 \(\mathcal{O}\left(\frac{1}{\eta} \log \frac{\eta + 1}{b_{\max} - R/L}\right)\)，说明调节步长参数 \(\eta\) 可以单调收紧约束违例；同时，算法关于最优静态分配策略的遗憾界以 \(\mathcal{O}\left(\frac{\lambda_T}{\eta} e^{\eta L T \max |b_{\ell j} - R/L|}\right)\) 增长。这在数学上证明了 \(\eta\) 是显式控制“约束满足率”与“分布保真度（KL 散度）”之间帕累托前沿的单一调节旋钮。

损失函数 / 训练策略¶

本方法完全属于推断期（Inference-time）引导技术，基础扩散模型（MDLM）无需针对约束属性进行任何微调或重训。输入只需预定义的单 token 贡献向量 \(b_{\ell j}\)（例如词表指示向量、原子量量化值或聚类流派标签比例）。当存在多个联合约束时（如同时约束分子量 \(\mathrm{MW} \ge 350\) 和杂原子数 \(\#(N+O) \ge 5\)），每个约束独立维护一组乘子 \((\lambda_i, \eta_i)\)，Logit 偏置直接线性叠加 \(\sum_i \lambda_i b_{\ell j}^{(i)}\)，原生支持多目标联合引导。

实验关键数据¶

主实验¶

论文在三个不同结构的离散模态上进行了全面验证：稀疏词汇计数的文本主题引导（TinyStories，要求包含 \(\ge 10\) 个海洋词汇）、密集数值约束的分子 SMILES 生成（ZINC20，要求分子量 \(\mathrm{MW} \ge 350\) 并保持化学有效性）、以及含噪间接聚类约束的音乐歌单补全（基于 LSH 音频聚类标签生成指定子流派）。基线包括无约束扩散、静态 Logit 偏置、分类器引导（D-CBG）和约束离散扩散（CDD）。

任务模态	评估指标	无约束基线	静态偏置 (Static)	分类器引导 (D-CBG)	约束扩散 (CDD)	SPDD (本文)
文本主题引导 (TinyStories)	Pass% (达标率) ↑	0.50%	81.90%	0.00%	38.00%	92.10%
	KL 散度 (失真度) ↓	—	0.706	0.127	0.221	0.420
分子药物生成 (ZINC20)	Pass% (\(\text{MW} \ge 350\)) ↑	29.40%	41.40%	42.80%	42.20%	48.50%
	KL 散度 (失真度) ↓	—	0.019	0.035	0.034	0.096
音乐歌单生成 (Playlist)	Pass% (子流派满足率) ↑	1.14%	1.72%	1.92%	9.18%	9.28%
	KL 散度 (失真度) ↓	—	0.003	0.115	0.779	0.355
推断相对耗时	采样耗时倍率 (Sample \(\times\))	1.00×	1.03×	1.09×	7.28×	1.03×
额外训练需求	训练成本 (Train)	无	无	需训练分类器	无	无

消融实验¶

针对文本主题生成实验（\(R=10\)，评估样本量 \(N=1000\)），下表展示了不同松弛量更新机制（Slack Mode）和更新步长 \(\eta\) 对达标率、困惑度（PPL）、2-gram 多样性（Dist-2）及 KL 散度的影响：

松弛量模式 (Slack Mode)	步长 \(\eta\)	海洋词均值 (Ocean)	达标率 (Pass%)	困惑度 (PPL) ↓	多样性 (Dist-2) ↑	KL 散度 ↓	说明
累积模式 (Accumulated)	0.1	6.5 ± 4.0	20.5%	25.5	0.185	0.281	极低扰动，温和爬升
累积模式 (Accumulated)	0.5	9.5 ± 2.6	46.9%	25.6	0.181	0.365	词数贴合目标 \(R=10\)
累积模式 (Accumulated)	1.0	10.5 ± 2.0	67.6%	25.7	0.180	0.402	约束满足率稳步提高
累积模式 (Accumulated)	2.0	11.0 ± 1.7	92.1%	25.3	0.175	0.420	兼顾高满足率与高保真
瞬时模式 (Instantaneous)	0.5	21.4 ± 4.4	100.0%	23.0	0.159	0.890	激进 bang-bang 控制，明显过冲
乐观预测模式 (OMD)	0.1	10.6 ± 5.4	53.1%	24.8	0.176	0.399	中等步长下优于瞬时与早期
乐观预测模式 (OMD)	0.5	21.3 ± 4.3	100.0%	23.2	0.158	0.898	利用预测消除过强扰动

关键发现¶

对稀疏关键词的压倒性优势：在文本主题实验中，D-CBG 依靠 Taylor 一阶泰勒展开近似分类器梯度，其梯度平滑效应无法在 8192 词表中将稀疏的 18 个目标词概率推至采样阈值以上（Pass% 为 0%）；而 SPDD 凭借闭式对偶指数加权，能精准拉高目标词得分，Pass% 突破 92.1%。
推断效率极大超越 CDD：CDD 在歌单与文本生成中虽然取得了一定约束效果，但每步多轮 Gumbel-softmax 优化使其推断时间暴增至 7.28 倍，且在歌单生成中引发严重的 token 多样性坍塌；SPDD 仅需计算标量乘法与加法，采样耗时仅为无约束基线的 1.03 倍，同时 KL 散度降低一半以上。
多约束扩展性完备：在分子双目标生成实验中（同时要求 \(\text{MW} \ge 350\) 且 \(\#(N+O) \ge 8\)），原本天然的相关性被打破，单目标与静态基线联合通过率暴跌至 17% 以下，而双乘子 SPDD 联合达标率依然维持在 34.4%，远超联合训练分类器的 D-CBG（25.1%）。

亮点与洞察¶

将在线凸优化与离散去噪动力学完美契合：论文敏锐抓住了掩码离散扩散“解掩码即固定”的单向吸收态特性，将其等价为在线资源分配问题，直接利用镜像下降得到 Logit 空间中的自适应乘子更新律，形式优雅且数学完备。
无需分类器与重训练的即插即用属性：摆脱了连续扩散模型中对可微分类器或 CFG 重训练的深度依赖，仅凭先验定义的属性打分向量即可实现跨模态（文本、分子、歌单）的强效约束，落地门槛极低。
零样本无缝迁移到超大模型：作者将 SPDD 直接应用于 8B 规模的开源离散语言扩散模型 LLaDA-8B 上，在不微调任何参数的前提下，仅通过解码期加性偏置即实现 100% 的稀疏关键词满足率，展现出强大的模型无关性。

局限与展望¶

吸收态不可逆性限制了后期纠错能力：当前算法建立在吸收态扩散之上，早期阶段一旦错误地固定了非目标 token，后续步骤便无法撤销，这使得累积松弛量更新高度依赖早期的合理决策。
对非可分全局属性的依赖较强：对于无法分解为 token 级累加的复杂全局性质（如药物分子的定量类药性 QED、文本连贯性等），目前需要构造近似的子序列或局部启发式打分，约束精度有所下降。
大模型长程重复风险：在零样本迁移至 LLaDA-8B 等大底座模型时，若缺乏温度调度或重复惩罚机制，强约束偏置容易诱发局部 token 连续重复的退化现象，需配合后处理工程技巧。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 首次将在线原始对偶镜像下降理论系统化引入离散掩码扩散解码，建立了优雅完备的理论体系。
实验充分度: ⭐⭐⭐⭐⭐ 跨越文本、分子、歌单三类异构离散模态，并包含 LLaDA-8B 迁移、双目标组合、深入理论假设验证及消融实验。
写作质量: ⭐⭐⭐⭐⭐ 逻辑严密，推导详实，算法框架与动机脉络叙述极为清晰。
价值: ⭐⭐⭐⭐⭐ 为离散扩散模型从理论走向实际可控生成提供了兼具计算轻量性与理论保障的基础性推断解码方案。