Elite-Weighted Supervised Fine-tuning for Goal-Directed Molecular Optimization
cs.LG
Submitted: 2026-08-31
Updated: 2026-08-31
Code: https://github.com/NVIDIA-BioNeMo/genmol
License: http://creativecommons.org/licenses/by/4.0/
The gist: Goal-directed optimization is essential for steering molecular generators to propose candidates with desired properties.
Terminology
Abstract
Goal-directed optimization is essential for steering molecular generators to propose candidates with desired properties. However, it is often implemented with policy-gradient reinforcement learning, which requires a generation-trajectory log-probability whose form depends on the model architecture and generation procedure. This makes an optimizer difficult to reuse across architectures and conditional generative designs. Supervised fine-tuning needs none of that machinery, but its update is driven by a fixed dataset, so the reward never enters the update. We introduce Elite-Weighted Supervised Fine-tuning (EW-SFT), which uses reward to guide elite selection of high-scoring molecules, and updates the model by its own pretraining loss on that set. Ablations show that reward information is passed primarily through elite selection, rather than through continuous weighting within the selected set. Because the update consumes only scored molecules and the model's native loss, the same rule applies across autoregressive, masked-diffusion, and discrete-flow generators, and across de novo, motif-extension, and linker-design tasks. Under a fixed budget of 3D shape alignment oracle calls on two kinase reference compounds, EW-SFT consistently outperforms the corresponding native optimizers. It further improves goal-directed optimization under a 2D similarity oracle on four held-out references and achieves comparable performance on a sample-efficiency benchmark without a trajectory-level RL formulation. These results demonstrate that EW-SFT is a unified and effective optimizer across molecular generators, design constraints, references, and oracles.
Sources
- Guiding Deep Molecular Optimization with Genetic Exploration
- What Matters In On-Policy Reinforcement Learning? A Large-Scale Empirical Study
- Structured Denoising Diffusion Models in Discrete State-Spaces
- Generative Flows on Discrete State-Spaces: Enabling Multimodal Flows with Applications to Protein Co-Design
- Implementation Matters in Deep Policy Gradients: A Case Study on PPO and TRPO
- Learning to Extend Molecular Scaffolds with Structural Motifs
- Simple and Effective Masked Diffusion Language Models
- Proximal Policy Optimization Algorithms
- Generating Focussed Molecule Libraries for Drug Discovery with Recurrent Neural Networks
- Genetic algorithms are strong baselines for molecule generation
- Advancing Ligand-based Virtual Screening and Molecular Generation with Pretrained Molecular Embedding Distance
- Hit and Lead Discovery with Explorative RL and Fragment-based Molecule Generation
- Graph-GRPO: Training Graph Flow Models with Reinforcement Learning
Related papers
- Polynomial-Augmented Neural Networks (PANNs) with Weak Orthogonality Constraints for Enhanced Function and PDE Approximation
- AIRL-S: Unifying Reinforcement Learning and Search-Based Test-Time Scaling via Adversarial Inverse Reinforcement Learning
- Transformers as Bayesian In-Context Experimenters: Smoothness-Adaptive Efficient ATE Estimation
- Convergence issues in Relational Concept Analysis based on AOC-posets
- Beliefs Beyond Posteriors: The Predictive Cost of ELBO Binding
- Understanding Diffusion Models via Ratio-Based Function Approximation with SignReLU Networks