Multi-Scale Protein Structure Modelling with Geometric Graph U-Nets

TL;DR

Geometric Graph U-Nets address the challenge of representing multi-scale protein structure by progressively coarsening and refining geometric graphs. The paper develops a formal GWL/IGWL expressivity framework showing pooling can preserve or boost discriminative power while enabling global context integration. Empirically, Geometric U-Nets outperform invariant and equivariant baselines on protein fold classification, demonstrating the value of multi-scale representations. The authors also provide open-source code and outline future directions for biomolecular complexes and multi-state design.

Abstract

Geometric Graph Neural Networks (GNNs) and Transformers have become state-of-the-art for learning from 3D protein structures. However, their reliance on message passing prevents them from capturing the hierarchical interactions that govern protein function, such as global domains and long-range allosteric regulation. In this work, we argue that the network architecture itself should mirror this biological hierarchy. We introduce Geometric Graph U-Nets, a new class of models that learn multi-scale representations by recursively coarsening and refining the protein graph. We prove that this hierarchical design can theoretically more expressive than standard Geometric GNNs. Empirically, on the task of protein fold classification, Geometric U-Nets substantially outperform invariant and equivariant baselines, demonstrating their ability to learn the global structural patterns that define protein folds. Our work provides a principled foundation for designing geometric deep learning architectures that can learn the multi-scale structure of biomolecules.

Figures (2)

• Figure 1: Geometric Graph U-Nets for multi-scale protein representation learning. (A) Protein function is governed by a structural hierarchy. Local secondary structures fold into tertiary domains, which assemble into quaternary complexes to perform complex functions like allosteric signaling. Standard GNNs struggle to capture these long-range interactions. (B) The Geometric Graph U-Net is a biologically inspired architecture that encodes a protein graph by progressive coarsening, moving from local representations to global ones. The decoder then refines this information, using residual connections to re-inject high-resolution local information at each level. Thus, the final representation integrates both global context and local geometric details.
• Figure 2: Comparison of Point Pooling and Sparse Pooling layers. Both pooling layers use farthest point sampling (FPS) to select supernode centers from the original graph. (Left) Point Pooling assigns all 1-hop neighbors of each supernode to its cluster, creating overlapping regions that capture local structural motifs. (Right) Sparse Pooling assigns each node to its nearest supernode, partitioning the graph into non-overlapping regions. Both layers aggregate features from assigned nodes to compute supernode representations.

Theorems & Definitions (16)

• Definition 2.1
• Definition 2.2
• Definition 2.3
• Definition 4.1
• Definition 4.2
• Definition 4.3
• Definition 4.4
• Proposition 4.1
• Theorem 4.1
• Lemma 4.1