E(n) Equivariant Network Block#

class EnEquivariantNetworkBlock(node_feature_dim, edge_feature_dim, pos_dim, hidden_dim=64, n_message_layers=2, n_update_layers=2, activation=<class 'torch.nn.modules.activation.SiLU'>, aggr='add', node_dim=-2, flow='source_to_target')[source]

Bases: MessagePassing

Implementation of the E(n) Equivariant Graph Neural Network block. This block is used to perform message-passing between nodes and edges in a graph neural network, following the scheme proposed by Satorras et al. in 2021. It serves as an inner block in a larger graph neural network architecture.

The message between two nodes connected by an edge is computed by applying a linear transformation to the sender node features and the edge features, together with the squared euclidean distance between the sender and recipient node positions, followed by a non-linear activation function. Messages are then aggregated using an aggregation scheme (e.g., sum, mean, min, max, or product).

The update step is performed by applying another MLP to the concatenation of the incoming messages and the node features. Here, also the node positions are updated by adding the incoming messages divided by the degree of the recipient node.

E(n) Equivariant Network Block#

This Page