Edifice.Graph.GINv2 (Edifice v0.2.0)

Copy Markdown View Source

GINv2: Graph Isomorphism Network with edge features (Hu et al., 2020).

Extends the original GIN architecture to incorporate edge features in the message passing step. For each edge (u, v), the message incorporates both the source node features and the edge features, allowing the network to learn from edge attributes such as bond types, distances, or relationship properties.

Architecture

Node Features  [batch, num_nodes, input_dim]
Adjacency      [batch, num_nodes, num_nodes]
Edge Features  [batch, num_nodes, num_nodes, edge_dim]
      |
      v
+----------------------------------------------+
| GINv2 Layer:                                  |
|   1. Edge messages: Dense(h_u + edge_proj)    |
|   2. Aggregate: sum over neighbors (via adj)  |
|   3. Combine: (1+eps)*h_v + aggregated        |
|   4. Transform: MLP(combined)                 |
+----------------------------------------------+
      |
      v
Node Embeddings [batch, num_nodes, hidden_dim]

Differences from GIN

  • Takes a third input "edge_features" with shape [batch, nodes, nodes, edge_dim]
  • Edge features are projected and combined with neighbor features before aggregation
  • More expressive for graphs with rich edge information (molecules, knowledge graphs)

Usage

model = GINv2.build(
  input_dim: 16,
  edge_dim: 4,
  hidden_dims: [64, 64],
  num_classes: 2
)

References

Summary

Types

Options for build/1.

Functions

Build a GINv2 model with edge features.

Get the output size of a GINv2 model.

Types

build_opt()

@type build_opt() ::
  {:activation, atom()}
  | {:dropout, float()}
  | {:edge_dim, pos_integer()}
  | {:hidden_dims, [pos_integer()]}
  | {:input_dim, pos_integer()}
  | {:num_classes, pos_integer() | nil}
  | {:pool, :mean | :sum | :max | nil}

Options for build/1.

Functions

build(opts \\ [])

@spec build([build_opt()]) :: Axon.t()

Build a GINv2 model with edge features.

Options

  • :input_dim - Input feature dimension per node (required)
  • :edge_dim - Edge feature dimension (required)
  • :hidden_dims - List of hidden dimensions for each layer (default: [64, 64])
  • :num_classes - If provided, adds a classification head (default: nil)
  • :dropout - Dropout rate (default: 0.0)
  • :activation - Activation for MLPs (default: :relu)
  • :pool - Global pooling for graph classification: :mean, :sum, :max (default: nil)

Returns

An Axon model with three inputs: "nodes", "adjacency", and "edge_features".

ginv2_layer(nodes, adjacency, edge_features, output_dim, opts \\ [])

@spec ginv2_layer(Axon.t(), Axon.t(), Axon.t(), pos_integer(), keyword()) :: Axon.t()

Single GINv2 layer with edge features.

Options

  • :name - Layer name prefix (default: "ginv2")
  • :dropout - Dropout rate (default: 0.0)
  • :activation - Activation for MLP (default: :relu)

output_size(opts \\ [])

@spec output_size(keyword()) :: pos_integer()

Get the output size of a GINv2 model.