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

## Abstract

Abstract Constructing transferable descriptors for conformation representation of molecular and biological systems finds numerous applications in drug discovery, learning-based molecular dynamics, and protein mechanism analysis. Geometric graph neural networks (Geom-GNNs) with all-atom information have transformed atomistic simulations by serving as a general learnable geometric descriptors for downstream tasks including prediction of interatomic potential and molecular properties.
However, common practices involve supervising Geom-GNNs on specific downstream tasks, which suffer from the lack of high-quality data and inaccurate labels leading to poor generalization and performance degradation on out-of-distribution (OOD) scenarios.
In this work, we explored the possibility of using pre-trained Geom-GNNs as transferable and highly effective geometric descriptors for improved generalization.
To explore their representation power, we studied the scaling behaviors of Geom-GNNs under self-supervised pre-training, supervised and unsupervised learning setups. We
find that the expressive power of different architectures can differ on the pre-training task. Interestingly, Geom-GNNs do not follow the power-law scaling on the pre-training task, and universally lack predictable scaling behavior on the supervised tasks with quantum chemical labels important for screening and design of novel molecules.
More importantly, we demonstrate how all-atom graph embedding can be organically combined with other neural architectures to enhance the expressive power. Meanwhile, the low-dimensional projection of the latent space shows excellent agreement with conventional geometrical descriptors.

## 1 Introduction

In silico molecular computation and simulation are indispensable tool in modern research for biology, chemistry, and material sciences capable of accelerating the discovery of new drugs and materials, as well as prediction of protein structures and functions. One of the key advance in molecular computation is the use of learnable geometrical descriptors for representing the atomic environment (Todeschini & Consonni, ; Mauri et al., ; Baig et al., ; Moriwaki et al., ).
Geometric Graph Neural Networks (Geom-GNNs) with all-atom resolutions can learn to construct system-specific molecular force fields by incorporating geometric inductive biases such as rotational invariance (Schütt et al., ; Gasteiger et al., ) and equivariance (Schütt et al., ; Batzner et al., ). These ML-based force fields can be seen as learnable geometric descriptors that map the molecular conformation to the potential function of a target.
While supervised-learning formulations have been widely explored, there is a
lack of high-quality data in biological and chemical domains (e.g., for accurate quantum chemical data beyond density functional theory (DFT) accuracy (Donchev et al., ; Peverati & Truhlar, ; Shu et al., ), long unbiased equilibrium protein trajectories with wide sampling of the conformation spaces (Lindorff-Larsen et al., ; Pan et al., ; Vander Meersche et al., ) and experimental protein structures (Liu et al., ; Burley et al., )).
Many scientific simulations involve learning-based dynamics require efficient and robust descriptors (Bonati et al., ; Mardt et al., ). With limited amount of training data, the resulting ML models can be strongly biased to the training domain leading to poor out-of-distribution (OOD) generalization.

Figure: Figure 1: Meta-architecture for Using Pre-trained Geom-GNNs as descriptors: The figure shows a framework where pre-trained Geom-GNNs act as local geometric descriptors to featurize residue-level conformations. Each window represents an atomic environment for residue feature extraction, defined by a user-defined context (illustrated here as a sliding window of nearest neighbors in sequence). In each window, atomic structures are treated as individual graphs and processed by the pre-trained Geom-GNN to extract atomic-level features, which are aggregated into residue-level representations or “tokens.” The architecture can employ self-attention (SA), multi-layer perceptron (MLP), or message passing mechanisms to enhance representational power. For graph-level tasks, the mixed tokens are pooled and input to a task-specific head for training and predictions.
Refer to caption: extracted/5961044/Figures/schema.png

In contrast to other domains such as text and vision, pre-trained models have been recognized as highly effective feature extractors (descriptors) trained with self-supervised objectives (Devlin et al., ; Radford et al., ; Kirillov et al., ).
Recently, denoising perturbed molecular conformations have been proposed as an effective objective to pre-train Geom-GNNs (Zaidi et al., ).
In this work, we first demonstrate that pre-trained Geom-GNNs act as zero-shot transfer learners, effectively representing molecular conformations even for out-of-distribution (OOD) scenarios.
We evaluate the performance of these pre-trained models by analyzing neural scaling laws (Kaplan et al., ; Hoffmann et al., ). 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. This paper is organized as follows: we first briefly introduce relevant backgrounds of pre-training and unsupervised tasks in molecular simulation with the architectures and datasets used. We then explain the scaling behaviors in three setups. We additionally show how those pre-trained graph embeddings could be easily combined with other architectures to enhance their expressive power. The low-dimensional projection of the latent graph feature space (Figure 2 and Appendix K) shows meaningful clustering of similar conformations. Our contributions can be summarized as follows:

- •
We demonstrate the use of pre-trained Geom-GNNs as zero-shot transfer learners for molecular conformation featurization. When representing residue conformations in peptides and proteins, graph embeddings match conventional rotamer features but possess better expressivity to describe the atomic environment flexibly.
The pre-trained embeddings excel at constructing kinetic models through the VAMPNet objective: a challenging unsupervised learning task that is crucial for enhanced sampling and probing energy landscapes in computational chemistry/biology. Compared to traditional pairwise distance features, the graph embeddings achieve significant performance gains up to 50% across complex peptide and protein systems.
- •
We systematically investigated the scaling behaviors of Geom-GNNs under the self-supervised, supervised and unsupervised setups. We found that Geom-GNNs 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. On supervised tasks with noisy data labels, their performance did not scale predictably, highlighting the need for high-quality labeled data.

## 2 Related works

Pretraining via Denoising is a novel pretraining strategy for equivariant GNN architectures, drawing inspiration from advancements in visual generative models (Sohl-Dickstein et al., ; Song et al., ; ). Zaidi et al. () introduced a method of corrupting input atomic coordinates of equilibrium molecular structures and training the equivariant GNN to denoise and restore the original coordinates. The denoising objective is defined as: $L_{\text{Denoising}}=\mathbb{E}_{r}\|\epsilon-\text{GNN}(\hat{r})\|^{2}$. Here, $\epsilon$ represents the noise added to the original coordinates $r$, producing the noised input $\hat{r}$. Wang et al. () expanded this method to include denoising of non-equilibrium structures. Further studies have investigated varying noise distributions such as sliced noise (Ni et al., ).

Graph neural scaling laws represent efforts to extend scaling laws from language modeling to the graph domain. Frey et al. () 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. () further advanced this area by demonstrating a predictable power-law relationship between supervised loss and dataset size using a topological GNN. More recently, Liu et al. () expanded the discussion to a parametric perspective, focusing on excessively large datasets such as PCQM4M (Nakata & Shimazaki, ; Hu et al., ) in a supervised setting without 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, ; Naritomi & Fuchigami, ; Scherer et al., ) and Variational Approach for Markov Process (VAMP) (Wu & Noé, ; Mardt et al., ), prediction of committer functions Strahan et al. (; ) in transition path theory (Weigend, ; Metzner et al., ), and the application of enhanced sampling methods (Bonati et al., ; Jung et al., ). These tasks require meticulously crafted descriptors that capture the structural nuances of the systems under study. Recent advancements have underscored the value of atomistic detail, leading to the integration of learnable Geom-GNNs (Klein et al., ; Huang et al., ; Liu et al., ) into aforementioned more advanced modeling scenarios. However, the scalability of these methods remains challenging due to the necessity of training graph representations on-the-fly and the demands of OOD exploration.

## 3 Experimental and dataset setup

We focus on two state-of-the-art architectures based on message-passing (MP): Equivariant Transformer (ET) (Thölke & De Fabritiis, ) and its extension, ViSNet (Wang et al., ), which incorporates higher body-order features. To investigate scaling effects, we separately studied effect of width, depth and aspect ratio of MP layers. Radius graph is the most common approach to construct the geometrical graph, where all nodes within the cutoff are connected by edges. We further analyze the effect of radius cutoff distance and refer readers to Appendix F for detailed model configurations and Table D for a summary of datasets used.

Pre-training setup.
For pre-training, we employed a vanilla denoising objective, where scaled noise was added directly to 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, ; Hu et al., ) and the training dataset of OrbNet-Denali (Christensen et al., ) (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 structures relaxed using DFT, thus providing equilibrium structures. Conversely, Denali contains multiple geometric conformations for each molecular topology, offering a diverse range of non-equilibrium structures.

Supervised downstream tasks setup.
The supervised downstream tasks to validate the learned representations encompass diverse molecular structures and conformations. For chemical structure diversity, we utilized the QM9 dataset (Ramakrishnan et al., ) for molecular property prediction, implementing an ID setup through random splitting (Random-QM9) and an OOD setup through scaffold splitting based on molecular motifs using RDKit (Landrum et al., ) (Scaffold-QM9). For conformational diversity, the xxMD dataset (Pengmei et al., ) was employed, which surpasses the MD17 benchmark (Chmiela et al., ) in sampling and computational techniques. This dataset features extensively sampled non-equilibrium conformations within reactive regions. We applied a temporal split (Temporal-xxMD) based on timesteps and a random split (Random-xxMD) as ID and OOD scenarios. Here, Random-xxMD represents a well-sampled condition, ensuring 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 approaches to extract pre-trained graph representations, following the methodology outlined in Mardt et al. () and Pengmei et al. (). For kinetic modeling tasks for characterizing non-reversible dynamics, we utilized molecular dynamics trajectories from three systems: alanine dipeptide (ala2), pentapeptides from MDshare (Nüske et al., ; Wehmeyer & Noé, ), and $\lambda$6-85 fast-folding protein (Bowman et al., ). We inferred and stored activations from the last MP layer of pre-trained all-atom Geom-GNN checkpoints. These features are used to optimize the network as shown in Figure 1 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 a previously developed protein-specific GNN (Wang et al., ) with activations from the pre-trained all-atom Geom-GNN. We then evaluated this enhanced model on a protein fold classification task (Hou et al., ).

## 4 Zero-shot transfer downstream tasks

Comparing supervised objectives, self-supervised objectives are emerging areas of researches in learning-based molecular dynamics simulations. As mentioned before, training graph representations on-the-fly are 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 simulation objectives.
This section is structured as follows: We begin by showcasing the application of zero-shot transferable Geom-GNN representations, integrated with other architectures to meet the VAMP objectives (Mardt et al., ), which aims to learn a low-dimensional representation that preserves the long-timescale kinetics of the molecular system. It is crucial in computational biology for enhanced sampling and understanding the underlying energy landscapes. We explored three systems with increased complexity up to 1258 heavy atoms: ala2, pentapeptide, and the $\lambda$6-85. Subsequently, we demonstrate how these features can be combined with protein-specific GNNs to improve the classification performance on protein folding, with analysis presented in Appendix K for brevity.

### 4.1 Kinetic modeling: VAMP

Figure: Figure 2: Visualization of learned singular vectors of Koopman operator onto the $\Psi-\Phi$ 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 J).
Refer to caption: extracted/5961044/Figures/ala2_coords.png

Initial experiments focused on the ala2 system, known for its metastable states characterized by two main backbone dihedrals, $\Psi$ and $\Phi$. 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 $\Psi-\Phi$ dihedrals. As illustrated in Figure 2, the network successfully identified the first learned singular vector as the slowest transition between positive and negative $\Psi$ angles. The subsequent vector characterized the transition between small and large $\Phi$ angles. Lower-order singular vectors identified transition states between the $\alpha$ and $\beta$ regions of the Ramachandran plot. These observations show that the graph representations capture structurally accurate and semantically meaningful information as illustrated in Figure 2, 14 and 17.

Koopman matrix quantifies the kinetic variance between instantaneous and time-lagged data (Appendix J). A minimal VAMP score of 1 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 3 shows the VAMP scores for models trained on pentapeptide trajectories across various output dimensions and lag times, following the PyEMMA setup Scherer et al. (). Furthermore, to interpret the learned coordinates, we visualized characteristic structures and their projections onto conventional physical collective variables in Figures 14, 15 and 16 (Appendix). These visualizations facilitate direct comparisons with established reference literature (Scherer et al., ; Bowman et al., ). 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.(^*^**[http://www.emma-project.org/latest/tutorials/notebooks/00-pentapeptide-showcase.html](http://www.emma-project.org/latest/tutorials/notebooks/00-pentapeptide-showcase.html))

**Table 1: Validation of VAMP-2 scores ($\uparrow$) for Ala2, Pentapeptide, and $\lambda$6-85 with PCQM-ViSNet with varying dimensions and model architectures, with a 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 post-coarse-graining. Lag times: Ala2 at 1 ps, Pentapeptide at 0.5 ns, and $\lambda$6-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 on an instance with 4 $\times$ NVIDIA A10G GPUs. We reproduced Mardt et al. () with the same setup, for ala2 and pentapeptide we consistently used the heavy atom pairwise distances, and we used pairwise distances of 80 $C_{\alpha}$ atoms for $\lambda$6-85.**
| Dimension | Ala2 | Pentapeptide | $\lambda$6-85 |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- |
| Token Mixing | Sum | Sum | MLP | SA | Sum | MLP | SA |
| Out. Dim. | 6 | 10 | 10 |  |  |  |  |
| $\#$Heavy Atoms | 10 | 44 (5 residues) | 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 | 7.01±0.21 |
| 128 | 4.72±0.01 | 6.41±0.07 | 6.91±0.32 | 7.11±0.43 | 5.17±0.08 | 8.35±0.36 | 7.80±0.18 |
| 256 | 4.70±0.01 | 6.10±0.22 | 7.67±0.08 | 7.04±0.21 | 4.75±0.26 | 8.52±0.14 | OOM |
| 384 | 4.70±0.01 | 6.62±0.16 | 7.71±0.23 | 7.26±0.11 | 5.81±0.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 all critical slow dynamics at this lower dimension. However, for more complex systems like Pentapeptide and $\lambda$6-85, which possess a 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 is 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., ). As shown in Table 1, this approach led to significant performance improvements, particularly for $\lambda$6-85, where applying the MLP-Mixer resulted in over a 46$\%$ increase in VAMP-2 scores compared to the baseline sum pooling method.

In comparison to the baseline VAMPnet model (Mardt et al., ), our approach using pre-trained graph features led to performance improvements of 20.4%, 50.0%, and 15.1% for Ala2, Pentapeptide, and $\lambda$6-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 $\lambda$6-85, the improvements gained from pre-trained features were not as 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 as refining the process of localizing graph embeddings for each residue in protein systems along molecular trajectories. Current methods, which treat the entire protein as a 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 our scope to biomolecules due to the lack of reference.

### 4.2 Protein representation learning: FOLD classification

Figure: 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. Random half of available trajectories are held for validation and the remaining are used for training.
Refer to caption: extracted/5961044/Figures/vamp_scan.png

As we explored the usage of pre-trained graph embedding to study the conformation space of peptides and proteins, we extend to explore the utility of such all-atom embedding in the protein representational learning. Depicted in Figure 1, modeling the dependence among multiple inputs require an active token mixer. MP-based GNNs themselves are also a type of token mixers. In protein-specific GNNs, residues are represented as nodes, and the features attached to each node are usually $C_{\alpha}$ positions and a set of geometrical descriptors for backbone or side-chain dihedral and torsion 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, particularly in environments with $\pi$-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 representation. For illustration, we naively concatenate projections of pre-inferred all-atom embedding to interaction blocks of Wang et al. () and test the downstream task performance on the FOLD dataset (Hou et al., ). Per Figure 20, we observe improvement of integrating such all-atom features into the ProNet without angle information. More importantly, the pre-trained graph embedding can match with conventional rotamer features in non-interacting cases as shown in Figure 17, 18 and 19, where we used a 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.

Figure: 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).
Refer to caption: extracted/5961044/Figures/scan_layers_converged.png

Effect of model width. We investigated the scaling behavior of ET and ViSNet on both PCQM and Denali datasets. For the denoising pre-training task, we 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 not strictly follow a power-law. (2) Rapid initial improvement: In low parameter regimes, model performance improves quickly as size increases. (3) Saturation effect: Beyond a certain threshold (e.g., after $10^{7}$ parameters), performance gains become marginal, with models often reaching a 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 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 in 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 models 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, we observe the checkpoints with similar amount of parameters and lower aspect ratio (width/depth) enjoy better pre-training performance per depicted in Figure 8 (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-GNNs critically defines the model’s effective receptive field by regulating the scope of atomic interactions during message passing. Figure 9 (Appendix) shows significant 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 a minimum at approximately 5 Å before 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.

## 6 Supervised downstream tasks

As described in the experimental setup, we additionally used two molecular benchmarks to verify the performance of learned features: QM9 for molecular properties prediction (chemical diversity) and xxMD for force 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 Scaffold-QM9 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 | $\mu$ (mD) | $\alpha$ ($ma^{3}_{0}$) | $\epsilon_{\text{HOMO}}$ (meV) | $\epsilon_{\text{LUMO}}$ (meV) | $\Delta\epsilon$ (meV) |
| --- | --- | --- | --- | --- | --- | --- |
| PT-ViSNet | 64 | 28.8 | 99.1 | 56.8 | 37.2 | 90.2 |
| 128 | 24.9 | 93.4 | 56.8 | 35.9 | 80.7 |  |
| 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 |  |
| ViSNet | 64 | 36.5 | 142.1 | 74.4 | 46.6 | 110.5 |
| 128 | 32.3 | 106.4 | 62.0 | 42.8 | 94.5 |  |
| 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 |  |
| PT-ET | 64 | 36.9 | 108.3 | 69.0 | 45.8 | 96.3 |
| 128 | 33.5 | 117.4 | 62.1 | 40.2 | 88.4 |  |
| 256 | 29.2 | 116.4 | 54.0 | 36.4 | 84.0 |  |
| 384 | 35.2 | 101.3 | 52.3 | 40.1 | 82.0 |  |
| 512 | 29.4 | 197.8 | 57.7 | 37.5 | 80.1 |  |
| ET | 64 | 35.9 | 104.2 | 72.7 | 47.5 | 101.9 |
| 128 | 32.3 | 98.4 | 65.4 | 44.6 | 96.2 |  |
| 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

Effect of Scaling and Pre-training: We evaluated the performance of both fine-tuned and from-scratch trained ViSNet and ET models across different capacities, as depicted in Table 2 Figure 10 (Appendix). We assessed model performance on a hold-out test set, leading to the following observations. Firstly, we noted a limited scaling behavior; while moderately larger models generally outperform smaller ones, the performance improvements do not predicatively scale with model size, indicating a 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 to ET, aligning with the pre-training results shown in Figures 6 and 7 (Appendix).

Effect of Data Split and Label Uncertainty. Figure 11 (Appendix) depicts the results from the Random-QM9 experiment. Typically, wider models exhibited improvements in Random-QM9, whereas scaffold-split models 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 transferability to downstream tasks, as 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 Appendix E and Zhao & Truhlar (); Zhang & Musgrave (); Mardirossian & Head-Gordon (), are markedly greater than the prediction errors, illustrating the inherent uncertainty in training and evaluation labels. And prediction error of B3LYP functional (Lee et al., ) used in QM9 dataset are usually at eV-scale. This error in ”ground truth” complicates model performance assessment and generalization interpretation, suggesting that training data uncertainty can largely contribute to the observed limitations in model scaling and OOD performance, beyond model limitations.

### 6.2 Conformational Diversity: xxMD

Building upon our analysis of chemical diversity, we now focus on conformational diversity within a 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 to 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 a 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., ), where pre-trained models exhibit early performance saturation. Chen et al. () discussed this issue for GIN models pretrained on molecular datasets at the $10^{5}$ sample scale. Our findings extend this concept to Geom-GNNs, demonstrating that ossification can occur even at the $10^{3}$ sample scale. Notably, we did not observe ossification effects in the QM9 dataset, despite its orders-of-magnitude-more samples, suggesting that ossification is related to both the number and the correlation among samples.

**Table 3: Performance of models with varying width and training strategy on xxMD-Temporal subsets. Best models are picked based on the force regression in terms of MAE. Only models pre-trained on Denali dataset can positively transfer, while ET receives more benefits from pre-training comparing to ViSNet.**
| Subset | Best Dim | Model | Energy (meV) | Force (meV/Å) |
| --- | --- | --- | --- | --- |
| ViSNet |  |  |  |  |
| Stilbene | 128 | PT-Denali | 321 | 137 |
| Azobenzene | 256 | No PT | 75 | 72 |
| Malonaldehyde | 256 | PT-Denali | 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 | 74 |

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 3 suggests that exposure to non-equilibrium conformations during pre-training enhances the model’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 atomic environment. 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-smoothing. Though there are benefits of scaling on supervised and unsupervised tasks, other aspects become more crucial, such as addressing data label uncertainty and employing active token mixing to mitigate the information bottleneck. 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, we only explored denoising pre-training objective, while future efforts include expanding the pre-training scope to include molecular topology and exploring co-pre-trained models for applications such as conditional molecule optimization, protein design, and protein-protein interaction.

## Appendix A Appendix

## Appendix B Equivariant graph neural networks

Equivariant Geom-GNNs are essential in chemistry and biology for leveraging spatial data like atomic positions 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))=T(f(x))$. Geom-GNNs using invariant scalars (Schütt et al., ; Gasteiger et al., ) can easily achieve roto-invariance, but they fall short in predicting general equivariant properties other than derivatives with respect to the input coordinates. Group-equivariant Geom-GNNs GNNs (Batzner et al., ; Batatia et al., ; Musaelian et al., ) utilize spherical harmonics and group representation theory. However, these architectures are expressive yet memory-intensive and computationally expensive. The emergence of vector-based equivariant Geom-GNNs GNNs, such as (Schütt et al., ; Thölke & De Fabritiis, ; Wang et al., ), has shown promising results while scaling more efficiently.

## Appendix C Neural scaling law on language

Neural Scaling Laws describe empirically-derived 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 et al. () identified key power-law relationships between the pre-training loss $L$ and several variables: $L\propto N^{-\alpha},\quad L\propto D^{-\beta},\quad L\propto C^{-\gamma}$ where $N$ is the number of model parameters, $D$ is the dataset size, and $C$ is the amount of compute. These laws have elucidated performance disparities between different neural architectures, such as Long Short-Term Memory (LSTM) networks (Hochreiter & Schmidhuber, ) and Transformers (Vaswani et al., ), particularly in terms of pre-training loss. Subsequent work by Hoffmann et al. () and Muennighoff et al. () further refined these relationships, considering factors like optimal batch size, compute-optimal training, and data constraint.

## Appendix D Dataset summary

**Table 4: Overview of datasets used for pretraining and downstream tasks, detailing their purpose, molecular domains, structural types (equilibrium vs. non-equilibrium), labeled properties, the number of samples, and the method of data splitting.**
| Purpose | Name | Domain | Structure | Labels | #samples | Split |
| --- | --- | --- | --- | --- | --- | --- |
| Pretraining | PCQM4Mv2(^†^††[https://ogb.stanford.edu/docs/lsc/pcqm4mv2/](https://ogb.stanford.edu/docs/lsc/pcqm4mv2/)) | Small molecules | Equilibrium | Conformation | 3,378,606 | N/A |
|  | OrbNet-Denali(^‡^‡‡[https://figshare.com/articles/dataset/OrbNet_Denali_Training_Data/14883867](https://figshare.com/articles/dataset/OrbNet_Denali_Training_Data/14883867)) | Molecules complexes | Non-equilibrium | Conformation | 2,383,351 | N/A |
| Downstream | QM9(^§^§§[https://pytorch-geometric.readthedocs.io/en/latest/generated/torch_geometric.datasets.QM9.html](https://pytorch-geometric.readthedocs.io/en/latest/generated/torch_geometric.datasets.QM9.html)) | Small molecules | Equilibrium | Conformation, quantum properties | 116,038 | Scaffold |
|  | xxMD(^¶^¶¶[https://github.com/zpengmei/xxMD](https://github.com/zpengmei/xxMD)) | Small molecules | Non-equilibrium | Conformation, potential, force | Varies | Temporal |
|  | Fold(^∥^∥∥[https://github.com/phermosilla/IEConv_proteins#download-the-preprocessed-datasets](https://github.com/phermosilla/IEConv_proteins#download-the-preprocessed-datasets)) | Proteins | Equilibrium | Conformation | 16,292 | Scaffold |
|  | Alanine Dipeptide(^**^****[https://markovmodel.github.io/mdshare/ALA2/#alanine-dipeptide](https://markovmodel.github.io/mdshare/ALA2/#alanine-dipeptide)) | Mini peptide | Non-equilibrium | Conformation | 3 x 250 ns trajectories | Random |
|  | Pentapeptide(^††^††††[https://markovmodel.github.io/mdshare/pentapeptide/#peptide](https://markovmodel.github.io/mdshare/pentapeptide/#peptide)) | Mini peptide | Non-equilibrium | Conformation | 25 x 500 ns trajectories | Trajectory |
|  | $\lambda$6-85(^‡‡^‡‡‡‡[https://exhibits.stanford.edu/data/catalog/mh050cw6709](https://exhibits.stanford.edu/data/catalog/mh050cw6709)) | Protein | Non-equilibrium | Conformation | 500 x trajectories with varying time | Trajectory |

## Appendix E Uncertainty in QM9 labels

Figure: Figure 5: HOMO-LUMO gap of ’gdb1’, ’gdb2’ and ’gdb3’ predictied by varying combinations of common basis sets and density functionals. Different combinations showed considerable variance.
Refer to caption: extracted/5961044/Figures/qm9_dft.png

As shown in Figure 5, we scanned combinations of five commonly used functionals (Zhao & Truhlar, ; Haoyu et al., ; Adamo & Barone, ; Lee et al., ) and basis sets (Weigend, ; Ditchfield et al., ; Dunning Jr, ), resulting in 25 distinct computational methods using PySCF package (Sun et al., ). The standard deviations of the HOMO-LUMO gap predictions were substantial: 12571$\pm$1775, 8411$\pm$1626, and 9333$\pm$1758 meV for gdb1, gdb2, and gdb3, respectively. Such uncertainty from DFT calculations has been widely studied in chemistry community (Zhang & Musgrave, ), especially the accuracy of B3LYP functional used in QM9 dataset is non-ideal per current standard (Zhao & Truhlar, ; Goerigk & Grimme, ; Mardirossian & Head-Gordon, ).

## Appendix F Hyperparameters summary

**Table 5: Hyperparameters used for different models used for pre-training.**
| Hyperparameter | ET and ViSNet |
| --- | --- |
| Hidden dimension | 64,128,256,384,512 |
| # MP layers | 2-10 |
| # Attention heads | 8 |
| # RBFs | 32 |
| Layernorm | 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), 0.2 (Denali) |

**Table 6: Hyperparameters used for different models in the supervised regression experiments using QM9 and xxMD datasets. In xxMD experiments, we set the weight on the force and energy as 100:1.**
| Hyperparameter | QM9 | xxMD |
| --- | --- | --- |
| Hidden dimension | 64,128,256,384,512 | 64,128,256,384 |
| # MP layers | 6 |  |
| # Attention heads | 8 |  |
| # RBFs | 32 |  |
| Layernorm | Max-Min, Whitened | Max-Min, None |
| Radius cutoff | 5 |  |
| Epochs | 500 | 200 |
| Early Stopping (epochs) | 10 | 20 |
| Batch size | 50 | 12 |
| Learning rate | 5e-5/1e-4/5e-4 (from-scratch) | 1e-4/5e-4 (from-scratch) |
| Optimizer | Adam |  |
| Scheduler | ReduceLROnPlateau |  |
| Scheduler patience (epoch) | 8 |  |

**Table 7: Hyperparameters used for different models in the VAMP experiments.**
| Hyperparameter | Ala2 | Pentapeptide | $\lambda$6-85 |
| --- | --- | --- | --- |
| Base model | ViSNet (PCQM) |  |  |
| Hidden dimension | 64,128,256,384 |  |  |
| Output dimension | 6 | 1-10 | 6 |
| # MP layers | 6 |  |  |
| # Attention heads | 8 |  |  |
| # RBFs | 32 |  |  |
| Layernorm | Max-Min |  |  |
| Radius cutoff | 5 |  |  |
| Token mixer | None | None | None, MLP-Mixer, Transformer |
| Epochs | 10 | 50 | 50 |
| Early Stopping (train,val batches) | 100,5 |  |  |
| Batch size | 5000 |  |  |
| Learning rate | 5E-04 | 2E-04 | 2E-04 |
| Optimizer | Adam |  |  |

**Table 8: Hyperparameters used for fold classification experiments.**
| Hyperparameter | ProNet-Amino-Acid/ViSNet |
| --- | --- |
| 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 |
| Coord noise | TRUE |
| Rep. noise | FALSE |
| Res. Mask | TRUE |
| Dropout | 0.4 (ProNet-AA), 0.5 (ProNet-Vis) |

## Appendix G Illustration of denoising results

Figure: Figure 6: Scaling behavior of ViSNet and ET models on PCQM and Denali datasets. The plots show the denoising Mean Squared Error (MSE) against the number of 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.
Refer to caption: extracted/5961044/Figures/pt_scaling.png

Figure: 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.
Refer to caption: extracted/5961044/Figures/scan_ar_e0.png

Figure: Figure 8: Influence of the aspect ratio on the self-supervised denoising pre-training task using the ViSNet model. With similar amount of parameters, pre-training performance improves as the aspect ratio decreases.
Refer to caption: extracted/5961044/Figures/scan_ar_param.png

Figure: Figure 9: Influence of the radius cutoff distance on 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 on both datasets demonstrate same optimal radius cutoff at 5 Å. Deviating further from the optimal cutoff, the penalty of models pre-trained on Denali dataset is more severe than the PCQM.
Refer to caption: extracted/5961044/Figures/scan_radius.png

## Appendix H Illustration of QM9 results

Figure: Figure 10: Performance comparison of ViSNet and ET models on Scaffold-QM9 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 to 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.
Refer to caption: extracted/5961044/Figures/qm9_dim.png

Figure: Figure 11: Illustration of the performance of PCQM pre-trained ViSNet models with varying width on Random and Scaffold-QM9. The plots indicate different scaling behaviors with respect to the model parameters. Random-QM9 can always benefit from scaling the model size while the Scaffold-QM9 cannot.
Refer to caption: extracted/5961044/Figures/qm9_split_dim.png

## Appendix I Illustration of xxMD results

Figure: Figure 12: Performance comparison of ViSNet models on the randomly split stilbene subset of xxMD-DFT 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-equilibrium structures). This experiment illustrates: (1) we observe an intersection point beyond $10^{3}$ samples which pre-training may actually hinder performance compared to training from scratch. This suggests a complex relationship between pre-training benefits and downstream dataset size. (2) in data-abundant regimes (with more than $10^{3}$ samples), larger models generally demonstrate superior performance, aligning with typical scaling 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-training datasets.
Refer to caption: extracted/5961044/Figures/xxmd_sti_rand_ef.png

Figure: Figure 13: Comparison of Energy MAE and Force MAE using Temporal-xxMD-DFT dataset across different systems for various architectures and training configurations. Each subplot represents a 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.
Refer to caption: extracted/5961044/Figures/xxmd_temp_ef.png

## Appendix J VAMPNet Objective and Illustration of more VAMP results

VAMPNet is a deep-learning-based objective that learns the projection from high-dimensional feature space to low-dimensional latent space while preserving the slow dynamics of a given time series data. Consider a molecular system characterized by a high-dimensional configuration space, $\mathcal{X}$. Let $\{x_{t}\}_{t=1}^{T}$ denote a trajectory from a molecular dynamics simulation, where $x_{t}\in\mathcal{X}$ represents the molecular configuration at time $t$. For general non-stationary and non-revsersible dynamics, the Koopman operator $\mathcal{K}_{\tau}$ describes the time evolution of an observable $\psi(x)$ over a time lag $\tau$:

$$ $\mathcal{K}_{\tau}\psi(x_{t})=\mathbb{E}[\psi(x_{t+\tau})\mid x_{t}].$ $$

To approximate the Koopman operator in practice, we used linear superposition ansatz as a finite set of basis functions (features).

$$ $f(x)\approx\sum_{i=1}^{N}c_{i}\chi_{i}(x),$ $$

where $c_{i}$ are coefficients, and $\chi_{i}(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 as:

$$ $K=C_{00}^{-1}C_{01}.$ $$

where the instantaneous covariance matrices $C_{00}$ and $C_{11}$, as well as the time-lagged covariance matrix $C_{01}$, are computed as follows:

$$ $C_{00}:=\frac{1}{T-\tau}\sum_{t=0}^{T-\tau}\left(\chi(t)-\mu_{0}\right)\left( \chi(t)-\mu_{0}\right)^{\top}$ $$

$$ $C_{11}:=\frac{1}{T-\tau}\sum_{t=\tau}^{T}\left(\chi(t)-\mu_{1}\right)\left( \chi(t)-\mu_{1}\right)^{\top}$ $$

$$ $C_{01}:=\frac{1}{T-\tau}\sum_{t=0}^{T-\tau}\left(\chi(t)-\mu_{0}\right)\left( \chi(t+\tau)-\mu_{1}\right)^{\top}$ $$

The mean vectors $\mu_{0}$ and $\mu_{1}$ are computed from all data excluding the last $\tau$ steps and the first $\tau$ steps of every trajectory, respectively:

$$ $\mu_{0}:=\frac{1}{T-\tau}\sum_{t=0}^{T-\tau}\chi(t)$ $$

$$ $\mu_{1}:=\frac{1}{T-\tau}\sum_{t=\tau}^{T}\chi(t)$ $$

To refine this approximation, the half-weighted Koopman matrix $\bar{K}$ is introduced, which normalizes the covariance matrices using their inverse square roots:

$$ $\bar{K}=C_{00}^{-1/2}C_{01}C_{11}^{-1/2}.$ $$

The matrix $\bar{K}$ is significant as it encodes the optimal reduced model for the dynamical system, particularly preserving the slowest modes. The singular value decomposition (SVD) of $\bar{K}$ is then performed to extract the primary components of the dynamics:

$$ $\bar{K}=U^{\prime}SV^{\prime\top},$ $$

where $U^{\prime}$ and $V^{\prime}$ represent the left and right singular vector matrices, and $S$ is 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 $\psi(t)$ and right singular functions $\phi(t)$:

$$ $\psi(t)=U^{\prime\top}C_{00}^{-1/2}(\chi(t)-\mu_{0}),\quad\phi(t)=V^{\prime \top}C_{11}^{-1/2}(\chi(t)-\mu_{1}).$ $$

The VAMP-2 score, which is a measure of the quality of the dynamical model in preserving the slow dynamics, is directly related to the singular values of the half-weighted Koopman matrix $\bar{K}$:

$$ $\text{VAMP-2}=\|C_{00}^{-1/2}C_{01}C_{11}^{-1/2}\|_{F}^{2}=\sum_{i=1}^{d} \sigma_{i}^{2}(\bar{K}),$ $$

where $\|\cdot\|_{F}$ denotes the Frobenius norm, and $\sigma_{i}(\bar{K})$ are the singular values of the half-weighted Koopman matrix $\bar{K}$. Then the VAMP-2 score can be used to optimize the network, which is equivalent to finding a representation $\psi$ such that the corresponding Koopman matrix captures the dominant dynamical modes.

Figure: Figure 14: Characteristic structures of pentapeptide in the trajectories visualized on the three-dimensional projection of the learned 6-state VAMP coordinates.
Refer to caption: extracted/5961044/Figures/pp.png

Figure: Figure 15: Three-dimensional projections of learned singular vectors in descending order from the 6-state VAMP model. Lag time: 0.5 ns.
Refer to caption: extracted/5961044/Figures/pp_cvs.png

Figure: Figure 16: Projection of the first VAMP singular vector (Lag time: 50 ns) on 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 al. () indicates 2-3 states on those common collective variables, and the results are comparable. We randomly drew 500 trajectories out of the original dataset to train the VAMP head.
Refer to caption: extracted/5961044/Figures/lamda6_coord.png

## Appendix K Illustration of FOLD classification results and embeddings

In this study, we explore whether graph embeddings can 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 deteriorating 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 HomologyTAPE dataset, as shown in Figures 17 and 18. 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 et al. (). Additionally, we noticed multiple preprocessing errors in the provided dataset when visualizing characteristic structures in the graph embedding space in Figure 19. We followed the Wang et al. () to use the pre-processed datasets from Hermosilla et al. () (available at [https://github.com/phermosilla/IEConv_proteins](https://github.com/phermosilla/IEConv_proteins)). These preprocessing errors are difficult to detect in feature spaces defined by dihedral and torsion angles, but are more apparent in the graph embedding space, which could negatively affect regression accuracy.

It is important to note that projections based on 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. It is still a 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 to untrained convolutional neural networks (CNNs) in vision tasks, as demonstrated by Ulyanov et al. (). However, the features derived from untrained networks are generally sparse and less interpretable.

Figure: Figure 17: Two-dimensional t-SNE visualization of residues in the validation subset of HomologyTAPE 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 a 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).
Refer to caption: extracted/5961044/Figures/tsne_proj.png

Figure: Figure 18: Two-dimensional t-SNE visualization of residues in the validation subset of HomologyTAPE 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 embedding from ViSNet with a 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).
Refer to caption: extracted/5961044/Figures/tsne_cluster.png

Figure: Figure 19: Visualization of characteristic structures for histidine (His) residues using two-dimensional t-SNE projections of pre-trained graph embeddings. The plot shows that structurally similar residues are 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.
Refer to caption: extracted/5961044/Figures/clustering.png

Figure: 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 architecture with the pre-trained GNN features for transfer (GNN-PRONET). It 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 transferability of pre-trained all-atom graph representations to a distinct architectural framework and task domain. We reproduced the original ProNet results with multiple runs to establish a baseline for comparison. There are minor discrepancies between our reproduced values and those reported in the original paper, presumably due to normal experimental variance with different hyperparameters. As illustrated in plots, both models stopped to improve at 256 width on the fold split, while increasing the model size can continue benefit the superfamily split.
Refer to caption: extracted/5961044/Figures/fold.png

## Peer Reviews

confidence: 4
contribution: 3
initial_rating: 5
presentation: 2
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?
soundness: 2
strengths: The paper continues an important line of work in exploring the utility of pre-training and scaling GNNs for molecular learning tasks. Unlike other dominant areas of deep learning, all-atom molecular representation learning relies on GNNs and does not directly benefit from advances in scaling sequence-based models.
The authors explore model size, aspect ratio, nearest neighbor cutoff radius, and architecture to provide a comprehensive look into the scaling behavior of molecular GNNs.
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.
weaknesses: Like most molecular GNN works, the authors are limited by the available evaluations (e.g., QM9). xxMD is an interesting evaluation, but these datasets are limited to a small set of specific molecules. QM9 with B3LYP in particular is not informative, and the authors might consider newer benchmarks like POLARIS or a subset of the Therapeutic Data Commons to strengthen their evaluations.
The paper does not offer a clear and concise summary of recommendations based on the empirical findings, which is essential to achieving the stated aim of inspiring practitioners to rethink GNN training.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
questions: 1. The paper claims to demonstrate zero-shot transfer learning using pre-trained Geom-GNNs. However, the described method seems closer to downstream task fine-tuning, given that a separate network is trained with the VAMP score objective. Could the authors clarify how this approach qualifies as zero-shot transfer rather than fine-tuning?

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

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

[1] Cui, Taoyong, et al. "Geometry-enhanced pretraining on interatomic potentials." Nature Machine Intelligence 6.4 (2024): 428-436.
[2] Zhou, Gengmo, et al. "Uni-mol: A universal 3d molecular representation learning framework." (2023).

---

confidence: 3
contribution: 3
initial_rating: 8
presentation: 4
questions: as Weaknesses.
soundness: 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.
2. **Clarity and Structure**: The article is well-organized, with clear presentation and summary of experiments and viewpoints, facilitating reader comprehension.
3. **Robust Experimentation**: The experimental design is thorough, effectively supporting the authors’ conclusions.
4. **Exploration of Zero-Shot Transfer Capability**: The investigation into the zero-shot transfer ability of Geom-GNNs is intriguing, with experiments indicating their potential as excellent zero-shot transfer learners.
5. **Pre-training Insights**: Through extensive denoising pre-training tasks, valuable experiences have been gained regarding the pre-training of Geom-GNNs, including aspects such as model width, depth, aspect ratio, and the cutoff radius in geometric atomic graph construction, providing rich guidance for pre-training.
6. **Advancement of Unified Architecture**: Given the widespread attention and efforts in the research of all-atom Geom-GNNs, this study effectively inspires researchers to reconsider the design of Geom-GNN architectures and the adjustment of training strategies, thereby promoting the development of a unified Geom-GNN architecture.
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.
weaknesses: 1. In Figure 6, which explores different model widths, even though the x-axis represents the total number of parameters (with model depth held constant), it would be more beneficial to indicate the model width for each point in the legend to enhance result presentation. Similarly, in Figures 4 and 7, which demonstrate the impact of model depth, using the legend to specify the exact number of layers might be more effective. In general, clear legends are always advantageous.
2. The comparison between models trained from scratch and those fine-tuned in Section 6.1 could be more comprehensive if extended to include model depth. Previous discussions (albeit in the context of pre-training) have presented certain viewpoints, and it is anticipated that these would also have significant effects during fine-tuning.

---

confidence: 3
contribution: 1
initial_rating: 3
presentation: 1
questions: - Why is the VAMP task considered zero-shot?
- Appendix D: Why is there a difference between datasets for pretraining and downstream? According to Table 2, QM9 there seem also experiments with pretraining. Are models in Table 3 not fine-tuned?
soundness: 2
strengths: relevant research topics in the paper:
- pretraining of (geometric) GNNs
- scaling behavior of GNNs
- oversmoothing, underreaching

experiments:
- experiments considered not only random splitting, but also scaffold and temporal splits
summary: In their work the authors empirically studied properties of pre-trained geometric GNNs (i.e., GNNs with coordinates attached).
They especially seemed to consider pretraining via denoising as introduced by Zaidi et al. (2022).
The authors consider several downstream tasks: molecular dynamics (kinetics modeling), fold classification, prediction of quantum-chemical properties, and, molecule conformations
They study properties such as power-law scaling and find that the geomtric GNNs do not follow that on the pre-training task.
The authors conclude that geometric GNNs "saturate early and are bounded by issues such as under-reaching and over-smoothing" and further that "other aspects become more crucial, such as addressing data label uncertainty and employing active token mixing to mitigate the information bottleneck."
weaknesses: In general the largest problem to my opinion is that the authors lack to specify a clear research goal. Methodologically there seems not too much novelty.
But if it is primarily an empirical research paper it seems very important to me that the research goal is clearly defined and dependent on that authors have
to reason why they are selecting the problems they look at and why they are selecting the methods they compare to. If a new benchmark dataset is used (as e.g., scaffold QM9), then empirical evaluations/comparisons should take into account at least some previously suggested methods (which are ideally complementary to each other to better see the potentials of each of the methods on the new benchmark and being able to compare to what authors say).
\
\
Details:
\
\
**hard to see the novelty in the paper**:\
Fact that pretraining is useful was already found e.g. by Zaidi et al. (2022).\
No new methodology seems to be suggested (or is "token mixing" the new method)?\
Study of scaling behavior might be interesting and to a certain degree novel for geometric GNNs, but no follow-up investigations seemed to be employed to draw formal connections to GNN properties like oversmoothing, underreaching, etc.
\
\
**paper is structured in a strange way, which makes it hard to follow the paper**:\
It is actually hard to understand what the research aim of the authors was.\
chapter 4, 5, and, 6, actually seem to be about empiricial experiment results.\
The setups are partly however already explained in chapter 3.\
chapter 4 and 6 show performances on problems\
chapter 5 however studies power-law behavior and other ablations.\
In sum it gets hard to follow the paper. Better reasoning why some experiments are applied at all (with respect to the general research goal) and why they are done as they are done is necessary (e.g., why which methods are selected for comparison or why which dataset is used).
\
\
**experiments**:\
although there are some good points as mentioned in strengths (such as the splits), it will get hard to understand how large the impact of pretraining really is, as the authors only test very few method once with pretraining and once without pretraining.\
\
also the there is no good argument in the paper, why the authors exactly compare to those methods they selected to compare to 
e.g.:\
for molecular dynamics, why don't they compare to Timewarp\
for QM9 and xxMD they could e.g. compare to "E(n) Equivariant Graph Neural Networks", "SE(3)-Transformers", DimeNet++, MACE, etc.\
An option to get more impression on the significance of the author's results would be to additionally compare to standard QM9, etc. (where there are also a lot of method comparisons out).
\
\
**minor points**:\
"token mixing" not defined, but heavily used\
grammatical errors/typos: "In silico molecular computation and simulation are indispensable tool in..."

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

Decision: Accept (Poster)