Equivariant Graph Neural Network (EGNN)

Unlike architectures such as Tensor Field Networks or SE(3)-Transformers, EGNNs achieve equivariance to rotations, translations, and reflections without relying on Clebsch-Gordan tensor products or spherical harmonics. The model maintains a set of scalar features (invariant) and vector features (equivariant) for each node. Equivariance is enforced by updating vector features as a weighted sum of relative position vectors, ensuring that all geometric transformations of the input produce corresponding transformations of the output. This design dramatically reduces computational overhead while preserving physical symmetry constraints.

Each node i in an EGNN is represented by two distinct feature channels:

• Scalar features (h_i): Standard invariant embeddings updated via message passing, encoding chemical element type, local environment, and bond information.
• Vector features (v_i): Equivariant directional embeddings that transform correctly under rotation. These are updated using the difference of positions between interacting nodes.

This dual representation allows the network to reason about both what an atom is and where it is in space, enabling accurate prediction of forces and dipole moments.

The EGNN message function incorporates interatomic distances directly into the edge computation. For each edge between nodes i and j:

• A message is computed from the scalar features of both nodes and the squared distance ||x_i - x_j||^2.
• Scalar features are updated via standard aggregation.
• Vector features are updated as: v_i ← v_i + Σ_j (x_i - x_j) · φ(m_ij), where φ is a learned function.

This formulation ensures that vector updates are linear combinations of relative position vectors, which is the mathematical condition for E(n) equivariance.

A defining characteristic of EGNNs is that atomic coordinates are updated at each layer using the learned equivariant vector features. The position update follows:

x_i ← x_i + v_i

This transforms the EGNN from a static feature extractor into a dynamic geometry refiner. The network learns to iteratively adjust 3D coordinates toward lower-energy conformations or correct binding poses. This mechanism is crucial for tasks like conformer generation, molecular dynamics acceleration, and binding pose prediction, where the final geometry is the primary output.

By avoiding higher-order tensor products, EGNNs achieve O(N) complexity in the number of atoms, compared to the O(N^2) or higher scaling of tensor product-based equivariant networks. Key efficiency factors:

• No spherical harmonic expansions beyond the input encoding.
• Vector features remain in 3D Cartesian space rather than higher-dimensional irreducible representations.
• Message computation uses only scalar distances, not full geometric tensors.

This linear scaling makes EGNNs practical for large biomolecular systems containing thousands of atoms, where tensor product methods become prohibitively expensive.

For molecular dynamics applications, EGNNs naturally predict conservative forces by computing the negative gradient of a learned potential energy surface. Because the architecture is E(n) equivariant:

• Predicted forces rotate correctly with the molecule.
• The force field is energy-conserving by construction when used with a scalar energy head.
• No additional constraints or post-processing are needed to enforce physical consistency.

This property makes EGNNs suitable as drop-in replacements for classical force fields in Neural Network Potentials (NNPs), enabling accurate and fast molecular simulations.