Mol-MoE: Training Preference-Guided Routers for Molecule Generation

Diego Calanzone, Pierluca D'Oro, Pierre-Luc Bacon

Track: Machine learning: computational method and/or computational results

Nature Biotechnology: Yes

Keywords: drug-design, multi-objective, molecule, sequence models, computational biology, language alignment, transformers

TL;DR: We frame drug-design as multi-objective language model alignment, we evaluate existing methods and we then introduce Mol-MoE, exhibiting superior performance in sample quality and steerability.

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.

Anonymization: This submission has been anonymized for double-blind review via the removal of identifying information such as names, affiliations, and identifying URLs.

Presenter: ~Diego_Calanzone1

Format: Yes, the presenting author will attend in person if this work is accepted to the workshop.

Funding: Yes, the presenting author of this submission falls under ICLR’s funding aims, and funding would significantly impact their ability to attend the workshop in person.

Submission Number: 43