Q-Steer: Action-Value Guidance for Molecular Policy Optimization

Q-Steer introduces a rollout-time action-value steering primitive for molecular policy optimization, designed to address the delayed-feedback problem in oracle-limited molecular generation where reward is only given after a complete molecule is generated. The method uses an offline-trained and frozen prefix-action value scorer called PAVS-Q, which estimates the downstream reward of taking a candidate next token under a partial SMILES prefix. During sampling, Q-Steer adds a normalized value bonus to the sampling logits: l̃t(a) = lt(a) + βQϕ(x<t, a), where Qϕ denotes the centered and scaled PAVS-Q score over valid actions at the current prefix, and β is a steering strength. The optimizer update rule and online oracle budget remain unchanged; the claim is fixed-online-oracle performance, not equal total compute.

PAVS-Q is the action-resolved counterpart of a prefix future-value prior. While a prefix-level value model Vϕ(x<t) estimates E[R(x)x<t] and ranks prefixes, PAVS-Q estimates Qϕ(x<t, a) ≈ E[R(x)x<t, xt = a], which ranks the next actions available from the same prefix. PAVS-Q is trained offline from scored molecular rollouts by creating supervised prefix-action examples DQ = (x<t, xt, y(x)): t = 1,..., T, where y(x) is the standardized terminal reward, and training by mean-squared error. The implementation uses a one-layer GRU prefix encoder, an action embedding, and an MLP head, with the tokenizer aligned with the online ACEGEN policy. The value model is frozen before online optimization begins and receives no molecules, rewards, or gradients from the evaluated online trajectory.

The Q-Steer method normalizes raw values into a local bonus: Q̄t(a) = (Qϕ(x<t, a) − µt)/(σt + ϵ), where µt is the mean and σt the standard deviation over valid actions, then forms a bounded value bonus bt(a) = clip(βQ̄t(a), −c, c), and adds it to base logits with invalid actions masked. In main experiments, β = 0.5 and c = 5.0. The completed molecule is scored once by the PMO oracle and handed back to the underlying optimizer exactly as in the baseline; for PPO, REINVENT, AHC, and hill-climbing, the optimizer-specific update rule is unchanged. The paper frames this as a practical finite-vocabulary value transform related to Doob h-transforms, but notes it should not be read as a theorem-level Doob transform since PAVS-Q is learned, task-specific, and used with normalization and clipping.

The evaluation uses PMO23 with a fixed 10,000-call online oracle budget per run, across two molecular language-model backbones (LSTM and GPT2) and four optimizer families (PPO, REINVENT, AHC, hill-climbing). The main factorial design contains 4 × 2 × 23 × 3 = 552 successful formal runs for LSTM and another 552 for GPT2, totaling 1104 successful formal PMO23 runs. Metrics include mean score, top-10 mean score, cumulative top-10 AUC, best score, validity, uniqueness, and runtime, with score summaries computed over valid unique molecules.

The main result shows that all eight backbone–optimizer cells have positive mean valid-unique reward deltas. In the controlled LSTM study, Q-Steer improves mean score on 78/92 optimizer–task comparisons: 19/23 tasks for PPO and hill-climbing, and 20/23 tasks for REINVENT and AHC. The LSTM macro mean-score gains range from +0.0343 to +0.0490. The matched GPT2 replication adds another 77/92 task wins, with positive macro gains for every optimizer family ranging from +0.0330 to +0.0440. Task-bootstrap confidence intervals for ∆Mean remain positive for all eight cells, and exact sign tests favor Q-Steer at p ≤ 0.001.

The mechanism controls are central to the story. Table 2 compares true Q-Steer to prefix-broadcast and shuffled-action controls in the LSTM PPO setting. True action-specific PAVS-Q improves mean score on 19/23 tasks with macro ∆Mean = +0.0343. Prefix-broadcast values, which preserve prefix quality but remove action discrimination, are nearly neutral: 12/23 wins and ∆Mean = +0.0022. Shuffled-action values, which preserve value magnitudes but assign them to the wrong tokens, hurt performance: 3/23 wins and ∆Mean = −0.0459. This ordering is qualitative: correct action values help, state-level broadcast values are almost inert, and wrong action assignments actively harm. A focused beta/shuffling diagnostic on six tasks with 1,500 online oracle calls per run shows true action-value steering wins 5/6 tasks at each tested strength, with macro mean gain increasing from +0.0586 at β = 0.25 to about +0.10 at β = 0.5 and 1.0, while shuffled action values lose 0/6 tasks and reduce mean score by −0.0624.

The consistent reward gain comes with a consistent diversity cost. Uniqueness decreases for every optimizer family and backbone. In LSTM, the macro uniqueness deltas are-0.2101 for PPO, -0.0176 for REINVENT, -0.0164 for AHC, and-0.0927 for hill-climbing; in GPT2, the corresponding deltas are smaller but still negative. The paper interprets Q-Steer as an exploitation knob, not a diversity-preserving exploration mechanism. Upper-tail metrics are mixed: top-10 wins are mixed for all families, best-score wins are weaker still, and cumulative top-10 AUC improves clearly for GPT2 rows and LSTM non-PPO optimizers but is approximately neutral for LSTM PPO.

The paper also reports runtime and practical cost considerations. Q-Steer adds model-side computation because candidate actions are scored by PAVS-Q during rollout. Wall-clock comparisons are implementation-dependent: PPO steering is slower in the reported runs, while some non-PPO guided runs are faster due to different accelerated infrastructure. The scientific comparison is fixed-online-oracle performance, not equal wall-clock or equal offline-pretraining cost. The offline PAVS-Q training cost is reported separately from the online PMO budget, analogous to using a pretrained molecular policy prior.

The paper concludes that Q-Steer turns delayed molecular rewards into local action guidance, improving average valid-unique score without changing the optimizer update or online oracle budget. It supports action-value steering as a reusable primitive for oracle-limited molecular policy optimization and motivates future work on combining it with exploration-preserving mechanisms such as novelty-aware penalties or adaptive β schedules. The limitations section notes that LSTM and GPT2 backbones are evaluated but not broad transformer scaling, uniqueness decreases across all settings, the method adds model-side computation, top-10 and best-score improvements are mixed, and mechanism controls are run only in LSTM PPO.

Improvements to AI Systems:

1. Add a frozen offline-trained action-value scorer (PAVS-Q) to any autoregressive generative policy to convert sparse, delayed terminal rewards into dense, per-token action guidance during sampling, without modifying the underlying optimizer or online reward budget. The improved system can generate molecules with higher average valid-unique reward (e.g., +0.034 to +0.049 mean score on PMO23) across diverse policy families (PPO, REINVENT, AHC, hill-climbing) and backbones (LSTM, GPT2).

2. Implement a normalized, clipped value-bonus steering mechanism (l̃t(a) = lt(a) + β·clip((Qϕ(x<t,a) − µt)/(σt + ε), −c, c)) to inject action-specific future-reward estimates into logits. The improved system can steer generation toward high-reward completions while preserving the base policy’s update rule, achieving positive mean-reward deltas on 155/184 optimizer–task comparisons across two backbones, with task-bootstrap confidence intervals excluding zero for all eight backbone–optimizer cells.

3. Use action-resolved value estimates instead of prefix-level or shuffled values to avoid inert or harmful guidance. The improved system can outperform prefix-broadcast (state-level) values (19/23 vs. 12/23 task wins) and avoid the degradation caused by shuffled-action values (3/23 wins, negative mean delta), ensuring that the steering signal is causally tied to the specific next token.

4. Enable exploitation-focused tuning via a single steering-strength hyperparameter (β) with bounded bonus (c = 5.0). The improved system can trade off exploration vs. exploitation: increasing β from 0.25 to 0.5–1.0 raises mean-score gains from +0.059 to +0.10 on focused tasks, while shuffled controls lose at all strengths—allowing practitioners to dial in reward maximization without retraining the policy or value model.

5. Provide a plug-and-play primitive for oracle-limited generation (e.g., drug discovery, material design) where each full sequence is expensive to evaluate. The improved system can reuse a single offline-trained PAVS-Q (trained once on scored rollouts) across multiple online optimizers and backbones, improving mean valid-unique scores without increasing the 10,000-call oracle budget, and without any online gradients or reward feedback to the value model.

6. Add a diversity-cost warning and mitigation pathway: since uniqueness decreases across all settings (e.g., −0.21 for LSTM PPO), the improved system can be combined with exploration-preserving mechanisms (novelty penalties, adaptive β schedules) to recover diversity while retaining the mean-score gains—enabling balanced optimization for tasks where both reward and structural novelty matter.

7. Enable faster convergence in non-PPO settings: because the value bonus is added at rollout time and the optimizer update is unchanged, the improved system can be integrated into hill-climbing or AHC pipelines with minimal code changes, and in some configurations even reduce wall-clock time due to accelerated infrastructure, while still improving final reward quality.