Title: Dynamic and Chemical Constraints to Enhance the Molecular Masked Graph Autoencoders
Venue: NeurIPS 2025
Accepted: yes
Keywords: molecule representation learning

graph auto-encoder

information bottleneck
Abstract: Masked Graph Autoencoders (MGAEs) have gained significant attention recently. Their proxy tasks typically involve random corruption of input graphs followed by reconstruction. However, in the molecular domain, two main issues arise: the predetermined mask ratio and reconstruction objectives can lead to suboptimal performance or negative transfer due to overly simplified or complex tasks, and these tasks may deviate from chemical priors. To tackle these challenges, we propose Dynamic and Chemical Constraints (DyCC) for MGAEs. This includes a masking strategy called GIBMS, which preserves essential semantic information during graph masking while adaptively adjusting the mask ratio and content for each molecule. Additionally, we introduce a Soft Label Generator (SLG) that reconstructs masked tokens as learnable prototypes (soft labels) rather than hard labels. These components adhere to chemical constraints and allow dynamic variation of proxy tasks during training. We integrate the model-agnostic DyCC into various MGAEs and conduct comprehensive experiments, demonstrating significant performance improvements. Our code is available at \url{https://github.com/forever-ly/DyCC}.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes improvements to Masked AutoEncoder-based (MAE) models applied to molecular graph tasks. While MAE is trained through masked reconstruction, its performance is highly sensitive to how the mask is selected—both in terms of region and ratio. Moreover, challenges remain regarding the informativeness of the masked regions and the uniqueness of reconstruction targets.\n\nTo address these issues, the authors propose a masking strategy based on the Information Bottleneck (IB) principle. Since IB identifies salient regions per sample based on compression objectives, it offers a principled approach to mitigate the above concerns. The paper formulates a self-supervised learning framework that incorporates Graph IB and further introduces an objective function that evaluates token labels using soft-labels rather than hard-labels.\n\nThe main experiments are conducted on MoleculeNet for property prediction tasks, demonstrating consistent performance improvements across three variants of Molecular Graph AutoEncoders (MGAEs) when trained with the proposed method.\n\nSTRENGTHS AND WEAKNESSES:\nPros\n- The proposed learning strategy is broadly applicable, regardless of the specific MGAE variant, indicating good generalizability.\n- Using IB to extract sample-specific salient regions and reflect them in the mask strategy is a reasonable and theoretically grounded approach.\n- According to the Appendix, the performance gains appear robust across different hyperparameter settings, which is a valuable practical property.\n\nCons\n- I have some concerns about the performance improvements reported in Table 1. Given the magnitude of the variance, can it be confidently stated that the proposed method yields a statistically significant \"boost\" in performance?\n\nIn summary, I like the proposed methodology, but am not confident in the strength of experimental supports.\n\nQUESTIONS:\n- In Appendix Table 5 (DTA regression experiments), it appears that models such as MolBERT and AttrMask were not evaluated using the proposed method. Was there any practical difficulty in applying the method to these models? If so, this could constitute a limitation of the proposed approach and should be mentioned explicitly.\n- One of the key advantages of the proposed method is its reduced reliance on arbitrary mask selection. Conversely, has there been any experimental verification of how much performance fluctuates in existing methods when masks are chosen arbitrarily? If such analysis exists, it would strengthen the motivation for your approach and help quantify its benefits more precisely.\n\nLIMITATIONS:\nYes (As far as I understand)", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors identify that pre-training masked graph autoencoders (MGAEs), commonly used to model moecules, lack the flexibility necessary to represent chemical information. This is because: 1) the pre-training masks a fixed number of tokens, 2) the tokens are chosen arbitrarily when some components of molecules clearly matter more than others, and 3) the tokens are decoded into specific discrete labels. To mitigate these issues, the authors introduce Dynamic Chemical Constraints (DyCC), a model-agnostic pre-training algorithm composed of two modules: a) graph-information bottleneck mask strategy (GIBMS), which identifies the most relevant tokens to mask using graph information theory, and b) Soft Label Generator, which replaces the discrete labels per token with a distribution over learnable prorotypes for each token. DyCC is used to pretrain a number of MGAEs, showing an improvement over regular pretraining in most cases.\n\nSTRENGTHS AND WEAKNESSES:\n## Strengths\n- The method is theoretically grounded. The GIBMS loss (equation 13) is derived from graph information bottleneck following detailed and well-justified steps. \n- Adapting GIB to the self-supervised setting is an interesting use case extending the applicability of the framework to various applications.\n- The method was thoroughly tested on a number of applications, and the results are reported transparently with detailed experimental procedures and error intervals.\n\n## Weaknesses\nMy main concerns with this work are with its **analysis of empirical results** and its **presentation**. I elaborate below on these two points and ask for related clarifications in **Questions**.\n* **Unsubstantiated claims in the empirical analysis (Section 4.2)**: I am unable to find evidence for the following claims in the paper:\n    - \"SimSGT-DyCC surpasses the \"no pretrain\" model by 10.4%.\" the average difference between the two models is actually 7.1% according to Table 1. If the authors are referring to the results for a specific dataset, this should be made clearer (I could not find any dataset with the exact difference reported). It would also be misleading to refer to a specific dataset instead of the average performance unless justified. \n    - \"\"substantial performance variations were observed among the original AttrMask, MoleBert, and SimSGT models due to differences in tokenizers.\": why is a difference of approx 2% (according to Table 1) considered substantial? \n    - \"The success rate of GIBMS was 74%, demonstrating the effectiveness of the GIBMS module.\": this is a promising result, but can you compare it to other baselines (in identifying relevant components of a graph) to get a sense of its importance?\n    - Figure 4 (left): can you report the results of a similar experiment on a larger scale? for example, you can compute the accuracy of identifying non-ring components as the most relevant elements in ring molecules. \n\n- **Presentation**: I found the way you motivated your approach confusing, especially in the introduction. If I understand your text correctly, your main concerns are that regular masking does not leverage relevant chemical information, both in the choice of tokens to mask (relevant tokens differ per molecule) and in the labels chosen when unmasking (many possible labels are equally good). The introduction  breaks down these concerns into irrelevant details (e.g. the 3 chemical constraints) especially since most are addressed by the same solution. I suggest either justifying your taxonomy through additional examples and by linking it to the part of the module addressing it, or simplifying the motivation in the introduction. \n\n- **Pointing out errors in the text**: \n    - paragraph \"Are both GIBMS and SLG necessary\" is redundant. This point was already made in the paragraph right before. I suggest moving the paragraph \"limitations\" in the appendix to the main text instead.\n     - equation 13 has an additional L_cl term, I think this is supposed to be the same as L_pred and therefore only one of the two should be mentioned in the equation?\n     - you are missing a legend in Figure 1. Also the colors of the plot lines changed between subfigures a and b in the same figure. \n     - line 36: change make => make\n     - line 176: typo near  \"To ensure\"\n     - I suggest introducing $Y^{sup}$ notation from the beginning, instead of replacing $Y$ later in the text.\n     - table 6 in the appendix missing bolding.\n\nQUESTIONS:\n- Can you report the training time as well in Appendix C.1?\n- In figure 3, you mention that SGL yields a more uniform distribution of tokens on the ZINC dataset, when the data has 95% of atoms with the 3 most common labels. Is a uniform distribution desired here?\n\nLIMITATIONS:\nYou could add additional limitations regarding the societal impact of your work.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 2, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe manuscripts presents two techniques to improve self-supervised learning on graphs. First, a masking strategy that learns an individual masking probability for each atom in a molecule based on its importance. Second, a soft label technique that replaces the learned atom embeddings by learned prototype embeddings and allows the model to effectively lower the loss for predicting atoms or atom environments that are similar to each other.\n\nSTRENGTHS AND WEAKNESSES:\nThe manuscripts provide two different techniques for improving self-supervised learning on graphs and demonstrate on benchmarks that \n1.) these techniques can be applied on different models\n2.) consistently improve the pretraining results for these models\n\nThe manuscript is well written. I am not able to verify the proofs in manuscript but the method sounds reasonable. Given that both techniques, node importance and soft labels, can be easily visualized and manually verified, I wished the manuscript would put more emphasize on discuss if the node importance or soft labels do make sense. For example, Figure 1a says that Br and Cl are interchangeable in this molecule, but there is no discussion or evaluation in the manuscript if the soft labels indeed cluster together atom environments that are easily exchangeable or have similar chemical properties. The same is true for the node importance prediction. In 4.3 it is mentioned that in 74% of the cases in a mutagenic dataset, 30% of the most important atoms were NH2 or NO2, which is chemical reasonable. However, molecules usually consist of a carbon backbone and 30% of the most important atoms in a molecule should already cover most of the non-carbon atoms. It is not that surprising that the method founds out that carbon atoms are not important. \nI also wonder if the atom importance is \"stable\" over multiple training runs, or if the method always selects a different set of atoms to be important.\n\nQUESTIONS:\n- I remember a similar paper by Liu et al. called \"Where to Mask: Structure-Guided Masking for Graph Masked Autoencoders\". They also developed a node importance strategy and I hoped to compare the benchmark results between both papers. I found that the numbers differ, though. For example, in Liu et al. the numbers for AttrMask are better than in this manuscript. The same is for the original manuscript of Hu et al. The manuscript mentions that the numbers of the baseline methods were taken from prior studies, is this also the case for AttrMask and, if so, why are the numbers different?\n\n- can you evaluate if the importance scores of the atoms are stable across different pretraining runs or if a different set of atoms is chosen to be \"important\" every time?\n\n- The soft label approach somewhat reminds me on vector quantization (VQ), which is also often used to bridge regression methods and discrete tokens. Can you elaborate why your softlabel approach works better and what makes it different?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 2, "CLARITY": 4, "IMPACT": 4, "ORIGINALITY": 2, "DATE": "2025-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces a new method of dynamic masking strategy to enhance molecular masked graph autoencoder. The authors first show that a static mask ratio and random graph corruption are suboptimal in molecular graph pretraining with masking-reconstruction loss. To address this issue, this paper utilize Graph Information Bottleneck to calculate an importance socre of each atom, based on which a masking probability is determined, so that a adaptive masking strategy can be applied to each molecule. Experiments are conducted on a dataset sourced from the 252 ZINC15 database and the proposed strategy are integrated into three different Masked Graph Autoencoders for molecular representation learning. Experimental results show that the proposed method enhance the performance of the three Masked Graph Autoencoders and the integrated models outperform existing a series of molecular pretraining methods.\n\nSTRENGTHS AND WEAKNESSES:\nStrength\n1.It is valuable to show that dynamic masking based on the importance score of each atom is superior than random masking with a fixed masking ratio in masked graph autoencoder pretraining. \n2.The paper is generally well written and easy to follow and the motivation is reasonable. \n3.The strategy of using information bottleneck to calculate an importance score of each atom is feasible. The contribution of using information bottleneck in an unsupervised scenario is valuable. \n\nWeakness\n1.The major weakness is that the baseline methods are weak and outdated. The proposed method is compared to quite a lot self-supervised graph pretraining models, but the results are from MoleBert [41], which is published in 2023. After the publication of [41], self-supervised pretraining for molecular representation learning has advanced rapidly and quite a lot new methods have been proposed. A search of papers citing [41] can find some of them. Here I list several papers that the authors may consider.\n\n[1] Hongxin Xiang, et al. An Image-enhanced Molecular Graph Representation Learning Framework. IJCAI 2024\n[2] Liang Wang, et al. MOLSPECTRA: PRE-TRAINING 3D MOLECULAR REPRESENTATION  WITH MULTI-MODAL ENERGY SPECTRA. ICLR 2025. \n[3] Yue Wan, et al. Multi-channel learning for integrating structural hierarchies into contextdependent  molecular representation. Nature Communications 2025. \n\n2.Relevant to the first point, this paper lacks a review of self-supervised molecular pre-training techniques. A comprehensive review of this field can place the current research in the correct context.\n\n3.The proposed method is only evaluated on one dataset. \n\n4.The is not comparison to conventional masking strategies with proper hyper-parameter selection.\n\nQUESTIONS:\n1. Please respond to the weakness points.\n2. Does the proposed method bring high time cost to pretraining?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 4, "CLARITY": 2, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-07-02"}