Active-GRPO: Adaptive Imitation and Self-Improving Reasoning for Molecular Optimization

Recent developments in large language models (LLMs) have demonstrated their potential in multi-step reasoning tasks, including mathematics, coding, and scientific analysis. Techniques such as chain-of-thought prompting, supervised fine-tuning on reasoning traces, and reinforcement learning with verifiable rewards have been employed to enhance reasoning capabilities. In scientific domains like chemistry and biology, LLMs are increasingly used for hypothesis generation, experimental design, and candidate screening. However, these tasks involve complex, constrained, multi-objective reasoning over structured chemical representations, which pose significant challenges. Existing training paradigms, including answer-only supervised fine-tuning and reward-based methods, often suffer from issues like reasoning collapse and sparse feedback, limiting their effectiveness. Reference-guided policy optimization (RePO) was introduced to address these issues by anchoring policy updates to dataset references, blending imitation and reward signals. Nonetheless, RePO’s reliance on static references makes it vulnerable to reference quality issues, creating performance ceilings that hinder further progress. This context motivates the development of more adaptive, self-improving training frameworks capable of overcoming the limitations of fixed references and sparse rewards.

The core challenge in instruction-conditioned molecular optimization is achieving robust, multi-objective reasoning that can adapt to varying reference quality and optimization potential. Traditional methods like supervised fine-tuning lack explicit reasoning and generalization, while reinforcement learning faces sparse reward problems, especially under strict similarity constraints. RePO mitigates some issues but is fundamentally limited by its static references, which can become stale or misaligned, capping the model’s performance. This static reference ceiling prevents models from surpassing the initial dataset’s quality, especially when references are noisy or suboptimal. The key problem is designing a training paradigm that dynamically adjusts its guidance based on the model’s evolving capabilities, allowing it to learn beyond the initial references and adapt to the quality and potential of each instance. Addressing this problem is crucial for advancing autonomous scientific reasoning and molecular design, enabling models to explore larger chemical spaces more effectively and efficiently.

The main innovation is the introduction of Active-GRPO, which embeds active decision-making into reference-guided policy optimization. It features two core mechanisms: (1) active imitate-reinforce, which adaptively blends imitation and reinforcement learning based on relative performance, allowing the model to decide when to imitate a reference and when to reinforce its own discoveries; (2) active referencing, which maintains a memory bank of the best generated molecules and replaces the static reference with the top-performing candidate, effectively upgrading the guidance target. This approach enables the model to break performance ceilings imposed by static references, adapt to varying reference quality, and learn beyond initial dataset constraints. The method extends PPO variants with reward normalization, KL penalties, and dynamic guidance weights, providing theoretical robustness and practical effectiveness for multi-objective molecular optimization.

• �� Input: task-specific instruction, original molecule, dataset reference molecule. • During training: generate candidate molecules and evaluate reward R. • Active imitate-reinforce: compare candidate performance to reference, adjust guidance weight β_{guide} accordingly—high when model lags, low when surpassing. • Active referencing: maintain a memory bank B(ci) per instance, initialized with dataset reference, and update it with higher-reward candidates surpassing a margin δ. • Replace reference with best candidate in bank, ensuring guidance remains informative. • Loss function combines PPO-based policy gradient with weighted imitation guidance. • The process iterates, dynamically balancing exploration and exploitation, with parameters tuned for stability and performance.

Experiments conducted on TOMG-Bench MOLOPT, focusing on properties LogP, MR, and QED. Baselines include RePO, GRPO, and variants like iterative SFT and offline-strengthened RePO. Hyperparameters such as guidance weight bounds (β_{min}, β_{max}), memory size, and promotion thresholds are tuned. Evaluation metrics include success rate (SR), Tanimoto similarity (Sim), and SR×Sim. Multiple seeds ensure statistical robustness. Ablation studies isolate the contributions of active mechanisms. Experiments test across different reference qualities and optimization headroom scenarios, assessing robustness and adaptability.

Active-GRPO consistently outperforms baselines, achieving an average SR×Sim of 0.1773, surpassing RePO (0.1665) and GRPO (0.0959). Improvements are statistically significant across all sub-tasks, with notable gains in high-headroom instances where static references are less informative. Ablation results confirm that both active mechanisms are essential; removing either reduces performance. The method maintains high success rates while balancing property improvements and structural similarity, demonstrating robustness across diverse scenarios. These results validate the effectiveness of adaptive, self-improving guidance in molecular optimization.

The framework is applicable to drug discovery, materials science, and any domain requiring multi-objective molecular design. It enables automated, efficient exploration of chemical space with minimal manual intervention. The method can be integrated into existing molecular generative pipelines, providing robust guidance even with noisy or low-quality references. Future integration with synthesis planning and multi-modal data can further enhance its utility, accelerating scientific innovation and industrial R&D.

Computational overhead due to maintaining and updating memory banks limits scalability. Sensitivity to hyperparameters like guidance weight bounds and promotion thresholds requires careful tuning. Performance may degrade in extremely high-dimensional chemical spaces or with highly noisy references. Extending to more complex property sets and real-world constraints remains challenging. Future work should focus on improving scalability, automating hyperparameter selection, and broadening applicability to diverse scientific tasks.