# One Transformer Can Understand Both 2D & 3D Molecular Data

arXiv: https://arxiv.org/abs/2210.01765 · HTML: https://mlsanity.com/paper/2210.01765/

## Key facts

- Cited by 77 papers in the map (data through 14 Oct 2026).
- Builds on 68 of its 102 references that are in the map.
- First posted on arXiv 4 Oct 2022.

## Builds on (most cited first, 30 of 68)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Learning Transferable Visual Models From Natural Language Supervision](https://mlsanity.com/paper/2103.00020/) (2021)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [Swin Transformer: Hierarchical Vision Transformer using Shifted Windows](https://mlsanity.com/paper/2103.14030/) (2021)
- [A Simple Framework for Contrastive Learning of Visual Representations](https://mlsanity.com/paper/2002.05709/) (2020)
- [Semi-Supervised Classification with Graph Convolutional Networks](https://mlsanity.com/paper/1609.02907/) (2016)
- [Graph Attention Networks](https://mlsanity.com/paper/1710.10903/) (2017)
- [Momentum Contrast for Unsupervised Visual Representation Learning](https://mlsanity.com/paper/1911.05722/) (2019)
- [Inductive Representation Learning on Large Graphs](https://mlsanity.com/paper/1706.02216/) (2017)
- [Layer Normalization](https://mlsanity.com/paper/1607.06450/) (2016)
- [How Powerful are Graph Neural Networks?](https://mlsanity.com/paper/1810.00826/) (2018)
- [Unsupervised Cross-lingual Representation Learning at Scale](https://mlsanity.com/paper/1911.02116/) (2019)
- [Neural Message Passing for Quantum Chemistry](https://mlsanity.com/paper/1704.01212/) (2017)
- [Playing Atari with Deep Reinforcement Learning](https://mlsanity.com/paper/1312.5602/) (2013)
- [A Comprehensive Survey on Graph Neural Networks](https://mlsanity.com/paper/1901.00596/) (2019)
- [Open Graph Benchmark: Datasets for Machine Learning on Graphs](https://mlsanity.com/paper/2005.00687/) (2020)
- [mT5: A massively multilingual pre-trained text-to-text transformer](https://mlsanity.com/paper/2010.11934/) (2020)
- [Graph Contrastive Learning with Augmentations](https://mlsanity.com/paper/2010.13902/) (2020)
- [Cross-lingual Language Model Pretraining](https://mlsanity.com/paper/1901.07291/) (2019)
- [Multilingual Denoising Pre-training for Neural Machine Translation](https://mlsanity.com/paper/2001.08210/) (2020)
- [Deep Networks with Stochastic Depth](https://mlsanity.com/paper/1603.09382/) (2016)
- [Strategies for Pre-training Graph Neural Networks](https://mlsanity.com/paper/1905.12265/) (2019)
- [E(n) Equivariant Graph Neural Networks](https://mlsanity.com/paper/2102.09844/) (2021)
- [Learning to Simulate Complex Physics with Graph Networks](https://mlsanity.com/paper/2002.09405/) (2020)

## Built on by (most cited first, 30 of 77)

- [Rethinking the Expressive Power of GNNs via Graph Biconnectivity](https://mlsanity.com/paper/2301.09505/) (2023)
- [Attending to Graph Transformers](https://mlsanity.com/paper/2302.04181/) (2023)
- [Developing ChemDFM as a large language foundation model for chemistry](https://mlsanity.com/paper/2401.14818/) (2024)
- [MUDiff: Unified Diffusion for Complete Molecule Generation](https://mlsanity.com/paper/2304.14621/) (2023)
- [Towards Foundational Models for Molecular Learning on Large-Scale Multi-Task Datasets](https://mlsanity.com/paper/2310.04292/) (2023)
- [Enabling Efficient Equivariant Operations in the Fourier Basis via Gaunt Tensor Products](https://mlsanity.com/paper/2401.10216/) (2024)
- [Fractional Denoising for 3D Molecular Pre-training](https://mlsanity.com/paper/2307.10683/) (2023)
- [GPS++: An Optimised Hybrid MPNN/Transformer for Molecular Property Prediction](https://mlsanity.com/paper/2212.02229/) (2022)
- [A Survey of Geometric Graph Neural Networks: Data Structures, Models and Applications](https://mlsanity.com/paper/2403.00485/) (2024)
- [Learning Joint 2D & 3D Diffusion Models for Complete Molecule Generation](https://mlsanity.com/paper/2305.12347/) (2023)
- [Learning Probabilistic Symmetrization for Architecture Agnostic Equivariance](https://mlsanity.com/paper/2306.02866/) (2023)
- [Automated 3D Pre-Training for Molecular Property Prediction](https://mlsanity.com/paper/2306.07812/) (2023)
- [Highly Accurate Quantum Chemical Property Prediction with Uni-Mol+](https://mlsanity.com/paper/2303.16982/) (2023)
- [Multimodal Molecular Pretraining via Modality Blending](https://mlsanity.com/paper/2307.06235/) (2023)
- [PowerGraph: A power grid benchmark dataset for graph neural networks](https://mlsanity.com/paper/2402.02827/) (2024)
- [Sliced Denoising: A Physics-Informed Molecular Pre-Training Method](https://mlsanity.com/paper/2311.02124/) (2023)
- [UniCorn: A Unified Contrastive Learning Approach for Multi-view Molecular Representation Learning](https://mlsanity.com/paper/2405.10343/) (2024)
- [GraphGPT: Generative Pre-trained Graph Eulerian Transformer](https://mlsanity.com/paper/2401.00529/) (2023)
- [Leveraging Biomolecule and Natural Language through Multi-Modal Learning: A Survey](https://mlsanity.com/paper/2403.01528/) (2024)
- [PRESTO: Progressive Pretraining Enhances Synthetic Chemistry Outcomes](https://mlsanity.com/paper/2406.13193/) (2024)
- [Pre-training with Fractional Denoising to Enhance Molecular Property Prediction](https://mlsanity.com/paper/2407.11086/) (2024)
- [A Survey of Graph Transformers: Architectures, Theories and Applications](https://mlsanity.com/paper/2502.16533/) (2025)
- [Triplet Interaction Improves Graph Transformers: Accurate Molecular Graph Learning with Triplet Graph Transformers](https://mlsanity.com/paper/2402.04538/) (2024)
- [An end-to-end attention-based approach for learning on graphs](https://mlsanity.com/paper/2402.10793/) (2024)
- [Learning a Fourier Transform for Linear Relative Positional Encodings in Transformers](https://mlsanity.com/paper/2302.01925/) (2023)
- [Let the Code LLM Edit Itself When You Edit the Code](https://mlsanity.com/paper/2407.03157/) (2024)
- [Geometry Informed Tokenization of Molecules for Language Model Generation](https://mlsanity.com/paper/2408.10120/) (2024)
- [Instructor-inspired Machine Learning for Robust Molecular Property Prediction](https://mlsanity.com/paper/2304.03906/) (2023)
- [MolSpectra: Pre-training 3D Molecular Representation with Multi-modal Energy Spectra](https://mlsanity.com/paper/2502.16284/) (2025)
- [Siamese Foundation Models for Crystal Structure Prediction](https://mlsanity.com/paper/2503.10471/) (2025)

## Abstract (arXiv)

Unlike vision and language data which usually has a unique format, molecules can naturally be characterized using different chemical formulations. One can view a molecule as a 2D graph or define it as a collection of atoms located in a 3D space. For molecular representation learning, most previous works designed neural networks only for a particular data format, making the learned models likely to fail for other data formats. We believe a general-purpose neural network model for chemistry should be able to handle molecular tasks across data modalities. To achieve this goal, in this work, we develop a novel Transformer-based Molecular model called Transformer-M, which can take molecular data of 2D or 3D formats as input and generate meaningful semantic representations. Using the standard Transformer as the backbone architecture, Transformer-M develops two separated channels to encode 2D and 3D structural information and incorporate them with the atom features in the network modules. When the input data is in a particular format, the corresponding channel will be activated, and the other will be disabled. By training on 2D and 3D molecular data with properly designed supervised signals, Transformer-M automatically learns to leverage knowledge from different data modalities and correctly capture the representations. We conducted extensive experiments for Transformer-M. All empirical results show that Transformer-M can simultaneously achieve strong performance on 2D and 3D tasks, suggesting its broad applicability. The code and models will be made publicly available at https://github.com/lsj2408/Transformer-M.