Title: UniIF: Unified Molecule Inverse Folding
Venue: NeurIPS 2024
Accepted: yes
Keywords: molecule learning

protein design

RNA design
Abstract: Molecule inverse folding has been a long-standing challenge in chemistry and biology, with the potential to revolutionize drug discovery and material science. Despite specified models have been proposed for different small- or macro-molecules, few have attempted to unify the learning process, resulting in redundant efforts. Complementary to recent advancements in molecular structure prediction, such as RoseTTAFold All-Atom and AlphaFold3, we propose the unified model UniIF for the inverse folding of all molecules.  We do such unification in two levels: 1) Data-Level: We propose a unified block graph data form for all molecules, including the local frame building and geometric feature initialization. 2) Model-Level: We introduce a geometric block attention network, comprising a geometric interaction, interactive attention and virtual long-term dependency modules, to capture the 3D interactions of all molecules. Through comprehensive evaluations across various tasks such as protein design, RNA design, and material design, we demonstrate that our proposed method surpasses state-of-the-art methods on all tasks. UniIF offers a versatile and effective solution for general molecule inverse folding.

UniIF: Unified Molecule Inverse Folding

Abstract

Molecule inverse folding has been a long-standing challenge in chemistry and biology, with the potential to revolutionize drug discovery and material science. Despite specified models have been proposed for different small- or macro-molecules, few have attempted to unify the learning process, resulting in redundant efforts. Complementary to recent advancements in molecular structure prediction, such as RoseTTAFold All-Atom and AlphaFold3, we propose the unified model UniIF for the inverse folding of all molecules. We do such unification in two levels: 1) DataLevel: We propose a unified block graph data form for all molecules, including the local frame building and geometric feature initialization. 2) Model-Level: We introduce a geometric block attention network, comprising a geometric interaction, interactive attention and virtual long-term dependency modules, to capture the 3D interactions of all molecules. Through comprehensive evaluations across various tasks such as protein design, RNA design, and material design, we demonstrate that our proposed method surpasses state-of-the-art methods on all tasks. UniIF offers a versatile and effective solution for general molecule inverse folding.

1 Introduction

Molecule inverse folding plays a pivotal role in drug and material design, enabling scientists to synthesize novel molecules with the desired structure. Previously, many studies focus on either macromolecules [19, 33, 21, 9, 17, 4, 10, 17, 8, 34] or small molecules [6, 26, 16, 28, 14, 31] separately, leaving the challenge of inverse folding general molecules. For example, the advanced small molecule model [6, 31] take atoms as basic units; the macromolecule models [8, 10] consider predefined microstructures (such as amino acids and nucleotides) as the basic units. Additionally, even for the same molecule, different models employ varying strategies to extract geometric features. Complementary to the great success of RoseTTAFold All-Atom [25] and AlphaFold3 [1] in molecular structure prediction, we propose a unified model, UniIF, for the inverse folding of all molecules.

By comparing small- and macro-molecules, we identify three challenges toward the unified model: (1) Unit Discrepancy : The macromolecules takes predefined microstructures (amino acids and nucleotides) as the basic units, while small molecules takes atoms as basic units. (2) Geometric Featurizer : Different studies employ various strategies for extracting geometric features from structures, such as distance, angles and tensor product; there are lack of unified featurization strategy. (3) System Size : The small-molecules allow the full attention transformer to learn long-term dependencies, but the quadratic computing cost limits the mechanism scalling up to macro-molecular systems. Alternatively, previous research use sparse

Unified Inverse Folding Model

Protein

RNA

Molecule

Structure

Tokens

GNN, which suffers from the limited local receptive field that causes over-smoothing and oversquashing [29]. In addition, developing a unified model working well for all molecules is challenging.

The unit discrepancy makes it challenge to adapt methods across small- and macro-molecules, which explains the divergence between these two research lines. As a solution, we propose a frame-based block to unify the representation of amino acids, nucleotides, and atoms: a group of atoms with varying size is treated as a block with fixed size. Each block includes decoupled equivariant basis and invariant features, generalizing the representation of AlphaFold2 and other small-molecule methods.

The geometric featurizer is necessary to capture the geometric interactions between blocks. We initialize the equivariant block basis using predefined rules or a learnable GNN layer, and then constructing invariant block features based on these basis. The key operation is to use local coordinates and dot product to capture the geometric interactions between virtual atoms. We reuse the featurizer in each model layer to interactively learn updated geometric features, where the concept of directed virtual atom is introduced to enhance the pairwise interactions. We show that the unified featurizer works well across protein design, RNA design, and material design.

We use sparse GNN to address the system size issue, while maintaining the ability to capture longterm dependencies. The transformer-style protein models like AlphaFold and RosettaFold require a substantial amount of GPU memory. Sparse GNNs, on the other hand, are criticized for their tendency to over-smooth and over-squash due to their limited local receptive field. To be efficient while preserving the ability to capture long-term dependencies, we introduce global virtual blocks. Each virtual block is connected to all real blocks, serving as an information exchange agent.

We conducted comprehensive experiments across various tasks, including protein design, RNA design, and material design, to demonstrate the effectiveness of UniIF. The results show that UniIF achieves state-of-the-art performance on all the tasks, which is non-trivial and may benefit the machine learning, drug discovery, and material science communities.

2 Related work

Unification. Unified molecular learning has attracted increasing attention in recent years. RoseTTAFold All-Atom (RFAA) [25] and AlphaFold3 [1] are two representative models that have achieved remarkable success in protein structure prediction. RoseTTAFold All-Atom uses an atombond graph for small molecules and a frame graph for macromolecules. AlphaFold3 uses a bi-level representations, i.e., atom representation and token representation, for all molecules. The token concept is requivalent to the block concept in this paper, which means a group of atoms, such as a amino acid or a nucleotide. GET [24] and EPT [20] are two recent models that use a block representation for both small and macromolecules and introduce a new equivariant transformer backbone. Unlike RFAA [25], which specifies a atom-bond graph for small molecules, our model employs a unified block graph for all molecule types and do not require the atom-bond graph. Our model also differs from AlphaFold3 [1], GET [24] and EPT [20] in the that we introduce the vector basis for each block.

Protein Inverse Folding. Recent research use k -NN graph to represent the 3D structure and employ graph neural networks for protein inverse folding. GraphTrans [19] uses the graph attention encoder and autoregressive decoder for protein design. GVP [21] proposes geometric vector perceptrons to learn from both scalar and vector features. GCA [33] introduces global graph attention for learning contextual features. In addition, ProteinSolver [32] is developed for scenarios where partial sequences are known while not reporting results on standard benchmarks. Recently, AlphaDesign [9], ProteinMPNN [4], ESMIF [17], LMDesign [40], KWDesign [8], VFN [27] achieves dramatic improvements. A benchmark [11] is proposed to comprehensively evaluate protein design models.

RNA Inverse Folding. RNA inverse folding is a challenging task due to the complex secondary structure and tertiary structure. Traditional methods [3] include colony optimization and constraint programming, in addition to adaptive walk, simulated annealing and Boltzmann sampling. Recent deep learning method RDesign [34] has achieved promising results and build a benchmark for AI researchers to follow-up. RiboDiffusion [18] use diffusion decoder to generate RNA sequences conditioned on the backbone structure embeddings.

Material Design. Deciding which chemical compositions are likely to form compounds is a critical task in material design [28, 26, 14]. An important application is to substitute lattice-site elements or ionic species within existing compounds that exhibit similar chemical behaviors. Wang et al. [36] successfully employed the elemental substitution method to discover 18,479 stable compounds out of a pool of 189,981 potential candidates. Recently, Jensen et al. [6] introduced a open dataset for material design and established a benchmark for evaluating deep learning models in this domain.

3 Method

3.1 Overall Framework

As shown in Fig. 2, we propose the unified model for general molecule inverse folding. The key insights include: (1) transforming all molecules into block graphs, where each block represents an amino acid, nucleotide, or atom; (2) proposing a geometric featurizer to initialize geometric node and edge features; and (3) introducing a new GNN layer with long-term dependencies to learn expressive block representations. Our unified model achieves competitive results across diverse tasks, including protein design, RNA design, and material design.

Micromolecules

Small molecules

Learned

Basis

Prior

Basis

Protein Design

RNA Design

Block Field Attention

Block Field Attention

Block Field Attention

Block Field Attention

Prediction Head

Block Graph Neural Network

Material Design

Tasks

Edge feature

Unified Pipeline for All Molecules

All Molecules

3.2 Block Graph

We introduce the block graph to represent all types of molecules, where the key insight is to transform irregular set of atoms (varying size) as regular block representation (fixed size).

carbon atom

direction

Rotation vector

Block

P

!

O5

!

C5

!

C4

!

O4

!

C1

!

C2

!

!

Base

O3

!

O2

!

O3

!"#

P

!$#

phosphorus atom

oxygen atom

carbon atom

chain direction

Block

RNA

nitrogen atom

alpha  carbon atom

carbon atom

oxygen atom

Base

C1

!

C2

!

O

!

𝑁

!

C2

!$#

𝑁

!"#

Block

chain direction

Protein

Small Molecules

alpha  carbon atom

Atom-based Block Representation. A block B = { ( x i , z i ) } |B| i =1 contains a set of atoms { x i , z i } |B| i =1 , where x i ∈ R 3 and z i ∈ R d represent the equivariant coordinate and invariant features, such as the atom type. The common block types include amino acids, nucleotides, and atoms, which are represented as B fold , B rna , and B smol , respectively. Formally, we write B fold = { ( x i , z i ) } i ∈V fold and B rna = { ( x i , z i ) } i ∈V rna , where V fold = { N , C 1 , C 2 , O } and V rna = { P , C 5 , C 4 , C 3 , C 2 , C 1 , O 5 , O 4 , O 3 , O 2 } are the sets of atoms for amino acids and nucleotides, respectively. For small molecules, each atom a represents a block B smol = { ( x a , z a ) } , where 1 ≤ a ≤ 118 . As the block size |B| varies for different types of blocks, the atom-based blocks representation could not directly be applied for unified modeling.

Frame-based Block Representation. We introduce frame-based block representation to unify the modeling of all molecules. A block B = ( F, f ) contains the equivariant frame F and invariant feature vector f ∈ R d . The local frame F ( R, t ) contains the axis matrix R = [ e 1 , e 2 , e 3 , · · · , e u ] and translation vector t . We set t as the coordinate of the representative atom, i.e., C 1 of macromolecules and the atom itself of small molecules. Following AlphaFold2, we consider the special case that R ∈ R 3 , 3 is orthogonal. However, additional experiments show that the model can also work well with non-orthogonal axis matrix. For macromolecules, the axis matrix R is predefined based on amino acids and nucleotides, while for small molecules,

Frame-based Block (fixed size)

Features

Used for Unified Modeling

Vector

Basis

we learn the axis matrix R as it does not have prior common structure patterns. The frame-based block representation decouples geometric information: (1) the local frame basis describe the equivariant pose; (2) the invariant feature vector could embed the atom type and invariant local structure patterns for different tasks. More importantly, the dimension of the frame-based block representation is fixed, which is beneficial for the unified modeling; we build block features in Sec. 3.3.

Frame-based Block Graph. Given a molecule M = {B s } n s =1 containing n blocks, we build the block graph G ( {B s } n s =1 , E ) using kNN algorithm. In the block graph, the s -th node is represented as B s = ( F s , f s ) , and the edge between ( s, t ) is represented as B st = ( F st , f st ) . The relative frame is defined as F st = F -1 s ◦ F t . Inspired by [22, 10, 27], we modify PiFold featurizer to initialize the geometric node feature f s and edge feature f s,t ; refer to Sec. 3.3.

Relation to Other Methods. The frame-based block is a generalized data form of AlphaFold2 and other methods. If R is required to be a rotation matrix, the frame-based block is equivalent to AlphaFold2's local frame; otherwise, it is equivalent to represent the atom as invariant feature h and equivalent vector x , similar to GVP [21] and DimeNet [12].

3.3 Block Graph Featurizer

Learning Local Frame. For small molecules, there is no predefined local frame, and we need to learn the local frame for each atom. Given the the molecule M = { ( x s , z s ) } |M| s =1 , we use a 1-layer of GNN to initialize the atom representation { z s } |M| s =1 ← BlockGAT ( { x s , z s } |M| s =1 ) , where the inital local frames are T ( I , 0 ) . The rotation vector r s of the s -th atom is constructed by message passing:

N s is the s -th atom's neighbor system. Let the direction ( r x , r y , r z ) = r s || r s || and magnitude θ = || r s || represent the rotation axis and angle, we compute the quanternion q s and rotation matrix R s :

Finally, the local frame of the s -th atom is T s ( R s , t s ) , where t s is the atom coordinate. Experiments show that learning rotation vectors consistently outperforms learning Schmidt-orthogonalized axises.

Node Geometric Feature. The invariant block feature captures the atom type and the local structure:

The inverse frame operation T -1 s project the equivalent global coordinates to the invariant local coordinate, i.e., x local = F -1 s ◦ x global = R T s ( x global -t s ) . We use MLP to embed atom type and local coordinates. All atom features in the same block are pooled to get the block feature f s .

Edge Geometric Feature. We initialize pairwise features following the principle that

Edge features capture the directed 3D interactions.

Instead of using mutually constructed distance and angle features, we concatenate the local coordinates of two blocks to fully describe their 3D positions. Given B s and B t with global coordinate matrices as X s ∈ R | v s | , 3 and X t ∈ R | v t | , 3 , the invariant edge features following s ← t direction is

where T -1 s = ( R T s , -R T s t s ) projects equivariant global coordinates to invariant local coordinates.

3.4 Block Graph Attention Module

Geo-Interaction

Gated Edge Attention

FFN

Edge Update

Block Graph Construction

Gated Edge Attention

Geo-Interaction

(b) Equi- virtual atoms

(c) Inv- relative coords

(a) Inv- edge feat

Real+Virtual Blocks

Frame-based Block Graph

Atom-based Block Graph

Global Virtual  Frames

Global Virtual Featurs

Invariant  Features

Real Features

Real Frames

Equivariant  Frames

(a)

(b)

(c)

(d)

Frame-based SE-(3) Module Design. Given the geometric transformation y = W x , we decompose W = R t Σ R T s using SVD and explain y = R t Σ R T s x as:

Projecting x as the local coordinate x local using the frame F s ( R s , 0 ) , i.e., x local = R T s x .

Updating local coordinates via gated attention, i.e., x local ← Σ x local .

Translating x local as the global coordinate using frame T t ( R t , 0 ) , i.e., y = R t x local .

If we parameterize Σ y as f θ ( y ) , and considers the effects of translation, i.e., T s ( R s , t s ) , T t ( R t , t t ) , the general principle of designing SE-(3) networks could be:

The well-known AlphaFold actually follows such a design, where they parameterize f θ as the IPA module. In this work, we replace f θ with an enhanced graph neural network:

where f ( l ) s and f ( l ) st represent the input node and edge features of the l -th layer. In Fig. 5, we show the design of the Block Graph Attention Module, consisting of three components: (1) geometric interaction extractor, (2) virtual block for long-term dependencies, and (3) the edge attention mechanism. We show the detailed design of the Block Graph Attention Module in the following sections.

Long-term Dependency via Virtual Blocks. The GNN is criticized by local receptive field, yielding the problem of over-smoothing and over-squashing [2, 5, 29, 13]. Transformers overcome these problems using direct paths between distant nodes, while suffering from the O ( n 2 ) computing cost. We introduce n ′ virtual blocks {B i } n + n ′ i = n as information agents for a graph. Each virtual block directly connects to all the real blocks, resulting in (2 · n · n ′ ) additional directed edges. As n ′ ≪ n , we claim that the computing cost is close to original GNN. All the virtual blocks, i.e., T n +1 ( R ′ , t ′ ) , T n +2 ( R ′ , t ′ ) , · · · , T n + n ′ ( R ′ , t ′ ) , share the same rotation R ′ and translation t ′ :

The invariant features of virtual blocks are different, as we hope they learn diverse interactions. We encode the index to initialize block features f ′ i = Embedding ( i ) for i ∈ { 1 , 2 , · · · , n ′ } .

Geometric Interaction Extractor. We enhance edge features with geometric interactions using the local coordinates of virtual inter-atoms and dot products of virtual intra-atoms. Previous works, such as PiFold [10], introduced virtual atoms in the featurizer to capture informative side-chain geometry beyond protein backbones, resulting in performance gains. VFN [27] extended this idea by allowing GNN layers to update the virtual atoms. However, these efforts are limited to learning virtual intra-atoms conditioned on node features. Instead, we propose virtual inter-atoms conditioned

Local Position Interaction

Dot Product Interaction

on edge features, allowing the same node to exhibit different virtual states specified by edges. Additionally, inspired by small molecule modeling, we use the dot product of virtual intra-atoms to capture angle information. We show the geometric interactions in Fig. 6, and formulate it as:

where T st = T -1 s ◦ T t = ( R -1 s R t , R -1 s ( t t -t s )) , q st = vec ( R st ) ∈ R 9 is the flatten rotation matrix of T st , and r st = || t s -t t || indicates the pairwise distance. All q st , r st an a st are invariant features. We highlight the difference to previous researches in color.

Gated Edge Attention. We modify PiFold's GNN to capture the geometric interactions when updating node features. For molecular design tasks, we find that aggregating edge features only leads to consistent performance gains. We understand this phenomenon as the model can pay more attention on learning 3D interactions under such a model design. In addition, we use a gated mechanism to control how the edge features are injected to node features. The gated edge attention module is:

where ⊙ is element-wise product operation, and σ ( · ) is the sigmoid function. We highlight the difference between PiGNN and the proposed module in color.

FFN & Edge Updating. Analogous to the transformer model, the FFN is a MLP. The edge updating layer remains the same as PiFold:

The proposed module could be equivalently implemented as a transformer module using matrix multplication. However, we find that padding proteins to the maximum length would grealy increase the computing cost in both GPU occupancy and runtime; We suggest using GNN without padding.

Regularization. We find that the proposed model fit training data better than PiFold and more likely to suffer from overfitting. To address this issue, we randomly drop out the nodes/edges with a probability of p to prevent overfitting. We find that controlling the dropout rate could result in models with different fitting abilities. The best performance is achieved when p = 0 . 05 .

4 Experiments

We show the effectiveness of UniIF via multiple inverse folding tasks and ablation studies. We briefly introduce molecular design tasks as follows:

Protein Design (T1): Designing protein sequences folding into the target structure.

RNA Design (T2): Designing RNA sequences folding into the target structure.

Material Design (T3): Discoverying stable composition from a known material structure.

4.1 Protein Design (T1)

Task Description Protein design aims to design protein sequences that fold into target structures. Given a protein backbone structure X = { X i ∈ R m, 3 : 1 ≤ i ≤ n } , where m is the maximum number of points belonging to the i -th residue, n is the number of residues and the natural proteins are composed by 20 types of amino acids, the goal is to learn a function F θ :

The parameters θ are learned by minimizing the cross-entropy loss, i.e., L ( F θ ( X ) , S ) = -∑ n i =1 log s i p (ˆ s i |X , θ ) . The task is challenging due to the combinatorial search space of amino acids and the complex relationship between sequence and structure.

Settings We evaluate of UniIF on the CATH4.3 dataset [30] following prior works [11, 8]. The dataset is split by the CATH topology classification code, yielding 16,631 training, 1,516 validation, and 1,864 testing samples. To assess generalization, we adopt a time-split strategy, considering the use of pretrained ESM2 models by some baselines, which risk data leakage. The time-split evaluation assigns data before a specific date to the training set and data after that date to the test set. For structural time-split evaluation, we use the CASP15 dataset [11], containing novel crystal structures not seen during training. For sequence time-split evaluation, we use the NovelPro dataset [8], which includes 76 protein sequences released within 30 days before November 23, 2023, with structures predicted by AlphaFold2. UniIF consists of 10 layers of BlockGAT with a hidden dimension of 128. It is trained using the Adam optimizer with a learning rate of 1e-3 and a batch size of 8 for 50 epochs.

Metrics & Baselines We report the median recovery rate of the top-1 predicted sequences, representing the percentage of correctly predicted residues. The ESM2-free baselines include StructGNN [19], GraphTrans [19], GCA [33], GVP [21], AlphaDesign [9], ProteinMPNN [4], and PiFold [10]. The ESM2-based baselines include LMDesign [17] and KWDesign [8]. While we prefer open-source baselines, we also re-implement VFN [27] for a comprehensive comparison.

Conclusion We provide results under different settings (with and without ESM2) and across diverse datasets (CATH4.3, CASP, NovelPro). Using a pure inverse folding model without ESM2, UniIF achieves the best performance on all datasets, demonstrating its effectiveness. Notably, UniIF outperforms the strong baseline PiFold with fewer learnable parameters. In time-split evaluations, UniIF surpasses all baselines, including ESM2-based methods, by a significant margin. On NovelPro, which features novel sequences, UniIF outperforms LMDesign and KWDesign that use ESM2 for sequence refinement. This indicates UniIF's superior generalizability, crucial for real-world

Table 1: Protein Design results. The best and suboptimal results are labeled with bold and underlined. "VFN" means that we replace the geometric interaction operation with VFN's operation [27]. "-GDP" means that we remove the geometric dot product features. "-EAttn" means that we replace the gated edge attention with PiGNN's attention module [10]. "-VFrame" means that we remove the global virtual frames.

applications. Ablation studies show that the proposed geometric featurizer, gated edge attention, and global virtual frame enhance performance. On CATH4.3, the overall improvement is slight due to strong baselines, but time-split evaluation highlights UniIF's superiority in generalization.

4.2 RNA Design (T2)

Task Description Similar to protein design, RNA design aims to design RNA sequences that fold into target structures. Specially, previous work [34] use the RNA secondary structure as additional input to guide the design process, since the tertiary structure is limited. In this work, we only use the tertiary structures as input for the reason of unification, which is more challenging than the baselines.

Datasets & Baselines We conduct experiments RNA on the dataset collected by RDesign [34], consisting of 2218 RNA tertiary structures, which are divided into training (1774 structures), testing (223 structures), and validation (221 structures) sets based on their structural similarity. Following RDesign's benchmark, baseline methods include SeqRNN, SeqLSTM, StructMLP, StructGNN, and StructGNN, GraphTrans [19], PiFold [10] and RDesign [34]. Given the small number of data samples, we report the median recovery and its standard deviation for theree independent runs.

Conclusion As shown in Table 2, UniIF achieves the best performance in all cases. The improvement is significant, as previous strong baselines like PiFold only excelled in protein design. To our knowledge, UniIF is the first model to achieve state-of-the-art performance in both protein and RNA design tasks, demonstrating its versatility and effectiveness. Compared to RDesign, which uses additional secondary structure features, UniIF relies solely on tertiary structure input and still performs better. UniIF successfully unifies the protein and RNA design processes, paving the way for a unified inverse folding model for protein-RNA complexes in future developments.

4.3 Material Design (T3)

Task Description Discovering stable atom compositions from known material structures is crucial for new material discovery [28, 26, 14]. This task is challenging due to the large composition space and the lack of large-scale data. Thanks to recent benchmark efforts [6], we can evaluate the performance of UniIF on this novel task.

Datasets & Baselines We evaluated UniIF on the CHILI-3K dataset [6], which consists of nanomaterial graphs derived from mono-metal oxides. The dataset includes 53 metallic elements and one non-metallic element (oxygen), comprising 3,180 graphs, 6,959,085 nodes, and 49,624,440 edges. Following the official benchmark, the dataset is randomly split into training (80%), validation (10%), and testing (10%) sets. Baselines include GCN [23], PMLP [39], GraphSAGE [15], GAT [35], GraphUNet [7], GIN [38], and EdgeCNN [37]. Experiments are repeated three times with different seeds, using early stopping with a patience of 50 epochs, and trained up to 1000 epochs.

Conclusion In Table 3, UniIF outperforms all baselines by a large margin. Ablation studies demonstrate the crucial role of the learned

local frame in enhancing interaction feature extraction. In addition, how to learn the local frame is also important. In the "- quat" ablation, we try to learn the x, y, and z axes with Householder orthogonalization directly, but found it less effective, with the recovery rate dropping from 75.3% to 65.2%. This highlights the value of the proposed local frame learning mechanism.

4.4 Case Study

In Fig. 7, we show the designed protein and RNA sequences.In addition, we use AlphaFold3 [1] to re-fold the designed sequences into structures. The ground truth (gray), PiFold (green), and UniIF (pink) structures are alinged and compared. We observe that UniIF improves both the recovery and RMSD of the designed protein and RNA, demonstrating its effectiveness in inverse folding tasks.

Rec:0.625

→

0.76  RMSD:8.57

→

8.32

PDB

RDesign

UniIF

PDB

PiFold

UniIF

PiFold

UniIF

Rec:0.565

→

0.65  RMSD: 0.341

→

0.257

Rec:0.5

→

0.73 RMSD:9.90

→

9.14

Rec:0.67

→

0.76 RMSD: 3.70

→

3.60

Rec:0.375

→

0.75  RMSD: 8.07

→

5.76

Rec:0.468

→

0.56  RMSD: 0.440

→

0.309

7CGK_B

4X4N_D

2ET8_A

UGUAAAUA

GCCGCGGCAGGUUCGAUCCUGCCGCGAUCG

GCGUCACACCGGUGAAGUCGCGCGUCACACCGGUGAAGUCGC

UGAAUUAU

ACGCGGAGAGGGUAAAACCCCUCUCCGAAU

GGGGCCCGGCCGAGAAGUCGCGCGCCUCUCGGGCGAAGCCCC

UGUAUUUA

GCGAUACGCCCGUGAUCUCGCGCGUGUCACCGGCGAAUUCGC

SEPPNLVQSDDPRAVALNYDNDYKKSNRIADNLPGLPPELQKHANCAHAVEDAEGATDEWIGCKLFPGKLVNRNGWCANFTVKKA

2CU7

PDB

GSSGSSGYSVKWTIEEKELFEQGLAKFGRRWTKISKLIGSRTVLQVKSYARQYFKNKVKCGLDKETPNQKTG

PiFold

GSSGSSGSEVRLTAEELELLEKGLAKFGSDLVKISKLIGKLTVDEVKSYAEQLEKEKAKSGKDSETPSSSTG

UnilF

PDB

MQDFRPGVYRHYKGDHYLALGLARADETDEVVVVYTRLYARAGLPMSTRLLRIWNETVDTGAGPQPRFAYVGHVTPE

5 Conclusion

We propose the first unified model, dubbled UniIF, for general molecule inverse folding. The key points include unifying the data representation, the featurizer and the model architecture without a drop in performance. Extensive experiments show that UniIF surpasses baseline methods on all tasks. Ablation studies reveal that the geometric interaction extractor, gated edge attention, and virtual long-term dependency modules contribute to performance gains. We believe that the proposed model can benefit multiple domains, such as machine learning, drug design, and material design.

References

Josh Abramson, Jonas Adler, Jack Dunger, Richard Evans, Tim Green, Alexander Pritzel, Olaf Ronneberger, Lindsay Willmore, Andrew J Ballard, Joshua Bambrick, et al. Accurate structure prediction of biomolecular interactions with alphafold 3. Nature, pages 1-3, 2024.

Uri Alon and Eran Yahav. On the bottleneck of graph neural networks and its practical implications. In International Conference on Learning Representations, 2020.

Alexander Churkin, Matan Drory Retwitzer, Vladimir Reinharz, Yann Ponty, Jérôme Waldispühl, and Danny Barash. Design of rnas: comparing programs for inverse rna folding. Briefings in bioinformatics, 19(2):350-358, 2018.

Justas Dauparas, Ivan Anishchenko, Nathaniel Bennett, Hua Bai, Robert J Ragotte, Lukas F Milles, Basile IM Wicky, Alexis Courbet, Rob J de Haas, Neville Bethel, et al. Robust deep learning based protein sequence design using proteinmpnn. bioRxiv, 2022.

Francesco Di Giovanni, Lorenzo Giusti, Federico Barbero, Giulia Luise, Pietro Lio, and Michael M Bronstein. On over-squashing in message passing neural networks: The impact of width, depth, and topology. In International Conference on Machine Learning, pages 78657885. PMLR, 2023.

Ulrik Friis-Jensen, Frederik L Johansen, Andy S Anker, Erik B Dam, Kirsten MØ Jensen, and Raghavendra Selvan. Chili: Chemically-informed large-scale inorganic nanomaterials dataset for advancing graph machine learning. arXiv preprint arXiv:2402.13221, 2024.

Hongyang Gao and Shuiwang Ji. Graph u-nets. In international conference on machine learning, pages 2083-2092. PMLR, 2019.

Zhangyang Gao, Cheng Tan, Xingran Chen, Yijie Zhang, Jun Xia, Siyuan Li, and Stan Z Li. Kw-design: Pushing the limit of protein deign via knowledge refinement. In The Twelfth International Conference on Learning Representations, 2023.

Zhangyang Gao, Cheng Tan, Stan Li, et al. Alphadesign: A graph protein design method and benchmark on alphafolddb. arXiv preprint arXiv:2202.01079, 2022.

Zhangyang Gao, Cheng Tan, and Stan Z. Li. Pifold: Toward effective and efficient protein inverse folding. In International Conference on Learning Representations, 2023.

Zhangyang Gao, Cheng Tan, Yijie Zhang, Xingran Chen, Lirong Wu, and Stan Z Li. Proteininvbench: Benchmarking protein inverse folding on diverse tasks, models, and metrics. Advances in Neural Information Processing Systems, 36, 2024.

Johannes Gasteiger, Janek Groß, and Stephan Günnemann. Directional message passing for molecular graphs. arXiv preprint arXiv:2003.03123, 2020.

Jhony H Giraldo, Konstantinos Skianis, Thierry Bouwmans, and Fragkiskos D Malliaros. On the trade-off between over-smoothing and over-squashing in deep graph neural networks. In Proceedings of the 32nd ACM International Conference on Information and Knowledge Management, pages 566-576, 2023.

Sean D Griesemer, Yi Xia, and Chris Wolverton. Accelerating the prediction of stable materials with machine learning. Nature Computational Science, 3(11):934-945, 2023.

Will Hamilton, Zhitao Ying, and Jure Leskovec. Inductive representation learning on large graphs. Advances in neural information processing systems, 30, 2017.

Geoffroy Hautier, Christopher C Fischer, Anubhav Jain, Tim Mueller, and Gerbrand Ceder. Finding nature's missing ternary oxide compounds using machine learning and density functional theory. Chemistry of Materials, 22(12):3762-3767, 2010.

Chloe Hsu, Robert Verkuil, Jason Liu, Zeming Lin, Brian Hie, Tom Sercu, Adam Lerer, and Alexander Rives. Learning inverse folding from millions of predicted structures. bioRxiv, 2022.

Han Huang, Ziqian Lin, Dongchen He, Liang Hong, and Yu Li. Ribodiffusion: Tertiary structure-based rna inverse folding with generative diffusion models. bioRxiv, pages 2024-04, 2024.

John Ingraham, Vikas K Garg, Regina Barzilay, and Tommi Jaakkola. Generative models for graph-based protein design. 2019.

Rui Jiao, Xiangzhe Kong, Ziyang Yu, Wenbing Huang, and Yang Liu. Equivariant pretrained transformer for unified geometric learning on multi-domain 3d molecules. arXiv preprint arXiv:2402.12714, 2024.

Bowen Jing, Stephan Eismann, Patricia Suriana, Raphael JL Townshend, and Ron Dror. Learning from protein structure with geometric vector perceptrons. arXiv preprint arXiv:2009.01411, 2020.

John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ronneberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Žídek, Anna Potapenko, et al. Highly accurate protein structure prediction with alphafold. Nature, 596(7873):583-589, 2021.

Thomas N Kipf and Max Welling. Semi-supervised classification with graph convolutional networks. arXiv preprint arXiv:1609.02907, 2016.

Xiangzhe Kong, Wenbing Huang, and Yang Liu. Generalist equivariant transformer towards 3d molecular interaction learning. arXiv preprint arXiv:2306.01474, 2023.

Rohith Krishna, Jue Wang, Woody Ahern, Pascal Sturmfels, Preetham Venkatesh, Indrek Kalvet, Gyu Rie Lee, Felix S Morey-Burrows, Ivan Anishchenko, Ian R Humphreys, et al. Generalized biomolecular modeling and design with rosettafold all-atom. Science, 384(6693):eadl2528, 2024.

Yue Liu, Tianlu Zhao, Wangwei Ju, and Siqi Shi. Materials discovery and design using machine learning. Journal of Materiomics, 3(3):159-177, 2017.

Weian Mao, Muzhi Zhu, Zheng Sun, Shuaike Shen, Lin Yuanbo Wu, Hao Chen, and Chunhua Shen. De novo protein design using geometric vector field networks. arXiv preprint arXiv:2310.11802, 2023.

Bryce Meredig, Ankit Agrawal, Scott Kirklin, James E Saal, Jeff W Doak, Alan Thompson, Kunpeng Zhang, Alok Choudhary, and Christopher Wolverton. Combinatorial screening for new materials in unconstrained composition space with machine learning. Physical Review B, 89(9):094104, 2014.

Khang Nguyen, Nong Minh Hieu, Vinh Duc Nguyen, Nhat Ho, Stanley Osher, and Tan Minh Nguyen. Revisiting over-smoothing and over-squashing using ollivier-ricci curvature. In International Conference on Machine Learning, pages 25956-25979. PMLR, 2023.

Christine A Orengo, Alex D Michie, Susan Jones, David T Jones, Mark B Swindells, and Janet M Thornton. Cath-a hierarchic classification of protein domain structures. Structure, 5(8):1093-1109, 1997.

Rafael Sarmiento-Perez, Tiago FT Cerqueira, Sabine Korbel, Silvana Botti, and Miguel AL Marques. Prediction of stable nitride perovskites. Chemistry of Materials, 27(17):5957-5963, 2015.

Alexey Strokach, David Becerra, Carles Corbi-Verge, Albert Perez-Riba, and Philip M Kim. Fast and flexible protein design using deep graph neural networks. Cell Systems, 11(4):402-411, 2020.

Cheng Tan, Zhangyang Gao, Jun Xia, and Stan Z Li. Generative de novo protein design with global context. arXiv preprint arXiv:2204.10673, 2022.

Cheng Tan, Yijie Zhang, Zhangyang Gao, Bozhen Hu, Siyuan Li, Zicheng Liu, and Stan Z Li. Hierarchical data-efficient representation learning for tertiary structure-based rna design. In The Twelfth International Conference on Learning Representations, 2023.

Petar Velickovic, Guillem Cucurull, Arantxa Casanova, Adriana Romero, Pietro Lio, Yoshua Bengio, et al. Graph attention networks. stat, 1050(20):10-48550, 2017.

Hai-Chen Wang, Silvana Botti, and Miguel AL Marques. Predicting stable crystalline compounds using chemical similarity. npj Computational Materials, 7(1):12, 2021.

Yue Wang, Yongbin Sun, Ziwei Liu, Sanjay E Sarma, Michael M Bronstein, and Justin M Solomon. Dynamic graph cnn for learning on point clouds. ACM Transactions on Graphics (tog), 38(5):1-12, 2019.

Keyulu Xu, Weihua Hu, Jure Leskovec, and Stefanie Jegelka. How powerful are graph neural networks? arXiv preprint arXiv:1810.00826, 2018.

Chenxiao Yang, Qitian Wu, Jiahua Wang, and Junchi Yan. Graph neural networks are inherently good generalizers: Insights by bridging gnns and mlps. arXiv preprint arXiv:2212.09034, 2022.

Zaixiang Zheng, Yifan Deng, Dongyu Xue, Yi Zhou, Fei Ye, and Quanquan Gu. Structureinformed language models are protein designers. bioRxiv, pages 2023-02, 2023.

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: Our experimental results support the claims.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [No]

Justification: The limitation is that we do not do wet-experiments. But we think this is out of scope for the AI paper.

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [NA]

Justification: Not applicable as the paper does not include theoretical results.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: We provided the model and training details in Section 3.1.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [No]

Justification: The code will be released upon acceptance.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: We have provided such information in the experimental section.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: We provided the standard deviation in the RNA design and material design tasks.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [No]

Justification: All experiments are conducted on an NVIDIA A100 with 80G memory. The longest training time is about 1 day.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification:

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [No]

Justification: This is a computational technical paper and does not have conducted realistic biological application.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification:

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification:

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA]

Justification: No new datasets are introduced in the paper.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification:

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [No]

Justification:

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes representing molecular systems as blocks of atoms and designs a neural network architecture that processes such representations. The method is applied for benchmark inverse folding tasks for protein/RNA. A materials-related task is also explored.\n\nSTRENGTHS:\n- The proposed method seems to perform well on various protein/RNA inverse folding benchmarks that focus on the recovery rate metric. (this reviewer is not experienced in protein/RNA design tasks, so it is hard for this reviewer to put the results in context.)\n\n- The idea of representing clusters of atoms as blocks is an interesting approach and is potentially useful for representing molecular systems across different scales. while such ideas have already been explored in several previous works as the authors noted, there are new elements such as vector basis for blocks.\n\nWEAKNESSES:\n- The protein inverse folding ablation seems to show the various introduced techniques do not influence model performance much. The insights of the models seem obscured from this perspective, especially considering block-based representations and architectures have already been explored in previous works (GET/EPT). The RNA design task does not include ablation, and this reviewer is not convinced by the materials-related task (see below). It is rather unclear what a reader can take away from this work in terms what techniques actually contribute to improved performance.\n\n- If \"unified\" is the key word of the method, perhaps the authors should include experiments where the model is trained on various data modalities and see if that improves performances on individual tasks.\n\n- The materials sections are rather strange in this paper. Inverse folding for proteins/RNAs can be understood as inferring the sequence from a given structure. The materials task is not introduced in great detail. What is being predicted? How is it \"inverse folding?\"\n\n- To follow up on that, the baselines that the proposed model is compared against are not what's considered state-of-the-art in materials representation. For materials, there are popular benchmarks such as Matbench (https://matbench.materialsproject.org/), JARVIS (https://pages.nist.gov/jarvis_leaderboard/), or Open catalyst (https://opencatalystproject.org/index.html), where results of SOTA models are readily available. This reviewer is not convinced by the materials-related results presented as significant details regarding the dataset/goals are missing, and the baselines are not really used for materials.\n\nQUESTIONS:\n- In the experiments, what exactly does with/without ESM mean? How is ESM used? Why can't the proposed method also become \"with ESM\"? Why is it valuable to consider \"without ESM\"? Are models \"with ESM\" more expensive or what?\n- Can the author clarify the materials task and why it is inverse folding? Can the authors compare to materials-focused SOTA methods in the space?\n\nLIMITATIONS:\nThe authors did not write about the limitations of the proposed method.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 1, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-08"}

comments: SUMMARY:
This paper propose a unified framework for inverse folding. Their model is applicable to small molecules, proteins and RNA.
More specifically, they propose a unified frame representation for amino acid, atom and nucleotide. They further propose a GNN-based model to learn the structure information. 
Results show that their model can outperform other baselines in protein design, RNA design and material design.


---

I changed my score to 5 after the discussion with other reviewers and AC.

STRENGTHS:
(1) The paper is clarified clearly. 
(2) The idea of unifying inverse folding for different types of molecules is interesting and reasonable. 
(3) Experimental results show their model work good in different design tasks.

WEAKNESSES:
Major points:

1. Some important experimental details are missing. For instance, while the paper proposes a unified inverse folding framework, it's unclear whether the model is trained on a combined dataset from different types of molecules. Specifically, the training sets for RNA design and material design are quite limited. If the model is trained with protein data, it would be beneficial to demonstrate whether incorporating data from other types improves performance.

2. Additionally, if I understand correctly, the paper does not explain the decoder part of the inverse folding model.

Minor points:

1. It would be helpful if the paper included a comparison of perplexity, as related works such as PiFold have used this metric.
   
2. The performance on protein design looks comparable with other models to me.

QUESTIONS:
(1) As illustrated in the previous part, do you train the model on three types of molecules?

(2) What decoder model is used?

LIMITATIONS:
The limitations are not discussed.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-09

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nPrevious approaches on molecule inverse folding (IF), which is crucial for drug and material discovery,  focus separately on either macro-molecules or small molecules, leading to the lack of a unified approach for different molecule types. To this end, the paper proposes UniIF to unify the IF from a data and model perspective. Specifically, by introducing a unified block for processing different molecules and a geometric block to capture their 3D interactions, UniIF shows improvements in various tasks across different molecules.\n\nSTRENGTHS:\n- The proposed UniIF is novel and effective for unifying IF across different molecules, including protein, RNA, and small molecules.\n- The UniIF can excel in protein, RNA, and material design, which is impressive. Extensive experiments have been conducted to provide justification for the effectiveness of the proposed method. \n- The paper is generally well-written, with clear illustrations and tables.\n\nWEAKNESSES:\n- The ablation experiments in Table 1 (-GDP), which remove the geometric dot product features, show miner drawbacks in proteins with longer sequences. This makes the geometric interaction extractor mainly capture the interaction of the virtual inter-atoms. The idea for capturing long-range dependency or interaction using virtual blocks has been studied in [1,2,3,4]; the authors could consider providing a discussion on how their proposed techniques differ from these referenced works.\n- In addition, the UniIF adopted both node and edge feature extractors as feature augmentation, which is model-agnostic. However, the comparison of all three tasks does not show the potential benefit of introducing such a featureizer. \n- The paper could provide further insight into the benefit of employing a unified framework for protein, RNA, and material design. For example, could the UniIF model different molecules in a unified latent space, allowing the researchers to investigate their underlying interactions?\n\n[1] Neural Message Passing for Quantum Chemistry. In ICML, 2017\n\n[2] An analysis of virtual nodes in graph neural networks for link prediction. In LoG, 2022\n\n[3] On the Connection Between MPNN and Graph Transformer. In ICML, 2023.\n\n[4] Neural Atoms: Propagating Long-range Interaction in Molecular Graphs through Efficient Communication Channel. In ICLR, 2024.\n\nQUESTIONS:\n1. Could you provide some discussions or experiments on the effectiveness of the proposed featureizer?\n2. Is it possible to train UniIF with mixed molecule datasets?\n3. For Equation 7, the SVD decomposition could be time-consuming when scaling the input size. Could the author explain the motivation for such a design and why a simple approach, like \"mean,\" does not work?\n\nLIMITATIONS:\nThe proposed UniIF offers a unified framework for learning different molecules, such as protein and RNA. However, it still models different molecules individually, hindering the proposed framework's potential impact.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-09"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper addresses the challenge of inverse folding, i.e. the design of novel molecules or macromolecules with specific desired 3D structure, with the goal of improving real-world drug and material design. The authors point out that this challenge has been addressed independently for different contexts, such as inverse folding for small molecules vs inverse folding for large macromolecules. They argue that this has resulted in redundant efforts. To address this issue, they propose a unified model for the inverse folding of all molecules. They demonstrate the efficacy of their novel unified approach across three benchmark molecular design tasks involving protein, RNA and material design.\n\nSTRENGTHS:\nThe authors make an interesting argument that inverse folding problems in different scientific domains should be addressed by a single approach. They present an elegant approach that unifies approaches across domains. Their results demonstrates improved performance as measured by sequence recovery at benchmark tasks involving protein, RNA and material design. The authors carry out ablation studies that provide some insight into the source of this improved performance.\n\nWEAKNESSES:\nWhile the unified approach is elegant, it is not clear (and the authors do not assert) that the unification is responsible for any performance improvements. However, it is easy for the reader to be left with the impression that the improved results stem from unification. The authors should make it clear that either (i) unification is satisfying but is not connected to performance improvements, or (ii) unification leads to performance improvements and explain clearly the data supporting this. \n\nThe recovery metric is confusing, since it is not clear that higher recovery actually corresponds to the ability to design novel molecules. Indeed, the recovery metric penalizes the model for novelty, which seems to contradict the stated goal ('enabling scientists to synthesize novel molecules with the desired structure'). This undermines the claim that the approach presented in the paper 'can benefit multiple domains, such as machine learning, drug design, and material design'. Specifically, rediscovering existing macromolecules does not benefit these domains. \n\nMoreover, even if the metric were aligned with the stated goals, it is not clear to this reviewer that existing inverse folding methods are not good enough to solve any real world application of inverse folding such as the stated goals described in the introduction to this paper of improvements in drug and material design. This leads me to wonder what problem this paper is trying to solve.\n\nWhile the paper focuses on maintaining the structure of the target molecules (or material), no attempt is made to ensure that the function is maintained. This does not bode well for applications such as drug design, that require that the designed molecule functions correctly. \n\nThe paper does not include any wetlab experiments, instead relying on performance at benchmark tasks to demonstrate improved performance. The authors argue that this is the only limitation of this paper, and that wetlab experiments are by definition out of scope for an AI paper. This raises the question of whether improvements on the benchmark tasks and associated metrics used in this paper corresponds to progress on the stated goals of the paper.\n\nQUESTIONS:\nIn the absence of wetlab experiments, please could the authors address this limitation by providing data that demonstrates that improved performance at the specific benchmark tasks in inverse folding cited in this paper results in improved performance at novel molecule design. This data can be extracted from previously published work, with appropriate citation. \n\nThere have been a series of recent works presenting improvements at inverse folding. These works, including this paper, make lofty claims about the impact that such improvements will have on drug and material design. Please could the authors point to any realized impact on drug and material design resulting from improvements at inverse folding as measured by performance on these benchmark tasks. \n\nIf such realized impact is difficult to find, please could the authors identify and address this limitation. \n\nPlease could the authors propose strategies to mitigate the risk that further improvements in performance on benchmark tasks, such as the improvements reported in this paper, will not result in realized impact in drug or material design. \n\nPlease could the authors propose strategies to mitigate the risk that the metric used in this work penalizes the model for the design of novel molecules, thereby undermining the stated goals. \n\nPlease could the authors identify a set of real world applications for inverse folding, describing the problem, how inverse folding will contribute to solving this problem, and why existing inverse folding methods are not good enough to solve each specific real world problem identified (NB inferior performance on benchmark tasks does not demonstrate that existing methods are not good enough).\n\nI note that the number of sequences for CATH 4.3 differs from the referenced papers - please could the authors explain this discrepancy. \n\nFor the time split protein tasks, it would be helpful to establish how similar each test sequence is to those used in training, both in terms of proximity in sequence space and in terms of 3D structure. Please then stratify performance as a function of each distance for each held out test set. It would be helpful to also compute these distances for the CATH topology-based split, to provide a watermark for comparison.\n\nLIMITATIONS:\nIn addition to the questions above, please could the authors describe limitations of their work that may prevent improvements at inverse design as measured by performance on benchmark tasks using a recovery metric from having impact on real-world challenges such as drug and material design.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-14"}