Geometric Self-supervised Pretraining on 3D Protein Structures Using Subgraphs

Introduces a geometric self-supervised pretraining task that predicts subgraph-centroid-to-global-centroid distances, replacing conventional node or edge masking. Prior graph pretraining largely relies on node/edge attribute masking or contrastive learning (e.g., GraphCL, DGI) requiring multiple views or augmentations; this work instead predicts the Euclidean distance between subgraph geometric centroids and the protein global centroid, discretizing the regression into 10 equal bins and optimizing with cross-entropy loss. The method is implemented on three backbones (ProNet, SchNet, GCN) and three Cα-based featurizations, compared against an edge-distance pretraining baseline within the ProteinWorkshop library, with means and standard deviations over three seeds.

Subgraph-distance pretraining consistently outperforms no pretraining and edge-distance pretraining on multiple downstream protein classification tasks. Relative to edge-distance pretraining, which also learns geometric/distance information, this task is generally higher on Fold, Superfamily, Family, and React; for example, SchNet improves Super-Family by 4.78% with ca_angles and 5.85% with ca_bb. Results come from the Fold dataset (16,712 proteins, 1,195 folds; 12,312 train / 736 validation; test sets of 718 Fold, 1,254 Superfamily, 1,272 Family) and the React dataset (37,428 proteins, 384 EC numbers; 29,215 train / 2,562 validation / 5,651 test), reported as mean and standard deviation over three runs.

Pretraining-task performance correlates with downstream task performance. The authors observe that models with higher pretraining accuracy under ca_angles and ca_bb featurizations generally perform better downstream, drawing an analogy to language modeling where pretraining performance indicates downstream performance. Based on accuracy and cross-entropy loss curves over 10 pretraining epochs for SchNet and GCN (Figures 2 and 3), with schnet_ca_bb showing the lowest final loss.

The method has low computational overhead and adapts to any backbone that can encode 3D protein structures. Subgraph representations are taken from the GNN's final-layer node representations, requiring only one forward pass per graph; subgraph computation can be done once as preprocessing and stored, eliminating the extra overhead. The paper provides a complexity discussion in the methods section and integrates the implementation into the ProteinWorkshop library with released source code.

The result applies to 3D protein graphs with Cα atoms as nodes, k=16 nearest-neighbor edges, amino-acid types as node features, and sequential distances as edge features, pretrained on 542k SwissProt structures from the AlphaFold Database; downstream evaluation is limited to Fold, Superfamily, Family, and React classification, with backbones limited to ProNet, SchNet, and GCN and featurizations limited to ca_base, ca_angles, and ca_bb. The authors state future work will explore different subgraph selection strategies and combinations with additional pretraining tasks, so this pretraining task can be directly used for protein representation learning that benefits from geometric and hierarchical structure awareness, and can serve as a starting point for pretraining other 3D graph models.

The correlation between pretraining-task and downstream performance is currently based on curve observations for SchNet and GCN over 10 epochs, and its stability across different data scales and longer training remains to be tested; subgraph selection is fixed to 2-hop ego networks centered on 10% of amino acids, and the effect of other subgraph selection strategies is not yet explored; performance beyond the classification tasks listed here, and the effect of combining with other pretraining tasks, are left as future work.