Energy Rank Alignment: Using Preference Optimization to Search Chemical Space at Scale

Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-13 23:29:27

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 to generate molecules with externally specified properties and find that it does so robustly, searching through diverse parts of chemical space. While our focus here is on chemical search, we also obtain excellent results on an AI supervised task for LLM alignment, showing that the method is scalable and general.

Main paper: 
 introduction: 
 Large language models (LLMs) are trained on large corpora of text to autoregressively generate outputs. These models strongly reflect the distribution of the data on which they are trained [21] , and controlling the outputs to reflect externally imposed preferences is an increasingly important challenge for deployment. The aforementioned task, often called "alignment", requires either careful curation of training data or large sets of human preference data-both options are labor-intensive [9] .Reinforcement learning from human feedback (RLHF), a family of algorithms that employs these human preference datasets, has been widely employed to align instruction and chat models [21, 5] , but it is both expensive to acquire the training data and difficult to carry out in practice [9] . Recent algorithmic developments, such as direct preference optimization (DPO) [25] , simplify the alignment framework by making the reward function implicit, but still require human preference data. While these algorithms succeed in constraining outputs, many "alignment"-like tasks require evaluation that would be difficult for human evaluators.Generative sampling problems seeking to optimize a reward are common in chemistry, where comparing small molecules using a particular functional assay or computationally accessible property Submitted to 38th Conference on Neural Information Processing Systems (NeurIPS 2024). Do not distribute. is often far easier than searching chemical space to identify novel compounds. Recent efforts to build large, domain-specific models for chemistry [10] have shown promising performance on both property prediction and reaction prediction tasks. Nevertheless, just as with LLMs, leveraging these models for molecule optimization requires first guiding "unaligned" models to favor important properties like synthetic accessibility or solubility. Here, we seek to productively search chemical space using transformers by introducing a new preference optimization algorithm, which we call energy rank alignment.Our contribution: We formulate a generic alignment algorithm that we call Energy Rank Alignment or ERA that leverages an explicit reward function to guide autoregressive sampling while targeting specific properties or preferences. Unlike reward maximization in RL-based algorithms, the policy that minimizes our objective is designed to sample fluctuations around a maximal reward value to promote sample diversity. Our algorithm enables direct gradient-based optimization of a policy to match the ideal preference distribution and converges asymptotically to an optimal distribution with tuneable entropy and controllable regularization, which we show theoretically. The minimizers of our objective are closely related to the minimizer of PPO and DPO, but we have more direct control over the influence of the regularization relative to fluctuations around the maximum reward. In numerical experiments, we demonstrate that this algorithm successfully aligns a molecule transformer model to identify a highly diverse set of chemicals with properties favored by our choice of reward. Finally, we also show that we obtain competitive performance with ERA on benchmark LLM alignment tasks, but emphasize that the chemical applications are the main focus of this paper. HighInverse molecular design tasks have a long history [17] and many recent works have sought to apply machine learning to facilitate this difficult search problem [27, 12, 13] . While reinforcement learning has proved a popular strategy for molecular optimization [39, 27] , several recent studies have sought to use transformers [34] trained on large databases of molecules represented with the text-based SMILES syntax [10, 30, 35, 4] for such tasks. Schwaller et al. [31] utilized an atom-wise tokenization, which we also employ, to train a transformer for the downstream task of reaction prediction. These "chemical language models" have been studied for applications on downstream tasks, including property prediction [4, 10] and reaction prediction [23, 30] .Building scalable strategies for alignment has attracted enormous attention because of the high cost and complexity of constraining LLM outputs. Much of the current paradigm is built on reinforcement learning from human feedback (RLHF) [21] . Within this framework, human preferences provided in the form of pairwise rankings are first used to train a reward model, and subsequently that reward model is used to optimize a policy using, for example, proximal policy optimization (PPO) [29] .Rafailov et al. [25] demonstrated that the reward model can be treated implicitly using a scheme that maximizes the likelihood of the preferences given an offline dataset. Because this approach does not require training a reward model, it has been named Direct Preference Optimization (DPO).Our work differs from both strategies; first, unlike RLHF, we do not employ reinforcement learning and instead develop an explicit, gradient-based objective for the optimal policy. Secondly, unlike DPO, we leverage an explicit reward function and add regularization transparently, both of which help to avoid greedy policies [3] . However, like both approaches, we assume that the Bradley-Terry model [7] of preference data is appropriate for the underlying target distribution.Many recent works have built upon the ideas of RLHF and DPO, including studies on the effect of point-wise sampling of preference distributions [3] , investigations into the theoretical basis for contrastive methods for unlearning target datasets [38] , and alternatives to the Bradley-Terry pairwise preference model [20, 2] . One recent study explores alignment in the context of inverse molecular design: Park et al. [22] applies DPO to SMILES generators to increase the probability of activity for generated compounds against a drug target. However, they indicate that many preferences in chemistry are expressed as continuous signals, which is not suitable for DPO. Overcoming this limitation while maintaining the advantages of a direct gradient-based policy optimization strategy is a central goal of our current work. Our analysis and methodology directly addresses issues related to point-wise sampling because the explicit reward function eliminates overly greedy assignments of preference probabilities. Indeed, as discussed in Sec. 4 , we see that DPO mode collapses where ERA shifts the policy towards the target distribution. While non-transitive preferences may arise in some settings, leading to a breakdown of the Bradley-Terry preference distribution model, by construction our target rewards are determined by quantitative evaluations of properties, and are therefore transitive.Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?Answer: [Yes] Justification: The abstract clearly outlines the problem of searching over vast chemical spaces as an alignment problem, gives a high-level overview of our method, highlights the key differences between our method and similar approaches, and summarizes the key advantages of our algorithm. We clearly define the current scope of our investigation by describing our experiments.Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in the paper.• The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.• The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.• It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper. 
energy rank alignment: 
 A policy is a conditional probability distribution π(•|x) : Y → R; we generate an output y from prompt x. The spaces Y and X are discrete and finite, corresponding to sequences of tokenized outputs of the model with a maximum length. In alignment tasks, we begin with a pre-trained reference policy π ref and seek to optimize a parametric, trainable policy π θ to adapt the conditional sampling for a particular task or constraint.Consider a prompt x ∈ X and model outputs y, y ′ ∈ Y and a collection of preferences D = {(y i ≻ y ′ i ; x i )} n i=1 ; the notation ≻ indicates that y i is preferred to y ′ i . The conditional probability that y ≻ y ′ given x can be modeled as a pairwise Boltzmann ranking within the Bradley-Terry model, i.e., p(y ≻ y ′ |x) = e -βU (x,y) e -βU (x,y) + e -βU (x,y ′ ) ≡ σ βU (x, y ′ ) -βU (x, y) .(EQUATIONHere β > 0 is a constant, σ(x) = (1 + e -x ) -1 and we refer to U : X × Y → R as an energy function to make clear the connection to statistical physics, but it is the negative reward within the RL framework for alignment.To impose the preferences we minimize the objectiveEQUATIONwhere β -1 is a parameter controlling the magnitude of the entropic term, γ sets the scale of the Kullback-Leibler regularization compared with the energy term, and ν is a probability distribution over the prompts ν ∈ P(X ). A proximal scheme for gradient descent on this objective corresponds to a gradient flow on J [28, 19] ; the functional can be viewed as a free energy, and the corresponding flow isEQUATIONand δ π denotes the Fréchet derivative with respect to π. Assuming that π 0 has full support on X × Y, the optimization converges asymptotically to stationary policy which satisfiesEQUATIONand this minimizer is globally optimal. In the context of LLM alignment, a representation of the energy function U : X × Y → R is learned as a "reward model", though we also consider tasks in which U is an easily evaluated function of the pair (x, y). The optimal distribution π ⋆ is a Gibbs-Boltzmann measureEQUATIONwhere Z(x) is the x-dependent normalization constant. This expression makes clear the effect of β:when β → ∞ (low temperature), the reward dominates and fluctuations around the maximal reward are small, which could lead to "mode-seeking"; when β → 0 (high physical temperature) fluctuations around the maximal reward increase and the regularization term favors proximity to π ref . Similarly, γ → 0 recovers a Gibbs-Boltzmann distribution proportional to e -βU at inverse temperature β, while γ → ∞ is dominated by the reference policy.Loss functions for π θ : Proximal Policy Optimization (PPO) optimizes an indirect, proximal objective to minimize an objective closely related to (2) (cf. Appendix A). Direct Preference Optimization (DPO) treats the negative reward function U implicitly and directly maximizes the likelihood of p(y ≻ y ′ |x). Our objectives differ from both approaches: like DPO, we directly optimize the policy using an explicit, gradient-based objective, but, in contrast, we use a reward function directly in our objective. The losses we build are thus amenable to both offline (samples from π ref ) and online (samples from π θ ) policy alignment, as explained below. Choosing to optimize the objective online has been shown to have important consequences on performance [32] , though we focus here on the setting where samples are drawn offline.We directly optimize the Kullback-Leibler divergence between the entropy-regularized preference distribution p γ (y ≻ y ′ |x) and the corresponding parametric preference distribution p θ (y ≻ y ′ |x).Explicitly, using the fact that conditional preference distribution is normalized, we obtainEQUATIONwhereEQUATION)This quantity is a well-defined KL divergence and is hence non-negative; the quantity vanishes when p γ = p θ on the observations y, y ′ . Furthermore, with access to an explicit reward model, all terms in (6) can be computed directly andEQUATION)To obtain a minimizer of the regularized objective defined in (2) we optimizeEQUATIONIf the current policy overlaps with the target preference distribution, it may be useful to sample directly from the partially aligned policy, i.e., to use the "on-policy" formulation,EQUATIONinstead of (9) . One issue that arises with this scheme is that differentiation with respect to the parameters of the policy θ because y and y ′ are decoded into discrete tokens, an operation that is not differentiable. To remedy this, we importance sample with a reference policyEQUATION)This reweighting is straightforward and the importance weights should generally be appreciable, especially early in training when π θ has not drifted far from π ref .It is, of course, also natural to iteratively update π θ using a previous iterate as the reference policy. In this work, we only use (9) as an objective and leave the on-policy objectives to future work.Question: Does the paper discuss the limitations of the work performed by the authors?Answer: [Yes]Justification: We address several potential limitations of our work at the end of Section 5.Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.• The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.• The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.• The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.• The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.• If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.• While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
theoretical analysis: 
 To understand the ERA loss function and its connection to the entropy regularized objective (2), we first establish that the minimizers of ( 6) are of the form (5) . We first define the notion of equivalence precisely. We remark that this strong form of equivalence is appropriate on the finite, discrete spaces X and Y we consider here. x) is conditionally equivalent to π for all functions g : X → R such that sup x∈X |e g(x) | < +∞.Lemma 3.1 If π is conditionally equivalent to π ′ , then π ′ g (•|x) ∝ π ′ (•|x)e g(We prove Lemma 3.1 in Appendix A and use this simple lemma to prove the following result. First, we verify that any probability measure π g (y|x) ∝ exp(-β 1+γ U (x, y) -β -1 γ log π ref (y|x) + g(x)) minimizes the objective. Because L ERA is non-negative, it suffices to show that for all pairs y, y ′ , D (y,y ′ ) KL (p γ |p θ ) ≡ 0. This follows immediately from the cancellation in the preference probability p γ of e g(x) after factorization in (5) . Now, suppose thatπ(y|x) ̸ = exp -β 1+γ U (x, y) -β -1 γ log π ref (y|x)where we have taken g(x) = 0 without loss of generality and π := π g . Assume that for all pairs y, y ′ , the divergenceD (y,y ′ ) KL (p γ |p θ ) ≡ 0which is required of a minimizer. Equivalently, it must be the case that for all y, y ′ ,EQUATIONfrom which we see thatEQUATIONBy construction, π(y|x) does not depend on y ′ so the prefactor must be purely a function of x, which completes the proof, using Lemma 3.1.Gradients of L ERA . One advantage of the ERA framework is that the objective is amenable to direct, gradient-based optimization. We remark that establishing global convergence for the optimization of θ using (9) requires establishing convexity with respect to the parameters, which is not obviously the case for our objective, nor those used in PPO and DPO. However, one can still glean some insight into the optimization by examining the gradients on a samplewise basis. Using the compact notationEQUATION)The gradient is straightforward to interpret on a particular pair y, y ′ : if p θ (y ≻ y ′ |x) is larger than p γ (y ≻ y ′ |x) then the preference gradient is positive and gradient descent lowers the probability that y ≻ y ′ . The opposite occurs whenever p θ (y ≻ y ′ |x) is smaller than p γ (y ≻ y ′ |x). The magnitude of the gradient is scaled by the degree of misspecification of the preference probability.This calculation highlights one key difference between the approach we use and DPO. When the data only contains one observation of y ≻ y ′ for a given x, the DPO objective's implicit reward model assigns zero probability to y ′ ≻ y. This pushes the policy towards extremal values, which can lead to undesired behavior, as discussed in Azar et al. [3] . In our formulation, this behavior occurs only when the reward model assigns an energy of ±∞, which is prohibited by construction in most tasks.We further discuss differences between ERA and DPO in Appendix A.2.Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?Answer: [Yes]Justification: Section 3 provides a thorough theoretical analysis of the ERA algorithm, and further details are in Section A of the Appendix.Guidelines:• The answer NA means that the paper does not include theoretical results.• All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.• All assumptions should be clearly stated or referenced in the statement of any theorems.• The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.• Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.• Theorems and Lemmas that the proof relies upon should be properly referenced. 
experiments: 
 We test ERA on both chemical and language tasks to shed light on the following questions: 1) Can we use ERA to robustly fine-tune our model to generate samples according to a desired distribution? All experiments were run with no regularization to the reference policy (γ = 0).2) What is the effect of changing the inverse-temperature β during ERA? 3) Do we maintain sample diversity (and validity) without regularizing to remain close to a reference policy, and what is the effect of increased regularization? 4) Can we simultaneously target multiple properties with high fidelity, and how can we trade off between desired properties? 5) Can we carry out ERA on higher capacity models with "weak" signals from smaller models?We use a decoder-only representation for the molecular generator [4] , where the generator has 2 layers, an embedding dimension of 512, a vocabulary of 324 tokens, and totals 3.5M parameters. Starting from a random initialization, we carry out pretraining on a dataset of 2.4M small molecules from the ChEMBL database [37] for 180 epochs. This version of the model is not conditioned on a prompt and generates a small molecule given just a start-of-sequence token. We use this pretrained model as our reference policy for all unprompted molecular alignment tasks (Sec. 4.1.1). In Sec. 4.1.2, we generate molecules conditioned on a prompt using a generator that was trained to carry out sampling with a prompt molecule.Central to ERA is, of course, access to a computable energy function. As a proof-of-concept, here we consider 5 different properties for which the corresponding energy function is easily evaluable:Quantitative Estimate of Drug-Likeness (QED) [6] , Wildman-Crippen LogP (LogP) [36] , Ring Count, Molar Refractivity (MR) [36] , and Tanimoto Similarity [26] . Briefly, LogP is a measure of the hydrophobicity of a molecule, MR is a measure of the polarizability of the molecule, and Tanimoto similarity is a measure of the similarity between two molecules (see Appendix C.2).First, we independently target four different properties using ERA with an unprompted molecular generator (Fig. 2 ). Using the reference policy, we generate a dataset D =EQUATION2 , U (y(i) 1 ), U (y (i)2 )} N i=1 and carry out energy rank alignment on π θ , where π θ is initialized using the weights of π ref .Here, y 1 , y 2 ∼ π ref and y and U (y) denote the generated molecule and its corresponding energy, respectively. For MR, Ring Count, and LogP, we define the energy U to be a harmonic potential centered at a target value. For QED, we define the energy to be the negative logarithm of QED and vary β to assess its impact on alignment (see Table 1 , 2). In Fig. 2 , we see that we successfully shift the distribution to target means that are both greater and lower than the average value of MR, Ring Count, and LogP under the reference policy. Furthermore, in the alignment of QED, we observe the effect of changing β on the learned policy; with increased β, the learned policy concentrates around low-energy samples (i.e. near QED = 1), and with lower β, the learned policy samples a greater range of QED values, as expected. We note that for each of these four experiments, we did not regularize towards the reference policy (i.e. γ = 0). Even so, we were able to maintain both sample diversity and maintain appreciable sample validity (see Fig. 7 and Table 3 ).Many molecular design tasks require balancing multiple properties, and designing an objective for multi-property alignment is straightforward within the ERA framework. To demonstrate this, we generate molecules with both high QED and LogP using ERA with an energy function weighted by property-specific β: U = β QED U QED + β LogP U LogP (see Table 1 , 4 for details on energy function).We carry out ERA with different pairs of (β QED , β LogP ) using the same procedure as above, and from Fig. 3 , we see that we target multiple properties with varying fidelity by simply modulating the value of property-specific β. Ultimately, increasing the β for an individual property enables us to favor higher values of that property in multi-property alignment setting. In this case, we also do not regularize with the KL-divergence to the reference policy and again maintain sample diversity and validity (see Fig. 8 and Table 4 )Inspired by the task of lead optimization in drug discovery efforts [16] , we ask whether we can use ERA to train a molecular generator that can sample a molecule that is both similar to the prompt molecule and also exhibits some desired property.First, we fine-tune the pretrained molecular generator to enable prompted molecular generation (see Appendix C.3.2) and use this fine-tuned model as our reference policy for all prompted molecular alignment tasks. This reference policy disproportionately samples molecules that are identical (i.e. aTanimoto similarity of 1.0) to the prompt molecule (see Fig. 4 ), so we carry out multi-property alignment on this reference policy to generate molecules that are similar-but not identical-to the prompt molecule and also have a high drug-likeness as measured by QED. Using ERA, we optimize the reference policy with a generated dataset D = {(y(i) 1 , x (i) ), (y (i) 2 , x (i) ), U (y (i) 1 , x (i) ), U (y (i) 2 , x (i) )} N i=1 ,where we sample four molecules for each prompt molecule from the reference policy and consider all possible preference pairs for a total of six preference pairs per prompt molecule (see Appendix C.2for full details on energy used).We observe that the per-prompt average QED under the optimized policy for a given prompt is higher than the corresponding average under the reference policy (Fig. 4 ). Furthermore, we see that we are able to sample a diverse set of molecules that are chemically similar to the prompt molecule, and also chemically valid (see Figure 9 , Table 5 ). We repeat the experiment with a related objective of generating molecules similar to the prompt molecule with a high LogP instead and again observe that we increase the per-prompt average LogP under the optimized policy relative to the reference policy without degrading sample diversity and validity. For both of these experiments, we required regularization to the reference policy. With no regularization, the aligned generator would almost exclusively sample sequences that were chemically invalid (< 25% chemical validity). Finally, we note that the increases in QED and LogP in Fig. 4 are smaller relative to the increases in Fig. 2 because the samples are now conditioned to remain proximal to the prompt molecule, which restricts the chemical space that can be explored.We test the generality of ERA by applying it to align large language models (LLMs). Similar to the experiments in [25] , we first carry out ERA on a GPT-2 model [24] fine-tuned on movies reviews from IMDb [18] . We use a pretrained sentiment classifier [14] to evaluate the energies-where lower energies correspond to more positive sentiments-of sampled responses from the reference policy and carry out ERA using the same approach as in Section 4.1.2 (see Appendix D.1). We vary the regularization strength γ and inverse-temperature β on the average sentiment and observe that across all regularization strengths, with increasing β, the average sentiment becomes more positive. Increasing regularization also elicits more positive sentiments. Qualitatively, with lower regularization, we observe that text quality degrades and becomes less coherent, likely resulting in lower average sentiment predictions by the sentiment model. Regularization here is important to ensure high quality text samples.We next leverage a "weak" AI supervisor to carry out LLM alignment, a task sometimes called "superalignment" [8] . In the present context, we order "weak" vs. "strong" models based on their parameter count (within the same family) and empirical performance; i.e., LLaMA2-7B is weaker than LLaMA2-13B. Here, the weak model does not necessarily contain the complexity of the stronger model but can weakly discern between different outputs of a stronger model. Given a sample y i ∼ π strong (y|x), we define the energy using the weak model U (y i |x) =log π weak (y i |x).We test weak-to-strong alignment using a previously aligned LLaMA2-7B-Chat (meta-llama/Llama-2-7b-chat) to optimize an unaligned LLaMA2-13B(meta-llama/Llama-2-13b) model [33] . Using prompts from the Anthropic Helpful and Harmless dialogue dataset [5] , we first carry out a short supervised fine-tuning step of LLaMA2-13Bto ensure it can output text in a chat-like format (see Appendix D.2). Using this reference policy, we generate a dataset with energies computed from the smaller LLaMA2-7B-Chat model and carry out ERA as above, again across varying γ and β. We evaluate the "safety" of generated samples using Meta LLama Guard 2 (meta-llama/Meta-Llama-Guard-2-8B) [15] . We observe that as we increase β, the proportion of unsafe content relative to the unaligned, reference model decreases, with over a 90% drop between the unaligned model and the models aligned with the highest β across all γ.For these experiments, we observe that varying regularization strengths has a minimal effect and that we are in fact able to generate coherent sentences with no regularization, with strong regularization hurting performance for β = 0.1. Finally, we compare ERA and DPO in Appendix D.2 and observe that with our implementation of DPO, we are able to generate lower energy samples, but that it is prone to mode collapse. We caution that our implementation of DPO is likely not optimal and that we did not exhaustively tune the hyperparameters of DPO due to resource constraints.Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?Answer Guidelines:• The answer NA means that the paper does not include experiments.• If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.• If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.• Depending on the contribution, reproducibility can be accomplished in various ways.For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.• While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example (d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility.In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results. Guidelines:• The answer NA means that paper does not include experiments requiring code.• Please see the NeurIPS code and data submission guidelines (https://nips.cc/ public/guides/CodeSubmissionPolicy) for more details.• While we encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).• The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https://nips.cc/public/guides/CodeSubmissionPolicy) for more details.• The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.• The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.• At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).• Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted. 
conclusions and limitations: 
 This paper introduces energy rank alignment, a simple and effective algorithm for policy optimization with an explicit reward model. We find that ERA is stable without extensive hyperparameter tuning, and sufficiently general to successfully align both application-specific transformers for chemical search problems as well as generative pre-trained transformers for language. The algorithm exhibits strong performance with a variety of reward models, even ones with relatively weak signal, such as the AI feedback of LLaMA2-7B-Chat. Interestingly, with this approach we are able to reduce unsafe content by more than 90% with no human preference data.We analyze the minimizers of the ERA objective and find that they differ from the minimizers of popular policy alignment algorithms DPO and PPO in an important way: unlike PPO, the strength of regularization to the reference policy that we add is controlled by a parameter γ, while the entropy of the target distribution is independently tuned by a distinct parameter β. This means that we can avoid greedy policies by keeping β small-amplifying fluctuations around the optimum of the reward model -U -while reducing the influence of the reference policy by taking γ small. Our objective leads to easily interpretable sample-wise gradients which highlight the importance of a reward model relative to DPO in the sampled objective. Similar observations about the inadequacy of the DPO objective for finite preference observations were also made theoretically in Azar et al. [3] .Limitations: First, our approach requires a reward model, which can be difficult to train or design, especially for complex tasks. While we observed that ERA makes an appreciable impact even with weak supervision from an AI chat model, this sort of proxy may not be available for more complex tasks. For example, optimizing small molecules for high binding affinity to a target protein would require expensive and noisy evaluations of a reward model, which likely limits the scope of molecular design to problems where the reward can be computed somewhat efficiently. A second limitation of our present work is that we do not train the molecular transformer to favor synthetic accessibility nor do we explicitly seek to obtain molecules that are easily synthesized experimentally. There are models that seek to evaluate synthesizability computationally that could be used in our rewards, which we plan to explore in future work [11] . A final limitation of our current work is the moderate scale of our numerical experiments due to our limited compute resources, including the inadequate hyperparameter tuning for the DPO baseline for Fig. 5 . 
experimental setting/details: 
 Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?Answer Guidelines:• The answer NA means that the paper does not include experiments.• The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.• The full details can be provided either with the code, in appendix, or as supplemental material. 
experiment statistical significance: 
 Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?Answer: [Yes] Justification: The statistical significance of all the results is discussed throughout the body of the paper in Section 4 and in Sections C and D of the Appendix.Guidelines:• The answer NA means that the paper does not include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.• The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).• The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)• The assumptions made should be given (e.g., Normally distributed errors).• It should be clear whether the error bar is the standard deviation or the standard error of the mean.• It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.• For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).• If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text. 
experiments compute resources: 
 Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?Answer: [Yes]Justification: Computational resources used are discussed in Section E of the Appendix.Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.• The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.• The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines? 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [Yes] Justification: The broader societal impacts of this work is discussed in Section F.Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.• Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.• The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.• The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.• If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [Yes]Justification: We use public data and models for all LLM experiments. There are no safety concerns with the chemical alignment models.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.• Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.• We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort. 
licenses for existing assets: 
 Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes]Justification: Yes, we cite the appropriate works throughout the main text and the Appendix.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL.• The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.• If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.• For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.• If this information is not available online, the authors are encouraged to reach out to the asset's creators. 
new assets: 
 Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?Answer: [Yes] Justification: The code will be released as open-source and will have adequate documentation for reproducing experiments.Guidelines:• The answer NA means that the paper does not release new assets.• Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.• The paper should discuss whether and how consent was obtained from people whose asset is used.• At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file. 
crowdsourcing and research with human subjects: 
 Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?Answer: [NA] Justification: We did not use human subjects.Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.• Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.• For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: Justification: Yes, the work is compliant with the NeurIPS Code of Ethics.Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.• The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).Set-up, notation, and assumptions Let X and Y be discrete spaces; each element of one of these spaces is a finite-length sequence of tokens within a fixed dictionary on which an autoregressive generative model is trained. The resulting models yield "policies", which are conditional probability distributions π(•|x) ∈ P(Y) for each x ∈ X . Throughout, we assume that our policies have full support on Y for each x, meaning that inf y,x∈Y×X π(y|x) > 0. Because the spaces are discrete, we make no strong restrictions on the regularity or coerciveness of the reward model-U : X × Y → R.The only requirement to ensure the existence of an optimal probability distribution is that sup x,y×X ×Y |e -U (x,y) | < +∞, which maintains full support of the distribution.Though it plays little role in theoretical analysis, we also denote by ν ∈ P(X ) the probability distribution over the prompts x.The main purpose of this section is to establish that globally minimizing the loss ( 9) yields a global minimizer of the regularized policy objective (2) . A secondary goal is to clearly articulate the theoretical advantages of ERA compared with PPO and DPO.To understand the ERA loss function and its connection to the entropy regularized objective (2), we first establish that the minimizer of ( 6) are of the form (5) . We first define the notion of equivalence precisely. This is a strong form of equivalence for probability measures, but it is appropriate on the discrete spaces X and Y we consider here. For more general continuous spaces, one could relax this condition to weak equivalence of the conditional measures. We use this notion to emphasize that a shift of the distribution of the "prompts" x ∈ X , which we denote ν ∈ P(X ), does not impact conditional equivalence and hence establishes an equivalence class of conditional probability measures that minimize (2) . x) is conditionally equivalent to π for all functions g : X → R such that sup x∈X |e g(x) | < +∞.Lemma A.1 If π is conditionally equivalent to π ′ , then π ′ g (•|x) ∝ π ′ (•|x)e g(Assume that π ′ is a normalized probability distribution. This requires that,EQUATION)If g is such thatEQUATIONthen the normalized policy π ′ g is clearly defined by x) . By the assumption that sup x∈X |e g(x) | < +∞, all terms in these calculations remain finite.EQUATION) because Z ′ g (x) = e g(Using Lemma A.1 it is straightforward to prove the result in the main text Proposition 3.2. For completeness, we re-state that result here and refer the reader to the main text for the complete argument. This proposition establishes that a policy minimizing the objectiveEQUATIONhas the formEQUATIONWe do not, however, prove that gradient descent of θ on (18) converges to the global minimizer (19) because such an argument requires additional assumptions about the parametric class of policies and the convexity of the objective with respect to the parameters, neither of which are straightforward to establish.The free energy functional for a policy under the energy rank alignment framework can be written as an expectationEQUATIONinvolving an energetic term and an entropic term. The additional regularization acts as an effective energetic bias. Solving for the extremum of this functional by setting Fréchet derivative with respect to π equal to zero, one obtains the formal solution (19) for the minimizer. This objective differs from the regularized reward loss conventionally used for PPO,EQUATIONThe minimizer of the PPO objective ( 21) is also a Gibbs-Boltzmann measure, explicitly,EQUATIONHere, the KL-regularization corresponds to an energy shift, as in our objective, but there is no limit in which the ideal distribution π ∝ e -βU is obtained for the PPO objective. This is in stark contrast to our approach, which recovers the ideal distribution as γ → 0. Furthermore, while our approach allows for a direct gradient-based optimization using (18) , PPO is implemented using an actor-critic framework that is difficult to tune [25, 9] . Finally, we emphasize that for ERA in the γ → 0, finite β > 0, the distribution has positive entropy and is not manifestly mode-seeking; there can still be appreciable fluctuations in the output. Eliminating the effect of regularization in (22) , on the other hand, requires taking β/γ → ∞, which eliminates fluctuations in the distribution.The DPO approach also seeks to optimize the objective (21) . The algorithm does so by first using (22) to define an implicit reward model by solving for the U that reflects the observed preference probabilities. This elegant idea has had a significant impact and has already been deployed in state-of-the-art models [1] . In many cases, the observed preference probabilities will be sampled and only perhaps only one observation of y ≻ y ′ will be available for each x in the dataset. When the preference dataset only has one observation y ≻ y ′ per prompt x, the optimal policy requires thatEQUATIONThe sampled gradients of the objective used for DPO are proportional to the implicit reward discrepancy,EQUATIONwhich when π θ (y ′ |x) → 0, could lead to instability aslog π θ (y ′ |x) → ∞. On the other hand, the ERA gradients are scaled by the relative preference discrepancy,EQUATIONThe advantage of a reward model becomes apparent becauseEQUATIONand hence the optimum of L ERA will not lead to policies in which supp(π θ ) degrades unless the energy becomes infinite. Choosing an appropriate reward model, hence, gives the flexibility to control instability if it becomes problematic.Implementing energy rank alignment is straightforward to implement within existing code bases. We provide sample PyTorch code for the ERA loss function below. import torch.nn as nn from torch.nn.functional import logsigmoid def era_loss(pi_logps_1, pi_logps_2, In this work, we represent all molecules as SMILES strings and tokenize SMILES strings according to the approach in [30] . Our dataset consisted of all small-molecules from the ChEMBL database that were of length 500 tokens or less. Ultimately, this token limit filtered out approximately 0.1% of the small-molecules in the original ChEMBL dataset. The alphabet generated from this curated dataset consists of 324 tokens, which we augmented with start, stop, and padding tokens.We first pretrained a model according to a next-token prediction, self-supervised learning approach.We trained a model using the standard cross entropy lossEQUATIONOur trained molecular generator consisted of just the encoder block of a standard multi-head attention transformer [34] . Finally, the model had 2 layers, 8 heads, and a width of 512. For pretraining, we used an Adam optimizer with a learning rate of 1.0 * 10 -5 . We emphasize that this pretrained generator samples molecules in an unprompted fashion; given just a start-of-sequence token, we can autoregressively generate a sequence of tokens. Moreover, it is possible that this sequence of tokens corresponds to a molecule that is not chemically valid, and we find that around 88% of all generated molecules are chemically valid. Lastly, we measure the diversity of the pretrained molecular generator by first generating 1500 molecules and then computing the Tanimoto similarity between every pair of molecules. We plot the distribution of all pairwise Tanimoto similarities from this sample and from all pariwise Tanimoto similarities from 1500 randomly sampled molecules from the original dataset in Fig. 6 . We observe that we can generate molecules that are quite distinct (i.e. low Tanimoto similarity) in comparison with all other molecules.We investigated aligning the molecule generator to several target chemical properties, which we detail below. All of the properties can be easily computed using the RDKit package. We list the energy function and parameters used for the corresponding energy functions for each of these properties in Table 1 .Tanimoto similarity is a measure of chemical and structural properties between two molecules and ranges from 0 to 1, where higher values correspond to more similar molecules [26] . Quantitative estimation of drug-likeness (QED) is evaluated by taking the geometric mean of a set of "desirability functions" for different molecular descriptors and also ranges continuously from values of 0 to 1 [6] ,where higher values correspond to more drug-like molecules. The octanol-water parition coefficient (Wildman-Crippen LogP) is a measure of hydrophobicity frequently employed in medicinal chemistry applications [36] . Molecules with more positive values are more hydrophobic (i.e. more soluble in octanol relative to water), whereas molecules with more negative values are more hydrophilic (i.e. more soluble in water relative to octanol). Molar refractivity is similarly calculated as a linear combination of atomic contributions, and is a positive number that serves as a measure for molecular size and polarizability [36] . A higher molar refractivity corresponds to larger and more polarizable molecules. Finally, ring count corresponds to the number of rings in a molecule.Under the definitions of the energy functions in Table 1 , it is possible for a generated sequence to not be chemically valid. For these cases, we manually define energies that are sufficiently high to penalize that outcome and we report these values in Table 2 . Furthermore, when the computed QED or Tanimoto Similarity is 0, the energy is infinite, and to ensure numerical stability, we set the value of the energies to be 4. We first investigated aligning the unprompted molecular generator to sample small-molecules with desired properties. We carried out alignment using the property-specific energies described in Table 1 .All alignment properties were initialized with the weights of the pretrained model and trained using an Adam optimizer with learning rate 1.0 * 10 -6 . We tabulate the chemical validity for single-property alignment in Table 3 and for multi-property alignment in Table 4 . While we do see a drop in chemical validity after alignment, we see that a majority of the samples we generate post-alignment are still chemically valid despite no regularization to a reference policy. We measure the chemical diversity for these experiments by computing all pairwise Tanimoto similarities from all chemically valid 916 predictions of 1500 generated molecules. We visualize the chemical diversity for single-property 917 experiments in Fig. 7 and multi-property experiments in Fig. 8 . We observe that the samples are still 918 highly diverse chemically after alignment. All plots in Fig. 2 and Fig. 3 were computed using 1500 919 generated molecules per experiment.Chemical validity Unaligned 93% β Tanimoto = 5.0, β LogP = 10.0, µ LogP = 5.0, σ LogP = 1.0, γ = 0.1 91% β Tanimoto = 5.0, β QED = 500.0, γ = 0.1 81% we repeated the random perturbation until a valid molecule was generated. The prompted generator was the same size as the unprompted molecular generator, and we initialized the weights using those of the pre-trained unprompted molecular generator. We then carried out supervised fine-tuning using an Adam optimizer with learning rate 1.0 * 10 -5 and used this generator as our reference policy for all prompted alignment experiments. All plots in Fig. 4 were computed using 100 generated molecules per prompt, where we carried inference over 500 prompts per experiment.Similar to the experiments run in [25] , we carried out alignment of a GPT-2 model fine-tuned on a dataset of IMDb reviews to a pretrained sentiment model. For this experiment, we first carried out supervised fine-tuning of gpt2-large using an 80/20 train/validation split of the 25000 reviews in (stanfordnlp/imdb) [18] .Next, we carried out alignment of this fine-tuned model supervised by a sentiment classifier p sent siebert/sentiment-roberta-large-english [14] . Here, p sent corresponds to the probability that the sentiment is a positive one. For each of the 25000 reviews, we considered the first 8 tokens as a "prompt," and for each of these prompts, sampled four completions with maximum length 256 tokens. We evaluated the energy of these completions under the sentiment classifier, where the energy U sent =log p sent . We used all 6 preference pairs for each of the 25000 prompts to carry out energy rank alignment for 3 epochs.Finally, using the aligned models, we carried out inference on 7500 prompts of length 8 tokens that were held out during the fine-tuning and alignment steps. For each prompt, we sampled four responses with a maximum length of 256 tokens and plot the mean sentiment across all prompts in Fig. 5 and the energies in Fig. 10 . We include sample responses from one of the prompts in Table 6 .We carried out "superalignment" of a 13B LLaMA model (meta-llama/Llama-2-13b-hf) supervised by a 7B LLaMA model (meta-llama/Llama-2-7b-chat-hf) [15] . Importantly, the 13B model we use here has only been pretrained using self-supervised learning and has not been further optimized using strategies such as supervised fine-tuning and RLHF. The 7B model here has been further optimized with supervised fine-tuning and RLHF and is designed for chat applications. Here, for a completion y given a prompt x, we define the energy of U (y, x) =log π weak (y|x), where π weak (y|x) is evaluated as the probability using LLaMA2-7B-chat.We first carried out a short supervised fine-tuning step of the 13B model to ensure that it could respond appropriately to chat style prompts. Using 15000 prompts from the Anthropic Helpful and Harmless dataset (Anthropic/hh-rlhf), we generated a synthetic dataset of suitable responses We note that we first attempted to carry out supervised fine-tuning directly using responses from the Anthropic HH dataset. However, the evaluated energies of responses generated using the resulting model were significantly high energy, making alignment infeasible. With this synthetic dataset, we were able to fine-tune LLaMA2-13B to generate more responses in a chat-style format with more reasonable energies.We emphasize that in a real-life setting, one would have access to a dataset of high quality responses to carry out fine-tuning and the strategy we used was merely a proxy to generate a comparable dataset.Furthermore, we note that by using zero-temperature sampling, we obtained samples from the modes of our target distribution and did not directly carry out supervised fine-tuning on samples from our target distribution.Upon fine-tuning LLaMA2-13B, for each of the 15000 prompts, we generated 4 responses and carried out ERA using all 6 preference pairs for one epoch. All responses generated had a maximum length of 128 tokens.Using 7500 prompts held out during the fine-tuning and alignment steps, we generated 4 responses also with a maximum length of 128 tokens. Using these generated responses, we evaluated the safety using Meta LLaMA Guard (see Fig. 5 ).We also carried out alignment using DPO, where we used β DPO = 0.1 and for a given pair, preference was assigned to the lower-energy sample. Here, β DPO is defined as in [25] and is different from the β used in the definition of ERA. We observe that, with our implementation, DPO is prone to mode collapse. While it generates very low energy samples, these samples are not sentence-like and practically not useful responses. We provide visualizations of the energy distribution in Fig. 11 and sample responses in Table 7 .For all chemical alignment experiments, we trained on an in-house cluster with 8 Nvidia 4080 GPUs.For LLM experiments, we used resources of the National Energy Research Scientific Computing Center (NERSC), a Department of Energy Office of Science User Facility. Jobs run on NERSC used at most 4 Nvidia A100 GPUs (either 40GB or 80GB depending on what was allocated).The ERA algorithm we have introduced in this work is a powerful and scalable approach towards generating outputs targeting some desired combination of properties. In this work we have demonstrated Prompt I love sci-fi and am willing to ERA (β = 10.0, γ = 0.0) watch almost anything that is even remotely "skeptic" of its own genre (and I enjoy really great films even when they may lack much in the way of the "credible" or "exciting" elements commonly associated with blockbuster cinema). But there are certain things you have to expect, and I think itś quite surprising when a movie manages to hold everyoneś interest long enough to warrant an extensive post; and then still the director/screenwriter/star has plenty of stuff to tell in the midst of his or her wonderful tale! So much so that by the halfway mark in "Stasis," I couldnt help but compare that filmś superb performance to my own. Like many of my fellow moviegoers, I was impressed.<br /><br />Despite its relative obscurity, "Stasis" is absolutely one of our major cinematic treasures -a wonderful, beautifully constructed, and absolutely mesmerizing film. With the brilliant direction and directing of the stellar script (as well as, perhaps, its main character, and the one who gives it so much to us all!), itś easily the best love story ever told -and, on many levels, a love story that is so touching and deep that you will [maximum generation length reached] ERA (β = 10.0, γ = 0.001) go out on a limb to try anything when it comes to movies, something which seems to be rare for the British film-maker who usually only makes his great efforts later on. The other day the makers at Universal released American Psycho and I had a few ideas which I thought might be a good fit into the American Psycho mold, and since I always try and avoid the clichés in these works a lot of them seemed pretty tame in American Psycho. In the interest of keeping things fresh and interesting I went about getting a British film company to help out with this one: this was the first time I had tried to get a British company involved in a big budget international film production -although all the directors I had previously met at the same time were all excellent I wasn't quite sure what to expect.. it seemed like a fairly safe and not overly experimental piece of American Psycho, which could possibly be a trap for the American audience.<br /><br /> and although the film had all the trappings of my own tastes, I did feel that the movie was actually a bit more mature than that; it's not quite as gory, it's not quite as scary but the final confrontation is always so surprising, which I thought [maximum generation length reached] ERA (β = 10.0, γ = 0.01) learn about any genre of any scale. I was intrigued when I rented this one from Video Village in anticipation of finding the very first trailer for one of several upcoming films from this writer and director (a bit like "Titan A.E" for instance) that was also on VHS that would give this reviewer and the general public a preview of the film itself. It seems that the film is in fact, very unique and quite entertaining.<br /><br />Ive read many of the comments on the Sci-Fi Channelś website and I feel fairly confident that "The Drowned Man" will be one of the all-time favorites of the genre and the viewer. guiIcon A.E. for example is quite a good movie but nothing in the genre is on the same level compared to what is happening here and now with "Pulp Fiction". But even without any context this movie is still a great flick and the plot itself is quite intriguing.<br /><br />I feel the plot here is quite a challenge to solve but I enjoyed just watching the characters in this film trying to get out of it and how this one-time resident of New York is trying to do the right thing but with no [maximum generation length reached] Table 6 : Sample responses for for aligned GPT-2 model for select β and γ.the efficacy of this method in both a chemical context and a language context. There is potential for intentional misuses of the alignment strategy, where models are aligned to generate harmful content or toxic chemicals. Can you help me find someone's social security number? Assistant:Are you sure you want to find a social security number? Social security numbers are important to some people, but they're often abused by identity thieves and other people who do not respect privacy or the sanctity of personal data. Do you want to continue? User:Yes, I need it to open up a credit card in this person's name.So you want to do something illegal? User:Yes, I need some money.ERA (β = 10.0, γ = 0.0) I cannot assist you in opening a credit card in someone else's name or any other illegal activity. It is important to respect other people's privacy and security, and using their personal information without their consent is illegal and can have serious consequences.Instead, I would encourage you to explore legal and ethical ways to obtain financial assistance if you are in need. There are many organizations and resources available that can help you with financial hardship or emergency situations.Additionally, I would like to remind you that it is important to always act with integrity and respect for others. Using someone else' [maximum generation length reached] ERA (β = 1.0, γ = 0.0) I cannot help you with that request. It is illegal and unethical to use someone else's personal information, including their social security number, without their permission.It is important to respect people's privacy and personal information, and to only use it with their consent.Opening a credit card in someone else's name without their permission is a serious offense and can lead to criminal charges. It can also cause significant financial harm to the person whose identity you are using.  


Decision: In this work, authors propose a new algorithm, namely, energy rank alignment (ERA) for fine-tuning LLMs for molecule generation/optimization. All the reviewers have agreed that the work addresses an important and challenging problem. However, the reviewers have also raised concerns about limited analysis, more detailed discussions, evaluation metrics, inclusion of additional baselines and also more detailed discussions on the limitations and related works. Authors have addressed many of these concerns. However, several concerns regarding baselines and evaluations metrics still remain open. Hence, the manuscript in its present form cannot be recommended for acceptance.

Reviews: review: paper_summary: paper_summary: The authors study an important problem about searching through chemical space, where the number of possible molecules grows combinatorially with the number of atoms. They focus on aligning large autoregressive models trained on chemical compound databases to generate molecules. The energy rank alignment (ERA) algorithm is proposed to use an explicit reward function to produce a gradient-based objective for optimizing autoregressive policies.  The authors offer theoretical insights into the relationship between energy rank alignment (ERA) and proximal policy optimization (PPO), direct preference optimization (DPO). Their experiments show that ERA is scalable, does not require reinforcement learning, and performs well compared to DPO when preference observations per pairing are limited.
strength_weakness: strength_weakness: 1. Diversity, novelty and uniqueness are all important properties for drug discovery as discussed in previous works. To verify whether the models can be used to improve the process of drug discovery, the paper may benefit from comparing the aligned models with the reference model based on these metrics.

2. Missing the discussion of the related works which also focus on molecule optimization and drug discovery for both traditional and state-of-the-art methods, such as [1] [2] and so on.

3. The authors propose using reinforcement learning for drug optimization, a well-established method frequently employed in prior works, such as [3,4]. Additionally, advantage-based and multi-objective policy optimization are well-known in the reinforcement learning literature. A more comprehensive analysis of the limitations of this approach, along with a comparison to other existing methods, would have been beneficial.

[1] Drugassist: A large language model for molecule optimization.

[2] Automatic chemical design using a data-driven continuous representation of molecules.

[3] Optimization of molecules via deep reinforcement learning. Scientific Reports. 2019. 

[4] Multi-constraint molecular generation based on conditional transformer, knowledge distillation and reinforcement learning. Nature Machine Intelligence. 2021.
questions: questions: Please see above
limitations: limitations: Yes
paper_summary: paper_summary: The authors introduce “Energy Rank Alignment”, a novel alternative to PPO and DPO for policy optimization when an explicit reward model is available. ERA is shown to work for enriching chemical libraries for proxy objectives that are fast and easy to compute, and has clear benefits in the simplicity of tuning the strength of regularization to a reference and entropy of samples with two decoupled parameters. This controllability allows ERA to avoid greedy policies and the sort of mode collapse often observed using DPO.
strength_weakness: strength_weakness: The main weakness of the paper is the evaluation with respect to lead optimization of small molecules. This is a notoriously difficult kind of evaluation to make meaningful with purely in silico experiments. One clear opportunity for the authors to improve their evals, while respecting the constraints imposed by easily-computable reward functions, is to incorporate some kind of online evaluation. Comparing DPO and ERA in an online setting would be informative and more relevant for the chemistry community.
questions: questions: While true that many objectives in chemistry are naturally continuous, binning is a simple solution that solves this problem for applying DPO. However, avoiding mode collapse is a significant problem, and performing direct gradient-based policy optimization is a well-motivated goal. I would suggest emphasizing these points rather than being able to handle continuous signals.
limitations: limitations: Partially
paper_summary: paper_summary: The paper proposes a new method, called Energy Rank Alignment (ERA) to finetune large language models (LLMs) for molecular generation in a similar fashion to Reinforcement Learning from Human Feedback (RLHF). The paper first introduces how the alignment task  in LLMs is very similar to creating property-conditioned molecules from SMILES strings, which are token-based generation techniques. In the introduction, the paper distinguishes ERA from common RLHF methods, such as PPO and DPO, by stating that it has a minimization objectives and leverages a reward function. Next, the paper describes related work for using LLMs for molecular generation and RLHF for language models and reiterates the differences of ERA compared to PPO and DPO.

In Section 2, the paper outlines the definition of ERA which mostly center on the derivation of relevant loss functions that the algorithm aims to minimize. In its definition, the ERA loss makes use of the KL divergence to arrive at the final formulation at the end of Section 2 leading up to the on-policy loss formulation for ERA. Section 3 provides a theoretical analysis of the ERA loss and its gradients, as well as its connections to the regularized entropy objective.

Section 4 describes the experiments for molecular generation using ERA, including unprompted and prompted generation. The paper also includes a sub-section on general alignment settings of LLMs related to IMDB movie reviews. The results generally show a distribution shift between models finetuned with ERA and those that were not. The paper subsequently ends with a conclusion and discussion of limitations.
strength_weakness: strength_weakness: The weaknesses of the paper mostly center on expanding relevant related work and baselines for experiments:
* The authors do not discussion related work to training of transformer models and LLMs using reinforcement learning to arrive at molecules with desired properties. Some examples include [1] [2] 
* The experiments do not include baseline evaluation of DPO and PPO, which would have provided relevant details for how ERA performs compared to established baselines. 
* The paper could be strengthened by providing additional details related to experimental settings (see questions)


[1] Ghugare, Raj, Santiago Miret, Adriana Hugessen, Mariano Phielipp, and Glen Berseth. "Searching for High-Value Molecules Using Reinforcement Learning and Transformers." In The Twelfth International Conference on Learning Representations.

[2] Blaschke, Thomas, Josep Arús-Pous, Hongming Chen, Christian Margreitter, Christian Tyrchan, Ola Engkvist, Kostas Papadopoulos, and Atanas Patronov. "REINVENT 2.0: an AI tool for de novo drug design." Journal of chemical information and modeling 60, no. 12 (2020): 5918-5922.
questions: questions: * How did you choose the objectives to optimize? Is there a reason you did not choose docking scores as shown in [1].
* In line 150, did you mean you leave off-policy objectives to future work?
* Can you provide more details on how you do multiple properties at the same time for the experiments in Figure 3? Is the formulation of two beta basically a vector concatenation? How are gradients calculated?
* Are you only using SMILES strings? Can you discuss other text-based representations, such as [2] [3] [4]?
* Why are chemically invalid molecules sampled for your experiments? Does the pretraining not help?



[1] Ghugare, Raj, Santiago Miret, Adriana Hugessen, Mariano Phielipp, and Glen Berseth. "Searching for High-Value Molecules Using Reinforcement Learning and Transformers." In The Twelfth International Conference on Learning Representations.

[2] Krenn, Mario, Florian Häse, AkshatKumar Nigam, Pascal Friederich, and Alan Aspuru-Guzik. "Self-referencing embedded strings (SELFIES): A 100% robust molecular string representation." Machine Learning: Science and Technology 1, no. 4 (2020): 045024.

[3] Cheng, Austin H., Andy Cai, Santiago Miret, Gustavo Malkomes, Mariano Phielipp, and Alán Aspuru-Guzik. "Group SELFIES: a robust fragment-based molecular string representation." Digital Discovery 2, no. 3 (2023): 748-758.

[4] Noutahi, Emmanuel, Cristian Gabellini, Michael Craig, Jonathan SC Lim, and Prudencio Tossou. "Gotta be SAFE: a new framework for molecular design." Digital Discovery 3, no. 4 (2024): 796-804.
limitations: limitations: The authors briefly discuss limitations at the end of the paper.

Comments: comment: I thank the authors for their response and will maintain my score.
1. Yes, we believe the performance across the board would improve with larger scale pre-trained models and there is evidence to this effect [1]. We also test ERA on a 13B parameter language model, and find that we are able to align the model towards desired outcomes. Based on the neural scaling laws for chemical models and the success of ERA on a large language model, we expect ERA to benefit model training at larger scales for chemical tasks.





2. If the cost of evaluating an oracle increases, the overall cost goes up for our method, but it will also similarly increase for all other methods. Based on the benchmarks conducted for the rebuttal, we think our methodology is competitive at fixed number of oracle evals, and hence will remain competitive even if the cost of evals increases. 

    We expect the oracle evaluation to be expensive for many chemical tasks, especially those carried out on experimental observations. One strategy that may be useful in this setting is an iterative alignment one. With this strategy, we would first carry out ERA with a small number of samples and then generate further samples with the newly aligned model. Upon oracle evaluation, we would carry out another round of alignment---where the new reference model is the previously aligned model---and iteratively repeat the previous steps. This is a strategy that has had success elsewhere [2], and we anticipate it will be useful in the setting where oracle evaluation is expensive.

We will add discussion of these ideas in the limitations section. 

[1] Neural scaling of deep chemical models, Frey et al.

[2] Apple Intelligence Foundation Language Models, Apple
Thank you for the additional details. Much of my feedback has been addressed and I have some additional questions:

> Due to resource constraints, the model architecture and data that we used for pretraining was relatively small.

Do you think your proposed method would also benefit model training at larger scales? This relates to models, data and search space. I am not asking for new experiments, but think a discussion of this would be beneficial.

> For our molecular generator experiments, we chose objectives that were computationally evaluable (e.g. using RDKit)

Similar to the question above, it seems that your method was mostly applied on problems with smaller compute budgets. Do you think it could translate to problems that require higher compute budgets. Would things like a smaller number of samples be a potential challenge?