Nat. Mach. Intell. 2026 | Saturn: 通过记忆机制调控实现高效样本利用的生成式分子设计

今天介绍的是发表在 Nature Machine Intelligence 上的一篇论文，题为“Sample-efficient generative molecular design using memory manipulation”。该研究主要围绕药物分子生成设计中的“样本效率问题”展开，重点探讨了如何在有限计算预算下高效生成具有目标性质的分子。作者提出了一种基于 Mamba 架构与记忆增强机制相结合的生成框架 Saturn，通过引入经验回放与数据增强策略，使模型能够在较少的评价次数下快速优化分子性质。研究表明，该方法在多参数优化任务中优于多种现有模型，甚至可以直接利用高精度的密度泛函理论（DFT）计算作为优化目标。该工作对于提升生成式分子设计的效率、推动高精度计算与人工智能方法的结合具有重要参考价值。

获取详情及资源:

0 摘要

近年来，用于药物发现的生成式分子设计已取得一系列实验验证成果。基于分子字符串表示的语言模型已成为最成功的模型架构之一。下游任务成功的关键在于计算预测器（即用于预测分子性质的计算模型）与目标终点（如结合亲和力）之间是否具有良好的相关性。为此，现有方法通常先采用成本更低、通量更高的替代预测器进行筛选，再对最有前景的候选分子子集使用高保真预测器进行评估。如果能够直接生成在高保真预测器（即计算成本更高但预测精度更高的模拟方法）评估下具有最优性质的分子，将有望显著提升生成式设计的效率并提高命中率。然而，当前模型的效率仍不足以支持这一目标，体现了样本效率方面的限制问题。近期，Mamba架构被提出作为广泛应用于大语言模型的Transformer架构的一种替代方案。已有研究已在从自然语言补全到生物基础模型等多种任务中验证了Mamba的性能。在此基础上，提出了一种名为Saturn的框架，用于展示Mamba架构在生成式分子设计中的应用。首先，阐明了结合数据增强的经验回放机制如何提升样本效率，以及Mamba架构如何进一步强化这一机制的效果。随后，在与药物发现相关的多参数优化任务中，证明了结合经验回放的Mamba模型优于16种现有模型，并且具备足够的样本效率，能够直接以密度泛函理论模拟作为高保真预测器进行优化。

1 引言

在过去一年半中，针对药物发现的生成式分子设计实验验证研究显著增加。最常见的范式是算法优化，其基本流程包括：生成分子、通过计算预测器评估其性质、更新模型并重复该过程。在缺乏实验验证的情况下，评估分子适用性的关键在于计算预测器与真实终点之间的相关性。通常，该设计过程采用漏斗式流程，即首先使用资源消耗较低的预测器筛选出最有潜力的候选子集，再通过计算成本更高的高保真预测器进行评估。一个典型示例是基于分子对接估计分子与蛋白质靶标的结合亲和力。经过调优的分子对接方法具备一定预测能力，并已成为实现实验验证的生成式设计流程中的常见组成部分。随后，从对接筛选出的候选分子中进一步采用更高保真的预测方法，尤其是分子动力学模拟，以获得更准确的结合亲和力估计。直接对高保真预测器进行优化有望显著提升生成分子的质量。然而，由于计算成本过高，这一策略在实际中难以实现，体现了样本效率问题。解决该问题的途径包括在不降低精度的前提下显著提升模拟效率，或提升生成模型的效率，使其在有限预测预算下实现优化。此外，定量构效关系模型也被广泛应用，该类模型具有较高预测精度，但其适用范围通常局限于训练数据分布内。

现有研究已采用多种模型架构开展生成式分子设计，包括基于简化分子输入线性表示的循环神经网络、Transformer、变分自编码器、对抗生成方法、基于图的模型、GFlowNet、遗传算法以及扩散或流模型等。长短期记忆网络也被用于遗传算法框架中的预测模型，用于生成自组装肽，并作为适应度函数。

尽管近年来三维分子生成逐渐成为研究趋势，基于语言的模型在样本效率方面仍具有显著优势，并能够满足如分子对接等依赖三维结构的优化目标。因此，文本仍然是生成式设计中最常用的分子表示方式之一，且相关模型架构已从循环神经网络和Transformer扩展至结构化状态空间序列模型。SMILES表示的非单射语法特性为生成设计提供了有利条件，即同一分子可对应至少N种不同的SMILES表示，这一过程称为SMILES增强、枚举或随机化。该机制可用于在低数据条件下进行模型预训练、提升样本效率，并支持基于单个正样本的迁移学习。

分子优化基准的提出进一步表明，提高样本效率在实际应用中具有重要意义。因此，近期研究通常在固定预测预算约束下进行模型评估，结果显示基于语言的模型仍表现优异。进一步地，AugmentedMemory方法通过结合经验回放与SMILES增强，有效提升了样本效率。在此基础上，提出了Saturn框架用于生成式分子设计。首先，对AugmentedMemory的作用机制进行了阐释，因为原始工作仅展示了其经验性能。随后，对包括循环神经网络、解码式Transformer以及Mamba模型在内的多种语言模型架构进行了系统评估。结果表明，Mamba能够通过强化数据增强与经验回放机制，在一定程度上通过策略性过拟合进一步提升样本效率。主要贡献如下：

图1｜基于 Mamba 架构的 Saturn 生成框架 所有生成的 SMILES 及其奖励在规范化处理后均存储于 oracle 缓存中。a. 可选择使用回放缓冲区作为父代种群应用遗传算法（GA）。b. 使用 Augmented Memory 对智能体进行多次更新。

2 结果和讨论

该节包含三个部分：首先通过一个简化的多参数优化任务分析AugmentedMemory的优化动力学；其次，在多参数优化分子对接任务中，与15种模型（包括基于数据集筛选的方法）进行基准比较；最后，以密度泛函理论电子性质作为高保真预测器进行直接优化。除密度泛函理论相关实验外，其余所有实验均在10个随机种子（0至9）下运行；由于计算成本较高，密度泛函理论实验仅使用3个随机种子（0至2）。在各项实验中，通过调整预测预算（即被评估的唯一分子数量）以实现与既有工作的公平比较，并在具体实验中明确说明。

2.1 解析AugmentedMemory的优化动力学

首先对AugmentedMemory的最优架构及其超参数进行分析。通过在循环神经网络、解码式Transformer和Mamba三种模型骨干结构上，改变批大小和数据增强轮数，对AugmentedMemory算法进行实验评估。随后，分析样本效率与多样性之间的权衡关系，并表明在化学空间中的局部开发机制对于提升样本效率具有重要作用。

实验细节：构建了一个简化的多参数优化任务，其目标为分子量小于350Da、环数不少于2，并最大化拓扑极性表面积。该优化任务要求生成具有杂原子富集环结构的分子，这类分子与训练数据存在显著差异，因此同时用于测试分布外优化能力。预测预算设为1000，模型基于ChEMBL33数据集进行预训练。

评价指标：样本效率通过产率和预测负担进行衡量。产率定义为生成且奖励值超过阈值的唯一分子数量，预测负担定义为生成N个满足奖励阈值的唯一分子所需的预测调用次数。本实验中奖励阈值设为0.7，此时生成的分子开始具有杂原子饱和环结构。大多数配置在预算范围内均能生成部分满足阈值的分子，从而可以进行统计分析。

不同模型骨干下AugmentedMemory的性能上限分析：AugmentedMemory是一种强化学习算法，通过对高奖励SMILES序列的重复学习提升样本效率。假设只要仍能生成新的高奖励SMILES序列，在批大小减小的情况下样本效率可以提高，但以多样性降低为代价。重复生成的SMILES不会引入额外预测开销，因为其奖励可通过SMILES标准化匹配从缓存中获取。通过网格搜索，在不同模型骨干下对批大小（64、32、16、8）和数据增强轮数（0至20）进行系统评估。结果表明，增加数据增强轮数通常可以提升样本效率指标，但该趋势在Mamba模型中更加稳定，表现出近似单调提升。在最具挑战性的预测负担指标中，Mamba在所有重复实验中均成功，而循环神经网络和Transformer分别仅在部分实验中成功。基于Mamba在样本效率上的优势，后续分析主要以其与循环神经网络进行对比。

样本效率与多样性权衡：在固定Mamba模型的条件下，进一步分析批大小、增强轮数、产率和多样性之间的关系。结果表明，增加增强轮数通常能够提高产率，但在较小批大小条件下表现不稳定，这是由于期望奖励的估计更加噪声化。在较小批大小（16和8）下，随着增强轮数增加，产率提升更加显著，但多样性显著下降。当批大小为8且增强轮数较高时，结果方差增大，多样性下降更加明显。综合来看，批大小为16、增强轮数为10的设置在样本效率、多样性与方差之间取得了较好的平衡。

进一步引入遗传算法作为生成过程的一部分，以研究其对多样性和样本效率的影响。在固定批大小为16并使用Mamba模型的条件下，通过网格搜索同时调整增强轮数和遗传算法的使用。结果表明，增加增强轮数仍能提升产率但降低多样性；引入遗传算法后，产率提升幅度减小，但多样性得到保持；若将更多预测预算分配给遗传算法，则产率反而下降。总体来看，遗传算法有助于在增强轮数增加的情况下维持多样性。

综合上述结果，较小批大小和较高增强轮数结合SMILES增强以及遗传算法，构成调控样本效率与多样性权衡的重要手段。在需要生成尽可能多高质量解的场景下，较小批大小和较高增强轮数更为有利；而在需要更高多样性的场景下，较大批大小、较低增强轮数以及引入遗传算法更为适合。在后续主要实验中未采用遗传算法。

模型能力对样本效率的影响：尽管Mamba与循环神经网络参数规模相近，但前者在预训练阶段收敛至更低的损失值。Mamba更倾向于生成重复的SMILES序列，表明其对数据分布存在一定程度的过拟合。在高增强轮数条件下，生成过程中最大条件token概率接近1，使得模型更可能重复生成相同SMILES序列。结果表明，更强的分布建模能力有助于提升样本效率。进一步实验表明，通过增加参数规模，循环神经网络和Transformer同样可以达到类似性能，因此该能力并非Mamba独有。该研究选择Mamba作为默认架构，是因为其在参数效率方面表现更优。

增强SMILES概率分布压缩机制：为解释AugmentedMemory的作用机制，设计了一个子实验。首先生成分子填满缓冲区，并在执行AugmentedMemory前后保存模型参数，同时记录所有增强后的SMILES序列。结果表明，经过AugmentedMemory后，模型更倾向于生成这些增强序列。对于原本概率较低的SMILES，其概率提升幅度更大；而对于原本概率较高的SMILES，由于softmax函数趋于饱和，其变化较小。在部分情况下，增强后某些SMILES的生成概率反而下降，这是由损失函数形式决定的。当增强序列与模型预测概率相等时，该项损失为零，其变化由批中其他序列决定。整体来看，AugmentedMemory通过调整概率分布，使增强SMILES序列的生成概率向最可能序列收敛。

Mamba的跳跃与局部开发行为：模型在训练过程中逐渐更倾向于生成缓冲区中的分子。基于此，提出Mamba具有“跳跃并局部开发”的行为特征，即通过生成某些SMILES表示形式，并对其进行微小token修改，从而在分子图层面实现局部探索。通过生成3000个分子并划分为多个子集，利用流形降维方法分析其分布轨迹。结果显示，Mamba在化学空间中的移动更加具有方向性，生成分子在局部区域内更为集中。进一步分析子集内和子集间的Tanimoto相似性表明，Mamba在子集内部具有更高相似性，而子集之间保持较低相似性。综合来看，在批大小为16、增强轮数为10的设置下，Mamba结合AugmentedMemory通过跳跃与局部开发机制有效提升了样本效率。该模型配置在后续中被定义为Saturn，并作为默认设置使用。

图2｜探索—利用权衡的调控及 Augmented Memory 的作用机制 a. 批大小与增强轮数对产率（yield）和 IntDiv1（多样性）的影响。总体而言，减小批大小并增加增强轮数可提升产率，但以牺牲多样性为代价。stdev，标准差。b. 在将批大小固定为16时，启用遗传算法（GA）可保持多样性，但随着增强轮数增加，产率提升幅度较小。c. 在强化学习（RL）训练轮次中的平均最大 token 概率。增强操作将智能体的动作分布推向 delta 分布（即将大部分概率质量集中于特定 token）。d. 上：运行 Augmented Memory 后，缓冲区中的 SMILES 的增强形式更容易被生成。回归线衡量在不同初始 NLL 下的 ΔNLL 变化，表明初始 NLL 越大（即原本概率越低），其概率提升越显著。下：当聚焦于初始 NLL ≤ 20 的 SMILES 时，这一关系有所减弱，但趋势仍然存在。两个子图中均标注了决定系数（R2）。e. 将 3,000 次 oracle 预算测试实验划分为每组 300 个 SMILES。对智能体在化学空间中的探索轨迹进行统一流形近似与投影（UMAP）嵌入（箭头表示每组的质心）。与 RNN（基线 Augmented Memory）持续进行全局采样不同，Mamba 表现出具有方向性的探索路径。f. Mamba 表现出“跳跃并局部开发”的行为，其组内 Tanimoto 相似性（上方数值）高于 RNN；下方数值表示组间相似性。

表1｜命中率（%）

2.2 Saturn框架的基准评估与优化能力提升

本节遵循Lee等提出的目标感知片段提取、组装与修饰模型的设定，在多参数优化分子对接任务上对Saturn框架进行基准测试。

实验细节：为实现与GEAM方法的严格对比，提取其预测器代码并在ZINC250k数据集上进行预训练。同时，其余模型基线亦来源于同一研究，并在相同数据集上训练。所有模型均针对GEAM提出的多参数优化目标函数进行优化：

R(x)=DS^(x)×QED(x)×SA^(x)∈[0,1]

其中，DS^表示归一化后的QuickVina2对接评分，SA^表示归一化后的合成可达性评分。按照相同设置，对五个蛋白靶点进行分子对接。实验中复现GEAM结果，并在预测预算为3000的条件下对Saturn在10个随机种子上进行评估。需要注意的是，GEAM在预训练阶段依赖已预计算对接值的标注数据，因此存在较高的前期计算成本。

评价指标：采用命中率作为主要指标，即同时满足以下条件的分子比例：对接评分优于已知活性分子的中位数，QED大于0.5，且SA小于5。同时定义新颖命中率，即在上述条件基础上，进一步要求与训练数据的最大Tanimoto相似度小于0.4。此外，引入更严格的评价指标，包括严格命中率和严格新颖命中率，分别要求QED大于0.7且SA小于3。尽管实际药物候选不必满足这些更严格标准，但该指标更能反映多目标优化能力。同时采用IntDiv1和#Circles指标衡量分子多样性。

Saturn与GEAM优于所有基线方法。通过对比随机采样方法与多种模型，结果表明仅AugmentedMemory、MOOD、GEAM及Saturn能够超越基线，其中GEAM与Saturn性能相近。在新颖命中率指标上，GEAM优于Saturn，后者生成分子更接近预训练数据分布。为满足新颖性约束，将任务划分为两个阶段：首先以Tanimoto不相似性为唯一优化目标进行生成，随后以该模型为初始点执行多参数优化任务。结果显示性能可恢复并达到GEAM水平，同时该过程计算成本较低，体现了方法的灵活性。

表2｜严格命中率（%）

Saturn在多参数优化中表现出更强能力。通过引入更严格的筛选条件（QED大于0.7且SA小于3）评估优化能力，结果表明GEAM的命中率显著下降，而Saturn基本保持稳定，说明其对多目标函数的优化更充分。同时，Saturn在更少的预测调用次数下即可获得满足严格条件的分子，尽管以降低多样性为代价。在部分靶点任务中，GEAM在多数重复实验中无法获得足够数量的满足严格条件的分子，而Saturn在所有实验中均成功实现。

进一步分析优化能力，将不同方法生成的分子合并并比较奖励分布。结果显示，GEAM生成分子的奖励分布与随机采样相近，而Saturn显著提升整体奖励水平。这是由于多参数优化目标同时考虑QED、SA及对接评分，GEAM虽在对接评分上优于随机采样，但在QED与SA方面表现较差。相比之下，Saturn能够更均衡地优化所有目标，从而使奖励分布整体向高值区域偏移。

图3｜在五个蛋白靶点上的奖励分布比较（ZINC 250k 随机采样、GEAM 生成及 Saturn 生成分子） 分子在十个随机种子（0–9，含）下汇总，未能完成对接的分子已被剔除。箱线图的须表示数据的完整范围，中间线为中位数（Med.），上下边界表示四分位距（IQR）。从每个汇总集合中选取奖励（Rew）最高的前两个分子进行展示。QED，药物相似性定量估计；SA，合成可及性评分；DS，对接评分。需要注意的是，这些分子对在 ECFP4 Tanimoto 相似性小于 0.4 的条件下仍表现出不相似性。

2.3 在密度泛函理论精度水平下直接优化电子性质

该节展示了利用Saturn框架直接优化高保真预测器的可行性。通过对生成分子进行密度泛函理论几何优化，以最高占据分子轨道与最低未占据分子轨道能隙作为目标性质，用于功能材料设计中的性能优化，例如半导体材料。优化目标为在有无同时约束合成可达性评分的条件下，最小化该能隙。为控制计算成本，实验采用3个随机种子并将预测预算限制为500。

结果展示了所有生成分子的整体分布，并按照生成顺序划分为不同阶段，以反映在严格预算约束下的优化过程。随着优化进行，生成分子的性质分布逐步向更低能隙区域移动，这一点可由标注的中位数能隙分子进一步验证。

进一步分析显示，能隙最低的分子在保持较低结构相似性的条件下，模型能够隐式引入参与共振与环共轭的带电原子，这与已有研究结果一致。同时，部分中性分子也呈现出高度共轭结构。结合整体分布可见，在未引入合成可达性约束时，生成分子的合成难度逐渐增加，说明该约束在生成过程中具有实际限制作用。

总体来看，结果表明在Saturn框架下，能够在密度泛函理论精度水平上直接对电子性质进行优化。

图4｜在 DFT 精度水平下的直接优化（oracle 预算为 500） 优化目标为仅最小化 HOMO–LUMO 能隙（右图），或同时最小化该能隙与 SA 评分（左图）。分子在三个随机种子（0–2，含）下汇总，等高密度图展示了从最初生成的 33% 分子到中间 33%，再到最后 33% 分子过程中 HOMO–LUMO 能量向更低值的转移。图中的点及标注分子表示各分段的 HOMO–LUMO 中位值，同时标注了首段与末段之间的 HOMO–LUMO 中位能量及 SA 评分。移除 SA 评分约束后，如预期所示，可获得更低的 HOMO–LUMO 值。从每个汇总集合中选取 HOMO–LUMO 能隙最低的两个分子进行展示。需要注意的是，这些分子对在 ECFP4 Tanimoto 相似性小于 0.4 的条件下仍表现出不相似性。

3 结论

该研究提出了Saturn框架，用于实现高样本效率的从头分子设计。该方法基于Mamba架构结合强化学习，并展示了其如何增强AugmentedMemory的效果。AugmentedMemory通过结合SMILES增强与经验回放来提升样本效率。进一步阐明了其作用机制，即提高对回放缓冲区中不同SMILES表示形式的生成概率。通过在多种模型架构（包括循环神经网络、解码式Transformer和Mamba）中分析样本效率与多样性之间的权衡关系，结果表明，具有更强分布学习能力的模型能够增强AugmentedMemory的效果。这一机制使生成分子之间更加相似，从而有利于提升样本效率，但会以降低多样性为代价。

Saturn框架能够直接针对高保真预测器进行优化，例如密度泛函理论。然而，该方法仍存在两个重要局限。首先，相较于该研究中采用的分子对接方法，分子力学泊松–玻尔兹曼表面积方法或自由能微扰方法在结合亲和力预测方面通常具有更高精度。因此，已有研究已尝试直接针对这些高保真预测器进行分子生成。尽管当前工作已在密度泛函理论场景中验证了该框架，未来仍需扩展至更高精度的结合亲和力预测方法。其次，生成分子的可合成性仍是限制因素。尽管部分实验在优化目标中引入了合成可达性评分，但通过显式逆合成分析可更准确评估分子可合成性，并引导模型生成更具可合成性的分子。已有后续研究将逆合成模型与Saturn框架结合，使得生成分子不仅可合成，还可满足特定反应路径约束，并优化合成路线长度。

AugmentedMemory算法通过提高回放缓冲区中分子的生成概率，实现局部开发。在该研究的案例中，该行为有助于提升样本效率。然而，该机制并不适用于所有优化任务，因为结构与性质之间的关系可能导致奖励信号出现非平滑变化。例如，在活性陡变现象中，微小结构变化即可引起性质的显著差异，此时生成相似分子反而可能不利。在此类情形下，更具探索性的策略可能更为有效。因此，一个重要的研究方向是系统分析Saturn框架的失效模式，并在面对不同优化目标时，实现对探索与开发之间权衡的自适应调控。相关研究已探索通过并行生成模型分别执行探索与开发任务，或通过多个协同模型共同搜索设计空间，以提升整体优化能力。