# Aligning Transformers with Continuous Feedback via Energy Rank Alignment

arXiv: https://arxiv.org/abs/2405.12961 · HTML: https://mlsanity.com/paper/2405.12961/

## Key facts

- Cited by 7 papers in the map (data through 14 Oct 2026).
- Placed in Preference optimization for molecular design, one of 10 directions in LLM alignment / RLHF / preference optimisation.
- Builds on 20 of its 53 references that are in the map.
- First posted on arXiv 21 May 2024.

## Placed in

- LLM alignment / RLHF / preference optimisation / [Preference optimization for molecular design](https://mlsanity.com/alignment/preference-optimization-for-molecular-design/)

## Builds on (most cited first, 20 of 20)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Training language models to follow instructions with human feedback](https://mlsanity.com/paper/2203.02155/) (2022)
- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [Direct Preference Optimization: Your Language Model is Secretly a Reward Model](https://mlsanity.com/paper/2305.18290/) (2023)
- [Training a Helpful and Harmless Assistant with Reinforcement Learning from Human Feedback](https://mlsanity.com/paper/2204.05862/) (2022)
- [A General Theoretical Paradigm to Understand Learning from Human Preferences](https://mlsanity.com/paper/2310.12036/) (2023)
- [Automatic chemical design using a data-driven continuous representation of molecules](https://mlsanity.com/paper/1610.02415/) (2016)
- [Open Problems and Fundamental Limitations of Reinforcement Learning from Human Feedback](https://mlsanity.com/paper/2307.15217/) (2023)
- [Negative Preference Optimization: From Catastrophic Collapse to Effective Unlearning](https://mlsanity.com/paper/2404.05868/) (2024)
- [ChemBERTa: Large-Scale Self-Supervised Pretraining for Molecular Property Prediction](https://mlsanity.com/paper/2010.09885/) (2020)
- [Therapeutics Data Commons: Machine Learning Datasets and Tasks for Drug Discovery and Development](https://mlsanity.com/paper/2102.09548/) (2021)
- [Molecular Transformer - A Model for Uncertainty-Calibrated Chemical Reaction Prediction](https://mlsanity.com/paper/1811.02633/) (2018)
- [Nash Learning from Human Feedback](https://mlsanity.com/paper/2312.00886/) (2023)
- [Preference Fine-Tuning of LLMs Should Leverage Suboptimal, On-Policy Data](https://mlsanity.com/paper/2404.14367/) (2024)
- [Sample Efficiency Matters: A Benchmark for Practical Molecular Optimization](https://mlsanity.com/paper/2206.12411/) (2022)
- [Beyond One-Preference-Fits-All Alignment: Multi-Objective Direct Preference Optimization](https://mlsanity.com/paper/2310.03708/) (2023)
- ["Found in Translation": Predicting Outcomes of Complex Organic Chemistry Reactions using Neural Sequence-to-Sequence Models](https://mlsanity.com/paper/1711.04810/) (2017)
- [Direct Preference-based Policy Optimization without Reward Modeling](https://mlsanity.com/paper/2301.12842/) (2023)
- [ACEGEN: Reinforcement learning of generative chemical agents for drug discovery](https://mlsanity.com/paper/2405.04657/) (2024)
- [Preference Optimization for Molecular Language Models](https://mlsanity.com/paper/2310.12304/) (2023)

## Built on by (most cited first, 7 of 7)

- [Unlocking Guidance for Discrete State-Space Diffusion and Flow Models](https://mlsanity.com/paper/2406.01572/) (2024)
- [Preference Learning Algorithms Do Not Learn Preference Rankings](https://mlsanity.com/paper/2405.19534/) (2024)
- [Reflections from the 2024 Large Language Model (LLM) Hackathon for Applications in Materials Science and Chemistry](https://mlsanity.com/paper/2411.15221/) (2024)
- [Boltzmann-Aligned Inverse Folding Model as a Predictor of Mutational Effects on Protein-Protein Interactions](https://mlsanity.com/paper/2410.09543/) (2024)
- [Pareto-Optimal Offline Reinforcement Learning via Smooth Tchebycheff Scalarization](https://mlsanity.com/paper/2604.13175/) (2026)
- [Training Large Language Models for Small-Molecule Design with Synthetic Task Scaling](https://mlsanity.com/paper/2609.04735/) (2026)
- [Generative modeling of intrinsically disordered protein regions by reinforcing sparse autoencoder features](https://mlsanity.com/paper/2610.02189/) (2026)

## Abstract (arXiv)

Searching through chemical space is an exceptionally challenging problem because the number of possible molecules grows combinatorially with the number of atoms. Large, autoregressive models trained on databases of chemical compounds have yielded powerful generators, but we still lack robust strategies for generating molecules with desired properties. 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 show theoretically that this algorithm is closely related to proximal policy optimization (PPO) and direct preference optimization (DPO), but has a minimizer that converges to an ideal Gibbs-Boltzmann distribution with the reward playing the role of an energy function. Furthermore, this 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. 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.