Reinforcement learning on structure-conditioned categorical diffusion for protein inverse folding

## Abstract

Abstract Protein inverse folding—that is, predicting an amino acid sequence that will fold into the desired 3D structure—is an important problem for structure-based protein design. Machine learning based methods for inverse folding typically use recovery of the original sequence as the optimization objective. However, inverse folding is a one-to-many problem where several sequences can fold to the same structure. Moreover, for many practical applications, it is often desirable to have multiple, diverse sequences that fold into the target structure since it allows for more candidate sequences for downstream optimizations. Here, we demonstrate that although recent inverse folding methods show increased sequence recovery, their “foldable diversity”—i.e. their ability to generate multiple non-similar sequences that fold into the structures consistent with the target—does not increase. To address this, we present RL-DIF, a categorical diffusion model for inverse folding that is pre-trained on sequence recovery and tuned via reinforcement learning on structural consistency. We find that RL-DIF achieves comparable sequence recovery and structural consistency to benchmark models but shows greater foldable diversity: experiments show RL-DIF can achieve an foldable diversity of 29% on CATH 4.2, compared to 23% from models trained on the same dataset. The PyTorch model weights and sampling code are available on GitHub.

## 1 Introduction

The task of designing sequences of amino acids that fold into a desired protein structure, also known as protein “inverse folding” (IF) Yue & Dill (); Ingraham et al. (), is a critical step for many protein design applications in areas such as therapeutics, biomaterials, and synthetic biology Ferruz et al. (); Cao et al. (); Dickopf et al. ().
Deep learning-based IF models are typically trained to recover the observed sequence $S$ of a protein, when conditioned on the corresponding backbone structure $X$.
Models are then evaluated on their ability to recover the original sequence (“sequence recovery”), consistency of the generated sequence’s (usually predicted) structure with the target structure (“structural consistency”), and the diversity of generated sequences (“sequence diversity”).
Sequence diversity is of particular interest since the inverse folding problem is a one-to-many mapping in which a single (input) structure could be formed by multiple (output) amino acid sequences.
For practical applications, it is often desirable to find these alternative sequences that can fold into the desired structure, since it enables greater breadth of options for critical downstream optimization steps such as improving stability, preventing aggregation, and reducing immunogenicity.
However, methods of increasing sequence diversity (e.g. raising the sampling temperature of an autoregressive model) may degrade the structural consistency of candidates.
Hence, a useful property of a protein inverse folding model is its “foldable diversity”—i.e. the diversity of generated sequences that fold into the target structure.

In this paper, we first characterize the trade-off between sequence diversity and structural consistency across multiple IF methods.
We use these observations to motivate foldable diversity as a new metric to measure performance of IF models.

We then report a new protein inverse folding model, RL-DIF (Figure S3).
RL-DIF is a categorical diffusion model pre-trained on sequence recovery and fine-tuned with reinforcement learning (RL) to optimize structural consistency.
The categorical diffusion component of our method is inspired by GradeIF Yi et al. () and PiFold Gao et al. (), and is trained on observed protein structures and sequences from the CATH 4.2 dataset.
We then fine-tune the diffusion model with denoising diffusion policy optimization (DDPO) Black et al. () to maximize the expected structural consistency of designed sequences.
We find that this two-phase training results in higher-quality designs while simultaneously ameliorating the tension between sequence diversity and structural consistency.

We evaluate state of the art IF models and RL-DIF on four benchmark datasets: CATH 4.2, TS50, TS500, and CASP15, and show that our model is able to balance sequence diversity and structural consistency resulting in an increase of foldable diversity ($29\%$ with RL-DIF compared to $23\%$ with the best prior model), whilst keeping structural consistency at competitive levels to SOTA inverse folding methods.

## 2 Preliminaries

### 2.1 Discrete denoising diffusion probabilistic models

The task of protein inverse folding is to design an amino acid sequence $S$ compatible with spatial coordinates $X\in\mathbb{R}^{N\times 3}$ representing the positions of the C$\alpha$ atoms in the backbone of a protein with $N$ residues.
We represent $S$ as a sequence of one-hot vectors with vocabulary $\mathcal{V}$ consisting of the 20 naturally-occurring amino acid species.
That is, $S\in\{0,1\}^{N\times|\mathcal{V}|}$ subject to $\sum_{j}S[i,j]=1,\forall i$.

As introduced in GradeIF Yi et al. (), we use conditional discrete denoising diffusion probabilistic models (D3PM) Austin et al. () to model $p(S|X)$.
Formally, we define a forward Markov diffusion process as random variables $S_{0},\dots,S_{T}$ related by:

$$ $S_{t}\sim q(S_{t}|S_{t-1},S_{0})\equiv\mathrm{Cat}(S_{t};p=S_{t-1}Q_{t})$ (1) $$

where $Q_{1},\dots,Q_{T}$ is a sequence of $|\mathcal{V}|\times|\mathcal{V}|$ transition matrices, and $S_{0}\equiv S$ is the sequence observed in nature.
Defining $\bar{Q}_{t}=\prod_{k=1}^{t}Q_{k}$, the posterior is:

$$ $q(S_{t-1}|S_{t},S_{0})=\mathrm{Cat}\left(S_{t-1};p=\frac{S_{t}Q_{t}^{\top} \odot S_{0}\bar{Q}_{t-1}}{S_{0}\bar{Q}_{t}S_{t}^{\top}}\right)$ (2) $$

It is standard to choose $Q$s such that $q(S_{T}|S_{0})$ is a stationary distribution $p(S_{T})$.
In this case, learning a reverse Markov process $p_{\theta}(S_{t-1}|S_{t};X)$ allows us to generate novel sequences by first sampling $S_{T}\sim p(S_{T})$ and then iteratively denoising samples with $p_{\theta}$.

### 2.2 Denoising diffusion policy optimization

Assume we can assign a scalar reward $R(\hat{S})$ for every sample $\hat{S}$ from a diffusion model $p_{\theta}$ .
DDPO Black et al. () treats the reverse denoising process as a $T$-step Markov decision process and defines a policy gradient on $p_{\theta}$ to maximize the expected reward $\mathcal{J}(\theta)$:

$$ $\mathcal{J}(\theta)=\mathbb{E}_{X\sim p(X),\hat{S}\sim p_{\theta}(\hat{S}|X)} \left[R(\hat{S})\right]$ (3) $$

where $p_{\theta}(S_{0}|X)=\mathbb{E}_{S_{1},\dots,S_{T}}p(S_{T})\prod_{t=1}^{T}p_{
\theta}(S_{t-1}|S_{t},X)$. The corresponding policy gradient $\nabla_{\theta}\mathcal{J}(\theta)$ is:

$$ $\nabla_{\theta}\mathcal{J}(\theta)=\mathbb{E}_{X\sim p(X),\hat{S}_{0},\dots, \hat{S}_{T}\sim p_{\mathrm{old}}}\left[\sum_{t=1}^{T}\frac{p_{\theta}(S_{t-1}| S_{t},X)}{p_{\mathrm{old}}(S_{t-1}|S_{t},X)}\nabla_{\theta}\log p_{\theta}(S_{ t-1}|S_{t},X)R(S_{0},X)\right]$ (4) $$

where the importance sampling ratio of the target and sampling policies $\frac{p_{\theta}(\cdot)}{p_{\mathrm{old}}(\cdot)}$ allows for multiple optimization epochs per sample.

### 2.3 Inverse folding metrics

To assess the quality of designed sequences, a number of metrics are used in the protein design literature.
Here, we focus our discussion on three commonly-used ones.
In these definitions, we assume we have a protein structure $X$ with $N$ residues, an observed sequence $S$, and designed sequences $\{\hat{S}^{1},\cdots,\hat{S}^{M}\}$.

- •
Sequence recovery: The average number of amino acid positions in agreement between the observed and proposed sequences:
$\textsc{Recovery}(\hat{S}^{j};S)=\frac{1}{N}\sum_{i=1}^{N}\mathbbm{1}\left[S[i
]=\hat{S}^{j}[i]\right].$
(5)
- •
Self-consistency TM score (sc-TM): The template modeling (TM) score Zhang & Skolnick () measures the similarity of two protein structures.
The score is 1 if they are identical and tends to 0 for dissimilar structures.
sc-TM is defined as Trippe et al. ():
$\text{sc-TM}(\hat{S}^{j};S)=\textsc{TM-Score}(\textsc{Fold}(\hat{S}^{j}),
\textsc{Fold}(S))$
(6)
where Fold is any protein folding algorithm such as AlphaFold2 Jumper et al. () or ESMFold Lin et al. ().
Note that we compare designs to $\textsc{Fold}(S)$ (not $X$) to reduce the effect of biases of the folding algorithm Gao et al. ().
- •
Sequence diversity: The average fraction of amino acids that differ between pairs of designs:
$\displaystyle\textsc{Diversity}(\{\hat{S}^{1},\dots,\hat{S}^{M}\})$
$\displaystyle=\dfrac{2}{NM(M-1)}\sum_{j=1}^{M}\sum_{k=1}^{j-1}\sum_{i=1}^{N}
\mathbbm{1}\left[\hat{S}^{j}[i]\neq\hat{S}^{k}[i]\right]$
$\displaystyle=\dfrac{2}{M(M-1)}\sum_{j=1}^{M}\sum_{k=1}^{j-1}d_{\mathrm{H}}(
\hat{S}^{j},\hat{S}^{k}).$
(7)
where $d_{\mathrm{H}}$ is the Hamming distance.
We note that sequence diversity alone is not a sufficient measure of a IF method’s quality, as it can be increased arbitrarily at the expense of sample quality (e.g. as measured by structural consistency).

## 3 Methods

### 3.1 Foldable diversity

In realistic protein design scenarios, it is necessary to have a diverse pool of candidates for synthesis and experimental characterization.
We are particularly interested in the sequence diversity of IF models when restricted to faithful structures, and so we propose foldable diversity (FD):

$$ $\displaystyle\textsc{Foldable Diversity}(\{\hat{S}^{1},\dots,\hat{S}^{M}\})=$ $\displaystyle\dfrac{2}{M(M-1)}\sum_{j=1}^{M}\sum_{k=1}^{j-1}\left(d_{\mathrm{H }}(\hat{S}^{j},\hat{S}^{k})\times\mathbbm{1}\left[\min\left\{\text{sc-TM}(\hat {S}^{j};S),\text{sc-TM}(\hat{S}^{k};S)\right\}>\text{TM}_{\text{min}}\right]\right)$ (8) $$

Foldable diversity only considers the diversity between pairs of sequences that are both structurally consistent with the input protein backbone.
Unlike sequence diversity, it is not possible to increase FD by sampling high-entropy sequences, without also ensuring they fold correctly.
We therefore argue that FD is a more grounded metric for evaluating IF methods than sequence recovery (which penalizes high-quality, highly-diverse sequences) and sequence diversity (which does not consider sequence quality).

Unless otherwise stated, we set $\text{TM}_{\text{min}}=0.7$.
Among observed protein structures, this threshold typically guarantees two proteins share a topological classification Xu & Zhang (); Wu et al. (), i.e. are structurally similar.
In this regime, FD combines sequence diversity and structural consistency into a single quality measure, if we assume that all sequences with sc-TM above $\text{TM}_{\text{min}}$ are “equally good”.

### 3.2 D3PM for inverse folding

As in Austin et al. (); Yi et al. (), we train a neural network $\hat{p}_{\theta}(S_{0}|S_{t})$ to recover the posterior by integrating over the vocabulary:

$$ $p_{\theta}(S_{t-1}|S_{t})\propto\sum_{v\in\mathcal{V}}q(S_{t}|S_{t-1},v)\hat{p }_{\theta}(v|S_{t})$ (9) $$

#### Model Architecture

We parameterize the network as a modified PiFold Gao et al. () architecture, adding multilayer perceptrons (MLPs) to process the partially-noised amino acid sequence and diffusion timestep.

Specifically, given a set of backbone coordinates $\mathbf{X}\in\mathbb{R}^{4N\times 3}$, we first construct a kNN graph $(k=30)$ between residues.
We then use the PiFold featurizer to define node and edge features $h_{V}$ and $h_{E}$.
These features include distances between atoms, dihedral angles, and direction vectors.

The denoising model is a function of $h_{P}$, $h_{E}$, the partially-denoised sequence $s_{t}$, and timestep $t$. We use the following architecture:

$$ $\displaystyle h^{\prime}_{V},h^{\prime}_{E}$ $\displaystyle=\text{MLP}(h_{V}),\text{MLP}(h_{E})$ $\displaystyle h_{o}$ $\displaystyle=\text{MLP}([s_{t},t])$ $\displaystyle h_{Vs}$ $\displaystyle=[h^{\prime}_{V},h_{o}]$ $\displaystyle h^{out}_{V},h^{out}_{E}$ $\displaystyle=(10\times\text{PiGNN})(h_{Vs},h^{\prime}_{E})$ $\displaystyle p(s_{t+1}|s_{t})$ $\displaystyle=\text{MLP}([h^{out}_{V},h_{Vs}])$ $$

where $[a,b]$ represents concatenation and PiGNN is the GNN layer introduced by PiFold.
Unless otherwise noted, we set all hidden layer sizes to the recommended values in Gao et al. ().

#### Training

We observed improved performance from using the full D3PM hybrid loss, as compared to the cross-entropy loss used by Grade-IF.
We use uniform transition matrices, since we did not observe a benefit from the Block Structured Matrices Henikoff & Henikoff () used in Grade-IF.

The model was trained on structure-sequence pairs from the CATH 4.2 dataset with the train/validation/test split curated by Ingraham et al. Ingraham et al. () based on CATH topological classifications.
This dataset and split is used by most prior IF models.
It comprises 18025 samples for training, 1637 for validation and 1911 for testing.

We used the Adam optimizer Kingma & Ba (), a learning rate of $10^{-3}$, an effective batch size of 64 (distributed over 4 Nvidia A10 GPUs), and 150 diffusion time steps.
We trained for 200 epochs.

### 3.3 RL-refinement of inverse folding models

As noted above, inverse folding is a one-to-many mapping problem, with potentially a large number of sequences satisfying a conditioned structure.
Furthermore, publicly-available protein structure data represents a sparse and heterogeneous sampling of the desired distribution $p(S_{0}=S|X)$.
This one-to-many mapping and data paucity are key challenges in the generalization of conditional generative models Yi et al. (); Dauparas et al. (); Hsu et al. ().
While collecting more data is possible, this is an expensive and slow process.

On the other hand, given a proposal sequence $\hat{S}$, we may verify its suitability for the conditioned structure $X$ by evaluating the self-consistency TM score.
We therefore propose a second phase of training, in which the inverse folding diffusion model is optimized for $\mathcal{J}(\theta)=\mathbb{E}_{\hat{S}\sim p_{\theta}}[\text{sc-TM}(\hat{S};S)]$.
In particular, we use the DDPO algorithm summarized in Section 2.2.

#### Training

During the second phase of training, we use the same training dataset described in Section 3.2.
Each training step takes as input a batch of 32 protein backbone structures.
First, we sample 4 sequences per structure from the diffusion model.
Raw rewards (sc-TM) are standardized $(\mu=0,\sigma=1)$ separately for each structure: each set of 4 sequences are shifted and scaled by their mean and standard deviation respectively.
Then, we perform minibatch gradient descent on the DDPO objective over the sample sequences (batch size of 32).

Unless otherwise specified, all RL models are trained for 1000 steps.
We used the Adam optimizer, a learning rate of $10^{-5}$ and an effective batch size of 32 (again over 4 Nvidia A10 GPUs).
As in Black et al. (), we used clipping to enforce a trust region for the policy update, with a clip value of 0.2.

#### Protein folding

During this phase of training, we sampled hundreds of thousands of sequences from the policy. Folding all of these sequences is computationally challenging, so we use ESMFold Lin et al. () instead of AlphaFold2 to strike a balance between speed and accuracy.
We leave it to future work to investigate more accurate folding algorithms Jumper et al. (); Wu et al. ().

We used the ESMFold implementation in the Huggingface Transformers library  and wrapped it in an MLflow v2.3.6 framework.
This was deployed on a Kubernetes cluster containing 20 Nvidia A10 GPUs, with an application load balancer to distribute traffic evenly.
This architecture was critical to enable efficient on-policy training, as folding generated sequences is far more computationally-intensive than updating the policy.

## 4 Related Works

#### Inverse folding models

A majority of existing IF methods utilize transformers (GraphTrans Ingraham et al. () and ESM-IF Hsu et al. ()) or graph neural networks in which nodes are amino acids, edges are defined between amino acids close together in the protein structure, and node and edge features are constructed from the protein structure backbone (ProteinMPNN Dauparas et al. (), PiFold Gao et al. (), Grade-IF Yi et al. ()).
The learning objective of these methods can be discriminative Gao et al. () or autoregressive Dauparas et al. (); Hsu et al. (); Ingraham et al. ().

More recently, GradeIF Yi et al. () introduced a diffusion-based IF method.
We note that the GradeIF results are competitive, but use solvent accessible surface area (SASA) features, which are strongly correlated with amino acid identity.
These additional features are not typically part of the IF task specification and their effect is studied in Appendix A.1.

KWDesign Gao et al. () fuses information from pre-trained protein structure and language models (GearNet Lopes & Costa (), ESM2Lin et al. (), and ESMIFHsu et al. ()) to improve amino acid representations and boost sequence recovery to 61%.

#### RL for biological sequence diffusion

SEIKO Uehara et al. () proposes a framework for online tuning of a diffusion model given a reward model and compares their method to DDPO and DPOKFan et al. ().
Among the evaluated problems is the design of green fluorescence protein (GFP) sequences.
We note that the IF task differs from GFP design in that the former is a conditional generation task, and the primary goal is to generalize to conditions (i.e. structures) not observed during training.

## 5 Experiments

In this section, we benchmark RL-DIF against SOTA models on different datasets and compare their foldable diversity and structural consistency (sc-TM).

### 5.1 Benchmarking Datasets

We benchmarked on the following datasets:

- •
CATH 4.2: This is identical to the test hold-out of the data described in Section 3.2. We further use the partitioning of the test data by Ingraham et al. () into a ”short” subset of all proteins shorter than 100 residues (94 proteins); a ”single” subset of all single chain proteins (102 proteins)
- •
TS50 and TS500: TS50 and TS500 Li et al. () are curated lists of proteins of size 50 and 500 respectively from the PISCES server Wang & Dunbrack ().
- •
CASP15: The CASP15 dataset comprises of 45 protein structures used to assess the quality of forward-folding models as these structures are de-novo protein structures.

In Appendix A.2, we assess the structural and sequence similarity between proteins in these datasets and the CATH 4.2 training set.
We find that TS50, TS500, and CASP15 have cross-split overlaps (as defined by SPECTRA Ektefaie et al. ()) ranging from 42-84%, indicating strong overlap with the training set.
We therefore include these datasets for consistency with previous work, but focus the comparison on CATH 4.2

### 5.2 Model sampling strategies

To sample diverse sequences from IF models, different strategies can be employed, depending on how the model parameterizes $p(S|X)$:

- •
Single-shot models: These models parameterize $p(S|X)$ by factorizing the joint distribution into $\prod_{i=1}^{N}p(S[i]|X)$. This distribution can be reshaped by a temperature parameter. A temperature of 0 always picks the highest-probability AA, and higher temperatures encourage greater diversity.
- •
Autoregressive models: AR models return a probability distribution over possible amino acids for a given residue, conditioned on already-sampled residues. As in the single-shot case, we can reshape the conditional distributions with temperature.
If an AR model is trained with randomly-permuted residues (instead of left-to-right), then we may also randomize the order in which we decode residues. This allows diverse sampling even at a temperature of 0.
- •
Diffusion models: Since a diffusion model iteratively maps a random vector to an amino acid sequence, we can introduce diversity by repeatedly sampling from a stationary distribution $p(S_{T})$. We do not reshape the model-predicted posterior $p_{\theta}(S_{t-1}|S_{t})$.

For the specific IF models evaluated in our study, we use the following settings:

- •
ProteinMPNN: An AR model, so we use temperature sampling with temperatures of $0.1$, $0.2$, and $0.3$. We also sample at a temperature of 0 0 with random decoding order.
- •
PiFold: A single-shot model, so we use temperature sampling with values $0.1$ and $0.2$.
- •
KWDesign: Another single-shot model, so we use the same settings as PiFold.
- •
ESMIF: An AR model. We follow the authors’ recommendation and sample with a temperature of 1.
- •
DIF-Only and RL-DIF: Diffusion models, so we sample from the uniform distribution $p(S_{T})$ and iteratively denoise sequences using the model.

### 5.3 Benchmarking and performance of IF models trained on CATH4.2 benchmark

#### Experimental Setup

We compare RL-DIF to ProteinMPNN Dauparas et al. (), PiFold Gao et al. (), and KWDesign Gao et al. ().
We also evaluate our model after the diffusion pre-training phase, before any RL-optimization (DIF-Only).
For each model and each benchmark dataset, we run the following pipeline:

- 1.
For each protein structure in the dataset, sample 4 sequences from the model.
- 2.
Among the 4 sampled sequences, compute the mean sequence recovery and sc-TM score.
- 3.
Compute the foldable diversity of the 4 sequences, as defined in Equation 8, with $\text{TM}_{\text{min}}=0.7$.

Model weights were not available for KWDesign for sampling, however they provided the probability distributions for CATH4.2 and CASP15, and we sampled from these to generate sequences.

Table: Table 1: Benchmarking results of ProteinMPNN, PiFold, KWDesign, DIF-Only, and RLDIF on CATH-single, CATH-short, CATH-all, TS50, TS500, and CASP15 reporting foldable diversity, sc-TM, and sequence recovery at various temperatures. T=X indicates temperature sampling occured at a temperature of X. RD indicates random decoding order. Best results for each dataset are bolded.

#### Results

Among the evaluated models, our proposed methods (DIF-Only and RL-DIF) achieve or match the highest foldable diversity on all benchmarks in Table 1.
In particular, we highlight the significantly-increased diversity (29% versus the next-best 23%) on the CATH-all dataset (which includes multi-protein complexes), with structural consistency only slightly degrading or improving.

Figure 2 demonstrates that the improvement in foldable diversity is minimally dependent on the choice of $\text{TM}_{\text{min}}$.
Despite our methods not always achieving the highest sc-TM score, they offer the highest diversity across a range relevant sc-TM thresholds.

Although our models do not have high sequence recoveries, we call attention to the argument put forth in Section 3.1: sequence recovery prefers sequences very close to the naturally-observed sequence, which is a small slice of the accessible design space.
We therefore focus our comparison on foldable diversity and structural consistency.

Across the various benchmarks, we observe a trend between DIF-Only and RL-DIF: RL-tuning consistently improves structural consistency, but frequently at a cost to diversity.
This is implies that the DDPO objective is being effectively optimized, but the entropy of the policy is decreasing during the second training phase.
We leave it to future work to incorporate stronger exploration strategies, which may reduce this effect.

To demonstrate the effect of high foldable diversity, we present example protein sequences in Figure 1.
All models inverse fold structurally self-consistent sequences, but with varying degrees of diversity.

Figure: Figure 1: Samples from RL-DIF, ProteinMPNN, and PiFold, conditioned on the same protein backbone (PDB: 5FLM.E). The diversity of amino acids at each position is colour-coded, ranging from dark-red for “all identical” to no colour for “all different”. While all models achieve high structural consistency (ESMFold-ed structures shown on the right), RL-DIF generates the most diverse set.
Refer to caption: x1.png

Figure: Figure 2: To evaluate the sensitivity of our analysis to the value of $\text{TM}_{\text{min}}$, we scan a range of values. We observe that DIF-Only and RL-DIF consistently perform the best, for thresholds $\geq 0.4$. Foldable diversity is computed across the ”all” CATH 4.2 test split using the definition in Section 3.1.
Refer to caption: extracted/5946617/figures/Effective_div_vs_TM_min.png

### 5.4 Comparison of RLDIF performance to ESM-IF

#### Experimental Setup

ESM-IF Hsu et al. () is another protein inverse folding model that was trained on CATH 4.3 and 12 million synthetic data samples generated by AlphaFold 2 Jumper et al. ().
While a large database of AlphaFold-ed structures is publicly available (2024) (), the ESM-IF paper did not release the IDs of the proteins that were trained on.
These IDs are essential in order to replicate their work and ensure no overlap between the training and benchmark datasets.

To attempt to replicate their work, we followed a similar strategy of curating a larger pre-training dataset orthogonal to all benchmarking datasets we consider.
This was done using Foldseek van Kempen et al. () to ensure no structure or sequence overlap, resulting in a dataset of 100,000 structures from the AlphaFold database.
We then retrained RL-DIF on this larger pretraining dataset, followed by the same 1000 steps of RL optimization.
We refer to these models as DIF-Only-100K and RL-DIF-100K.
In order to faciltate future work we release the IDs of the entries used to train RL-DIF-100K so future work can train larger pre-trained protein inverse folding models.

#### Results

We found that, as compared to RL-DIF, RL-DIF-100K improves foldable diversity across all datasets.
As demonstrated in Table 2, DIF-Only-100K and RL-DIF-100K approach (and in the case of TS50, surpass) ESM-IF’s performance, despite using 70 times less data and 30 times fewer parameters.

Given the increased performance from expanding the pre-training dataset from 18K (CATH training set) to 100K, we leave it to future work to explore the effect of further dataset scaling.
Ideally, we would have liked to extend the comparison to use 12M structures as done by the ESM-IF authors, but found this computationally intractable with our resources (in particular, Foldseek has linear memory complexity).
In Figure S4, we show all considered IF models, comparing their foldable diversity and amount of training data.

Table: Table 2: Benchmarking results of ESM-IF, DIF-Only-100k, and RL-DIF-100K models on various datasets reporting foldable diversity, sc-TM, and sequence recovery. * indicates both results are statistically equivalent (p-value $>$ 0.05)

### 5.5 The effect of reinforcement learning on DIF-Only performance

To investigate the effect of reinforcement learning on DIF-Only performance we ran our RL for 4000 steps, taking the model at every 1000 steps and benchmarking its performance for all metrics. We found average sc-TM had the largest improvement in the first 1000 steps of RL before stagnating after step 2000 (Table 3). Foldable diversity continued to decrease with modest gains in structural consistency. These results support our decision to run RL for 1000 steps.

Table: Table 3: Effect of RL-tuning of DIF-Only-100K, as measured by performance on the TS50 dataset.

## 6 Conclusion

Although sequence recovery, sequence diversity, and structural consistency are commonly used to evaluate protein IF models, here we demonstrate that those metrics alone do not necessarily capture the ability of the model to generate multiple sequences that fold into the desired structure.
For many practical applications, foldable diversity is a useful metric since it gives users multiple “shots-on-goal” for filtering or optimizing designs on criteria beyond structural consistency.
We also present new inverse folding models: DIF-Only and RL-DIF.
We find that among evaluated methods, these models achieves highest foldable diversity and sc-TM in the CATH4.2, CASP15, TS50, and TS500 datasets.

Nonetheless, we note some limitations of RL-DIF. First is the exclusive use of ESMFold as the folding model. Alternative models such as AlphaFold2 Jumper et al. (), OpenFold Ahdritz et al. (), or OmegaFold Wu et al. () have demonstrated superior folding accuracy and could potentially enhance our results. An ensemble of these models could also be employed to leverage the specific strengths and mitigate the weaknesses of each individual model.
Additionally, instead of limiting our sampling to four iterations, we could continuously sample until achieving four samples with a sc-TM greater than 0.7. While this approach would likely improve model accuracy, it also significantly increases computational cost. Consequently, we imposed a limit of four samples to manage computational resources. Finally, we could leverage entropy bonus methods used in RL to further increase sequence diversity and exploration of the policy.

### Code Availability

RL-DIF model weights and sampling code is currently available at [https://github.com/flagshippioneering/pi-rldif](https://github.com/flagshippioneering/pi-rldif).

## Appendix A Appendix

### A.1 GradeIF Results

When exploring the performance of GradeIF we found the
solvent accessible surface area (SASA) feature is calculated utilizing side chain information. When plotting the SASA values per amino acid type, there is a clear separation between some amino acids which can inflate model performance (Figure S1). When we retrained GradeIF with recommended hyperparameters without (Figure 2(a)) and with SASA (Figure 2(b)) we noticed a 23% decrease in performance. With this in mind, we decided to pursue utilizing PiFold Gao et al. () as the underlying architecture for the categorical diffusion.

Figure: Figure S1: SASA feature value for every amino acid in CATH4.2 training set proteins.
Refer to caption: x2.png

Figure: (a) GradeIF performance before removal of surface features. The top half of the plot shows train (blue) and val (orange) loss. The bottom half of the plot shows test perplexity (orange) and sequence recovery (blue).
Refer to caption: extracted/5946617/figures/GradeIF_Before_Removal.png

### A.2 SPECTRA and Foldseek

To understand how model performance changes as a function of test protein similarity to the CATH4.2 training set, we utilize the spectral framework of model evaluation (SPECTRA) Ektefaie et al. (). We utilized Foldseek van Kempen et al. () to determine if two proteins were similar. Two proteins are similar if e-value returned by Foldseek is greater than 1e-3 or the sequence similarity is greater than 0.3. Given two datasets, cross-split overlap is defined as the proportion of proteins that are similar.

Figure: Listing 1: Commands used to run Foldseek

Using this procedure we evaluated the structural sequence similarity between each benchmarking set and the CATH 4.2 training set. We found TS50, TS500, and CASP15 datasets had high levels of similarity to the CATH 4.2 train dataset (Table S1).

Table: Table S1: CATH 4.2 all, CATH 4.2 single, CATH 4.2 short, TS50, TS500, and CASP15 test set cross-split overlap with CATH 4.2 train set. Cross-split overlap is defined as the proportion of proteins in dataset that had similar sequence or structure to a protein in the CATH 4.2 train set (Appendix A.2).

### A.3 Supplementary tables and figures

Table: Table S2: Benchmarking results of ESMIF and augmented DIF-Only and RLDIF models on CATH 4.2 test hold-out structures reporting foldable diversity, sc-TM, and sequence recovery.

Figure: Figure S3: The overall framework for RL-DIF. Training phase 1: RL-DIF is pretrained to generate amino acid sequences conditioned on protein backbone structures using a categorical diffusion objective. Training phase 2: RL-DIF is refined to maximize the expected structural consistency of its generated sequences.
Refer to caption: x3.png

Figure: Figure S4: The effect of data scale on IF model performance. In general, we observe a positive relationship between additional data and foldable diversity, though this is not readily observable for structural consistency (indicated by color and labels). Note that, in the case of RL-tuned models, we include the sequences sampled during on-policy training in the dataset size.
Refer to caption: extracted/5946617/figures/cath4.2.png

Figure: Figure S5: Sequence designs from 3 different IF models for a short single alpha-helix peptide structure (PDB: 2X7R.B)
Refer to caption: x4.png

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
questions: The expectation after eq (3) should really be a sum

q(S_t|S_{t-1},v) should be q(S_{t-1}|S_t,v) in eq (9) 

In the model architecture, line 183, do you mean p(s_0|s_t) rather than p(s_{t+1}|s_t) as currently written? 

Table 1 would be easier to read if summarized in terms of curves. For example, calculate diversity within all generated samples from a method that pass a sc-TM threshold and then vary this threshold.
soundness: 3
strengths: The paper is clearly written and the evaluations are fairly extensive, involve a number of relevant baselines, and seem conducted with care.
summary: The paper focuses on inverse folding, i.e., recovering a residue sequence that would fold into a given 3D backbone structure. The key motivation for the paper is specifically to improve diversity of generated sequences. The method consists of a discrete diffusion model that is then fine-tuned using RL (policy gradient) to improve foldability (sc-TM). Evaluation includes a number of baseline methods, including ProteinMPNN, PiFold, etc.
weaknesses: The proposed method is largely a direct combination of available techniques. Discrete diffusion component seems to be a variant of GradeIF/D3PM with small changes, and the RL fine-tuning is based on the rather straightforward policy gradient method (e.g.) from Black et al. I don't really see generalizable methodological innovations. Of course, it's fine to mix and match methods with adjustments to produce a well-performing technique but then the paper rests on its empirical results. 

I think it would be better to replace eq (8) with a measure calculated only within foldable sequences so as to separate foldability from diversity (so they can be assessed separately and together). Diversity was taken as the key motivation for the paper but the two notions (diversity and foldability) are now mixed together in the proposed metric. E.g., the downward slope in Figure 2 is in part just due to the fact that higher folding threshold automatically decreases the value of the current metric. Figure 2 can also give a wrong impression that RL-diff increases diversity since the metric can be improved by increasing foldability instead. In fact, the reward used for fine-tuning the diffusion model is sc-TM so tailored to increase foldability (not control diversity).

---

confidence: 4
contribution: 3
initial_rating: 5
presentation: 3
questions: Could the authors could provide (A) more information (plots) on the trade off between structural similarity and sequence diversity along the RL training trajectory (b) More information on how hyper parameter can be used to relax structural similarity constant to increase diversity (c) exploration of regularization strategies like entropy based regularization to increase the diversity of sequences generated about the gains shown. (d) validation of structural similarity of generated sequences via alpha fold.
soundness: 3
strengths: Paper tables an important problem in inverse folding with an interesting solution in the combination of diffusion model and RL. It's an exciting contribution that could open up a new set of metrics for inverse folding models. The diversity of inverse folding models is a known problem in practice, but few direct attacks have been taken. The authors develop an interesting approach that shows improvement on sequence diversity metrics relative to base line models.
summary: This paper introduces RL-DIF, a novel reinforcement learning-driven model for protein inverse folding, addressing the need for generating diverse ensembles of amino acid sequences that fold into specific, target 3D protein structures via inverse folding. Authors introduce a framework they call RL-DIF that applies a categorical diffusion approach pre-trained on sequence recovery, followed by reinforcement learning (RL) fine-tuning to optimize structural consistency. This two-phase training strategy enables RL-DIF to improve "foldable diversity"—the diversity of sequences that maintain the correct fold—while retaining sequence recovery and structural accuracy. In tests on CATH 4.2 and other datasets, RL-DIF achieved up to 29% foldable diversity, outperforming prior models, which peaked at 23%.
Key contributions of the paper include:
RL-DIF Model: A diffusion-based model refined with denoising diffusion policy optimization to balance diversity and structural fidelity.
Benchmarking Results: Demonstrating RL-DIF’s superior foldable diversity on multiple datasets.
Methodological Advancements: Introducing foldable diversity as a new metric to assess inverse folding model quality.
The RL-DIF model code and PyTorch weights are available in HF.
weaknesses: The paper feels close to the acceptance threshold for me. However,I felt the authors could do a bit more. The increase in seq diversity generated by RL-DIF is modest in some cases. The authors could provide (A) more information on the trade off between structural similarity and sequence diversity along the training trajectory (b) More information on how hyper parameter can be used to relax structural similarity constant to increase diversity (c) exploration of regularization strategies like entropy based regularization to increase the diversity of sequences generated about the gains shown. (d) validation of structural similarity of generated sequences via alpha fold.

---

confidence: 5
contribution: 2
initial_rating: 5
presentation: 3
questions: - How are the temperature values for the pretrained models selected? In Table 1, T=0.1 is much better than T=0.2 and T=0.3 for all baselines. How do they perform with a smaller T, eg. T=0.05 or T=0.01?
- The reviewer understands that it is computationally expensive to train RL-DIF on the same amount of data as ESM-IF. For a rough comparison, how is the performance of ESM-IF trained on the same dataset as RL-DIF-100K?
soundness: 2
strengths: - The proposed method combines recent advances in discrete diffusion and reinforcement learning alignment methods to improve the sequence diversity of protein inverse folding, an essential task in protein biology.
- The paper is easy to follow, and the experiments on various protein datasets show the improvement in diversity with the proposed method.
summary: The paper presents RL-DIF, a discrete diffusion model trained for protein inverse folding and further fine-tuned on structural consistency. RL-DIF achieves comparable sequence recovery and structural consistency to existing methods, while having a better foldable diversity, beneficial to downstream optimizations.
weaknesses: - The novelty of RL-DIF is limited. Several existing papers utilize discrete diffusion models for protein inverse folding, for example, [1] and [2]. The idea of using structural consistency to fine-tune the inverse folding model has been explored in ESM3[3], although with language model as the base model.
- The motivation for fine-tuning the model with structural consistency to improve the sequence diversity is not very clear. This is also reflected in the experiments in Table 1, where DIF-Only achieves higher diversity in most datasets except CATH-short and CATH-single. 
- The Foldable Diversity is only calculated based on 4 generated sequences for each protein structure, which can have high variance. It is better to use a higher number of generated sequences, and report the standard deviation of the results across random seeds. Besides, it is also helpful to provide the ratio of the generated sequences that satisfy the sc-TM score constraint to give some idea of to what degree the Foldable Diversity is affected by the structural consistency.
- The model is compared to only non-diffusion based inverse folding baseline methods. A comparison with diffusion-based methods [1,2] is missing.
- Hamming distance is used to measure sequence diversity. However, it may not be able to capture the high-level diversity between sequences, which is more related to protein functions and more useful in practice. Thus, the distance calculated in a pretrained embedding space (eg. ESM embedding) could be better.

[1] Generative Flows on Discrete State-Spaces: Enabling Multimodal Flows with Applications to Protein Co-Design. ICML 2024.

[2] Unlocking Guidance for Discrete State-Space Diffusion and Flow Models. arXiv 2024.

[3] Simulating 500 million years of evolution with a language model. bioRxiv 2024.

Meta-review: This paper constructs a discrete diffusion model combined with RL fine tuning to solve the inverse protein folding problem, with a particular focus on high "foldable diversity": sequences that are dissimilar to each other but nevertheless fold into the correct target. The reviewers shared concerns about a lack of methodological novelty (L2jQ, 6Vm7), and concerns about the mixing of foldability and diversity (L2jQ) when these could be separated and analyzed as a trade-off, which 6Vm7 also points out in the author discussion period. Reviewer R2ES similarly asks for data on the trade off between structural similarity and sequence diversity. Ultimately, ignoring for the moment novelty concerns, which are often vague and hard to tackle concretely, I think the paper would be much improved by taking to hear the unanimous point about trade-offs in these metrics, and plotting these metrics as axes in a 2D plot against each other rather than reporting or even combining them.

Decision: Reject