Graph Neural Network (GNN)

A type of Neural Network that is the standard for Deep Learning on non-Euclidean data.

How it works

1. The Embedding Space

Every node in the graph starts with a raw feature vector $x_i$ (e.g., Atom Type, Charge). The GNN converts this discrete vector into a continuous Hidden State (embedding) $h_i^{(k)}$, where k is the layer number.

• $h_i^{(0)}=x_i$ (Initial state is just the raw features).

2. The Cycle: Message, Aggregate, Update

For every layer in the network, the following three steps happen in parallel for every node in the graph.

1. Message (Gather)

Every node looks at its neighbors. The neighbors transform their own information into a “message.”

• GCNConv: $\frac{1}{\sqrt{d_i{d}_j}}\cdot x_j$
  • $x_j$: The node features (e.g., Atom: Carbon).
• GATConv: $Message={\alpha }_{ij}\cdot x_j$
  • $\alpha$: Importance score
• GINEConv: $Message=MLP(x_j+e_{ij})$
  • $e_{ij}$: The edge features (e.g., Bond Type: Double).

2. Aggregate (Pool)

The central node collects all messages from its neighbors. Since a node can have 1 neighbor or 100, we must use a function that can handle variable input sizes (permutation invariant).
$m_i^{agg}={\sum }_{jϵN(i)}{m}_{ij}$

• Common aggregators: Sum, Mean, Max.
• Analogy: The Carbon atom sums up the “pull” from all surrounding atoms to understand its local electronic environment.

3. Update (Combine)

The node takes the aggregated message and combines it with its own previous state to calculate its new state.

• Standard: $h_u^{(k+1)}=\mathrm{UPDAT}{\mathrm{E}}^{(\mathrm{k})}\left(h_u^{(k)},M_u^{(k)}\right)$
  • $h_u^{(k+1)}$: The new feature representation (embedding) for node $u$ at the next iteration (layer) $k+1$
  • $h_u^{(k)}$: The current feature representation of node $u$ from the previous iteration $k$.
  • $M_u^{(k)}$: The aggregated message from the neighborhood
• GraphSAGE: $h_u^{(k+1)}=LINEAR(h_u^{(k)}||M_u^{(k)})$
• GIN/GINE: $h_u^{(k+1)}=MLP((1+ϵ)\cdot h_u^{(k)}+M_u^{(k)})$
  • ϵ: A learnable or fixed scalar that weights the importance of the central node relative to its neighbors.
• Usually involves a Linear Layer or a GRU.
• Analogy: The Carbon atom realizes, “I was neutral before, but because I’m surrounded by Oxygen, I am now partially positive (δ+).“

3. The “Receptive Field” (Reaching Further)

This is the most critical concept for architecture design.

• 1 Layer: A node knows its immediate neighbors (1-hop).
• 2 Layers: A node knows its neighbors’ neighbors (2-hops).
• K Layers: A node knows information from K steps away.
  Rule of Thumb: For small organic molecules, 3–6 layers is usually sufficient.

4. Permutation Invariance

The GNN operates on the structure (who is connected to whom) and sums the messages (A+B = B+A), the output is identical regardless of the input order. If you have a benzene ring, it doesn’t matter if you list the atoms as [C1, C2, C3...] or [C3, C1, C2...] in your data file.

• Standard Neural Networks (MLPs) cannot do this. They treat [C1, C2] and [C2, C1] as different inputs.

5. Final Readout

After $K$ layers, every node holds a vector $h_i^{(K)}$ containing information about itself and its context within the graph. To get a single prediction for the molecule (e.g., Solubility), you pool them together:
$\hat{y}=GlobalPool(\{h_1^{(K)},...,h_N^{(K)}\})$

Layers

Types:

1. GCNConv (Graph Convolutional Network)

The “Hello World” of GNNs. It averages neighbor features.

• Mechanism:
  ${\mathbf{x}}_i^{\prime }={\mathbf{\Theta }}^{\top }{\sum }_{j\in \mathcal{N}(i)\cup \{i\}}\frac{e_{j,i}}{\sqrt{{\hat{d}}_j{\hat{d}}_i}}{\mathbf{x}}_j$
  • It sums neighbor features ($x_j$).
  • It normalizes by node degrees ($d_j$,$d_i$) so high-degree nodes don’t explode the values.
• Pros: Very fast, computationally cheap.
• Cons: It treats every neighbor equally. It cannot distinguish between a neighbor that is critical (e.g., a functional group) and one that is irrelevant. It also struggles to distinguish distinct graph structures (isomers) if the averages look the same.

2. GraphSAGE (Sample and Aggregate)

Designed for massive graphs (like social networks), but useful in chemistry for its flexibility. It separates the “Aggregation” step from the “Update” step.

• Mechanism:
  ${\mathbf{x}}_i^{\prime }={\mathbf{W}}_1{\mathbf{x}}_i+{\mathbf{W}}_2\cdot {\text{aggregate}}_{j\in \mathcal{N}(i)}({\mathbf{x}}_j)$

• Aggregators: Unlike GCN (which forces a weighted sum), SAGE allows you to choose:

  • mean: Similar to GCN.
  • max: Takes the max value of neighbors (good for detecting sharp features like “is there a Halogen nearby?”).
  • lstm: Uses an LSTM to process neighbors (rarely used due to slowness).

• Key Feature: It is Inductive. It can generate embeddings for nodes it has never seen before, making it great for dynamic datasets.

3. GATConv (Graph Attention Network)

This introduces the Attention Mechanism to graphs. Instead of a fixed weight (like GCN), the model learns a weight ${\alpha }_{ij}$ for every edge.

• Mechanism:

• The Attention Score (${\alpha }_{ij}$): The model looks at Node i and Neighbor j and calculates a score: “How important is j to i?”
• Relevance: In chemistry, this allows the model to “focus” on chemically active parts of the molecule (e.g., a hydroxyl group) while ignoring the carbon backbone, even if they are both neighbors.

4. GINEConv (Graph Isomorphism Network + Edge Features)

This is theoretically the most powerful GNN for distinguishing different graph structures (graph isomorphism). Standard GCNs can fail to distinguish simple isomers.

• Mechanism:
  ${\mathbf{x}}_i^{\prime }=h_{\Theta }\left((1+ϵ)\cdot {\mathbf{x}}_i+{\sum }_{j\in \mathcal{N}(i)}{\mathbf{x}}_j\right)$
  • MLP ($h_{\Theta }$): Unlike GCN which applies a linear matrix, GIN applies an entire Neural Network (MLP) after aggregation. This preserves non-linear patterns.
  • Epsilon ($ϵ$): A learnable parameter that weighs the central node vs. its neighbors.
• GINEConv (GIN with Edges):
  • Standard GIN ignores edge features (bond types).
  • GINE adds the edge attributes to the neighbor features before aggregation.

Pooling Layers

Once the nodes have messaged each other, you need to condense the whole graph into one vector to predict an output.

global_mean_pool

Returns the average of all vectors.

• Problem: Two small molecules might have the same average as one giant molecule.

global_add_pool

Returns the sum of all vectors.

• Benefit: Captures “size” information (molecular weight correlates with solubility). Preferred for Chemistry.

global_max_pool:

Returns the maximum value across all vectors.

• Benefit: Detects the presence of a specific “feature” (e.g., “If any atom is highly polar, the whole molecule is polar”).