Constrained Decoding for Diffusion Language Models via Efficient Inference over Finite Automata¶
会议: NeurIPS 2026
arXiv: 2607.07026
领域: LLM 其他
关键词: 离散扩散语言模型, 约束解码, 有穷自动机, 概率图模型, 对数深度树采样
一句话总结¶
针对离散扩散语言模型多位置并行去噪导致传统自回归约束解码失效的瓶颈，本文将有穷自动机视作链式概率图模型，提出了首个精确可解且支持对数并行深度的有穷自动机约束解码算法，实现 100% 语法合规并在下游任务中取得显著精度跃升。
研究背景与动机¶
大语言模型在工具调用（JSON 结构化函数调用）、数据库查询（SQL 语法与 Schema 约束）以及规划推理等真实场景部署时，严格遵循语法或格式约束是系统可靠运转的生命线。在自回归生成（Autoregressive LLMs）体系下，约束解码技术（如 Outlines、XGrammar、Guidance 等）已十分成熟：解码沿从左至右单向推进，在每个时间步只需根据已生成的前缀状态，将不合法的候选词概率掩码（mask out）置零即可保证局部有效性。
然而，近年来作为自回归强劲竞争对手的离散扩散语言模型（Diffusion Language Models，如 LLaDA、Dream、Mercury 等）彻底打破了这一“单向自回归因果前缀”的假设。扩散语言模型在每个去噪时间步通过全分解的均值场分布（mean-field distribution）对全序列中多个未决位置同时独立采样。由于绝大多数结构化约束在全局跨度上产生强耦合，在每个位置上独立做局部合法性掩码根本无法保证整体合规——各个位置局部合法的 token 组合起来依然可能产生完全崩溃的非法序列（例如在浮点数约束下，第一位选小数点、第二位也选小数点，组合成非法的两个点）。此前针对扩散模型的尝试要么退化为效率极低的拒绝采样，要么依赖近似的启发式检验，无法在有限计算步数内保证 100% 的严格约束满足。
解决扩散模型约束解码的核心矛盾在于：模型内部去噪预测是完全因子分解的无状态独立分布，而外部约束则是具有全局长程依赖的硬性判别器；若想在每一步扩散采样中精确求解受限后验，必须找到一种既能保留约束支撑集、又能高效与均值场相乘并完成归一化采样的可解概率结构。本文的核心 idea 是：将有穷自动机（FA）等价建模为链式图模型 / 隐马尔可夫模型，利用算术线路深度压缩思想将 \(O(L)\) 的后向归一化与前向采样重构为 \(O(\log L)\) 的二叉前缀-后缀状态树，实现每一步去噪对受约束均值场后验的精确、快速且保证 100% 合规的并行采样。
方法详解¶
整体框架¶
本文提出的约束解码算法在扩散语言模型的每一个逆向去噪时间步 \(t \to s\) 介入执行。扩散模型原本的均值场前向预测 \(p_\theta(x^0 \mid x^t) = \prod_{i=1}^L p_\theta(x_i^0 \mid x^t)\) 给出全序列各位置的词表概率分布，而待满足的结构语法被编译为（非确定或确定型）有穷自动机 \(\mathcal{M} = (\mathcal{S}, \mathcal{V}, \delta, s_0, \mathcal{F})\)。算法将该自动机展开为序列长度 \(L\) 的隐马尔可夫链式图模型，状态转移矩阵编码合法语法流转，发射分布通过均值场概率实施逐位置重加权。在去噪步中，算法先通过自底向上的跨步转移乘法计算归一化因子，随后自顶向下分治采样隐状态序列，最后全并行发射出完全合规的序列 \(\hat{x}^0\)，并基于约束边际概率计算置信度引导重掩码。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["扩散噪声状态 x^t 与目标约束 FA"] --> B["均值场发射重加权<br/>构建受限链式图模型 g(x^0)"]
B --> C["对数深度树采样<br/>自底向上聚合 2^k 跨步矩阵与自顶向下分治"]
C --> D["约束边际概率重掩码<br/>按边际置信度选取确认子集并迭代至 x^0"]关键设计¶
1. 均值场发射重加权：构建精确受限图模型 为了在每个扩散步中采样完全满足约束的序列 \(\hat{x}^0 \sim p_\theta(x^0 \mid x^t, \mathcal{C}) \propto p_\theta(x^0 \mid x^t) \cdot \mathbb{1}[x^0 \in \mathcal{C}]\)，算法将有穷自动机诱导为长度为 \(L\) 的链式概率图模型 \(p_\mathcal{M}(x_{1:L}^0)\)。其隐变量取自自动机状态转移或边集，合法路径的支撑集恰好覆盖约束语言 \(\mathcal{L}(\mathcal{M})\)。由于扩散模型均值场先验 \(p_\theta(x^0 \mid x^t)\) 具有完全位置因子分解的特性，它与链式图模型 \(p_\mathcal{M}\) 的乘积仍然维持严格的链式图模型结构，无需引入额外的交叉高维因子。具体机制上，算法保持自动机的状态转移矩阵完全不变，仅将各时间步由状态转移边 \(z_i=e\) 发射到 token \(x_i=v\) 的未归一化发射因子乘上均值场预测值： $\(g(x_i = v \mid z_i = e) = p_\mathcal{M}(x_i = v \mid z_i = e) \cdot p_\theta(x_i = v \mid x^t)\)$ 通过在状态空间聚合边集合，每个位置的状态转移矩阵简化为 \(M_t(s, s') = \sum_{v \in \mathcal{V}} g(s_t = s', x_t = v \mid s_{t-1} = s)\)。该重加权过程将模型预测与硬性约束完美耦合，且构造所得的未归一化联合分布在整个长度为 \(L\) 的接受路径空间上支持精确全局归一化，从原理上根除了局部独立掩码导致的冲突错误。
2. 对数深度树采样：从线性链式推断到并行分治 在传统的链式隐马尔可夫模型中，计算归一化后向消息 \(\beta_t(s)\) 及前向祖先采样需要沿着序列执行 \(L\) 步串行迭代。由于去噪过程本身需要进行数十乃至数百步反向扩散，若每次去噪内部再串行迭代 \(L\) 步，串行深度将在 GPU 上造成严重的内核启动与同步开销。受算术线路深度压缩理论启发，本文利用马尔可夫链的条件独立性：给定区间左右端点状态 \(z_\ell\) 与 \(z_r\)，以中点状态 \(z_m\)（通常 \(m = \lfloor (\ell+r)/2 \rfloor\)）为条件时，左半段 \(x_{\ell:m-1}\) 与右半段 \(x_{m:r-1}\) 严格条件独立： $\(p(z_m \mid z_\ell, z_r) \propto p(z_m \mid z_\ell) \cdot p(z_r \mid z_m)\)$ 据此，算法划分为两个并行阶段：在自底向上的预计算阶段，以二分倍增方式利用矩阵乘法在 \(\log_2 L\) 层内并行计算所有跨越 \(2^k\) 步的状态转移概率矩阵；在自顶向下的采样阶段，先从联合边际中采样边界状态 \((z_1, z_{L+1})\)，随后递归地对每层中点状态 \(z_m\) 进行类别采样。中点固定后，左右两半子问题在 GPU 上完全并发递归求解，最终在基准层并行采样所有位置的 token。该设计将推断深度由 \(O(L)\) 锐减至 \(O(\log L)\)，将时间复杂度转化为可在 GPU 张量核心上极高并发度运行的稠密矩阵运算。
3. 约束边际概率重掩码：动态对齐解码轨迹 离散扩散模型在每个时间步必须决定将哪些位置的新预测 token 固定（unmask），哪些位置重新遮盖（re-mask）。常规扩散模型采用无约束的原始均值场概率 \(p_\theta(x_i^0 \mid x^t)\) 作为置信度评分（例如 Dream 选最小熵，LLaDA 选最大概率）。然而当原始均值场分布偏离合法约束空间时，无约束的高置信度位置往往诱导解码走向无法挽回的死胡同。本文利用图模型的前向-后向或者二叉树消息传递副产物，直接推导出每个位置在约束后验下的精确边际概率分布 \(p_\theta(x_i^0 = \hat{x}_i \mid x^t, \mathcal{C})\)。以此作为置信度指标（Mar），让扩散模型在决定保留 token 时能够统筹兼顾“模型自身偏好”与“全局约束合规性”的双重打分，从根本上校准去噪顺序与生成轨迹。
实验关键数据¶
主实验¶
论文在 Dream-7B 与 LLaDA-8B 两大代表性扩散模型上全面评估，涵盖函数调用（xLAM、BFCL Live/Non-Live，包含 JSON 与 Python 两种格式）、规划任务（Sudoku、Countdown）以及代码与数学推理（Spider Text-to-SQL、GSM-Symbolic）。对比基线为无约束模型（Base），评估贪心解码（Greedy, T=0）与随机采样（Sampling, T=1）两种模式。下表汇总了核心任务在 Dream-7B 上的表现，以及模型原生的约束满足率（Constraint Satisfaction, CS）：
| 任务 / 数据集 | 解码方式 | 无约束满足率 (CS) | 无约束基线 (Base) | 本文方法 (Ours) | 绝对提升 (Gain) |
|---|---|---|---|---|---|
| BFCL-Live (JSON) | Greedy | 86.7% | 63.9% | 71.5% | +7.6% |
| BFCL-Live (JSON) | Sampling | 31.2% | 22.3% | 69.0% | +46.7% |
| BFCL-Live (Python) | Greedy | 36.0% | 22.4% | 69.7% | +47.3% |
| BFCL-Live (Python) | Sampling | 4.4% | 2.8% | 67.8% | +65.0% |
| Sudoku (4x4) | Greedy | 38.0% | 36.4% | 92.4% | +56.0% |
| Sudoku (4x4) | Sampling | 7.6% | 4.7% | 91.9% | +87.2% |
| Countdown | Greedy | 53.6% | 53.6% | 100.0% | +46.4% |
| Countdown | Sampling | 35.2% | 35.2% | 100.0% | +64.8% |
| Spider Text-to-SQL | Greedy | 72.2% | 52.7% | 53.0% | +0.3% |
| Spider Text-to-SQL | Sampling | 19.3% | 15.5% | 52.1% | +36.6% |
注：本文方法（Ours）在所有任务与模式下的约束满足率均为 100.0%。在 BFCL 上与自回归模型对比中，Dream-7B + 本文方法在 Non-Live (Python) 取得 78.3%，逼近 Llama-3.1-8B 的 84.0%；Live (Python) 取得 69.7%，与 Llama-3.1-8B 的 70.8% 处于同等水平。
消融实验¶
论文针对重掩码置信度来源（均值场置信度 Mf vs. 受约束边际概率置信度 Mar）进行了消融对比，并评估了链式串行推断与对数深度树采样的计算延迟及吞吐量（在 A6000 GPU 上针对 200 个 BFCL 样本，Batch Size=1 测延迟，不同 Batch Size 测吞吐）：
| 配置 / 变体 | BFCL-L (Py) 精度 | 单样本延迟 (s) | 相对开销 (Overhead) | 吞吐量 (tok/s, bs=1) | 吞吐量 (tok/s, bs=4) |
|---|---|---|---|---|---|
| 无约束基线 (Unconstrained) | 22.4% | 24.26 ± 0.05 | 基准 (0%) | 10.55 ± 0.02 | 11.17 ± 0.05 |
| 本文链式变体 (Ours chain) | 78.5% | 51.92 ± 0.14 | +114% | 4.93 ± 0.01 | 8.62 ± 0.05 |
| 本文树采样变体 (Ours tree) | 79.2% | 25.35 ± 0.04 | +4% | 10.10 ± 0.02 | 10.44 ± 0.03 |
| 仅过滤未改置信度 (Mf 消融) | 61.5% | — | — | — | — |
| 边际概率引导重掩码 (Mar 完整) | 69.7% | — | — | — | — |
关键发现¶
- 随机采样崩塌的本质是约束失范：无约束扩散模型在温度 \(T=1\) 的随机采样下精度大幅崩塌（如 Dream BFCL-Live Python 从 22.4% 跌至 2.8%），主因是无约束采样导致合规率急剧下跌至 4.4%。本文方法在保证 100% 结构合规的前提下，使随机采样恢复到 67.8% 的高水平，彻底消除了扩散模型随机采样的脆弱性。
- 格式鸿沟是结构缺陷而非语义缺陷：基线模型在从 JSON 切换到更严格的 Python 括号与命名参数格式时性能腰斩（Dream 从 63.9% 跌至 22.4%），但引入自动机约束后 Python 格式精度恢复至 69.7%，证明扩散模型具备充分的语义理解能力，只是缺乏结构生成的局部协调机制。
- 对数深度树采样以微小开销破除延迟瓶颈：串行链式推断虽提升了精度，但因 GPU 串行 Kernel 启动和同步等待导致延迟增加 114%；而树采样算法利用矩阵乘法替代迭代，将额外时延压缩至仅 4%，吞吐量基本保持无损。
- 小步数去噪鲁棒性显著增强：当扩散去噪步数从 256 压缩至 16 时，基线模型精度由于并行度提升产生大量冲突而急剧恶化，而本文方法能始终维持合法语法拓扑，表现出极强的步数压缩鲁棒性。
亮点与洞察¶
- 将图模型理论引入扩散离散推断：敏锐地捕捉到均值场先验的因子分解特性与链式图模型乘积依然保持链式结构这一数学优雅性，使得连续全局约束在离散扩散步内可以被闭式精确求解。
- 算术线路深度压缩思想的落地：没有止步于传统动态规划或前向后向算法，而是利用中点条件独立性将 \(O(L)\) 序列推断改造为 \(O(\log L)\) 二叉树并行分治，完美契合了现代张量处理器的架构特性。
- 边际概率重掩码指导扩散演化：指出约束解码不能仅充当最后的“过滤器”，其后验边际概率应该反哺并主导扩散去噪轨迹的演进，这一洞察对未来受控扩散生成具有普适指导意义。
局限与展望¶
- 形式化表达能力受限于正则语言：当前推断算法严格建立在有穷自动机（DFA / NFA）基础之上，无法原生支持包含任意深度嵌套括号的上下文无关文法（CFG）或富属性语法。
- 复杂自动机的状态空间膨胀：对于像 Spider 这样包含大量数据表名和列名跨度的复杂文法，确定型自动机状态数呈组合爆炸，必须回退至包含约 10 万条边的大规模 NFA，对内存带宽构成一定压力。
- 未来方向：探索高阶图模型或有界下推自动机（bounded pushdown automata）向扩散模型的延伸，以及将该算法与推测解码（speculative decoding）相结合以进一步提升极端并发下的效率。
相关工作与启发¶
- vs DINGO (Suresh et al., NeurIPS 2025): DINGO 是首个探索扩散模型有穷自动机约束的工作，但其仅局限于块状贪心解码和确定型自动机的动态规划 MAP 推断；本文是严格的全面推广，原生支持 NFA、随机后验精确采样、任意去噪调度与对数级树推断。
- vs Lookahead-then-Verify (Zhang et al., 2026) & Mündler et al. (ICLR 2026): 这类工作尝试通过拒绝采样与前瞻检验来处理上下文无关文法（CFG），但由于扩散模型单步生成空间极大，在有限计算预算下接受率极低且无法提供 100% 合规保证；本文在自动机范畴内实现了数学上保证满足且确定性终止的精确推断。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ 首次建立离散扩散模型均值场与有穷自动机图模型的精确乘积表征，并提出对数深度树采样算法。
- 实验充分度: ⭐⭐⭐⭐⭐ 覆盖两大主流扩散大模型、五大代表性任务、六个主流基准，详细对比了贪心与随机采样以及延迟吞吐开销。
- 写作质量: ⭐⭐⭐⭐⭐ 问题定位敏锐，图模型建模与算法推导严密自洽，理论复杂度与工程实现权衡清晰。
- 价值: ⭐⭐⭐⭐⭐ 为离散扩散语言模型填补了工业级部署不可或缺的高效可靠结构化约束解码空白。