Generative model alignmentGenerative AI applications are now routinely fine-tuned with feedback-based RL during post-training, which has been crucial for the success of applications in the natural language domain49,50. Here, we provide a brief overview of the classical model training pipeline, which consists of three stages: SFT, feedback collection and RL fine-tuning51,52.SFTHere, we define SFT as continuing the original pretraining objective on high-quality samples from a downstream labeled dataset. In this work for example, for SFT on the ‘Megascale’ dataset, high-quality samples are sequences which achieve superior stability to the native sequence. During the SFT stage we use a dataset of input-output pairs, \({\mathcal{D}}={\{{\bf{x}},{\bf{y}}\}}_{i = 1}^{N}\) and train a model pθ with next-token maximum likelihood loss on high-quality pairs, where yi is the ith token of the output y$${{\mathscr{ \mathcal L }}}_{{p}_{\theta }}=-{{\mathbb{E}}}_{{\bf{x}},{\bf{y}}\sim{\mathscr{ \mathcal D }}}\left[\sum _{i=1}^{| {\bf{y}}| }\mathrm{log}{p}_{\theta }(\,{y}_{i}| {\bf{x}},{{\bf{y}}}_{ < i})\right].$$(1)Feedback stageDuring the feedback stage we use the base model or the SFT model pSFT and a set of inputs to sample a number K > 0 of candidate completions y(i), …, y(K) per prompt x from a dataset of preferences \({{\mathcal{D}}}_{{\rm{pref}}}\). Feedback can be provided in several ways:(1)Prior RLHF work has used ranking type of comparisons in which the candidate outputs are ranked \({{\bf{y}}}^{({i}_{1})}\succ {{\bf{y}}}^{({i}_{2})}\succ \ldots \succ {{\bf{y}}}^{({i}_{K})}\), with pairwise rankings being a popular choice. To use this data from the RL stage, we need to extract an evaluation score from it. We usually assume a parameterization of the preference distribution, such as the Plackett-Luce53,54 choice model (or the Bradley–Terry55 for pairwise data). In the most general case, given the permutation \(\tau :{{\bf{y}}}^{(1)},\ldots ,{{\bf{y}}}^{(K)}\to {{\bf{y}}}^{({i}_{1})},\ldots ,{{\bf{y}}}^{({i}_{k})}\), we have$$P(\tau | {{\bf{y}}}^{(1)},\ldots ,{{\bf{y}}}^{(K)},{\bf{x}})=\prod _{k=1}^{K}\frac{\exp (r({\bf{x}},{{\bf{y}}}^{({i}_{k})}))}{{\sum }_{j=k}^{K}\exp (r({\bf{x}},{{\bf{y}}}^{({i}_{\!j})}))}.$$(2)We can then form ulate this as a maximum likelihood loss over the reward model$${{{ \mathcal L }}}_{r}=-{{\mathbb{E}}}_{{\bf{x}},{{\bf{y}}}^{({i}_{1})},\ldots ,{{\bf{y}}}^{({i}_{K})}\sim {{{ \mathcal D }}}_{\mathrm{pref}}}\sum _{k=1}^{K}\left[r\left({\bf{x}},{{\bf{y}}}^{({i}_{k})}\right)-{\mathrm{log}}\sum _{j=k}^{K}\exp \left(r\left({\bf{x}},{{\bf{y}}}^{({i}_{\!j})}\right)\right)\right].$$(3)(2)In the case of pairwise rankings yw ≻ yl, this just reduces to the standard binary logistic regression loss,$${{{ \mathcal L }}}_{r}=-{{\mathbb{E}}}_{{\bf{x}},{{\bf{y}}}^{w},{{\bf{y}}}^{l}\sim {{{ \mathcal D }}}_{\mathrm{pref}}}\left[{\mathrm{log}}\sigma (r({\bf{x}},{{\bf{y}}}^{w})-r({\bf{x}},{{\bf{y}}}^{l}))\right],$$(4)where σ(⋅) is the sigmoid function. Notice that if we are provided with a complete ranking over K candidates then we can construct \(\left(\begin{array}{c}K\\ 2\end{array}\right)\) pairwise rankings, which was used in prior works22. These can still be easily optimized together since we still require a single forward pass through the reward model.(3)Alternatively a numerical score, ri (such as used in a Likert scale) could be provided per candidate input-output pair (x, y(i)).RL stageDuring the RL stage we use a reference model pref, sometimes initialized as pSFT from the SFT stage, and optimize the following RL problem:$${\mathop{\max }\limits_{{p}_{\theta }}}\,\underbrace{{{\mathbb{E}}}_{{\bf{x}} \sim {\mathcal{D}},{\bf{y}} \sim {p}_{\theta }({\bf{y}}| {\bf{x}})}[r({\bf{x}},{\bf{y}})]}_{{\rm{maximize}}\,{\rm{score}}\,{\rm{of}}\,{\rm{generations}}}-\underbrace{\beta {{\mathbb{D}}}_{{\rm{KL}}}[{p}_{\theta }({\bf{y}}| {\bf{x}})| | {p}_{{\rm{ref}}}({\bf{y}}| {\bf{x}})]}_{{\rm{distributional}}\,{\rm{penalty}}}.$$(5)We want to optimize a model pθ to maximize the score, while also maintaining a distribution close to the reference model. The second terms aims to keep the model close to the distribution, which was used to collect the feedback data and prevent out-of-distribution issues56.Optimization algorithmsTwo main algorithms have emerged for the optimization of the objective in equation (5): (1) classical RL algorithms, such as PPO57, which use on-policy rollouts from the model; and (2) direct alignment algorithms, such as DPO16, which train fully-offline using only feedback data. Although successful in large scale-systems, such as GPT-4, the PPO pipeline requires a very accurate score model that can be queried in a continuous manner. Prior work has directly trained a separate score network rθ(x, y) using the feedback data with the objective in equation (4) and directly applied PPO to the problem in equation (5) using the trained score model. However, such a pipeline can be quite complex, unstable and computationally expensive58,59. DPO16 is an alternative to this algorithmic pipeline, which we present below.DPOThe DPO algorithm is based on the observation that the objective in equation (5) has a closed form solution of the type$${p}^{* }({\bf{y}}| {\bf{x}})=\frac{1}{Z({\bf{x}})}{p}_{\mathrm{ref}}({\bf{y}}| {\bf{x}}){{\rm{e}}}^{\frac{1}{\beta }r({\bf{x}},{\bf{y}})},$$(6)where$$Z({\bf{x}})=\sum _{{\bf{y}}}{p}_{\mathrm{ref}}({\bf{y}}| {\bf{x}}){{\rm{e}}}^{\frac{1}{\beta }r({\bf{x}},{\bf{y}})}$$(7)is an intractable partition function. The key insight of the DPO algorithm is the model-reward equivalency, that is,$$r({\bf{x}},{\bf{y}})=\beta \log \frac{{p}^{* }({\bf{y}}| {\bf{x}})}{{p}_{{\rm{ref}}}({\bf{y}}| {\bf{x}})}+\beta \log Z({\bf{x}}).$$(8)Now, if instead of training a score model rθ, we directly parameterize the generative model pθ, we can use the preference data to directly optimize pθ, essentially marginalizing the RL loop.For the ranked objective, if we substitute the above expression into the Plackett–Luce reward model in equation (3), we obtain the following objective:$$\begin{array}{l}{{\mathcal{L}}}_{{p}_{\theta }}=-{{\mathbb{E}}}_{{\bf{x}},{{\bf{y}}}^{({i}_{1})},\ldots ,{{\bf{y}}}^{({i}_{K})} \sim {{\mathcal{D}}}_{{\rm{pref}}}}\mathop{\sum }\limits_{k=1}^{K}\left[\beta \log \frac{{p}_{\theta }({{\bf{y}}}^{{i}_{k}}| {\bf{x}})}{{p}_{{\rm{ref}}}({{\bf{y}}}^{{i}_{k}}| {\bf{x}})}\right.\\\quad\,\,\,\,\left.-\log \mathop{\sum }\limits_{j=k}^{K}\exp \left(\beta \log \frac{{p}_{\theta }({{\bf{y}}}^{({i}_{j})}| {\bf{x}})}{{p}_{{\rm{ref}}}({{\bf{y}}}^{({i}_{j})}| {\bf{x}})}\right)\right].\end{array}$$(9)For the paired objective, If we instead use the pairwise objective in equation (4), we obtain the canonical form of the DPO objective$${{\mathcal{L}}}_{r}=-{{\mathbb{E}}}_{{\bf{x}},{{\bf{y}}}^{w},{{\bf{y}}}^{l} \sim {{\mathcal{D}}}_{{\rm{pref}}}}\log \sigma \left(\beta \log \frac{{p}_{\theta }({{\bf{y}}}^{w}| {\bf{x}})}{{p}_{{\rm{ref}}}({{\bf{y}}}^{w}| {\bf{x}})}-\beta \log \frac{{p}_{\theta }({{\bf{y}}}^{l}| {\bf{x}})}{{p}_{{\rm{ref}}}({{\bf{y}}}^{l}| {\bf{x}})}\right).$$(10)Here, due to the homogeneity of both the Plackett–Luce and Bradley–Terry model (that is, comparing several options for the same context), the untractable partition term Z(x) cancels out and we are left with tractable, fully offline objectives.Weighted DPO algorithmPrior methods, including DPO16, have all focused on extracting a per-sample score r(x, y) from rankings. However, if we have some external evaluation score, that is, data \({\{{\bf{x}},{{\bf{y}}}_{1},r({\bf{x}},{{\bf{y}}}_{1})\}}_{i=1}^{K}\), then this step is redundant. Given numerical scores, we could train a score model, using a regression-based framework and then use the standard PPO pipeline, but this approach inherits all the shortcomings of the standard RLHF pipeline. We could use the numerical scores to induce a ranking over candidates yi and the apply the Plackett–Luce or Bradley–Terry DPO objectives in equations (9) and (10), respectively; however that faces two drawbacks:1.Plackett–Luce and Bradley–Terry both assume a particular parametric form of the ranking likelihood, that is, the rankings are inherently probabilistic, which would induce additional noise in the optimization procedure if we generate rankings based on those choice models.2.We can also generate rankings based on raw scores similarly to prior works60. However, note that any monotone transformation would still induce the same relative rankings. That is, there might be actual relative information in the score which could be lost in such a scenario.Weighted regression-based methods61,62,63 are a good fit for this problem; however, they have been suboptimal in the large generative model setting and do not yield significant improvement64. Instead, we will aim to develop a DPO type of algorithm that can efficiently utilize the numerical scores of the candidate sequences. Given an x and candidate sequences \({\{{{\bf{y}}}^{(i)},r({\bf{x}},{{\bf{y}}}^{(i)})\}}_{i = 1}^{K}\), consider the parameterized reward$${r}_{\theta }({\bf{x}},{{\bf{y}}}^{(i)})=\beta\,\mathrm{log}\frac{{p}_{\theta }({{\bf{y}}}^{(i)}| {\bf{x}})}{{p}_{\mathrm{ref}}({{\bf{y}}}^{(i)}| {\bf{x}})}+\beta \mathrm{log}{Z}_{\theta }({\bf{x}}).$$(11)That is a parameterized distribution pθ(y∣x) induces a reward function rθ(x, y), such that the model pθ is the optimal solution for the corresponding RL problem in equation (5). Now, instead of aligning the generative model pθ with some sampled rankings or preferences, we will instead minimize a distance function of the type D(r(x, y), rθ(x, y)); that is, we will match the true score samples with the implicit ones. A remaining issue is that we still have the constraining partition function Zθ(x) (because it depends on pθ), which we would like to eliminate. When we consider$${w}_{\theta }^{(i)}=\frac{\exp\left({r}_{\theta }\left({\bf{x}},{{\bf{y}}}^{(i)}\right)\right)}{{\sum }_{i}\exp\left({r}_{\theta }\left({\bf{x}},{{\bf{y}}}^{(i)}\right)\right)}=\frac{\exp\left(\beta \log \frac{{p}_{\theta }({{\bf{y}}}^{(i)}| {\bf{x}})}{{p}_{ref}({{\bf{y}}}^{(i)}| {\bf{x}})}\right)}{{\sum }_{i}\exp \left(\beta \log \frac{{p}_{\theta }({{\bf{y}}}^{(i)}| {\bf{x}})}{{p}_{ref}({{\bf{y}}}^{(i)}| {\bf{x}})}\right)},$$(12)notice that the partition function Zθ(x) was eliminated. In the same manner we have that$${w}^{(i)}=\frac{\exp\left(r\left({\bf{x}},{{\bf{y}}}^{(i)}\right)\right)}{{\sum }_{i}\exp\left(r\left({\bf{x}},{{\bf{y}}}^{(i)}\right)\right)}.$$(13)Notice now that w = (w(1), …, w(K)) and \({{\bf{w}}}_{\theta }=({w}_{\theta }^{(1)},\ldots ,{w}_{\theta }^{(K)})\) are both valid probability distributions. In fact, the probability distributions resemble the Boltzmann distribution from statistical mechanics65$$p(i)=\frac{\exp\left(-\frac{{E}_{i}}{{k}_{B}T}\right)}{{\sum }_{i}\exp\left(-\frac{{E}_{i}}{{k}_{B}T}\right)},$$(14)where p(i) is the probability of the system being in a state i with energy Ei, and kBT is the scaled absolute temperature of the system. Substituting r(x, y(i)) for the negative energy of the state −Ei, and considering kBT, a parameter which allows us to scale the numerical scores r(x, y(i)), the weighted DPO objective is simply minimization of the Kullback–Leibler (KL) divergence \({{\mathbb{D}}}_{KL}[{\bf{w}}| {{\bf{w}}}_{\theta }]\) between the Boltzmann-like distributions parameterized by the generative model pθ and numerical scores from the dataset \({\mathcal{D}}\). With some simple algebra, this yields:$${{\mathscr{ \mathcal L }}}_{{p}_{\theta }}=-{{\mathbb{E}}}_{{\mathscr{ \mathcal D }}}\left[\sum _{i}{w}^{(i)}\beta \mathrm{log}\frac{{p}_{\theta }({{\bf{y}}}^{(i)}| {\bf{x}})}{{p}_{\mathrm{ref}}({{\bf{y}}}^{(i)}| {\bf{x}})}-\mathrm{log}\sum _{i}\exp \left(\beta \mathrm{log}\frac{{p}_{\theta }({{\bf{y}}}^{(i)}| {\bf{x}})}{{p}_{\mathrm{ref}}({{\bf{y}}}^{(i)}| {\bf{x}})}\right)\right].$$(15)In the case where we have only two candidate sequences y(1) and y(2), this can be simplified to$${{\mathscr{ \mathcal L }}}_{{p}_{\theta }}=-{{\mathbb{E}}}_{{\mathscr{ \mathcal D }}}\left[{w}^{(i)}\mathrm{log}\sigma \left(\beta \mathrm{log}\frac{{p}_{\theta }({{\bf{y}}}^{(1)}| {\bf{x}})}{{p}_{\mathrm{ref}}({{\bf{y}}}^{(1)}| {\bf{x}})}-\beta \mathrm{log}\frac{{p}_{\theta }({{\bf{y}}}^{(2)}| {\bf{x}})}{{p}_{\mathrm{ref}}({{\bf{y}}}^{(2)}| {\bf{x}})}\right)+\right.$$(16)$$\left.{w}^{(2)}\mathrm{log}\sigma \left(\beta \mathrm{log}\frac{{p}_{\theta }({{\bf{y}}}^{(2)}| {\bf{x}})}{{p}_{\mathrm{ref}}({{\bf{y}}}^{(2)}| {\bf{x}})}-\beta \mathrm{log}\frac{{p}_{\theta }({{\bf{y}}}^{(1)}| {\bf{x}})}{{p}_{\mathrm{ref}}({{\bf{y}}}^{(1)}| {\bf{x}})}\right)\right].$$(17)That is, this is essentially a weighted version of the DPO algorithm. One interpretation is that this objective is the expected DPO objective under p(y(1) ≻ y(2)∣x) = w(1) and p(y(2) ≻ y(1)∣x) = w(2). These objectives are theoretically optimal in the offline regime as they are both mode-seeking and use a ‘negative gradient’ formulation64.ImplementationAlgorithm 1 Compute sequence log-likelihood 1: Input: sequence y, encoded structure h 2: Output: log-likelihood \(\log {p}_{\theta }({\bf{y}}| {\bf{h}})\) 3: // Get logits via teacher forcing 4: O ← TransformerDecoder(y[0: |y|−1], h) // Logits for next-token prediction 5: // Convert logits to probabilities 6: P ← Softmax(O) // Token probabilities over vocabulary 7: // Compute log-likelihood as negative cross-entropy 8: l ← − CrossEntropy(P, y[1: |y|]) // Per-position negative loss 9: \(\log {p}_{\theta }({\bf{y}}| {\bf{h}})\leftarrow {\sum }_{t}{\ell }_{t}\) // Sum over sequence positions10: return \(\log {p}_{\theta }({\bf{y}}| {\bf{h}})\)Algorithm 2 DPO for structure-conditioned protein language model (pairwise) 1: Input: structure \({\bf{X}}\in {{\mathbb{R}}}^{L\times 3}\), preferred sequence yw, dispreferred sequence yl, reference log-likelihoods \(\log {p}_{{\rm{ref}}}({{\bf{y}}}_{w}| {\bf{X}})\) and \(\log {p}_{{\rm{ref}}}({{\bf{y}}}_{l}| {\bf{X}})\), KL penalty coefficient β 2: Output: DPO loss \({{\mathscr{ \mathcal L }}}_{\mathrm{DPO}}\) 3: // Encode structure once for both sequences 4: h ← GVPEncoder(X) 5: // Compute policy log-likelihoods for both sequences 6: Ow, Ol ← TransformerDecoder(yw, h), TransformerDecoder(yl, h) // Logits 7: \(\mathrm{log}{p}_{\theta }({{\bf{y}}}_{w}| {\bf{X}}),\mathrm{log}{p}_{\theta }({{\bf{y}}}_{l}| {\bf{X}})\leftarrow \,\mathrm{ComputeLogLikelihood}\,({{\bf{y}}}_{w},{{\bf{O}}}_{w}),\,\)\(\mathrm{ComputeLogLikelihood}\,({{\bf{y}}}_{l},{{\bf{O}}}_{l})\) 8: // Compute pairwise DPO loss 9: \({r}_{\theta }\leftarrow \log {p}_{\theta }({{\bf{y}}}_{w}| {\bf{X}})-\log {p}_{\theta }({{\bf{y}}}_{l}| {\bf{X}})\) // Policy log-ratio10: \({r}_{\rm{ref}} \leftarrow {\rm{log}} \,\,{p}_{\rm{ref}} ({\bf{y}}_{w} |{\bf{X}}) - {\rm{log}}\,{p}_{\rm{ref}}({\bf{y}}_{l}|{\bf{X}})\) // Reference log-ratio11: \({{{ \mathcal L }}}_{\mathrm{DPO}}\leftarrow -\mathrm{log}\sigma (\beta ({r}_{\theta }-{r}_{\mathrm{ref}}))\)12: return \({{\mathcal{L}}}_{{\rm{DPO}}}\)Algorithm 3 DPO for structure-conditioned protein language model (ranked) 1: Input: structure \({\bf{X}}\in {{\mathbb{R}}}^{L\times 3}\), ranked sequences {y1, …, yK} where y1 ≻ y2 ≻ … ≻ yK, reference log-likelihoods \({\{\log {p}_{{\rm{ref}}}({{\bf{y}}}_{k}| {\bf{X}})\}}_{k = 1}^{K}\), KL penalty coefficient β 2: Output: DPO loss \({{\mathcal{L}}}_{{\rm{DPO}}}\) 3: // Encode structure once for all sequences 4: h ← GVPEncoder(X) 5: // Compute policy log-likelihoods for all sequences 6: for k = 1 to K do 7: Ok ← TransformerDecoder(yk, h) // Logits 8: \(\log {p}_{\theta }({{\bf{y}}}_{k}| {\bf{X}})\leftarrow \,\text{ComputeLogLikelihood}\,({{\bf{y}}}_{k},{{\bf{O}}}_{k})\) 9: end for10: // Plackett–Luce ranking loss11: \({{\mathcal{L}}}_{{\rm{DPO}}}\leftarrow 0\)12: for k = 1 to K − 1 do13: \({r}_{k}\leftarrow {\beta}\,({\mathrm{log}}\,{p}_{\theta } ({{\bf{y}}}_{k}| {\bf{X}})- \mathrm{log}\,{p}_{\mathrm{ref}} ({{\bf{y}}}_{k}| {\bf{X}}))\) // Reward for rank k14: \({s}_{k}\leftarrow \log \mathop{\sum }\nolimits_{j = k+1}^{K}\exp \left(\beta (\log {p}_{\theta }({{\bf{y}}}_{j}| {\bf{X}})-\log {p}_{{\rm{ref}}}({{\bf{y}}}_{j}| {\bf{X}}))\right)\)15: \({{\mathscr{ \mathcal L }}}_{\mathrm{DPO}}\leftarrow {{\mathscr{ \mathcal L }}}_{\mathrm{DPO}}-\mathrm{log}\sigma ({r}_{k}-{s}_{k})\)16: end for17: return \({{\mathcal{L}}}_{{\rm{DPO}}}\)Algorithm 4 DPO for structure-conditioned protein language model (weighted) 1: Input: structure \({\bf{X}}\in {{\mathbb{R}}}^{L\times 3}\), sequences \({\{{{\bf{y}}}_{k}\}}_{k = 1}^{K}\), reference log-likelihoods \({\{\log {p}_{{\rm{ref}}}({{\bf{y}}}_{k}| {\bf{X}})\}}_{k = 1}^{K}\), preference targets \({\bf{w}}\in {{\mathbb{R}}}^{K}\), KL penalty coefficient β, temperature T 2: Output: DPO loss \({{\mathcal{L}}}_{{\rm{DPO}}}\) 3: // Encode structure once for all sequences 4: h ← GVPEncoder(X) 5: // Compute policy log-likelihoods for all sequences 6: for k = 1 to K do 7: Ok ← TransformerDecoder(yk, h) // Logits 8: \(\log {p}_{\theta }({{\bf{y}}}_{k}| {\bf{X}})\leftarrow \,\text{ComputeLogLikelihood}\,({{\bf{y}}}_{k},{{\bf{O}}}_{k})\) 9: end for10: // Compute reward-weighted cross-entropy loss11: \({\bf{r}}\!\leftarrow\!\beta {\left[\log {p}_{\theta }({{\bf{y}}}_{1}| {\bf{X}})\!-\!\log {p}_{{\rm{ref}}}({{\bf{y}}}_{1}| {\bf{X}}),\!\ldots\! ,\log {p}_{\theta }({{\bf{y}}}_{K}| {\bf{X}})\!-\!\log {p}_{{\rm{ref}}}({{\bf{y}}}_{K}| {\bf{X}})\right]}^{\top }\)12: \({{\mathcal{L}}}_{{\rm{DPO}}}\leftarrow -\mathop{\sum }\nolimits_{k = 1}^{K}{w}_{k}\log \frac{\exp ({r}_{k}/T)}{\mathop{\sum }\nolimits_{j = 1}^{K}\exp ({r}_{\!j}/T)}\) // Cross-entropy with temperature scaling13: return \({{\mathcal{L}}}_{{\rm{DPO}}}\)DatasetsMegascaleThe Megascale dataset (v2 4/20/2023) served as the experimentally measured stability dataset from which we aligned the pretrained ESM-IF1. After filtering out data marked as unreliable for machine learning by Megascale authors, the dataset contained approximately 660,000 variants for 405 protein domains. The variants consisted of insertions, deletions, single and double substitutions. Similar to prior methods13, insertions and deletions were not considered as DPO training was designed to compare stabilities across the same structure; however, due to the robust nature of the model we were able to train and model with the double-mutations in the dataset. We additionally removed duplicate sequences. FoldSeek23 clustering with the ‘easy-cluster algorithm was performed on all PDBs remaining in the Megascale dataset with an alignment threshold of 50%. From this, clusters were randomly assigned to respective train (90%), validation (5%) and test (5%) datasets to ensure generalizability of the model across different structures. Because the model is trained with sets of sequences for a common structure, either the ΔG (absolute) or ΔΔG (relative) of sequences were used for training. While these should be identical in theory after accounting for native ΔG, due to noise in the data the two metrics are different. We found negligible difference in loss and validation accuracy training with either metric, so all models were trained on ΔG as more data points had ΔG measurements. For SFT, all variants with ΔΔG