Beyond Atoms: Enhancing Molecular Pretrained Representations with 3D Space Modeling

TL;DR

SpaceFormer introduces a principled 3D space modeling framework for molecular pretrained representations by discretizing the surrounding 3D space into grid cells, applying grid sampling/merging for efficiency, and using efficient 3D positional encodings (RoPE-based directions and RFF-based distances). Coupled with MAE pretraining, SpaceFormer demonstrates strong gains over prior atom-centric 3D MPR models across 15 downstream tasks with limited data, ranking first on 10 tasks and top-2 on 14. Ablation studies confirm the importance of grid discretization, adaptive merging, 3D PE, and MAE, while analyses highlight favorable scaling properties and the benefit of modeling space beyond atoms. The work suggests a new paradigm for molecular representation learning and offers potential for applying 3D space modeling to larger biological complexes like proteins.

Abstract

Molecular pretrained representations (MPR) has emerged as a powerful approach for addressing the challenge of limited supervised data in applications such as drug discovery and material design. While early MPR methods relied on 1D sequences and 2D graphs, recent advancements have incorporated 3D conformational information to capture rich atomic interactions. However, these prior models treat molecules merely as discrete atom sets, overlooking the space surrounding them. We argue from a physical perspective that only modeling these discrete points is insufficient. We first present a simple yet insightful observation: naively adding randomly sampled virtual points beyond atoms can surprisingly enhance MPR performance. In light of this, we propose a principled framework that incorporates the entire 3D space spanned by molecules. We implement the framework via a novel Transformer-based architecture, dubbed SpaceFormer, with three key components: (1) grid-based space discretization; (2) grid sampling/merging; and (3) efficient 3D positional encoding. Extensive experiments show that SpaceFormer significantly outperforms previous 3D MPR models across various downstream tasks with limited data, validating the benefit of leveraging the additional 3D space beyond atoms in MPR models.

Figures (4)

• Figure 1: Overview of SpaceFormer. For clarity, we illustrate the model in 2D space. SpaceFormer begins by discretizing the 3D cuboid containing molecules into grid cells. Grid sampling/merging strategy is then applied to reduce the number of input cells for efficiency. Note that exact atomic positions are still encoded after discretization owing to various positional encoding (PE). To efficiently encode pairwise positional relationships, SpaceFormer utilizes 3D Directional PE with RoPE and 3D Distance PE with Random Fourier Features.
• Figure 2: Log-scale valid loss of Uni-Mol when adding virtual points. The blue horizontal line represents the original Uni-Mol model, while the orange curve represents the Uni-Mol performance with respect to the number of virtual points.
• Figure 3: Adaptive grid merging. Dark cells represent merged cells, and $L$ denotes the number of merging steps.
• Figure 4: Efficiency comparison of SpaceFormer and Uni-Mol with different numbers of cells/points.