Aligning Transformers with Continuous Feedback via Energy Rank Alignment
Authors: Shriram Chennakesavalu
Frank Hu
Sebastian Ibarraran
Grant M. Rotskoff
Venue: NeurIPS 2025 poster ()
Keywords: transformers
alignment
preference optimization
protein language models
chemical language models
property optimization
Primary area: machine_learning_for_sciences

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.

Aligning Transformers with Continuous Feedback via Energy Rank Alignment
Shriram Chennakesavalu Department of Chemistry Stanford University Stanford CA, 94305
Frank Hu Department of Chemistry Stanford University Stanford CA, 94305
thriramc@stanfore ord edu
frankhu@stanford.
Sebastian Ibarraran
Grant M. Rotskoff
Department of Chemistry Stanford University Stanford CA, 94305
Department of Chemistry Stanford University Stanford CA, 94305 notskoff@stante edu
Abstract
Searching through chemical space an exceptionally challenging problem because the number of possible molecules grows combinatorially with the number of atoms. Large, autoregressive models trained on databases chem- ical compounds have yielded powerful generators, but still lack robust strategies for generating molecules with desired properties. This molecular search problem closely resembles the "alignment" problem for large lan- guage models, though for many chemical tasks we have a specific and easily evaluable reward function. Here, we introduce algorithm called energy rank aligmment (ERA) that leverages an explicit reward function to produce gradient-based objective that We show theoretically that this algorithm is closely related to proximal pol- icy optimization (PPO) and direct preference optimization (DPO), but has minimizer that converges to an ideal Gibbs-Bolizmann 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 DPO when the number of preference observations per pairing is small. We deploy this approach to align molecular transformers and protein language models generate molecules and protein sequences, respectively, with externally specified properties and find that it does so robustly, searching through diverse parts of chemical space.
optimize autoregressive policies.
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 Ouyang al. [2022], and controlling the outputs 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 labor-intensive Casper et al. [2023]. Reinforcement learning from human feedback (RIHF), family of algorithms that employs these human preference datasets, has been widely employed to align instruction and chat models Ouyang al. [2022], Bai et al. [2022], but it is both expensive to
options are
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

acquire the training data and difficult to carry out in practice Casper et al. [2023]. Recent algorithmic developments, such as direct preference optimization (DPO) Rafailov al. [2023], 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" tasks require evaluation that would be difficult for human evaluators.
Generative sampling problems seeking optimize reward are common in chemistry, where comparing small molecules using a particular functional assay or computationally accessible property often far easier than searching chemical space to identify novel compounds. Recent efforts to build large, domain-spectie 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.
models for chemistry Chithrananda et al. [2020] have shown promising
Our contribution: We formulate a generic alignment algorithm that we call Energy Rank Aligument 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 maximal reward value to promote sample diversity. Our algorithm enables direct gradient-basee based optimization of 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 molecular transformer model to identify highly diverse set of chemicals with properties favored by our choice of reward. Finally, demonstrate that ERA is able to align a protein language model to generate mutated protein sequences with desirable properties according to computational reward model.
Target Property
Generated Molecules
Z-le
10
High Ring Count
oro
π ref
10
TT. y x
Hydrop hydrophobicitt
10
High Polarinablitt
10
High Drug Sikenace
milel
10-2
y x
Figure 1: Energy rank alignment (ERA) enables targeting low-energy, high-reward regions with controllable fluctuations. Optimal policy approaches Boltzmann distribution with low regularization (7 + and reference policy with high regularization ( + oo) (left). Aligned models can be used to sample molecules with desired chemical properties (right).
1.1 Related Work
Inverse molecular design tasks have long history Lindsay al. [1993] and many recent works have sought to apply machine learning to facilitate this difficult search problem Sanchez- Lanchez-Lengeling and Aspuru-Guzik [2018], Gromski et al. [2019], Gómez Bombarlll et al. [2018]. While reinforcement learning has proved a popular strategy for molecular optimization Zhou et al. [2023], Sanchez- Lengeling and Aspuru-Guzik [2018], several recent studies have sought to use transformers Vaswani et al. [2017] trained on large databases of molecules represented with the text-based SMILES syntax Chithrananda et al. [2020], Schwaller et al. [2018], Wang et al. [2019], Bagal al. [2022] for such tasks. Schwaller et al. [2019] 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 Bagal et al. [2022], Chithrananda et al. [2020] and reaction prediction Pesciullesi et al. [2020], Schwaller et al. [2018].
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 built on reinforcement learning from human feedback (RLHF) Ouyang et al. [2022]. 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) Schulman al. [2017]. Rafailov al. [2023] 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, 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- objective for the optimal policy. Secondly, unlike DPO, we leverage an explicit reward function and add regularization transparently, both which help to avoid greedy policies Azar et al. [2023]. However, like both approaches, we assume that the Bradley-Tery Terry model Bradley and Terry [1952] 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 Azar et al. [2023]. investigations into the theoretical basis for contrastive methods for unlearning target datasets Zhang et al. [2024], and alternatives to the Bradley-Teree pairwise preference model Munos et [2023], An et al. [2023]. One recent study explores alignment in the context of inverse molecular design: Park et al. [2023] 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 continuous signals, which is not suitable for DPO. Overcoming this limitation while maintaining the advantages of a direct gradient-based policy optimization strategy is central goal of our current work. Our analysis and methodology directly addresses issues related to point-wise sampling because the explicit reward function climinates 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.
2 Energy rank alignment
policy is conditional probability distribution <(-|ar): : y R; we generate an output y from prompt æ. The spaces ) 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 Rref and seek to optimize a parametric, trainable policy no to adapt the conditional sampling for a particular task or constraint.
Consider prompt x E X' and model outputs y, v E )' and collection of preferences D = {(v, > y æx) ) the notation indicates that y, is preferred y/. The conditional probability that > y given can be modeled as a pairwise Boltzmann ranking within the Bradley-Tery model, i...
(U..))
p (z) (0 (( (3U y²) 3U v))
(1)
Here 8 O is a constant, (z) = (1 + - and we refer to U X x y 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 objective
J(x) Enaae If (((()))dr(()a))+ - / (1
(dg)
where 3-1 is a parameter controlling the magnitude of the entropic term, 7 sets the scale of the Kulbback-Leiler regularization compared with the energy term, and is probability distribution
3

---

over the prompts v E P(X). A proximal scheme for gradient descent on this objective corresponds to gradient flow on J Santambrogio [2017], Maas [2011]: I]: the functional can be viewed as free energy, and the corresponding flow is
TCT = (Te (
(3)
and ox denotes the Fréchet derivative with respect to π. Assuming that To has full support on X' x J/ the optimization converges asymptotically to stationary policy which satisfies
[m, =0 = T. x 0ETIE 147 U+ log Seef
(4)
and this minimizer is globally optimal. In the context of LLM alignment, representation of the energy function X' 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 (æ, y). The optimal distribution π. is Bubbs-Boltzmann measure
B
n. x)
exp
(U(x, y) - 3- log Teet (b)()))))) 11+
(5)
where (ax) is the ae-depenenn normalization constant. This expression makes clear the effect of B: when 3 - oo (low temperature), the reward dominates and fluctuations around the maximal reward small, which could lead to "mode-seeking"; when 3 -> (high physical temperature) fluctuations around the maximal reward increase and the regularization term favors proximity to Teef- Similarly, distribution proportional to e -B0 at inverse temperature B, while > 2o is dominated by the reference policy.
y + 0 recovers
Loss functions for me: Proximal Policy Optimization (PPO) optimizes indirect, proximal objective to minimize an objective closely related to (2) (cf. Appendix 3, A). Direct Preference Optimization (DPO) treats the negative reward function U implicitly and directly maximizes the likelihood of p(y > z)) Our objectives differ from both approaches: like DPO, we directly optimize the policy using explicit, gradient-based objective, but, in contrast, function directly in our objective. The losses we build are thus amenable to both offline (samples from Tref) and online (samples from Th) policy alignment, explained below. Choosing optimize the objective online has been shown to have important consequences on performance Tajwar et al. [2024]. though we focus here the setting where samples are drawn offline.
use reward
We directly optimize the Kullback-Leible divergence between the antropy-regulane preference distribution p., (y > y'(=:) and the corresponding parametric preference distribution pe (y > y [zz). Explicitly, using the fact that conditional preference distribution is normalized, we obtain
D(y'') (pylpo) (3 y (/) log KT.
(v > y' [x)
y|æ) log pr > y/z)
+ py
poly y [a))
po (y'' > (r)
(y > y' [x)
1 - Py (y > y [x)
= =v 3/( log
+ 13 y v'(x)) log 1 poly > V
poly > y [a:)
where
( B (U(z,'') (n. + B 1 log 1+
Tref ( Trot (y' x>)
Py:
This quantity is a well-defined KL divergence and is hence non-negative; the quantity vanishes when P-y = Po the observations y, y' Furthermore, with access to explicit reward model, all terms
in (6) can be computed directly and
To ( (x)
me (z)
po yy u'|=²) No ( Toly"
log
(8)
To (y a.)
To obtain a minimizer of the regularized objective defined in (2) we optimize
C CERA (mee) Ee D
(y.v') (p-l
(9)
If the current policy overlaps with the target preference distribution, may be useful to sample directly from the partially aligned policy, i.e., to use the "on-policy" formulation,
CERA (TO) EE D.yy ~wo ( D
(..'") KL (p- (po)
(10)

---

instead of (9). One issue that arises with this scheme is differentiation with respect to the parameters of the policy 0 because y and y' are decoded into discrete tokens, an operation that differentiable. To remedy this, we importance sample with reference policy
CERA (TT) = ..EDDEB D. /
T Toly
(y.v') DKL (pylpo)) lpo). xt)
(11)
Trof (3/ a Toor (
This reweighting is straightforward and the importance weights should generally be appreciable, especially early in training when To has not drifted far from Tref- is, of course, also natural to iteratively update Te using a previous iterate the reference policy. In this work, we only use (9) as an objective and leave the on-policy objectives to future work. For an ablation of the parameters of ERA and direct comparison task performace DPO, Section C.3.1.
3 Theoretical Analysis
To understand the ERA loss function and its connection to the entropy regularized objective (2), we begin by establishing that the minimizers of (9) are of the form (5). We first define the notion of equivalence precisely.
Definition 3.1 The conditional probability measures T[- x)) and n' (-1x)) are conditionallyy equivalent Vx E X, π and N' are such that supyey (yyl - T (1 æ) = 0.
(x)
We remark that this strong form of equivalence is appropriate on the finite, discrete spaces X' and y we consider here.
Lemma 3.1 If is conditionally equivalent to n', then Ty((1x>) x T((+2)))(()) is conditionally equivalent to π for all functions g : X E such that sup=ex le +00.
We prove Lemma 3.1 in Appendix A and use this simple lemma to prove the following result.
Proposition 3.2 Suppose r(-(a)) E P()) and that supp(r) supp(mret). Let 0, y 0 and that the reward model is such that supp,yexxy Je- U(a.))) conditionally equivalent to me.
< +00. Then, the minimizer of CRRA is
x exp(- 14y (U(x. y) 8- log (y)())) + g())) minimizes the objective. Because CERA is non-negative, suffices show that for all pairs y, y'. DIL.
First, we verify that any probability measure Tg (z)(z)
8
(1 (p,,)) O. This follows immediately from the cancellation in the preference probability p., of (=())) after factorization (5).
=2) where have taken g(x) = O without loss of generality and π := T.. Assume that for all pairs y', the divergence
Now, suppose that (yylz) exp (- (U(x,)) 4) 3 1 log Tee (y)(y))))
DRi (ww') KI, (pylPo) = 0 which is required of a minimizer. Equivalently, it must be the case that for all y, V.
π x))
n.. (3) T- (- (z)
(v((x)
(12)
y ( (3/(m))
T (V'
(y)ar) (y(
from which we see that
n [x) 1+ (U (m. ) 8 by log (e)r
(U.())))) 8-1 log
(13)
By construction, (yy(=) does not depend on y' so the prefactor must be purely function of æ. which completes the proof, using Lemma 3.1. A detailed theoretical analysis of the ERA objective provided in Appendix A. This analysis compares the ERA loss with the DPO and PPO objectives and demonstrates that ERA yields correct global rankings in the low data limit, while DPO maximizes pairwise margin.
5

---

0.0150
Unaligned
Unaligned
Aliund (a 50)
Aigned (j 1)
0.0125
Aligned (a 180)
0.3
Aligned (M 8)
0.0100
0.2-
0.0075
0.0050
0.0025
0.0000
0.0
100
200
300
100
500
15
Molar Refractivity
Ring Count
Unaligned
Unaligned
0.20
Aligned (8 = 5.0)
Aligned (fx 0)
Aligned (i = 10.0)
Aligned (μ 7.5)
Aligned (j3 20.0)
0.15
0.10
1
0.05-
0.00
0.0
0.2
0.4
0.6
0.8
-10
10
Quantitative Estimate of Drug-Likenes Likeness
Wildman Crippen LogP
Figure 2: Unprompted molecular generator alignment. Distributions of different chemical properties for molecules sampled from aligned and unaligned policies. The center the harmonic potential, μt varied for MR (3 1.0), Ring Count (8 1.0), and LogP (3 10.0), while 3 is varied for QED, All experiments were run with no regularization to the reference policy (r = 0).
Experiments
We test ERA on both chemical and language tasks to shed light on the following questions: 1) Can ERA robustly fine-tune model to generate samples according desired distribution? 2) What is the effect of changing the inverse-temperature / during ERA? 3) Do we maintain sample diversity (and validity) without regularizing remain close 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?
4.1 Generating small molecules with desired properties
We use decoder-only representation for the molecular generator Bagal al. [2022], where the generator has 2 layers, embedding dimension of 512, vocabulary of 324 tokens, and totals 3.5M parameters. Starting from random initialization, we carry out pretraining on dataset of 2.4M small molecules from the ChEMBL. database Zdrazil et al. [2024] for 180 epochs. For sampling from our molccular generator, we use top-k sampling with l and a sampling temperature of T in all experiments for consistency. This version of the model is not conditioned a small molecule (represented as a SMILES string) given just a start-of-ser of-equeeeeeeeece token.
prompt and generates
Central to ERA is, of course, access to computable energy function. As proof-of-concee we first consider 5 different properties for which the corresponding energy function easily evaluable: Quantitative Estimate of Drug-Likenen Likeness (QED) Bickerton et al. [2012], Wildman Wildman-Crrppen Crippen LogP (LogP) Wildman and Crippen [1999]. Ring Count, Molar Refractivity (MR) Wildman and Crippen [1999]. and Tanimoto Similarity Rogers and Tanimoto [1960] (Section 4.1.1, 4.12). Briefly, LogP is a measure of the hydrophobicity a molecule, MR a measure of the polarizability of the molecule, and Tanimoto similarity is measure of the similarity between two molecules (see Appendix C.2). We

---

then investigate ERA in more challenging context: generating molel-molscules with high predicted bioactivity for two kinases (Section 4. 4.1.3).
4.1.1
Unprompted molecular alignment RDKit oracles
First, we independently target four different properties using ERA with unprompted molecular generator (Fig. 2). Using the pretrained model as our reference policy, we generate dataset
D = -yy (i) (8) (y (8)) (8) ) and carry out energy rank alignment on TTe where Te is U2
initialized using the weights of Tref- Here, y1, U2 ~ Wref and y and (t) 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 3 to assess its impact on alignment (see Tables S2, S3). 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 3 on the learned policy; with increased 3, the learned policy concentrates around low-energy samples (i.e. near QED 1). and with lower 3, the learned policy samples greater range of QED values, as expected. We note that for each of these
four experiments, we did regularize towards the reference policy (i.e. y = 0). Even so, we were able to maintain both sample diversity and appreciable sample validity (see Fig. S8 and Table S6).
1.0, Agus 1.0
Prgg 1.0, AyaD 10.0
BLa 1.0, Aquu 20.0
Arage 1.0, 9ugD 100.0
1.00
0.75
0.50
QED
18:00
2Loe 5.0, BOUD= = 1.0
L 3.D. AOED 10.0
LLy 5.0, BOUD 20.0
Brg D.D, foED 100,0
0.50
0.25
6.00
Wildman-Crippen Log P
Figure 3: Unprompted multi collii-propety molecular generator alignment. 2D histograms of LogP versus QED for different combinations property-specifiy / illustrating a clear trade-off when performing multi-property alignment. Relative increases in / for a given property target higher values for that property. All experiments were run with no regularization to the reference policy (7 0).
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 3: U RQEDUQQE + BtogpU/in (see Tables S2, S7 for details on the energy function). We carry out ERA with different pairs of (BQED: B1o0g)) 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 croperty-specific B. Ultimately, increasing the for an individual property enables to favor higher values of that property in multi- puuti-propprtt 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. S9 and Table S7).
4.1.2 Prompted molecular alignment on RDKit oracles
Inspired by the task of lead optimization in drug discovery efforts Keserü and Makara [2009], we ask whether we can use ERA to train a molecular generator that can sample molecule that 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.3) and use this fine-tuned model reference policy for all prompted molecular alignment tasks. This reference
7

---

policy disproportioattly samples molecules that are identical (i.e. a Tanimoto similarity of 1.0) to the prompt molecule (see Fig. S10), so we carry out mulli-properry alignment on this reference policy to generate molecules that are similar- but not identical to the prompt molecule and also have high drug-likeness as measured by QED. Using ERA, we optimize the reference policy (8) x ()), (y2 (8) ()), U(y) (8) ).)), U (4)
with a generated dataset D = {(yl
. ) where we sample four molecules for each prompt molecule from the reference policy and consider all possible
preference pairs for total of six preference pairs per prompt molecule (see Appendix C.2 for full
details on the energy used).
We observe that the per-prompt average QED under the optimized policy for given prompt is higher than the corresponding average under the reference policy (Fig. S10). Furthermore, see that are able to sample a diverse set of molecules that are chemically similar to the prompt molecule, and also chemically valid (see Fig. S11, Table S8). 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-prompll 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 the reference policy (7 > 0). 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. S10 are smaller relative to the increases in Fig. because the samples are now conditioned to remain proximal to the prompt molecule, which restricts the chemical space that can be explored.
4.1.3 Unprompted molecular alignment on protein-ligante docking oracles
GSK33 top-100
JNK3 top-100
mean score
IntDiv
mean score
IntDiv
ERA
0.996 1 0.000
0.219 1 0.002
0.987 1 0.001
0.264 1 0.005
MoIRL-MGPT
1.000 = 0.000
0.362 0.015
0.961 0.010
0.372 + 0.025
GFlowNet
0.649 + 0.072
0.715 + 0.104
0.437 0.219
0.716 + 0.145
GraphGA
0.919 = 0.016
0.365 = 0.024
0.875 + 0.025
0.380 + 0.015
JT-VAE
0.235 + 0.083
0.770 + 0.067
0.159 + 0.040
0.781 + 0.127
REINVENT
0.965 1 0.011
0.308 + 0.035
0.942 + 0,019
0.368 + 0.021
Table I: Mean scores and internal diversities (IntDiv) of experiments on GSK33 and JNK3 tasks averaged across 5 random seeds. For each task, 20k molecules were sampled, and metrics were computed on top-100 scoring valid, novel and unique molecules filtered from the initial 20K samples (i.c. molecules not in dataset and molecules not previously sampled). Compared to sfatc-of-the-art methods, ERA samples more diverse molecules with higher predicted docking scores. Results for compared methods are reproduced from Hu et al. [2023].
We next investigate the performance of ERA in designing compounds that have high predicted docking scores for the kinases JNK3 and GSK333 For each of these targets, we use an in silico oracle that predicts docking scores, ranging from 0 to 1, where a higher value corresponds to stronger predicted score Sun et al. [2017]. Using only data from ChemBL, we first carry out short supervised fine-tuning step on all molecules in ChemBL with an oracle score above 0.5 (7386 molecules for JNK3 and 43381 for GSK3,9). Using this fine-tuned model as our reference policy, we then earry out alignnment using ERA (8=100 and >=0), where we use a comparably high /3 to target molecules with high activity. As with the QED alignment runs in Section 4. 1.1, we define the energy for this task as the negative logarithm of the oracle score.
From the aligned models, we sample 20000 molecules (see Fig. S12) and tabulate metries of the 100 performing molecules (see Table 1). We note that the molecules in the top-100 are both novel and unique after filtering to exclude any molecules that are present in the ChemBL dataset and any repeated molecules. For GSK38, our mean score is marginally lower than the best performing method but the diversity in sampled molecules significantly higher (i.e. lower IntDiv). For JNK3 our mean score is significantly higher than the best performing method and the diversity in sampled molecules is higher than any method. The inference costs are notably low for our approach; sampling 20000 molecules and filtering takes only minutes on single GPU.
8

---

We additionally measure sample efficiency using the top- 10 AUC metric Gao et al. [2022], Bou et al. [2024], which is the area under the curve (AUC) of the mean property value of the top-10 performing molecules versus the number of oracle calls (see Fig. S13 and Table S9). We likewise only include novel, unique, and valid molecules in this analysis. We observe that we are able to generate novel and unique high-scoring molecules, with high sample efficiency especially in comparison to existing state-of-the-att methods such as REINVENT, GraphGA, PPO, and PPOD Gao et al. [2022], Bou et al. [2024]. Ultimately, high sample efficiency is crucial in settings where evaluation is expensive, which will generally be true for most real- world chemical and biological tasks (e.g. wet-lab experiment). Finally, we also perform Glide Standard Precision Friesner al. [2004] docking the top-scoring molecules according to the oracles (score of 1.0) against their respective receptors. We observe that the diverse set of sampled molecules exhibit chemically plausable docked poses obtained from physics-based docking approach (Fig. 4).
GSK33
JNK3
Visualization of three generated ligands docked against the GSK3,3 kinase target (top) and three generated ligands docked against the JNK3 kinase target (bottom). each case, these were the three molecules with the best (most negative) Glide Standard Precision docking scores and oracle scores of 1.0.
Figure 4:
4.2 Directed evolution of proteins with ERA
1.0
0.8
All Sequences
0.6
Unaligned (ESM3)
Aligned 10)
Aligned [6 1)
0.4
Aligned (ā = 0.1)
0.0
-40
-30
-00
- 10 0
10 20
EV Mutation Score
Figure 5: Alignment of ESM3-1.44 with 3=0, 0.1, 1.0, 10.0 and 7=0.001 on the task of maximizing EVmutation score. Positions 182, 183, 184, and 186 of the TrpB parent sequence were masked and ESM3-1.4B predicted amino acids at those sites. The distribution of the EVmutation scores for generated sequences shifts significantly as increased.
We also consider the performance of ERA in a large-molecule setting, namely ML-guided directed evolution proteins. Directed evolution campaigns aim to optimize protein sequence toward some desired property of interest via iterative mutagenesis, library screening, and selection of best variants

---

Wang et al. [2021]. This has become widely used methodology in protein engineering but comes with key limitations. The inherently iterative nature of directed evolution campaigns can lead to costly and time-consuming experimental campaigns, and meaningfully understanding the effects of protein mutations on protein activity can often be difficult. These challenges have led the application of machine learning methods to more efficiently guide directed evolution campaigns Yang et al. [2024] Shanker et al. [2024]. Given the success of ERA in guiding the optimization of small-molecules using SMILES language representation, we examined whether ERA could be used to optimize large protein molecules using a protein language (i.e. primary sequence) representation.
There has been significant recent effort to design and train large protein language models (PLMs) Lin et al. [2023]. Hayes et al. [2024]. Furthermore, these models have demonstrated remarkable capabilities across number of protein tasks Widatalla et al. [2024]. Shanker et al. [2024]. As such, we decided to use the state-of-the-art ESM3-1 1.BB Hayes et al. [2024] as our pretrained model, for which we carried out alignment using ERA. Despite the multimodal nature of ESM3, here, we only focus on generating
Seennary-squence--aea based representations of proteins.
We consider directed evolution of the B-subunit of tryptophansymthase (TrpB) from Thermotoga maritima, an enzyme that catalyzes tryptophan production Buller et al. [2015]. Here, we seek to evolve the protein to increase its evolutionary fitness. In this work, we do not have access to experimental validation and so we evaluate the fitness of sequences using the computationally evaluable EVmutation score, an oracle that is predictive of variant sequence's 's performance relative to the parent sequence in its native function Hopf et al. [2017].
As in other directed evolution campaigns for the TrpB protein Yang et al. [2023], we consider mutating four different sites to one of the 20 standard amino acids. We randomly sampled 512 mutated sequences, emulating a random mutagenesis experiment. Using ESM3-1.4B as our reference model, we carry out alignment using ERA with various 8 = (0.1, 1.0, 10.0) and 7 = 0.001 and plot the results in Fig. 5 (see Appendix D for more details). We observe that with higher 3, we are able to sample mutants with the highest possible EVmutation score in single round of alignment. These results are promising for the application of ERA in directed evolution campaigns and future work will focus on the guidance of wet-lab directed evolution campaigns in conjunctions with muuti-round, on-policy ERA.
5 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 application-specific transformers for chemical search problems and protein language models. The algorithm exhibits strong performance with a variety of reward models, even ones with relatively weak signal. 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 the reference policy that we add is controlled by a parameter 7, while the entropy of the target distribution is independently tuned by a distinct parameter 3. This means that we can avoid greedy policies by keeping / small amplifying fluctuations around the optimum of the reward model -hile reducing the influence of the reference policy by taking y small. Our objective leads to easily interpretable sample-wiss gradients which highlight the importance of reward model relative to DPO in the sampled objective.
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, 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 Coley et al. [2018].
10

---

References
@I Meta. Llama model card. 2024. URL main/MODEL_ CRRDDD md.
G. An, J. Lee, X. Zuo, N. Kosaka, K.-M. Kim, and H. O. Song. Direct Preference- based Policy Optimization without Reward Modeling. Advances in Neural Information Processing Systems, 36: 70247-70266, Dec. 2023.
M. G. Azar, M. Rowland, B. Piot, D. Guo, D. Calandricllo, M. Valko, and R. Munos. A General Theoretical Paradigm to Understand Learning from Human Preferences, Nov. 2023,
V. Bagal, R. Aggarwal, P. K. Vinod, and U. D. Priyakumar. MoIGPT: Molecular Generation Using a Transformer-Decodee Model. Journal of Chemical Information and Modeling, 62(9):2064-2076 May 2022. ISSN 1549-9596.
doi: 1..0///cs.jiiiii 1c00600.
Y. Bai, A. Jones, K. Ndousse, A. Askell, A. Chen, N. DasSarma, D. Drain, S. Fort, D. Ganguli, T. Henighan, N. Joseph, S. Kadavath, J. Kernion, T. Conerly, S. El-Showk, N. Elhage, Z. Hatfield- Dodds, D. Hernandez, T. Hume, S. Johnston, S. Kravec, L. Lovit,, N. Nanda, C. Olsson, D. Amodei, T. Brown, J. Clark, S. McCandlish, C. Olah, B. Mann, and J. Kaplan. Training a Helpful and Harmless Assistant with Reinforcement Learning from Human Feedback, Apr. 2022.
G. R. Bickerton, G. V. Paolini, J. Besnard, S. Muresan, and A. L. Hopkins. Quantifying the chemical beauty of drugs. Nature Chemistry, 4(2):90-98, Feb, 2012. ISSN 1755-4330, 1755-4349. doi: 10.1038/nchem.122 1243.
A. Bou, M. Thomas, S. Dittert, C. Navarro, M. Majewski, Y. Wang. S. Patel, G. Tresadern, M. Ahmad, V. Moens, W. Sherman, S. Sciabola, and G. D. Fabritiis. Acegen: Reinforcement learning of generative chemical agents for drug discovery. Journal of Chemical Information and Modeling, 64 (15):5900-5911, 2024. doi: 10.1022/aacsjcim4440009999 URL htttp::/ppbbl acs. org/ /do/100. 1021/acs. jtcaaacss.c.m...moccmmmm...c...a 4c00895.
R. A. Bradley and M. E. Terry. Rank analysis of incomplete block designs: the method of paired comparisons. Biomeirika, 39(3/4):324-33) 1952. ISSN 0006-3444. doi: 10.2307/2334029,
A. R. Buller, S. Brinkmann-Chen, D. K. Romney, M. Herger, J. Murciano-Calles, and F. H. Arnold. Directed evolution of the tryptophan synthase 3-subunit for stand-alone function recapitulates allosteric activation. Proceedings of the National Academy of Sciences, 112(47): 2015, doi: 11.77//paa.151144111111 pnas. 1516401112.
URL https: /wwww//wwww///ooo/////// pnas wwtpooo////////////////// 1073/
Casper, X. Davies, C. Shi, K. Gilbert, Scheurer, J. Rando, R. Freedman, Korbak, D. Lindner, P. Freire, T. Wang, S. Marks, C.-R. Segerie, M. Carroll, A. Peng, P. Christoffersen, M. Damani, S. Slocum, Anwar, A. Siththaranjan, M. Nadeau, F. Michaud, Pfau, D. Krasheninnikov, X. Chen, L. Langosco, P. Hase, E. Biyık, A. Dragan, D. Krueger, D. Sadigh, and D. Hadfield- Menell. Open problems and fundamental limitations of reinforcement learning from human feedback, Sept. 2023.
A. Chen, S. Malladi, L. H. Zhang, X. Chen, Q. Zhang, R. Ranganath, and K. Cho. Preference learning algorithms do not learn preference rankings. In Advances Neural Information Processing Systems, volume 38. Curran Associates, Inc., 2024.
Chithrananda, G. Grand, and B. Ramsundar. ChemBERTa: Large-Scale Self-Supervised Pretraining for Molecular Property Prediction. In Machine Learning for Molecules Workshop at NeurIPS, 2020.
C. W. Coley, L. Rogers, W. H. Green, and K. F. Jensen. SCScore: from a Reaction Corpus. Journal of Chemical Information and Modeling, 58(2):252-261, Feb. 2018. ISSN 1549-9596. doi: 10.102 /acc..cci.....OUU0 7b00622.
Synthetic Complexity Learned
R. A. Friesner, J. L. Banks, R. B. Murphy, T. A. Halgren, J. Klicic, D. T. Mainz, M. P. Repasky, Shaw, P. Francis, and P. Shenkin. Glide: A new approach for rapid, accurate docking and scoring. 1. method and assessment of docking accuracy. Journal of Medicinal Chemistry, 47(7), 2004. doi: https://doi.or/// 1021m/00000/ URL. https:/ / //bbbs acs. wet..ogg//oooo org/doi/100 1021/jm030640..
E. H. Knoll, M. Shelley, J. K. Perry,

---

W. Gao, T. Fu, and J. S. andConnor W. Coley. Sample efficiency matters: A benchmark for practical molecular optimization. Advances in Neural Information Processing Systems, volume 36. Curran Associates, Inc., 2022.
P. S. Gromski, A. B. Henson, J. M. Granda, and L. Cronin. How to explore chemical space using algorithms and automation. Nature Reviews Chemistry, 3(2):119-128, 2019.
R. KömezzBombarrlli,, N. Wei, D. Duvenaud, J. M. Hernández-Lobato, B. Sánchez-Lengeling, T. D. Hirzel, R. P. Adams, and A. Aspuru-Guzik, Automatic
D. Sheberla, J. chemical design using data-driven continuous representation of molecules. ACS Central Science, 4(2):268 (22) Feb. 2018. ISSN 2374-7943. doi: 10. 10.00//accccetscc..oo wtt.eemtttccccco5722.... 7b00572.
T. Hayes, R. Rao, H. Akin, N. J. Sofroniew, D. Oktay, Z. Lin, R. Verkuil, V. Q. Tran, J. Deaton, M. Wiggert, R. Badkundri, Shafkat, Gong, A. Derry, R. Molina, N. Thomas, Y. Khan, C. Mishra, C. Kim, L. J. Bartie, M. Nemeth, P. D. Hsu, T. Sercu, S. Candido, and A. Rives. Simulating 500 million years of evolution with language model, July 2024.
T. A. Hopf, J. B. Ingraham, F. J. Poelwijk; C. P. I. Schärfe, M. Springer, C. Sander, and D. S. Marks. Mutation effects predicted from sequence co-variation. Nature Biotechnology, 35(2): 2017. doi: 10.1038/nbt.376. URL https:/ /wwp nature com/ articceseibttttteeeeeeeeeeeeeeee 3769.
X. Hu, G. Liu, Zhao, and H. Zhang. De novo drug design using reinforcement learning with dynamic vocabulary. In Advances in Neural Information Processing Systems, volume 37. Curran Associates, Inc., 2023.
K. Huang, T. Fu, W. Gao, Y. Zhao, Y. Roohani, J. Leskovec, C. W. Coley, C. Xiao, J. Sun, and M. Zitnik. Therapeutics data commons: Machine learning datasets and tasks for drug discovery and development. Proceedings of Neural Information Processing Systems, NeurlPS Datasets and Benchmarks, 2021.
G. M. Keserü and G. M. Makara. The influence of lead discovery strategies on the properties of drug candidates. Nature Reviews Drug Discovery, 8(3):203 212, Mar. 2009. ISSN 1474-1776, 1784. doi: 10. m0.8/////2...
Z. Lin, H. Akin, R. Rao, B. Hic, Z. Zhu, W. Lu, N. Smetanin, R. Verkuil, O. Kabeli, Y. Shmueli, A. Dos Santos Costa, M. Fazel-Zarandi, T. Sercu, S. Candido, and A. Rives. Evolutionary- scale prediction of atomic-level protein structure with a language model. Science, 379(6637); 1123-1130, Mar. 2023. ISSN 0036-8075, 1095-9203. doi: 10.1 1255 URL https://sa /ttw//wwwwo/
http::://www..ceeneeee.e www.scciecc.............e...ee.c wrg/ 2. ade2574.
R. K. Lindsay, B. G. Buchanan, E. A. Feigenbaum, and J. Lederberg. Dendral: A case study of the first expert system for scientific hypothesis formation. Artificial Intelligence, 21(2)2209-261, June 1993. ISSN 00043702, doi: 10. 0000600043377229999000666
J. Maas. Gradient flows of the entropy for finite Markov chains. Journal of Functional Analysis, 261 (8):2250-222) 2292, Oct. 2011. ISSN 0022- 1236. doi: if0.6/ijfa..2111.000
R. Munos, M. Valko, D. Calandriello, M. G. Azar, M. Rowland, Z. D. Guo. Y. Tang, M. Geist, T. Mesnard, A. Michi, M. Selvi, S. Girgin, N. Momchey, O. Bachem, D. J. Mankowitz, D. Precup, and B. Piot. Nash Learning from Human Feedback, Dec. 2023.
I. Ouyang, J. WI, X. Jiang, D. Almeida, C. Wainwright, P. Mishkin, C. Zhang, S. Agarwal, K. Slama, A. Ray, J. Schulman, J. Hilton, F. Kelton, L. Miller, M. Simens, A. Askell, P. Welinder, P. F. Christiano, J. Leike, and R. Lowe. Training language models to follow instructions with human feedback. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh, editors, Advances in Neural Information Processing Systems, volume 35, pages 27730-27744 Curran Associates, Inc., 2022.
R. Park, R. Theisen, N. Sahni, M. Patck, A. Cichońska, and R. Rahman. Preference Optimization for Molecular Language Models, Oct. 2023.
G. Pesciullesi, P. Schwaller, T. Laino, and J.-L. Reymond. Transfer learning enables the molecular transformer to predict regio- and stereoselective reactions on carbohydrates. Nature Communica- tions, 1187448874 Sept. 2020. ISSN 2041-1723. doi: 10. 10.103/641446-0200 -00--7---------17-
12

---

R. Rafailov, A. Sharma, E. Mitchell, C. D. Manning. S. Ermon, and C. Finn. Direct preference opti- mization: Your language model is secretly reward model. In A. Oh, Neumann, A. Globerson, K. Saenko, M. Hardt, and S. Levine, editors, Advances in Neural Information Processing Systems, volume 36, pages 53728-53741. Curran Associates, Inc., 2023.
D. J. Rogers and T. T. Tanimoto. A Computer Program for Classifying Plants: The computer is programmed to simulate the taxonomic process of comparing each case with every other case. Science, 132(3434):1115-11 132.3434.1115.
Oct. 1960. ISSN 0036-8075, 1095-9203. doi: 10. 2..1116//cienc.
B. Sanchez-Lengeling and A. Aspuru-Guzik. Inverse molecular design using machine learning: Generative models for matter engineering. Science, 36((6400)360-336, July 2018. doi: 10. 126/ science.nat2663.
Santambrogio. (Euclidean, Metric, and Wasserstein| | gradient flows: An overview. Bulletin of Math- ematical Sciences, 7(1):87-154, Apr. 2017. ISSN 1664-3615. doi: 10. 10.10077113333-01--0000-1-
Schulman, Wolski, P. Dhariwal, Radford, and O. Klimov. Proximal Policy Optimization Algorithms, Aug. 2017.
Schwaller, T. Gaudin, D. Lányi, C. Bekas, and T. Laino. "Found in Translation": Predicting outcomes of complex organic chemistry reactions using neural Chemical Science, 9(28):6091-6098, 2018. doi: 00.1039/C8SC02339E.
models.
P. Schwaller, T. Laino, T. Gaudin, P. Bolgar, C. A. Hunter, C. Bekas, and A. A. Lee. Molecular transformer: A model for uncertainty-callivatee chemical reaction prediction. ACS Central Science, 5(9):1572-5883, Sept. 2019. ISSN 2374-7943, 2374-7951.
doi: 10. 10.100//acxceemssc.000007777
V.R. Shanker, T. J. Bruun, B. L. Hie, and P. S. Kim. Unsupervised evolution of protein and antibody complexes with a structure-informed language model. Science, 385(6704):46-53, 2024. doi: adk8946.
http://ulul//11.1//ssoiocee/o URL https: http///wwwwincee////////// org/ borg//oi//ai orgg//oiffll//0oooo./
J. Sun, N. Jeliazkov, V. Chupakhin, J.-F. Golib-Dzib, O. Engkvist, L. Carlsson, J. Wegner, H. Ceulemans, Georgiev, V. Jeliazkov, N. Kochev, T. J. Ashby, and H. Chen. Excape-db: integrated large scale dataset facilitating big data analysis in chemogenomies. Journal of Cheminformaties, 917)), 2017. doi: https://doi.org/11.. 100111666113321117--00000 //jcheminf biomedcentrall com/articles/10....
URL https:
Tajwar, A. Singh, A. Sharma, R. Rafailov, J. Schneider, T. Xie, S. Ermon, C. Finn, and A. Kumar. Preference Fine- Tuning of LLMs Should Leverage Suboptimal, On-Policy Data, Apr. 2024.
A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. Kaiser, and I. Polosukhin. Attention is All you Need. In Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017.
S. Wang, Y. Guo, Y. Wang, H. Sun, and J. Huang. SMILES-BERT: Large Scale Unsupervised Pre- TeeeTTainng for Molecular Property Prediction. In Proceedings of the 10th ACM International Conference Bioinformatics, Computational Biology and Health Informatics, BCB '19, pages 429-436, New York, NY, USA. Sept. 2019. Association for Computing Machinery. ISBN 978-1- 4503-6666-3, doi: 011114//33773993321111
Y. Wang, P. Xue, M. Cao, T. Yu, S. T. Lane, and H. Zhao. Directed evolution: Methodologics and applications. Chemical Reviews, 121(20): 2021. doi: URL https: //ppss////bbu/// acs wwtp////ffll//l.oo//// 1021/ chemrev. 1c00260.
cc000ev.1c000600
T. Widatalla, R. Rafailov, and B. Hie. Aligning protein generative models with experimental fitness via direct preference optimization, May 2024.
S. A. Wildman and G. M. Crippen. Prediction of Physicochemical Parameters by Atomic Contribu- tions. Journal of Chemical Information and Computer Sciences, 39(5):868-8)) 873, Sept. 1999. ISSN 0095-2338, 1520-5144 142. doi: 10.1021/c19903071.
13

---

J. Yang, J. Ducharme, K. E. Johnston, F.-Z. Li, Y. Yue, and F. Н. Amold. Decoil: Optimization of degenerate codon libraries for machine surning-assister protein engineering. ACS Synthetic Biology, 12(8):2444 2454, 2023. doi: 10.1021/40 (002/aassyybooo3000000 URL https: //ppss///uuuu... acs.
htggdoi/full//// /ttp///ss/....aaaaaooooooooo 1c00000
Yang, R. G. Lal, J. C. Bowden, R. Astudillo, M. A. Hameedi, S. Kaur, M. Hill, Y. Yue, and F. H.
Amold. Active learning-assisten directed evolution, July 2024.
B. Zdrazil, E. Felix, F. Hunter, E. J. Manners, J. Blackshaw, S. Corbett, M. de Veij. H. loannidis, D. M. Lopez, J. F. Mosquera, M. P. Magarinos, N. Bosc, R. Arcila, T. Kizilören, A. Gaulton, A. P. Bento, M. F. Adasme, P. Monecke, G. A. Landrum, and A. R. Leach. The ChEMBL Database in 2023: A drug discovery platform spanning multiple bioactivity data types and time periods. Nucleic Acids Research, 52((1):D11 180-D1 192, Jan. 2024. ISSN 0305-1048, 1362-4962. doi: 10.1093/nar/gee 1004
R. Zhang, L. Lin, Y. Bai, and S. Mei. Negative Preference Optimization: From Catastrophic Collapse to Effective Unlearning, Apr. 2024.
Z. Zhou, J. Liu, C. Yang, J. Shao, Y. Liu, X. Yue, W. Ouyang, and Y. Qiao. Beyond One- PreePreference- Fits-All Alignment: Multi-Objective Direct Preference Optimization, Dec, 2023.
14

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
The abstract clearly outlines the problem searching over vast chemical spaces as an alignment problem, gives 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.
Justification:
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 long it is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
We address several potential limitations of our work the end of Section
Guidelines:
The answer NA means that the paper has 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 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 few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or speech-to-text system might not be used reliably provide closed captions for online lectures because it fails 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 fairesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged
the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
15

---

Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
Justification: Section 2 provides a thorough theoretical analysis of the ERA algorithm, and further details are Sections A and B 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 cross- referenced.
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 provide 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.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental 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)?
Question:
Answer: [Yes]
Justification: Between Section (Experiments) and Sections D and E of the Appendix (Details for molecular generator experiments and Details for PLM experiments, respectively). we describe the experimental procedure used to obtain each result.
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 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 dataset and/or model, the authors should describe the steps taken make their results reproducible 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 be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to model. In general. releasing code and data 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 large language model), releasing of a model checkpoint, other means that are appropriate to the research performed.
specific model and empirical evaluation, may
While NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm. (b) If the contribution is primarily the architecture clearly and fully.
model architecture, the paper should describe
(c) the contribution is a new model (c.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open- open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the partieular way they provide for reproducibiliity In the case of closed-source models, it may be that access to the model is limited in
16

---

some way (e.g., to registered users), but it should be possible for other researchers have some path reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification:
We will release the code as open source upon submission of the paper.
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/
Souiides/ooeSuumisseee
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "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 NeurlPS code and data submission guidelines (https:
for more details.
The authors should provide instructions on data access and preparation, including how 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 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) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: The exact settings and hyperparameters used for training and testing are provided in descriptions of the experiments in Section 3 and in Sections D and E of the Appendix.
Guidelines:
The answer NA means that the paper does 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.
7. 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 3 and in Sections D and E of the Appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
17

---

The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall
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 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 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 figures symmetric error bars that would yield results that are out range (e.g. negative 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.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: Computational resources used are discussed in Section F 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 compute required for each of the individual experimental runs as well 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 into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https //teer//s.....eeeee cc ppppulic/tticssuuiiiiiiep
Answer: [Yes]
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 cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is a special consid- eration due to laws or regulations their jurisdiction).
10. 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 G.
Guidelines:
18

---

The answer NA means that there is no societal impact of the work performed.
If the authors answer NA No, they should explain why their work has societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deploymentt 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 intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (c.g., gated release models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [Yes]
Justification: We use public data and models for all PLM 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 challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
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?
Question:
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 URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
19

---

For scraped data from a particular source (e.g., website). the copyright and terms of service of that souree should be provided.
assets are released, the license, copyright information, and terms use in the package should be provided. For popular datasets, paperswi thcode com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of 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 encouraged to reach out to the asset's 's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
The code will be released as open-source and will have adequate documentation for reproducing experiments.
Justification:
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/model 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 anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. 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 details about compensation (if any)?
Answer: [NA]
Justification:
We did 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 contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
We did not use human subjects in this work.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
20

---

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the 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.
Declaration of LLM usage
Question: Does the paper describe the usage of LLMS if it is important, original, non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, formatting purposes and does not impact the core methodology. scientific rigorousness, or originality of the research, declaration is not required.
Answer: [NA]
Justification:
We did not use any I.LMs in this study.
Guidelines:
The answer NA means that the method development in this research does not involve LLMs as any important, original, or non-standard components.
Please refer to our LLM policy (https: //teur/ps....eeee for what should or should be described.
202M/LLM)
21

---

A Detailed Theoretical Analysis
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 conditional probability distributions क ([-1az) E 'P()) for each æ E X'. Throughout, we assume that our policies have full support on y for each æ, meaning that inf y,f [ty(z) > O. Because the spaces are discrete, we make no strong restrictions on the regularity or coerciveness of the reward model -U X X y + R. The only requirement distribution is that supm.yxxxx> ee (V( < +oo, which maintains full support of the distribution.
ensure the existence of an optimal probability
Though it plays little role in theoretical analysis, also denote by v P(2)) the probability distribution over the prompts æ.
Goals of the analysis presented here The main purpose of this section to establish that globally minimizing the loss (9) yields 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 (19) are of the form (5). We first define the notion of equivalence precisely.
Definition A.1 The conditional probability measures at- [z) and x' (+(x:) in P(Y) conditionally equivalent va E X, π and n.' are such that supyey [ [x) x' (((y))) [>) = O.
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" æ E X, which we denote v E P())), does not impact conditional equivalence and hence establishes an equivalence class of conditional probability measures that minimize (2).
Lemma A.1 If π is conditionally equivalent to n', then (-(r))))) xx > < (((()))))00)) is conditionally equivalent to a for all functions g : X + R such that sup |cala)|| < +00.
Assume that T' is normalized probability distribution. This requires that,
Z'(x)
Σ (1/+)) =
(14)
VEY
If g is such that
2z(x) ((l +
(15)
then the normalized policy T'' is clearly defined by
1
N'' x(
(16)
4y(x))
because (x) c9(4), By the assumption that supzex |e012)| < 400, all terms in these calcula- tions remain finite.
Using Lemma A. it is straightforward to prove the result in Proposition 3.2. For completeness, we re-state that result here and refer the reader to Appendix 3 for the complete argument.
Proposition A.2 Suppose (-(wx) E P(Y) and that supp(z) = supp(Treet). Let B > 0, v 2 and that the reward model is such that supa le U(a < +00. Then, the minimizer of CRRA conditionally equivalent to n..
22

---

This proposition establishes that policy minimizing the objective
CERA (TO) = BEED By mot (122)
(uw')
(palpo);
KL
log To (
x))
po :-
(3' (x))
(17)
8
Teet (3
(U(I, - (Ux v)) + + t yyoog
Py == d
1 1+7
Teee (3/ xx)
has the form
3 11+ (U(x, -B log oe (1/(z))
T. (v(z)) -- (x)) exp
(18)
We do not, however, prove that gradient descent of 0 on (17) converges to the global minimizer (18) 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.
A.1 Loss functions for :ea
Proximal Policy Optimization (PPO) optimizes an indirect, proximal objective to minimize an objective closely related to (2) (ef. Appendix A). Direct Preference Optimization (DPO) treats the negative reward function U implicitly and directly maximizes the likelihood of ply > y v/|z). Our objectives differ from both approaches: like DPO, directly optimize the policy using an explicit, gradient-based objective, but, in contrast, we use reward function directly in our objective. The losses build are thus amenable both offline (samples from Tref) and online (samples from xo) policy alignment, as explained below. Choosing to optimize the objective online has been shown have important consequences on performance Tajwar al. [2024], though we focus here on the setting where samples are drawn offline.
We directly optimize the Kullback-Leibler divergence between the entropy reguopy-regullrrized preference distribution P. (y > y' x)) and the corresponding parametric preference distribution pe (y > y'(x).
Explicitly, using the fact that conditional preference distribution is normalized, obtain
(' (p>lpo) (3 1/(xz) log KL.
( > y' [x)
+ +y'
(s/ > y log
Py P2 > y/x)
Po > y xr)
po Po > y/æ=)*
(r > y' [x))
1 - Pr > y [x)
=p > 1 log
+
p3 y ) log
py > y [x)
1 po > y [x)
(19)
where
B (U(z,y') Uax y)) + B 8 log
Tref (y)e((
py:0
(20)
1+
Tref (y'(x:)
This quantity is well-defined KL divergence and is hence non-negative; the quantity vanishes when Py = on the observations y'. Furthermore, with access to an explicit reward model, all terms
in (19) can be computed directly and
To (x)
mo (z(zr)
po y (ar)
log no (o (x))
(21)
no ( To(y'
To obtain a minimizer of the regularized objective defined in (2) optimize
C ERA (TO) = Exw B D (y.v') KL
(p-lPe);
(22)
If the current policy overlaps with the target preference distribution, may be useful to sample directly from the partially aligned policy, i.e., to use the "on-policy" formulation,
ERR (TT)) = E
D (y.v') (p-lPe) KL
(23)
instead of (9). One issue that arises with this scheme differentiation with respect to the parameters of the policy 0 because y and y are decoded into discrete tokens, operation that not differentiable. To remedy this, we importance sample with a reference policy
( Fo (y) (y'la))) D (y.v')
CRA (TO) = Em- B... o (y)l) on
KL (p)l
lpo).
(24)
Thet ( () Feet (y (x)
23

---

This reweighting is straightforward and the importance weights should generally be appreciable, especially early training when To has not drifted far from Tref- is, of course, also natural to iteratively update xo using a previous iterate as the reference policy. In this work, we only use (9) as an objective and leave the on-policy objectives future work.
A.2 Gradients of CERA.
One advantage of the ERA framework is that the objective amenable to direct, gradient- based optimization. We remark that establishing global convergence for the optimization of 0 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 still glean some insight into the optimization by examining the gradients on a samplewise basis. Using the compact notation po > y [)) de and P. (y > 0)(=)=0 O..
1- o+ 0+ 1 00 00
V Vo ERA EE Ey Tee
Veog.
(25)
The gradient is straightforward interpret on particular pair y. y': if Po y'(x) is larger than Py (y > y y'(xr) then the preference gradient is positive and gradient descent lowers the probability that y'. The opposite occurs whenever po (y y' (a:) is smaller than py(y > y (z)). The magnitude of the gradient is scaled by the degree of misspecification
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 æ, the DPO objective's '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. [2023]. In our formulation, this behavior occurs only when the reward model assigns energy of co, which is prohibited by construction in most tasks. We further discuss differences between ERA and DPO in Appendix A.4.
A.3 Comparison with PPO Objective
The free energy functional for a policy under the energy rank alignment framework can be written as an expectation
JERA[T] = Exri [U(C, U( ()) + 3-1 (1++)) 7) x (zy(x) x) log Tret (y)lo)d
(26)
involving 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 (18) for the minimizer. This objective differs from the regularized reward loss conventionally used for PPO,
(T)o(z) = Eige [ U(x ()d + + / log
<gyxx) [)or(y di du(()z))
((( (( - -1 DKL ( (re [
(27)
Ez
The minimizer of the PPO objective (27) is also Gibbs-Boltzmann measure, explicitly,
3 U(rr, y) + log Trer
n x exp
(PPO)
(28)
Here, the KL-regularization corresponds to an energy shift, as in our objective, but there is no limit in which the ideal distribution π -U is obtained for the PPO objective. This is stark contrast to our approach, which recovers the ideal distribution as 7 - 0. Furthermore, while our approach allows for a direct gradient- based optimization using (17), PPO implemented using an actor-critio framework that is difficult to tune Rafailov et al. [2023], Casper et al. [2023]. Finally, we emphasize that for ERA the ; + O, finite 3 > O, the distribution has positive entropy and is not manifestly mode-seeking; there can still be appreciable fluctuations in the output. Eliminating the effect of
SU
regularization in (28), on the other hand, requires taking B/7 - oc, which eliminates fluctuations in the distribution.
24

---

A.4
Comparison with DPO Objective
The DPO approach also seeks to optimize the objective (27). The algorithm does so by first using (28) to define an implicit reward model by solving for the U that reflects the observed preference probabil- ities. This elegant idea has had significant impact and has already been deployed in state-of-the -hrt models AI@ Meta [2024]. In many cases, the observed preference probabilities will be sampled and only perhaps only one observation of y > y will be available for each æ in the dataset. When the preference dataset only has one observation > y' prompt æ, the optimal policy requires that
DPO (y)x)
y xx) 0. and x
(29)
The sampled gradients of the objective used for DPO are proportional to the implicit reward discrep- ancy,
To(y'( log To ar) y. y. .z) =a (B ( log (rel(yy(
To ax) log
Vo DDO
(30)
Tref (yee((y)3) x))
To To (y' xx)
which when ( [z) - could lead instability as - log No (yy (x) > co. On the other hand, the ERA gradients are scaled by the relative preference discrepancy,
1 - o. > y/(x) (
0.(y > y (t)
VoL FRA (U,v',
Vooa > y l/)a
(31)
ao y')
oo v'(x))
The advantage of a reward model becomes apparent because
a. - (y'x) (y v'(z) (1+
[U(E) y') U2x y)) + y log
zee (3)
(y' |x))
(32)
and hence the optimum of CRRA will not lead to policies in which supp( degrades unless the energy becomes infinite. Choosing an appropriate reward model, hence, gives the flexibility control instability if it becomes problematic.
DPO
Chemical Validity
0.8
ERA
60
0.6
0.4
20
0.2
0.0
1000
2000
3000
1000
2000
3000
Checkpoint
Figure S6: Comparison of DPO (3pPo=0.1) and ERA (8 20.0 and 7 0.0) the task of maximizing QED (Section 4.1.1) a) Chemical validity degrades significantly for DPO and slightly degrades for ERA as function of checkpoints. b) At each checkpoint, 20k molecules are sampled and mean QED of valid molecules plotted with shading corresponding to standard deviation. Both DPO and ERA alignment runs were trained with the same hyperparameters and dataset, with checkpoints saved every 100 epochs.
We carry out an online evaluation of DPO and ERA on the task of generating amall-molecule with high QED, as in Figure 2. We align DPO using the same dataset and hyperparameters as ERA with @DPO = 0.1 and train for thousands of checkpoints (over GPU-hours). We load intermediate checkpoints for both the DPO and ERA (BERA- 20.0, TERA 0.0) runs and carry out inference (Figure S6). We observe that the first saved checkpoint of the DPO alignment run, the model generates molecules with high QED scores but with low validity (~20%)). However, upon further
25

---

training, the chemical validity of further checkpoints drops to 0% for the remaining runs, despite the overall DPO training and validation losses still dropping.
With ERA, we see that we are able to similarly sample high QED small-molecules with reasonably high chemical validity (~85%). While the validity does drop subsequent checkpoints, it does not do so precipitously. Moreover, the ERA-based alignment had no regularization (? = 0), and in the main text, we document how having non-zero 7 can enable increases in chemical validity (Figure S10). Finally, we note that we did not extensively tune the hyperparameters for DPO, and is possible that a different set of hyperparameters would elicit a more desired outcome; however, the lack of meaningful regularization in DPO Azar et al. [2023] and its performance degradation in online metries has been well-documented Chen al. [2024].
B ERA implementation
Implementing energy rank alignment is straightforward to implement within existing code bases. We provide sample PyTorch code for the ERA loss function below.
26

---

import torch. mn as mn from torch. mn. functional import logsigmoid
def era_ loa_
pi_logps ref ref_logps_1, ref_ _ogps_2, energies_1, energies_2, beta, gamma):
pi_logps_1: logprob under policys model of first sequence in pair (B,)
pi_logps_2: logprob under policys model of second sequence in pair (B,)
ref ref_logps_i: logprob under reference model of first sequence in pair (B,) ref ref_logps_2 2: logprob under reference model of second sequence in pair (B,) energies_1: energies of first sequence in pair (B,) energies_2: energies of second sequence in pair (B,) beta: inverse temperature
gamma: regularization controlling strength of KL penalty
beta_prime (beta / (1 gamma))
gamma prime (gamma (1 gamma))
logp Logrigmoiddpollicy__pppppp_
policy_logps_yye
logp_prime primee =
logppsssy_yyos_y volicy_logps..
logp_stad = logsigmoid (-beta_prime * (energies_y? energies_y1)
gaamaa_prime (ref_logps_y² ref ref_logps_y1)
Logp_star_pring = logsigmoidd-bbta_pprin * (energies_yl energies_y2)
gamma_prime (ref_logps_)
ref ref_ _og_loggs_y)))
era_loss (torch. exp (logp_ star) (logp_stee star - logp)
+ torch. exp (llog__sta_ppriinnn. * (logp_sta__prine logp-prime))
return loss. mean
C Details for molecular generator experiments
C.1 Pretraining details
In this work, we represent all molecules as SMILES strings and tokenize SMILES strings according the approach in Schwaller al. [2018]. 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 from this curated dataset consists of 324 tokens, which we augmented with start, stop, and padding tokens.
in the original ChEMBL dataset. The alphabet generated
We first pretrained a model according to a next-token prediction, self-supervised learning approach. We trained model using the standard cross entropy loss
T
LCE Σ
log po (2 x1)))
(33)
Our trained molecular generator consisted of just the encoder block of a standard multi-head attention transformer Vaswani et al. [2017]. Finally, the model had layers, heads, and width of 512. For pretraining, we used an Adam optimizer with learning rate of 1.0 10-5, We emphasize that this pretrained generator samples molecules in an unprompted fashion; given just a sfart-of-sequence token, we can autoregressively generate sequence of tokens. Moreover, it is possible that this sequence of tokens corresponds to molecule that 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 I 1500 molecules and then computing the Tanimoto similarity between every pair of molecules. We plot the distribution of all pairwise Tanimoto similarities from
27

---

ChEMBL Dataset
10.0
Molecular Generator
7.5
5.0
2.5
0.0
0.00
0.25
0.50
0.75
1.00
Tanimoto Similarity
Figure S7: Chemical diversity of samples from training dataset and from unprompted molecular generator (unaligned) as measured by pairwise Tanimoto similarities. Lower Tanimoto similarities correspond to more chemically dissimilar molecules.
Property name (f)
Energy function (U)
Tanimoto similarity
U log([(y)
QED
U = log(f(y)
Docking Oracles (GSK33 and JNK3)
U (zg(
Wildman-Crippen LogP
U (S(y) -1)/20²
Molar refractivity
= (f(y) - 220
Ring count
(f(y) -2/
Table S2: Definitions of energy functions (in reduced units) used for each of the five chemical properties investigated in this work. Here refers to the generated molecule.
this sample and from all pariwise Tanimoto similarities from 1500 randomly sampled molecules from the original dataset in Fig. S7. We observe that we can generate molecules that are quite distinct (i.e. low Tanimoto similarity) in comparison with all other molecules.
C.2 Chemical properties
We investigated aligning the molecule generator to several target chemical properties, which we detail below. All of the properties can be casily computed using either the RDKit package the tdc Huang et al. [2021] package. We list the energy function and parameters used for the corresponding energy functions for each of these properties in Table S2.
Tanimoto similarity is measure of chemical and structural properties between two molecules and ranges from 0 to 1, where higher values correspond to more similar molecules Rogers and Tanimoto [1960]. Quantitative estimation of drug-like likenesssss (QED) evaluated by taking the geometric mean of set of "desirability functions" for different molecular descriptors and also ranges continuously from values of to Bickerton et [2012], where higher values correspond more drug-like molecules. The octanol-water parition coefficient (Wildman-Crippp frequently employed in medicinal chemistry applications Wildman and Crippen [1999]. Molecules with more positive values are more hydrophobic (i.c. more soluble in octanol relative to water), whereas molecules with more negative values are more hydrophilic (i.c. more soluble water relative to octanol). Molar refractivity is similarly calculated and is positive number that serves as a measure for molecular size and polarizability Wildman and Crippen [1999]. A higher molar refractivity corresponds to larger and more polarizable molecules. Finally, ring count corresponds to the number of rings in a molecule.
LogP) is a measure of hydrophobicity
linear combination of atomic contributions,
Under the definitions of the energy functions in Table $2, 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 report these values in Table S3, Furthermore, when the computed
28

---

Property name (f)
Energy
Tanimoto similarity
10
QED
4.5
Docking Oracles (GSK33 and JNK3)
4.5
Wildman-Crippen LogP
300
Molar refractivity
400
Ring count
70
Table S3: Property-specific energy values (in reduced units) used to treat chemically invalid sequences.
QED or Tanimoto Similarity is O, the energy infinite, and to ensure numerical stability, we set the value of the energies to be 4.5 and 10 respectively. Finally, in the prompted molecular generator experiments Section 4.1.2, we assign an energy of 3.5 to the setting where Tanimoto similarity between the generated and prompt molecule is 1.0 (i.e they are the same) in order to penalize this outcome. Here, all energy and 3 values are reported in reduced units.
C.3 Molecular alignment details
C.3.1 ERA ablations and comparison to DPO
8
7
Validity
Top 100 Mean
Top 100 Diversity
10
0.01
72.00
0.932
0.129
0.1
74.23
0.933
0.131
100
0.01
74.71
0.934
0.129
100
0.1
75.14
0.935
0.131
Table S4: Descriptive statisties of 10,000 generated SMILES after aligning the molecular transformer model with ERA for QED maximization across several values B and 7/.
DPO results
ADPO
Validity (%) f
Top-100 mean f
Top- 100 diversity 4
100
86.91
0.936
0.131
80.92
0.943
0.134
82.23
0.943
0.156
0.1
54.46
0.947
0.174
ERA results (y = 9)
BERA
Validity (%) f
Top-100 mean 1
Top- 100 diversity |
0.1
84.78
0.939
0.149
85.73
0.936
0.145
10
79.07
0.932
1,141
100
88.83
0.941
0..45
Table S5: Comparison of DPO and ERA results across different / values on the task of QED maximization.
To evaluate the effects of the parameters 3 and the performance of ERA, assess the algorithm's performance on the task QED maximization by ablating 3 and y in Table S4. We observe that higher 33 leads to greater dominance of the reward and increasing values for the Top-100 mean QED whereas higher 7 leads to higher validity due to the increased regularization against the reference policy, both consistent with the minimizer presented in 5.
To assess ERA against DPO, we evaluate each method on their performance on the task QED maximization. We note that the DPO objective does not contain explicit reward component and instead only maximizes the margins between preferred and dis-preferred samples. Additional care needs to be taken when comparing DPO and ERA hyperparameters as 3ERA of (ppo due to

---

Moler Refractivity
Ring Count
10
Unxlaped
Uralgped
Abgred (x-
Abgpell (e-
Allged ( 190)
Atgel gr-
0..0
0.5
0.0
0.5
QED
Wildman Wildmen-Crippen Log P
10
Uaslagasi
Xlizsal (5) 5.)
Akgped (j 10
10
Allgeel
7.5
Altyoed (8 2010)
000
0.3
0.0
0.5
Tanimoto Similarity
Figure S8: Chemical diversity of samples from unprompted molecular generator after alignment as measured by pairwise Tanimoto similarities. (See Fig. 2, Section 4.1.1)
difference in definition. Specifically, the minimizer of the DPO objective is:
ADPO y) cxp(- + log Tref)
Whereas for ERA, the minimizer is
-BERA r(x, 3) + Y log Tref 1++
TERA (z,g) 8) x exp
1+7
The formulation of the DPO objective does not permit explicitly scaling the contribution of the regularization term against the prior while ERA permits this. This also implies that
1+7
BERA 3DPO
Furthemore, take be relatively large since as 7 + XC, 14y 1. In our comparisons, choose y 9. Table S5 contains descriptive statisties of 10,000 generated SMILES after aligning the molecular transformer model for QED maximization using either DPO or ERA across several matched values of (3. It is arranged such that the corresponding rows can be compared directly against each other. We did not scan any 33 < 0.1 since the validity in those cases was low for both methods. Due to parameter matching, we expect ERA and DPO be similar here: validity is comparable, with ERA surpassing DPO at higher SERA. as well as the top-100 average QED, ERA consistently generates diverse molecules across values of BERA as opposed DPO, as evidenced by the lower similarity among the top- 100 generations for BERA 10 and SERA 100. This again reflects the fact that the entropy-regulated ERA objective promotes diverse generations while effectively avoiding greedy policies as opposed to DPO, which focuses only on maximizing the margins between preferred and dispreferred sequences. This highlights a key advantage of ERA since good molecular diversity is a key consideration for chemical and drug discovery tasks.
C.3.2 Unprompted molecular generation (RDKit oracles)
We first investigated aligning the unprompted molecular generator sample small-molecules with desired properties. We carried out alignment using the property-specific energies described in Table S2. All alignment properties were initialized with the weights of the pretrained model and trained using an Adam optimizer with learning rate 1.0 x 10-6, We tabulate the chemical validity for pungle-propert alignment Table S6 and for multi-property alignment in Table S7. While do see a drop in chemical validity after alignment, we see that majority of the samples we generate post- 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 predictions of 1500 generated molecules. We visualize the chemical diversity for experiments in Fig. S8 and multi-properts experiments in Fig. $9. We observe that the samples are still highly diverse chemically after alignment. All plots in Fig. and Fig. were computed using 1500 generated molecules per experiment.
30

---

Property name
Hyperparameters
Chemical validity
Unaligned
N/A
88%
Molar Refractivity
8 1.0, Jr 50, 10, 0.0
82%
Molar Refractivity
1.0, μr 180, 0 = -
74%
Ring Count
3 1.0, Jr 1, 1
84%
Ring Count
1,O, = 8, 0 1.0, 7 0.0
59%
LogP
8 10.0, / 2.5,
-1.0, 0.0
74%
LogP
10.0, / --=7.4 7.5, 0 = =100, 0.0
63%
QED
3 5.0, 9 0.0
54%
QED
10.0, 7 0.0
66%
QED
3 20.0, 7 0.0
0.0
65%
Table S6: Percentage of generated sequences that were chemically valid for samples from unprompted molecular generator after alignment. (See Fig. 2, Section 4.1.1).
Unaligned
Aligned
Pug 1.0, Aarn 1.0
8 P
1.0, &gen 10.0
100, Booo 20.0
Shgp 1.0, Porn 100.0
BuuP AQEE 1.0
B 5.0, 8ggD
BLN 5.0, REED 20.0
Bus 5.0, AQED
100.0
Tanimoto Similarity
Figure S9: Chemical diversity of samples from unprompted molecular generator after multi culti-property
alignment measured by pairwise Tanimoto similarities. (See Fig. 3, Section 4.1.1).
Hyperparameters Unaligned
Chemical validity
88%
iQED 1.0, Logg 1.0, PLogP 7.5, LogP 1.0, 7 0.0
60%
BQED = 1.0, Logg = 10.0, PLogP = 7.5, LLgg = 1.0, 7 0.0
67%
BOED 1.0, Sragp 20.0, PLogg 7.5, TLogP 1.0. 7= 0.0
68%
BQED 1.0, &aggP = 100.0, MLogP 7.5, FLogg = =1.0, 0.0
63%
BQED 5.0, Saggr 1.0, /oogP -7.,, GLogP = 1.0, -0,7=0.00
64%
AQED = -5.0, Bagr = 10.0. PLagr = =7.5, oLog = -1.0, 7 0.0
62%
BOED 5.0, SLogp 20.0, #LogP 7.5, LLgP 1.0, 7 0.0
62%
BQED 5.0. SLogP = 100.0, PLogP 7.5, oLog 1.0. 7- 0.0
68%
Table S7: Percentage of generated sequences that were chemically valid for samples from unprompted molecular generator after multi multi-property alignment. (See Fig. 3, Section 4.1.1).
31

---

Multienergy Alignment (QED and Tanimoto)
Multienergy Aligmment (LogP and Tanimoto)
Unadigned Aligned
7.5
Unaligned Aligned
1.5
10-
10-
5.0
1.0
S
2.5
0.0
0.0
0.0
Tanimoto Similarity
(E FoeD (y)z))
Tamimoto Similarity
(E Logp ( 2 Stgpp (ykr))
Ex Bon((yy())
Figure S10: Prompted multi-property molecular generator alignment. From left to right: Tanimoto similarities computed between the prompt and sampled molecules for both aligned and unaligned policies (QED and Tanimoto aligmment), per-promp. difference in the average QED under aligned and unaligned policies (QED and Tanimoto alignment), Tanimoto similarities computed between the prompt and sampled molecules for both aligned and unaligned policies (LogP and Tanimoto alignment), and per-prompt difference in the average LogP under aligned and unaligned policies (LogP and Tanimoto alignment). With alignment, we target higher QED and LogP values, while still sampling molecules chemically similar but not identical to the prompt molecule.
Prompt Molecule
Generated Molecules
Prompt Molecule
Generated Molecules
Jot gr oo co
DRepy prope PRO6
2 Joues SPayes gyeuo
Figure S11: Sample molecules from prompted molecular generator after alignment experiments: QED and Tanimoto (left) and LogP and Tanimoto (right). With alignment, generated molecules are diverse, while still chemically similar to prompt molecule.
C.3.3 Prompted molecular generation (RDKit oracles)
Next, we generate small-molecules with desired properties conditioned on a prompt, where the prompt is itself another molecule. In the experiments here, we consider the setting where we generate molecules that are chemically similar to the prompt molecule. With this in mind, we first earry out a fine-tuning step using synthetic dataset 'D = { y Vi),... (æ Vn)) ) where æ corresponds to the SMILES string of a prompt molecule and y corresponds the SMILES string of the conditionally generated molecule. To curate this dataset, we consider all molecules in our original filtered ChEMBL dataset to be a prompt molecules and for each prompt molecule æs, we generate a response molecule y, by simply perturbing a random token from x;. If the perturbed sequence was chemically invalid, we repeated the random perturbation until a valid molecule was generated. The prompted generator was the same size as the unprompted molecular generator, and 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 and used this generator our reference policy for all prompted alignment experiments. All plots in Fig. S10 were computed using 100 generated
Hyperparameters Unaligned
Chemical validity
93%
BTanimoto 5.0, 3.og 10.0, MgogP = 5.0, T1.ogP 1.0, 0.1 @Tanimoto = 5.0, BQEn = 500.0, 7 = 001
91%
81%
Table S8: Percentage of generated sequences that were chemically valid for samples from prompted molecular generator after multi-property alignment. (See Fig. S10, Section 4.1.2).
32

---

molecules per prompt, where we carried inference over 500 prompts per experiment. We tabulate the chemical validity of the prompted generator in Table S8.
C.3.4 Unprompted molecular generation (protein-ligand docking oracles)
GSK33 top-10 AUC
JNK3 top- 10 AUC
ERA
0.985 - 0.001
0.989 = 0.002
REINVENT
0.865 = 0.043
0.783 0.023
GraphGA
0.788 = 0.070
0.553 0.136
PPO
0.90 = 0.02
0.80 0.04
PPOD
0.92 + 0.02
0.87 0.02
Table S9: Top- 10 AUC scores on GSK33 and JNK3 tasks averaged across 5 random seeds. Compared to tfate-of-the-as methods as reported in Gao et al. [2022] and Bou et al. [2024], ERA has higher sampling efficiency. Results for compared methods are reproduced from Gao et al. [2022] and Bou et al. [2024]
For the work here, we used computational oracle that predicts the docking score for two kinases, JNK3 and GSK33, where these oracles were defined using the tde package. We first carried out a supervised fine-tuning step using all molecules in ChemBL with an oracle score above 0.5. For the JNK3 model, we carried out fine-tuning on a dataset of size 7386 molecules, and for GSK33, we carried out fine-tuning on dataset of size 43381 using the Adam optimizer with a learning rate of 1.0 * 10-5 From these fine-tuned models, we sampled 40000 molecules, evaluated the oracle scores, and used this dataset to carry out alignment using all possible pairs of molecules. For any invalid molecule, we assign an energy of 4.5. All alignment runs were done using the Adam optimizer with a learning rate of 1.0 10 We observed sampled validity of 74% on the model aligned for GSK33 and a sample validity of 93.6% for the model aligned for JNK3.
We compute metrics on the top-100 novel and unique molecules on 20000 sampled molecules (see Figure $12) from the aligned models. We compute the mean score and the internal diversity score (IntDiv) Hu et al. [2023] computed according to the following
Iiv(A) := A[([A] - 1)
Σ
T(æ,y),
(34)
Ax.
where A is a set of compounds and T represents the Tanimoto similarity. Lower IntDiv scores correspond to more diverse molecules.
We additionally compute sample efficiency using the top- 10 AUC metric (see Table S9), which is the area under the curve (AUC) of the mean property value the top- 10 performing molecules versus the number of oracle calls (see Figure S13). We note that once we reach top- 10 average of 1.0, we do not make further oracle calls as subsequent oracle calls will not change the top- 10 average and will artificially inflate the AUC.
C.3.5 Analysis of invalid generations
Especially when the regularization = O, we see that ERA reduces validity of the generated SMILES. There some limitations imposed by the fact that most of our downstream characterizations cannot be run on invalid SMILES string because of the failure of RDKit to parse. To assess the failure modes of these invalid generations, we computed two measures of the diversity: an estimate of the Levenshtein distance and an estimate of the Shannon entropy for both valid and invalid molecules.
Table S10 demonstrates properties of the valid and invalid generations from model aligned using ERA for QED maximization for 3 = 100.0, 7 = 0.1. Notably, the Shannon entropy of generated tokens and the mean Levenshtein distance between all of the generated SMILES is higher for the invalid outputs than the valid ones. These metrics suggest that there in fact less clustering around specific region of chemical space in these failure cases.
For the same QED maximization run (3 100.0, 7 = 0.1) as above, we found a significant failure mode to be that of missing tokens to close branches (67.6% of all invalid generations). Furthermore, failure to properly close rings also occurs in a signifcant amount of invalid generations (30.5%). One potential failure mode that seems to be relatively insignificant is that of failing to predict the stop token
33

---

JNK3
GSK33
Unaligned
Unaligned
Aligned
Aligned
0.6
0.8
0.2
0.4
0.6
0.8
1.0
Oracle Scores
Figure S12: Distribution of GSK3, and JNK3 oracle scores sampled from unaligned reference model and aligned model (3 100.0, y 0.0). 20k molecules were sampled from each model and only oracle scores of valid molecules are plotted.
JNK3
GSK33
Average Top-10 Score
1.00-
1.00
0.95
0.95
0.90-
0.85
0.90
0.80
0.75
0.85
0.65
0.80
500
1500
250
500
Oracle Call Number
Figure S13: The average score of t0p-10 performing valid, novel, and unique molecules function of the number of oracle calls made to the aligned models. Scores are computed using the JNK3 and GSK33 oracles, respectively, for five different random seeds. Samples that invalid, present in the dataset, or already previously sampled are discarded and do not count towards an oracle call.
before hard limit of 500 tokens, which occurs in fewer than 1 1% of invalid generations. Beyond this analysis, visual inspection of several failure cases does not appear to reveal any particular motif that is repeated among these examples beyond the model attempting to generate highly branched molecules or molecules with several rings.
Generations
Shannon Entropy Levenshtein Distance
Valid
2.2574
39.8225
Invalid
2.3177
56. 16.6
Table $10: Shannon entropy and Levenshtein distances for valid and invalid generations from QED maximization molecular generator alignment experiment with B = 100.0, 7 0.1.
34

---

D Details for protein language model experiments
We consider mutating the TrpB protein at 4 sites (positions 182, 183, 184, and 186) all of the 20 standard amino acids and compute the EVMutation score Hopf et al. [2017] for all 20 = 160000 sequences (see Yang et [2023] for dataset). We used the standard tokenization scheme in ESM3 and during training and inference masked all non-sequence tasks Hayes et al. [2024]. ESM3 is bidirectional transformer and so to generate mutant sequences, we first pass the full sequence with the mutant sites masked and use the model to "unmask" these four sites simultaneously. For the sequence track, ESM3 also contains additional tokens representing non-standard residues and special cases, and we consider sequences "invalid" if the generated token does not correspond to one of the 20 standard amino acid sequences.
When sampling mutant sequences, the native ESM3-1 1BM3-1.4BB ,4B model does not generate diverse set of sequences and so we first synthesize an initial dataset of 512 random mutant sequences as would normally be done in random mutagenesis experiment. We did not carry out any supervised fine- tuning step and here considered the energies to be the negative of the EVMutation score and evaluated log probabilities the pretrained ESM3-1.4B model. We then carried alignment using ERA, where our dataset consisted of all possible unique pairs these 512 sequences. For the experiments here, we used the RMSProp optimizer with learning rate of 1.0 10-5. We plot the EVMutation score of 512 generated sequences across various /3 values and y 0.001 in Figure 5. Finally, we note that we did not sample any "invalid" sequences defined above.
E
Computational resources
For all chemical alignment experiments, we trained on an in-house cluster with Nvidia 4080 GPUs. For ESM3 experiments, used resources of the National Energy Research Scientific Computing Center (NERSC), Department of Energy Office of Science User Facility. Jobs run on NERSC used most Nvidia A100 GPUs (either 40GB or 80GB depending what allocated).
F Societal and broader impacts
The ERA algorithm we have introduced in this work is powerful and scalable approach towards gen- erating outputs targeting some desired combination properties. In this work we have demonstrated the efficacy of this method in both a chemical context and a language context. There is potential for intentional misuses of alignment strategy, where models are aligned to generate harmful content or toxic chemicals.
35


Reviews:
[{"review_id": "SINtQT4TnX", "reviewer": "Reviewer_L7xi", "summary": "The 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.", "questions": "See weakness", "limitations": "Yes", "rating": 2, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 2, "originality": 2, "strengths_and_weaknesses": "Strengths:\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.", "was_revised": true, "final_justification": "My concerns have increased after reviewing the authors’ response. Respectfully, I am unsure whether the authors have a deep enough understanding of RL; some design choices may appear interesting but could actually harm performance instead or introduce unnecessary complexity. As they mention, “ERA works entirely offline”, which raises concerns about whether they fully understand the environment and how to take advantage of the given environment. It is also not fair to compare against DPO in a setting where an explicit reward function is available. A thorough understanding of each component in the PPO (and DPO) objective, as well as the potential empirical effects of altering them, is essential.", "rebuttal": "In response to the questions raised in your review, we explain below both experimental and theoretical characterization of ERA to clarify the relation to existing RL techniques.\n\n> “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.”\n\nWe agree this has conceptual similarities to the entropic terms that are often used in reinforcement learning algorithms, like REINFORCE and PPO. The advantages of ERA include: it is a direct gradient-based procedure that works entirely offline, providing an important advantage for practical applications, and we do not need complicated protocols, like replay buffers or ad hoc objectives to promote sample diversity.\n\n>“With an explicit reward function, the necessity of sampling fluctuations around the maximum reward is not well-justified.”\n\nIt is well-known that the reward functions used for molecular optimization are susceptible to reward hacking in which all the probability associated with a given task concentrates on a single motif. Just as in classical RL approaches, adding an ad hoc entropy term can mitigate this undesired behavior. **The theoretical justification for adding an entropy term is to increase sample diversity in a controllable fashion, ensuring near optimality, but avoiding greedy optimization.** When the number of observed comparisons is small, DPO maximizes the margin between pairs of samples. In our experiments, this leads to both low validity of the generated SMILES and low diversity because the probability concentrates on the “winning” molecules and diversity collapses:\nWhen the preference dataset only has one observation $\\boldsymbol{y}\\succ \\boldsymbol{y}'$ per prompt $\\boldsymbol{x}$, the optimal policy requires that\n\\begin{equation}\n    \\pi_{\\star}^{\\textrm{DPO}}(\\boldsymbol{y}|\\boldsymbol{x}) = 1 \\quad \\textrm{and} \\quad \\pi_{\\star}^{\\textrm{DPO}}(\\boldsymbol{y}'|\\boldsymbol{x}) = 0.\n\\end{equation}\nThe sampled gradients of the objective used for DPO are proportional to the implicit reward discrepancy, \n\\begin{equation}\n    \\nabla_{\\boldsymbol{\\theta}} \\hat{\\mathcal{L}}^{\\textrm{DPO}}(\\boldsymbol{y}, \\boldsymbol{y}', \\boldsymbol{x}) = \\sigma \\left( \\beta^{-1}\\gamma \\left[\\log \\frac{\\pi(\\boldsymbol{y}'|\\boldsymbol{x})}{\\pi_{\\rm ref}(\\boldsymbol{y}'|\\boldsymbol{x})} - \\log \\frac{\\pi(\\boldsymbol{y}|\\boldsymbol{x})}{\\pi_{\\rm ref}(\\boldsymbol{y}|\\boldsymbol{x})} \\right] \\right) \\nabla_{\\boldsymbol{\\theta}} \\log \\frac{\\pi(\\boldsymbol{y}|\\boldsymbol{x})}{\\pi(\\boldsymbol{y}'|\\boldsymbol{x})},\n\\end{equation}\nwhich when $\\pi(\\boldsymbol{y}'|\\boldsymbol{x}) \\to 0$, could lead to instability as $-\\log \\pi(\\boldsymbol{y}'|\\boldsymbol{x}) \\to \\infty.$\nOn the other hand, the ERA gradients are scaled by the relative preference discrepancy. The advantage of a reward model becomes apparent because the optimum of $\\mathcal{L}^{\\textrm{ERA}}$ will not lead to policies in which $\\textrm{supp}(\\pi)$ degrades unless the energy becomes infinite. This could alternatively be implemented by taking $\\beta\\to\\infty$. Choosing an appropriate reward model, hence, gives the flexibility to control instability if it becomes problematic. \n\n> “Comparison to PPO: The paper does not clearly articulate ERA’s advantages over PPO when using an explicit reward function.”\n\nThe purpose of ERA is to optimize with a molecular reward function while preserving sample diversity. We compare ERA to several reinforcement learning strategies throughout, which optimize objectives that are commonly used in PPO. \n\n - REINVENT uses a policy gradient algorithm on the PPO objective described below. Our benchmarks consistently outperform REINVENT.\n - MolRL-MGPT uses multi-agent reinforcement learning with an objective identical to the one used in REINVENT. The objective requires manually balancing several terms to mitigate overfitting. **The results of these comparisons can be seen in Table 1 of our manuscript.**\n\nThis objective we introduce differs from the regularized reward loss conventionally used for PPO: the minimizer of the PPO objective is also a Gibbs-Boltzmann measure, explicitly, \n\\begin{equation}\n    \\pi_{\\star}^{(\\textrm{PPO})} \\propto \\exp \\left[ -\\frac{\\beta}{\\gamma} U(\\boldsymbol{x},\\boldsymbol{y}) + \\log \\pi_{\\rm ref}(\\boldsymbol{y}|\\boldsymbol{x}) \\right].\n\\end{equation}\nHere, the KL-regularization corresponds to an energy shift, as in our objective, but there is no limit in which the ideal distribution $\\pi\\propto e^{-\\beta U}$ is obtained for the PPO objective. \nThis is in stark contrast to our approach, which recovers the ideal distribution as $\\gamma\\to0.$\nFurthermore, while our approach allows for a direct gradient-based optimization using the direct estimator introduced in the paper, PPO is implemented using an actor-critic framework that is difficult to tune [cf. e.g., rafailov_direct_2023, casper_open_2023].\n\n> “The paper lacks an ablation study comparing ERA, PPO, and DPO under explicit reward conditions, limiting insights into their relative performance.”\n\nWe provide here an ablation study of the $\\beta$ and $\\gamma$ hyperparameters involved in ERA:\n\n| $\\beta$  | $\\gamma$ | Validity (%) $\\uparrow$ | Top-100 mean $\\uparrow$ | Top-100 diversity $\\downarrow$ |\n|---------|----------|--------------------------|--------------------------|-------------------------------|\n| 1       | 0        | 28.7                     | 0.598                    | 0.122                         |\n| 1       | 0.01     | 30.34                    | 0.611                    | 0.117                         |\n| 1       | 0.1      | 36.6                     | 0.685                    | 0.118                         |\n| 10      | 0        | 72.25                    | 0.930                    | 0.139                         |\n| 10      | 0.01     | 72.0                     | 0.932                    | 0.129                         |\n| 10      | 0.1      | 74.23                    | 0.933                    | 0.131                         |\n| 100     | 0        | 74.31                    | 0.934                    | 0.142                         |\n| 100     | 0.01     | 74.71                    | 0.934                    | 0.139                         |\n| 100     | 0.1      | 75.14                    | 0.935                    | 0.134                         |\n\nThe table above contains descriptive statistics of 10,000 generated SMILES after aligning the molecular transformer model for QED maximization across several values of β and γ. For both β and γ, increasing its value led to an improvement in validity and performance in maximizing QED. These results demonstrate clear degradation in the model performance with certain choices of β (below 10), but these are not necessarily the case across all tasks. Table S4 demonstrates how choices of β that were not appropriate for QED, like β=1.0, can be beneficial for other tasks such as MR and Ring Count. This study is being replicated for other optimization tasks, and the results can be made available during the open discussion period.\n\nFor DPO, we note that the DPO objective does not contain an explicit reward and instead only maximizes the margins between preferred and dis-preferred samples. Additional care needs to be taken when comparing DPO and ERA hyperparameters because $\\beta_{DPO} \\neq \\beta_{ERA}$ due to a difference in definition. Specifically, the minimizer of the DPO objective is \n$$\n\\pi_{DPO}(x, y) \\propto \\exp(-\\beta_{DPO}^{-1}r(x, y) + \\log \\pi_{ref})\n$$\nWhereas for ERA, the minimizer is\n$$\n\\pi_{ERA}(x, y) \\propto \\exp\\left(\\frac{-\\beta_{ERA}}{1 + \\gamma}r(x, y) + \\frac{\\gamma}{1 + \\gamma}\\log \\pi_{ref}\\right)\n$$\nThe formulation of the DPO objective does not permit explicitly scaling the contribution of the regularization term against the prior while ERA permits this. This also implies that \n$$\n\\beta_{ERA} = \\frac{1 + \\gamma}{\\beta_{DPO}}\n$$\nWe take $\\gamma$ to be relatively large since as $\\gamma \\to \\infty$, $\\frac{\\gamma}{1+\\gamma} \\to 1$. In our comparisons, we choose $\\gamma = 9$. The tables below are arranged such that the corresponding rows can be compared directly against each other. We did not scan any $\\beta < 0.1$ since the validity in those cases was low for both methods:\n\n**DPO results**\n| $\\beta_{DPO}$ | Validity (%) $\\uparrow$ | Top-100 mean $\\uparrow$ | Top-100 diversity $\\downarrow$ |\n| -------- | -------- | -------- | -------- |\n| 100      | 86.91     | 0.936     | 0.131\n| 10       | 80.92     | 0.943     | 0.134\n| 1        | 82.23     | 0.943     | 0.156\n| 0.1      | 54.46     | 0.947     | 0.174\n\n**ERA results ($\\gamma = 9$)**\n| $\\beta_{ERA}$ | Validity (%) $\\uparrow$ | Top-100 mean $\\uparrow$ | Top-100 diversity $\\downarrow$ |\n|---------------|--------------------------|--------------------------|---------------------------------|\n| 0.1           | 84.78                    | 0.939                    | 0.149                           |\n| 1             | 85.73                    | 0.936                    | 0.145                           |\n| 10            | 79.07                    | 0.932                    | 0.141                           |\n| 100           | 88.83                    | 0.941                    | 0.145                           |\n\nThe tables above contain descriptive statistics of 10,000 generated SMILES after aligning the molecular transformer model for QED maximization using either DPO or ERA across several matched values of $\\beta$. Due to parameter matching, we expect ERA and DPO to be similar here: validity is comparable, with ERA surpassing DPO at higher $\\beta_{ERA}$, as well as the top-100 average QED. **ERA consistently generates diverse molecules across values of $\\beta_{ERA}$ as opposed to DPO**, cf. lower similarity among the top-100 generations for $\\beta_{ERA} = 10$ and $\\beta_{ERA} = 100$ because ERA objective promotes diverse generations, even without a strong prior."}, {"review_id": "JAGGFjyo9T", "reviewer": "Reviewer_zXeq", "summary": "This 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.", "questions": "1. **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?", "limitations": "- **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.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "**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.", "was_revised": true, "final_justification": "I appreciate the authors’ solid rebuttal addressing all reviewers’ comments. I am happy to keep my original score.", "rebuttal": "Thank you for the thorough review and positive assessment of our work. We respond in detail below:\n\n> 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\nWe agree that this should be specified clearly in the paper. In order to ensure that our experiments focus on the hyperparameters of ERA and in all experiments we use the fairly standard practice of top-k sampling with $k=5$ and do not adjust the temperature $T=1.$\n\n> 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\nIndeed, especially when the regularization $\\gamma=0$, we see that ERA reduces validity of the generated SMILES. There are some limitations imposed by the fact that most of our downstream characterizations cannot be run on invalid SMILES strings because of the failure of RDkit to parse. However, in light of this interesting suggestion, we computed two measures of the diversity: an estimate of the Levenshtein distance and an estimate of the Shannon entropy for both valid and invalid molecules. These results are summarized in the table below:\n\n| Generations | Shannon Entropy | Levenshtein Distance |\n|-------------|------------------|----------------------|\n| Valid       | 2.2574           | 39.8225              |\n| Invalid     | 2.3177           | 56.1066              |\n\nThe table above demonstrates properties of the valid and invalid generations from a model aligned using ERA for QED maximization for $\\beta$=100.0, $\\gamma$=0.1. Notably, the Shannon entropy of generated tokens and the mean Levenshtein distance between all of the generated SMILES is higher for the invalid outputs than the valid ones. These metrics suggest that there is in fact less clustering around one specific region of chemical space in these failure cases.\n\n\n> 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\nFor the same QED maximization run ($\\beta$=100.0, $\\gamma$=0.1) as above, we found a significant failure mode to be that of missing tokens to close branches (67.6% of all invalid generations). Furthermore, failure to properly close rings also occurs in a signifcant amount of invalid generations (30.5%). One potential failure mode that seems to be relatively insignificant is that of failing to predict the stop token before our hard limit of 500 tokens, which occurs in fewer than 1% of invalid generations. Beyond this analysis, visual inspection of several failure cases does not appear to reveal any particular motif that is repeated among these examples beyond the model attempting to generate highly branched molecules or molecules with several rings.\n\n\n>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\n| $\\beta$  | $\\gamma$ | Validity (%) $\\uparrow$ | Top-100 mean $\\uparrow$ | Top-100 diversity $\\downarrow$ |\n|---------|----------|--------------------------|--------------------------|-------------------------------|\n| 1       | 0        | 28.7                     | 0.598                    | 0.122                         |\n| 1       | 0.01     | 30.34                    | 0.611                    | 0.117                         |\n| 1       | 0.1      | 36.6                     | 0.685                    | 0.118                         |\n| 10      | 0        | 72.25                    | 0.930                    | 0.139                         |\n| 10      | 0.01     | 72.0                     | 0.932                    | 0.129                         |\n| 10      | 0.1      | 74.23                    | 0.933                    | 0.131                         |\n| 100     | 0        | 74.31                    | 0.934                    | 0.142                         |\n| 100     | 0.01     | 74.71                    | 0.934                    | 0.139                         |\n| 100     | 0.1      | 75.14                    | 0.935                    | 0.134                         |\n\n\nThe table above contains descriptive statistics of 10,000 generated SMILES after aligning the molecular transformer model for QED maximization across several values of β and γ. For both β and γ, increasing its value led to an improvement in validity and performance in maximizing QED. These results demonstrate clear degradation in the model performance with certain choices of β (below 10), but these are not necessarily the case across all tasks. Table S4 in the appendix demonstrates how choices of β that were not appropriate for QED, like β=1.0, can be beneficial for other tasks such as MR and Ring Count. This study is being replicated for other optimization tasks discussed in the paper, and the results can be made available during the open discussion period at the reviewers' request. \n\n>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\nThis is an excellent suggestion, we have not implemented it because our models were still state-of-the-art in sample efficiency even accounting for the invalid generations, but it is certainly a good idea and one that we will investigate."}, {"review_id": "wRIRXdWku2", "reviewer": "Reviewer_7xzk", "summary": "The 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.", "questions": "1. 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.", "limitations": "Yes.", "rating": 3, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 4, "significance": 3, "originality": 3, "strengths_and_weaknesses": "The 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.", "was_revised": true, "final_justification": "I maintain my recommendation for rejection, as the rebuttal did not adequately address my concerns.", "rebuttal": "> However, 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\n\nFor DPO, we note that the DPO objective does not contain an explicit reward component and instead only maximizes the margins between preferred and dis-preferred samples. Additional care needs to be taken when comparing DPO and ERA hyperparameters because $\\beta_{DPO} \\neq \\beta_{ERA}$ due to a difference in definition. Specifically, the minimizer of the DPO objective is \n$$\n\\pi_{DPO}(x, y) \\propto \\exp(-\\beta_{DPO}^{-1}r(x, y) + \\log \\pi_{ref})\n$$\nWhereas for ERA, the minimizer is\n$$\n\\pi_{ERA}(x, y) \\propto \\exp\\left(\\frac{-\\beta_{ERA}}{1 + \\gamma}r(x, y) + \\frac{\\gamma}{1 + \\gamma}\\log \\pi_{ref}\\right)\n$$\nThe formulation of the DPO objective does not permit explicitly scaling the contribution of the regularization term against the prior while ERA permits this. This also implies that \n$$\n\\beta_{ERA} = \\frac{1 + \\gamma}{\\beta_{DPO}}\n$$\nFurthemore, we take $\\gamma$ to be relatively large since as $\\gamma \\to \\infty$, $\\frac{\\gamma}{1+\\gamma} \\to 1$. In our comparisons, we choose $\\gamma = 9$. The tables below are arranged such that the corresponding rows can be compared directly against each other. We did not scan any $\\beta < 0.1$ since the validity in those cases was low for both methods:\n\n\n**DPO results**\n| $\\beta_{DPO}$ | Validity (%) $\\uparrow$ | Top-100 mean $\\uparrow$ | Top-100 diversity $\\downarrow$ |\n| -------- | -------- | -------- | -------- |\n| 100      | 86.91     | 0.936     | 0.131\n| 10       | 80.92     | 0.943     | 0.134\n| 1        | 82.23     | 0.943     | 0.156\n| 0.1      | 54.46     | 0.947     | 0.174\n\n**ERA results ($\\gamma = 9$)**\n| $\\beta_{ERA}$ | Validity (%) $\\uparrow$ | Top-100 mean $\\uparrow$ | Top-100 diversity $\\downarrow$ |\n|---------------|--------------------------|--------------------------|---------------------------------|\n| 0.1           | 84.78                    | 0.939                    | 0.149                           |\n| 1             | 85.73                    | 0.936                    | 0.145                           |\n| 10            | 79.07                    | 0.932                    | 0.141                           |\n| 100           | 88.83                    | 0.941                    | 0.145                           |\n\n\n\nThe tables above contain descriptive statistics of 10,000 generated SMILES after aligning the molecular transformer model for QED maximization using either DPO or ERA across several matched values of $\\beta$. We note that the validity in generations is comparable, with ERA maintaining consistently high generation validity and surpassing DPO at higher $\\beta_{ERA}$, as well as the average QED among the top-100 generations. ERA also more consistently generates diverse molecules across values of $\\beta_{ERA}$ as opposed to DPO as evidenced by the lower similarity among the top-100 generations for $\\beta_{ERA} = 10$ and $\\beta_{ERA} = 100$. This again reflects the fact that the entropy-regulated ERA objective promotes diverse generations while effectively avoiding greedy policies as opposed to DPO, which focuses only on maximizing the margins between preferred and dispreferred sequences. This highlights a key advantage of ERA since good molecular diversity is a key consideration for chemical and drug discovery tasks.\n\n> In 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\nOur submission contains several quantifications of the superiority of ERA compared to reinforcement learning strategies used in the literature. In our benchmarks, ERA outperforms RL-based strategies, including REINVENT and MolRL-mGPT, on the task of generating ligands for protein docking against GSK3$\\beta$ and JNK3. These are techniques that use an objective similar to the one used in PPO. For completeness, we reproduce Table 1 of our submission here. The best numbers are bolded:\n\n| Method     | GSK3$\\beta$ mean score $\\uparrow$ | GSK3$\\beta$ IntDiv $\\downarrow$ | JNK3 mean score $\\uparrow$ | JNK3 IntDiv $\\downarrow$ |\n|:---------- | --------------------------------- |:-------------------------------:|:--------------------------:|:------------------------:|\n| ERA        | 0.996 $\\pm$ 0.000                 |        **0.219 $\\pm$ 0.002**        |     **0.987 $\\pm$ 0.001**      |    **0.264 $\\pm$ 0.005**     |\n| MolRL-MGPT | **1.000 $\\pm$ 0.000**             |        0.362 $\\pm$ 0.015        |     0.961 $\\pm$ 0.010      |    0.372 $\\pm$ 0.025     |\n| GFlowNet   | 0.649 $\\pm$ 0.072                 |        0.715 $\\pm$ 0.104        |     0.437 $\\pm$ 0.219      |    0.716 $\\pm$ 0.145     |\n| GraphGA    | 0.919 $\\pm$ 0.016                 |        0.365 $\\pm$ 0.024        |     0.875 $\\pm$ 0.025      |    0.380 $\\pm$ 0.015     |\n| JT-VAE     | 0.235 $\\pm$ 0.083                 |        0.770 $\\pm$ 0.067        |     0.159 $\\pm$ 0.040      |    0.781 $\\pm$ 0.127     |\n| REINVENT   | 0.965 $\\pm$ 0.011                 |        0.308 $\\pm$ 0.035        |     0.942 $\\pm$ 0.019      |    0.368 $\\pm$ 0.021     |\n\n\nFurthermore, we have added explicit comparisons to DPO as shown above, emphasized our previous comparisons to state of the art RL strategies for molecular optimization, and added additional parameter ablations to study the role of the inverse temperature parameter $\\beta$ and the prior coupling strength $\\gamma$. We emphasize that DPO constructs the reward function implicitly by design and we view ERA as the natural extension of DPO to the setting in which an explicit reward function is available."}, {"review_id": "joh5zdvY3f", "reviewer": "Reviewer_EGLr", "summary": "In 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.", "questions": "Can the authors elaborate on the “design decisions” of the algorithm in a more compelling way? (See “Weaknesses” above for context.)", "limitations": "If the costs are a limitation of the algorithm, I suggest that the authors add a discussion in this section.", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 3, "originality": 3, "strengths_and_weaknesses": "**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.", "was_revised": true, "final_justification": "- The paper’s significance and novelty are good\n- Clarity issues have been addressed well during the rebuttal\n- The authors have also clarified their algorithm’s costs\n- After seeing other reviews and the authors’ responses, I don’t see any significant issues anymore", "rebuttal": "> 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> Can the authors elaborate on the “design decisions” of the algorithm in a more compelling way? (See “Weaknesses” above for context.)\n\nWe agree that the derivation and motivation of our approach would benefit from a more thorough explanation in the main text. If accepted, we will expand Section 2 so that it includes a systematic derivation of the approach and clear articulation of the advantages over alternative approaches. \n\nBriefly, the main difference in the goal of ERA is to leverage an explicit reward function to balance sample diversity and \"greedy\" optimization to maximize rewards. In both PPO and DPO, the prior term cannot be independently tuned at fixed temperature. More mathematically, the minimizer of the DPO objective is, \n$$\n\\pi_{\\textrm{DPO}}(x,y) \\propto \\exp \\left( \\beta_{\\rm DPO}^{-1} r(x,y) + \\log \\pi_{\\rm ref} \\right), \n$$\nwhereas the minimizer of the ERA objective is\n$$\n\\pi_{\\textrm{ERA}}(x,y) \\propto \\exp \\left( \\frac{\\beta_{\\rm ERA}}{1+\\gamma} r(x,y) + \\frac{\\gamma}{1+\\gamma}\\log \\pi_{\\rm ref} \\right).\n$$\nIn DPO, the magnitude of the regularization to the prior cannot be tuned independently of $\\beta$. In particular, this means that the only way to change the impact of the prior is to effectively lower the sampling temperature and hence lower sample diversity. \nOn the other hand ERA allows these parameters to be tuned independently, for example, $\\gamma\\to 0$ enables a distribution with no prior regularization while preserving sample diversity.\n\nThe inclusion of an explicit reward model is the second important aspect of ERA. The explicit reward matters most when the number of observations of pairwise comparisons is small. **Stated simply, DPO overly penalizes molecules with similar reward values by maximizing the margin between them.**\n\nFor example, if the comparison of two unique molecules shows that $\\boldsymbol{y} \\succ \\boldsymbol{y}'$, i.e., $\\boldsymbol{y}$ is preferred to $\\boldsymbol{y}'$, then the DPO requires that the optimal policy requires that\n\\begin{equation}\n    \\pi_{\\star}^{\\textrm{DPO}}(\\boldsymbol{y}|\\boldsymbol{x}) = 1 \\quad \\textrm{and} \\quad \\pi_{\\star}^{\\textrm{DPO}}(\\boldsymbol{y}'|\\boldsymbol{x}) = 0.\n\\end{equation}\n**This is true even if both molecules have similar, high reward values.**\n\nThis can also lead to instability: the sampled gradients of the objective used for DPO are proportional to the implicit reward discrepancy, \n\\begin{equation}\n    \\nabla_{\\boldsymbol{\\theta}} \\hat{\\mathcal{L}}^{\\textrm{DPO}}(\\boldsymbol{y}, \\boldsymbol{y}', \\boldsymbol{x}) = \\sigma \\left( \\beta^{-1}\\gamma \\left[\\log \\frac{\\pi(\\boldsymbol{y}'|\\boldsymbol{x})}{\\pi_{\\rm ref}(\\boldsymbol{y}'|\\boldsymbol{x})} - \\log \\frac{\\pi(\\boldsymbol{y}|\\boldsymbol{x})}{\\pi_{\\rm ref}(\\boldsymbol{y}|\\boldsymbol{x})} \\right] \\right) \\nabla_{\\boldsymbol{\\theta}} \\log \\frac{\\pi(\\boldsymbol{y}|\\boldsymbol{x})}{\\pi(\\boldsymbol{y}'|\\boldsymbol{x})},\n\\end{equation}\nwhich when $\\pi(\\boldsymbol{y}'|\\boldsymbol{x}) \\to 0$, leads to instability as $-\\log \\pi(\\boldsymbol{y}'|\\boldsymbol{x}) \\to \\infty.$ On the other hand, the ERA gradients are scaled by the relative preference discrepancy,\n\\begin{equation}\n    \\nabla_{\\boldsymbol{\\theta}} \\mathcal{L}^{\\textrm{ERA}}(\\boldsymbol{y}, \\boldsymbol{y}', \\boldsymbol{x}) = \\left( \\frac{1-\\sigma_{\\star}(\\boldsymbol{y}\\succ\\boldsymbol{y}'|\\boldsymbol{x})}{1-\\sigma_{\\boldsymbol{\\theta}}(\\boldsymbol{y}\\succ\\boldsymbol{y}'|\\boldsymbol{x})} - \\frac{\\sigma_{\\star}(\\boldsymbol{y}\\succ\\boldsymbol{y}'|\\boldsymbol{x})}{\\sigma_{\\boldsymbol{\\theta}}(\\boldsymbol{y}\\succ\\boldsymbol{y}'|\\boldsymbol{x})} \\right) \\nabla_{\\boldsymbol{\\theta}} \\sigma_{\\boldsymbol{\\theta}}(\\boldsymbol{y}\\succ\\boldsymbol{y}'|\\boldsymbol{x}).\n\\end{equation}\nThe advantage of a reward model becomes apparent because \n\\begin{equation}\n    \\sigma_{\\star}(\\boldsymbol{y}\\succ\\boldsymbol{y}'|\\boldsymbol{x}) = p_{\\gamma}(\\boldsymbol{y}\\succ \\boldsymbol{y}'|\\boldsymbol{x}) = \\sigma \\left(  \\frac{\\beta}{1+\\gamma} \\left[ (U(\\boldsymbol{x},\\boldsymbol{y}') - U(\\boldsymbol{x},\\boldsymbol{y})) +  \\beta^{-1}\\gamma \\log \\frac{\\pi_{\\rm ref}(\\boldsymbol{y}|\\boldsymbol{x})}{\\pi_{\\rm ref}(\\boldsymbol{y}'|\\boldsymbol{x})} \\right] \\right)\n\\end{equation}\nand hence the optimum of $\\mathcal{L}^{\\textrm{ERA}}$ will not lead to policies in which high reward samples are suppressed by the model. This could alternatively be implemented by taking $\\beta\\to\\infty$. Choosing an appropriate reward model, hence, gives the flexibility to control instability if it becomes problematic.\n\n\n> 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\n\nIndeed, costs are extremely important in this context. In our submission, we discuss two types of costs: first the computational costs associated with inference for molecular generation, which are low:\n> Line 237: The inference costs are notably low for our approach; sampling 20000 molecules and filtering takes only minutes on a single GPU.\n\nThe more relevant costs for molecular optimizations are typically the number of oracle calls, which might be experimentally or computationally expensive to evaluate. We quantified this in our submission (Fig. S13) by looking at top-10 statistics as a function of number of oracle calls. The following Table (original submission Table S7) demonstrates that the sample efficiency of ERA is superior to REINVENT and GraphGA.\n\n\n|          | GSK3β top-10 AUC | JNK3 top-10 AUC |\n|----------|------------------|------------------|\n| **ERA**  | **0.985 ± 0.001** | **0.989 ± 0.002** |\n| REINVENT | 0.865 ± 0.043 | 0.783 ± 0.023 |\n| GraphGA  | 0.788 ± 0.070 | 0.553 ± 0.136 |\n\n\n\nIn our submission, we discussed this Table S7 and Figure S13 in lines 240-248. \nThe top-10 AUC has previously been proposed as a key metric of performance relative to cost and we believe that our experiments indicate that ERA is state-of-the-art on this metric."}]

Author rebuttal:

Meta-review: This paper introduces energy rank alignment, an approach for to use explicit reward function to produce a gradient-based objective for auto-regressive policies; the authors show how this can be used in transformer-based generative models in molecular applications. The paper had extensive experimentation and evaluation in the initial submission. The most critical issues the reviewers raised concerned the comparisons to several baselines, ablations and presentation, all of which were thoroughly addressed during. the discussion phase. The paper presents an interesting new method, which may have impact beyond molecular machine learning applications in the future.
Decision: Accept (poster)