Title: 3D-GSRD: 3D Molecular Graph Auto-Encoder with Selective Re-mask Decoding
Venue: NeurIPS 2025
Accepted: yes
Keywords: Self-Supervised Learning

3D Molecule Pretraining

Molecular Property Prediction

3D Masked Graph Modeling

3D Auto-Encoder
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.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper addresses a key concern in masked graph models: an overly expressive decoder can undermine the encoder, leading to suboptimal representations for downstream tasks. Thus, enhancing reconstruction quality under a restricted decoder is crucial.\n\nSTRENGTHS AND WEAKNESSES:\nPros:\n\n1.The empirical analysis in Section 5 is well-executed and informative.\n\n2.The authors have released the code, promoting transparency and reproducibility.\n\n3.Investigating potential issues with decoders in masked graph modeling (MGM) is a valuable and timely research direction.\n\nCons:\n\n1.The method description lacks clear logic and motivation, making it difficult to understand the rationale behind the design choices.\n\n2.It remains unclear whether the 2D-to-3D distillation is intended to transfer knowledge from 3D to 2D. The authors claim this prevents structural leakage in the 2D encoder. However, as stated in the introduction, GNNs inherently reveal structural information. Thus, the motivation for distillation is ambiguous, and the distilled knowledge appears to implicitly encode structural information.\n\n3.Figure 3 is unclear and unintuitive—at first glance, it is not obvious that SRD represents a distillation mechanism. Is it trained jointly with the reconstruction module? This should be clarified.\n\n4.In the related work section, classifying GraphMVP and 3D InfoMax as 2D molecular graph pretraining methods may be inaccurate, as both are fundamentally 3D-based approaches.\n\n5.The two challenges outlined in the introduction have already been discussed in prior work, such as SimSGT. It is unclear what novel perspective or contribution this paper offers in comparison.\n\n6.The baseline methods used in the experiments are somewhat outdated, with most comparisons limited to works from around 2023.\n\n7.The paper appears to introduce a new 3D encoder, 3D-ReTrans, which contradicts the claim of encoder-agnosticism. What happens if existing encoders are used instead? How does this encoder differ from prior works? Additionally, the reported performance differs from results in the original papers:\n\n[1] Spherical Message Passing for 3D Molecular Graphs\n\n[2] ComENet: Towards Complete and Efficient Message Passing for 3D Molecular Graphs\n\n[3] A New Perspective on Building Efficient and Expressive 3D Equivariant Graph Neural Networks\n\nQUESTIONS:\nSee above.\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 2, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-06-08"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces a novel masked modeling approach for 3D molecular pre-training, together with a dedicated molecule transformer architecture designed to complement this pre-training method. The central contribution is a 3D molecular autoencoder pre-training scheme that addresses the issue of 2D structural information “leakage” to the molecule decoder, which has limited the effectiveness of previous masked molecule modeling frameworks. To tackle this, the authors decouple the 2D graph encoder from the 3D graph encoder, providing masked atoms with 2D structure information exclusively through a separate encoder. This 2D encoder is aligned with the 3D encoder via contrastive learning, which both distills 3D information into the 2D representation and enforces consistency between the modalities. In addition, the paper proposes a new molecule transformer architecture for 3D encoding, inspired by TorchMD-Net but enhanced with 3D data augmentation strategies to promote equivariant learning—an approach that is, to my knowledge, novel in this context.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths**\n\n1. The authors cleverly noticed the 2D graph structure leakage to the molecule decoder, which is a rather novel observation. Addressing this issue, they used a separate 2D graph encoder to introduce 2D structure information, and apply a structure-unaware molecule decoder. This 2D graph encoder is trained by contrastive learning with the 3D encoder. Overall, I feel the observed challenge is insightful, and the proposed solution is intricate but makes sense.\n2. The method is empirically validated with strong performance on relevant benchmarks, demonstrating improvements over competitive baselines. The experimental section (Section 5) provides supporting analysis for key design choices.\n3. The paper is generally well-written. The technical motivation, architectural decisions, and experimental results are clearly articulated, making the work accessible to the community.\n\n\n**Weakness**\n\n1. The paper would benefit from more detailed analysis and visualization of the learned 2D encoder. For example, does the 2D encoder capture the necessary structural information required by the decoder, and how does this representation evolve during contrastive training?\n2. The 2D-from-3D distillation approach bears some resemblance to the methodology in GraphMVP. A clearer discussion of similarities, differences, and relative advantages compared to GraphMVP would strengthen the paper.\n3. While the proposed molecule transformer is motivated as an advance over relational transformers, its specific improvements over TorchMD-Net are less clear. Further ablation studies isolating the contributions of the new architectural elements (e.g., 3D data augmentations) would be helpful.\n\nQUESTIONS:\n1. Can you report or discuss the computational efficiency of the proposed masked modeling method and the new transformer architecture, in comparison to prior approaches?\n2. A straightforward baseline for the 2D encoder would be to use structural embeddings, as in the GraphGPS paper. Could the authors provide an ablation study comparing their contrastive alignment strategy to such alternatives?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-25"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper addresses a core dilemma in 3D molecular modeling using re-mask decoding: whether or not to leak 2D information into the decoder. The authors propose a clean solution by introducing a student–teacher framework that selectively reveals only the 2D information already learned by the encoder to the decoder. This effectively mitigates the 2D leakage issue. Experimental results on downstream 3D molecular tasks demonstrate the effectiveness of the proposed approach.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths**:\n\n1, The paper is well-written, with a clear motivation and compelling empirical evidence (e.g., Figure 2) supporting the proposed approach. And all the key claims made early on are later validated through ablation studies.\n\n2, This paper actually effectively addresses two critical challenges: (1) the pre-training vs. downstream mismatch in 3D molecular modeling, resolved by reusing the re-mask technique from 2D models; and (2) the inherent conflicts in re-mask decoding. Both issues are highly relevant to 3D molecular representation learning, and the proposed methods and insights are likely to be influential for the community.\n\n**Weaknesses**\n\n1, My main concern lies in the magnitude of the technical contribution. While the analysis are strong, the technical novelty appears limited. Although the authors claim to use a graph auto-encoder, the model architecture is in fact very similar to transformer-based models already used in the 3D molecular domain, such as Uni-Mol and Mol-AE, both of which are already cited in the paper. In particular, the design is highly similar to Mol-AE, with almost identical architecture and atom drop strategy — the only difference lies in the positional encoding used in the decoder.\n\n2, The use of a distilled PE does not fully eliminate 2D leakage. For example, in an image–text CLIP procedure, we cannot assume that aligned embeddings are completely free of modality-specific information. Additionally, in this paper’s setting, the information of dropped atoms is not present in the 3D encoder output, so their 2D structure is inevitably leaked to the decoder by 2D PE.\n\n3, There exists a much simpler solution to eliminate 2D leakage altogether. During downstream evaluation, you could directly use the 3D embedding combined with the 2D PE as the final molecular embedding—potentially eliminating the need for the distillation design entirely.\n\nQUESTIONS:\nCould the authors conduct an additional experiment based on Table 1? I'm particularly curious about the performance when using SRD without distillation, but instead representing the molecule with 3D embeddings combined with 2D positional encoding for downstream tasks.\n\nIt would also be helpful to try this setup on your current with-distillation model as well. Since it doesn’t require redoing pre-training, this experiment should be feasible to complete during the rebuttal period.\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 4, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-07-03"}