Papers
Topics
Authors
Recent
Gemini 2.5 Flash
Gemini 2.5 Flash
124 tokens/sec
GPT-4o
8 tokens/sec
Gemini 2.5 Pro Pro
47 tokens/sec
o3 Pro
5 tokens/sec
GPT-4.1 Pro
38 tokens/sec
DeepSeek R1 via Azure Pro
28 tokens/sec
2000 character limit reached

An Equivariant Pretrained Transformer for Unified 3D Molecular Representation Learning (2402.12714v2)

Published 20 Feb 2024 in cs.LG and physics.chem-ph

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 in a specific domain, either proteins or small molecules, missing the opportunity to leverage cross-domain knowledge. To mitigate this gap, we introduce Equivariant Pretrained Transformer (EPT), an all-atom foundation model that can be pretrained from multiple domain 3D molecules. Built upon an E(3)-equivariant transformer, EPT is able to not only process atom-level information but also incorporate block-level features (e.g. residuals in proteins). Additionally, we employ a block-level denoising task, rather than the conventional atom-level denoising, as the pretraining objective. To pretrain EPT, we construct a large-scale dataset of 5.89M entries, comprising small molecules, proteins, protein-protein complexes, and protein-molecule complexes. Experimental evaluations on downstream tasks including ligand binding affinity prediction, protein property prediction, and molecular property prediction, show that EPT significantly outperforms previous state-of-the-art methods in the first task and achieves competitively superior performance for the remaining two tasks. Furthermore, we demonstrate the potential of EPT in identifying small molecule drug candidates targeting 3CL protease, a critical target in the replication of SARS-CoV-2. Among 1,978 FDA-approved drugs, EPT ranks 7 out of 8 known anti-COVID-19 drugs in the top 200, indicating the high recall of EPT. By using Molecular Dynamics (MD) simulations, EPT further discoveries 7 novel compounds whose binding affinities are higher than that of the top-ranked known anti-COVID-19 drug, showcasing its powerful capabilities in drug discovery.

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

Summary

  • The paper introduces a novel transformer architecture that leverages E(3) equivariance and block-level denoising to unify geometric learning across multi-domain 3D molecules.
  • It enhances atom-level features by integrating residue-level contexts and distance matrices, leading to improved predictions in ligand binding and molecular properties.
  • Experimental results demonstrate that EPT outperforms state-of-the-art methods, setting a new benchmark for unified cross-domain molecular representation learning.

Equivariant Pretrained Transformer for Multi-Domain 3D Molecular Learning

In the recent surge of advancements in deep learning applications for scientific research, the accurate representation and understanding of molecular structures have emerged as a pivotal area, particularly due to its relevance in drug discovery, materials science, and biochemistry. Traditional approaches often constrain their focus to specific domains, either proteins or small molecules, neglecting the enriching potential of cross-domain knowledge sharing. Addressing this research gap, the newly introduced Equivariant Pretrained Transformer (EPT) proposes a novel framework designed for unified geometric learning across different domains of 3D molecular structures.

Introduction to EPT

EPT stands out by its innovative employment of a block-enhanced representation technique that enriches each atom’s context by aligning atom-level and residue-level features. Built upon a transformer architecture, EPT integrates E(3) equivariance, enabling it to capture the 3D structure more accurately than traditional methods. A notable breakthrough in this research is the block-level denoising pretraining task, which allows for a more nuanced understanding of the complex hierarchical geometry inherent in 3D molecules.

Experimental Evaluation

EPT's performance was rigorously tested against a variety of benchmarks in the fields of ligand binding affinity prediction, molecular property prediction, and protein property prediction. The results affirm EPT’s capacity to significantly outperform state-of-the-art methods in affinity prediction, while achieving comparable or superior performance in other tasks, evidencing its robust applicability across different molecular domains.

Technical Insights and Analysis

The paper offers profound technical insights into the components contributing to EPT's performance. Key findings include:

  • Enhancing atom features with block-level information slightly elevates performance, suggesting the effectiveness of incorporating broader atom context.
  • The integration of distance matrices and edge features into the attention mechanism underpins the model's ability to encapsulate diverse interatomic relations effectively.
  • The block-level denoising strategy adopted by EPT positively impacts its ability to capture broader molecular dynamics, as illustrated by its superior performance across various molecular benchmarks.

Implications and Future Directions

The development of EPT initiates a promising direction towards the creation of generalizable and accurate models for molecular representation learning. It paves the way for further exploration into the benefits of cross-domain knowledge transfer and the development of unified models capable of understanding the universal principles governing molecular structures. Future research could focus on enhancing the scalability of EPT to larger molecular systems and extending its applicability to encompass more diverse scientific domains.

In conclusion, the Equivariant Pretrained Transformer (EPT) sets a new benchmark in the field of 3D molecular learning, demonstrating an unprecedented level of performance across multiple domains. Its innovative approach to pretraining and geometric representation holds promising potential for revolutionizing how molecular systems are modeled and understood in computational research.

File Document Download Save Streamline Icon: https://streamlinehq.com PDF File Document Download Save Streamline Icon: https://streamlinehq.com Markdown