Aligning Chemical and Protein Language Models with Continuous Feedback using Energy Rank Alignment

Published: 06 Mar 2025, Last Modified: 26 Apr 2025GEMEveryoneRevisionsBibTeXCC BY 4.0
Track: Machine learning: computational method and/or computational results
Nature Biotechnology: Yes
Keywords: alignment, molecular transformer, protein language model, drug discovery, directed evolution
TL;DR: We develop a highly scalable algorithm for alignment of autoregressive generative models that does not require reinforcement learning and performs well for a variety of molecular optimization tasks.
Abstract: Large, autoregressive models trained on databases of chemical compounds and biomolecules have yielded powerful generators, but we still lack robust strategies for controlled generation. This molecular search problem closely resembles the ``alignment'' problem for large language models, though for many chemical tasks we have a specific and easily evaluable reward function. Here, we introduce an algorithm called energy rank alignment (ERA) that leverages an explicit reward function to produce a gradient-based objective that we use to optimize autoregressive policies. We deploy this approach to align molecular transformers and protein language models to generate molecules and protein sequences, respectively, with externally specified properties and find that it does so robustly, searching through diverse parts of chemical space. The algorithm is highly scalable, does not require reinforcement learning, and performs well relative to DPO when the number of preference observations per pairing is small.
Anonymization: This submission has been anonymized for double-blind review via the removal of identifying information such as names, affiliations, and identifying URLs.
Presenter: ~Sebastian_Ibarraran1
Format: Yes, the presenting author will attend in person if this work is accepted to the workshop.
Funding: No, the presenting author of this submission does *not* fall under ICLR’s funding aims, or has sufficient alternate funding.
Submission Number: 4
Loading