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

TL;DR

RL-DIF is presented, a categorical diffusion model for inverse folding that is pre-trained on sequence recovery and tuned via reinforcement learning on structural consistency that achieves comparable sequence recovery and structural consistency to benchmark models but shows greater foldable diversity.

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.

Figures (7)

• 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.
• 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 \ref{['sec:eff_diversity']}.
• Figure S1: SASA feature value for every amino acid in CATH4.2 training set proteins.
• Figure S2: GradeIF performance before and after removal of surface features.
• 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.