Q-Steer：分子策略优化的动作价值指导

# Q-Steer：面向分子策略优化的动作价值引导

来源：https://arxiv.org/html/2607.26391
Xinyu Wang¹ Jinbo Bi¹ Minghu Song² ¹康涅狄格大学计算机科学与工程系 ²合肥综合性国家科学中心健康与医学研究所 ¹美国康涅狄格州斯托斯，邮编 06269 ²中国合肥 230601

###### 摘要

Oracle 受限的分子优化仅在生成完整分子后给予奖励，而每次 rollout 需要许多局部下一个 token 决策。这种延迟反馈接口使得分子策略优化变得短视：优化器可以知道某个分子是好的，却不知道是哪些中间动作使其变好。我们引入了 Q-Steer，一种用于分子语言模型的 rollout 时动作价值引导原语。Q-Steer 使用一个离线训练并冻结的前缀-动作价值评分器 PAVS-Q，该评分器在给定部分 SMILES 前缀下估计采取候选下一个 token 的下游奖励，然后将归一化的价值奖励添加到采样 logits 中。优化器的更新规则和在线 oracle 预算保持不变；其声称的是固定在线 oracle 性能下的提升，而非总计算量相等。在 PMO23 上，使用固定 10,000 次调用的在线预算，跨越两个分子语言模型主干和四个优化器的完整因子研究表明，Q-Steer 在所有八个主干-优化器组合中提高了平均有效唯一分子分数，宏观平均分数增益在 +0.033 到 +0.049 之间，每个组合获得 18-20 个任务胜场。机制控制实验表明，动作身份信息至关重要：前缀广播价值几乎呈中性，而打乱的动作价值会损害性能。这些结果支持 Q-Steer 作为一种可重用的 rollout 时动作价值包装器，在不改变在线 oracle 预算的情况下，跨优化器家族和策略主干提高平均分子优化奖励。

## 1 引言

目标导向的分子生成通常是在严格的 oracle 预算下进行评估的：算法只能查询属性 oracle 有限次数，而每次查询都需要一个完整的有效分子。这使得分子策略优化成为一个延迟反馈序列问题。一个 SMILES 策略必须一次选择一个 token，但驱动学习的奖励仅在终止后到达。像 REINVENT 和 PPO 这样的强化学习方法可以优化这种接口（Olivecrona 等人，2017（https://arxiv.org/html/2607.26391#bib.bib2）；Blaschke 等人，2020（https://arxiv.org/html/2607.26391#bib.bib3）；Schulman 等人，2017（https://arxiv.org/html/2607.26391#bib.bib1）），但它们的 rollout 策略仍然是局部短视的：在某个前缀处，生成器通常缺乏关于哪个下一个 token 会导致高分分子的动作级信息。

改进这一设置的常规方法是改变生成器、改变优化器或以不同方式使用 oracle 预算。我们研究了一种更模块化的替代方案：保持分子优化器和在线 oracle 预算固定，但使每次 rollout 的动作不那么短视。如果一个冻结的辅助模型能够估计每个候选下一个 token 的未来奖励，那么优化器可以在局部获取前瞻性指导，同时在对完整分子评分后仍像以前一样进行更新。

我们提出 Q-Steer，一种用于分子策略优化的 rollout 时动作价值引导原语。其引导器是 PAVS-Q，一个前缀-动作价值评分器。它是 PAVS 风格前缀未来价值先验的动作对应版本：它不仅估计部分分子是否有前景，还从该前缀估计哪个下一个 token 是有前景的。给定一个前缀 \(x_{<t}\) 和一个候选 token \(a \in \mathcal{V}\)，PAVS-Q 估计采取动作 \(a\) 所得的下游分子奖励。在采样过程中，Q-Steer 通过归一化的价值奖励移动优化器的动作 logits：
\[
\tilde{\ell}_t(a) = \ell_t(a) + \beta Q_\phi(x_{<t}, a),
\]
其中 \(\ell_t(a)\) 是原始采样 logits，\(Q_\phi\) 表示在当前前缀有效动作上居中并缩放的 PAVS-Q 分数，\(\beta\) 是引导强度。干预刻意很小。在每对比较中，策略主干、优化器更新规则、PMO oracle 和 10,000 次调用的在线预算都保持固定；只有 rollout 分布收到动作级的未来价值引导。

该设计刻意将主张与更广泛的架构或优化器改变分开。Q-Steer 不是事后重排序器：它在每个 token 采样之前发挥作用。它不是新的 PPO 损失或新的 REINVENT 目标：完整的分子和奖励仍返回给相同的优化器更新。它也不旨在作为多样性机制。所支持的主张更窄且可测试：离线价值先验可以在相同的在线 oracle 预算下提高有效唯一分子的平均分数。这不是一个总计算量相等的声称；离线 PAVS-Q 训练另行披露，冻结的引导器不会从评估的在线轨迹中接收任何分子、奖励或梯度。

我们在两个分子语言模型主干和四个优化器家族上使用完整的 PMO23 因子研究来检验这一主张。在相同的 10,000 次调用在线预算下，添加 Q-Steer 在所有八个主干-优化器组合中提高了平均有效唯一分子分数，每个组合获得 18-20 个任务胜场。机制控制实验完成了闭环：前缀广播价值几乎呈中性，打乱的动作价值损害性能，而 beta 扫描显示了观察到的奖励-多样性权衡。由此得出的结论并非 Q-Steer 普遍改善每个指标。相反，它是一个可重用的利用原语：它使用动作级未来值集中搜索，在降低唯一性的同时提高平均奖励，并将 top-k 发现结果保留为混合状态。

参见图注

图 1：概念概述。基线分子优化器在采样完整的 SMILES 字符串后收到终端 oracle 反馈，使得前缀动作局部短视。Q-Steer 仅改变 rollout 分布：相同的优化器对相同的有效动作集提出 logits，一个离线训练的冻结 PAVS-Q 引导器对这些候选动作进行评分，Q-Steer 在完整分子由 PMO oracle 评分一次之前添加方程 8-9 中的受限 logit 奖励。

#### 贡献。

我们做出三项贡献。首先，我们将 Q-Steer 表述为一个 rollout 时动作价值包装器：一个冻结的 PAVS-Q 引导器对候选下一个 token 进行评分并倾斜采样 logits，而不改变优化器更新或在线 oracle 预算。其次，我们在一个跨越四个优化器家族和两个主干的配对 PMO23 因子研究中评估该包装器，展示了在所有八个主干-优化器组合中一致的平均有效唯一奖励增益。第三，我们提供了机制和局限性证据：动作身份很重要，打乱的价值有害，更强的引导以奖励换取唯一性，而上尾发现结果是混合的而非普遍改善。

## 2 相关工作

#### 用于分子生成的强化学习。

分子 RL 方法优化一个生成器对抗属性 oracle，通常使用 SMILES 语言模型作为策略（Olivecrona 等人，2017（https://arxiv.org/html/2607.26391#bib.bib2）；Blaschke 等人，2020（https://arxiv.org/html/2607.26391#bib.bib3）；Gottipati 等人，2020（https://arxiv.org/html/2607.26391#bib.bib5））。REINVENT 风格的方法在提高高分分子可能性的同时，将智能体正则化到先验，而 PPO 风格的方法则在生成的轨迹上优化一个裁剪的策略梯度目标（Schulman 等人，2017（https://arxiv.org/html/2607.26391#bib.bib1））。增量爬山和相关的似然比更新提供了更简单的利用导向替代方案，反复将生成偏向于高分样本（Blaschke 等人，2020（https://arxiv.org/html/2607.26391#bib.bib3））。这些方法之所以灵活，是因为 oracle 可以编码相似性、QSAR 活性、MPO 目标或结构约束。它们共同的局限性是延迟反馈：奖励附加于一个完整分子，而非产生该分子的局部 token 动作。Q-Steer 旨在解决这一局部决策问题，同时保持优化器固定。

#### 价值引导与信用分配。

价值函数是强化学习中延迟奖励的标准解决方案，但使用方式很重要。状态价值可以判断前缀是否有前景，但它不能区分来自该前缀的下一个 token 动作。终端奖励模型或重排序器可以对完整分子进行评分，但它在采样过程中不会干预。Q-Steer 使用一个前缀-动作价值，这更接近于分子级别的动作价值引导对应物：它估计在特定部分分子下哪个下一个 token 是有前景的。包含前缀广播和打乱动作的控制实验正是为了直接验证这一区别。

#### 解码时引导。

诸如 PPLM、GeDi、DExperts 和 FUDGE 等语言模型引导方法使用辅助模型或判别器在生成时修改 token 概率（Dathathri 等人，2020（https://arxiv.org/html/2607.26391#bib.bib7）；Keskar 等人，2019（https://arxiv.org/html/2607.26391#bib.bib8）；Krause 等人，2021（https://arxiv.org/html/2607.26391#bib.bib9）；Liu 等人，2021（https://arxiv.org/html/2607.26391#bib.bib10）；Yang 和 Klein，2021（https://arxiv.org/html/2607.26391#bib.bib11））。Q-Steer 共享了在不重训基座生成器的情况下改变采样分布的想法，但在目标和信号上有所不同：引导器是一个任务特定的分子前缀-动作价值模型，受控变量是在固定 oracle 预算下的在线 PMO 奖励，而非文本属性。与通用属性引导不同，引导器根据预测的下游 PMO 价值对候选分子动作进行排序，并通过配对的优化器运行而非无条件生成质量进行测试。这种联系激励了将 \(\beta\) 视为一个引导强度参数而非一个新的优化器目标。

#### Oracle 预算基准。

Oracle 受限的分子优化基准强调固定预算，使得不仅报告最终最佳分子还报告分布质量和搜索动态变得重要（Brown 等人，2019（https://arxiv.org/html/2607.26391#bib.bib4）；Gao 等人，2022（https://arxiv.org/html/2607.26391#bib.bib6））。平均分数、top-10 奖励、累积 top-10 AUC 和唯一性衡量了优化的不同方面。一种方法可能提高平均奖励，同时降低多样性，或未能改善早期 top-k 发现。我们的评估遵循这一更广泛的视角，而非优化单一指标。

## 3 方法

### 3.1 延迟奖励分子优化

令 \(\pi_\theta\) 为一个分子语言模型策略，作用于 token 序列 \(x = (x_1, \ldots, x_T)\)，每个 token 从受 SMILES 有效性掩码约束的有限词汇表中采样。一个黑盒 oracle 仅在完整分子生成并验证后返回一个标量奖励 \(R(x)\)。在在线 oracle 预算 \(B\) 下，一个分子优化器引发一系列 rollout 分布并接收最多 \(B\) 个终端分数。其通用目标可写为
\[
\max_\theta \; \mathbb{E}_{x \sim \pi_\theta} [R(x)],
\]
但策略仍然必须在每个前缀 \(x_{<t}\) 处选择局部动作 \(a = x_t\)，然后才能观察到 \(R(x)\)。这种不匹配正是 Q-Steer 所针对的信用分配瓶颈。我们不替换更新 \(\pi_\theta\) 的优化器；相反，我们只修改用于收集完整分子的采样分布。

### 3.2 从 PAVS 到 PAVS-Q

对于延迟反馈生成，有用的对象是部分序列上的未来价值先验。一个前缀级价值模型（我们称之为 PAVS）估计
\[
V_\phi(x_{<t}) \approx \mathbb{E}[R(x) \mid x_{<t}].
\]
这对于判断部分分子是否有前景是有信息量的，但它本身并不能告诉生成器选择哪个下一个 token。如果将相同的标量 \(V_\phi(x_{<t})\) 添加到每个有效动作的 logit 中，那么该前缀处动作的相对概率不变。

PAVS-Q 是该先验的动作解析版本。它估计在特定前缀下采取特定候选动作的价值：
\[
Q_\phi(x_{<t}, a) \approx \mathbb{E}[R(x) \mid x_{<t}, x_t = a].
\]
因此，PAVS 提供了未来价值框架，PAVS-Q 提供了 token 级动作区分，而 Q-Steer 是使用这些动作价值来引导采样的 rollout 时包装器。区别很重要：\(V(x_{<t})\) 对前缀进行排序，而 \(Q(x_{<t}, a)\) 对来自同一前缀的下一个可用动作进行排序。

### 3.3 训练前缀-动作价值模型

PAVS-Q 是从已评分分子 rollout 中离线训练的。给定一个评分分子 \(x = (x_1, \ldots, x_T)\) 及其终端奖励 \(R(x)\)，我们创建监督式前缀-动作示例
\[
\mathcal{D}_Q = \{(x_{<t}, x_t, y(x)) : t = 1, \ldots, T\},
\]
其中 \(y(x)\) 是 PMO 任务的标准终端奖励。价值模型通过均方误差训练，
\[
\min_\phi \; \mathbb{E}_{(x_{<t}, a, y) \sim \mathcal{D}_Q} (Q_\phi(x_{<t}, a) - y)^2.
\]
在我们的实现中，PAVS-Q 使用一个单层 GRU 前缀编码器、一个动作嵌入和一个 MLP 头。检查点通过保留的前缀-动作验证损失选择，然后在在线优化开始前冻结。分词器和动作词汇与在线 ACEGEN 策略对齐，因此由 PAVS-Q 评分的动作 ID 与优化器采样的动作 ID 相同。这种对齐是必要的：基于不同分词化的价值模型将无法定义有效的动作级 logit 奖励。

对于多组件 PMO 目标，PAVS-Q 可以在应用任务特定目标后估计标量奖励，或估计根据 PMO 目标组合的奖励组件。实验使用任务特定的冻结检查点。这种离线价值模型训练是一个额外的先验构建成本，类似于使用预训练的分子策略先验。它不计入 10,000 次调用的在线 PMO 预算，并且评估的在线轨迹中生成的分子不会用于更新 PAVS-Q。

### 3.4 Q-Steer：Rollout 时动作价值引导

在前缀 \(x_{<t}\) 处，基座优化器在有效下一个 token 动作 \(a \in \mathcal{A}(x_{<t})\) 上产生 logits \(\ell_t(a)\)。Q-Steer 对相同的候选动作集评估 PAVS-Q，并将原始值转换为局部归一化奖励，
\[
\bar{Q}_t(a) = \frac{Q_\phi(x_{<t}, a) - \mu_t}{\sigma_t + \epsilon}, \quad \mu_t = \frac{1}{|\mathcal{A}_t|} \sum_{a \in \mathcal{A}_t} Q_\phi(x_{<t}, a),
\]
其中 \(\sigma_t\) 是有效动作上的标准差，\(\mathcal{A}_t = \mathcal{A}(x_{<t})\)。然后我们形成一个有界价值奖励，
\[
b_t(a) = \text{clip}(\beta \bar{Q}_t(a), -c, c),
\]