Equivariant Pretrained Transformer for Unified Geometric Learning on Multi-Domain 3D Molecules (2402.12714v1)
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.
- Geom, energy-annotated molecular conformations for property prediction and molecular generation. Scientific Data, 9(1):185, 2022.
- Learning protein sequence embeddings using information from structure. arXiv preprint arXiv:1902.08661, 2019.
- The protein data bank. Nucleic acids research, 28(1):235–242, 2000.
- The role of ai in drug discovery: challenges, opportunities, and strategies. Pharmaceuticals, 16(6):891, 2023.
- Language models are few-shot learners. Advances in neural information processing systems, 33:1877–1901, 2020.
- Prottrans: Toward understanding the language of life through self-supervised learning. IEEE transactions on pattern analysis and machine intelligence, 44(10):7112–7127, 2021a.
- Prottrans: Towards cracking the language of lifes code through self-supervised deep learning and high performance computing. IEEE Transactions on Pattern Analysis and Machine Intelligence, pp. 1–1, 2021b. doi: 10.1109/TPAMI.2021.3095381.
- Artificial intelligence for synthetic biology. Communications of the ACM, 65(5):88–97, 2022.
- Protein-ligand binding representation learning from fine-grained interactions. arXiv preprint arXiv:2311.16160, 2023a.
- Fractional denoising for 3d molecular pre-training. In International Conference on Machine Learning, pp. 9938–9961. PMLR, 2023b.
- Deciphering interaction fingerprints from protein molecular surfaces using geometric deep learning. Nature Methods, 17(2):184–192, 2020.
- Self-supervised pocket pretraining via protein fragment-surroundings alignment. arXiv preprint arXiv:2310.07229, 2023.
- Fast and uncertainty-aware directional message passing for non-equilibrium molecules. arXiv preprint arXiv:2011.14115, 2020.
- Simple gnn regularisation for 3d molecular property prediction and beyond. In International Conference on Learning Representations, 2021.
- Intrinsic-extrinsic convolution and pooling for learning on 3d protein structures. In International Conference on Learning Representations, 2020.
- Equivariant diffusion for molecule generation in 3d. In International conference on machine learning, pp. 8867–8887. PMLR, 2022.
- Ogb-lsc: A large-scale challenge for machine learning on graphs. arXiv preprint arXiv:2103.09430, 2021.
- Energy-motivated equivariant pretraining for 3d molecular graphs. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 37, pp. 8096–8104, 2023.
- Unsupervised protein-ligand binding energy prediction via neural euler’s rotation equation. arXiv preprint arXiv:2301.10814, 2023.
- Equivariant graph neural networks for 3d macromolecular structure. arXiv preprint arXiv:2106.03843, 2021.
- Deepaffinity: interpretable deep learning of compound–protein affinity through unified recurrent and convolutional neural networks. Bioinformatics, 35(18):3329–3338, 2019.
- Bert: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of NAACL-HLT, pp. 4171–4186, 2019.
- Semi-supervised classification with graph convolutional networks. arXiv preprint arXiv:1609.02907, 2016.
- End-to-end full-atom antibody design. arXiv preprint arXiv:2302.00203, 2023a.
- Generalist equivariant transformer towards 3d molecular interaction learning. arXiv preprint arXiv:2306.01474, 2023b.
- Denoising diffusion probabilistic models on SO(3) for rotational alignment. In ICLR 2022 Workshop on Geometrical and Topological Representation Learning, 2022. URL https://openreview.net/forum?id=BY88eBbkpe5.
- xformers: A modular and hackable transformer modelling library. https://github.com/facebookresearch/xformers, 2022.
- Equiformer: Equivariant graph attention transformer for 3d atomistic graphs. arXiv preprint arXiv:2206.11990, 2022.
- Pre-training molecular graph representation with 3d geometry. In International Conference on Learning Representations, 2021.
- Molecular geometry pretraining with se (3)-invariant denoising distance matching. arXiv preprint arXiv:2206.13602, 2022.
- A group symmetric stochastic differential equation model for molecule multi-modal pretraining. In International Conference on Machine Learning, pp. 21497–21526. PMLR, 2023.
- One transformer can understand both 2d & 3d molecular data. arXiv preprint arXiv:2210.01765, 2022.
- Deepdta: deep drug–target binding affinity prediction. Bioinformatics, 34(17):i821–i829, 2018.
- Pytorch: An imperative style, high-performance deep learning library. Advances in neural information processing systems, 32, 2019.
- Frame averaging for invariant and equivariant network design. In International Conference on Learning Representations, 2021.
- Accelerating materials discovery using artificial intelligence, high performance computing and robotics. npj Computational Materials, 8(1):84, 2022.
- Improving language understanding by generative pre-training. 2018.
- Language models are unsupervised multitask learners. OpenAI blog, 1(8):9, 2019.
- Quantum chemistry structures and properties of 134 kilo molecules. Scientific data, 1(1):1–7, 2014.
- Evaluating protein transfer learning with tape. Advances in neural information processing systems, 32, 2019.
- A generalist agent. Transactions on Machine Learning Research, 2022.
- E (n) equivariant graph neural networks. In International conference on machine learning, pp. 9323–9332. PMLR, 2021.
- Equivariant message passing for the prediction of tensorial properties and molecular spectra. In International Conference on Machine Learning, pp. 9377–9388. PMLR, 2021.
- Schnet–a deep learning architecture for molecules and materials. The Journal of Chemical Physics, 148(24), 2018.
- Multi-scale representation learning on proteins. Advances in Neural Information Processing Systems, 34:25244–25255, 2021.
- Generative modeling by estimating gradients of the data distribution. Advances in neural information processing systems, 32, 2019.
- Soper, D. E. Classical field theory. Courier Dover Publications, 2008.
- 3d infomax improves gnns for molecular property prediction. In International Conference on Machine Learning, pp. 20479–20502. PMLR, 2022.
- Torchmd-net: equivariant transformers for neural network based molecular potentials. arXiv preprint arXiv:2202.02541, 2022.
- Tensor field networks: Rotation-and translation-equivariant neural networks for 3d point clouds. arXiv preprint arXiv:1802.08219, 2018.
- Atom3d: Tasks on molecules in three dimensions. arXiv preprint arXiv:2012.04035, 2020.
- Alphafold protein structure database: massively expanding the structural coverage of protein-sequence space with high-accuracy models. Nucleic acids research, 50(D1):D439–D444, 2022.
- Attention is all you need. Advances in neural information processing systems, 30, 2017.
- Learning hierarchical protein representations via complete 3d graph networks. In International Conference on Learning Representations (ICLR), 2023.
- The pdbbind database: methodologies and updates. Journal of medicinal chemistry, 48(12):4111–4119, 2005.
- Lm-gvp: an extensible sequence and structure informed deep learning framework for protein property prediction. Scientific reports, 12(1):6832, 2022.
- De novo design of protein structure and function with rfdiffusion. Nature, 620(7976):1089–1100, 2023.
- Discovering the representation bottleneck of graph neural networks from multi-order interactions. arXiv preprint arXiv:2205.07266, 2022.
- Geodiff: A geometric diffusion model for molecular conformation generation. arXiv preprint arXiv:2203.02923, 2022.
- Unified molecular modeling via modality blending. arXiv preprint arXiv:2307.06235, 2023.
- Pre-training via denoising for molecular property prediction. arXiv preprint arXiv:2206.00133, 2022.
- Protein representation learning by geometric structure pretraining. arXiv preprint arXiv:2203.06125, 2022.
- Pre-training protein encoder via siamese sequence-structure diffusion trajectory prediction. In Annual Conference on Neural Information Processing Systems, 2023.
- Uni-mol: a universal 3d molecular representation learning framework. 2023.