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

arXiv: https://arxiv.org/abs/2410.17173 · HTML: https://mlsanity.com/paper/2410.17173/

## Key facts

- Cited by 6 papers in the map (data through 14 Oct 2026).
- Placed in RL for energy, networks and operations, one of 15 directions in Deep reinforcement learning methods.
- Builds on 6 of its 33 references that are in the map.
- First posted on arXiv 22 Oct 2024.

## Placed in

- Deep reinforcement learning methods / [RL for energy, networks and operations](https://mlsanity.com/deeprl/rl-for-energy-networks-and-operations/)

## Builds on (most cited first, 6 of 6)

- [Training Diffusion Models with Reinforcement Learning](https://mlsanity.com/paper/2305.13301/) (2023)
- [Generative Flows on Discrete State-Spaces: Enabling Multimodal Flows with Applications to Protein Co-Design](https://mlsanity.com/paper/2402.04997/) (2024)
- [Diffusion probabilistic modeling of protein backbones in 3D for the motif-scaffolding problem](https://mlsanity.com/paper/2206.04119/) (2022)
- [Unlocking Guidance for Discrete State-Space Diffusion and Flow Models](https://mlsanity.com/paper/2406.01572/) (2024)
- [PiFold: Toward effective and efficient protein inverse folding](https://mlsanity.com/paper/2209.12643/) (2022)
- [Graph Denoising Diffusion for Inverse Protein Folding](https://mlsanity.com/paper/2306.16819/) (2023)

## Built on by (most cited first, 6 of 6)

- [Iterative Distillation for Reward-Guided Fine-Tuning of Diffusion Models in Biomolecular Design](https://mlsanity.com/paper/2507.00445/) (2025)
- [BioNeMo Framework: a modular, high-performance library for AI model development in drug discovery](https://mlsanity.com/paper/2411.10548/) (2024)
- [Generative Modeling in Protein Design: Neural Representations, Conditional Generation, and Evaluation Standards](https://mlsanity.com/paper/2603.26378/) (2026)
- [SurfDesign: Effective Protein Design on Molecular Surfaces](https://mlsanity.com/paper/2606.07567/) (2026)
- [ABOPD: Antibody CDR Design via On-Policy Distillation](https://mlsanity.com/paper/2607.18835/) (2026)
- [MORAD: Reinforcement Learning with Multi-Objective Rewards for RNA Inverse Design](https://mlsanity.com/paper/2610.13305/) (2026)

## Abstract (arXiv)

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.