Pushing the Limits of All-Atom Geometric Graph Neural Networks: Pre-Training, Scaling, and Zero-Shot Transfer

Authors: Zihan Pengmei, Zhengyuan Shen, Zichen Wang, Marcus D. Collins, Huzefa Rangwala
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
Keywords: Geometric Graph Neural Networks, Self-supervised Pre-training, Scaling, Zero-shot Transfer, Molecular Representation

Abstract:
The ability to construct transferable descriptors for molecular and biological systems has broad applications in drug discovery, molecular dynamics, and protein analysis. Geometric graph neural networks (Geom-GNNs) utilizing all-atom information have revolutionized atomistic simulations by enabling the prediction of interatomic potentials and molecular properties. Despite these advances, the application of all-atom Geom-GNNs in protein modeling remains limited due to computational constraints. In this work, we first demonstrate the potential of pre-trained Geom-GNNs as zero-shot transfer learners, effectively modeling protein systems with all-atom granularity. Through extensive experimentation to evaluate their expressive power, we characterize the scaling behaviors of Geom-GNNs across self-supervised, supervised, and unsupervised setups. Interestingly, we find that Geom-GNNs deviate from conventional power-law scaling observed in other domains, with no predictable scaling principles for molecular representation learning. Furthermore, we show how pre-trained graph embeddings can be directly used for analysis and synergize with other architectures to enhance expressive power for protein modeling.

Decision: accept (poster)
Reviewer scores: 6, 6, 8, 6

Reviews and discussion:
<thread>
creation date:
1729954335659

thread_type:
official_review

<thread_object>
creator: reviewer_5lxw

<content>
summary:
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."

strengths:
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

weaknesses:
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..."

questions:
- 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?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730535279718

thread_type:
official_review

<thread_object>
creator: reviewer_i3nd

<content>
summary:
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.

strengths:
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.

weaknesses:
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.

questions:
as Weaknesses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730579990907

thread_type:
official_review

<thread_object>
creator: reviewer_kcdu

<content>
summary:
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.

strengths:
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.

weaknesses:
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).

questions:
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?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730645400667

thread_type:
official_review

<thread_object>
creator: reviewer_gnd7

<content>
summary:
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.

strengths:
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.

weaknesses:
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.

questions:
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?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732144449571

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Global Rebuttal

comment:
We sincerely thank the reviewers for their thoughtful and constructive feedback, which has greatly contributed to improving the quality and clarity of our manuscript. Below, we summarize the major updates and changes made in response to the reviewers’ comments:

### Summary of Changes and Key Clarifications:

1. **Research Question and Contributions:**
   - We emphasized our central research question:  
     **"Are pre-trained all-atom geometric graph neural network (GNN) representations transferable to protein modeling, and how expressive are they?"**
   - To address this, we studied the scaling behaviors of state-of-the-art geometric GNNs in unsupervised, self-supervised, and supervised setups and demonstrated their transferability to proteins with all-atom resolution.

2. **Figure Revisions:**
   - Based on the feedback from **Reviewer i3nd**, we revised Figures 4, 6, and 7 for improved clarity by annotating legends with specific model depths and widths where relevant. These updates are reflected in the manuscript.

3. **Restructuring the Introduction and Abstract:**
   - Per suggestions from multiple reviewers, we restructured the introduction and abstract sections to better highlight our research question, contributions, and the relevance of this study.

4. **Additional Experimental Results:**
   - Responding to feedback from **Reviewer i3nd**, we added experimental results comparing the effect of model depth and aspect ratio during fine-tuning on QM9. The results, summarized in **Appendix L, Table 10 (Page 32)**, show a general trend where pre-trained deeper models perform better than shallower ones, aligning with findings in prior literature (e.g., Li et al. 2024).  

5. **Clarifications on Zero-Shot Transfer and Pre-Training Objectives:**
   - We addressed concerns from **Reviewer KCDU** and **Reviewer 5LXW** about the definition of zero-shot transfer learning. Our setup involves pre-training on small molecules and directly transferring embeddings to infer atomistic representations for proteins, which is then combined with separate tasks. We clarified why this qualifies as zero-shot embedding inference, distinct from fine-tuning.
   - We elaborated on the rationale behind selecting coordinate denoising as the pre-training objective for its simplicity and effectiveness.

6. **Benchmarks and General Evaluations:**
   - We addressed concerns about dataset diversity by discussing scaffold splitting in QM9 and incorporating non-equilibrium datasets like xxMD. These ensure rigorous evaluation of generalization capabilities.
   - Additional data on xxMD datasets comparing pre-trained and non-pre-trained models were provided in **Appendix L, Table 9 (Page 31)** to better support claims about model performance across dimensions and molecular systems.

### Final Remarks

We hope these revisions and clarifications address the reviewers' concerns and further demonstrate the novelty and importance of our work. We remain grateful for the reviewers' insightful feedback, which has strengthened our manuscript significantly.

Thank you for your time and consideration.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732252575444

thread_type:
official_comment

<thread_object>
creator: area_chair_8tuf

<content>
comment:
Hi reviewers,

The authors have posted their rebuttals. Could you please check their responses and engage in the discussions? Please also indicate if/how their responses change your opinions.

Thanks,

AC
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper studies if the pre-trained all-atom geometric GNN representations are transferable to protein modeling and their expressivity.
To answer this, the paper studies the scaling behaviors of state-of-the-art geometric GNNs in unsupervised, self-supervised, and supervised setups and demonstrates their transferability to proteins with all-atom resolution. The authors explore many aspects of geometric GNNs, like model size, aspect ratio, nearest neighbor cutoff radius, architecture, and transferability among different data types, which is comprehensive and distinguishes the work from other studies. Overall, I think the work is interesting, novel, and useful to the field of geometric GNNs. Thus, an acceptance is recommended.