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
Venue: ICLR 2025 Poster (2025)
Keywords: Geometric Graph Neural Networks
Self-supervised Pre-training
Scaling
Zero-shot Transfer
Molecular Representation
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
TL;DR: 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.

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.

Published a conference paper at ICLR 2025
PUSHING THE LIMITS OF ALL-ATOM GEOMETRIC GRAPH NEURAL NETWORKS: TRAINRRIN SCALING AND ZERO-
SHOT TRANSFER
Zihan Pengmei*
Zhengyuan Shen Amazon Web Services
Zichen Wang
The University of Chicago zpengmei www.cchccc..... edu
Amazon Web Services
donshen@amazon.c com
2 aichewane amazon..aaaaoooommmm.oommoo.o.. com
Marcus Collins
Huzefa Rangwala
Amazon
Amazon Web Services rhuze saazzefaaaazzo..cmm com
col collmre @ amazon. com
ABSTRACT
The ability to construct transferable descriptors for molecular and biological systems has broad applications in drug discovery, molecular dynamics, and protein analysis. Geomet- ric 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-GNN in protein mod- eling remains limited due to computational constraints. In this work, we first demonstrate the potential of pre-trained Geom-GNNe 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-GNN 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 archi- tectures to enhance expressive power for protein modeling.
INTRODUCTION
In silico molecular computation and simulation are indispensable tool in modern research for biology, chem- istry, and material sciences capable of accelerating the discovery of new drugs and materials, as well as predic- tion of protein structures and functions. One of the key advance in molecular computation is the use of learn- able geometrical descriptors for representing the atomic environment (Todeschini & Consonni 2008 Mauri let al. 2017 Baig al.) 2018 Moriwaki et al., 2018). Geometric Graph Neural Networks (Geom-Giie GNNs) with all-atom resolutions can learn to construct system-specific molecular force fields by incorporating geometric in- ductive biases such as rotational invariance (Schütt et 2018 Gasteiger et al. 2020) and equivariance (Schütt) et al. 2021 Batzner et al. 2022). These ML-based force fields can be seen as learnable geometric descriptors that map the molecular conformation to the potential function of target. These capabilities position Geom- GNNs as powerful tools for mapping molecular conformations to target functions, such as potential energy surfaces. However, most applications remain restricted to quantum chemistry tasks, with limited exploration in protein modeling. And such all-atom information can be crucial in modeling protein-protenn interaction (PPI), protein-ligand binding, and protein reaction kinetics.
In contrast to other domains such text and vision, where pre-trained models have been recognized as highly effective feature extractors trained with self-supervised objectives (Devlin et al., 2019 Radford et al. 2021 Kirillov et al.) )023). the potential of pre- trained Gooss for protein modeling remains underexplored. Recent adnances such as denoising pre- traitrainng, have shown promise in enhancing the generalization of Geom-
Work dome during an internship at Amazon Weh Service.

---

Published conference paper at ICLR 2025
GNNs across diverse molecular systems (Zaidi et al. 2022). Building on these developments, we first demon- strate that pre-trained Geom-GNNs serve as effective zero-shot learners for protein systems, extracting rich, transferable representations when pretrained on small-molecule datasets, thereby addressing the chal- lenges discussed earlier. We evaluate the performance of these pre-trained models by analyzing neural scaling laws (Kaplan et al. 2020 Hoffmann et al., 20022). Specifically, we studied pre-training graph embeddings via characterizing the graph neural scaling laws in the self-supervised learning pre-training, supervised regression and unsupervised feature extraction setups.
Our results reveal that, contrary the power-law scaling observed in other domains, Geom-GNNs ex- hibit unique scaling patterns characterized by early saturation and diminishing returns. Furthermore, we explore how pre-trained all-atom graph embeddings can be integrated with other neural architectures to improve their expressive power in complex tasks such as protein folding and kinetic modeling, We additionally show how those pre-trained graph em- beddings could be easily combined with other archi- tectures to enhance their expressive power for pro- tein modeling. The low-dimensional projection of the pre-trained graph embedding space (Figure and Appendix K shows meaningful clustering of similar conformations. These findings not only consolidate our understanding of Geom-GNN scaling principles but also pave the way for their broader application in molecular and protein sciences. Our contributions can be summarized as follows:
Pre-trateed Geom-GNN
Token Mixing (SAMMPP/PP)
Pooling
All residues tokens
Task Hoad
Figure aeta-architecture for Using Pre-trained Geom-GNNs as descriptors: The figure shows framework where pre-trained Geom-GNNs act as lo- cal geometric descriptors featurize residue-level con- formations. Each window represents an atomic envi- ronment for residue feature extraction, defined by user-defined context (illustrated here as a sliding win- dow of nearest neighbors in sequence). In each win- dow, atomic structures are treated as individual graphs and processed by the pre-trained Geom-GNN tract atomic-level features, which are aggregated into residue-level representations or "tokens." " The archi- tecture can employ self-attention (SA), multi-layer per- ceptron (MLP), or message passing mechanisms to en- hance representational power. For graph-level tasks, the mixed tokens are pooled and input to a task-specific head for training and predictions.
We first demonstrate the use of pre-trained Geom- GNNs as zero-shot transfer learners for protein fcaturization when pre-trained on small-molccule datasets. When representing residue conformations in proteins with nearest neighbors, graph embed- dings match conventional rotamer features but pos- sess better expressivity in describing the atomic en- vironment flexibly. The pre-trained graph embed- dings excel at constructing kinetic models through the VAMPNet objective: a challenging unsuper- vised learning task that is crucial for enhanced sam- pling and probing energy landscapes in computa- tional chemistry and biology. Compared to tradi- tional pairwise distance features, the graph embeddings achieve significant performance gains up to 50% across complex peptide and protein systems.
ex-
We systematically investigated the scaling behaviors of Geom-GNNs under the self-supervised, supervised and unsupervised setups. We found that Geom-GNNY do not exhibit typical power-law scaling during self- supervised pre-training. suggesting distinct scaling principles for molecular representations compared to areas like language modeling. Under more rigorous setups requiring out-o--distribution (OOD) generalization, we found there is no predictable scaling behavior on molecular property and force field prediction tasks, in contrast to earlier discoveries. In the transfer learning setups for protein modeling, also found there is no predictable scaling behavior in kinetic model construction and fold classification.
2 RELATED WORKS
Pretraining via Denoising a novel pretraining strategy for equivariant GNN architectures, drawing inspira-
tion from advancements in visual generative models (Sohl-Dickstein et al. 2015 Song et al. 2020a b). Zaidi
al. (2022) introduced a method of corrupting input atomic coordinates of equilibrium molecular structures
2

---

Published a conference paper at ICLR 2025
and training the equivariant GNN to denoise and restore the original coordinates. The denoising objective is defined as: Denoising E, l|e [F Here, represents the noise added to the original coordinates r, producing the noised input f. Wang et al. 2023) expanded this method to include denoising of non equiliiiuuuum structures. Further studies have investigated varying noise distributions such as sliced noise (Ni et al. 2023).
Graph neural scaling laws represent efforts to extend scaling laws from language modeling to the graph domain. Frey al.) (023) demonstrated the power-law relationship between supervised loss, model size N, and dataset size D when training a Geom-GNN as a potential function. [Chen et al. (022) further advanced this area by demonstrating a predictable power-law relationship between supervised loss and dataset size using a topological GNN. More recently, [Lu et al. (2024) expanded the discussion to parametric perspective, focusing
on excessively large datasets such as PCQM4M (Nakata & Shimazaki 2017
Hu al. 2020) in a supervised
setting without pre- pre-training.
Learning- based molecular dynamics In molecular simulations, addressing complex tasks beyond basic molecular property predictions involves techniques such as kinetic modeling using time-lagged independent
component analysis (TICA) (Molgedey& & Schuster] [994| Naritomi Fuchigami] 2011 [Scherer et al.] 2015)
and Variational Approach for Markov Process (VAMP) (Wu & Noē, 2020; Mardt et al. 2018), prediction of
committer functions Strahan et al.] (2023a/b) by in transition path theory (Weigend, 2006 Metzner et al.) 2009).
and the application of enhanced sampling methods (Bonati et al. 202T; Jung et al.) 2023). These tasks require meticulously crafted descriptors that capture the structural nuances of the systems under study, Recent ad- vancements have underscored the value of atomistic detail, leading to the integration of learnable Geom-
KKlein et al., 2024 Huang et al., [2244] Liu et al. 2023) into aforementioned more advanced modeling sce- narios. However, the scalability of these methods remains challenging due to the necessity of training graph representations on-the-fly.
3 EXPERIMENTAL AND DATASET SETUP
We focus on two otate-of-the-art architectures based on message-passing (MP): Equivariant Transformer (ET) (Thölke & De Fabritiis 2022) and its extension, ViSNet (Wang et al.. 2024), which incorporates higher body- order features. To investigate scaling effects, we separately studied effect of width, depth and aspect ratio of MP layers. We further analyze the effect of radius cutoff distance and refer readers to Appendix [F for detailed model configurations and Table D for summary of datasets used.
Pre-training setup. For pre-training, we employed a vanilla denoising objective, where scaled noise was added directly the input atomic coordinates, and the network was tasked with predicting this added noise. After pre-training, the denoising head was discarded. We utilized two distinct datasets for pre-training: PCQM4Mv2
(Nakata & Shimazaki 2017; Hu al.] (020)) and the training dataset of OrbNet-Denali (Christensen et al., 2021) (hereafter referred to as PCQM and Denali). These datasets represent equilibrium and non-equilibrium molecular structures, respectively. In PCQM. each sample represents a unique molecular topology, with struc- tures relaxed using DFT. thus providing equilibrium structures. Conversely, Denali contains multiple geometric conformations for each molecular topology, offering diverse range of non-equilibrium structures.
Supervised downstream tasks setup. The supervised downstream tasks to validate the learned representa- tions encompass diverse molecular structures and conformations. For chemical structure diversity, we utilized the QM9 dataset (Ramakrishnan et al., 2014) for molecular property prediction, implementing an ID setup through random splitting (Random-Qul QM9) and an OOD setup through scaffold splitting based on molecular motifs using RDKit (Landrum et al.) 2013) (Scaffold-QM9) For conformational diversity, the xxMD dataset (Pengmei et al., 2024a) was employed, which surpasses the MD17 benchmark (Chmicla et al. 2017) in sam- pling and computational techniques. This dataset features extensively sampled non-equilibrium conformations within reactive regions. We applied a temporal split caamporal-xxMID based on timesteps and a random split (Random-xxMD) as ID and OOD scenarios. Here, EDD represents well-sampled condition, ensur- ing even representation of conformational space across training and testing sets. Conversely, Temporal-xxMD presents a more challenging scenario with less time correlation and poorer sampling, demanding significant extrapolation capabilities.
Zero-shot transfer setup. Beyond conventional supervised learning. we explored zero-shot transfer ap- proaches to extract pre-trained graph representations, following the methodology outlined in/Mardt et al.] (2018) and Pengmei et al. (2024b). For kinetic modeling tasks for characterizing non-reversible dynamics, we utilized
3

---

Published conference paper ICLR 2025
molecular dynamics trajectories from three systems: alanine dipeptide (ala2), pentapeptides from MDshare (Nüske et al., 2017] Wehmeyer & Noë 2018), and X6-85 fast-folding protein (Bowman et al. 2011). We in- ferred and stored activations from the last MP layer of pre-trained all-atom Geom-GNN checkpoints. These features are used to optimize the network shown in Figure [ with the VAMP score. To further assess the versatility of pre-trained representations, we investigated their potential to enhance other GNN architectures. Specifically, we augmented previously developed protein-specific GNN (Wang al. (022) with activations from the pre-trained all-atom Geom-GNN. We then evaluated this enhanced model on a protein fold classifica- tion task (Hou et al. 2018).
4 ZERO-SHOT TRANSFER DOWNSTREAM TASKS
Comparing supervised objectives, self-supervised objectives are emerging areas of researches learning-based molecular dynamics simulations. As mentioned before, training all-atom graph representations on-the-fly subject to scalability and generalization issues. We propose to address these challenges by pre-trained graph embeddings including atom-level information, thereby separating the training of Geom-GNNs from the sim- ulation objectíves. This section is structured as follows: We begin by showcasing the application of zero- shot transferable Geom-Gi Nroo-GNNNTT representations, integrated with other architectures to meet the VAMP objectives (Mardt et al. 2018), which aims to learn a low-dimensional representation that preserves the singttimsccale kinetics of the molecular system.
crucial in computational biology for enhanced sampling and under- standing the underlying energy landscapes. We explored three systems with increased complexity up to 1258 heavy atoms: ala2, pentapeptide, and the A6-85. Subsequently, we demonstrate how these features can com- bined with protein-specifee GNNs improve the classification performance protein folding, with analysis presented in Appendix for brevity.
4.1
KINETIC MODELING: VAMP
-1.0
-34
(degree)
150
155
150
-100
-150
-1
-1.0
100
Phil (s)
Figure 2: Visualization of learned singular vectors of Koopman operator onto the v + dihedral angle space of ala2 using pre-trained ViSNet embedding with 64 width. Those singular vectors describe the slow modes of the underlying dynamics (Appendix
Initial experiments focused on the ala2 system, known for its metastable states characterized by two main backbone dihedrals, A and Φ. We input the atomic numbers and Cartesian coordinates of ten heavy atoms into the pre-trained ViSNet and extracted graph-level embedding. A separate VAMP head was then trained to learn the kinetic model. The learned coordinates were projected onto a two-dimensional space constructed by V & dihedrals. As illustrated Figure 2 the network successfully identified the first learned singular vector as the slowest transition between positive and negative V angles. The subsequent vector characterized the transition between small and large & angles. Lower-order singular vectors identified transition states between the and

---

Published a conference paper at ICLR 2025
regions of the Ramachandran plot. These observations show that the graph representations capture structurally accurate and semantically meaningful information illustrated in Figure 14 and 17
Koopman matrix quantifies the kinetic variance between instantaneous and time-lagged data (Appendix J.. A minimal VAMP score of I corresponds to equilibrium, while the maximum score is theoretically determined by the dimensionality of the Koopman operator constructed, with the highest possible score being the output dimension plus one. However, the actual VAMP score varies depending on the characteristics of the molecular trajectories and the selected lag time. Figure S shows the VAMP scores for models trained on pentapeptide tra- jectories across various output dimensions and lag times, following the PyEMMA setup Scherer et al. (2015). Furthermore, to interpret the learned coordinates, we visualized characteristic structures and their projections onto conventional physical collective variables in Figures 14 13 and [[(pppendix) These visualizations facil- itate direct comparisons with established reference literature (Scherer et al., 2015) Bowman et al.) 2011). Our observations reveal that increasing the output dimension to 10, across all five tested lag times, does not plateau in performance when compared to the linear TICA method, which utilizes backbone torsion angles. This demonstrates that using graph features allow for the higher resolution comparing to conventional descriptors
Table 1: Validation of VAMP-2 scores (1) for Ala2, Pentapeptide, and X6-85 with PCQM-V with varying dimensions and model architectures, with constant batch size of 5000, Scores reflect the performance with and without ('Sum' refers to non-learnable sum. pooling) residue-level token mixing in graph embedding vectors wwacooarr--ggannnng
Lag times: Ala2 ps, Pentapeptide at 0.5 ns, and X6-85 at 25 ns. Except Ala2, all systems are split by trajectories, where half of trajectories are reserved for validation. 'OOM' refers to out-of-memory at this combination of dimension and batch size an instance with 4 NVIDIA A10G GPUs. We reproduced Mardt et al. (2018) with the same setup, for ala2 and pentapeptide we consistently used the heavy atom pairwise distances, and we used pairwise distances of 80 C. atoms for A6-85.
Dimension
Ala2
Pentapeptide
X6-85
Token Mixing
Sum
Sum
MLP
SA
Sum
MLP
SA
Out. Dim.
6
10 44 (5 residues)
10
Heavy Atoms
10
1258 (80 residues)
64
4.71=0.01
5.63+0.37
7.00=0.47
4.81+1.52
4.27=0.06
8.49#0.24
0.01
128
4.72+0.01
6.41#0.07
6.91=0.32
7.11+0.43
5.17=0.08
8.35+00.3
7.80+0. 18
256
4.70=0.01
6. 6.10+0.22
7.67=0.08
7.04+0.21
4.75=00.26
8.52+0.14
OOM
384
4.70+0.01
6.6240.16
7.71+0.23
7,2640.11
5.8140.18
OOM
OOM
Mardt et al.
3.92+0.47
5.14+0.25
7.40=0.40
In terms of scaling feature dimensions, We observed no significant performance gains for Ala2 when increasing the feature dimension beyond 64, as the model effectively captured critical slow dynamics at this lower dimension. However, for more complex systems like Pentapeptide and X6-85, which possess greater number of degrees of freedom, scaling the embedding dimension provided clear performance improvements. This indicates that the effectiveness of scaling is contingent upon the complexity of the system and whether the network has reached its representational limit. Other than scaling, another key factor affecting performance the use of token mixing techniques, where residue-level embeddings were generated from pre-trained all- atom graph embeddings and subsequently mixed using higher-level operators such as the MLP-Mixer and self-attention mechanisms (Pengmei et al. 2024b). As shown in Table I this approach led significant performance improvements, particularly for X6-85, where applying the MLP-Mixer resulted in over 46% increase in VAMP-2 scores compared to the baseline sum pooling method.
In comparison to the baseline VAMPnet model (Mardt et al.) 2018), our approach using pre-trained graph features led to performance improvements of 20.4%, 50.0%, and 15.1 % for Ala2, Pentapeptide, and X6-85, respectively. In smaller systems such as the Pentapeptide, early-stage information compression can occur, and token mixing methods alleviate this limitation. For larger systems like X6-85, the improvements gained from pre-trained features were not substantial as for smaller systems. This indicates that future work should focus on developing more expressive models capable of generalizing global relationships from local tokens, as well
hctp:// http///ww.mmmmm///////
ide- showease.html

---

Published conference paper ICLR 2025
as refining the process of localizing graph embeddings for each residue protein systems along molecular trajectories. Current methods, which treat the entire protein single graph, may suffer from over-smoothing, leading to the mixing of global and local features at each node. By addressing these challenges, more effective residue-level token mixing and introducing global features could be achieved to improve overall performance. In principle, this approach can be extended to study phase diagrams and free energy surfaces of polymeric and other complex chemical systems but we limit scope to biomolecules due the lack of reference.
4.2 PROTEIN REPRESENTATION LEARNING: FOLD CLASSIFICATION
As we explored the usage of pre-trained graph embedding study the conformation space peptides and proteins, we extend to explore the utility of such all-atom embed- ding the protein representational learning. Depicted in Figure [ modeling the dependence among multiple inputs require an active token mixer. MP-based GNNs themselves ves are also type of token mixers. In protein-specific GNNs, residues are represented as nodes, and the features attached to each node are usually Ca positions and a set of geometri- cal descriptors for backbone or side-chain dihedral and tor- sion angles ("rotamers"). While these features are effective for describing secondary structures, they are inadequate for describing protein tertiary structures, meaning the lack of direct modeling of interactions between side-chains, partic- ularly in environments with x-stacking between aromatic residues. More importantly, torsion angles do not capture hydrogen bonding between distant residues, salt bridges or disulfide bonds, which are vital for the tertiary structure identification. The direct modeling of the possible side- chain with solvent, membrane and ligand interaction is also missing, which can be addressed by Geom-GNN represen- tation. For illustration, of pre-inferred all-atom embedding to interaction blocks of Wang et al. (2022) and test the downstream task performance the FOLD dataset (Hou et al.) 2018). Per Figure 20) we observe improvement of integrating such all-atom features into the ProNet without angle information. More importantly, the pre-trained graph em- bedding can match with conventional rotamer features in non-interacting cases shown in Figure 17 18 and
Lag=0,1ns
Lag=0.2ns
Validation VAMP-2 Score
Lag=0.5ms
Lag=1.0ns
lag=2.0ns
2-
10
Output Dimension
Figure 3: Validation VAMP-2 scores using pre- trained ViSNet embeddings with 256 length (no token mixer) across various output dimensions and lag times for the pentapeptide system. Ran- dom half of available trajectories held for val- idation and the remaining are used for training.
naively concatenate projections
19 where we used sliding window on the backbone sequence. And it remains future efforts to curate such benchmarks that test how side-chains interact with environments.
5 SELF-SUPERVISED PRE-TRAINING TASKS
As we have studied a few preliminary application of pre-trained graph embedding, we wonder if the power of features given by Geom-GNNs with varying architectures and configurations. We approach this by first varying the network width with constant depth and cutoff radius, followed by separate assessments of depth and cutoff radius impacts. We also analyzed the effects of aspect ratio while maintaining a consistent total parameter count considering the specific challenges of MP-based architectures, including under-reaching and over-smoothing.
Effeet of model width. We investigated the scaling behavior ET and ViSNet on both PCQM and Denali datasets. For the denoising pre- pre-training task, measured performance using the mean square error (MSE) between the predicted and normalized noise. We varied model width and allowing each model to train 10 epochs. We report both the first epoch loss (where each data point is seen once) and the converged loss as parameter-limited and data-limited scenarios. Figure 6 illustrates findings, revealing: (1) Deviation from power-law: The scaling relationship between pre-training loss and model size does strictly follow a power- law. (2) Rapid initial improvement: In low parameter regimes, model performance improves quickly as size increases. (3) Saturation effect: Beyond certain threshold (c.g., after 10" parameters), performance gains

---

Published a conference paper at ICLR 2025
PCOM Converged
Denall Converged
10
Layers
23.
14
210
Loss
1."
126
1.84
432
1.44
432
64
360
2444
3GC
160
250
125
180
2x100
376
304
10
312
512
200
101
10'
Number of Parameters
Number of Parameters
Figure 4: Scaling behavior of ViSNet depth for the converged results of training on PCQM and Denali datasets. The plot shows the relationship between the number of layers and the pre-training loss, illustrating the initial rapid improvement followed by diminishing returns as depth increases. We additionally show the results of the first epoch loss in Figure 7 (Appendix).
become marginal, with models often reaching point where further increases in model size yield diminishing returns, deviating from the fitted power-law line.
Effect of Model Depth. The number of MP layers determines the receptive field size, as each layer gathers information from adjacent nodes, enabling data propagation from distant nodes. In experiments, we varied the network's 's width and depth to assess their impacts on pre-training loss, as depicted in Figures 4 and 7 Our results show significant performance gains with increased depth the initial layers, with diminishing returns beyond 6 layers. Notably, scaling behaviors differ between the PCQM and Denali datasets. The PCQM dataset shows saturation in both initial and converged losses, whereas in the Denali dataset, deeper and narrower mod- els tend to surpass wider and shallower ones, likely due to the larger average graph sizes (PCQM: 29.5 nodes, Denali: 45.0 nodes). Larger graphs benefit from more MP steps, but the over-smoothing issue inherent in MP- based GNNs limits the advantages of greater depth. With deeper networks, the influence of depth diminishes, and the model's performance becomes more reliant on increasing parameter counts. This shift suggests that optimal depth achieves maximum performance before capacity expansion becomes more effective.
Effect of aspect ratio. To further discern the effect of scaling the model depth, we fix the total parameter count by varying the model width to control the aspect ratio. By eliminating the parameterization effect, observe the checkpoints with similar amount of parameters and lower aspect ratio (width/depth) enjoy better pre-training performance per depicted in Figure (Appendix). In other words, deeper models are favored, though such improvement diminished as the depth is saturated.
Effect of Cutoff Radius. The cutoff radius in Geom-GNN: critically defines the model's effective receptive field by regulating the scope of atomic interactions during message passing. Figure ( nificant performance fluctuations at different radii. On the PCQM and Denali datasets, the pre-training loss for a 6-layer ViSNet model displays a U-shaped curve, decreasing to minimum at approximately 5 Å be- fore rising again up to 9 Á. This pattern indicates an optimal cutoff radius that effectively integrates local and global information, with more pronounced effects in the larger, more complex Denali dataset. In this dataset, losses increase sharply beyond the optimal radius, indicating the importance of distinguishing relevant local interactions from misleading long-range correlations.
(Appendix) shows sig-
SUPERVISED DOWNSTREAM TASKS
As described in the experimental setup, we additionally used two molecular benchmarks to verify the perfor- mance of learned features: QM9 for molecular properties prediction (chemical diversity) and xxMD for force
7

---

Published conference paper ICLR 2025
field prediction (conformational diversity). Our analysis covers two scenarios: transferring weights from pre- trained model and fine-tuning all parameters, and training models with identical configurations from scratch.
Table 2: Performance of ViSNet and ET models on the Scaffol--MM. dataset with varying model widths. The table shows the mean absolute error (MAE) vs model width, comparing pre-trained (fine-tuned) models using PCQM dataset with those trained from scratch. All models are trained without the auxiliary denoising loss.
Model
Dim
μ (mD)
a (mag)
chооo (meV)
CLUMO (mcV)
Ae (mc
64
28.8
56.8
37.2
90.2
128
24.9
93.4
56.8
35.9
80.7
PT-ViSNet
256
25.1
110.8
51.7
32.3
77.8
384
23.2
91.2
47.1
34.8
77.0
512
26.5
90.5
54,9
35.5
71.0
64
36.5
142.1
74.4
46.6
110.5
128
32.3
106.4
62.0
42.8
94.5
ViSNet
256
36.3
101.5
55.0
41.8
92.5
384
34.3
115.2
60.9
44.2
96.2
512
35.9
139.9
60.0
52.5
123.6
64
36.9
108.3
69.0
45.8
96.3
128
33.5
117.4
62.1
40.2
88.4
PT-ET
256
29.2
116.4
54.0
36.4
84.0
384
101.3
52.3
40.1
82.0
512
29.4
197.8
57.7
37.5
80.1
64
35.9
104.2
72.7
47.5
101.9
128
32.3
98.4
65.4
44.6
96.2
ET
256
31.3
127.5
65.1
43.2
100.9
384
34.8
116.6
60.6
44.0
98.7
512
53.8
178.7
83.7
53.3
97.7
6.1 CHEMICAL DIVERSITY: QM9
We evaluated the performance of both fine-tuned and from-scratch trained ViSNet and ET models across different capacities, as depicted in 2a [0] (Appendix). We assessed model performance
Effeet of Scaling and Pre-training:
hold-out test set, leading to the following observations. Firstly, noted limited scaling behavior; while moderately larger models generally outperform smaller ones, the performance im- provements do predicatively scale with model size, indicating manifestation of generalization error due to model over-parameterization Secondly, models initialized with pre-trained weights consistently outperform their counterparts trained from scratch across both architectures and all model sizes, highlighting the value of the pre-training approach for downstream tasks. Lastly, regardless of the splitting strategy and training approach, ViSNet demonstrates superior performance compared ET, aligning with the pre-training results shown in Figures 6 and 7 (Appendix).
Effect of Data Split and Label Uncertainty. Figure IT (Appendix) depicts the results from the Random QM9 experiment. Typically, wider models exhibited improvements in Random-QN QM9,, whereas scaffold-spli mod- els struggled, highlighting challenges in domain generalization. Increased model parameters might suggest overfitting rather than true generalization enhancement. Pre-trained models, exposed to a broader molecular structure range in the PCQM dataset (including QM9 molecules), showed superior transfcrability to down- stream tasks, reflected in the performance disparity between pre-trained and from-scratch models discussed earlier. Additionally, evaluation errors are significantly attributed to label uncertainty. For instance, variations in HOMO-LUMO gap predictions among common DFT methods for three QM9 molecules, detailed in Ap- pendix [ and Zhhao & Truhlar (2005); ZZhang Musgrave (2007); MMrdirossian & Head- Gordonn (2017), are markedly greater than the prediction errors, illustrating the inherent uncertainty in training and evaluation la- bels. And prediction error of B3LYP functional (Lee et al.) 1988) used in QM9 dataset are usually at eV-scale. This error in "ground truth" complicates model performance assessment and generalization interpretation, sug- gesting that training data uncertainty can largely contribute to the observed limitations in model scaling and OOD performance, beyond model limitations.

---

Published conference paper ICLR 2025
6.2 CONFORMATIONAL DIVERSITY: XXMD
Building upon analysis of chemical diversity, now focus conformational diversity within single stilbene molecule using the Random-xXMD. We begin our analysis by controlling the number of samples used to supervise the network. We prepared models with varying widths and training strategies examine dataset and parameter scaling effects, as well as the transferability of embeddings learned from equilibrium and non-equilibrium structures. As illustrated in Figures 12] (Appendix), Our findings highlight three key trends: (1) Pre-training shows diminishing returns beyond a certain dataset size, suggesting nuanced balance between pre-training depth and dataset scale. (2) In data-rich scenarios, larger models consistently outperform. (3) In data-limited situations, only models pre-trained on the equilibrium PCQM dataset effectively enhance performance. These observations echo the phenomenon of "parameter ossification" mentioned in pre-training literature (Hernandez et al. 2021) 1), where pre-trained models exhibit early performance saturation. [Cen et al. (2024) discussed this issue for GIN models pretrained on molecular datasets at the 10 sample scale. Our findings extend this concept Geom- GNNs, demonstrating that ossification can occur even at the 10 sample scale. Notably, we did not observe ossification effects the QM9 dataset, despite its orders-ot soaerssof-mmme aof---manitdee-eeee samples, suggesting that ossification is related to both the number and the correlation among samples.
Performance of models with varying width and training strategy xxMD xxMD-Temporal subsets. Best models are picked based on the force regression in terms MAE, Only models pre-trained on Denali dataset can positively transfer, while ET receives more benefits from pre-training comparing ViSNet.
Table 3:
Subset
Best Dim
Model
Energy (meV)
Force (/A)
ViSNet
Stilbene
128
PT-Denali
321
137
Azobenzene
256
No PT
75
72
Malonaldehyde
256
PT--eenlli
95
142
Dithiopehene
256
No PT
87
55
ET
Stilbene
384
PT-Denali
348
142
Azobenzene
64
PT-Denali
127
98
Malonaldehyde
384
PT-Denali
144
157
Dithiopehene
64
PT-Denali
78
Next, we examined the performance on the Temporal-xxMD dataset, using all four systems: azobenzene, stilbene, malonaldehyde, and dithiophene. Figures 13 (Appendix) illustrates the results for energy and force regressions. The trends observed in this experiment differ notably from our earlier findings: (1) pre-trained embeddings from the Denali dataset demonstrate superior transfer compared to those from PCQM. Table suggests that exposure to non-equilibrium conformations during pre-training enhances the model's 's ability to generalize to the more diverse conformational space in the xxMD dataset. When both the pre-training and downstream datasets lack coverage of non-equilibrium conformations, the network struggles to extrapolate effectively. However, networks that have learned non-equilibrium representations from the Denali dataset show improved transfer to reactive, highly non-equilibrium conformations. (2) we observe that pre-training does not universally improve performance. Specifically, while ET models consistently benefit from pre-training on either Denali or PCQM datasets, ViSNet models do not always show positive transfer. Greater benefit from pre-training for ET is attributed to its lower baseline performance, allowing more room for improvement since it has not yet reached the underlying DFT uncertainty.
7 CONCLUSION
In short, this work explored the application of Geom-GNNs from a new perspective of represenation learning where pre-trained all-atom Geom-GNNs are effective zero-shot transfer learners that can faithfully describe complicated protein environment with all-atom granularity. To study the expressive power of Geom-GNNs with varying architectures and configurations, we characterize the scaling behaviors of Geom-GNNs on self- supervised, supervised and unsupervised tasks. Deviating from common neural scaling power-laws, Geom- GNNs saturate early and are bounded by issues such as under-reaching and over-smo. Though there are

---

Published a conference paper ICLR 2025
benefits of scaling on supervised and unsupervised tasks, other aspects become more crucial, such as address- ing data label uncertainty and crafting optimal architectures for protein modeling with pre-trained all-atom embeddings. We hope our work can inspire to rethink how should Geom-GNNs be trained and applied on a broader spectrum of problems. Limited by the scope of the paper, only explored denoising pre-training objective, while future efforts include expanding the pre-training scope to include molecular topology and ex- ploring co-pre-traiiee preettrined models for applications such conditional molecule optimization, protein design, and mrotein-protcio interaction.
REFERENCES
Carlo Adamo and Vincenzo Barone. Toward reliable density functional methods without adjustable parameters: The pbe0 model. The Journal of chemical physics, 110(13)x118-61100 1999.
Mohammad H Baig. Khurshid Ahmad, Gulam Rabbani, Mohd Danishuddin, and Inho Choi. Computer aided drug design and its application to the development of potential drugs for neurodegenerative disorders. Cur- rent neuropharmacology. 16(6):740-748, 2018.
Ilyes Batatia, David P Kovacs, Gregor Simm, Christoph Ortner, and Gábor Csányi Mace: ariant message passing neural networks for fast and accurate force fields. Advances in Neural Information Processing Systems, 35:1 1 1436, 2022.
Higher order equiv-
Simon Batzner, Albert Musaelian, Lixin Sun, Mario Geiger, Jonathan P Mailoa, Mordechai Kombluth, Nicola Molinari, Tess E Smidt, and Boris Kozinsky. E -3ququurarra graph neural networks for data-efficient and 2022.
accurate interatomic potentials. Nature communications,
Luigi Bonati, GiovanniMaria Piccini, and Michele Parrinello. Deep learning the slow modes for rare events sampling. Proceedings of the National Academy of Sciences, 1 118(44):e21 13533118, 2021.
Gregory Bowman, Vincent Voelz, and Vijay S Pande. Atomistic folding simulations of the five-helix bundle protein X6- 85. Journal of the American Chemical Society, 133(4):664 667, 2011.
Dingshuo Chen, Yanqiao Zhu, Jieyu Zhang, Yuanqi Du Zhixun Li, Qiang Liu, Shu Wu, and Liang Wang- Uncovering neural scaling laws molecular representation learning. Advances Neural Information Pro- cessing Systems, 36, 2024.
Stefan Chmiela, Alexandre Tkatchenko, Huziel E Sauceda, Igor Poltavsky, Kristof T Schütt, and Klaus-Robert Müller. Machine learning of accurate berggy-conservnng molecular force fields. Science advances, 3(5): e1603015, 2017.
Anders Christensen, Sai Krishna Sirumalla, Zhuoran Qiao, Michael B O'Connor, Daniel GA Smith, Feizhi Ding, Peter J Bygrave, Animashree Anandkumar, Matthew Welborn, Frederick R Manby, al. Orbnet denali: A machine learning potential for biological and organic chemistry with semi-empirical cost and dit accuracy. The Journal of Chemical Physics, 155(20), 2021.
Jacob Devlin, Ming-Wei Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. 2019. URL
wbtps///..g//b///0...o. 04805
RHWJ Ditchfield, Warren J Hehre, and John A Pople. Self-consistent molecular-orbice methods. ix. an ex- tended gaussian-type basis for molecular-orbite studies of organic molecules. The Journal of Chemical Physics, 54(2):724-728, 1971.
Thom H Dunning Jr. Gaussian basis sets for use in correlated molecular calculations. i. the atoms boron through neon and hydrogen. The Journal of chemical physics, 90(2):1007-1023, 1989.
Nathan Frey, Ryan Soklaski, Simon Axelrod, Siddharth Samsi, Rafael Hoomez-oombrrelli Connor W Coley, and Vijay Gadepally. Neural scaling of deep chemical models. Nature Machine Intelligence, 5(11):1297- 1305, 2023.
Johannes Gasteiger, Janek GroB, and Stephan Günnemann. Directional message passing for molecular graphs arXiv preprint XXXx:2003.001233 2020.
10

---

Published conference paper ICLR 2025
Lars Goerigk and Stefan Grimme. A thorough benchmark of density functional methods for general main group thermochemistry, kinetics, and noncovalent interactions. Physical Chemistry Chemical Physics, 13 (14):6670-6688, 2011.
S Yu Haoyu, Xiao He, Shaohong L Li, and Donald G Truhlar. Mn15: A kohn-sham global-hybrid exchange- correlation density functional with broad accuracy for multi-reference and single-reference systems and noncovalent interactions. Chemical science, 7(8):5032-5051, 2016.
Pedro Hermosilla, Marco Schäfer, Matëj Lang, Gloria Fackelmann, Pere Pau Vázquez, Barbora Kozlíková, Michael Krone, Tobias Ritschel, and Timo Ropinski. Intrinsic-extrinsic convolution and pooling for learning on 3d protein structures. International Conference on Learning Representations, 2021.
Danny Hernandez, Jared Kaplan, Tom Henighan, and Sam McCandlish. Scaling laws for transfer. arXiv preprint arXiv:2102.01299 2021.
Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, 9():1733- 1997.
Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, et al. Training compute-optimal large language models. arXiv preprint arXiv:2203. 2022.
Jie Hou, Badri Adhikari, and Jianlin Cheng. Deepsf: deep convolutional neural network for mapping protein sequences to folds. Bioinformatics, 349)):1255-1333, 2018
Weihua Hu, Matthias Fey, Marinka Zitnik, Yuxiao Dong, Hongyu Ren, Bowen Liu, Michele Catasta, and Jure Leskovec. Open graph benchmark: Datasets for machine learning on graphs. Advances in neural information processing systems, 33:22118 10-21132233 2020.
Ying Huang, Huiling Zhang, Zhenli Lin, Yanjie Wei, and Wenhui Xi. Revgraphvamp: A protein molecular simulation analysis model combining graph convolutional neural networks and physical constraints. bioRxiv, pp. 2024-03, 2024.
Hendrik Jung, Roberto Covino, A Arjun, Christian Leitold, Christoph Dellago, Peter Bolhuis, and Gerhard Hummer. Machine-guided path sampling to discover mechanisms of molecular self-organization. Nature Computational Science, 3(4):334-344 34544-35,2220 2023.
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint arXiv:2001.08361 2020.
Alexander Kirillov, Eric Mintun, Nikhila Ravi, Hanzi Mao, Chloe Rolland, Laura Gustafson, Tete Xiao, Spencer Whitehead, Alexander C Berg, Wan-Yen Lo, et al. Segment anything. In Proceedings of the IEEF/CVF International Conference on Computer Vision, pp. 4015-4026, 2023.
Leon Klein, Andrew Foong, Tor Fjelde, Bruno Mlodozeniec, Mare Brockschmidt, Sebastian Nowozin, Frank Noé, and Ryota Tomioka. Timewarp: Transferable acceleration of molccular dynamics by learning time- coarsened dynamics. Advances in Neural Information Processing Systems, 36, 2024.
Greg Landrum et al. Rdkit: A software suite for cheminformatics, computational chemistry, and predictive modeling. Greg Landrum, 8311.10):5211, 2013.
Chengteh Lee, Weitao Yang, and Robert G Parr. Development of the colle-salvetti correlation-energy formula into functional of the electron density. Physical review B, (7():7855 1988.
Yunyang Li, Yusong Wang, Lin Huang, Han Yang, Xinran Wei, Jia Zhang, Tong Wang, Zun Wang, Bin Shao, and Tie-Ye Yan Liu. Long-short-range message-passing A physics-informed framework capture non-local interaction for scalable molecular dynamics simulation. In The Twelfth International Conference on Learning Representations, 2024. URL httpp:://pperreview net formmm/rrmmm/m 2tt

---

Published a conference paper at ICLR 2025
Bojun Liu, Mingyi Xue, Yunrui Qiu, Kirill A Konovalov, Michael S O' Connor, and Xuhui Huang. Graphvamp- nets for uncovering slow collective variables of self-assembly dynamics. The Journal of Chemical Physics, 159(9), 2023.
Jingzhe Liu, Haitao Mao, Zhikai Chen, Tong Zhao, Neil Shah, and Jiliang Tang. Neural scaling laws on graphs. arXiv preprint arXiv:2402.0200 2024.
Shengchao Liu, Hongyu Guo, and Jian Tang. Molecular geometry pretraining with se (3)-invariant denoising distance matching. arXiv preprint arXiv:2206. 13602, 2022.
Narbe Mardirossian and Martin Head-Gordon. Thirty years of density functional theory in computational chemistry: an overview and extensive assessment of 200 density functionals. Molecular physics, 115(19): 2315-2372, 2017.
Andreas Mardt, Luca Pasquali, Hao Wu, and Frank Noé. Vampnets for deep learning of molecular kinetics. Nature communications, 9(1):5, 2018.
Andrea Mauri, Viviana Consonni, Roberto Todeschini, et al. Molecular descriptors. In Handbook of computa- tional chemistry, pp. 2065-2093. Springer International Publishing, 2017.
Philipp Metzner, Christof Schütte, and Eric Canden-Eijnded Transition path theory for markov jump processes. Multiscale Modeling& Simulation, 7(3):1192 1219, 2009.
Lutz Molgedey and Heinz Georg Schuster. Separation of a mixture of independent signals using time delayed correlations. Physical review letters, 72(23):3634, 1994.
Hirotomo Moriwaki, Yu-Shi Tian, Norihito Kawashita, and Tatsuya Takagi. Mordred: a molecular descriptor calculator. Journal of cheminformatics, 10:1-14, 2018.
Niklas Muennighoff, Alexander Rush, Boaz Barak, Teven Le Scao, Nouamane Tazi, Aleksandra Piktus, Sampo Pyysalo, Thomas Wolf, and Colin A Raffel. Scaling lata-constrained language models. Advances Neural Information Processing Systems, 36, 2024.
Albert Musaelian, Simon Batzner, Anders Johansson, Lixin Sun, Cameron J Owen, Mordechai Kornbluth, and Boris Kozinsky. Learning local equivariant representations for large-scale atomistic dynamics. Nature Communications, 14(1):579, 2023.
Maho Nakata and Tomomi Shimazaki. Pubchemqe project: a large-scale first-principles electronic structure database for data-driven chemistry. Journal of chemical information and modeling, (3(6)):300-10000 1308, 2017.
Yusuke Naritomi and Sotaro Fuchigami. Slow dynamics in protein fluctuations revealed by time-structure based independent component analysis: the case of domain motions. The Journal of chemical physics, 134 (6), 2011.
Yuyan Ni, Shikun Feng, Wei- Wei-Ying Ma, Zhi-Ming Ma, and Yanyan Lan. Sliced denoising: A physics-informed molecular pre- pre-trainin method. arXiv preprint arXiv:2311.02122 2023.
Feliks Nüske, Hao Wu, Jan-Hendrik Prinz, Christoph Wehmeyer, Cecilia Clementi, and Frank Noé. Markov state models from short non- equuiiuiiirr simulations -nnalysss and correction of estimation bias. The Jour- nal of Chemical Physics, 146(9), 2017.
Zihan Pengmei, Junyu Liu, and Yinan Shu. Beyond md17: the reactive xxmd dataset. Scientific Data, 11(1): 222, 2024a.
Zihan Pengmei, Chatipat Lorpaiboon, Spencer C. Guo, Jonathan Weare, and Aaron R. Dinner. geom2vec: pretrained gnns as geometric featurizers for conformational dynamics. arXiv preprint arXiv:2409. 19838, 2024b. URL http::tttt//ttiieee org. abs /20l///a..///..../// 19838
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal. Girish Sas- try, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. In International conference on machine learning, pp. 8768- PMLR, 2021.
12

---

Published a conference paper ICLR 2025
Raghunathan Ramakrishnan, Pavlo Dral, Matthias Rupp, and O Anatole Von Lilienfeld. Quantum chemistry structures and properties of 134 kilo molecules. Scientific data, 1(1):1-7, 2014.
Martin K Scherer, Benjamin Trendelkamp- Fabian Paul, Guillermo Pérez-Hermández, Moritz Hoff- mann, Nuria Plattner, Christoph Wehmeyer, Jan-Hendrik Prinz, and Frank Noé. Pyemma 2: A software package for estimation, validation, and analysis of markov models. Journal of chemical theory and compu- tation, 1 11(11 1):55)555555444 2015.
Kristof Schütt, Oliver Unke, and Michael Gastegger. Equivariant message passing for the prediction of tensorial properties and molecular spectra. In International Conference Machine Learning, pp. 9377-9388. PMLR, 2021.
Kristof Schütt, Huziel E Sauceda, P-J Kindermans, Alexandre Tkatchenko, and K-R Müller. Schnet-a deep learning architecture for molecules and materials. The Journal of Chemical Physics, 148(24), 2018.
Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamies. In International conference on machine learning, pp. 2256-2265. PMLR, 2015.
Jaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. arXiv preprint arXiv:2010.02502, 2020a.
Yang Song, Jaseha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:2011.13456, 2020b.
John Strahan, Justin Finkel, Aaron R Dinner, and Jonathan Weare. Predicting rare events using neural networks and short-traeectory data. Journal of computational physics, 488: 12152, 2023a.
John Strahan, Spencer C Guo, Chatipat Lorpaiboon, Aaron R Dinner, and Jonathan Weare. Inexact iterative numerical linear algebra for neural network-based spectral estimation and rare-event prediction. The Journal of Chemical Physics, 159(1), 2023b.
Qiming Sun, Xing Zhang, Samragni Banerjee, Peng Bao, Mare Barbry, Nick S Blunt, Nikolay Bogdanov, George H Booth, Jia Chen, Zhi-Hao Cui, et al. Recent developments in the pyscf program package. The Journal of chemical physics, 153(2), 2020.
Philipp Thölke and Gianni De Fabritiis. Torchmd-net: equivariant transformers for neural network based molecular potentials. arXiv preprint arXiv:2202.02533 2022.
Roberto Todeschini and Viviana Consonni. Handbook of molecular descriptors. John Wiley & Sons, 2008.
Dmitry Ulyanov, Andrea Vedaldi, and Victor Lempitsky. Deep image prior. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 9446-9454, 2018.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
Limei Wang, Haoran Liu, Yi Liu, Jerry Kurtin, and Shuiwang Ji. Learning hierarchical protein representations via complete 3d graph networks. arXiv preprint arXiv:2207.12600, 2022.
Yusong Wang, Tong Wang, Shaoning Li, Xinheng He, Mingyu Li, Zun Wang, Nanning Zheng, Bin Shao, and Tie-Yan Liu. Enhancing geometric representations for molecules with equivariant vector-scalar interactive message passing. Nature Communications, 15(1):313, 2024.
Yuyang Wang, Changwen Xu, Zijie Li, and Amir Barati Farimani. Denoise pretraining on nonequilibrium molecules for accurate and transferable neural potentials. Journal of Chemical Theory and Computation, 19 (15):5077-5087, 2023.
Christoph Wehmeyer and Frank Noe. Time-lagged autoencoders: Deep learning of slow collective variables for molecular kinetics. The Journal of chemical physics, 148(24), 2018.
13

---

Published conference paper ICLR 2025
Florian Weigend. Accurate coulomb-fiting basis sets for to m. Physical chemistry chemical physics, 8()): 1057-1065, 2006.
Hao Wu and Frank Noé. Variational approach for learning markov processes from time series data. Journal of Nonlinear Science, 30(1):23-66, 2020.
Sheheryar Zaidi, Michael Schaarschmidt, James Martens, Hyunjik Kim, Yee Whye Tch, Alvaro Sanchez- Gonzalez, Peter Battaglia, Razvan Pascanu, and Jonathan Godwin. Pre-training via denoising for molecular property prediction. arXiv preprint arXiv:2206.00/33. 2022.
Gang Zhang and Charles B Musgrave. Comparison of dft methods for molecular orbital eigenvalue calcula- tions. The journal of physical chemistry A, 111(8):1554-151, 2007.
Yan Zhao and Donald G Truhlar. Benchmark databases for nonbonded interactions and their use to test density functional theory. Journal of Chemical Theory and Computation, 133):415-42,, 2005.
Yan Zhao and Donald G Truhlar. The m06 suite of density functionals for main group thermochemistry. thermo- chemical kinetics, noncovalent interactions, excited states, and transition elements: two new functionals and systematic testing of four m06-class functionals and other functionals. Theoretical chemistry accounts, 120:215-241, 2008.
14

---

Published conference paper ICLR 2025
A
APPENDIX
B EQUIVARIANT GRAPH NEURAL NETWORKS
Equivariant Geom-GNNs are essential in chemistry and biology for leveraging spatial data like atomic posi- tions while maintaining rotational and translational invariance and equivariance. Equivariance ensures that a function f transforms its output consistently with its input transformation, expressed mathematically as f(T(=))
(x)). Geom- Boom-GNNs using invariant scalars (Schütt al. 2018 Gasteiger et al. 2020 can
easily achieve roto-invariance, but they fall short in predicting general equivariant properties other than deriva- tives with respect the input coordinates. Group-equivariant DeoomGNNss GNNs (Batzner et al. 2022 Batatia et al. 2022 Musaclian et al. 2023) utilize spherical harmonics and group representation theory. However, these architectures are expressive yet memory-inenssee and computationally expensive. The emergence of vector- based equivariant Geom-GNN GNNs, such as (Schütt et al. 2021 Thölke & De Fabritiis, 2022 Wang et al.) 2024), has shown promising results while scaling more efficiently.
C NEURAL SCALING LAW ON LANGUAGE
Neural Scaling Laws describe pepirrccally-derivee power-law relationships between model performance and various scaling factors. These laws have been instrumental in understanding and predicting the behavior of large language models. Kaplan al. (2020) identified key power-law relationships between the pre-training loss L and several variables: L X N
< D- L x C-? where N is the number of model parame- ters, D is the dataset size, and the amount of compute. These laws have elucidated performance disparities between different neural architectures, such Long Short-Term Memory (LSTM) networks (Hochreiter & Schmidhuber 1997) and Transformers (Vaswani al. 2017), particularly in terms of pre-training loss. Sub- sequent work by (Hoffman et al. (2022) and Muennighofff et al. (2024) further refined these relationships, considering factors like optimal batch size, compute-optimal training, and data constraint.
D DATASET SUMMARY
pretraining and downstream detailing their purpose, ular domains, structural types (equilibrium vs. non equilibruu), labeled pruperties, the number of sam- ples, and method data splitiing
Table Overview
Purpose
Domain
Structure
#samples
Split
Pretraining
PCQM4M<
Small
Equilibrium
Conformation
3,378,606
N/A
molecules Molecules
OrhNei
Non-
Conformation
2,383,351
N/A
Denal
compleses
equilibrium
Downstream
QM&
Small malecules
Equilibrium
Conformation, quantum properties Conformation, po- tential, force Conformation Conformation
116,038
Scaffold
Non-
Temporal
molecules
equilibrium
Folg
Proteins
Equilibrium
16,292 ns tra- jectories
Scaffold Random
Alanine
Mini poptide
Dipoptid
equilibrium
Pentapepti
Mini peptide
Non-
Conformation
25 500 ns tra- jectories
Trajectory
equilibrium
X6-8
Protcin
Non-
Conformation
trajecto-
Trajectory
equilibrium
ries with vary- ing
h:tps://ogh. http://wwww...tttt//////////////
/pylocch geometr hedocs, http:///wooooooo...rir.r//////rr datasse QMt.htni
bc htmo:://qithbb.. nom/ httmm//ppngmmm//mmm /xx0 NXX
b http:///tcccbb.com///eeeetooooo
. git.h..hh......... htt.hh/....oosoe/eeaeeeaaaa hntpaa/reeeeeeeeaaaaaeeeeeeeeee. oteppopoceeeeeeeeeeeeeeeeeeee detasets
h http::/:aeeeetec..t....eceeeeete
b bctps: http:///xxhoooee stanford..ar........... http:///wac.teeeeeaeeeaaaaaaaaaa
15

---

Published conference paper ICLR 2025
E UNCERTAINTY IN QM9 LABELS
HOMO-LUMD Gap (meV)
1
gdb
6th
11000
131104
1:30104
9:54993
3334103
1.20993
700
30104
S0114
Ic134
8.7e103
1996
222
-22-04
13
24
854+90
0S6+23
7.le000
8:2
23000
Lueeas
Lence
1.0404
.2000
2.00
10:
1.0004
10133
10:044
15
2.5ew0k
144004
1.5e04
1.3+344
1.0-04
1ie-c4
9000
1.2xx44
1.0+24
2.3+00
3000
610
635
Figure 5: HOMO-LUMO gap of 'gdbl', 'gdb2' and 'gdb3' predictied by varying combinations of common basis sets and density functionals. Different combinations showed considerable variance.
As shown
Figure 5 we scanned combinations five commonly used functionals (Zhao Truhlar 2008
Haoyu
al.,
2016 Adamo & Barone 1999; Lee et al.] 1988) and basis sets (Weigend) [2006] Ditchfield
et al.
1971
Dunning Jr 1989), resulting in 25 distinct computational methods using PySCF package (Sun
al. 2020).
The standard deviations of the HOMO-LUMO gap predictions were substantial: 12571=
16226 and 9333 1758 meV for gdb1, gdb2, and gdb3, respectively. Such uncertainty from DFT calcu- lations has been widely studied in chemistry community (Zhang & Musgrave (007), especially the accuracy of B3LYP functional used in QM9 dataset is non-ideal per current standard (Zhao & Truhlar 2005 Goerigk & Grimme 2011 Mardirossian & Head-Gordon 2017).
16

---

Published conference paper ICLR 2025
F HYPERPARAMETERS SUMMARY
Table 5: Hyperparameters used for different models used for pre-training.
Hyperparameter
ViSNet
Hidden dimension
64, 14.12222566 25,,25566
# MP layers
2-10
# Attention heads
8
# RBFs
32
Layemorm
Max-Min (ViSNet), Whitened/None (ET)
Radius cutoff
3-9
Epochs
10
Batch size
400
Learning rate
0.0005
Optimizer
Adam
Noise level
0.05 (PCQM), (Denali)
Table Hyperparameters used for different models in the supervised regression experiments using QM9 and xxMD datasets. In xxMD experiments, set the weight the force and energy 100:1.
Hyperparameter
QM9
xxMD
Hidden dimension
14,112,,56,,844512
64,122,256
MP layers
Attention heads
8
# RBFs
32
Layernorm
Max-Min, Whitened
Max-Min, None
Radius cutoff
Epochs
500
200
Early Stopping (epochs)
10
20
Batch size
50
12
Learning
5c-5/le-4/5e-4 (from-scratch)
1c-4/5c-4 (from-scratch)
Optimizer
Adam
Scheduler
Reducel_ROnPlateen
Scheduler patience (epoch)
8
Table 7: Hyperparameters used for different models in the VAMP experiments.
Hyperparameter
Ala2
Pentapeptide
X6-85
Base model
ViSNet (PCQM)
Hidden dimension
64, 128,25688888
Output dimension
6
1-10
6
MP layers
Attention heads
# RBFs
Layernorm
Max-Min
Radius cutoff
Token mixer
None 10
None 50
None, MLP-Mixer, Transformer
Epochs
50
Early Stopping (train,val batches)
100,5 5000
Batch size
Learning rate
5E-04
2B-04
2E-04
Optimizrr
Adam
17

---

Published conference paper ICLR 2025
Table 8: Hyperparameters used for fold classification experiments.
Hyperparameter
ProNetAAmino-AcddViiNee
Hidden dimension
64,128,256,384
# MP layers
2
Attention heads
8
RBFs
32
Layernorm
Max-Min
Radius cutoff
5 (ViSNet), 10 (ProNet)
Epochs
400
Batch size
32
Learning rate
0.0005
Optimizer
Adam TRUE FALSE
Coord noise
Rep. noise
Res. Mask
TRUE
Dropout
0.4 (ProNet-AA), 0.5 (ProNet- Vis)
G ILLUSTRATION OF DENOISING RESULTS
ViSket
TroND
Conwerged
Denaii
Converged
2.w
22
Dowleng
2x
10
14x
12x
150
500 102
200
300
320
400
Model Width
1.5
25
Figure 6: Scaling behavior of ViSNet and ET models on PCQM and Denali datasets. The plots show the denoising Mean Squared Error (MSE) against model width and colored with the number model parameters. Results are presented for both the first epoch and converged performance on each dataset. Solid lines with markers represent observed data, while dashed lines indicate power-law fits.
18

---

Published a conference paper at ICLR 2025
PCOM First Enoch
Denali First Epoch
20-10-
Layers
210
216
22x10-1
356
256
64
1.
2x30
160
1:0
360
323
312
1/1
128
10
14"
128
250
384
3R
512
200
11
101
10'
Number of Parameters
Number Parameters
Figure 7: Scaling behavior of ViSNet depth for the first epoch results of training on PCQM and Denali datasets. The plot shows the relationship between the number of layers and the pre-training loss, illustrating the initial rapid improvement followed by diminishing returns as depth increases.
POOM Firtt Epoch
Donai Finst Esodh
20"16"
20""10
23-10
210
216
22"10"
200
27"10"
120
184
104
128
128
26"
112
1.8x10"
250 312
244100
112
120
250
258
224
255
2:3-10 10
120
256
312
250
354
22"10"
10'
612
Number of Farameters
1.-10
2''100'
10
16'
1
Aepect milo
Axpect
INCTUM Con wwrged
Denaai Converged
20
25"100
18x 10"
216
250
10"
23x10"
432
15x
184
Cose
144
3ED
20"10"
128
160
360
120
312
112
120
312
225
2013
2
258
12x
224
240
268
334
10010'
512
10'
16'
10
Sopect rolo
Nepoot rofo
Figure 8: Influence of the aspect ratio the self-supervised denoising pre- pre-training task using the ViSNet model. With similar amount of parameters, pre-training performance improves as the aspect ratio decreases.
19

---

Published conference paper ICLR 2025
First Epoch Training Loss
Converged Training Loss
0.24
0.28
Epoch
PCQM Converged
Denali First Epoch
Denali Converged
0.22
0.26
Training Loss
Training Loss
0.20
0.22
0.18
0.20
0.14
3
5
9
3
Radius Cutoff (A)
Radius Cutoff (A)
Figure 9: Influence of the radius cutoff distance the self-supervised denoising pre-training task using the ViSNet model. Orange and blue curve refer to models trained using Denali and PCQM datasets, respectively. Both the first epoch and the converged loss are reported. Models pre-trained both datasets demonstrate same optimal radius cutoff at 5 Á. Deviating further from the optimal cutoff, the penalty of models pre-trained on Denali dataset more severe than the PCQM.
H ILLUSTRATION OF QM9 RESULTS
Dipole moment (D)
Isotoplc polarizabiity (mağ)
Eucuo: HOMO enengy (meV)
200
199
45
MAVE
NAE
80
100
30
100
100
200
230
400
500
100
203
500
400
500
100
200
300
400
500
Feetlrre dimansion
Fealure dimension
Feature dimansion
Eurso: energy (mel)
Gap (meV)
110
45
35
80
70
100
200
300
500
100
400
500
Feature dimension
Feature dimension
Figure 10: Performance comparison of ViSNet and ET models Scaffold-QMP dataset with varying model widths. The plot shows Mean Absolute Error (MAE) against model width, comparing pre-trained (fine-tuned) models with those trained from scratch. Dashed lines refer non pre-trained models and solid lines refer to pre-trained models using PCQM dataset. ViSNet consistently outperforms ET. Pre-trained models show better performance than their from- scratch counterparts.
20

---

Published conference paper ICLR 2025
CM9 Taiqetca: pca (amm
D20 Target bootooo pdattablly ima]l
CV0 Turgel Aucees HCVO enery Imal)
sufted
17.4
150
104
130
100
200
300
200
Fasture divansionn
Fastua Erwedonn
Fasture
Turge
73
48
Feskure derorrnd
Pallan
Figure 11: Illustration of the performance of PCQM pre-trained ViSNet models with varying width Ran- dom and Scaffold-QM9. The plots indicate different scaling behaviors with respect to the model parameters. Random-QMP can always benefit from scaling the model size while the Scaffold-QM9 cannot.
I ILLUSTRATION OF XXMD RESULTS
PROMPPT Dan 64
POOK-PT Dan 128
POOM-PT Din 258
POOM-PT Dan 564
PT Din 64
Denal FT Din 120
...
Denali PT Dim 250 Nc-PT Dan 250
DenaliFF Din 304
No-PT Din 64
No Din 129
-a-
Nc.PT Dim 384
Energy MAE (mev)
MIE (melol))
10
10
10'
Data Points
Figure 12: Performance comparison ViSNet models on the randomly split stilbene subset xxMD-DFF for energy and force prediction. The plot shows the Mean Absolute Error (MAE) of potential energy against the number of training samples for models with varying widths [64, 128, 256, 384]. Three training scenarios are compared: from-scratch, pre-trained on PCQM (Equilibrium structures), and pre-trained on Denali (Non- cquilibrium structures). This experiment illustrates: (I) we observe an intersection point beyond 10x samples which pre-training may actually hinder performance compared to training from scratch. This suggests a com- plex relationship between pre- pre-trainin benefits and downstream dataset size. (2) in data-abundant regimes (with more than 10 samples), larger models generally demonstrate superior performance, aligning with typical scal- ing laws in machine learning. (3) in data-sparse scenarios, only models pre-trained on the PCQM dataset show positive transfer to downstream tasks, highlighting the difference of learned embedding from equilibrium and non-equilibrium pre- rre-training datasets.
21

---

Published conference paper ICLR 2025
Energy MAE Syytem a70
Energy MAE SaeEE-Syyem...o
Energy MAE Systemt mal
Energy MAE System
422
180
200
-022
Energy MEE (mm44)
120
Erange MAC (mra))
180
Erengy MME (1ri))
393
100
Energy
210
302
140
100
120
100
100
300
80
80
303
130
200
400
200
500
100
160
200
Foalure
dirensionn
MAE System
System:
Bystem:
9ystam
130
170
120
140
951
Foroo MKE (meeo)))
Fooce MMKE (ma4iA))
242
M6E (two)(A)
110
120
992
223
142
140
200
200
ViSket
PSMM
PT ToolD ETiDenal) PT.bxwt EYYPONN)
SNesDeroll
MoXE
Figure 13: Comparison of Energy MAE and Force MAE using Cemporal-xxMD-JUU dataset across different systems for various architectures and training configurations. Each subplot represents different system ('azo': azobenzene, 'sti': stilbene, 'mal': malonaldehyde, and 'dia': dithiophene), with consistent color coding for architectures and distinct markers indicating different training methods and pre-training datasets. Top row shows Energy MAE, while the bottom row shows Force MAE, demonstrating model performance variation with feature dimensions.
J VAMPNET OBJECTIVE AND ILLUSTRATION OF MORE VAMP RESULTS
VAMPNet is a deep-learning-bant based objective that learns the projection from high-dimensional feature space -e-dimmnssionll latent space while preserving the slow dynamics of given time series data. Consider molecular system characterized by hhigh-hiimnnn digh-diimnsional configuration space, X' Let {2,][-1 denote trajectory from molecular dynamics simulation, where xa E X' represents the molecular configuration time For general non-stationary and non-revsersible dynamics, the Koopman operator Kr, describes the time evolution of observable (xx) over a time lag T:
(+)(z)) E(v(z(+7) mz).
To approximate the Koopman operator in practice, used linear superposition ansatz finite set of basis functions (features).
N
f(x)
Σ cxi(1),
where e; are coefficients, and (x) are basis functions or features that depend on the conformational degrees of freedom of the system. Then, the finite-dimensional approximation of the Koopman matrix K is defined
K CTo Cun-
22

---

Published conference paper ICLR 2025
where the instantaneous covariance matrices Coo and CI1, as well as the time-lagged covariance matrix Coi. computed as follows:
Coo
(x(1) po) (x(() pD)
T
1=0
T
Cn = T - Σ
(x(t) - (x(() - -1)
Co
Σ (x(t) no) 7) -)
The mean vectors 1l and ja computed from all data excluding the last steps and the first steps of every trajectory, respectively:
N4
x(t)
Σ ((t)
To refine this approximation, the half-weighted Koopman matrix K introduced, which normalizes the co- variance matrices using their inverse square roots:
K = 0o Cm C--1/2 Co
The matrix is significant as encodes the optimal reduced model for the dynamical system, particularly preserving the slowest modes. The singular value decomposition (SVD) of K is then performed extract the primary components of the dynamics:
=
where U and V" rrepresent the left and right singular vector matrices, and the diagonal matrix of singular values.
To derive the singular functions, the input feature vectors are subsequently projected onto the singular vectors, yielding the left singular functions (t) and right singular functions é(t):
(t) = U C (x(t) mo).
<(t) V"T
V' 1T (x(() m1).
The VAMP-2 score, which is a measure of the quality of the dynamical model in preserving the slow dynamics, directly related to the singular values of the half- weighted Koopman matrix K:
VAMP-2 ||Co0
1/2 Co 1/2 1/2
Σ of
where || || denotes the Frobenius norm, and az R) are the singular values of the half half-weighted Koopman matrix K. Then the VAMP-2 score can be used to optimize the network, which is equivalent to finding representation & such that the corresponding Koopman matrix captures the dominant dynamical modes.
23

---

Published conference paper ICLR 2025
15
-15
CV3
Component
-5
10
15
5
CV2
10
CV1
Figure 14: Characteristic structures pentapeptide in the trajectories visualized on the three-dimensional projection of the learned 6-state VAMP coordinates.
K ILLUSTRATION OF FOLD CLASSIFICATION RESULTS AND EMBEDDINGS
In this study, explore whether graph emheddings effectively represent local all-atom structures. To achieve this, we fragmented the proteins into individual residues while retaining a sliding window of 3/2 residues based on the sequence to provide contextual information. We then input the corresponding atomic coordinates into the Geom-GNNs to obtain the embedding feature vectors. This approach also helps mitigate the over-smoothing problem, where global information is mistakenly mixed into local features, thus deteriorat- ing the quality of local embeddings.
To visualize the features learned by Geom-GNNs, we generated two-dimensional t-SNE embeddings of the latent space for both pre-trained and untrained networks using the validation set of the Homology TAPE dataset, as shown in Figures[T]]and]]]] [8] While equivariant features provide higher resolution, they are often challenging to reduce in dimensionality with standard techniques like principal component analysis (PCA) or t-SNE. To enhance visualization, we regulated the scalar channels of the network during denoising pre-training, following the approach of Liu al. (2022). Additionally, we noticed multiple preprocessing errors in the provided dataset when visualizing characteristic structures in the graph embedding space Figure 19 We followed the Wang et al. (2022) to the pre-processed datasets from Hermosilla et al. (2021) (available at hntps:: //qithub. com/ www.phtweeeeeeoccccoooooooeeeeee poperoos rmos wwwssmllamIaaaa...m. _ These preprocessing errors are difficult to detect in feature spaces defined by dihedral and torsion angles, but more apparent in the graph embedding space, which could negatively affect regression accuracy.
It important note that projections based dihedral and torsion angles should not be directly compared with graph embeddings. While these angular features are effective for characterizing secondary structures, they lack all-atom resolution, are coarse-grained, and do not account for environmental effects. is still missing piece of study to define some general emperical descriptors that can roughly describe the interactions between residues and their side-chains. Our comparisons serve to confirm consistency and agreement between the different representations, Interestingly, even untrained Geom- GNNs are able to capture patterns in the dataset, similar untrained convolutional neural networks (CNNs) in vision tasks, as demonstrated by [Uyanov et al. 2018). However, the features derived from untrained networks are generally sparse and less interpretable.
24

---

Published conference paper ICLR 2025
20
15.0
12.5
10
15
10
5
CV3
Component
CV3
10.0
Component
-5
-5
5.0
15
-15
15
15
10
0.0
CV2
CV2
10
CV1
15
20
CV1
15
0
15
15
10
15
10
10
CV3
Component
CV3
Component
5
5
100
-15
-15
15
10
100
10
10
10
CV1
CV1
0
15
6
15
10
10
CV3
Component
Component
5
10
-10
-15
-15
+
15
-4
10
CV2
5
10
CV1
20
Figure 15: Three-dimensional projections learned singular vectors in descending order from the 6-state VAMP model. Lag time: 0.5 ns.
25

---

Published conference paper ICLR 2025
27.5
24
25.0
Tro22-1y333 Distance (Angstrom)
1.0
Radius of Gyration (Angstrom)
1.0
22.5
20
0.5
20.0
17.5
15.0
16
-0.5
-0.5
12.5
10.0
-1.0
12
-1.0
RMSD (Angstrom)
RMSD (Angstrom)
1.0
24
0.9
1.0
1.0
Radius of Gyration (Angstrom)
22
0.5
20
0.5
0.7
18
0.0
0.0
0.6
-0.5
-0.5
24
-1.0
-1.0
10
15
20
10.0
12.5
15.0
17.5
20.0
22.5
25.0
27.5
RMSD (Angstrom)
Trp22-Tyr33 Distance (Angstrom)
0.9
1.0
0.9
0.8
0.8
0.6
0.5
0.5
0.4
0.4
10.0 12.5 15.0 17.5 20.0
22.5 25.0
27.5
16
18
20
24
Trp22-Tyr33 Distance (Angstrom)
Radius of Gyration (Angstrom)
Figure 16: Projection of the first VAMP singular vector (Lag time: 50 ns) typical collective variables such as radius of gyration, structural rooted mean squared difference (RMSD) to the crystal structure, fraction of native contacts (Q), pair-wise distance of Trp22 and Tyr33. [Bowman et (20TT) indicates 2-3 states on those common collective variables, and the results are comparable. We randomly drew 500 trajectories out of the original dataset train the VAMP head.
26

---

Published conference paper ICLR 2025
60
80
40
2nd dimensson
2nd dinersso
0
-20
2nd
-20
-40
-40
-40
-60
-60
-50
-75
D
25
-50
-75
75
50
-50
-75
75
1st dimensior
1st dimensioo
1st d dimensson
75
50
2nd dinensioon
2nd dimension
20
-25
-2
-50
-40
-60
-80
1st dimensior
dimension
1st d mimession
40
40
2nd dimension
2nd dimensioon
20
2nd dimension
-20
-20
-20
-40
-40
-40
-40
-20
0
20
40
-40
-20
0
20
40
60
-40
-20
20
40
1st dimensior
dimension
1st mimession
2nd mimessoon
20
2nd dimension
20
mension
-20
-20
-40
-40
-40
1st dimensior
dimension
1st memesson
Figure 17: Two-dimensional (-SNE visualization of residues in the validation subset of CaanologyTAPE dataset for the folding classification task. Points are colored according to the one-dimensional t-SNE projection of the angles. Columns represent: (Left) backbone dihedral and sidechain torsion angles, (Middle) pre-trained graph embedding from ViSNet with width of 128, and (Right) untrained graph embedding from the same model. Rows correspond to different amino acids: (1) Histidine (His), (2) Aspartic Acid (Asp), (3) Cysteine (Cys), and (4) Methionine (Met).
27

---

Published conference paper ICLR 2025
60
60
40
2nd dinensson
2nd dinersion
20
2nd dimension
0
-20
-20
-40
-40
-40
-60
-60
-50
-75
D
25
-50
-75
75
50
-50
-75
75
1st dimensior
1st dimension
1st d diimensoon
75
50
2nd dinensioon
2nd dimension
20
-25
-25
-20
-50
-40
-60
-80
1st dimensior
dimension
1st d memession
40
40
40
2nd dimension
2nd dimensioon
20
2nd dimension
-20
-20
-20
-40
-40
-40
-40
-20
0
20
40
-40
-20
20
40
60
-40
-20
20
40
1st dimensior
dimension
1st memsssnn
2nd dimessoon
20
2nd dimension
20
mension
-20
-20
-40
-40
-40
1st dimensior
dimension
1st mmmessio
Figure 18: Two-dimensional t-SNE visualization of residues in the validation subset of HomologyTAie TAPE dataset for the folding classification task. Points are colored according to DBSCAN clustering labels of the angles. Columns represent: (Left) backbone dihedral and sidechain torsion angles, (Middle) pre-trained graph embed- ding from ViSNet with width of 128, and (Right) untrained graph embedding from the same model. Rows correspond to different amino acids: (() Histidine (His), (2) Aspartic Acid (Asp), (3) Cysteine (Cys), and (4) Methionine (Met).
28

---

Published conference paper ICLR 2025
20
2nd t-SEN dimension
20
440
-60
-60
40
1st t-SNE cimension
Figure 19: Visualization characteristic structures for histidine (His) residues using two-dimensional t-SNE projections of pre-trained graph embeddings. The plot shows that structurally similar residues clustered together, while multiple outliers, likely due to preprocessing errors, are also present. Similar preprocessing errors were observed for other types of residues as well.
29

---

Published conference paper ICLR 2025
FOLD split
Superfamily split
50
49
66
48
64
Accuracy (%)
47
Accuracy (%)
62
46
60
44
43
AA-PRONET
42
56
GNN-PRONET
100
200
300
400
100
200
300
400
Feature dimension
Feature dimension
Figure 20: Folding classification accuracy with the fold and superfamily split. For pronet architecture, we restrict ourselves to use amino-acid-level features (AA-PRONET), whereas we supplement the original ar- chitecture with the pre-trained GNN features for transfer (GNN-PRONET), is important to note that our preliminary implementation is not compatible with the data augmentation techniques employed in the original ProNet paper. Our primary objective is not to establish a new state-of-the-art model, but rather to analyze the transfcrability of pre-trained all-atom graph representations to distinct architectural framework and task domain. We reproduced the original ProNet results with multiple runs to establish a baseline for comparison. There minor discrepancies between reproduced values and those reported in the original paper, presum- ably due to normal experimental variance with different hyperparameters. As illustrated in plots, both models stopped to improve at 256 width the fold split, while increasing the model size can continue benefit the superfamily split.
30

---

Published conference paper ICLR 2025
L ADDITIONAL RESULTS XXMD-TEMPORKNE AND QM9 BENCHMARKS
Table 9: Performance comparison across different dimensions and training setups for ViSNet and ET models on four subsets of xxMD-DFT using temporal split.
Molecule
Model
Dim
No PT
PT Denali
PT PCQM
E
F
E
F
E
F
ViSNet
64
91
95
79
83
99
99
128
82
85
73
83
82
90
256
75
72
85
92
99
94
384
104
94
93
124
Azobenzene
ET
64
121
133
127
98
160
122
128
135
118
145
104
187
135
256
224
131
176
108
162
115
384
113
172
106
200
ViSNet
64
332
148
368
148
315
147
128
173
321
137
361
256
302
140
335
157
330
139
384
352
157
361
154
313
144
Stilbene
ET
64
406
174
349
153
321
147
128
350
172
355
148
298
144
256
164
317
149
331
384
379
154
348
142
346
144
ViSNet
64
162
99
156
99
151
128
94
157
98
145
117
160
256
78
147
95
142
112
157
384
184
193
95
164
225
290
Malonaldehyde
64
188
124
165
113
128
136
183
117
163
120
180
256
147
174
155
162
165
171
384
154
179
144
157
173
176
ViSNet
64
141
79
137
79
78
78
128
128
127
111
71
128
79
256
75
57
76
384
127
90
102
83
74
164
Dithiopehene
ET
64
117
82
78
74
110
80
128
110
82
92
82
100
88
256
122
116
81
105
384
126
97
83
144
M GENERAL RECOMMENDATIONS FOR SCALING AND EVALUATING GEOMETRIC GNNS
To facilitate the training and evaluation of geometric GNNs, provide integrated summary of the obser- vations drawn from experiments on self-supervised pre-training, unsupervised transfer, and supervised fine- tuning. These findings offer practical guidance for practitioners aiming to scale and optimize GNN architec- tures effectively.
From the perspective of pre-training, observed consistent decrease in training loss the depth of the net- work increased. However, these benefits diminished once a critical depth was reached. For instance, a ViSNet with six layers achieved training loss comparable to that an eight-layer ViSNet with the parameter count matched. This plateau effect suggests that, while increasing depth initially enhances the network's 's learning capacity, over-scaling yields diminishing returns. In molecular property prediction tasks, as demonstrated in
31

---

Published conference paper ICLR 2025
Table 10: Comparison of éLUMO and Ae tasks in Scaffold-QM9 across different Layer-Dimensions with various pre-training setups.
Layer-Dim
Setup
Metrics
CLUMO
Ac
6L-128
No PT
42.8
94.5
PT on PCQM
35.9
80.7
5L-144
No PT
46.1
104.8
PT on PCQM
39.5
88.3
4L-160
No PT
45.6
98.5
PT PCQM
39.8
96.7
31-184
No PT
44.5
98.4
PT on PCQM
44.1
103.0
2L-216
No PT
43.5
97.6
PT on PCQM
46.6
96.0
increased depth consistently improved prediction accuracy, reaffirming its importance for certain downstream tasks. Nonetheless, in force field prediction tasks, as highlighted by Li et al. (2024) (Figure 2). the advantage of scaling depth similarly diminished beyond a certain point, albeit with less controlled parameter setup. An intriguing phenomenon emerged concerning the radius graph cutoff. Increasing the cutoff distance initially improved performance but eventually led deterioration, forming characteristic U-shaped perfor- mance curve (Figure 99. This aligns with findings from Li et al. (2024), suggesting that excessively large cutoff radii may introduce noise from irrelevant long-range interactions, thereby reducing model effectiveness.
Table 10
In terms of parameter scaling, models with width of 256 384 provided an optimal balance between compu- tational cost and performance across various tasks, as illustrated Figures TO 12 13, and 20 Wider models generally showed only marginal gains relative to their increased computational overbead, indicating diminish- ing returns for excessive parameter expansion. when evaluating the resulting models. Improper splits can lead overfitting and hinder the generalization of learned representations. For example, scaffold splits are more reliable than random splits for assessing the transferability of molecular representations. Moreover, architecture-Ieee leveel improvements remain necessary. For protein kinetic modeling, incorporating additional Transformer or MLP-Mixer layers significantly improved performance (Table Similarly, in force field prediction tasks involving larger molecules, adding long-range interaction modules further enhanced model predictions if applicable, as evidenced by Li et al. (2024).
critical carefully consider the dataset splitting strategy
32


Reviews:
[{"review_id": "8zKhX9ZdtU", "reviewer": "Reviewer_GND7", "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.", "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?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 3, "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.\nThe 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.\nThe 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "R6tfZ9qbLB", "reviewer": "Reviewer_KcDu", "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.", "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?\n\n2. Limited Pre-Training Approaches:\nThe 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?\nAdditionally, could the authors hypothesize the potential impact of such an extended pre-training approach on scaling behavior? \n\n3. 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?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. Demonstrates substantial performance improvements in kinetic modeling on VAMP score metrics by utilizing denoising pretraining techniques.\n2. Applies the self-supervised pretraining approach to a variety of downstream tasks, successfully proving its effectiveness across applications.\n3. 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.\n2. 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.\n3. 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.\n\n[1] Cui, Taoyong, et al. \"Geometry-enhanced pretraining on interatomic potentials.\" Nature Machine Intelligence 6.4 (2024): 428-436.\n[2] Zhou, Gengmo, et al. \"Uni-mol: A universal 3d molecular representation learning framework.\" (2023).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "j2CpFh5XDZ", "reviewer": "Reviewer_i3nd", "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.", "questions": "as Weaknesses.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "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.\n2. **Clarity and Structure**: The article is well-organized, with clear presentation and summary of experiments and viewpoints, facilitating reader comprehension.\n3. **Robust Experimentation**: The experimental design is thorough, effectively supporting the authors’ conclusions.\n4. **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.\n5. **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.\n6. **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.\n2. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "5w3V5g9rTX", "reviewer": "Reviewer_5LXW", "summary": "In their work the authors empirically studied properties of pre-trained geometric GNNs (i.e., GNNs with coordinates attached).\nThey especially seemed to consider pretraining via denoising as introduced by Zaidi et al. (2022).\nThe authors consider several downstream tasks: molecular dynamics (kinetics modeling), fold classification, prediction of quantum-chemical properties, and, molecule conformations\nThey study properties such as power-law scaling and find that the geomtric GNNs do not follow that on the pre-training task.\nThe 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.\"", "questions": "- Why is the VAMP task considered zero-shot?\n- 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?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 1, "contribution": 1, "strengths": "relevant research topics in the paper:\n- pretraining of (geometric) GNNs\n- scaling behavior of GNNs\n- oversmoothing, underreaching\n\nexperiments:\n- 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.\nBut 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\nto 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).\n\\\n\\\nDetails:\n\\\n\\\n**hard to see the novelty in the paper**:\\\nFact that pretraining is useful was already found e.g. by Zaidi et al. (2022).\\\nNo new methodology seems to be suggested (or is \"token mixing\" the new method)?\\\nStudy 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.\n\\\n\\\n**paper is structured in a strange way, which makes it hard to follow the paper**:\\\nIt is actually hard to understand what the research aim of the authors was.\\\nchapter 4, 5, and, 6, actually seem to be about empiricial experiment results.\\\nThe setups are partly however already explained in chapter 3.\\\nchapter 4 and 6 show performances on problems\\\nchapter 5 however studies power-law behavior and other ablations.\\\nIn 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).\n\\\n\\\n**experiments**:\\\nalthough 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.\\\n\\\nalso the there is no good argument in the paper, why the authors exactly compare to those methods they selected to compare to \ne.g.:\\\nfor molecular dynamics, why don't they compare to Timewarp\\\nfor QM9 and xxMD they could e.g. compare to \"E(n) Equivariant Graph Neural Networks\", \"SE(3)-Transformers\", DimeNet++, MACE, etc.\\\nAn 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).\n\\\n\\\n**minor points**:\\\n\"token mixing\" not defined, but heavily used\\\ngrammatical errors/typos: \"In silico molecular computation and simulation are indispensable tool in...\"", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)