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

AI summary

Overview

Research area: Computational RNA biology and generative machine learning — specifically structure-based RNA inverse folding (designing an RNA sequence that folds into a given 3D structure), combining latent diffusion models (LDMs) with reinforcement learning (RL).

Technical level: Advanced. Understanding the method requires familiarity with diffusion/denoising models, DDIM sampling, policy-gradient reinforcement learning (PPO), and RNA structural metrics (secondary structure, minimum free energy, RMSD, LDDT).

Scope: The paper proposes SOLD (Step-wise Optimization of Latent Diffusion Model), an RL framework that refines a pre-trained latent diffusion model to directly optimize non-differentiable RNA structural objectives, and benchmarks it against physics-based, heuristic, and deep learning baselines on two test datasets.

What This Paper Is About

Most RNA inverse folding methods are trained mainly for sequence recovery — matching the target sequence nucleotide by nucleotide — which does not guarantee that the designed sequence actually achieves desired structural properties. SOLD addresses this by pairing a latent diffusion model built on RNA-FM embeddings with a reinforcement learning stage that directly optimizes structural objectives (secondary structure consistency, minimum free energy, and LDDT) that cannot be optimized by ordinary gradient descent. The goal is to generate RNA sequences that balance sequence naturalness with structural fidelity.

Key Contributions

1. A latent diffusion model that uses pre-trained language-model embeddings. SOLD encodes RNA-FM embeddings of shape (L, 640) into a compressed latent space (L, D) with an MLP encoder, performs denoising with a GVP-GNN plus Diffusion Transformer (DiT) network, and decodes back to a (L, 4) nucleotide probability distribution, improving sequence recovery over one-hot sequence diffusion.

2. The first integration of reinforcement learning into a latent diffusion model for RNA inverse folding, enabling optimization of complex structural objectives that existing diffusion-based methods cannot handle because they are non-differentiable.

3. A step-wise RL optimization algorithm. Rather than sampling full denoising trajectories, SOLD randomly samples a timestep, predicts the clean latent (or the intermediate latent with k = 1) in a single step, decodes it, and evaluates a reward — achieving faster convergence than trajectory-based RL methods such as DDPO and DPOK.

4. A piecewise long-term/short-term reward scheme that uses short-term rewards at early (high-noise) denoising steps and long-term rewards at later steps, combined with a KL-constrained clipped PPO objective (clip range 0.0001).

Main Findings

• LDM alone beats prior methods on 1D metrics. Before any RL fine-tuning, SOLD's LDM achieved the highest Sequence Recovery and NT Recovery on both SOLD TEST (0.5728 / 0.5034) and CASP15 TEST (0.5462 / 0.5473), compared with RhoDesign (0.2734 / 0.2859 and 0.2606 / 0.2575), RDesign (0.4457 / 0.3966 and 0.3264 / 0.3251), gRNAde (0.5108 / 0.4890 and 0.5097 / 0.5149), RiboDiffusion (0.5125 / 0.4416 and 0.5388 / 0.3871), and DRAKES-Pretrain (0.4524 / 0.4088 and 0.3357 / 0.3374).

• Single-objective RL fine-tuning improves on the LDM baseline and competitors. For MFE (lower is better), SOLD reached -19.7428 on SOLD TEST and -68.2100 on CASP15 TEST, versus DRAKES-Pretrain (-12.5123 / -52.6827), DRAKES (-14.2374 / -61.0354), LDM (-13.1519 / -52.7387), DDPO (-18.7498 / -63.9567), and DPOK (-17.4660 / -67.7949).

• Gains also hold for SS and LDDT. SS (higher is better): SOLD 0.7551 / 0.7010 on SOLD TEST / CASP15 TEST versus LDM 0.7274 / 0.5543, DDPO 0.7595 / 0.6649, and DPOK 0.7511 / 0.6303. LDDT: SOLD 0.6384 / 0.3548 versus LDM 0.6184 / 0.3237, DDPO 0.6286 / 0.3406, and DPOK 0.6329 / 0.3351.

• Step-wise optimization is dramatically faster per epoch. Average training time per epoch: for MFE, DDPO 5953 s, DPOK 7677 s, SOLD 256 s; for SS, 6190 s, 7330 s, and 263 s; for LDDT, 14000 s, 14200 s, and 6900 s. The authors attribute the smaller LDDT speedup to the cost of structure prediction.

• Weighted multi-objective optimization improves structural fidelity. Under equal weighting of SS, MFE, and LDDT (with MFE mapped to (0,1) via exp(1/(MFE - 1/4))), SOLD on SOLD TEST achieved Sequence Recovery 0.5732, MFE -16.8611, SS 0.7601, RMSD 11.8612, and LDDT 0.6360, compared with LDM's 0.5728, -13.3275, 0.7269, 12.5732, and 0.6178. On CASP15 TEST, SOLD reached 0.5888, -64.0375, 0.6957, 26.8422, and 0.3680. DRAKES recorded a higher SS on SOLD TEST (0.7691) than SOLD (0.7601), while SOLD led on the other metrics.

• A riboswitch case study succeeded. For a TPP-specific riboswitch example (PDB: 3D2V), SOLD designed a sequence folding into the target structure whereas, per the authors, other methods produced conformations distant from the goal.

• Latent dimension ablation shows a sharp dependence. Test-set recovery rose with latent dimension D: 0.5161 (D = 8), 0.8746 (D = 16), 0.9768 (D = 32), 0.9955 (D = 64), and 0.9989 (D = 128). The main model uses D = 32.

• Dataset scale. After clustering with PSI-CD-HIT (sequence threshold 0.3) and US-align (structural threshold 0.45), 8222 structures were used: 7067 for pre-training, 389 for RL fine-tuning, and 766 for SOLD TEST, with CASP15 TEST as an independent benchmark. All experiments ran on a single A100 GPU.

Methodology in Plain English

The system has two stages. In the first, the researchers take an RNA sequence's representation from RNA-FM — a large pre-trained RNA language model whose embeddings carry co-evolutionary information — and compress it from 640 dimensions down to a small latent vector using a three-layer MLP encoder. A denoising network made of a 4-layer GVP-GNN (node hidden dimension 512, edge hidden dimension 128, with dihedral angle features) plus an 8-layer Transformer (embedding size 512, 16 attention heads, dropout 0.2) learns to predict the clean latent from a noised version, conditioned on the 3D backbone geometry. A three-layer MLP decoder then turns that latent into probabilities over the four nucleotides. Training minimizes a combination of mean squared error on the latent and cross-entropy on the decoded sequence.

The second stage is where the paper diverges from prior work. Instead of running the full diffusion trajectory during training, SOLD samples a random timestep, makes a single-step DDIM prediction of the clean latent, decodes it to a sequence, and scores it with a reward function. The reward is computed directly by ViennaRNA (for secondary structure and minimum free energy) and RhoFold (for 3D structure metrics), so no separately trained, differentiable reward model is needed. Because single-step predictions are unreliable at high noise levels, the authors add a short-term reward (using k = 1 to predict the intermediate latent) and blend it with the long-term reward through time-dependent weights: short-term only at early steps, long-term only at late steps. Model parameters are then updated with PPO using a clipped surrogate objective plus a KL penalty keeping the policy close to the reference model. At inference time, the model reverts to full trajectory denoising to preserve output quality.

Why This Matters

Impact on research. The paper argues that training RNA design models on sequence recovery alone leaves structural objectives unaddressed, and that RL can fill that gap for non-differentiable metrics. It also shows that a latent diffusion model operating on RNA-FM embeddings outperforms diffusion directly in one-hot sequence space, and that single-step RL optimization can be far cheaper than trajectory-based RL. This provides a modular template — the reward component is described as fully pluggable — for aligning generative molecular models to structural objectives.

Real-world applications (as framed by the authors):

• RNA therapeutics — designing sequences with controlled folding behavior for therapeutic use.
• Gene regulation — rationally designed riboswitches that can precisely control mRNA translation, illustrated by the TPP-specific riboswitch case study.
• Synthetic biology — producing RNA sequences that satisfy stringent structural constraints for engineered biological systems.
• Molecular design pipelines — substituting more accurate structure-prediction or energy-evaluation tools into the reward slot as they become available.

Industry relevance. The reward-model-free design reduces computational overhead, and the reported training speed advantage of SOLD over DDPO and DPOK (for example, 256 s versus 5953 s and 7677 s per epoch for MFE) points to lower cost for iterative RNA design workflows. The code is released at https://github.com/darkflash03/SOLD.

Future Directions

• Expand high-quality RNA structural data. The authors state that SOLD's performance is constrained by limited availability of high-quality RNA structural data and list dataset expansion as future work.
• Improve reward evaluation. ViennaRNA and RhoFold introduce approximation errors that may affect optimization accuracy; the authors propose refining rewards with better predictors.
• Study how 1D, 2D, and 3D metrics interact. The paper notes it has not extensively explored how these metrics coordinate, and calls for synergistic multi-scale optimization.
• Generalizability and robustness. The authors target broader robustness and applicability across diverse RNA design challenges, with the modular reward component as the mechanism for swapping in improved models.

Target Audience

This paper is most useful to machine learning researchers working on generative models for molecular design, computational biologists and bioinformaticians working on RNA structure and inverse folding, and practitioners in RNA therapeutics or synthetic biology who need designed sequences to satisfy structural constraints rather than only sequence-identity objectives. Readers should already be comfortable with diffusion models, reinforcement learning from policy gradients, and standard RNA structural metrics.

Authors’ 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.