🤖 machine learning

On the Design Space of Discrete Diffusion Online Adaptation for Molecular Optimization

本文提出了一种用于分子优化的离散扩散模型综合在线自适应框架，该框架集成了采集、奖励塑造、模型去偏、重放和有效性控制，旨在受限的预言机预算下超越离线微调和推理时搜索基准。

原作者： Trevor Chen, Ariel Dai, Jason Yang, Riccardo De Santi, Daniel Khalil, Wenda Chu, Nate Gruver, Pranav Murugan, Alexander F. G. Goldberg, Maruan Al-Shedivat, Yisong Yue

发布于 2026-07-07

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Trevor Chen, Ariel Dai, Jason Yang, Riccardo De Santi, Daniel Khalil, Wenda Chu, Nate Gruver, Pranav Murugan, Alexander F. G. Goldberg, Maruan Al-Shedivat, Yisong Yue

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下，你有一位大师级厨师（AI 模型），他通过学习海量的食谱积累了多年的烹饪经验。这位厨师知道如何制作数百万种不同的料理，从简单的吐司到复杂的舒芙蕾。这就是“预训练先验”（pretrained prior）。

然而，你并不只是想要任何一道菜。你有一个非常具体且昂贵的目标：你想要世界上最好吃的一道菜，但你支付给美食评论家（Oracle/预言机）进行品鉴和评分的预算非常有限。你负担不起让评论家品尝这位厨师可能做出的每一道菜。

这篇论文研究的是如何教导这位厨师，让他不再随机制作料理，而是开始专注于你想要的特定类型的食物，并且尽可能减少对评论家品鉴次数的需求。作者称之为**“在线自适应”（Online Adaptation）**。

以下是他们成功的“食谱”的简单拆解，使用了论文中的类比：

问题所在：“随机猜测”陷阱

如果你只是要求厨师做 1,000 道菜，然后让评论家品鉴其中最好的 10 道，你很可能会陷入困境。厨师会一直制作那些“还不错”的菜（比如标准的意面），因为那是他们最擅长的，但他们可能会错过那道虽然有点怪异且具有风险、但却极其美味的顶级佳肴。

论文研究了一个循环：厨师制作料理，评论家品鉴其中的几道，然后厨师根据反馈学习，以便在下一轮做得更好。但运行这个循环有很多种方式，作者想知道：哪些特定的技巧组合在一起效果最好？

成功食谱的 5 种配料

作者测试了一个拥有五种特定工具（旋钮）的“厨房”，并发现将它们全部结合使用会产生最好的结果，尤其是在处理小分子（如新药）时。

“赌徒的选择”（汤普森采样 / Thompson Sampling）
- 类比： 厨师并不仅仅只要求评论家去品鉴那些厨师“认为”最好的菜，厨师还会挑选一些具有不确定性的菜。也许厨师不确定香辣咖喱是否可行，但它可能非常惊艳。
- 结果： 这防止了厨师陷入反复制作同样“稳妥”的意面这种死循环。它迫使团队去探索那些高风险、高回报的想法。
“孤注一掷”的专注（CVaR 奖励塑造 / CVaR Reward Shaping）
- 类比： 通常，你可能会试图提高“平均水平”。但在这里，目标是找到那道“完美的菜”。这个工具告诉厨师：“不要在意平均水平；忽略那些平庸的菜。把所有的精力都放在让前 10% 的菜变得更出色上。”
- 结果： 它推动厨师去追求“头奖”，而不是满足于一份 B+ 等级的餐点。
“反群体思维”（密度熵正则化 / Density Entropy Regularization）
- 类比： 厨师天生喜欢做那些他们最熟悉的菜（“流行”的模式）。这个工具就像一个严厉的经理，说：“别再做那道流行的意面了！我知道你能做，但我们需要尝试那些冷门、不那么常见的食谱，以寻找隐藏的珍品。”
- 结果： 它迫使厨师离开舒适区，去探索那些他们平时会忽略的厨房角落。
“记忆库”（经验回放缓冲 / Replay Buffer）
- 类比： 如果厨师只学习“当下”制作的菜，他们可能会忘记几轮之前做过的一道好菜。这个工具保留了一份至今为止发现的最佳菜品的“精彩集锦”，并将其重新混合到训练中。
- 结果： 它稳定了学习过程，确保厨师在尝试新事物时不会忘记过去的成功发现。
“安全网”（有效性惩罚 / Validity Penalty）
- 类比： 在寻找完美菜肴的过程中，厨师可能会尝试用纯空气或不可食用的石头来“做菜”（无效分子）。这个工具会对任何不是真实、可食用菜肴的行为给予巨大的“差评”。
- 结果： 它防止厨师在不可能的想法上浪费时间，确保每一次尝试都是真实的、可烹饪的分子。

重大发现：小分子 vs. 蛋白质

论文发现，这种“完整食谱”在小分子（如新药）方面表现得异常出色。

原因： 厨师原始的知识（先验）是非常广泛且通用的。为了找到一种伟大的新药，厨师必须做出一个巨大的跨越，远离他们通常烹饪的内容。上述工具帮助他们在安全的情况下完成这个大跨越。

然而，对于蛋白质（如构建特定的酶）来说，结果则没那么显著。

原因： 厨师已经非常专业了。他们是针对那一类蛋白质进行专门训练的，所以“伟大的”菜品离他们已知的领域很近。他们不需要做出巨大的跨越，因此“反群体思维”和“孤注一掷”这些工具就不那么必要了。

最终结论

作者将他们的“在线自适应”循环与另外两种常见方法进行了对比：

离线微调（Offline Fine-tuning）： 用一大堆数据教导厨师一次，然后让他们直接进入厨房。
推理时搜索（Inference-Time Search）： 要求厨师一次性生成 1,000 道菜，并在中间不进行任何学习的情况下选出最好的那一个。

胜出者： “在线自适应”循环（五种工具的食谱）获胜。

它使用更少的评论家调用次数（Oracle calls）找到了更好的分子。
它在计算时间（GPU 小时）方面更高效。
当最优解远离厨师最初所知的内容时，它的威力尤其强大。

简而言之： 要寻找最好的新分子，不要只是随机猜测或仅仅学习一次。相反，要使用一个智能循环：探索风险想法、只关注顶尖表现、强迫 AI 离开舒适区、记住过去的成功，并保持一切的有效性。这种组合是发现高回报分子的最高效方式。

技术摘要：关于离散扩散模型在线优化分子设计的探索空间

问题陈述

分子优化通常始于一个预训练的生成模型（例如离散扩散模型），该模型捕捉了有效化学结构的广泛先验。然而，测试时的实际目标并非从该先验中采样，而是识别能够最大化特定任务奖励函数 $r(x)$ （例如结合亲和力或蛋白质适应度）的分子，且使用的是有限的预言机（oracle）预算。用于评估奖励的预言机通常成本高昂（例如，湿实验测定或昂贵的对接流程）。

核心挑战在于在线模型适配：如何引导预训练生成器 $p_{\theta_0}$ 向高奖励区域移动，同时保持有效性和多样性。这涉及一个紧密耦合的循环，系统必须决定：

采集（Acquisition）： 在给定的预言机预算下，选择哪些候选分子进行评估。
适配（Adaptation）： 如何将观察到的奖励转化为模型更新。
反馈重用（Feedback Reuse）： 在训练中保留哪些历史数据。
先验去偏（Prior Debiasing）： 在不导致有效性崩溃的前提下，使生成器偏离其预训练模式有多远。

虽然这些组件在孤立状态下已被研究过，但目前尚不清楚当它们被集成到完整的在线适配循环中时，是会产生互补作用、干扰作用还是变得冗余。

方法论

作者提出了一个用于离散扩散模型在线适配的可控“支架（harness）”模块化研究。该循环（图 1）在采样候选分子、选择一批次进行预言机评估以及微调生成器之间交替进行。该研究通过九个任务（六个小分子结合亲和力任务和三个蛋白质适应度任务）评估了五个特定的“插件式”设计选择（表 1）。

关键设计组件

转向（Steering/Acquisition）：
- 汤普森采样（Thompson Sampling）： 作者没有使用基于代理模型均值的贪婪选择，而是使用汤普森采样来选择候选者。这通过采样奖励估计值 $\hat{r}_i = \mu_i + \epsilon_i \sigma_i$ 来平衡探索与利用，其中 $\sigma_i$ 代表集成差异（不确定性）。这激励了对即使均值预测较低但具有高不确定性的候选者进行评估。
适配（奖励塑造与去偏）：
- CVaR 奖励塑造： 为了专注于高奖励尾部而非平均水平，训练奖励使用条件风险价值（CVaL）进行塑造。奖励被转换为 $\tilde{r}(x) = (r(x) - Q_\tau)_+ / (1-\tau)$ ，其中 $Q_\tau$ 是 $\tau$ -分位数（设置为 0.8）。这使梯度更新集中在表现最好的样本上。
- 密度熵正则化（模型去偏）： 为了防止模型坍缩到先验中已有的高密度模式，添加了一个正则项： $\log \tilde{r}_{debias}(x) = \log r(x) - \gamma \log p_{\theta_t}(x)$ 。由于计算离散扩散模型的精确边际似然 $\log p_{\theta_t}(x)$ 是难以处理的，作者使用了一种基于扩散 ELBO（单次传递近似）的高效蒙特卡洛估计器来近似负对数似然。
目标失配修正：
- 重放缓冲区（Replay Buffer）： 为了防止随着采样分布的变化而丢失早期发现的高奖励结果，使用了具有优先级剔除功能（优先移除低奖励样本）的固定容量 FIFO 缓冲区。通过分层采样抽取小批量数据，以确保高奖励离群值不会被稀释。
- 无效输出惩罚： 为了维持有效性（例如有效的 SMILES 字符串），对于解析失败（例如通过 RDKit）的候选分子，分配一个固定的负惩罚 $r_{inv}$ 。当探索产生无效结构时，这会将生成器拉回到有效的分子流形。

实验设置

任务： 针对六个蛋白质靶点的小分子结合优化（使用 GenMol 作为生成器，FlashAffinity/Boltz-2 作为预言机）以及三个蛋白质家族（GB1, CreiLOV, TrpB）的蛋白质适应度优化。
基线： 全量在线循环与以下模型进行对比：
- 离线微调（在静态数据集上训练）。
- 推理时搜索（束搜索 Beam Search、MCTS、Best-of-N 等）配合固定生成器。
- 消融版本的在线循环（每次移除一个组件）。
指标： 性能通过最佳发现奖励、Top-k 均值和归一化奖励进行衡量，并针对累计预言机调用次数和 GPU 时数进行评估。次要指标包括有效性、多样性（Tanimoto）、QED 和合成可及性（SA）。

关键结果

1. 组件的互补性

研究发现，所提出的组件提供了互补的收益，特别是在高奖励候选者远离预训练先验的小分子任务中。

小分子： 采集（汤普森采样）、CVaR 塑造和密度熵正则化带来了最明显的奖励增益。汤普森采样优化了预言机调用的分配；CVaR 将更新聚焦于奖励尾部；而去偏操作则推动探索远离过度占据的先验模式。
蛋白质适应度： 消融效应通常较弱。蛋白质的家族特定先验已经覆盖了高适应度区域，使得大规模分布偏移变得不那么关键。然而，汤普森采样通过改进预言机分配，仍然提供了持续的收益。

2. 稳定性机制

重放缓冲区和有效性惩罚的作用是稳定器而非主要的奖励提升器。重放缓冲区使梯度与评估数据的时序分布保持一致，防止了早期高奖励发现的丢失。有效性惩罚防止了“奖励黑客行为（reward hacking）”，即模型生成高分但化学无效的字符串，如果不加惩罚，这种失效模式会迅速发生。

3. 在线适配 vs. 推理时搜索

在匹配的预言机调用次数和 GPU 时数统计下，全量在线循环显著优于推理时搜索基线（如束搜索、MCTS）。

效率： 在小分子任务中，在线循环达到最佳搜索基线最终得分所需的预言机调用次数大约仅为后者的 1/58。在蛋白质任务上，加速比为 16 倍。
原因： 搜索方法将预算消耗在从固定生成器生成和评分候选分子上。相比之下，在线适配利用预言机反馈来更新生成器本身，从而使每次预言机调用的价值产生复利。搜索方法消耗大量计算资源进行生成，导致在固定的墙钟时间内几乎没有时间进行梯度步的微调。

4. 分布偏移分析

当最优解需要相对于预训练先验进行大幅度偏移时，性能增益最为显著。

密度熵正则化是主要的驱动力，它推动生成器的质量远离高密度先验区域，向低密度、高奖励区域移动。
CVaR 和汤普森采样根据高奖励候选者相对于先验的位置，有条件地促进这种偏移。
蛋白质任务显示出较弱的分布偏移，因为家族特定先验已经集中在接近高适应度的区域。

意义与主张

本文提出了一个用于高效反馈分子优化的实用方案：结合采集、奖励塑造、去偏、重放和有效性控制的在线微调。

创新性： 本工作在一个统一的离散扩散框架内整合并系统地评估了这些组件，证明了它们并非冗余，而是相互增强的。
效率： 所提方法优于离线微调和推理时搜索基线，特别是在预言机昂贵且目标分子相对于预训练模型属于分布外（OOD）的情况下。
鲁棒性： 该方案被证明在不同的微调目标（DDPP-LB 和 VIDD）和不同的领域（小分子 vs. 蛋白质）中都是鲁棒的，尽管增益幅度有所不同。

作者得出结论，虽然这些组件在小分子设计（先验较广且目标较远）中最关键，但该框架为在有限反馈的科学发现任务中引导离散生成模型提供了一种通用的方法。