VN-EGNN: Vector–Node Equivariant GNN

• VN-EGNN is a vector–node equivariant graph neural network that integrates virtual nodes and multiple vector channels to capture global and latent geometric features.
• It enhances standard EGNN by addressing oversquashing and improving message passing, making it effective for protein binding site identification and modeling dynamics.
• Empirical benchmarks show VN-EGNN achieves state-of-the-art performance with minimal runtime overhead, proving its practical benefits in complex spatial tasks.

VN-EGNN, or "Vector–Node Equivariant Graph Neural Network," encompasses two related but distinct architectures extending the E(n)-Equivariant Graph Neural Network (EGNN) framework: (1) an E(3)-equivariant GNN with a global set of virtual nodes for protein binding site identification (Sestak et al., 2024), and (2) an E(n)-equivariant GNN endowing each node with multiple equivariant vector channels to increase its expressive power for general physical systems modeling (Levy et al., 2023). Both variants address fundamental limitations of the standard EGNN by enhancing message passing and geometric representation, while preserving equivariance.

1. Extension of the EGNN Paradigm

The foundational EGNN architecture, introduced by Satorras et al., operates on spatial graphs where node representations consist of coordinates $x_i\in\mathbb R^n$ and hidden features $h_i\in\mathbb R^d$, with message passing steps designed for E(n)-equivariance. A standard EGNN layer updates nodes using neighbor-relative messages: $$m_{ij} = \phi_e(h_i, h_j, \Vert x_i - x_j \Vert^2), \quad m_i = \sum_{j \in \mathcal N(i)} m_{ij}$$

$$h_i' = \phi_h(h_i, m_i), \quad x_i' = x_i + \sum_{j \in \mathcal N(i)} (x_i - x_j)\phi_x(m_{ij})$$

These update rules guarantee equivariance to Euclidean transformations (rotations, translations, reflections), rendering the architecture suitable for modeling spatial physical systems (Levy et al., 2023).

VN-EGNN generalizes this principle in two directions:

• By introducing virtual nodes with learnable coordinates and feature embeddings, to capture non-local geometric entities such as protein binding pockets (Sestak et al., 2024).
• By upgrading each node's coordinate to a set of $m$ equivariant "vector channels," enabling richer latent geometric representations per node (Levy et al., 2023).

2. Virtual-Node Augmented VN-EGNN for Binding Site Identification

The first VN-EGNN variant targets protein binding site identification, where proteins are modeled as spatial graphs with nodes (e.g., residues or atoms) linked by spatial proximity: $$\mathcal{P} = (\{(x_i, h_i)\}_{i=1}^N, \mathcal{E})$$, with edge set $\mathcal{E}$ connecting $k$-nearest neighbors.

A small, global set of $K$ virtual nodes $(z_k, v_k)$ is added, each with coordinates $h_i\in\mathbb R^d$0 and embeddings $h_i\in\mathbb R^d$1. These virtual nodes are connected to all physical nodes but not to each other, ensuring that any physical node can exchange information globally via a two-hop path (physical → virtual → physical). The virtual nodes are designed to migrate in coordinate space toward the centers of binding pockets and to accumulate pocket-specific features during inference (Sestak et al., 2024).

Each VN-EGNN layer consists of a three-phase heterogeneous message passing procedure:

1. Atom→atom: Standard EGNN-style neighbor updates among physical nodes with current atomic geometry and features.
2. Atom→virtual: Physical nodes update virtual node states using the latest atomic representations.
3. Virtual→atom: Updated virtual nodes broadcast their information back to the physical nodes.

The embedding and coordinate update rules for each phase are defined analogously to EGNN but use separate MLPs for each phase. This heterogeneous message passing ensures that virtual nodes have access to the latest atomistic features and vice versa after each layer.

3. Vector-Channel VN-EGNN for General Physical Systems

This second VN-EGNN variant is a minimal, computationally efficient extension of EGNN, endowing each node with $h_i\in\mathbb R^d$2 coordinate-like "vector channels": $h_i\in\mathbb R^d$3 Group actions (e.g., rotations in $h_i\in\mathbb R^d$4) act on the spatial dimension, not across channels. The update steps become channelwise:

• Compute channel-wise relative displacements $h_i\in\mathbb R^d$5
• Calculate channel norms for each channel
• Message passing MLPs take concatenated features $h_i\in\mathbb R^d$6, $h_i\in\mathbb R^d$7, channel norms, and edge features
• Update $h_i\in\mathbb R^d$8 using coordinate-mixing weights $h_i\in\mathbb R^d$9 and matrix multiplication over channels

The hidden-feature update remains as in standard EGNN. All updates ensure E(n)-equivariance, and by using $$m_{ij} = \phi_e(h_i, h_j, \Vert x_i - x_j \Vert^2), \quad m_i = \sum_{j \in \mathcal N(i)} m_{ij}$$0 at input/output boundaries, the original EGNN interface is preserved (Levy et al., 2023).

This structure allows each node to carry multiple latent vector fields (e.g., one channel for position, others for angular momentum, spin, etc.), significantly boosting expressivity in representing complex dynamical systems.

4. Training Objectives and Binding Site Readout

In the binding site VN-EGNN (Sestak et al., 2024), a multi-term loss combines:

• Segmentation loss (node-level Dice or binary cross-entropy) to label each node as pocket or non-pocket:

$$m_{ij} = \phi_e(h_i, h_j, \Vert x_i - x_j \Vert^2), \quad m_i = \sum_{j \in \mathcal N(i)} m_{ij}$$1

where $$m_{ij} = \phi_e(h_i, h_j, \Vert x_i - x_j \Vert^2), \quad m_i = \sum_{j \in \mathcal N(i)} m_{ij}$$2 for final-layer features $$m_{ij} = \phi_e(h_i, h_j, \Vert x_i - x_j \Vert^2), \quad m_i = \sum_{j \in \mathcal N(i)} m_{ij}$$3.

• Binding-site-center loss: Assigns each ground-truth pocket center to its closest predicted virtual node coordinate $$m_{ij} = \phi_e(h_i, h_j, \Vert x_i - x_j \Vert^2), \quad m_i = \sum_{j \in \mathcal N(i)} m_{ij}$$4:

$$m_{ij} = \phi_e(h_i, h_j, \Vert x_i - x_j \Vert^2), \quad m_i = \sum_{j \in \mathcal N(i)} m_{ij}$$5

• Self-confidence calibration: A small MLP on each $$m_{ij} = \phi_e(h_i, h_j, \Vert x_i - x_j \Vert^2), \quad m_i = \sum_{j \in \mathcal N(i)} m_{ij}$$6 predicts a confidence $$m_{ij} = \phi_e(h_i, h_j, \Vert x_i - x_j \Vert^2), \quad m_i = \sum_{j \in \mathcal N(i)} m_{ij}$$7, trained against the true spatial error between predicted and actual pocket centers.

Predictions are read out as both nodewise pocket probabilities and a set of $$m_{ij} = \phi_e(h_i, h_j, \Vert x_i - x_j \Vert^2), \quad m_i = \sum_{j \in \mathcal N(i)} m_{ij}$$8 candidate pocket centers. After mean-shift clustering, the $$m_{ij} = \phi_e(h_i, h_j, \Vert x_i - x_j \Vert^2), \quad m_i = \sum_{j \in \mathcal N(i)} m_{ij}$$9 highest-confidence predictions are retained as binding site centers.

5. Empirical Performance and Benchmarks

The virtual-node VN-EGNN sets state-of-the-art benchmarks on established protein binding site datasets, outperforming prior methods such as EquiPocket, Fpocket, and P2Rank (Sestak et al., 2024). DCC (distance-to-known-center) and DCA (distance-to-closest-atom) success rates at 4Å are summarized as follows:

Dataset	VN-EGNN DCC	EquiPocket DCC	VN-EGNN DCA	EquiPocket DCA
COACH420	0.605 (±0.009)	0.423	0.750 (±0.008)	0.656
HOLO4K	0.532 (±0.021)	0.337	0.659 (±0.026)	0.662
PDBbind2020	0.669 (±0.015)	0.545	0.820 (±0.010)	0.721

This demonstrates robust gains, especially under strong domain shifts. The generalized vector-channel VN-EGNN also achieves improved accuracy across multiple physical modeling tasks, including solar-system N-body forecasting, charged-particle interactions, and molecular property prediction (QM9), with minimal added runtime and parameter cost. Optimal channel count $$h_i' = \phi_h(h_i, m_i), \quad x_i' = x_i + \sum_{j \in \mathcal N(i)} (x_i - x_j)\phi_x(m_{ij})$$0 is task-dependent, and for $$h_i' = \phi_h(h_i, m_i), \quad x_i' = x_i + \sum_{j \in \mathcal N(i)} (x_i - x_j)\phi_x(m_{ij})$$1, parameter and runtime inflation are modest (<10%) (Levy et al., 2023).

6. Architectural and Implementation Choices

Key architectural considerations for the virtual-node VN-EGNN (Sestak et al., 2024):

• Number of layers: 5 complete VN-EGNN layers (each comprising AA→AV→VA steps)
• Virtual nodes: $$h_i' = \phi_h(h_i, m_i), \quad x_i' = x_i + \sum_{j \in \mathcal N(i)} (x_i - x_j)\phi_x(m_{ij})$$2 initialized on a Fibonacci-lattice sphere of radius matching protein extent, with random rotation per sample to enforce invariance
• Feature dimension: 100, with pre-trained ESM-2 embeddings linearly projected
• Activation/Normalization: SiLU, layer normalization, dropout
• Optimizer: AdamW at $$h_i' = \phi_h(h_i, m_i), \quad x_i' = x_i + \sum_{j \in \mathcal N(i)} (x_i - x_j)\phi_x(m_{ij})$$3 with scheduler, Huber loss for coordinate prediction, coordinate normalization by divisor 5
• Clustering: Mean-shift, to collapse redundant virtual node predictions

Vector-channel VN-EGNNs are implemented so that first/last layers have $$h_i' = \phi_h(h_i, m_i), \quad x_i' = x_i + \sum_{j \in \mathcal N(i)} (x_i - x_j)\phi_x(m_{ij})$$4, ensuring drop-in EGNN compatibility; hidden layers promote to $$h_i' = \phi_h(h_i, m_i), \quad x_i' = x_i + \sum_{j \in \mathcal N(i)} (x_i - x_j)\phi_x(m_{ij})$$5 channels. MLPs for message passing and channel mixing are parameter-shared; increases in parameter count and computation are $$h_i' = \phi_h(h_i, m_i), \quad x_i' = x_i + \sum_{j \in \mathcal N(i)} (x_i - x_j)\phi_x(m_{ij})$$6 per layer but negligible for practical $$h_i' = \phi_h(h_i, m_i), \quad x_i' = x_i + \sum_{j \in \mathcal N(i)} (x_i - x_j)\phi_x(m_{ij})$$7 (Levy et al., 2023).

7. Theoretical and Practical Implications

VN-EGNN addresses oversquashing in deep GNNs by integrating virtual nodes, which can efficiently accumulate and disseminate global geometric information without deep message-passing chains. In the general physical modeling case, vector channels can be understood as "latent vector fields" per node, facilitating the modeling of multi-body interactions and vector-valued observables. The approach preserves E(n)-equivariance by design, ensuring physical symmetries are respected throughout learning and inference.

A plausible implication is that these extensions can be broadly adopted in molecular and physical sciences whenever geometric entities or higher-order interactions must be encoded efficiently. Benchmarks suggest that even modest increases in vector channels or the inclusion of virtual nodes can deliver substantial empirical gains for suitably complex spatial prediction tasks (Sestak et al., 2024, Levy et al., 2023).