Equivariant Pretrained Transformer for Unified Geometric Learning on Multi-Domain 3D Molecules

Rui Jiao,Xiangzhe Kong,Ziyang Yu,Wenbing Huang,Yang Liu

2024-02-20

Abstract:Pretraining on a large number of unlabeled 3D molecules has showcased superiority in various scientific applications. However, prior efforts typically focus on pretraining models on a specific domain, either proteins or small molecules, missing the opportunity to leverage the cross-domain knowledge. To mitigate this gap, we introduce Equivariant Pretrained Transformer (EPT), a novel pretraining framework designed to harmonize the geometric learning of small molecules and proteins. To be specific, EPT unifies the geometric modeling of multi-domain molecules via the block-enhanced representation that can attend a broader context of each atom. Upon transformer framework, EPT is further enhanced with E(3) equivariance to facilitate the accurate representation of 3D structures. Another key innovation of EPT is its block-level pretraining task, which allows for joint pretraining on datasets comprising both small molecules and proteins. Experimental evaluations on a diverse group of benchmarks, including ligand binding affinity prediction, molecular property prediction, and protein property prediction, show that EPT significantly outperforms previous SOTA methods for affinity prediction, and achieves the best or comparable performance with existing domain-specific pretraining models for other tasks.

Machine Learning,Chemical Physics

What problem does this paper attempt to address?

The problem addressed in this paper is how to perform unified geometric learning on multi-domain 3D molecules to overcome the limitations of existing pre-trained models that focus only on specific domains.

The paper introduces a new framework called Equivariant Pretrained Transformer (EPT), aimed at integrating geometric learning of small molecules and proteins. Current pre-trained models typically target either proteins or small molecules, while EPT enhances the representation methodology to capture a broader context of each atom, enabling unified modeling across domains. It utilizes E(3) equivariance to accurately represent 3D structures and employs block-level pre-training tasks to enable the model to recognize translations and rotations perturbations for each block, enhancing its modeling capabilities for complex hierarchical molecular geometry.

Experimental results demonstrate that EPT performs remarkably well in various benchmark tests such as ligand binding affinity prediction, molecular property prediction, and protein property prediction, outperforming previous state-of-the-art (SOTA) methods and achieving performance comparable to or better than domain-specific pre-trained models in other tasks. This indicates that EPT effectively utilizes cross-domain knowledge to improve the model's performance and generalization ability.