Mol-MoE: Training Preference-Guided Routers for Molecule Generation
Diego Calanzone (1, 2), Pierluca D'Oro (2), Pierre-Luc Bacon (1, 2) (1) Universite de Montreal, (2) Mila Quebec AI Institute arXiv: https://arxiv.org/abs/2502.05633
Abstract: Recent advances in language models have enabled framing molecule generation as sequence modeling. However, existing approaches often rely on single-objective reinforcement learning, limiting their applicability to real-world drug design, where multiple competing properties must be optimized. Traditional multi-objective reinforcement learning (MORL) methods require costly retraining for each new objective combination, making rapid exploration of trade-offs impractical. To overcome these limitations, we introduce Mol-MoE, a mixture-of-experts (MoE) architecture that enables efficient test-time steering of molecule generation without retraining. Central to our approach is a preference-based router training objective that incentivizes the router to combine experts in a way that aligns with user-specified trade-offs. This provides improved flexibility in exploring the chemical property space at test time, facilitating rapid trade-off exploration. Benchmarking against state-of-the-art methods, we show that Mol-MoE achieves superior sample quality and steerability.
How to use this model
This LM is fine-tuned to generate molecules in the SMILES format wrt. desired properties.
For unconditioned SMILES generation, use the BOS token <s>.
For conditioned generation, you can target the following properties: JNK3, DRD2, GSK3B, CYP2D6, CYP2C19.
This model is a fine-tuned version of LLaMa 3.2 1B through two stages:
Fine-tuning on ~3.5M molecules extracted from: ZINC 250K, MOSES, CHEMBL
RLHF-tuning using instruction fine-tuning on 5 distinct reward signals.
The detailed pipeline we followed is reported in the original paper: "Rewards-in-Context: Multi-objective Alignment of Foundation Models with Dynamic Preference Adjustment" Yang et al. 2024 [1]
