[Paper Review] Generalist Equivariant Transformer Towards 3D Molecular Interaction Learning

Xiangzhe Kong, Wenbing Huang+1 more|arXiv (Cornell University)|Jun 2, 2023

Computational Drug Discovery Methods4 citations

TL;DR

This paper proposes the Generalist Equivariant Transformer (GET), a unified E(3)-equivariant model that represents diverse 3D molecular complexes as geometric graphs of sets—capturing both block-level (e.g., residues, nucleobases) and atom-level structures. GET uses bilevel attention and equivariant modules to preserve fine-grained information across variable-sized blocks, achieving state-of-the-art performance in cross-domain molecular interaction prediction, including zero-shot RNA/DNA-ligand affinity prediction.

ABSTRACT

Many processes in biology and drug discovery involve various 3D interactions between molecules, such as protein and protein, protein and small molecule, etc. Given that different molecules are usually represented in different granularity, existing methods usually encode each type of molecules independently with different models, leaving it defective to learn the various underlying interaction physics. In this paper, we first propose to universally represent an arbitrary 3D complex as a geometric graph of sets, shedding light on encoding all types of molecules with one model. We then propose a Generalist Equivariant Transformer (GET) to effectively capture both domain-specific hierarchies and domain-agnostic interaction physics. To be specific, GET consists of a bilevel attention module, a feed-forward module and a layer normalization module, where each module is E(3) equivariant and specialized for handling sets of variable sizes. Notably, in contrast to conventional pooling-based hierarchical models, our GET is able to retain fine-grained information of all levels. Extensive experiments on the interactions between proteins, small molecules and RNA/DNAs verify the effectiveness and generalization capability of our proposed method across different domains.

Motivation & Objective

• To address the lack of a unified, generalist representation for diverse 3D molecular interactions across proteins, small molecules, and RNA/DNAs.
• To overcome limitations of domain-specific models and hierarchical pooling-based approaches that lose fine-grained structural information.
• To develop a single, equivariant model capable of capturing both domain-specific hierarchies and universal interaction physics.
• To enable zero-shot transfer learning across molecular domains, particularly for RNA/DNA-ligand binding affinity prediction.
• To preserve both block-level specificity and atom-level shareability in a unified geometric graph representation.

Proposed method

• Proposes a geometric graph of sets representation: a top-level graph of molecular blocks (e.g., residues, nucleobases) with each block containing a set of atoms.
• Designs a bilevel attention mechanism that jointly attends over sparse block-level and dense atom-level representations to update atomic features.
• Employs E(3)-equivariant feed-forward and layer normalization modules that are invariant to permutations within each block and equivariant to 3D rotations, translations, and reflections.
• Ensures permutation invariance within blocks and size-agnostic operation across blocks of varying sizes.
• Uses a unified transformer architecture with E(3)-equivariant modules to process the bilevel geometric graph without hierarchical pooling.
• Trains the model end-to-end on mixed datasets to enable transfer learning and zero-shot generalization across domains.

Experimental results

Research questions

• RQ1 Can a single, unified model effectively represent and learn from diverse 3D molecular complexes across proteins, small molecules, and RNA/DNAs?
• RQ2 Does a bilevel geometric graph of sets representation preserve both block-level specificity and atom-level shareability better than existing single-level or hierarchical approaches?
• RQ3 Can an E(3)-equivariant transformer retain fine-grained structural information without relying on pooling-based aggregation?
• RQ4 To what extent can a unified model generalize to unseen domains, such as RNA/DNA-ligand binding, via zero-shot transfer learning?
• RQ5 How does the proposed GET compare to domain-specific, single-level, and hierarchical models in predicting molecular interaction affinities?

Key findings

• GET achieved the highest Pearson correlation (0.697 ± 0.003) and Spearman correlation (0.555 ± 0.008) on protein-protein affinity prediction, outperforming all baselines.
• On ligand binding affinity prediction, GET achieved the lowest RMSE (1.327 ± 0.005) and highest Pearson (0.620 ± 0.004) and Spearman (0.611 ± 0.003) correlations, with strong zero-shot performance on RNA/DNA-ligand systems.
• The model with mixed pretraining (GET-mix) achieved state-of-the-art results across all benchmarks, demonstrating strong generalization and transferability.
• GET outperformed hierarchical models that use pooling, which tend to lose fine-grained information, by retaining both atom-level and block-level representations.
• The ablation study confirmed that the bilevel attention and E(3)-equivariant design are critical for performance, especially in cross-domain transfer.
• GET demonstrated strong zero-shot generalization on RNA/DNA-ligand binding affinity prediction, suggesting it captures universal interaction physics beyond domain-specific patterns.