Search-Augmented Masked Diffusion Models for Constrained Generation¶

会议: NeurIPS 2026
arXiv: 2602.02727
领域: 图像生成
关键词: 离散扩散模型, 约束生成, 神经符号搜索, 局部搜索精化, 掩码扩散

一句话总结¶

本文提出 SearchDiff，一种无需训练的神经符号推理框架，通过在掩码离散扩散的反向去噪步中嵌入候选搜索采样（CSS）与局部邻域精化，将非可微黑盒约束评估直接注入反向转移核，实现高保真度且严格满足约束的离散序列生成。

研究背景与动机¶

离散扩散模型（Discrete Diffusion Models）作为自回归生成的重要替代范式，通过在离散状态空间中迭代去噪损坏的序列，打破了传统左向右单向因果生成的顺序依赖，展现了卓越的长程依赖捕获能力与全局结构建模优势。这一全局视角在小分子设计、生物大分子（多肽与 tRNA）合成以及符号逻辑推理等科学计算任务中极具吸引力，因为分子合成可行性、生物折叠构象或布尔命题的真值往往取决于整条序列的全局互作，无法在生成前缀的局部截断处被准确判定。

然而，标准离散扩散模型的核心训练目标是最大化训练数据的对数似然，本质上是在拟合无约束的数据经验分布，并不具备在推理阶段强制施加硬性约束或优化非可微物理属性的内在机制。现有的受控扩散方案主要依赖分类器引导（Classifier-Based Guidance, CBG）或无分类器引导（Classifier-Free Guidance, CFG），通常要求约束函数可微或需要显式微调辅助网络；投影扩散方法（Projected Diffusion）虽然尝试将预测分布投影到可行集，但强依赖平滑梯度信号并易破坏原始扩散轨迹。更重要的是，在实际药物化学或生物计算管线中，核心物理化学评估器（如 RDKit 计算合成可达性 SA 与类药性 QED、ViennaRNA 计算二级结构折叠能）均为离散、黑盒且高度不可微的模拟器，使得传统基于梯度的引导机制彻底失效。

为了打破“梯度依赖”与“模型再训练”的固有局限，本文发现掩码扩散（Masked Diffusion）独特的双向并行去噪机制为序列级的全局结构修复提供了天然接口。本文的切入角度是：不再强求将黑盒约束硬编码进扩散模型的连续权重或梯度反向传播中，而是在反向去噪的每一步中，直接利用去噪器生成的全局干净序列预测作为高概率先验，在其离散邻域内利用非可微黑盒违规函数进行前向启发式局部搜索，进而反向修正当前步的转移核。核心 idea：将离散启发式搜索无缝交织进掩码扩散的反向去噪步中，通过候选搜索采样与局部邻域编辑在无梯度黑盒评估下生成低违规候选，并据此构造指导性反向转移核驱动采样向可行流形收敛。

方法详解¶

整体框架¶

SearchDiff 的核心工作流程建立在掩码扩散语言模型（MDLM）的反向去噪链之上。反向过程从全掩码序列 $\bm{x}_T = (\mathtt{M}, \dots, \mathtt{M})$ 开始，在时间步 $t = T, \dots, 1$ 逐步推进。在每个反向去噪步 $t$，系统由三阶段闭环构成：首先，神经网络去噪器 $\bm{x}_\theta(\bm{x}_t, t)$ 输入当前带掩码状态 $\bm{x}_t$，输出对未损坏干净序列的全局分类概率预测；随后，搜索算子将该预测作为先验候选分布，依次执行候选搜索采样（Candidate Search Sampling, CSS）和局部搜索精化（Local Search Refinement），借助黑盒违规函数 $V(\cdot)$ 寻找既贴合模型先验又极小化约束违规度的离散候选序列 $\bar{\bm{x}}_t$；最后，修改后的反向转移核（Modified Reverse Kernel）依据 $\bar{\bm{x}}_t$ 保持已去噪 token，并在未去噪位置依据掩码调度概率性地决定是否解掩码并固化为搜索得到的 token。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    In["当前潜变量状态 x_t 与时间步 t"] --> A["去噪器全局预测<br/>神经网络输出干净序列概率分布"]
    A --> B["候选搜索采样 CSS<br/>并行采样 M 个候选并按违规度粗筛"]
    B --> C["局部搜索精化<br/>Hamming-1 邻域贪心编辑极小化违规"]
    C --> D["修改后反向转移核<br/>依据精化候选执行随机解掩码"]
    D --> Out["下一阶段潜变量状态 x_{t-1}"]

关键设计¶

1. 候选搜索采样（CSS）：基于去噪器先验的低违规初始化 针对标准随机解码在复杂约束流形上面临初始搜索点随机、违规严重的问题，CSS 充分利用了神经网络去噪器在时间步 $t$ 的后验预测能力。去噪器 $\bm{x}_\theta(\bm{x}_t, t) \in \Delta^{|\mathcal{V}| \times L}$ 输出了全序列每个位置关于干净 token 的分类概率分布。在反向步 $t$，已去噪的位置（即 $x_t^i \neq \mathtt{M}$）保持锁定，模型仅对当前仍处于掩码状态的坐标从该分类分布中独立采样，构成完整的离散候选。给定采样预算 $M$，CSS 独立抽取 $M$ 个离散候选序列 $\bm{x}_t^{(1)}, \dots, \bm{x}_t^{(M)}$，并调用非可微约束违规评价函数 $V(\cdot)$ 评估违规得分，选取其中违规度最低的序列作为精化阶段的初始锚点： $$ \bm{x}t^{\mathrm{css}} \in \arg\min). $$ 这一设计使得搜索算法无须从完全盲目的先验出发，而是站在预训练扩散模型学到的高置信度数据流形上进行初筛，在保证多样性与保真度的同时，大幅缩短了后续局部优化到达可行域所需的步数。} V(\bm{x}_t^{(j)

2. 局部搜索精化：非可微约束下的 Hamming 邻域贪心编辑 尽管 CSS 提供了优质起点，但采样所得序列仍可能存在微小的局部结构缺陷（例如分子出现未闭合环、RNA 局部碱基配对缺失或数独行冲突）。针对传统基于梯度的引导无法直接处理离散非平滑目标的问题，SearchDiff 在每个去噪步 $t$ 显式构建了步长条件化的搜索问题 $\mathcal{S}_t = \langle \mathcal{X}_t, \mathcal{A}_t, \mathrm{Succ}_t, V \rangle$。其动作空间 $\mathcal{A}_t(\bm{x})$ 仅作用于当前时间步允许编辑的坐标，执行单 token 替换操作。定义当前候选序列的步长条件化 Hamming-1 邻域为： $$ \mathcal{N}_t(\bm{x}) := \left{ \mathrm{Succ}_t(\bm{x}, (i, v)) : (i, v) \in \mathcal{A}_t(\bm{x}) \right}. $$ 精化过程从 $\bm{x}^{(0)} = \bm{x}_t^{\mathrm{css}}$ 出发，执行最多 $R$ 轮贪心邻域探索。在每一轮 $r$，算法枚举或评估邻域内所有编辑变体，选取全局违规函数 $V(\bm{x}) = \sum_k \lambda_k \nu_k(\bm{x})$ 下降最多的候选。若存在严格改进（即 $\min_{\tilde{\bm{x}} \in \mathcal{N}_t(\bm{x}^{(r)})} V(\tilde{\bm{x}}) < V(\bm{x}^{(r)})$），则更新候选状态，否则提前触发终止并输出精化结果 $\bar{\bm{x}}_t = \bm{x}^{(r)}$。该机制完全摆脱了对目标函数解析梯度的依赖，可无缝对接任何黑盒物理模拟器或 SAT 求解器。

3. 搜索引导的反向转移核：将离散优化注入扩散轨迹 在生成过程中，如果仅在采样完成后对最终样本做后处理（Post-hoc Repair），往往会因为早期不可逆的解掩码错误导致搜索陷入局部死胡同。为使每一步搜索所获得的结构修正确实指引后续轨迹，SearchDiff 重新形式化了反向状态转移条件概率。传统掩码扩散在反向步通过加权结合全掩码 token $\mathtt{M}$ 与神经网络连续预测分布 $\bm{x}_\theta$ 来决定下一步 token。SearchDiff 将这一连续后验预测显式替换为经由 CSS 与局部精化得到的确定性离散候选 $\bar{\bm{x}}_t$ 所定义的单点分布（One-hot distribution）$\delta_{\bar{x}_t^i}$。反向转移核 $\bar{p}_\theta(x_{t-1}^i \mid \bar{\bm{x}}_t, \bm{x}_t)$ 定义为： $$ \bar{p}{\theta}(x^{i}} \mid \bar{\bm{x}{t}, \bm{x} \mathrm{Cat}(x^{i}}) = \begin{cases{t-1}; \bar{x}^{i}}), & \text{if } x^{i{t} \neq \mathtt{M}, \[6pt] \mathrm{Cat}!\left(x^{i}}; \gamma_{t} \mathtt{M} + \eta_{t} \, \delta_{\bar{x}^{i{t}}\right), & \text{if } x^{i}, \end{cases} $$ 其中系数 } = \mathtt{M$\gamma_t = \frac{1 - \alpha_{t-1}}{1 - \alpha_t}$ 与 $\eta_t = \frac{\alpha_{t-1} - \alpha_t}{1 - \alpha_t}$ 遵循原始掩码扩散时间调度。已解掩码的 token 保持确定性继承，而仍处于掩码的位置则以概率 $\eta_t$ 固化为搜索所得的高质量离散 token $\bar{x}_t^i$，以概率 $\gamma_t$ 维持掩码留待后续步精化。这一闭环构造保证了中间步搜索积累的约束满足信息能够有机地传递到下一个潜在状态，实现了轨迹级的约束引导。

实验关键数据¶

主实验¶

论文在小分子设计、多肽设计、tRNA 设计以及符号逻辑推理（数独与 Boolean 3-SAT）五大基准任务上进行了严格评测，基线涵盖标准无约束扩散（MDLM）、各类引导扩散方法（CBG、CFG、CDD、TreeG-SC）以及受控自回归模型（AR FUDGE、PPLM、CFG）。

任务 / 数据集	评测指标	SearchDiff (本文)	次优基线	提升幅度 / 对比分析
小分子 (QM9, QED 最大化)	合法分子数 (Admissible Mol) ↑	497.6 ± 4.5	397.8 ± 12.9 (MDLM)	较基线提升 25.1%，较次优引导 AR PPLM (142.0) 提升 3.5×
小分子 (QM9, QED 最大化)	平均类药性 (Mean QED) ↑	0.77 ± 0.00	0.62 ± 0.00 (UDLM CFG)	显著高于所有扩散与自回归引导基线 (+0.15)
小分子 (SA ≤ 4 且 QED 最大化)	严格达标分子数 (SA ≤ 4) ↑	470.2 ± 7.8	150.2 ± 17.0 (TreeG-SC)	相比 TreeG-SC 提升 3.13×，较 MDLM (119.2) 提升近 4×
多肽设计 (GRAMPA, 3重约束)	全约束满足数 (All constraints) ↑	1024.0 ± 0.0	867.2 ± 4.5 (TreeG-SC)	达到 100% 完美满足（基线 MDLM 仅 530.8）
tRNA 设计 (GtRNAdb, 3重拓扑约束)	全结构满足数 (All constraints) ↑	719.0 ± 105.1	25.8 ± 1.4 (TreeG-SC)	相比最佳基线取得超 27× 的飞跃式提升
符号推理 (Sudoku 9×9)	解题准确率 (Accuracy %) ↑	96.5%	93.2% (MDM 6M)	较无约束扩散提升 3.3%，大幅击败同规模自回归 AR (31.25%)
符号推理 (Boolean 3-SAT, n=7)	求解准确率 (Accuracy %) ↑	76.0%	9.6% (MDM 6M)	相比无约束扩散提升近 7.9× (9.6% → 76.0%)

消融实验¶

针对搜索组件（CSS 候选数 $M$）与局部搜索精化（Local Search, LS）在不同作用时机下的贡献，以 QM9 分子生成（约束 SA ≤ 4 且最大化 QED）为基准展开消融：

局部搜索配置 (LS)	CSS 候选数 ($M$)	满足 SA ≤ 4 分子数 ↑	平均 QED ↑	单分子耗时 (s)	架构消融说明
完全不激活 (No LS)	1	110.0	0.48	0.67 (MDLM基线)	纯随机解码，缺乏约束修复能力
完全不激活 (No LS)	32	143.2	0.62	1.88	仅靠增加采样预算对硬约束改进有限
完全不激活 (No LS)	128	145.0	0.65	6.00	算力边际递减，仍无法突破流形分布壁垒
仅在最后一步运行至收敛 (Last step)	1	419.2	0.66	-	事后单步修复易受早期错误不可逆影响
仅在最后一步运行至收敛 (Last step)	32	456.2	0.72	-	相比全过程去噪仍有明显性能差距
仅在最后一步运行至收敛 (Last step)	128	500.2	0.73	-	虽有提升但计算开销集中于最终步
全过程所有步激活 (All steps, 本文)	1	433.8	0.67	3.66	仅配置单候选 + 局部搜索即超过 128 候选无 LS
全过程所有步激活 (All steps, 本文)	32	470.2	0.73	10.68	精度与推理效率的最佳平衡点（默认配置）
全过程所有步激活 (All steps, 本文)	128	514.2	0.75	-	极限性能配置，产生最多高质量兼顾分子

关键发现¶

局部搜索是性能飞跃的核心驱动力：消融数据表明，仅靠扩充 CSS 候选集（即使规模拉至 128）在 SA 达标数上仅能从 110 提升至 145；而引入单候选的局部搜索（$M=1$ + All steps LS）即可瞬间将达标数推高至 433.8。这证明了离散编辑算子能高效翻越连续采样难以跨越的局部能量势垒。
轨迹级渐进修复压倒事后补救：在 Sudoku、Boolean SAT 以及分子生成中，将搜索限制在“最后一步”（Last step）的表现普遍弱于全流程搜索（如 SAT 准确率全步搜索达 76.0%，而最后一步仅有 11.4%）。由于扩散模型在去噪前期的决策会锁定全局拓扑，全过程介入能从早期阻断错误路径的扩散。
跨越“分布陷阱”，发掘强分布外（OOD）高价值解：在极端约束（SA ≤ 3 且 QED ≥ 0.6）测试中，常规基准只能落入训练集高频分布陷阱（满足分子数仅 1~2 个），而 SearchDiff 激增至 560.6 个，且生成的分子类药性最高值甚至打破了 QM9 训练集上限，验证了推理期搜索具备强大的向低密度高价值区域拓展的能力。

亮点与洞察¶

训练无关与非可微友好：SearchDiff 无须为每项新增约束额外训练判别网络，也无须设计可微替代目标，任何能够提供违规评分的第三方代码库（如 RDKit、Biopython、ViennaRNA、Z3）均可即插即用。
神经概率先验与符号搜索的有机解耦：扩散模型负责提供宏观语境与全局语义的高概率建议（粗粒度生成），启发式搜索负责在微观邻域中执行严谨的规则校验与修复（细粒度纠错），各展所长。
推理解码算力可扩展性（Test-Time Scaling）：实验展示了随着 CSS 采样数与局部搜索深度的增加，生成质量呈现清晰的单调递增规律，为在不改变模型架构的前提下以“推理换质量”提供了可行路径。

局限与展望¶

推理时间延迟增加：相较于纯无约束去噪（QM9 上约 0.67 秒/分子），SearchDiff 在默认配置下（$M=32$, 全步搜索）耗时约为 10.68 秒，当约束函数涉及昂贵的生化折叠模拟（如 tRNA）时延迟会进一步累加。
邻域评估的算力敏感性：当前局部精化主要依赖 Hamming-1 邻域的完全或贪心枚举，面对长序列大词表时邻域规模激增，亟需引入基于梯度的重要度剪枝或分层邻域采样以降低评测开销。
长程耦合约束的局部最优困境：在布尔 SAT 等极度崎岖的组合优化曲面上，简单的贪心下降仍有陷入局部极值的风险，未来可结合模拟退火（Simulated Annealing）或 MCTS 进一步增强跳脱能力。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 巧妙结合掩码扩散模型与离散局部搜索，首创了以无梯度方式将黑盒约束融入反向转移核的理论与算子。
实验充分度: ⭐⭐⭐⭐⭐ 涵盖分子、多肽、tRNA 3项生命科学逆向设计与数独、3-SAT 2项符号推理任务，消融与对比极度详实。
写作质量: ⭐⭐⭐⭐⭐ 概念清晰，形式化定义（搜索问题元组、修改后反向转移核）严密，图文指引直观。
价值: ⭐⭐⭐⭐⭐ 为科学生成、AI 药物与分子设计、离散约束控制等高壁垒方向提供了一种通用、实用且高性能的基石级推理范式。