Title: Aligning Transformers with Continuous Feedback via Energy Rank Alignment
Venue: NeurIPS 2025
Accepted: yes
Keywords: transformers

alignment

preference optimization

protein language models

chemical language models

property optimization
Abstract: 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.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn materials and drug discovery, we want to find molecules that maximize some desired properties, such as docking score and solubility. Such processes amount to searching the vast chemical space of molecules. \n\n\nThis paper draws a similarity between search in chemical space and language model (LM) alignment. Indeed, alignment in LMs amounts to “biasing” a search policy towards sequences with high rewards or desired properties. Specifically, the authors propose the energy rank alignment (ERA) algorithm that leverages an explicit reward function that can be used to optimize autoregressive policies in the chemical space.\n\n\nERA draws similarity to standard reinforcement learning from human feedback algorithms such as PPO and DPO. Moreover, it is shown to be effective in searching through diverse parts of chemical space.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths**\n\n\n1. The application of LM alignment in molecule generation is very interesting.\n\n\n2. The proposed method seems to be effective for aligning molecular LM.\n\n\n**Weaknesses**\n\n\n1. I think the main issue with the paper is its presentation (Sec. 2). All the equations are simply stated, without much step-by-step derivation. This makes it hard to understand the motivation behind the different quantities presented. E.g. in Eq. 2, the authors simply stated “To impose the preferences, we minimize the objective: …” — it would be much better if the authors first derive/motivate each term (e.g. why the integrals?) in Eq. 2. Note that, this is not just about Eq. 2; I highly encourage the authors to revisit Sec. 2 as a whole.\n\n\n2. Connected to the point above, the discussion in Sec. 2 lacks theoretical justification/analysis. It would be great if the authors can show some theoretical results there. I know that some results might have been deferred to the appendix, but I encourage the authors to state some of them in the main text.\n\n\n3. Presentation of the results can also be significantly improved. E.g., it is quite hard to understand what Fig. 2 is trying to communicate, esp. since it is quite far away from Sec. 3.2, where it is being discussed.\n\n\n4. Last but definitely not least, costs are very important in materials/drug discovery. The authors need to discuss them. For example, through a wall-clock runtime vs. performance plot, and show that the proposed algorithm is on the Pareto frontier.\n\nQUESTIONS:\nCan the authors elaborate on the “design decisions” of the algorithm in a more compelling way? (See “Weaknesses” above for context.)\n\nLIMITATIONS:\nIf the costs are a limitation of the algorithm, I suggest that the authors add a discussion in this section.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 2, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper introduces an algorithm called Energy Rank Alignment to adapt general molecular generative models to specific design objectives. The method formulates a direct optimization objective as an energy function derived from task-specific reward functions, and theoretically proves its convergence to the ideal Gibbs–Boltzmann distribution. Experimental results demonstrate the effectiveness of the proposed strategy across diverse scenarios, including both small molecule generation and protein sequence design.\n\nSTRENGTHS AND WEAKNESSES:\nThe paper is well-written, featuring an intuitive introduction and detailed methodological description. The proposed approach presents a degree of novelty and is supported by experiments that demonstrate its effectiveness in guiding molecular generation toward desired objectives.\n\nHowever, a major concern is the lack of comparison with other established post-adaptation strategies, such as fine-tuning, Direct Preference Optimization (DPO), and Proximal Policy Optimization (PPO). While the authors show that their method can effectively optimize generation outputs based on task-specific rewards, they do not provide empirical evidence of its advantages over these alternative techniques. Including such comparisons would significantly strengthen the claim of superiority and clarify the practical benefits of the proposed method.\n\nIn addition, the inclusion of quantitative metrics, rather than relying solely on visual figures, would provide a clearer assessment of the model’s performance. Further analysis and ablation studies of the proposed method would also be valuable to better understand its behavior and effectiveness.\n\nQUESTIONS:\n1. Figure 1 is not referenced in the main text. It is recommended that the authors cite and briefly discuss it to clarify its relevance.\n\n2. The authors are encouraged to provide more intuition and explanations for each term in the key equations, particularly Equations (2)–(5), to improve readability and accessibility.\n\n3. The font size used in the figures is relatively big and may hinder readability. It is suggested that the authors adjust the font size for better clarity.\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 3, "CLARITY": 4, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces Energy Rank Alignment (ERA), a simple and effective method to fine-tune generative models so that they produce molecules or protein sequences with better properties, like higher drug-likeness or stronger binding affinity. Instead of using complex reinforcement learning or relying on human-labeled preferences, ERA takes advantage of a known reward function (e.g., QED, docking score) to decide which outputs are better. It compares pairs of samples generated by a pretrained model and teaches a new model to prefer the higher-scoring ones.\n\nA key strength of ERA is that it works entirely offline — it doesn’t need to generate new data during training, making it much more efficient and practical for real-world use. The authors apply ERA to several tasks: molecule generation without prompts, lead optimization conditioned on a given molecule, docking-based bioactivity optimization, and protein sequence design. Across all tasks, ERA improves the quality of generated samples while keeping a good balance between validity and diversity.\n\nIn short, ERA offers a general-purpose, easy-to-apply method for aligning generative models with specific goals in molecular and protein design, backed by both theory and strong experimental results.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths:**\n\n1. **Clear motivation and well-scoped goal**: The paper addresses the problem of aligning molecular and protein generative models toward desirable properties — a key challenge in inverse design. The scope is clearly defined and practically relevant.\n\n2. **Elegant and general algorithm (ERA)**: The proposed Energy Rank Alignment (ERA) framework is simple, principled, and broadly applicable. It sidesteps the complexities of reinforcement learning (e.g., PPO) and avoids the need for human-labeled preferences (as in DPO), while still offering strong theoretical guarantees. Its ability to decouple reward optimization (via β) and reference regularization (via γ) provides flexible control over the exploration–stability tradeoff.\n\n3. **Fully offline training**: A major strength of ERA is its ability to work with fixed datasets and an explicit reward function, enabling completely offline training. This is especially useful in domains where querying the environment is expensive (e.g., docking or fitness evaluation). The method never samples from the evolving model during training, which distinguishes it clearly from on-policy methods.\n\n4. **Strong performance across domains**: The framework is applied to four distinct settings — unprompted molecule generation, prompted molecule optimization, protein design, and docking — and consistently shows performance improvements over pretrained baselines. This demonstrates ERA’s robustness and cross-domain applicability.\n\n5. **Theoretical clarity**: The paper offers a well-grounded derivation showing that ERA minimizes a KL divergence between the model’s predicted preference distribution and a softmax over rewards. This theoretical foundation lends credibility and interpretability to the method.\n\n6. **Reasonable diversity and sample quality**: The model maintains chemical diversity and high reward without collapsing to a narrow output space. It performs comparably or better than other baselines such as REINVENT, MolRL-MGPT, and GFlowNet, with better sample efficiency in many cases.\n\n---\n\n**Weaknesses:**\n\n1. **Validity tradeoff not analyzed**: Although validity is reported (and remains fairly high), it consistently drops after alignment — particularly in high-reward targets like LogP. The authors do not discuss the cause of this drop. A likely explanation is that ERA pushes the model toward high-reward but syntactically fragile or chemically unstable regions of the molecule space. This tradeoff between reward-seeking and validity deserves deeper analysis or mitigation strategies.\n\n2. **No ablation on β and γ**: The method introduces two key hyperparameters — β for entropy shaping and γ for regularization toward the reference policy — but their effects are not thoroughly explored. A hyperparameter ablation would help guide users and support reproducibility.\n\n3. **Random sampling strategy under-specified**: While ERA relies on sampling multiple outputs \\( y, y' \\) from the reference model for a fixed input \\( x \\), the paper does not clarify what sampling strategy is used (e.g., temperature, top-k, top-p). Since this randomness directly affects diversity and the quality of preference pairs, it would be helpful to include this detail or discuss its impact.\n\nQUESTIONS:\n1. **Sampling randomness during training**: The ERA framework relies on sampling multiple outputs \\( y, y' \\) from the reference model \\( \\pi_{\\text{ref}}(y|x) \\) to construct preference pairs. However, the paper does not specify the sampling strategy (e.g., temperature, top-k, top-p) used during this process. Could the authors clarify what sampling method was used, whether temperature was tuned, and whether different sampling settings significantly affect performance, diversity, or training stability?\n\n2. **Validity drop after alignment**: In several tasks, the validity of generated SMILES strings drops after applying ERA (e.g., LogP↑ task). Could the authors provide more insight into this phenomenon? Specifically, do invalid generations tend to cluster around certain reward-optimized motifs or SMILES patterns? Is there evidence that the model is being pushed toward syntactically fragile or chemically implausible regions of molecular space?\n\n3. **Analysis of failure cases**: It would be helpful to understand the nature of invalid or low-reward generations. Could the authors provide examples or visualizations of failure cases (e.g., invalid SMILES, chemically unstable structures) and discuss what types of prompts or optimization targets tend to induce them?\n\n4. **Ablation of β and γ hyperparameters**: The ERA algorithm introduces two key hyperparameters — β for entropy shaping and γ for reference policy regularization — yet there is limited analysis of their effects. Could the authors provide a more detailed ablation study or discussion showing how varying these parameters affects the tradeoffs between reward, validity, and diversity across tasks? Additionally, is there any task-specific guidance on selecting appropriate values?\n\n5. **Potential for constrained decoding**: Given that validity tends to decline slightly under strong reward optimization, have the authors considered integrating grammar-based decoding, constrained decoding, or validity filters during generation? If so, what tradeoffs were observed? If not, could this be a promising direction to mitigate invalid generations while preserving ERA’s reward-guided behavior?\n\nLIMITATIONS:\n- **Lack of ablation on key hyperparameters (β, γ):** The method introduces two central hyperparameters — β for controlling entropy (sharpness of reward preference) and γ for regularization toward the reference policy — but does not offer a detailed ablation or sensitivity analysis. Understanding how these values affect reward optimization, validity, and diversity is critical for reproducibility and for guiding practitioners applying ERA to new tasks.\n\n- **Drop in chemical validity post-alignment not analyzed:** While validity is reported throughout, it often decreases after ERA alignment, especially in tasks like LogP optimization. The paper does not investigate this effect or analyze whether ERA leads the model to explore syntactically fragile or chemically unstable regions of molecule space.\n\n- **Sampling randomness during training is under-specified:** Since ERA relies on comparing multiple samples \\( y, y' \\) drawn from the reference model for a given input \\( x \\), the quality and diversity of these samples depend heavily on the sampling procedure. However, the paper does not specify whether temperature, top-k, or nucleus sampling was used, nor does it analyze the impact of sampling strategy on training quality.\n\n- **No demonstration of on-policy or adaptive extension:** While the method is theoretically amenable to on-policy training (via importance weighting or active sampling), all experiments are conducted in the fully offline setting. It would be valuable to demonstrate the effectiveness or limitations of on-policy extensions in sparse reward regimes such as docking or protein design.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper presents Energy Rank Alignment (ERA), a new algorithm for optimizing autoregressive models to efficiently explore chemical space and generate molecules with desired properties. Unlike RLHF, which depends on costly human preferences, or DPO, which uses implicit rewards, ERA employs an explicit reward function for gradient-based optimization. It aligns molecular generators to produce molecules with specific properties while ensuring diversity, converging to a Gibbs-Boltzmann distribution. Applied to small-molecule and protein sequence optimization, ERA demonstrates scalability and efficiency, advancing drug discovery and protein engineering.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n1. Energy Rank Alignment (ERA) presents an innovative gradient-based optimization approach that uses an explicit reward function, eliminating the need for expensive reinforcement learning (e.g., RLHF) or implicit reward models (e.g., DPO).\n2. The paper establishes a strong theoretical basis, demonstrating ERA’s convergence to a Gibbs-Boltzmann distribution and linking it to PPO and DPO.\n3. ERA achieves higher sample efficiency than baselines like REINVENT and GraphiCA, effectively generating diverse molecules with desired properties.\n\nWeaknesses:\n1. Clarity Issue: Line 15’s claim that ERA “does not require reinforcement learning” is confusing, as introducing sample fluctuations around the maximal reward resembles RL techniques.\n2. Unclear Need for Fluctuations: With an explicit reward function, the necessity of sampling fluctuations around the maximum reward is not well-justified.\n3. Comparison to PPO: The paper does not clearly articulate ERA’s advantages over PPO when using an explicit reward function.\n4. Missing Ablation Study: The paper lacks an ablation study comparing ERA, PPO, and DPO under explicit reward conditions, limiting insights into their relative performance.\n\nQUESTIONS:\nSee weakness\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 2, "REVIEWER_CONFIDENCE": 4, "CLARITY": 2, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-07-03"}