Structure-based RNA Design by Step-wise Optimization of Latent Diffusion Model

TL;DR

This work tackles RNA inverse folding by designing sequences that realize target 3D structures, addressing limitations of physics-based and 2D-structure methods. It introduces SOLD, a latent diffusion model that leverages pre-trained RNA-FM embeddings and backbone geometry, combined with a novel step-wise reinforcement learning protocol that optimizes multiple structural objectives ($SS$, $MFE$, $LDDT$) without differentiable reward models. The approach uses a single-step DDIM-based refinement during training and a piecewise reward scheme with PPO optimization, achieving faster convergence and superior multi-objective performance compared to state-of-the-art baselines. Empirical results on diverse RNA targets demonstrate strong gains in sequence fidelity and structural accuracy, highlighting SOLD's potential for RNA therapeutics and synthetic biology, with room for improvement in data diversity and reward evaluation accuracy.

Abstract

RNA inverse folding, designing sequences to form specific 3D structures, is critical for therapeutics, gene regulation, and synthetic biology. Current methods, focused on sequence recovery, struggle to address structural objectives like secondary structure consistency (SS), minimum free energy (MFE), and local distance difference test (LDDT), leading to suboptimal structural accuracy. To tackle this, we propose a reinforcement learning (RL) framework integrated with a latent diffusion model (LDM). Drawing inspiration from the success of diffusion models in RNA inverse folding, which adeptly model complex sequence-structure interactions, we develop an LDM incorporating pre-trained RNA-FM embeddings from a large-scale RNA model. These embeddings capture co-evolutionary patterns, markedly improving sequence recovery accuracy. However, existing approaches, including diffusion-based methods, cannot effectively handle non-differentiable structural objectives. By contrast, RL excels in this task by using policy-driven reward optimization to navigate complex, non-gradient-based objectives, offering a significant advantage over traditional methods. In summary, we propose the Step-wise Optimization of Latent Diffusion Model (SOLD), a novel RL framework that optimizes single-step noise without sampling the full diffusion trajectory, achieving efficient refinement of multiple structural objectives. Experimental results demonstrate SOLD surpasses its LDM baseline and state-of-the-art methods across all metrics, establishing a robust framework for RNA inverse folding with profound implications for biotechnological and therapeutic applications.

Figures (11)

• Figure 1: The LDM encodes RNA-FM embeddings into a latent representation, performs denoising with GVP-GNN + DiT blocks, and decodes the refined latent into RNA sequences consistent with structural constraints.
• Figure 2: An overview of the SOLD. SOLD utilizes long-term and short-term reward feedback to directly optimize trained latent diffusion model in a random denoising step.
• Figure 3: rewards for SOLD, DPOK, and DDPO with MFE, SS, and LDDT as reward objectives: (a) MFE, (b) SS, (c) LDDT.
• Figure 4: Comparison of rna design methods for example (PDB: 3D2V), ground-truth structure (gold), SOLD (blue)
• Figure 5: Sequence length distribution across Pre-training, RL Fine-tuning, and SOLD TEST datasets.