NeurIPS 2026 Paper Enables Training-Free Black-Box Reward Steering of Discrete Diffusion Models

新智元报道

近两年，离散扩散模型（Discrete Diffusion）正在成为生物序列生成的重要工具。与自回归模型逐个token从左到右生成不同，掩码离散扩散（Masked Discrete Diffusion）从一条全部被遮盖的序列出发，逐步「揭开」每个位置上的碱基或氨基酸。

它天然适合DNA、蛋白质这类有限字母表的数据，也已经在调控元件设计、蛋白质逆折叠等任务上展现出不错的生成质量。

但在真实的科研场景中，仅仅「从预训练分布中采样」并不够。做增强子设计的同学希望序列在目标细胞中活性更高；做蛋白质设计的同学希望序列更稳定、并能折叠回指定结构。

换句话说，我们真正需要的是：让一个已经训练好的生成模型，朝着某个任务奖励（reward）去生成。

这件事听起来很自然，做起来却并不容易。微调模型可以提升奖励，但需要额外训练、往往还需要奖励可微；而现有的免训练方法要么同样依赖梯度，要么依赖代价随轨迹长度或搜索宽度增长的粒子/搜索过程。很多生物学奖励，例如Enformer预测器或结构预测打分，本质上是「黑盒」，只能查询分数，难以反传梯度。

能否在不改动模型参数、不需要奖励梯度的前提下，只修改推理阶段的采样规则，就把一个冻结的离散扩散模型适配到任意黑盒奖励上？

在NeurIPS 2026发表的一篇工作中，研究人员给出了一个答案：把「奖励适配」看作一个分布搬运问题，并利用随机过程中的经典工具Doob h-变换，推导出它在掩码离散扩散上的形式。

在此基础上，研究人员设计了一个只需奖励值、无需训练的高效采样器D-DOIT（Discrete Doob-Oriented Inference-time Transformation）。

在HepG2增强子设计上，高预算的D-DOIT-Q预测活性达到6.19，超过需要微调的DRAKES（5.61）；在蛋白质逆折叠上，D-DOIT的结构–稳定性联合成功率为68.1%，高于无需微调的SMC（63.6%）和TDS（62.9%）。

论文链接：****https://arxiv.org/abs/2610.04938

代码链接：github.com/WeiminWu2000/D-DOIT

如何让冻结的离散扩散模型「听奖励的话」？

先回顾掩码离散扩散的生成过程。给定长度为N的序列，前向过程以一定概率把每个token替换成掩码符号m，并且一旦被遮盖就不会恢复；反向过程则从全掩码状态出发，在每一步利用模型的预测头（prediction head）给出每个被遮盖位置的「干净token」分布，并按照转移规则揭开其中一部分位置。

图1：标准掩码离散扩散的生成过程。从全掩码序列出发，每一步按转移规则独立揭开若干位置，已揭开的位置不再改动，直到得到完整序列。

从上到下是同一条序列在不同时刻的样子：中间那一步只在掩码位置上采样，已经揭开的位置原样带到下一行。

记token字母表为𝒱、掩码符号为m。下式中的α是一个干净token到第t步仍未被遮盖的概率，它随t增大而下降，在起点处接近0，此时序列几乎全是掩码。前向过程对每个位置独立地做：

反向过程从全掩码状态出发。预测头μ给出每个被遮盖位置的干净token后验，把它代入解析的反向后验，就得到从第t步到第t−1步的单步转移（逐坐标j独立）：

三行分别对应：这一步仍然保持掩码、被揭开成某个具体token、以及已经揭开的位置原样保留。整条反向链在各坐标上因子化：

上面的式子里已经写着这个过程的两个特点：一是每个位置独立揭开（最后一式的连乘），局部、快速，但对奖励毫无感知；二是揭开不可逆（单步转移的第三行），一旦某个位置被填上某个碱基，后续步骤就无法修改。

训练式方法：效果好，但代价高

一类方法通过强化学习或奖励微调直接更新骨干模型（如DRAKES）。它们可以获得很强的奖励提升，但需要额外的训练开销、超参数调节，通常还要求奖励可微。每换一个任务或奖励，就得重新训练一次。后面表格里的CFG（Classifier-Free Guidance，无分类器引导）也属于这一类。

免训练方法：要么要梯度，要么开销随轨迹增长

另一类方法保持模型冻结，在推理阶段做引导。分类器引导（Classifier Guidance，CG）需要奖励梯度；序贯蒙特卡罗（Sequential Monte Carlo，SMC）和扭曲扩散采样器（Twisted Diffusion Sampler，TDS）等粒子方法通过轨迹似然比重加权，代价与轨迹历史相关；软价值解码（Soft Value-Based Decoding，SVDD）等方法基于软价值估计在候选状态中做选择。

它们各有优势，但缺少一个统一、精确的目标分布视角来回答：引导之后，到底想从什么分布里采样？每一步的转移概率应该被修正成什么样？

奖励倾斜分布

离散 Doob h-变换

想要的分布目标

冻结的模型本身已经定义了一个分布，记作p，它描述了什么样的序列看起来是自然的。想要的不是它，但也不是「奖励最高的那一条」：后者往往钻了预测器的空子，根本不像真实序列。

真正想要的在两者之间，既保持模型眼里的自然，又偏向高奖励。把这个愿望写成公式，就是奖励倾斜分布（reward-tilted distribution）：

这里β是温度参数，分母是归一化常数。这个目标把预训练模型当作序列先验，同时偏好高奖励的输出：β越小，越激进地追求奖励；β越大，越贴近原模型，从而保留序列的「自然性」。

这对生物序列尤其重要：不希望为了刷高预测分数，就生成一堆不像真实增强子的序列。

先看连续情形：Doob h-变换是一个加性漂移

Doob h-变换是随机过程里的经典工具：把一个马尔可夫过程「条件化」到某个关于未来的事件上，得到的仍然是一个马尔可夫过程，只是多了一项由h决定的修正。这里的h就是那个未来事件在当前状态下的条件期望，即终点奖励权重的条件期望。

在连续扩散里，反向采样是一条随机微分方程（Stochastic Differential Equation，简称 SDE），Doob h-变换给出的修正加在漂移项上：

等价地说，采样用的score从∇logp变成∇log p+ ∇log h。

离散Doob h-变换：乘法重加权，而非加性漂移

回到掩码离散扩散。同样是Doob h-变换，在这里的形式却完全不同。研究人员先定义离散Doob h-函数：给定当前部分去噪的状态，未来最终奖励权重的条件期望：

它可以理解为当前这个「中间状态」的未来价值（reward-to-go）：从这里继续用原模型去噪，最终能拿到多高的奖励。研究人员证明（Theorem 2.1），只要把每一步的反向转移核按「下一状态与当前状态的h值之比」进行重加权，

那么变换后的反向链在终点的边缘分布恰好就是奖励倾斜分布，而模型参数完全不需要改动。

回头对照上一节：连续情形下的Doob修正是加在score上的一个漂移项；而掩码离散扩散的状态空间是类别型的，生成通过不可逆的「揭开」转移完成，根本不存在欧氏意义下的score场。

因此，在离散情形下，Doob变换必须作用在转移核层面：对每一个候选下一状态做乘法重加权。

D-DOIT

让精确变换变得可执行

理论上的精确变换需要知道每个候选下一状态的未来价值h，而精确计算它需要从每个候选出发做大量的未来去噪rollout，这在计算上是不可行的。为了让价值h可算，D-DOIT利用了两个近似和一个可选模块。

图2：D-DOIT概览。左：标准掩码扩散每步独立揭开各位置，便宜但对奖励无感知；右：D-DOIT的一步引导，即提出M个候选（图中M=3）、用预测头一次前向补全并打分，再按h值采样下一状态；下：可选的后期best-of-K精炼（图中K=4）。

近似一｜预测头h-估计器：一次前向，看见终点

掩码扩散模型的预测头本来就会为每个被遮盖的位置给出「干净token」的分布。研究人员利用这一点：对每个候选下一状态，固定已揭开的位置，用预测头一次性把剩余掩码位置补全成一条完整序列，再把它送入奖励函数打分，得到这个候选的h值估计（补全时用温度γ重整预测头）：

这样做有两个好处：

其一，避免了嵌套rollout（rollout指从当前状态把去噪过程一路跑到底，得到一条完整序列），每个候选只需一次前向；

其二，奖励只在补全后的序列上被「查询」，不需要可微、也不需要反传梯度，任何黑盒打分器都可以直接接入。

近似二｜有限候选的Doob加权选择

在每个引导步，D-DOIT先从（温度为τ的）模型反向核中采样M个候选下一状态，这M个候选共享同一次前向；随后对每个候选做一次补全与打分，得到每个候选的h值估计；最后按照与这些估计值的ω次方成正比的概率，从候选中采样出下一状态：

这里有两个容易混淆的温度，它们作用在同一个预测头上，职责却不同，两者都是对预测头做幂次重整。τ用在提议上，控制M个候选的多样性：

γ用在补全上，控制用来打分的那条完整序列有多尖锐：

当τ=1时，这就是对Doob变换核的有限候选自归一化近似。参数ω控制「跟随价值的锐度」，与β合起来等价于一个有效逆温度ω/β。每个引导步的代价是1+M次模型前向和M次奖励查询。

可选模块｜后期best-of-K：对冲不可逆的揭开决策

由于揭开不5可逆，早期的错误决定无法被修正。D-DOIT可选地在去噪后期某个分裂点之后，将每条轨迹复制成K份独立继续，最终保留奖励最高的那条。由于分裂只发生在末段，K倍的代价只在后期支付，而不是整条轨迹都乘以K。

分裂点放在末段主要有两方面的考虑，一方面是多样性：去噪早期还有大量位置没有揭开，采样本身就会把轨迹带向差别很大的方向，这时再复制K份，增加的多样性有限；

到了后期，大部分token已经固定，剩下的随机性很小，轨迹几乎注定收敛到一个结果，分裂才真正补上了缺少的变化。另一方面是算力：分裂点越靠后，每份副本需要重算的步数越少，K倍开销只落在末段，而不是整条轨迹。

两个工作配置：LBOK-S与Q

两个配置使用同一个采样器，采用不同的预算分配（下面为HepG2设置下的每条序列）：

• D-DOIT-LBOK-S：小候选预算（M=9）+后期best-of-K（K=4），约1,970次模型前向，运行时间接近TDS（约1.16×），面向实用的低预算场景；
• D-DOIT-Q：把更多推理预算投入到候选评估与精炼（M=52，K=8），约54,272次模型前向，用于探索冻结骨干模型的质量上限。

实验结果

研究人员在三个生物序列设计任务上评估D-DOIT，所有任务中的掩码离散扩散骨干模型都保持冻结。

HepG2增强子设计：免训练也能比肩训练式方法

在DRAKES HepG2增强子基准上，目标是生成200bp（碱基对）、在HepG2细胞中预测活性高且统计特征自然的增强子序列。评估指标包括预测活性（Pred-Activity）、染色质可及性（ATAC-Acc）、3-mer相关性与JASPAR motif相关性。

图3：HepG2增强子设计结果。左：预测活性，为中位数；右：染色质可及性，为均值。

图4：HepG2生成序列的统计自然性。左：3-mer相关性，可以为负；右：JASPAR motif相关性。

几个值得关注的结果：

• D-DOIT-Q在全部四项指标上超过所有免训练基线，预测活性达到6.19；并在Pred-Activity、3-mer Corr和JASPAR Corr上超过需要微调的DRAKES，仅在ATAC-Acc上略低（91.7%vs.92.5%）。
• 运行时间接近TDS的D-DOIT-LBOK-S，在预测活性上与TDS持平（4.56 vs. 4.58）的同时，把ATAC-Acc从40.0%提升到65.0%，JASPAR Corr从0.798提升到0.948。这说明提升并不只是预测分数被刷高，序列也更像真实的高活性增强子。
• 在相近的实际耗时下，D-DOIT-Q的四项指标均值都优于SVDD-PM；与等时的Frozen Best-of-N相比，在活性、可及性和3-mer相关性上更好，而后者在JASPAR相关性上更高。

细胞类型特异性增强子设计：更严格的可控生成

在Ctrl-DNA基准上，任务要求序列在目标细胞系中活跃、在非目标细胞系中沉默（覆盖HepG2、K562、SK-N-SH三个细胞系），以特异性ΔR（目标细胞得分减去非目标细胞平均得分）衡量。

图5：Ctrl-DNA细胞类型特异性设计结果。左：特异性ΔR；右：motif相关性。Frozen Best-of-708是与D-DOIT-Q等耗时的端点搜索对照。

在扩散引导类基线中，CG、SMC、TDS几乎无法推动特异性奖励，而D-DOIT明显做到了：LBOK-S在接近TDS耗时下将ΔR从0.002提升到0.033，同时Motif Corr从0.573提升到0.648；D-DOIT-Q进一步把ΔR推到0.119，代价是motif保持度有所下降。

在这个任务上，等时的Frozen Best-of-N同样能获得大部分增益。D-DOIT-Q与Frozen Best-of-708的均值差距很小。

蛋白质逆折叠：跨模态迁移

为了验证方法在不同序列空间和奖励下的通用性，研究人员将D-DOIT用于DRAKES蛋白质逆折叠基准：

给定目标骨架结构，生成氨基酸序列，同时评估预测稳定性（Pred-ddG，预测的折叠自由能变化）与结构一致性（scRMSD，self-consistency RMSD，自洽均方根偏差），以两者同时达标的比例作为成功率。

图6：DRAKES蛋白质逆折叠结果。左：稳定性与结构同时达标的联合成功率；右：预测稳定性。

D-DOIT取得了68.1%±3.5%的联合成功率，高于SMC（63.6%）和TDS（62.9%），并拥有最高的预测稳定性；相对的TDS在结构保真度上略好。配置（M=52，K=8，每个候选补全1次）在测试评估前就已固定。

鲁棒性

• 补全质量：每个候选用1次补全已足够实用；增加到2、3次可以把最终引导奖励从5.272提升到5.764和5.978，但会带来41%和81%的额外耗时。单步补全与多步补全得到的h值分布非常接近。
• 奖励预测器鲁棒性：把引导用的奖励换成Enformer、CNN或MLP预测器，D-DOIT相对无引导采样都能在四项指标上带来提升（Pred-Activity分别达到4.491、2.203、1.544，冻结模型仅为0.371），并且这些提升在独立的留出评估器上依然保持。

局限与展望

这项工作也有明确的边界：

• 近似带来的偏差：温度缩放、权重锐化与best-of-K选择都会使实际采样器偏离精确目标；后期分支也无法修正已经揭开的token。
• 湿实验验证：目前的评估依赖预测性的生物学基准，最终仍需要实验验证。
• 推理成本：最高质量的配置需要相当可观、且每次生成都要重复支付的推理开销，如何在奖励与算力之间权衡仍是一个实际问题。

展望未来，研究人员认为Doob式的转移核重加权为「用序列级奖励引导离散生成模型」提供了一个通用框架。它不局限于DNA和蛋白质：只要是掩码离散扩散、能对完整序列打分，就可以采用同样的思路。

进一步的，如何学习更准确的h-估计器、如何在揭开不可逆的约束下做更好的回溯与修正、如何与训练式方法结合实现摊销，都是值得继续探索的方向。

结语

回到开头的问题：能否只修改推理阶段的采样规则，就让一个冻结的离散扩散模型适配任意黑盒奖励？D-DOIT给出的回答是：可以，而且有一个精确的理论对象指引该怎么做：离散Doob h-变换。

它告诉我们，引导不是在离散空间里「加一个方向」，而是按每个候选的未来价值对转移概率进行重加权。

参考资料：

https://arxiv.org/abs/2610.04938

秒追ASI

⭐点赞、转发、在看一键三连⭐

点亮星标，锁定新智元极速推送！

跳转微信打开