Research map / 3D vision / 3D generation
Beyond Atoms: Enhancing Molecular Pretrained Representations with 3D Space Modeling
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes SpaceFormer, a Transformer-based molecular pretrained representation that discretizes the 3D space around molecules into grid cells (with grid sampling/merging, linear-complexity RoPE and random-Fourier-feature positional encodings, and MAE pretraining) and beats 3D atom-based MPR baselines on 15 limited-data tasks.
Problem
Existing 3D molecular pretrained representation models treat molecules as discrete atom sets and ignore the surrounding 3D space, which may limit representation quality for downstream property prediction with limited labeled data.
Approach
Discretize the whole 3D cuboid around a molecule into grid cells (each atom or empty cell becomes a token), reduce cell count via random grid sampling and adaptive grid merging, encode relative positions with linear-complexity 3D RoPE and random Fourier features, and pretrain with masked autoencoding over cells that predicts occupancy and atom type/offset.
Key results
- SpaceFormer achieves best results on 10 of 15 downstream tasks and top-2 on 14 of 15, e.g. HOMO MAE 0.0042 vs 0.0050 for the runner-up Mol-AE.
- Adding randomly sampled virtual points to Uni-Mol steadily lowers validation loss and improves some downstream tasks, but SpaceFormer retains a clear gap over Uni-Mol with virtual points.
- Adaptive grid merging and grid sampling cut the number of input cells by roughly an order of magnitude while keeping performance; MAE pretraining outperforms atom-only masking (S3 vs S5).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use SpaceFormer rather than the strongest baseline | Molecular property prediction · HOMO, LUMO, GAP, E1-CC2, E2-CC2, f2-CC2, Dipmom, aIP, HLM, Solu | MAE | Helped |
| Use SpaceFormer rather than the strongest baseline | Molecular property prediction · f1-CC2, D3_disp_corr, MME | MAE | Hurt |
| Use SpaceFormer rather than the strongest baseline | Molecular property prediction · BBBP, BACE | AUC | Hurt |
| Reduce non-atom cells from 1000 to 200 (ablation) | SpaceFormer · HOMO | MAE: 0.0042 -> 0.0047 | Hurt |
| Remove non-atom cells (ablation) | SpaceFormer · HOMO | MAE: 0.0042 -> 0.0051 | Hurt |
| Scale model width by 4x (ablation) | SpaceFormer · HOMO | MAE: 0.0042 -> 0.0046 | Hurt |
| Replace MAE pretraining with denoising (ablation) | SpaceFormer · HOMO, LUMO, E1-CC2, E2-CC2 | MAE: HOMO 0.0042 -> 0.0044; LUMO 0.0040 -> 0.0041; E1-CC2 0.0058 -> 0.0059; E2-CC2 0.0074 -> 0.0075 | Hurt |
| Add 10 virtual points (ablation) | Uni-Mol · HOMO | MAE: 0.0052 -> 0.0052 | No clear effect |
| Add 10 virtual points (ablation) | Uni-Mol · LUMO, E1-CC2, E2-CC2 | MAE: LUMO 0.0060 -> 0.0058; E1-CC2 0.0067 -> 0.0063; E2-CC2 0.0080 -> 0.0078 | Helped |
| Remove RFF while keeping RoPE (ablation) | SpaceFormer · HOMO, LUMO, E1-CC2 | MAE: HOMO 0.0042 -> 0.0044; LUMO 0.0040 -> 0.0043; E1-CC2 0.0058 -> 0.0059 | Hurt |
| Remove RFF while keeping RoPE (ablation) | SpaceFormer · E2-CC2 | MAE: 0.0074 -> 0.0073 | Helped |
| Remove both RoPE and RFF (ablation) | SpaceFormer · HOMO, LUMO, E1-CC2, E2-CC2 | MAE: HOMO 0.0042 -> 0.0047; LUMO 0.0040 -> 0.0048; E1-CC2 0.0058 -> 0.0066; E2-CC2 0.0074 -> 0.0078 | Hurt |
Surprising findings
- Randomly sampled virtual points, which carry no physical information, improve a pure 3D MPR model's pretraining loss and some downstream results.
- Scaling SpaceFormer width 4x without non-atom cells (S4) still underperforms the grid-cell model on most tasks, suggesting gains are not just from extra FLOPs.
Limitations
- Efficiency remains higher than atom-based methods despite the sampling and merging optimizations.
- Evaluated only on small organic molecules; extension to proteins and complexes is left to future work.
Experimental setup
Datasets: Uni-Mol pretraining dataset (19M molecules) GDB-17 subsets (20K and 21.7K) cata-condensed polybenzenoid hydrocarbons (8.7K) MoleculeNet BBBP and BACE Biogen ADME (HLM, MDR1-MDCK ER, Solubility)
Benchmarks: HOMO LUMO GAP E1-CC2 E2-CC2 f1-CC2 f2-CC2 Dipmom aIP D3_disp_corr HLM MME
Metrics: MAE AUC
Baselines: GROVER GEM 3D Infomax Uni-Mol Mol-AE
Builds on: Uni-Mol Mol-AE RoPE FlashAttention masked autoencoder
Closest papers by meaning
- 3D Denoisers Are Good 2D Teachers: Molecular Pretraining via Denoising and Cross-Modal Distillation (AAAI 2025 · similarity 0.90)
- Swallowing the Bitter Pill: Simplified Scalable Conformer Generation (ICML 2024 · similarity 0.88)
- Learning Over Molecular Conformer Ensembles: Datasets and Benchmarks (ICLR 2024 · similarity 0.88)
- GeoMFormer: A General Architecture for Geometric Molecular Representation Learning (ICML 2024 · similarity 0.88)
- Multimodal Molecular Pretraining via Modality Blending (ICLR 2024 · similarity 0.88)
- nach0-pc: Multi-task Language Model with Molecular Point Cloud Encoder (AAAI 2025 · similarity 0.88)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.