3D-GSRD: 3D Molecular Graph Auto-Encoder with Selective Re-mask Decoding

TL;DR

3D-GSRD tackles the challenge of learning molecular representations from 3D data by introducing Selective Re-mask Decoding (SRD) that remasks only 3D-relevant information while preserving 2D structural context. The approach uses a 3D Relational-Transformer encoder and a structure-independent Transformer decoder, with a 2D Graph Position Encoder distilled from the 3D encoder to prevent leakage. The distillation and SRD together shift the encoder's focus toward molecular representation learning, achieving state-of-the-art results on 7 of 8 MD17 targets and strong performance on QM9. The method demonstrates robust 3D/2D integration and generalizes SRD across encoders, with released code for reproducibility.

Abstract

Masked graph modeling (MGM) is a promising approach for molecular representation learning (MRL).However, extending the success of re-mask decoding from 2D to 3D MGM is non-trivial, primarily due to two conflicting challenges: avoiding 2D structure leakage to the decoder, while still providing sufficient 2D context for reconstructing re-masked atoms. To address these challenges, we propose 3D-GSRD: a 3D Molecular Graph Auto-Encoder with Selective Re-mask Decoding. The core innovation of 3D-GSRD lies in its Selective Re-mask Decoding(SRD), which re-masks only 3D-relevant information from encoder representations while preserving the 2D graph structures. This SRD is synergistically integrated with a 3D Relational-Transformer(3D-ReTrans) encoder alongside a structure-independent decoder. We analyze that SRD, combined with the structure-independent decoder, enhances the encoder's role in MRL. Extensive experiments show that 3D-GSRD achieves strong downstream performance, setting a new state-of-the-art on 7 out of 8 targets in the widely used MD17 molecular property prediction benchmark. The code is released at https://github.com/WuChang0124/3D-GSRD.

Figures (8)

• Figure 1: Illustration of 3D MGM with re-mask decoding and 2D structure leakage.
• Figure 2: Reconstruction loss across two pretraining settings. We compare settings using frozen encoders (3D-ReTrans) and structure-independent (Transformer) versus structure-dependent (3D-ReTrans) decoders.
• Figure 3: Overview of 3D-GSRD. It contains three key elements: (1) a 3D-ReTrans encoder; (2) the SRD that re-masks only 3D-relevant information from the encoder representations while preserving its 2D structure information via 2D-from-3D distillation; (3) a structure-independent decoder.
• Figure 4: Illustration of 3D-ReTrans. (a) 3D-ReTrans is constructed by stacking multiple 3D Relational-Attention and 3D Update Layers. (b) 3D Relational-Attention that processes both atom-wise and pair-wise features. (c) 3D Update Layer that includes a residual connection.
• Figure 5: Reconstruction loss across four pretraining settings.