Source-linked AI summary
ATOM3D: Tasks On Molecules in Three Dimensions
Raphael J. L. Townshend, Martin Vögele, Patricia Suriana, Alexander Derry, Alexander Powers, Yianni Laloudakis, Sidhika Balachandar, Bowen Jing, Brandon Anderson, Stephan Eismann, Risi Kondor, Russ B. Altman, Ron O. Dror
TL;DR
Molecular learning has lacked systematic 3D benchmarks and unified data tools despite the importance of molecular geometry. ATOM3D addresses this with standardized datasets, multiple 3D model classes, and an open-source toolkit, finding broad gains over lower-dimensional methods while highlighting architecture-dependent performance and benchmark scope limits.
Problem
Systematic evaluation of 3D molecular learning methods and standardized tools for processing molecular structures were lacking, despite the importance of 3D geometry.
Method
ATOM3D curates benchmark datasets, implements graph, 3D convolutional, and equivariant methods, and provides tools for dataset processing, training, and evaluation.
Results
3D molecular learning yields gains over 1D and 2D methods on nearly all tasks formulated in lower dimensions, with the best architecture depending on the task.
Takeaways & Limitations
ATOM3D provides benchmarks and tools intended to support further development of machine learning methods for 3D molecular data.
Takeaways & Limitations
ATOM3D is biased toward biomolecules with solved structures, so benchmark performance may not generalize to intrinsically disordered or transmembrane proteins.
Abstract
from arXiv · showhide
Computational methods that operate on three-dimensional molecular structure have the potential to solve important questions in biology and chemistry. In particular, deep neural networks have gained significant attention, but their widespread adoption in the biomolecular domain has been limited by a lack of either systematic performance benchmarks or a unified toolkit for interacting with molecular data. To address this, we present ATOM3D, a collection of both novel and existing benchmark datasets spanning several key classes of biomolecules. We implement several classes of three-dimensional molecular learning methods for each of these tasks and show that they consistently improve performance relative to methods based on one- and two-dimensional representations. The specific choice of architecture proves to be critical for performance, with three-dimensional convolutional networks excelling at tasks involving complex geometries, graph networks performing well on systems requiring detailed positional information, and the more recently developed equivariant networks showing significant promise. Our results indicate that many molecular problems stand to gain from three-dimensional molecular learning, and that there is potential for improvement on many tasks which remain underexplored. To lower the barrier to entry and facilitate further developments in the field, we also provide a comprehensive suite of tools for dataset processing, model training, and evaluation in our open-source atom3d Python package. All datasets are available for download from https://www.atom3d.ai .
1 Introduction
Three-dimensional molecular shape is relevant to biology, chemistry, and drug discovery, but prior benchmarks largely used 1D or 2D representations. ATOM3D addresses this gap by standardizing datasets and benchmarking 3D molecular learning methods.
- 3D atomistic geometry captures molecular shape more directly than linear sequences or chemical bond graphs.The passage connects 3D structure to questions in drug discovery, molecular design, and fundamental biology.
- Prior benchmarks such as MoleculeNet and TAPE examined diverse molecular tasks using lower-dimensional representations.
- The atom is proposed as a machine learning datatype shared across proteins, DNA, RNA, and drugs, with common physical symmetries.
- ATOM3D creates curated benchmark datasets and processes them into a standardized format for atomistic machine learning.
2 Related Work
Prior work established graph, convolutional, and equivariant neural networks as approaches for learning from molecular geometry. These methods differ in how they encode spatial structure, symmetries, and physical transformations.
- Three-dimensional molecular learning has recently shown impact in protein structure prediction, equilibrium-state sampling, and RNA structure prediction.
- Graph neural networks encode pairwise connectivity and have been applied broadly to atomic systems and molecular prediction tasks.
- Three-dimensional convolutional networks capture complex molecular geometries but incur significant computational expense.
- Equivariant neural networks transform their outputs in well-defined ways when geometric transformations are applied to their inputs.Their neurons are designed to learn rules resembling physical interactions.
3 Datasets for 3D Molecular Learning
ATOM3D assembles standardized datasets spanning small molecules, proteins, protein interactions, RNA, and ligand binding. The tasks cover diverse molecular properties, structural predictions, and interaction problems, while some datasets are newly constructed.
- The benchmark suite spans structural biophysics and medicinal chemistry across multiple molecule types, with standardized formats and leakage-conscious splits.The datasets require no specialized libraries and are accompanied by best practices and dataset-specific details.
- Table 2 schematically summarizes task inputs using protein, small-molecule, and RNA labels, with interaction lines and residue markers.New datasets are indicated in bold.
- 3.1 Small Molecule Properties (SMP): QM9 provides 134,000 small organic molecules with ground-state geometries and calculated energetic, electronic, and thermodynamic properties.ATOM3D predicts molecular properties from each molecule’s ground-state structure.
- The datasets include protein interaction, residue-level amino-acid, mutation-stability, ligand-affinity, and protein-state tasks.These tasks address interaction sites, local structural environments, mutation effects, binding strength, and ligand-favored protein states.
- Protein and RNA structure-prediction datasets use CASP structural models and FARFAR2-generated RNA models to predict GDT_TS and RMSD, respectively.
4 Benchmarking Setup
ATOM3D benchmarks graph, voxel-based 3D convolutional, and SE(3)-equivariant networks for molecular learning. The implementations use task-appropriate representations and output mechanisms for whole systems or individual atoms and residues.
- The benchmark compares graph networks, three-dimensional convolutional networks, and equivariant networks as major classes of atomistic deep-learning methods.
- Graph neural networks: GNNs represent atoms as graph nodes, connect atoms within 4.5 Å, and use inverse-distance edge weights with ε = 10^-5.Five graph-convolution layers are followed by normalization, ReLU, fully connected layers, and dropout.
- Three-dimensional convolutional networks: 3DCNNs discretize molecular space into 1 Å voxels with one-hot atom-type channels, using 1.3 Å resolution for PSR and RSR.
- Equivariant networks: ENN models represent atoms by coordinates and atom types in SE(3)-equivariant networks, so rotational augmentation is unnecessary.The architectures use four layers of covariant neurons with the Clebsch–Gordan transform as nonlinearity.
5 Benchmarking Results
ATOM3D benchmarks show that 3D molecular methods generally outperform lower-dimensional approaches across small-molecule, biopolymer, joint, and structure-ranking tasks, while architecture choice affects which tasks benefit most.
- 3D methods improve performance across all targets in the small-molecule SMP dataset relative to 2D representations.
- 3D methods outperform 1D methods on all biopolymer-only datasets: PIP, RES, and MSP.
- 3D representations yield improved or comparable performance on the joint protein–small-molecule datasets LBA and LEP.
- Structure-ranking tasks require 3D information because identical sequences and bond graphs cannot distinguish alternative biopolymer shapes; 3D methods achieve competitive or better results.
- 3D methods hold state-of-the-art performance on the majority of ATOM3D tasks.
- Different tasks benefit from different architectures: 3DCNNs generally outperform GNNs on data-rich biopolymer tasks, whereas GNNs and ENNs become competitive or superior as datasets shrink.
- Different tasks benefit from different architectures: SMP is an exception: GNNs and ENNs improve performance there, consistent with small molecules having less complex geometry and targets depending on precise neighboring-atom positions.
- Different tasks benefit from different architectures: Equivariant networks show promise, achieving state-of-the-art performance on SMP, but the evaluated implementation had substantial scaling limitations on other tasks.
6 Conclusion
ATOM3D introduces benchmark datasets and computational tools for 3D molecular learning, showing performance gains over lower-dimensional methods while highlighting architecture-dependent strengths. The authors also identify scope and realism limitations and point to further applications and dataset expansion.
- ATOM3D provides benchmark datasets and computational tools for building machine learning models for 3D molecular data.
- Across nearly all tasks that can be formulated in lower dimensions, 3D molecular learning improves performance over 1D and 2D methods.
- Architecture selection is critical: 3DCNNs, GNNs, or ENNs may be most appropriate depending on the underlying data structure.
- ATOM3D datasets are biased toward biomolecules with solved structures, so benchmark performance may not generalize to intrinsically disordered or transmembrane proteins.
- Several benchmarks use artificial task formulations that differ from typical biological, chemical, and drug-design settings, despite their usefulness for evaluating representations and methods.
- Future directions include virtual screening, pose prediction, conformational ensembles, and additional datasets and tasks from the research community.
- The framework is intended to lower entry barriers, encourage algorithms for 3D atomistic data, and promote a structural-biology and medicinal-chemistry research paradigm.
- Using complex molecular systems can make interpretation and error assessment more difficult, with potential consequences for medicine and patient treatment.
C Working with 3D Molecular Data Using ATOM3D
ATOM3D provides guidelines and computational tools for working with its datasets and curating new ones, with continued development planned for the atom3d package.
- ATOM3D supplies high-level dataset guidelines and atom3d tools for working with existing datasets and curating new ones.
- The authors plan to continue developing and supporting the package’s functionality.
C.1 Assembling New Datasets
ATOM3D assembles standardized molecular datasets from public repositories while emphasizing structural coverage, data quality, and the limits of static experimental structures.
- atom3d mines and converts common molecular formats, including PDB, SDF, and XYZ, into a standardized dataset format.
- Available structural databases do not capture the full diversity of biological macromolecules or potential drug molecules.
- Molecular conformations vary with thermal fluctuations, molecular interactions, disordered regions, and environmental factors such as pH.
- Experimental and computational structures contain uncertainties that can produce unrealistic conformations or missing and duplicate atoms.
- Hydrogen atoms are often absent from structural data and may need to be added, requiring consistent preprocessing decisions.
C.2 Developing and Benchmarking New Algorithms
The toolkit separates dataset preparation from algorithms and standardizes evaluation, while cautioning that structural, labeling, and benchmark-design uncertainties affect performance interpretation.
- Dataset preparation should be separated from algorithm implementation so processed structures can be reused across multiple algorithm runs.
- Evaluation should use field-standard metrics and report established metrics alongside justified new metrics to preserve comparability.
- Structural and label uncertainties mean near-perfect performance is rarely attainable for tasks involving complex biological macromolecules.
- Test-set reuse can make measured improvements reflect minor dataset specifics rather than broadly generalizable progress.
- ATOM3D anticipates adding new tasks and datasets as the field evolves and existing benchmarks become saturated.
D Dataset Preparation
ATOM3D prepares task-specific molecular datasets in a standardized hierarchical format, with dataset-specific processing, filtering, splits, and labels supporting diverse benchmark tasks.
- Dataset Preparation: Each example is stored as an ensemble containing task-specific subunits, allowing atomistic structures to be extracted and transformed for different models.
- Dataset Preparation: The standardized LMDB representation stores atoms as rows in data frames while preserving biomolecular hierarchy and associated labels or metadata.
- QM9: QM9 contains molecular properties including dipole moment, orbital energies, thermochemical quantities, and heat capacity, for which benchmark metrics are reported.
- QM9: QM9 processing excludes 3,054 geometrically inconsistent molecules and 1,398 molecules that RDKit cannot process so models use the same data.
- Protein Interface Prediction: Protein interface data define positive amino-acid pairs from bound structures using a 6 Å heavy-atom neighborhood and map them to unbound structures.
- Protein Datasets: Training datasets are filtered by protein content, length, resolution, experimental method, and interaction criteria before being split for learning.
D.5 Ligand Binding Affinity (LBA)
Ligand binding affinity benchmarks combine curated protein–ligand structures with experimentally measured affinity data, while using task-specific splits and computationally constrained molecular inputs.
- Dataset and Splits: PDBBind provides 4,852 refined protein–ligand complexes filtered by structural quality and invalid ligand-bonding information.
- Dataset and Splits: An additional 60% sequence-identity split contains 3,678 training, 460 validation, and 460 test complexes.
- Model Inputs: The ENN input is reduced to selected heavy elements and binding-pocket atoms within 6 Å of the ligand, capped at 600 total atoms for memory reasons.
- Efficacy Prediction: Active-versus-inactive ligand efficacy data comprise 527 ligands across 27 proteins, with labels drawn from agonist, activator, and inactivator annotations.
- Efficacy Prediction: Docked complexes are split by protein target to test generalization across proteins, using prepared structures and default Glide SP docking.
- Efficacy Prediction: Efficacy prediction excludes information about other ligands at the same protein, although methods using such information may be more accurate when available.
E Task-Specific Experimental Details
ATOM3D benchmarks 3D molecular learning with simple, task-specific architectures spanning convolutional, graph, and equivariant networks. Baselines are selected to represent strong 1D or 2D alternatives where appropriate, while training procedures address task-specific outputs and class imbalance.
- Architectures: Benchmark architectures include 3DCNNs, GNNs, and Cormorant equivariant networks, with simple but robust configurations rather than exhaustive hyperparameter tuning.Learning rate, epochs, and 3DCNN grid size or resolution were tuned using grid search.
- 3DCNN representation: 3DCNN inputs are discretized into task-specific 3D voxel grids, generally at 1 Å resolution, with larger structures requiring 1.3 Å resolution for GPU memory.Grid radii range from 7.5 Å to 40 Å across tasks.
- Training: Binary tasks use class-weighted binary cross-entropy with resampling for imbalance, whereas regression tasks use mean-squared error.Models are trained with Adam or AMSgrad, and the best validation checkpoint is evaluated on the held-out test set.
- GNN representation: GNNs connect atoms within 4.5 Å, weight edges by interatomic distance, and encode heavy atoms with one-hot features.For residue-specific tasks, the embedding of the target amino acid’s Cα atom is extracted instead of summing all node embeddings.
- Baselines: Baselines are chosen as state-of-the-art or near-state-of-the-art methods using 1D or 2D representations, subject to comparable implementations or reported splits.Examples include molecular GNNs and N-Gram Graph XGB for molecular properties, BIPSPI for protein interfaces, TAPE for residue identity, and DeepDTA for binding affinity.
G State-Of-The-Art Methods
ATOM3D compares its models with available state-of-the-art methods using task-appropriate datasets and splitting criteria. The comparisons show state-of-the-art performance in several tasks, while highlighting dataset-standardization and comparability limits in others.
- QM9: QM9 comparisons use reported state-of-the-art results, with message passing, continuous-filter convolutional, and Cormorant networks among the strongest methods.These methods have reached excellent performance on QM9, making improvements difficult.
- PIP: PIP comparisons use BIPSPI with structural and sequence-based amino acid features rather than its sequence-only baseline.The baseline choice changes according to whether the comparison targets non-3D or broader state-of-the-art performance.
- RES: 0.425 was reported for a CNN trained on a balanced residue-environment dataset, whereas 0.572 was reported for a model trained on an unbalanced dataset.ATOM3D restricts training and evaluation to a balanced downsampled subset and uses only half of those environments for fair replicated comparisons.
- LBA: Ligand-binding-affinity comparisons use X-score at a 30% sequence-identity split because the PDBBind core-set standard measures in-distribution rather than generalizable scoring ability.The authors note that the X-score comparison is imperfect because its training procedure substantially reduces the training set.
- State-of-the-art comparisons: 3DCNN, GNN, and ENN models achieve state-of-the-art performance on several tasks, while competing methods for SMP, PIP, and LBA also use 3D geometry.Table 7 marks cases where exact training data differ despite matching splitting criteria.