Pushing the Limits of All-Atom Geometric Graph Neural Networks: Pre-Training, Scaling, and Zero-Shot Transfer
Geometric Graph Neural Networks;Self-supervised Pre-training;Scaling;Zero-shot Transfer;Molecular Representation
Yes
No ethics review needed.
Many of the evaluations are on datasets for specific molecules, which are very useful for understanding specific model behavior, but should be complemented by more general evaluations, especially in the context of examining scaling behavior. Can the authors comment on or provide additional evidence that these specific, bespoke evaluations are connected to more general relationships between pre-training setups and downstream evaluations?
The paper continues an important line of work in exploring the utility of pre-training and scaling GNNs for molecular learning tasks. Unlike other dominant areas of deep learning, all-atom molecular representation learning relies on GNNs and does not directly benefit from advances in scaling sequence-based models.
The authors explore model size, aspect ratio, nearest neighbor cutoff radius, and architecture to provide a comprehensive look into the scaling behavior of molecular GNNs.
In this paper the authors extend previous work on scaling laws for all-atom graph neural networks applied to self-supervised and supervised training. They investigate aspects of pre-training task choice, different downstream evaluations, GNN model size and aspect ratio, as well as the radial cutoff for constructing nearest neighbor graphs.
Like most molecular GNN works, the authors are limited by the available evaluations (e.g., QM9). xxMD is an interesting evaluation, but these datasets are limited to a small set of specific molecules. QM9 with B3LYP in particular is not informative, and the authors might consider newer benchmarks like POLARIS or a subset of the Therapeutic Data Commons to strengthen their evaluations.
The paper does not offer a clear and concise summary of recommendations based on the empirical findings, which is essential to achieving the stated aim of inspiring practitioners to rethink GNN training.
Yes
No ethics review needed.
1. The paper claims to demonstrate zero-shot transfer learning using pre-trained Geom-GNNs. However, the described method seems closer to downstream task fine-tuning, given that a separate network is trained with the VAMP score objective. Could the authors clarify how this approach qualifies as zero-shot transfer rather than fine-tuning?

2. Limited Pre-Training Approaches:
The pre-training in this work is restricted to coordinate denoising. Given that prior work has successfully used a combination of node mask prediction and coordinate denoising for improved performance. Specifically, how might adding a node mask objective influence for the molecular kinetics tasks? Would it enhance the model’s ability to generalize across different molecular conformations?
Additionally, could the authors hypothesize the potential impact of such an extended pre-training approach on scaling behavior? 

3. While the effectiveness of ET and ViSNet is demonstrated on several tasks, the study lacks comparisons with invariant feature-based networks (e.g., SchNet, DimeNet, GemNet) and tensor product-based networks (e.g., Equiformer, MAC). Could the authors provide insights into how their method might perform relative to these alternative architectures?
1. Demonstrates substantial performance improvements in kinetic modeling on VAMP score metrics by utilizing denoising pretraining techniques.
2. Applies the self-supervised pretraining approach to a variety of downstream tasks, successfully proving its effectiveness across applications.
3. Examines scaling laws in both standard equivariant models and pre-trained equivariant models, finding that even pre-trained models diverge from typical neural scaling laws due to issues like under-reaching and over-smoothing. The author suggests that, while scaling models has its benefits for supervised and unsupervised tasks, it may be more effective to focus on addressing data label uncertainty and using active token mixing to mitigate information bottlenecks.
The paper investigates denoising pretraining and potential scaling laws for geometric graph neural networks (GNNs) on supervised tasks. These GNNs are pre-trained on noisy coordinate data to learn how to denoise and reconstruct the original coordinates. The effectiveness of this approach is tested on various downstream applications, including molecular kinetics, fold classification, and energy and force prediction. Additionally, the paper examines the scaling behavior of these models and highlights specific limitations in supervised prediction tasks.
1. The author claims to be the first to demonstrate the pre-trained Geom-GNNs’ capability as zero-shot transfer learners. However, I would consider this more appropriately described as downstream task fine-tuning. In zero-shot learning, a model test on data source is directly applied to unseen data class without additional training (e.g., training on English and French, then testing on Chinese without further adjustments). Unlike this approach for molecular kinetics, the paper involves training a separate network with the VAMP score objective.
2. The pretraining methods in this work are limited to coordinate denoising. Other approaches [1,2] that leverage both node mask prediction and coordinate denoising have already proven effective.
3. Although the paper demonstrates the effectiveness of ET and VisNet on several tasks, it does not include evaluations on invariant feature-based networks (such as SchNet, DimeNet, or GemNet) or tensor product-based networks like Equiformer and MAC.

[1] Cui, Taoyong, et al. "Geometry-enhanced pretraining on interatomic potentials." Nature Machine Intelligence 6.4 (2024): 428-436.
[2] Zhou, Gengmo, et al. "Uni-mol: A universal 3d molecular representation learning framework." (2023).
Yes
No ethics review needed.
as Weaknesses.
1. **Innovation**: The study presents a novel perspective, exploring an area that remains under-researched (to my knowledge), offering significant insights for the development of Geom-GNNs.
2. **Clarity and Structure**: The article is well-organized, with clear presentation and summary of experiments and viewpoints, facilitating reader comprehension.
3. **Robust Experimentation**: The experimental design is thorough, effectively supporting the authors’ conclusions.
4. **Exploration of Zero-Shot Transfer Capability**: The investigation into the zero-shot transfer ability of Geom-GNNs is intriguing, with experiments indicating their potential as excellent zero-shot transfer learners.
5. **Pre-training Insights**: Through extensive denoising pre-training tasks, valuable experiences have been gained regarding the pre-training of Geom-GNNs, including aspects such as model width, depth, aspect ratio, and the cutoff radius in geometric atomic graph construction, providing rich guidance for pre-training.
6. **Advancement of Unified Architecture**: Given the widespread attention and efforts in the research of all-atom Geom-GNNs, this study effectively inspires researchers to reconsider the design of Geom-GNN architectures and the adjustment of training strategies, thereby promoting the development of a unified Geom-GNN architecture.
This work presents a novel investigation into whether pre-trained Geom-GNNs (Graph Neural Networks for Conformational Molecules) possess efficient and transferable geometric representation capabilities, particularly in addressing the low generalization of models typically trained on specific tasks. The authors also aim to introduce “Neural Scaling Laws” to summarize the performance behavior of these pre-trained Geom-GNNs. However, it is unfortunate that the experimental results indicate that Geom-GNNs do not adhere to power-law scaling laws and fail to demonstrate predictable scaling behavior across various supervised tasks. Furthermore, the findings reveal that the all-atom embedding graph representations derived from Geom-GNNs exhibit excellent expressive capabilities, suggesting that Geom-GNNs can function effectively as zero-shot transfer learners.
1. In Figure 6, which explores different model widths, even though the x-axis represents the total number of parameters (with model depth held constant), it would be more beneficial to indicate the model width for each point in the legend to enhance result presentation. Similarly, in Figures 4 and 7, which demonstrate the impact of model depth, using the legend to specify the exact number of layers might be more effective. In general, clear legends are always advantageous.
2. The comparison between models trained from scratch and those fine-tuned in Section 6.1 could be more comprehensive if extended to include model depth. Previous discussions (albeit in the context of pre-training) have presented certain viewpoints, and it is anticipated that these would also have significant effects during fine-tuning.
Yes
No ethics review needed.
- Why is the VAMP task considered zero-shot?
- Appendix D: Why is there a difference between datasets for pretraining and downstream? According to Table 2, QM9 there seem also experiments with pretraining. Are models in Table 3 not fine-tuned?
relevant research topics in the paper:
- pretraining of (geometric) GNNs
- scaling behavior of GNNs
- oversmoothing, underreaching

experiments:
- experiments considered not only random splitting, but also scaffold and temporal splits
In their work the authors empirically studied properties of pre-trained geometric GNNs (i.e., GNNs with coordinates attached).
They especially seemed to consider pretraining via denoising as introduced by Zaidi et al. (2022).
The authors consider several downstream tasks: molecular dynamics (kinetics modeling), fold classification, prediction of quantum-chemical properties, and, molecule conformations
They study properties such as power-law scaling and find that the geomtric GNNs do not follow that on the pre-training task.
The authors conclude that geometric GNNs "saturate early and are bounded by issues such as under-reaching and over-smoothing" and further that "other aspects become more crucial, such as addressing data label uncertainty and employing active token mixing to mitigate the information bottleneck."
In general the largest problem to my opinion is that the authors lack to specify a clear research goal. Methodologically there seems not too much novelty.
But if it is primarily an empirical research paper it seems very important to me that the research goal is clearly defined and dependent on that authors have
to reason why they are selecting the problems they look at and why they are selecting the methods they compare to. If a new benchmark dataset is used (as e.g., scaffold QM9), then empirical evaluations/comparisons should take into account at least some previously suggested methods (which are ideally complementary to each other to better see the potentials of each of the methods on the new benchmark and being able to compare to what authors say).
\
\
Details:
\
\
**hard to see the novelty in the paper**:\
Fact that pretraining is useful was already found e.g. by Zaidi et al. (2022).\
No new methodology seems to be suggested (or is "token mixing" the new method)?\
Study of scaling behavior might be interesting and to a certain degree novel for geometric GNNs, but no follow-up investigations seemed to be employed to draw formal connections to GNN properties like oversmoothing, underreaching, etc.
\
\
**paper is structured in a strange way, which makes it hard to follow the paper**:\
It is actually hard to understand what the research aim of the authors was.\
chapter 4, 5, and, 6, actually seem to be about empiricial experiment results.\
The setups are partly however already explained in chapter 3.\
chapter 4 and 6 show performances on problems\
chapter 5 however studies power-law behavior and other ablations.\
In sum it gets hard to follow the paper. Better reasoning why some experiments are applied at all (with respect to the general research goal) and why they are done as they are done is necessary (e.g., why which methods are selected for comparison or why which dataset is used).
\
\
**experiments**:\
although there are some good points as mentioned in strengths (such as the splits), it will get hard to understand how large the impact of pretraining really is, as the authors only test very few method once with pretraining and once without pretraining.\
\
also the there is no good argument in the paper, why the authors exactly compare to those methods they selected to compare to 
e.g.:\
for molecular dynamics, why don't they compare to Timewarp\
for QM9 and xxMD they could e.g. compare to "E(n) Equivariant Graph Neural Networks", "SE(3)-Transformers", DimeNet++, MACE, etc.\
An option to get more impression on the significance of the author's results would be to additionally compare to standard QM9, etc. (where there are also a lot of method comparisons out).
\
\
**minor points**:\
"token mixing" not defined, but heavily used\
grammatical errors/typos: "In silico molecular computation and simulation are indispensable tool in..."
This paper explores using pre-trained geometric graph neural networks as effective zero-shot transfer learners for molecular conformation representation in out-of-distribution scenarios, and their scaling behavior under various setups.
@inproceedings{
anonymous2024pushing,
title={Pushing the Limits of All-Atom Geometric Graph Neural Networks: Pre-Training, Scaling, and Zero-Shot Transfer},
author={Anonymous},
booktitle={Submitted to The Thirteenth International Conference on Learning Representations},
year={2024},
url={https://openreview.net/forum?id=4S2L519nIX},
note={under review}
}
Constructing transferable descriptors for conformation representation of molecular and biological systems has been a long-standing challenge in drug discovery, learning-based molecular dynamics, and protein mechanism analysis. Geometric graph neural networks (Geom-GNNs) with all-atom information have transformed atomistic simulations by serving as a general learnable geometric descriptors for downstream tasks such as interatomic potential and molecular property prediction. However, common practices involve supervising Geom-GNNs on specific downstream tasks, which suffer from the lack of high-quality data and inaccurate labels, potentially leading to poor generalization and performance degradation on out-of-distribution (OOD) scenarios, especially with quantum chemical data. In this work, we explored the possibility of using pre-trained Geom-GNNs as transferable and highly effective geometric descriptors for improved generalization. To explore their representation power, we studied the scaling behaviors of Geom-GNNs under self-supervised pre-training, supervised and unsupervised learning setups. We found the expressive power of different architectures can differ on the pre-training task. Interestingly, Geom-GNNs do not follow the power-law scaling on the pre-training task, and universally lack predictable scaling behavior on the supervised tasks with quantum chemical labels. More importantly, we demonstrate how all-atom graph embedding can be organically combined with other neural architectures to enhance the expressive power. Meanwhile, the low-dimensional projection of the latent space shows excellent agreement with conventional geometrical descriptors.
I certify that there is no URL (e.g., github page) that could be used to find authors’ identity.
I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.
Geometric Graph Neural Networks Self-supervised Pre-training Scaling Zero-shot Transfer Molecular Representation
I certify that there is no acknowledgement section in this submission for double blind review.
/pdf/80c45a708014eea392ed8db0e423c825ba02b4d0.pdf
applications to physical sciences (physics, chemistry, biology, etc.)
I understand the reciprocal reviewing requirement as described on https://iclr.cc/Conferences/2025/CallForPapers. If none of the authors are registered as a reviewer, it may result in a desk rejection at the discretion of the program chairs. To request an exception, please complete this form at https://forms.gle/Huojr6VjkFxiQsUp6.
I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2025/AuthorGuide.
Pushing the Limits of All-Atom Geometric Graph Neural Networks: Pre-Training, Scaling, and Zero-Shot Transfer
ICLR 2025 Conference Submission
ICLR.cc/2025/Conference/Submission