Technical Summary: PRIME (Protein Representation via Physics-Informed Multiscale Equivariant Hierarchies)

Problem Statement

Proteins are inherently multiscale physical systems where functional properties emerge from coordinated structural organization ranging from atomic interactions to global fold topology. Existing protein representation learning methods typically operate at a single structural resolution (e.g., sequence, atomic, or surface level) or treat different structural sources as parallel modalities integrated via latent alignment or late-stage fusion. These approaches fail to explicitly model the hierarchical relationships between structural scales (e.g., how atoms form residues, which form secondary structures, which form the global fold). Consequently, they treat distinct spatial resolutions as independent semantic views rather than nested levels within a shared physical hierarchy, limiting their ability to capture interactions that emerge across resolutions.

Methodology

The authors introduce PRIME, a unified framework that models proteins as a nested family of five physically grounded structural graphs. The framework operates through two primary stages: hierarchical graph construction and multiscale message passing.

1. Hierarchical Graph Construction

PRIME defines a hierarchy $\mathcal{G} = \{G^{(0)}, G^{(1)}, G^{(2)}, G^{(3)}, G^{(4)}\}$ spanning five levels from fine to coarse resolution:

• Level 0 (Surface): Nodes are triangular faces of the molecular surface mesh.
• Level 1 (Atom): Nodes are heavy atoms.
• Level 2 (Residue): Nodes are amino acid residues.
• Level 3 (Secondary Structure): Nodes are secondary structure elements (SSEs: helices, strands, loops).
• Level 4 (Protein): A single node representing the entire protein.

Connectivity and Assignment:

• Edges: The finest level ($G^{(0)}$) uses $k$-nearest neighbors ($k=8$) on face centroids. Coarser levels derive edges deterministically from the finest level via partition matrices $\Pi^{(\ell)}$, where $A^{(\ell)} = (\Pi^{(\ell)})^\top A^{(\ell-1)} \Pi^{(\ell)}$. This results in weighted adjacency matrices reflecting the number of fine-scale connections between coarse units, encoding spatial proximity without explicit coordinate-based edge construction at every level.
• Assignment Maps: Transitions between levels are governed by deterministic, physics-informed assignment operators $\pi^{(\ell)}$ derived from biochemical principles (e.g., Euclidean distance for surface-to-atom, covalent bonding for atom-to-residue, DSSP algorithm for residue-to-SSE).
• Node Features:
  • Surface: Handcrafted geometric descriptors (area, centroid distance) + embeddings from a pretrained Equivariant Mesh Neural Network (EMNN).
  • Atom: Element type/backbone membership + embeddings from a pretrained Equivariant Graph Neural Network (EGNN).
  • Residue: Amino acid identity, backbone angles, curvature, and torsional angles.
  • SSE: One-hot encoding of type and normalized length.
  • Protein: Global representation from a pretrained ESM-2 model.

2. Multiscale Message Passing

PRIME employs a stack of $N$ hierarchical layers to propagate information bidirectionally. Each layer performs three sequential operations:

1. Intra-level Message Passing: Refines representations at each level independently using graph convolutions with pre-layer normalization, residual connections, and feedforward networks.
2. Bottom-Up Aggregation: Propagates fine-grained details upward by aggregating representations from level $\ell-1$ to $\ell$ using $(\Pi^{(\ell)})^\top$, followed by gated fusion.
3. Top-Down Contextual Refinement: Broadcasts coarse-level context downward from level $\ell$ to $\ell-1$ using $\Pi^{(\ell)}$, followed by symmetric gated fusion.

The gated fusion mechanism selectively incorporates context $c$ into the current representation $z$ via a learnable gate $g = \sigma(W_g [z \| c])$. The final representation is read out from a designated level $\ell^*$ (typically the residue level) for downstream tasks. The framework is designed to be SE(3)-invariant end-to-end, as all input features and operations are coordinate-free or utilize invariant scalar embeddings from equivariant encoders.

Key Contributions

• Unified Framework: PRIME proposes a single framework that models proteins as a nested family of five physically grounded graphs, connected by deterministic assignment operators, enabling explicit modeling of hierarchical physical organization.
• Bidirectional Information Exchange: The method facilitates information flow via bottom-up aggregation (fine-to-coarse) and top-down refinement (coarse-to-fine), allowing each level to encode both its own modality-specific information and contextual signals from all other scales.
• Physics-Informed Inductive Bias: Unlike data-driven graph coarsening, PRIME leverages known biochemical and geometric constraints to define the hierarchy, providing a principled basis for multiscale representation learning.
• Adaptive Readout: The framework supports adaptive cross-attention readouts, allowing the model to autonomously identify and weight the most task-relevant structural resolutions at prediction time.

Experimental Results

PRIME was evaluated on four benchmarks from ProteinWorkshop: Fold Classification, Reaction Class Prediction, Gene Ontology (GO) Prediction, and Protein-Protein Interaction (PPI) Site Prediction.

• Fold Classification: PRIME achieved state-of-the-art accuracy across all splits. Notably, on the challenging Superfamily and Fold splits, it outperformed the strongest geometric GNN baseline (EGNN) by margins of 13.80 and 18.30 points, respectively.
• Reaction Class Prediction: PRIME achieved 84.10% accuracy, surpassing all baselines including ESM (83.10%), suggesting that explicit structural hierarchy provides a superior signal for enzyme function prediction compared to sequence-only models.
• Gene Ontology: The model achieved the best performance on Molecular Function (MF) and Cellular Component (CC) ontologies, which are closely tied to structural motifs and global folds. Performance on Biological Process (BP) was lower, likely due to the annotation's reliance on evolutionary context rather than intrinsic structure.
• PPI Site Prediction: PRIME achieved competitive ROC-AUC scores, demonstrating applicability to node-level tasks, though it did not top the rankings, indicating potential for specialized architectural choices for precise interface prediction.

Ablation Studies:

• Removing any single structural level consistently degraded performance, confirming that each level contributes complementary, non-redundant information.
• The removal of the Secondary Structure (SSE) level caused the most significant drop (e.g., Fold split accuracy dropped from 59.78% to 16.61%), highlighting the centrality of SSE topology in fold discrimination.
• Adaptive cross-attention analysis revealed that the model autonomously assigns higher attention weights to SSE levels for fold classification and adjusts weights based on the task (e.g., higher weight on atomic/residue levels for Reaction Class), validating the necessity of a unified multiscale design.

Significance

The paper claims that PRIME provides a principled foundation for protein representation learning by explicitly accounting for the hierarchical organization of protein structure. By moving beyond flat or parallel multimodal representations, PRIME captures interactions that emerge across structural resolutions. The results demonstrate that modeling proteins as a nested hierarchy of physically grounded graphs yields generalizable representations that outperform existing unimodal and multimodal approaches, particularly in tasks requiring the discrimination of structurally similar folds or the prediction of enzyme functions governed by active site geometry. The framework's ability to autonomously identify task-relevant resolutions further underscores the value of maintaining a full multiscale hierarchy rather than committing to a single fixed resolution.