Amortized Molecular Optimization via Group Relative Policy Optimization¶

会议: NeurIPS 2026
arXiv: 2602.12162
OpenReview: https://openreview.net/forum?id=BueyrCFsnx
代码: https://github.com/Hash-hh/AMORTIX/
领域: 强化学习
关键词: 分子优化, 摊销强化学习, 结构约束设计, 图Transformer, 组相对策略优化

一句话总结¶

针对结构约束分子优化中起始母核难度异质性导致强化学习基线方差崩溃的痛点，本文提出基于图Transformer的摊销生成策略 AMORTIX，通过群组相对策略优化（GRPO）消除跨母核方差，实现推理阶段零预言机调用的单次前向高质量分子修饰。

研究背景与动机¶

在现代计算机辅助药物研发中，先导化合物优化、骨架跃迁修饰、片段生长以及连接子设计等关键任务，本质上都需要在严格保留特定起始子结构（如母核、活性药效团）的前提下对分子进行精准编辑与修饰。为了引导生成分子具备理想的生物活性、药物亲和力与成药性，算法通常需要依赖计算预言机（Oracle）进行反馈评估。然而，目前占主导地位的“实例优化”（Instance Optimization）范式——包括基于图遗传算法（如 Mol GA）和引导式离散扩散模型（如 GenMol）——将每个分子的优化都视作一次孤立的从零搜索过程，针对每个给定的起始骨架都需要反复调用数千次预言机评分。当候选骨架库规模庞大，或者预言机评估成本高昂（如高精度分子动力学模拟或自由能微扰 FEP）时，这种推理成本随输入结构数量线性暴增的模式极难规模化。

理论上，“摊销优化”（Amortized Optimization）提供了一条优雅的解决路径：通过在离线阶段充分训练一个条件策略网络，在推理时便能对任意未知的新骨架直接进行单次前向采样，将巨大的搜索开销完全摊销在离线阶段，实现推理期零预言机调用。然而，现存的摊销式生成方法在面对严格的硬性拓扑约束时普遍泛化能力受限；更严峻的是，当尝试在多样化骨架上通过强化学习进行端到端优化时，策略训练往往会剧烈震荡甚至崩溃。这种失败的核心根源在于“起始分子结构的优化难度存在极端异质性”：部分母核处于化学空间中天然容易修饰出高分子的区域，而另一些紧凑或高度受限的母核即便找到最优解其绝对得分依然偏低。如果采用全局基线（Global Baseline）或常规 Critic 价值网络，奖励信号的标度差异将产生极高方差的优势估计（Advantage），使得策略更新被“简单骨架”完全主导，无法对复杂困难骨架建立有效的梯度指引。

针对这一本质瓶颈，本文将大语言模型推理任务中的群组相对策略优化思想迁移至结构约束分子生成领域，构建了无需价值网络的直接强化学习框架。核心 idea：构建严格保证原子价键化学合法性的解码器图Transformer策略，并在训练时对每个起始骨架独立采样完成轨迹群组，以组内均值实施群组相对优势归一化（GRPO），彻底解耦骨架内在难度差异对奖励信号的干扰，实现高泛化性、零推理开销的端到端摊销分子优化。

方法详解¶

整体框架¶

AMORTIX 将分子图的逐步修饰与构建过程建模为一个离散的马尔可夫决策过程（MDP）。给定任意非空的初始分子子图（从单原子到复杂药物骨架），策略网络按步骤向当前子图中添加新原子或在已有原子间添加化学键，直至生成终止标志。整个框架由三大部分紧密衔接：首先是以原子与键为基本单元的解码器图Transformer骨干，配合原子价键动态掩码，确保每一步动作天然符合化学价态；其次是基于 ChEMBL 大规模数据集的自监督模仿预训练，赋予策略化学合理性先验；最后是基于无 Critic 的群组相对策略优化（GRPO）微调阶段，通过按骨架独立成组计算相对优势，引导策略探索高奖励化学空间。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    In["输入起始分子骨架 S0<br/>(非空子图 / 骨架 / 单原子)"] --> NodeEnc["原子状态动态嵌入<br/>(元素类型 + 当前度数 + 决策阶段指示)"]
    NodeEnc --> GTBackbone["图Transformer主干网络<br/>(无位置编码 + 键级注意力偏置)"]
    GTBackbone --> ActionDec["分层动作解码与价键动态掩码<br/>(L0操作类型 → L1目标原子 → L2键级指定)"]
    ActionDec --> ValidCheck{"生成是否终止 /<br/>达到最大原子数?"}
    ValidCheck -- 否: 拓展分子图 --> NodeEnc
    ValidCheck -- 是: 获得完整分子 G --> RolloutGroup["群组轨迹采样与预言机评估<br/>(每骨架采样 G 条轨迹计算奖励)"]
    RolloutGroup --> GRPOCalc["组内相对优势归一化与策略更新<br/>(减去组内均值，单步梯度回传)"]
    GRPOCalc --> Out["优化后分子生成策略<br/>(推理期零预言机调用)"]

关键设计¶

1. 分层动作分解与价键动态掩码：保证单步动作低复杂度与拓扑合法性

在离散分子图生成中，直接预测全局图结构或在巨大组合动作空间中采样极易导致训练发散并生成大量违背化学常理的无效图。为了降低动作空间的组合复杂度，AMORTIX 将每一步的图拓展动作分解为三层阶梯式决策：Level 0 为操作选择（从 23 种包含电荷与手性的原子词表中选取新原子添加、指定现有原子以新增化学键、或执行 Stop 终止生成）；Level 1 为目标选择（若 Level 0 为添加新原子，则在现有图中选择其连接锚点；若 Level 0 为修改现有原子，则指定新键的终点原子）；Level 2 为键级指定（确定新生成的化学键为单键、双键或三键）。在每个决策子层级，模型引入基于 RDKit 原子价态规则的严格动态掩码 \(\mathcal{M}(s_t)\)，将任何会导致原子超过最大共价容量的动作概率直接置零。通过这种机制，整个采样过程被完全限定在化学合法空间内，无需额外的非法惩罚项，保证了生成分子百分之百的化学拓扑有效性（Chemical Validity）。

2. 键级拓扑感知图Transformer架构：保持图置换等变性与上下文捕获

为了有效处理分子图的无序性质并对局部键型结构保持敏感，AMORTIX 采用仅解码器（decoder-only）的图Transformer架构，将分子表示为无序的原子集合，彻底摒弃破坏置换对称性的绝对位置编码。模型在原子输入特征中融合元素类型、当前连接度数以及指示当前处于 Level 0 还是 Level 1 决策的动作状态嵌入（Action-State Embeddings），使网络在级联决策过程中能够动态“聚焦”于被选中的活性位点。为了在全局全连接自注意力中保留分子的真实拓扑连接信息，模型将化学键级作为偏置项直接注入注意力矩阵中。原子 \(i\) 与原子 \(j\) 之间的注意力权重计算如下：

\[A_{ij} = \text{softmax}\left(\frac{Q_i K_j^T}{\sqrt{d_k}} + \phi(e_{ij})\right)\]

其中 \(\phi(e_{ij})\) 为根据两原子间实际键级 \(e_{ij}\) 映射的可学习标量偏置。此外，模型引入一个全局全连接的“虚拟超级节点”（Virtual Node），用于聚合整图上下文表征，并在 Level 0 动作头中直接用于判断整图级决策（如输出 Stop 标志或新增原子的类型）。

3. 群组相对优势归一化（GRPO）：消除骨架异质性带来的梯度高方差

传统强化学习方法（如带全局基线的 REINFORCE 或带 Critic 价值网络的 PPO）在处理多样化骨架修饰时，会因为不同母核固有的成药难度差异而产生尺度不一的奖励波动。全局移动平均基线无法辨别一个低分是源于策略在简单骨架上的失误，还是在困难骨架上的高质量极限发挥；而为图结构训练一个准确的价值网络本身极难收敛。AMORTIX 抛弃了独立的 Critic 网络，在每个训练批次中对 \(B\) 个起始骨架中的每一个独立采样 \(G\) 条完整生成轨迹 \(\mathcal{O}_i = \{O_{i,1}, \dots, O_{i,G}\}\)，并通过预言机评估得到标量奖励 \(R(O_{i,j})\)。每个样本的优势函数 \(A_{i,j}\) 直接通过相对于该骨架内部组均值 \(\mu_i\) 进行中心化计算：

\[A_{i,j} = R(O_{i,j}) - \mu_i, \quad \mu_i = \frac{1}{G}\sum_{k=1}^G R(O_{i,k})\]

与标准 GRPO 相比，本文做出了两个针对化学分子特性的关键设计调整：第一，明确省略除以组标准差的操作，防止某些极易或极难的骨架因所有样本得分极度相近（标准差趋近于 0）而在反向传播中被赋予不成比例的巨大梯度权重；第二，明确省略轨迹长度归一化，避免惩罚那些只需较少步数就能修饰出高分结构的紧凑生成路径。参数梯度的聚合形式为：

\[\nabla_\theta \mathcal{J}(\theta) \approx \frac{1}{BG} \sum_{i=1}^B \sum_{j=1}^G A_{i,j} \sum_{t=1}^{T_{i,j}} \nabla_\theta \log \pi_\theta(a_t \mid s_{<t})\]

由于分子轨迹采样速度极快，模型在每个批次上仅进行单次梯度更新，彻底免去了 PPO 复杂的近端裁剪操作与价值网络对齐负担。

损失函数 / 训练策略¶

整个训练流程分为两阶段： 1. 化学合理性自监督预训练：在 ChEMBL 35 数据库过滤出的 150 万分子上，将每个分子分解为规范的单原子逐步生长轨迹，采用最大似然估计（Teacher Forcing）优化交叉熵损失 \(\mathcal{L}_{\text{MLE}}(\theta) = -\mathbb{E}[\sum_{t=0}^T \log \pi_\theta(a_t^* \mid s_t^*)]\)，使模型在进入强化学习前建立符合实际药物化学直觉的合理分子分布先验。 2. GRPO 强化学习微调：采用 Adam 优化器，学习率设为 \(1 \times 10^{-4}\)，梯度裁剪 \(L_2\text{-norm} = 1.0\)。在每个 epoch 中采样 10 个起始结构（9 个真实骨架 + 1 个单碳原子），每个骨架采样 16 条轨迹，单批总规模约 160 条轨迹，进行单步梯度回传，总预言机预算严格控制在 50,000 次内。

实验关键数据¶

主实验¶

在目标导向的骨架修饰任务中，从 ZINC-250k 数据集中通过 Murcko 骨架分解与 Butina 聚类（Tanimoto 阈值 0.4）严格切分出拓扑结构完全外推（OOD）的 500 个测试骨架，分别在单靶点（GSK3β、JNK3）及激酶多参数优化（Kinase MPO：集成 GSK3β、JNK3 活性分类器与 QED、合成可及性 SA）上进行评估。

方法类型	方法	GSK3β (↑)	JNK3 (↑)	Kinase MPO (↑)	训练预言机调用数	测试期单样本调用数
摊销优化	LibINVENT	\(0.410 \pm 0.115\)	\(0.260 \pm 0.071\)	\(0.455 \pm 0.006\)	约 50,000	0
摊销优化	GraphXForm	\(0.064 \pm 0.008\)	\(0.024 \pm 0.003\)	\(0.409 \pm 0.003\)	280,000	0
摊销优化	DrugEx v3	\(0.194 \pm 0.004\)	\(0.356 \pm 0.003\)	\(0.354 \pm 0.003\)	640,000	0
摊销优化	AMORTIX (本文)	\(\mathbf{0.775 \pm 0.005}\)	\(\mathbf{0.796 \pm 0.004}\)	\(\mathbf{0.619 \pm 0.004}\)	约 50,000	0
实例优化	Mol GA	\(0.735 \pm 0.014\)	\(0.659 \pm 0.011\)	\(0.461 \pm 0.013\)	0	约 500
实例优化	GenMol	\(0.700 \pm 0.012\)	\(0.470 \pm 0.015\)	\(0.479 \pm 0.012\)	0	约 333

在标准 PMO 基准（22 项从头设计任务，严格限制 10,000 次预言机预算）中，AMORTIX 作为实例优化器运行，累积 Top-10 AUC 达到 16.433，超越了所有摊销策略基线（如 f-RAG 16.301、Genetic GFN 16.078、REINVENT 15.003），在全场所有算法中位列第二（仅次于预训练规模达十亿级分子的离散扩散模型 GenMol 17.540）。

消融实验¶

在 Kinase MPO 任务下对 AMORTIX 的核心组件及不同强化学习范式进行深入消融与对比，揭示组相对归一化对稳定训练的关键作用。

配置 / 变体	Kinase MPO 总分	GSK3β 得分	JNK3 得分	QED	SA (越小越易合成)	优势信号标准差 \(\sigma_A\)
AMORTIX (GRPO完整版)	\(\mathbf{0.619 \pm 0.004}\)	\(\mathbf{0.743 \pm 0.007}\)	\(\mathbf{0.588 \pm 0.011}\)	\(0.381 \pm 0.004\)	\(3.114 \pm 0.075\)	\(\approx 0.025\) (平稳)
AMORTIX-REINFORCE (全局基线)	\(0.475 \pm 0.120\)	\(0.352 \pm 0.293\)	\(0.285 \pm 0.205\)	\(0.505 \pm 0.223\)	\(3.178 \pm 2.068\)	较 GRPO 高 1-2 个数量级 (剧烈震荡)
AMORTIX-PPO (带值网络Critic)	\(0.552 \pm 0.015\)	\(0.606 \pm 0.011\)	\(0.478 \pm 0.039\)	\(0.380 \pm 0.025\)	\(3.304 \pm 0.329\)	较 GRPO 高 1-2 个数量级 (平台期高方差)
AMORTIX-DeNovo (无骨架条件输入)	\(0.397 \pm 0.005\)	\(0.071 \pm 0.012\)	\(0.030 \pm 0.004\)	\(0.710 \pm 0.014\)	\(3.007 \pm 0.059\)	- (骨架约束坍缩)

关键发现¶

实例优化器在多目标生物活性上的妥协退化：细分指标（Appendix Table 8）显示，Mol GA 和 GenMol 在 MPO 任务中总分较低的根源在于它们倾向于“钻空子”，过度优化容易提升的启发式指标（Mol GA 的 QED 达 0.81，GenMol 的 QED 达 0.88、SA 达 2.93），但在稀疏的生物靶点流形上严重受挫（GSK3β/JNK3 均低于 0.24）。相比之下，AMORTIX 在保持合理成药性的前提下，成功突破了生物活性流形，GSK3β 达到 0.743，JNK3 达到 0.588。
摊销成本回收拐点极其显著：AMORTIX 约 50,000 次的固定预言机训练成本，在测试集处理到第 100 个骨架时便被 Mol GA 的累积开销追平，在第 150 个骨架时被 GenMol 追平。超过 150 个分子后，AMORTIX 的摊销算力优势呈现线性持续扩大。在单卡 H100 上，GenMol 对单个骨架进行 10 轮迭代优化需要约 7 分钟，而 AMORTIX 批量前向推理耗时仅为毫秒级。
少样本前药修饰规则的跨分子迁移能力：在仅使用 4 个前药范式母核（吗啡、GABA、去甲槟榔次碱、阿司匹林）微调后，面对 52 个未见过的药物分子，AMORTIX 在 31%（16个）的分子上自主生成了结构纯净的前药（保持母核骨架不变，所有新增键均为可水解的酯键或氨基甲酸酯键），63% 的分子具备至少一个可水解手性位点，平均血脑屏障穿透率（BBB）预测值较原始母药提升 +0.26，并成功在天冬氨酸、甘氨酸等分子上实现了多位点复合保护。

亮点与洞察¶

将组相对优势估计引入分子强化学习：敏锐指出了分子结构修饰中“起始骨架固有难度不均”这一常被忽视的统计学陷阱，巧妙利用组内均值中心化消除了跨骨架方差，完全省去了庞大且难以收敛的 Critic 网络。
显式价键掩码与图Transformer协同：在保持置换等变性与键级拓扑敏感度的同时，通过动态价态掩码确保了每步生成的绝对合法性，避免了传统 RL 依靠惩罚项导致策略不敢离开预训练分布的弊端。
高质量单次前向直出替代迭代采样：颠覆了扩散模型或遗传算法必须依赖多轮推理期搜索的沉重负担，为超大规模虚拟筛选和先导化合物库的高通量修饰提供了切实可行的工业级推理范式。

局限与展望¶

局限于 2D 拓扑结构：当前生成机制仅在 2D 分子图上运行，尚未显式引入三维构象动力学或立体化学空间坐标，对于高度依赖空间构象匹配的靶点口袋结合亲和力建模存在局限。
单连通子图约束与不可逆性：目前的动作掩码机制假设输入为一个单独连通的骨架，无法直接原生支持非连通子图（如两端片段连接 Linker 设计）；且动作空间仅支持增加原子和键，不支持原子删除或骨架骨折替换（Scaffold Hopping）。
训练期预言机依赖仍存：虽然推理开销完全归零，但策略微调阶段仍需数万次预言机调用，未来可探索将摊销策略与实例搜索相结合的混合范式（由摊销策略快速输出近优起点，再由局部搜索精修）。

评分¶

新颖性: ⭐⭐⭐⭐☆ 成功将 GRPO 迁移并针对性适配至结构约束分子生成，解决了骨架难度异质性的核心痛点。
实验充分度: ⭐⭐⭐⭐⭐ 涵盖标准 PMO 基准、严格 OOD 聚类切分的激酶多目标优化、消融分析以及极具化学实际意义的前药迁移案例。
写作质量: ⭐⭐⭐⭐⭐ 逻辑严密，图表详实，问题定位清晰，对各种优化范式的权衡分析十分透彻。
价值: ⭐⭐⭐⭐⭐ 彻底打通了结构约束高通量先导化合物修饰的零推理开销路径，对计算药物化学和AI制药具有重大实用价值。