AI810 Blog Post (20255250) [gRNAde - Geometric Deep Learning for 3D RNA Inverse Design]

Chaitanya K. Joshi, Arian R. Jamasb, Ramon Viñas, Charles Harris, Simon V. Mathis, Alex Morehead, Rishabh Anand, Pietro Liò (ICLR 2025)

Motivation and Background

Designing RNA sequences that fold into a desired structure is a classic inverse folding problem in computational biology. Traditional methods often focus on matching a given secondary structure (2D base-pair pattern) and typically consider only a single target conformation. However, real RNA molecules are highly dynamic: a single sequence can adopt multiple distinct 3D conformations, which is crucial for functions like riboswitches and ribozymes. Prior to gRNAde, the use of modern machine learning for RNA 3D design lagged behind protein design, due to limited RNA 3D data and the complexity of RNA flexibility. gRNAde (geometric RNA design) addresses this gap by explicitly accounting for RNA 3D structure and conformational diversity in sequence design. This aligns with recent trends in geometric deep learning that leverage graph neural networks (GNNs) for biomolecular structures, now extending those advances from proteins to the RNA world.

The key motivation is to enable multi-state RNA design: instead of designing a sequence for just one rigid structure, design sequences that can stably fold into multiple specified shapes. This reflects real use-cases like riboswitches, which toggle between shapes, or RNAs that need to maintain structure under different conditions. Prior physics-based tools (like Rosetta) and learning models were limited to single states or 2D structure targets. gRNAde’s innovation is a learning-based pipeline that takes one or more RNA backbone conformations as input and produces nucleotide sequences that are predicted to fold into all those conformations. By incorporating 3D geometry and multiple conformers, gRNAde aims to improve design accuracy (recovering native sequences, achieving functional folds) and broaden the design space beyond static structures.


Methodology

Multi-State Geometric RNA Design Pipeline

gRNAde introduces a multi-stage model comprising a geometric multi-graph encoder and an autoregressive decoder. The overall goal is to estimate the probability of a sequence given the target RNA backbone structure(s), and then to sample sequences that maximize this probability. Formally, for a target RNA of length L with one or more backbone conformations, gRNAde models the conditional distribution and generates a sequence $x = (x_1, x_2, \dots, x_L)$ nucleotide by nucleotide as:

\[P(x_1, \dots, x_L \mid \text{backbone}(s)) = \prod_{i=1}^{L} P(x_i \mid \text{backbone}(s), x_1, \dots, x_{i-1}),\]

decoding from the 5’ end to 3’ end. This auto-regressive factorization allows the model to incorporate already-designed prefix bases when predicting the next base.

Geometric Multi-Graph Neural Network Encoder

To featurize the 3D RNA backbones, gRNAde represents each RNA conformation as a geometric graph. Nodes correspond to nucleotides (located at the C4’ atom of the sugar), and edges capture local and spatial relationships between nucleotides. Notably, the authors use a 3-bead coarse-grained model: each nucleotide is represented by three key atoms (P, C4’, and the base’s N1/N9) to capture the backbone geometry in sufficient detail. Edges are formed between each nucleotide and its $k$ nearest neighbors in 3D (with $k=10$ in the implementation) based on C4’ distances. This builds a graph where spatially proximal nucleotides (which may interact or stack) are connected, beyond just the sequential neighbors. Rich geometric features are assigned to nodes and edges: for example, each node gets orientation vectors (pointing along the backbone direction and from sugar to base) akin to features used in protein GNNs. Edges similarly carry geometric features like the 3D displacement (as a unit vector), distance encoded by radial basis functions, and even the along-backbone distance encoded sinusoidally. This geometric featurization ensures the model can learn the relationship between sequence and 3D shape.

Crucially, when multiple backbone conformations are provided, gRNAde does not simply merge them or choose one – instead, it constructs a “multi-graph” representation that preserves each conformation’s distinct information. If there are M conformations given, gRNAde builds M separate graphs (each with the same set of node indices, corresponding to the same nucleotide positions across conformers). These graphs are then stacked along a new axis to form a unified data structure for message passing. In practice, this means the node feature matrices from each conformation (size $L \times d$ for $L$ nodes and feature dimension $d$) are stacked into a tensor of shape $L \times M \times d$, and similarly for edge features. All graphs share the same node indexing (1 through L for the nucleotides), and an edge exists in the unified graph if it exists in any conformation’s graph (the union of edges). This clever construction allows a single GNN to process all conformations simultaneously: essentially, the model sees a tensor of node features where the second dimension corresponds to different conformers.

The multi-state GNN encoder then performs message passing on this unified multi-graph. Intuitively, one can think of it as M parallel GNNs (one per conformation) that share weights and exchange information only by virtue of ultimately being merged in the output. The architecture is carefully designed to be equivariant and invariant to various symmetries: it is invariant to the order in which conformations are presented (shuffling the conformers shouldn’t change the encoding) and to rotations/translations of the 3D coordinates (a physical requirement), as well as the usual GNN invariance to node permutations. To achieve this, gRNAde’s message-passing layers treat the conformation axis as a set dimension. For example, the GNN uses SE(3)-equivariant operations (specifically a Graph Vector Perceptron or GVP GNN layer) that handle scalar and vector features, ensuring that rotating a backbone appropriately rotates those internal vector features while leaving scalars unchanged. During a message-passing update, a node’s features for each conformer are updated independently based on that conformer’s edges – the stacking ensures no accidental mixing of information between different conformations during the update step. This way, each conformation’s local structural context contributes to the node’s representation without “bleeding over” into other conformations, which guarantees the encoder respects the symmetry of swapping conformers.

After several GNN layers (gRNAde uses 3 encoder layers in experiments), the model obtains M feature vectors for each nucleotide (one from each conformation’s branch). Now, to produce a single consolidated representation per nucleotide, the encoder applies a conformation order-invariant pooling: essentially a symmetric function (like elementwise average or sum) across the M conformer feature vectors. gRNAde uses a simple average pooling, which introduces no extra parameters. The result is a single enriched feature vector for each nucleotide position, encoding information from all the provided 3D shapes. Importantly, because of the invariance, providing the conformations in a different order would yield the same pooled result. This pooled set of node embeddings (length *L) now encapsulates the structural constraints of potentially multiple backbone states.


Autoregressive Sequence Decoder

Given the pooled node representations from the encoder, gRNAde employs an autoregressive decoder to construct the RNA sequence one nucleotide at a time. The decoder is also implemented as GVP-based GNN layers operating on the nucleotide graph, but now the task is to output a probability distribution over the four bases {A, U, G, C} for each position in sequence order. The decoding proceeds iteratively from the 5’ end (position 1) to the 3’ end (position L): at each step i, the decoder is conditioned on the previously decided nucleotides $x_1,\dots,x_{i-1}$ (for i = 1, no prior context) and the encoder’s structural embeddings, and predicts the probability of each nucleotide at position i. In practice, this is done by masking unknown future positions and iteratively unmasking them in sequence order. The decoder uses 3 GNN layers as well (mirroring the encoder) and at each position incorporates the already-decoded neighbors as features (often via one-hot encoding of known bases in the node feature). This approach is analogous to how autoregressive protein design models (like Ingraham et al. 2019 or ProteinMPNN) operate on graphs.

One advantage of this decoder design is flexibility: while the default is 5’→3’ sequential decoding, the authors note that unordered or masked decoding schemes could also be used (similar to non-sequential strategies in protein design). For example, one could design all nucleotides in parallel or in random order, or even fill in specific parts (if some nucleotides are pre-defined). In their experiments, however, sequential decoding was effective and straightforward.


Experimental Results and Key Contributions

gRNAde was evaluated on both standard benchmarks and new tasks, showing strong performance improvements over prior methods. The contributions and findings of the paper can be summarized as follows:

In addition to the above, gRNAde introduced a new RNA design dataset derived from the RNAsolo database, comprising 3,764 unique RNA sequences each with 3D structure(s) (over 11,000 total structures). The dataset was split in multiple challenging ways (by conformational diversity, by number of conformers, by sequence dissimilarity) to rigorously evaluate generalization. All these contributions together demonstrate gRNAde as a comprehensive platform for RNA design, pushing the frontier of what is achievable with geometric deep learning on RNA.


Limitations and Future Directions

While gRNAde represents a significant advance, there are some limitations and assumptions in its current form. First, gRNAde treats the RNA backbone structure(s) as given inputs – it does fixed-backbone design only. This means it cannot suggest new 3D shapes on its own; it needs an initial model or experimentally determined structure to design against. In practice, this is often fine (one might have a target shape in mind from experiment or imagination), but future work could integrate structure prediction and design, allowing suggestions of novel shapes alongside sequences. The authors themselves hint that the 3-bead graph representation could be used for RNA backbone generation in the future, which would move toward de novo RNA 3D design.

Another limitation is data scarcity and bias: RNA 3D structures in the PDB are relatively few and often biased toward certain motifs (e.g., many tRNAs or riboswitches). gRNAde’s training data, though larger than past works, is still thousands of sequences – orders of magnitude smaller than protein structure datasets. The model might struggle with very large RNAs (there were only a handful of >1000-nt examples) or with RNA topologies that were absent in training. Moreover, gRNAde’s performance drops if asked to design for a structure far outside its training distribution (though the paper showed decent generalization via sequence identity splitting). As more RNA structures become available (e.g., via experimental methods or computational predictions), gRNAde could be retrained or scaled up, which would likely improve its coverage of RNA structural space.

The multi-state design ability is novel, but comes with assumptions: it assumes the set of input conformations is representative of the important states of the RNA. If the RNA has other unseen states or intermediate folding forms, the designed sequence might still fail. There’s also an implicit assumption that the sequence should stabilize all provided conformations, which might trade off stability in each individual state. In practice, designing for multiple states can be at odds with optimizing one state fully – gRNAde addresses this with an ML approach, but it might sometimes output a compromise sequence that is not optimal for any single state. Future work might explore weighting different states or incorporating kinetic/folding pathway information, so that certain conformations are prioritized or the transition barriers are considered.

From a model perspective, gRNAde currently uses relatively shallow GNN layers (3 encoder + 3 decoder) and an average pooling. It’s possible that deeper networks or more sophisticated pooling (learned pooling or attention across conformers) could capture interactions between conformations better. Additionally, the decoder could potentially be improved by non-sequential strategies or incorporating global sequence constraints (like enforcing a certain base composition or motif presence). The authors note that techniques from protein design like masked token decoding could be applied – this could allow partial sequence specification or inpainting (filling in missing pieces of an RNA given the rest), which would increase gRNAde’s utility in practical design scenarios.

Lastly, while gRNAde showed promising lab validation, a 50% success rate means there is still room to grow. The failures might be due to subtle effects not captured by the model (ion interactions, non-canonical base pairs, tertiary contacts beyond 10Å that were not encoded, etc.). Integrating some physics-informed features or energy terms (perhaps via an energy-based reranker or incorporating known thermodynamic parameters as features) might further improve the success for difficult targets. Nonetheless, gRNAde provides a powerful foundation on which such future improvements can be built.


Conclusion and Impact

gRNAde exemplifies how geometric deep learning can push forward the field of RNA design, analogous to how protein design has benefited from GNN and transformer models. By leveraging 3D information and embracing RNA’s intrinsic flexibility, it achieves state-of-the-art results and introduces capabilities (like multi-state design and rapid scoring of variants) that were previously out of reach for RNA engineers. This work sits at the intersection of structural biology and machine learning, and it underscores a broader trend: ML models are now sophisticated enough to handle the complexities of RNA structure (including pseudoknots and alternative folds) which used to be exclusively the domain of physics-based models.

From a practical standpoint, gRNAde could be a game-changer for researchers designing RNA-based therapeutics and nanotechnology. For example, one could envision designing a riboswitch that must adopt two distinct conformations – gRNAde provides a tool to tackle that directly. It also offers a new paradigm for multi-objective design (designing sequences that meet multiple structural criteria). The methods introduced, such as the multi-graph GNN encoder, could be generalized to other problems where inputs are an ensemble of related structures (not just RNA – perhaps protein loops in different conformations, or small molecules in multiple poses). In summary, gRNAde’s comprehensive approach to RNA inverse folding demonstrates the power of modern ML in solving long-standing biological design problems, and it opens the door for more dynamic, ensemble-aware design techniques in the future.