Research map / Graph learning / GNN architectures
Improving Equivariant Graph Neural Networks on Large Geometric Graphs via Virtual Nodes Learning
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes FastEGNN, which adds a small ordered set of learnable virtual nodes to EGNN, regularized by an MMD loss, so that sparse or edge-dropped large geometric graphs can still be modeled accurately.
Problem
Equivariant GNNs scale quadratically on large geometric graphs when fully connected, and accuracy drops when sparsified to local graphs for speed.
Approach
Learn an ordered set of E(3)-equivariant virtual nodes that act as global message relays, with channel-specific message passing for distinctiveness and an E(3)-invariant MMD loss aligning virtual and real coordinates for distributedness.
Key results
- On Protein Dynamics, FastEGNN-<3,1.00> with all edges dropped reaches MSE 1.88 vs 2.25 for EGNN with full edges.
- On Water-3D, FastEGNN reaches MSE 3.40e-4 vs 6.00e-4 for EGNN.
- On N-body, FastEGNN yields 28% MSE improvement over EGNN (1.12e-2 vs 1.41e-2) with 0.35 relative inference time at 0.75 edge drop.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| add 3 learned virtual nodes with no edge dropping | EGNN · N-body system | MSE (×10−2): 1.41 -> 1.12 ± 0.04 | Helped |
| add 3 learned virtual nodes with no edge dropping | EGNN · Protein Dynamics | MSE: 2.25 -> 1.82 ± 0.02 | Helped |
| add 3 learned virtual nodes with no edge dropping | EGNN · Water-3D | MSE (×10−4): 6.00 -> 2.49 ± 0.14 | Helped |
| drop 75% of edges while using 3 virtual nodes | FastEGNN · N-body system | Relative Time: 1.00 -> 0.35 | Helped |
| drop all edges while using 3 virtual nodes | FastEGNN · Protein Dynamics | MSE: 2.25 -> 1.88 ± 0.01 | Helped |
| drop all edges while using 3 virtual nodes | FastEGNN · N-body system | MSE (×10−2): 1.12 ± 0.04 -> 9.52 ± 0.02 | Hurt |
| treat virtual nodes as permutation-equivariant global nodes instead of an ordered set (ablation) | FastEGNN · Protein Dynamics | MSE: 1.863 -> 1.888; 1.878 -> 1.967; 1.876 -> 2.121 | Hurt |
| remove the MMD loss (ablation) | FastEGNN · Protein Dynamics | MSE: 1.863 -> 1.882; 1.878 -> 1.883; 1.876 -> 1.923 | Hurt |
| vary the number of virtual nodes from 1 to 3 (ablation) | FastEGNN · N-body system | MSE (×10−2) | Mixed |
| vary the number of virtual nodes from 1 to 3 (ablation) | FastEGNN · Protein Dynamics and Water-3D | MSE | Mixed |
| add virtual node learning to RF | RF · Protein Dynamics | MSE: 2.25 -> 2.07 | Helped |
| add virtual node learning to TFN | TFN · Protein Dynamics | MSE: 2.25 -> 1.84 | Helped |
Surprising findings
- Naively treating virtual nodes as permutation-equivariant global nodes performs worse than the ordered virtual nodes design.
- EGNN with all edges dropped performs close to linear dynamics, while FastEGNN with all edges dropped still beats full-edge EGNN.
Limitations
- Virtual-node count and MMD hyperparameters (C, sigma, lambda) must be tuned per dataset; 1 node best on N-body, 3 on protein and water.
- Fully dropping edges on N-body (p=1.00) degrades MSE to 9.52e-2 for FastEGNN-<3,1.00>, so benefits are dataset dependent.
Experimental setup
Datasets: N-body system Protein MD (AdK) Water-3D
Benchmarks: N-body System Protein MD Water-3D
Metrics: MSE relative inference time
Baselines: EGNN TFN Radial Field SchNet MPNN Linear dynamics SGNN global node MEAN
Builds on: EGNN MMD Villar et al. equivariant point cloud framework
Closest papers by meaning
- Universally Invariant Learning in Equivariant GNNs (NeurIPS 2025 · similarity 0.89)
- Are High-Degree Representations Really Unnecessary in Equivariant Graph Neural Networks? (NeurIPS 2024 · similarity 0.88)
- On Explaining Equivariant Graph Networks via Improved Relevance Propagation (ICML 2025 · similarity 0.88)
- On the Expressive Power of Sparse Geometric MPNNs (ICLR 2025 · similarity 0.87)
- On the Generalization of Equivariant Graph Neural Networks (ICML 2024 · similarity 0.87)
- E(n) Equivariant Topological Neural Networks (ICLR 2025 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.