Pushing Biomolecular Utility-Diversity Frontiers with Supergroup Relative Policy Optimization¶

会议: NeurIPS 2026
arXiv: 2605.08659
OpenReview: https://openreview.net/forum?id=T7f7FXqdfT
代码: https://github.com/IDEA-XL/SGRPO
领域: 计算生物
关键词: 生物分子生成, 强化学习后训练, 超组相对策略优化, 效用多样性前沿, 集合级多样性

一句话总结¶

针对生物分子生成中强化学习优化效用导致多样性坍缩的瓶颈，提出超组相对策略优化（SGRPO），通过同条件超组采样与留一法（Leave-One-Out）贡献再分配，直接将集合级多样性转化为个体候选的策略梯度优势，显著拓宽小分子与蛋白质生成的效用-多样性 Pareto 前沿。

研究背景与动机¶

在现代计算生物学与药物设计中，预训练生成模型常通过强化学习（RL）后训练来对齐特定功能目标（如下游属性评分、受体结合对接打分或蛋白质折叠稳定性）。然而，单纯最大化单分子效用奖励往往会导致模型概率质量过度集中在少数局部的结构家族或高分化学型上。这种模式坍缩在生物制药场景下极其危险：药物化学家和结构生物学家需要涵盖不同骨架、理化特性的多样化候选分子以规避下游实验毒性与成药失败风险，生成系统的真实价值取决于在保持高效用的同时所能探索的候选空间广度，即“效用-多样性 Pareto 前沿”。

现有解决多样性缺失的方法主要依赖历史样本惩罚（如基于记忆库或化学骨架过滤降权与过去样本相似的生成物）或熵正则化。但这些策略存在根本性矛盾：多样性本质上是“当前生成集合”的集合级属性（Set-level Property），而传统策略优化是作用于单次采样（Rollout-level）的更新机制。依靠历史记忆库不仅引入了间接代理偏差，还极易对真正有用的高密度合理模式产生过度惩罚，甚至导致后训练中的分布漂移；而在多受体或条件生成任务中，跨条件的历史记忆更会混淆不同设计靶点的内在难度差异。

本文的切入角度是：不再借助历史记忆库等间接代理，而是将“当前生成集合的真实多样性”直接作为一等公民奖励信号引入策略优化，并通过统计上保持无偏的集合内个体贡献重分配机制，让具体的生成候选承担相应的策略梯度信用。核心 idea：构建同条件超组（Supergroup）采样机制评估集合级多样性，并利用留一法（Leave-One-Out）边际多样性贡献将集合奖励重分配至个体样本，在 GRPO 架构下直接向外拓展生物分子的效用-多样性 Pareto 前沿。

方法详解¶

整体框架¶

SGRPO（Supergroup Relative Policy Optimization）是一个与生成器架构、单体效用函数以及多样性度量解耦的通用强化学习后训练框架。给定生成条件 $\mathcal{C}$（在从头小分子或蛋白质生成中为空，在受体口袋驱动生成中为靶点口袋坐标与残基序列），算法在当前策略下采样一个由多个候选子组构成的同条件“超组”（Supergroup），分别评估单样本效用与子组集合多样性；随后在同条件子组间计算相对多样性优势，并按留一法边际贡献将多样性信号再分配给子组内的各个候选，最后与个体效用线性组合后执行无 Critic 的组相对策略更新。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入设计条件 C<br/>口袋结构 / 序列提示"] --> B["同条件超组采样<br/>采样 M 个子组各含 K 个候选"]
    B --> C["子组级多样性评估与同条件相对对比<br/>计算组内多样性及组间相对优势"]
    C --> D["留一法多样性贡献重分配<br/>依据边际增益将组奖励分配给个体"]
    D --> E["超组相对优势计算与策略更新<br/>组合效用-多样性奖励并执行 PPO 裁剪更新"]
    E --> F["输出拓展前沿的模型<br/>高效用且富有多样性的分子/蛋白质"]

关键设计¶

1. 同条件超组采样与组相对多样性信号：彻底隔离跨条件的任务难度混淆

针对不同生物靶点或条件固有难度差异巨大、跨条件奖励不可比的问题，SGRPO 在每个条件 $\mathcal{C}$ 下独立采样 $M$ 个候选子组 $G_1, \dots, G_M$（每个子组包含 $K$ 个独立采样 rollout），构成包含 $N = MK$ 个候选的同条件超组 $\mathcal{S}(\mathcal{C})$。每个候选分子 $x_{m,i}$ 计算个体效用 $r_{m,i} = r(x_{m,i}, \mathcal{C})$，每个子组 $G_m$ 计算集合级多样性得分 $R_m = D(G_m)$（如分子指纹 Tanimoto 距离或序列 Levenshtein 距离）。为了在子组之间建立基线对比，SGRPO 采用超组内部的留一式组相对多样性优势： $$A^{\mathrm{grp}}_{m} = R_{m} - \frac{1}{M-1}\sum_{h \neq m} R_{h} = \frac{M}{M-1}(R_{m} - \bar{R})$$ 其中 $\bar{R}$ 为该超组内所有子组的平均多样性。该设计确保仅在相同输入条件下比较哪些候选组更具多样性，从理论上证明了在平衡均匀划分下平均小规模子组多样性是全样本多样性的无偏估计，避免了跨口袋或跨任务的基线漂移。

2. 留一法集合多样性符号感知重分配：将集合级信号转化为个体可行动信用

由于策略梯度的反向传播必须作用于各个具体的候选 rollout 上，必须将子组多样性标量 $R_m$ 精确归因到组内各样本。SGRPO 为子组 $G_m$ 中的每个候选 $x_{m,i}$ 计算留一法多样性边际贡献 $c_{m,i} = D(G_m) - D(G_m \setminus \{x_{m,i}\})$，即剔除该候选后子组多样性的净损失。随后在组内进行标准化得到 $z_{m,i} = (c_{m,i} - \bar{c}_m) / (\sigma(c_{m,\cdot}) + \zeta)$，并构造正负符号感知的 softmax 权重向量： $$w^{\pm}_{m,i} = K \cdot \frac{\exp(\pm z_{m,i} / \tau_c)}{\sum_{j=1}^K \exp(\pm z_{m,j} / \tau_c)}$$ 进而计算再分配后的多样性奖励： $$\widetilde{R}_{m,i} = R_m + [A^{\mathrm{grp}}_m]_+ (w^+_{m,i} - 1) - [-A^{\mathrm{grp}}_m]_+ (w^-_{m,i} - 1)$$ 这种符号感知机制保证了逻辑自洽：当子组表现出高于平均的多样性（$A^{\mathrm{grp}}_m > 0$）时，超额的正向奖励重点分配给对多样性贡献最大的样本；而当子组多样性不足（$A^{\mathrm{grp}}_m < 0$）时，惩罚重点集中在多样性贡献最小的冗余样本上。数学上严格保证了组内平均多样性奖励守恒，即 $\frac{1}{K}\sum_{i=1}^K \widetilde{R}_{m,i} = R_m$。

3. 多目标奖励组合与组内留一优势裁剪更新：端到端扩展 Pareto 前沿

在获得个体化多样性奖励后，SGRPO 将单体效用 $r_{m,i}$ 与再分配多样性奖励 $\widetilde{R}_{m,i}$ 线性加权组合为总奖励 $\widehat{r}_{m,i} = (1 - \lambda) r_{m,i} + \lambda \widetilde{R}_{m,i}$，其中 $\lambda \in [0, 1]$ 平衡效用与多样性的相对权重。随后在整个超组内计算留一基线优势： $$A_{m,i} = \frac{MK}{MK - 1} (\widehat{r}_{m,i} - \bar{r}_{\mathcal{S}})$$ 并利用 PPO 裁剪代理目标结合相对于参考策略 $\pi_{\mathrm{ref}}$ 的 KL 散度约束更新策略参数 $\theta$。对于离散扩散模型（如 GenMol / GenMol-P），SGRPO 与 Coupled-GRPO 的互补掩码采样概率对数估计相结合；对于自回归序列模型（如 ProGen2），则直接使用原生 token 似然构建重要性权重，实现了对不同生成架构的无缝兼容。

损失函数 / 训练策略¶

策略优化最小化带 KL 惩罚的裁剪替代损失： $$\mathcal{L}_{\mathrm{SGRPO}}(\theta) = -\mathbb{E}\left[\min\left(\rho_{m,i} A_{m,i}, \operatorname{clip}(\rho_{m,i}, 1-\epsilon, 1+\epsilon) A_{m,i}\right)\right] + \beta\,\mathbb{E}\left[\mathrm{KL}(\pi_{\theta}(\cdot\mid\mathcal{C}) \parallel \pi_{\mathrm{ref}}(\cdot\mid\mathcal{C}))\right]$$ 其中 $\rho_{m,i} = \pi_{\theta}(x_{m,i}\mid\mathcal{C}) / \pi_{\theta_{\mathrm{old}}}(x_{m,i}\mid\mathcal{C})$ 为重要性采样比率。在训练时，小分子实验采用 8 组各 64 个样本（共 512 rollouts），口袋分子实验采用 8 组各 24 个样本（共 192 rollouts），蛋白质实验采用 8 组各 12 个样本（共 96 rollouts），训练超参保持与标准 GRPO 对齐，确保在不增加总采样预算的前提下纯粹依靠更优的优势估计带来性能收益。

实验关键数据¶

主实验¶

实验在从头小分子生成（GenMol）、口袋条件小分子生成（GenMol-P）及从头蛋白质设计（ProGen2）三个任务上开展，通过对解码温度和随机性参数扫描绘制 Pareto 曲线，使用超体积（Hypervolume, HV ↑，基准点设为各方法观测极小值）、理想点欧氏距离（Distance to Ideal Point, DIP ↓，理想点为 $(1,1)$）和加权 Tchebycheff 亏损（R2 Indicator, R2 ↓）衡量综合前沿质量。

数据集 / 任务	评估指标	本文 (SGRPO)	之前SOTA / 强基线 (GRPO)	提升 / 优势对比
从头小分子生成 (GenMol)	HV ↑ DIP ↓ R2 ↓	0.0758 ± 0.0029 0.2551 ± 0.0020 0.0955 ± 0.0004	0.0703 ± 0.0029 0.2677 ± 0.0013 0.1017 ± 0.0005	DIP 降低 0.0126，R2 显著降低；综合指标超越 DMB-GRPO (DIP 0.2581) 与 Mem-GRPO (0.2696)
口袋驱动分子设计 (CrossDocked2020)	HV ↑ DIP ↓ R2 ↓	0.0630 ± 0.0033 0.3797 ± 0.0022 0.1807 ± 0.0007	0.0098 ± 0.0016 0.7459 ± 0.0153 0.3836 ± 0.0078	HV 暴涨 6.4 倍；DIP 降低接近一半（0.7459 → 0.3797），完全克服单靶点对接坍缩
从头蛋白质设计 (ProGen2)	HV ↑ DIP ↓ R2 ↓	0.3687 ± 0.0086 0.3538 ± 0.0048 0.1693 ± 0.0036	0.2140 ± 0.0054 0.6519 ± 0.0058 0.3023 ± 0.0034	HV 提升 72.3%，DIP 降低 0.2981，大幅超越熵正则化 ER-GRPO (HV 0.3128)

消融实验¶

在基于 GenMol 的从头小分子设计任务上对核心组件进行隔离消融（均在相同解码扫描区间下测试）：

配置	超体积 HV ↑	理想点距离 DIP ↓	说明
Full SGRPO	0.0758	0.2551	完整模型：包含超组相对多样性奖励与留一法符号感知信用重分配
w/o 留一法信用重分配 (仅超组多样性均匀分配)	0.0715	0.2640	仅将子组多样性标量直接赋给样本，无法辨析组内样本的真实多样性贡献
w/o 多样性奖励 (退化为 Coupled-GRPO)	0.0703	0.2677	完全去除集合多样性反馈，仅依靠 QED 与 SA 效用优化，Pareto 前沿最差

关键发现¶

结构坍缩最严重的靶向口袋设计中收益最大：在特定受体口袋中，对接打分对局部最佳构象具有极强正反馈，标准 GRPO、DMB-GRPO 和 ER-GRPO 均因过度追逐极高 docking 得分而造成骨架极度匮乏（DIP 高达 0.67~0.75）。SGRPO 明确奖励同一口袋下的差异化分子，将 DIP 降至 0.3797，前沿扩展极其显著。
训练轨迹保持多峰分布探索：通过 UMAP 降维追踪蛋白质生成序列空间动态发现，标准 GRPO 在 20 步到 100 步迅速收缩至单一狭窄簇；Memory-GRPO 虽强迫避开已有样本但导致整体分布漂移至不合理的边远区域；SGRPO 则在早期即探索多个高密度有潜力的模态，并在后续训练中同步保持这些模态的精细优化。
小分组规模与权重鲁棒性优异：子组大小在极小的 $K=4$ 时即可显著超过标准 GRPO 基线，无需膨胀采样总预算；在多样性权重 $\lambda \in [0.1, 0.9]$ 的广泛区间内，超体积 HV 均稳定优于纯效用优化，并在 $\lambda=0.5$ 处达到峰值。

亮点与洞察¶

将无形的集合属性精准拆解为个体梯度：利用留一法计算剔除单个样本后的多样性落差，并配合正负符号敏感的加权分配，在严格保持期望守恒的同时优雅地赋予每个生成候选正确的策略梯度，是强化学习优化非加性集合指标的标准范式。
同条件比较剔除任务内在方差：生物分子条件生成中，不同受体口袋或设计提示的结合难度天差地别，SGRPO 坚持超组内相对归一化，杜绝了把“任务本身简单”误判为“生成策略优异”的偏差。
计算代价的高度复用：子组内部计算留一法贡献时完全复用成对相似度矩阵（Pairwise Similarity Matrix），使得留一法边际计算的时间复杂度从重新计算的 $O(K^3)$ 骤降至矩阵局部求和的 $O(K^2)$，工程落地极其轻量。

局限与展望¶

计算复杂度随子组大小二次增长：虽然可复用相似度矩阵，但在需要极大子组规模（如 $K > 100$）以精确刻画超复杂组合空间时，计算所有样本对的生物指纹或序列编辑距离开销依然不可忽略。
多目标多样性融合的权衡：目前各任务的多样性度量主要采用单一指标（小分子采用 Morgan-Tanimoto 距离，蛋白质采用 Levenshtein 距离），未能同时覆盖三维形状多样性、骨架跃迁多样性或电荷表面多样性等更复杂的物理化学特征。
未来方向：探索低秩核近似或层次聚类以加速大规模子组的多样性评估，并将超组相对优化机制拓展至多智能体协同设计及更多三维结构生成任务中。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 首次提出直接面向集合级多样性的超组相对策略优化与留一法符号感知信用重分配机制。
实验充分度: ⭐⭐⭐⭐⭐ 涵盖从头小分子、靶向口袋设计和从头蛋白质三大代表性生物场景，贯穿离散扩散与自回归模型，评估严谨全面。
写作质量: ⭐⭐⭐⭐⭐ 理论推导清晰完备，算法逻辑与生物化学问题结合紧密，图表信息量大且自洽。
价值: ⭐⭐⭐⭐⭐ 为计算生物学与药物设计生成模型的 RL 后训练提供了极具通用性与即插即用价值的标准优化范式。