Chapter 4: Advanced GNN Techniques - State-of-the-Art Architectures and Interpretability
Grasp the intuition behind equivariant models such as SchNet/NequIP/MACE and understand what kinds of problems they are effective for. You will also understand the trade-off between computational cost and accuracy.
π‘ Supplement: Higher accuracy tends to mean higher cost. Estimate the required accuracy and computational resources, and advance gradually.
Learning Objectives
By reading this chapter, you will be able to: - Understand hierarchical representation learning through graph pooling - Implement advanced GNNs that utilize edge features - Master SchNet and DimeNet, which consider 3D geometric information - Understand the principles of equivariant GNNs (E(3)-equivariant) - Visualize the basis of predictions with GNNExplainer
Reading time: 20-25 min Code examples: 8 Exercises: 3
4.1 Graph Pooling: Hierarchical Representation Learning
4.1.1 What Is Graph Pooling?
Graph pooling is a technique that reduces the number of nodes while preserving graph structure, thereby learning hierarchical representations.
Importance: - π Multi-stage feature extraction: Learn features in the order local β intermediate β global - π Reduced computational cost: Improve computational efficiency by reducing the number of nodes - π― Selection of important nodes: Automatically identify atoms/structures important for prediction
Representative methods: 1. Top-K Pooling: Select the top K nodes by score 2. SAGPooling: Self-Attention Graph Pooling (compute importance via an attention mechanism) 3. DiffPool: Differentiable soft clustering
4.1.2 Implementing Top-K Pooling
import torch
import torch.nn.functional as F
from torch_geometric.nn import GCNConv, TopKPooling, global_mean_pool
from torch_geometric.data import Data, DataLoader
class GNN_with_Pooling(torch.nn.Module):
"""
GNN using Top-K Pooling
Architecture:
- GCN layer β Pooling β GCN layer β Global Pool β Fully connected layer
"""
def __init__(self, num_node_features, num_classes, hidden_channels=64, pool_ratio=0.5):
super().__init__()
# Block 1: GCN + TopKPooling
self.conv1 = GCNConv(num_node_features, hidden_channels)
self.pool1 = TopKPooling(hidden_channels, ratio=pool_ratio)
# Block 2: GCN + TopKPooling
self.conv2 = GCNConv(hidden_channels, hidden_channels)
self.pool2 = TopKPooling(hidden_channels, ratio=pool_ratio)
# Block 3: GCN
self.conv3 = GCNConv(hidden_channels, hidden_channels)
# Fully connected layers
self.lin1 = torch.nn.Linear(hidden_channels, hidden_channels // 2)
self.lin2 = torch.nn.Linear(hidden_channels // 2, num_classes)
def forward(self, x, edge_index, batch):
# Block 1
x = F.relu(self.conv1(x, edge_index))
x, edge_index, _, batch, _, _ = self.pool1(x, edge_index, None, batch)
# Block 2
x = F.relu(self.conv2(x, edge_index))
x, edge_index, _, batch, _, _ = self.pool2(x, edge_index, None, batch)
# Block 3 (no pooling)
x = F.relu(self.conv3(x, edge_index))
# Global pooling
x = global_mean_pool(x, batch)
# Fully connected layers
x = F.relu(self.lin1(x))
x = F.dropout(x, p=0.3, training=self.training)
x = self.lin2(x)
return x
# Instantiate the model
model = GNN_with_Pooling(
num_node_features=7,
num_classes=1,
hidden_channels=64,
pool_ratio=0.5 # Reduce nodes to 50%
)
print("===== Top-K Pooling GNN =====")
print(model)
print(f"\nNumber of parameters: {sum(p.numel() for p in model.parameters()):,}")
# Test with sample data
x = torch.randn(20, 7) # 20 nodes, 7-dimensional features
edge_index = torch.randint(0, 20, (2, 40))
batch = torch.zeros(20, dtype=torch.long)
with torch.no_grad():
out = model(x, edge_index, batch)
print(f"\nInput: {x.shape[0]} nodes")
print(f"Output: {out.shape}")
4.1.3 SAGPooling (Self-Attention Graph Pooling)
from torch_geometric.nn import SAGPooling
class GNN_with_SAGPool(torch.nn.Module):
"""
GNN using SAGPooling (learn node importance with an attention mechanism)
"""
def __init__(self, num_node_features, num_classes, hidden_channels=64, pool_ratio=0.5):
super().__init__()
# GCN layer
self.conv1 = GCNConv(num_node_features, hidden_channels)
# SAGPooling (learnable attention mechanism)
self.pool1 = SAGPooling(hidden_channels, ratio=pool_ratio)
# Block 2
self.conv2 = GCNConv(hidden_channels, hidden_channels)
self.pool2 = SAGPooling(hidden_channels, ratio=pool_ratio)
# Block 3
self.conv3 = GCNConv(hidden_channels, hidden_channels)
# Fully connected layers
self.lin1 = torch.nn.Linear(hidden_channels, hidden_channels // 2)
self.lin2 = torch.nn.Linear(hidden_channels // 2, num_classes)
def forward(self, x, edge_index, batch):
# Block 1 (GCN + SAGPooling)
x = F.relu(self.conv1(x, edge_index))
x, edge_index, _, batch, perm1, score1 = self.pool1(
x, edge_index, None, batch
)
# Block 2
x = F.relu(self.conv2(x, edge_index))
x, edge_index, _, batch, perm2, score2 = self.pool2(
x, edge_index, None, batch
)
# Block 3
x = F.relu(self.conv3(x, edge_index))
# Global pooling
x = global_mean_pool(x, batch)
# Fully connected layers
x = F.relu(self.lin1(x))
x = F.dropout(x, p=0.3, training=self.training)
x = self.lin2(x)
return x, (perm1, score1, perm2, score2) # Also return the importance scores
# Usage example
model_sag = GNN_with_SAGPool(num_node_features=7, num_classes=1)
with torch.no_grad():
out, (perm1, score1, perm2, score2) = model_sag(x, edge_index, batch)
print("\n===== SAGPooling =====")
print(f"First pooling: {x.shape[0]} nodes β {perm1.shape[0]} nodes")
print(f"Importance scores: {score1[:5].squeeze()}") # Scores of the top 5 nodes
4.1.4 Comparison of Pooling Methods
import matplotlib.pyplot as plt
import numpy as np
# Performance comparison of each pooling method (mock data)
pooling_methods = {
'No Pooling': {'MAE': 0.35, 'Time': 42.3, 'Memory': 1200},
'Top-K Pooling': {'MAE': 0.32, 'Time': 38.5, 'Memory': 980},
'SAGPooling': {'MAE': 0.28, 'Time': 45.8, 'Memory': 1050},
'DiffPool': {'MAE': 0.25, 'Time': 62.1, 'Memory': 1800},
}
fig, axes = plt.subplots(1, 3, figsize=(15, 4))
# MAE comparison
methods = list(pooling_methods.keys())
mae_values = [pooling_methods[m]['MAE'] for m in methods]
axes[0].bar(methods, mae_values, color=['gray', 'steelblue', 'forestgreen', 'coral'])
axes[0].set_ylabel('MAE (eV)', fontsize=12)
axes[0].set_title('Prediction Accuracy (lower is better)', fontsize=13)
axes[0].tick_params(axis='x', rotation=15)
axes[0].grid(True, alpha=0.3, axis='y')
# Computation time comparison
time_values = [pooling_methods[m]['Time'] for m in methods]
axes[1].bar(methods, time_values, color=['gray', 'steelblue', 'forestgreen', 'coral'])
axes[1].set_ylabel('Training Time (s)', fontsize=12)
axes[1].set_title('Computational Cost', fontsize=13)
axes[1].tick_params(axis='x', rotation=15)
axes[1].grid(True, alpha=0.3, axis='y')
# Memory usage comparison
memory_values = [pooling_methods[m]['Memory'] for m in methods]
axes[2].bar(methods, memory_values, color=['gray', 'steelblue', 'forestgreen', 'coral'])
axes[2].set_ylabel('Memory Usage (MB)', fontsize=12)
axes[2].set_title('Memory Efficiency', fontsize=13)
axes[2].tick_params(axis='x', rotation=15)
axes[2].grid(True, alpha=0.3, axis='y')
plt.tight_layout()
plt.show()
4.2 GNNs Considering 3D Geometric Information: SchNet
4.2.1 Principles of SchNet
SchNet (SchΓΌtt et al., 2017) is a continuous-filter convolutional GNN for 3D molecular representation learning that considers interatomic distances.
Features: - π Uses 3D coordinates: Directly inputs interatomic distances - π Continuous filters: Encodes distances with Gaussian basis functions - π Rotational invariance: Predictions invariant to 3D rotation
Equation: $$ h_i^{(t+1)} = h_i^{(t)} + \sum_{j \in \mathcal{N}(i)} W(r_{ij}) \odot h_j^{(t)} $$
Here, $W(r_{ij})$ is a function of the interatomic distance $r_{ij}$ (a continuous filter).
4.2.2 Implementing SchNet
import torch
import torch.nn as nn
from torch_geometric.nn import SchNet
# Use SchNet from PyTorch Geometric
model_schnet = SchNet(
hidden_channels=128,
num_filters=128,
num_interactions=6, # Number of message passing steps
num_gaussians=50, # Number of Gaussian basis functions
cutoff=10.0, # Cutoff distance (Γ
)
max_num_neighbors=32,
readout='add' # Global pooling (sum)
)
print("===== SchNet =====")
print(model_schnet)
print(f"\nNumber of parameters: {sum(p.numel() for p in model_schnet.parameters()):,}")
# Sample data (methane molecule: CH4)
# C: (0, 0, 0), H: 4 vertex positions
z = torch.tensor([6, 1, 1, 1, 1]) # Atomic numbers (C=6, H=1)
pos = torch.tensor([
[0.0, 0.0, 0.0], # C
[1.09, 0.0, 0.0], # H1
[-0.36, 1.03, 0.0], # H2
[-0.36, -0.51, 0.89], # H3
[-0.36, -0.51, -0.89] # H4
], dtype=torch.float)
batch = torch.zeros(5, dtype=torch.long)
# Forward pass (energy prediction)
with torch.no_grad():
energy = model_schnet(z, pos, batch)
print(f"\nInput: {z.shape[0]} atoms (methane molecule)")
print(f"Predicted energy: {energy.item():.4f} eV")
4.2.3 Training SchNet (QM9 Dataset)
from torch_geometric.datasets import QM9
from torch_geometric.loader import DataLoader
# Load the QM9 dataset (disable logging)
import warnings
warnings.filterwarnings('ignore')
dataset = QM9(root='./data/QM9')
# Set only internal energy (U0) as the target variable
target_idx = 7 # Index of U0
for data in dataset:
data.y = data.y[:, target_idx:target_idx+1]
# Data split
train_dataset = dataset[:10000]
test_dataset = dataset[10000:11000]
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)
# Device setup
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model_schnet = model_schnet.to(device)
# Training preparation
optimizer = torch.optim.Adam(model_schnet.parameters(), lr=0.001)
criterion = torch.nn.MSELoss()
def train_schnet(model, loader, optimizer, criterion, device):
model.train()
total_loss = 0
for data in loader:
data = data.to(device)
optimizer.zero_grad()
# SchNet uses atomic numbers (z) and coordinates (pos)
out = model(data.z, data.pos, data.batch)
loss = criterion(out, data.y)
loss.backward()
optimizer.step()
total_loss += loss.item() * data.num_graphs
return total_loss / len(loader.dataset)
# Training loop (simplified version)
print("\n===== Start SchNet Training =====")
for epoch in range(1, 21):
train_loss = train_schnet(model_schnet, train_loader, optimizer, criterion, device)
if epoch % 5 == 0:
print(f"Epoch {epoch:03d}, Train Loss: {train_loss:.4f}")
print("Training complete!")
4.3 DimeNet: A Direction-Aware GNN
4.3.1 Features of DimeNet
DimeNet (Directional Message Passing Neural Network) considers not only interatomic distances but also bond angles.
Key elements: - π Relationship among three atoms: the angle $\theta_{ijk}$ of i-j-k - π― Spherical harmonics: encode angles - π¬ High accuracy: outperforms SchNet on QM9
Equation (simplified): $$ m_{ij} = \sum_{k \in \mathcal{N}(j)} W(\theta_{ijk}, r_{ij}, r_{jk}) h_k $$
4.3.2 Using DimeNet
from torch_geometric.nn import DimeNet
# Instantiate the DimeNet model
model_dimenet = DimeNet(
hidden_channels=128,
out_channels=1,
num_blocks=6,
num_bilinear=8,
num_spherical=7,
num_radial=6,
cutoff=5.0,
max_num_neighbors=32,
envelope_exponent=5,
num_before_skip=1,
num_after_skip=2,
num_output_layers=3
)
print("===== DimeNet =====")
print(f"Number of parameters: {sum(p.numel() for p in model_dimenet.parameters()):,}")
# Forward pass with sample data
with torch.no_grad():
energy = model_dimenet(z, pos, batch)
print(f"\nPredicted energy (DimeNet): {energy.item():.4f} eV")
4.3.3 SchNet vs DimeNet Performance Comparison
import pandas as pd
import matplotlib.pyplot as plt
# QM9 benchmark results (literature values)
results = {
'Model': ['GCN', 'SchNet', 'DimeNet', 'DimeNet++'],
'U0 MAE (meV)': [230, 14, 6.3, 4.4],
'HOMO MAE (meV)': [190, 41, 27, 23],
'LUMO MAE (meV)': [200, 34, 20, 19],
'Params (M)': [0.5, 3.0, 2.0, 2.1]
}
df = pd.DataFrame(results)
fig, axes = plt.subplots(1, 2, figsize=(14, 5))
# MAE comparison (U0)
axes[0].bar(df['Model'], df['U0 MAE (meV)'], color=['gray', 'steelblue', 'forestgreen', 'coral'])
axes[0].set_ylabel('MAE (meV)', fontsize=12)
axes[0].set_title('Internal Energy (U0) Prediction Accuracy', fontsize=13)
axes[0].set_ylim(0, 250)
axes[0].grid(True, alpha=0.3, axis='y')
# Parameter count comparison
axes[1].bar(df['Model'], df['Params (M)'], color=['gray', 'steelblue', 'forestgreen', 'coral'])
axes[1].set_ylabel('Number of Parameters (millions)', fontsize=12)
axes[1].set_title('Model Size', fontsize=13)
axes[1].grid(True, alpha=0.3, axis='y')
plt.tight_layout()
plt.show()
print("===== QM9 Benchmark =====")
print(df.to_string(index=False))
4.4 Equivariant GNNs (E(3)-Equivariant)
4.4.1 What Is Equivariance?
Equivariance is the property whereby a transformation of the input (rotation, translation) is reflected in the output as the same transformation.
Mathematical definition: $$ f(R \cdot x) = R \cdot f(x) $$
Here, $R$ is a rotation matrix and $x$ is the 3D coordinates.
Importance: - π Compliance with physical laws: independent of the molecule's orientation - π― Data efficiency: no need for rotational augmentation - π Generalization performance: high accuracy even for orientations outside the training data
4.4.2 An Example of Equivariant GNN: NequIP
NequIP (Neural Equivariant Interatomic Potentials) is a GNN with E(3) equivariance.
Features: - Equivariant message passing via tensor products - Angle encoding via spherical harmonics - Optimal for learning force fields
flowchart TD
A[3D Atomic Coordinates] --> B[E3 Equivariant Embedding]
B --> C[Tensor Product Message Passing]
C --> D[Spherical Harmonic Filter]
D --> E[Equivariant Update]
E --> F[Energy / Force Prediction]
style A fill:#e3f2fd
style B fill:#fff3e0
style C fill:#f3e5f5
style D fill:#e8f5e9
style E fill:#fff9c4
style F fill:#ffccbc
4.4.3 Verifying Equivariance
import torch
import numpy as np
def rotate_coordinates(pos, axis='z', angle=np.pi/4):
"""
Rotate coordinates
Parameters:
-----------
pos : torch.Tensor (num_atoms, 3)
Atomic coordinates
axis : str
Rotation axis ('x', 'y', 'z')
angle : float
Rotation angle (radians)
Returns:
--------
rotated_pos : torch.Tensor (num_atoms, 3)
Rotated coordinates
"""
cos_a = np.cos(angle)
sin_a = np.sin(angle)
if axis == 'z':
R = torch.tensor([
[cos_a, -sin_a, 0],
[sin_a, cos_a, 0],
[0, 0, 1]
], dtype=torch.float)
elif axis == 'y':
R = torch.tensor([
[cos_a, 0, sin_a],
[0, 1, 0],
[-sin_a, 0, cos_a]
], dtype=torch.float)
else: # 'x'
R = torch.tensor([
[1, 0, 0],
[0, cos_a, -sin_a],
[0, sin_a, cos_a]
], dtype=torch.float)
return pos @ R.T
# Rotate the methane molecule
pos_original = torch.tensor([
[0.0, 0.0, 0.0],
[1.09, 0.0, 0.0],
[-0.36, 1.03, 0.0],
[-0.36, -0.51, 0.89],
[-0.36, -0.51, -0.89]
], dtype=torch.float)
pos_rotated = rotate_coordinates(pos_original, axis='z', angle=np.pi/2)
# Predict with SchNet (verify rotational invariance)
model_schnet.eval()
z = torch.tensor([6, 1, 1, 1, 1])
batch = torch.zeros(5, dtype=torch.long)
with torch.no_grad():
energy_original = model_schnet(z, pos_original, batch)
energy_rotated = model_schnet(z, pos_rotated, batch)
print("===== Verifying Rotational Invariance =====")
print(f"Predicted energy for original coordinates: {energy_original.item():.4f} eV")
print(f"Predicted energy for rotated coordinates: {energy_rotated.item():.4f} eV")
print(f"Difference: {abs(energy_original.item() - energy_rotated.item()):.6f} eV")
if abs(energy_original.item() - energy_rotated.item()) < 1e-4:
print("β
Rotational invariance is satisfied!")
else:
print("β Rotational invariance is incomplete.")
4.5 Attention Mechanisms and Transformer Integration
4.5.1 Graph Attention Networks (GAT)
GAT intensively learns important nodes through an attention mechanism.
Computation of attention coefficients: $$ \alpha_{ij} = \frac{\exp(\text{LeakyReLU}(a^T [Wh_i | Wh_j]))}{\sum_{k \in \mathcal{N}(i)} \exp(\text{LeakyReLU}(a^T [Wh_i | Wh_k]))} $$
from torch_geometric.nn import GATConv
class GAT_Model(torch.nn.Module):
"""
Graph Attention Network
"""
def __init__(self, num_node_features, num_classes, hidden_channels=64, heads=8):
super().__init__()
# GAT layers (multi-head attention mechanism)
self.conv1 = GATConv(num_node_features, hidden_channels, heads=heads, dropout=0.2)
self.conv2 = GATConv(hidden_channels * heads, hidden_channels, heads=heads, dropout=0.2)
self.conv3 = GATConv(hidden_channels * heads, hidden_channels, heads=1, concat=False, dropout=0.2)
# Fully connected layers
self.lin1 = torch.nn.Linear(hidden_channels, hidden_channels // 2)
self.lin2 = torch.nn.Linear(hidden_channels // 2, num_classes)
def forward(self, x, edge_index, batch, return_attention_weights=False):
# GAT layer 1
x, attn1 = self.conv1(x, edge_index, return_attention_weights=True)
x = F.elu(x)
# GAT layer 2
x, attn2 = self.conv2(x, edge_index, return_attention_weights=True)
x = F.elu(x)
# GAT layer 3
x = self.conv3(x, edge_index)
x = F.elu(x)
# Global pooling
x = global_mean_pool(x, batch)
# Fully connected layers
x = F.relu(self.lin1(x))
x = F.dropout(x, p=0.3, training=self.training)
x = self.lin2(x)
if return_attention_weights:
return x, (attn1, attn2)
else:
return x
# Instantiate the model
model_gat = GAT_Model(num_node_features=7, num_classes=1, heads=8)
print("===== Graph Attention Network =====")
print(model_gat)
print(f"\nNumber of parameters: {sum(p.numel() for p in model_gat.parameters()):,}")
4.5.2 Visualizing Attention Weights
import matplotlib.pyplot as plt
import networkx as nx
def visualize_attention(edge_index, attention_weights, node_labels=None, figsize=(10, 8)):
"""
Visualize attention weights on the graph
Parameters:
-----------
edge_index : torch.Tensor (2, num_edges)
Edge indices
attention_weights : torch.Tensor (num_edges, heads)
Attention weights
node_labels : list
Node labels (e.g., atomic symbols)
"""
# Create a NetworkX graph
G = nx.Graph()
num_nodes = edge_index.max().item() + 1
G.add_nodes_from(range(num_nodes))
# Add edges and attention weights
for i in range(edge_index.size(1)):
src, dst = edge_index[:, i].tolist()
weight = attention_weights[i].mean().item() # Average over multi-heads
G.add_edge(src, dst, weight=weight)
# Layout
pos = nx.spring_layout(G, seed=42)
# Draw
fig, ax = plt.subplots(figsize=figsize)
# Draw edges (thickness = attention weight)
edges = G.edges()
weights = [G[u][v]['weight'] for u, v in edges]
weights_normalized = [w / max(weights) * 10 for w in weights]
nx.draw_networkx_edges(G, pos, width=weights_normalized, alpha=0.6, ax=ax)
# Draw nodes
nx.draw_networkx_nodes(G, pos, node_size=800, node_color='lightblue', ax=ax)
# Labels
if node_labels:
labels = {i: node_labels[i] for i in range(num_nodes)}
else:
labels = {i: str(i) for i in range(num_nodes)}
nx.draw_networkx_labels(G, pos, labels, font_size=12, ax=ax)
ax.set_title('Attention Weight Visualization (thicker line = higher attention)', fontsize=14)
ax.axis('off')
plt.tight_layout()
plt.show()
# Usage example (sample data)
edge_index_sample = torch.tensor([[0, 1, 1, 2, 2, 3, 3, 0],
[1, 0, 2, 1, 3, 2, 0, 3]], dtype=torch.long)
attention_weights_sample = torch.rand(8, 8) # 8 edges Γ 8 heads
node_labels_sample = ['C', 'H', 'H', 'H']
visualize_attention(edge_index_sample, attention_weights_sample, node_labels_sample)
4.6 GNNExplainer: Interpretability of Predictions
4.6.1 What Is GNNExplainer?
GNNExplainer is a technique for explaining the basis of GNN predictions.
Main functions: - π Identification of important substructures: which atoms/bonds contributed to the prediction - π Visualization: display on the graph as an attention map - π― Improved reliability: explainable AI rather than a black box
Principle: Find the important subgraph $G_S$ via the following optimization problem: $$ \max_{G_S} \text{Mutual Information}(Y, G_S) $$
4.6.2 Implementing GNNExplainer
from torch_geometric.explain import Explainer, GNNExplainer as GNNExplainerAlgo
# Use the trained model
model_gat.eval()
# Configure GNNExplainer
explainer = Explainer(
model=model_gat,
algorithm=GNNExplainerAlgo(epochs=200),
explanation_type='model',
node_mask_type='attributes',
edge_mask_type='object',
model_config=dict(
mode='multiclass_classification',
task_level='graph',
return_type='raw',
),
)
# Generate an explanation for a sample graph
x_sample = torch.randn(10, 7)
edge_index_sample = torch.randint(0, 10, (2, 20))
batch_sample = torch.zeros(10, dtype=torch.long)
# Generate the explanation
explanation = explainer(x_sample, edge_index_sample, batch=batch_sample)
print("===== GNNExplainer =====")
print(f"Node importance: {explanation.node_mask}")
print(f"Edge importance: {explanation.edge_mask}")
# Visualize importance
import matplotlib.pyplot as plt
fig, axes = plt.subplots(1, 2, figsize=(14, 5))
# Node importance
axes[0].bar(range(len(explanation.node_mask)), explanation.node_mask.detach().numpy())
axes[0].set_xlabel('Node ID', fontsize=12)
axes[0].set_ylabel('Importance', fontsize=12)
axes[0].set_title('Node Importance (higher = more contribution to prediction)', fontsize=13)
axes[0].grid(True, alpha=0.3, axis='y')
# Edge importance
axes[1].bar(range(len(explanation.edge_mask)), explanation.edge_mask.detach().numpy())
axes[1].set_xlabel('Edge ID', fontsize=12)
axes[1].set_ylabel('Importance', fontsize=12)
axes[1].set_title('Edge Importance (higher = more contribution to prediction)', fontsize=13)
axes[1].grid(True, alpha=0.3, axis='y')
plt.tight_layout()
plt.show()
4.6.3 Real-World Applications
# Identify important substructures in molecular toxicity prediction
# Example: toxicity evaluation of a benzene ring
# The GNN explains "which part contributes to toxicity"
def explain_toxicity(model, smiles, explainer):
"""
Explain a molecule's toxicity prediction
Parameters:
-----------
model : torch.nn.Module
Trained GNN model
smiles : str
SMILES string
explainer : Explainer
GNNExplainer
Returns:
--------
explanation : Explanation
Importance mask
"""
from rdkit import Chem
# Convert SMILES to a graph
mol = Chem.MolFromSmiles(smiles)
# ... graph conversion process ...
# Generate the explanation
# explanation = explainer(x, edge_index, batch)
# Identify important functional groups
# important_atoms = torch.where(explanation.node_mask > 0.5)[0]
print(f"SMILES: {smiles}")
print(f"Toxicity prediction: {'High' if predicted_toxicity > 0.5 else 'Low'}")
print(f"Important atoms: {important_atoms.tolist()}")
return explanation
# Usage example (conceptual)
# explanation = explain_toxicity(model, "c1ccccc1", explainer)
4.7 Chapter Summary
What You Learned
-
Graph Pooling - Top-K Pooling: selection of the top K nodes by score - SAGPooling: learnable pooling via an attention mechanism - Improved prediction accuracy through hierarchical representation learning
-
GNNs Considering 3D Geometric Information - SchNet: uses interatomic distances via continuous-filter convolution - DimeNet: also considers bond angles (SOTA performance) - MAE 4-6 meV on QM9 (state-of-the-art)
-
Equivariant GNNs - E(3) equivariance: invariance to rotation and translation - NequIP: optimal for force-field learning - High-accuracy prediction that complies with physical laws
-
Attention Mechanisms - GAT: intensively learns important nodes via multi-head attention - Improved interpretability through attention weight visualization - Integration with Transformers
-
Interpretability - GNNExplainer: explanation of the prediction basis - Identification of important substructures - Building reliable AI systems
Key Takeaways
- β Graph pooling improves both computational efficiency and accuracy
- β Using 3D information (distances, angles) dramatically improves prediction accuracy
- β Equivariance is the key to building physically correct models
- β Attention mechanisms improve interpretability
- β GNNExplainer makes it possible to explain "why this prediction"
To the Next Chapter
Chapter 5 covers real-world applications and career paths: - Catalyst design (OC20 Challenge) - Crystal structure prediction (CGCNN, Matformer) - Materials screening (Materials Project integration) - Industrial application cases - Career paths for GNN experts
Chapter 5: Real-World Applications and Careers β
Exercises
Problem 1 (Difficulty: medium)
Explain the difference between Top-K Pooling and SAGPooling, and propose in what situations each method should be used.
Hint
Compare them from the perspectives of learnability and computational cost.Sample Solution
**Top-K Pooling**: - **Feature**: Learns node scores and selects the top K (fixed ratio) - **Computational cost**: Low (simple sort operation) - **Learning**: Learns only the score function **SAGPooling (Self-Attention Graph Pooling)**: - **Feature**: Dynamically learns node importance via an attention mechanism - **Computational cost**: Somewhat higher (attention mechanism computation) - **Learning**: Learns including the attention weights (more flexible) **Guidelines for choosing**: | Situation | Recommended Method | Reason | |-----|----------|-----| | Few data (<1000) | Top-K Pooling | Fewer parameters, less prone to overfitting | | Many data (>10000) | SAGPooling | Attention mechanism can learn complex patterns | | Limited computational resources | Top-K Pooling | Low computational cost | | Interpretability is important | SAGPooling | Attention weights can visualize important nodes | | Maximum accuracy needed | SAGPooling | More flexible learning is possible | **Implementation example**:# Selection based on the amount of data
if len(dataset) < 1000:
pooling = TopKPooling(hidden_channels, ratio=0.5)
else:
pooling = SAGPooling(hidden_channels, ratio=0.5)
Problem 2 (Difficulty: hard)
Explain, using equations, why SchNet has rotational invariance.
Hint
Use the fact that interatomic distances are invariant to rotation.Sample Solution
**Proof of SchNet's Rotational Invariance**: **Premise**: - Let the molecule's 3D coordinates be $\mathbf{r}\_i$ (position vector of atom $i$) - Let the rotation matrix be $R$ ($R^T R = I$, $\det(R) = 1$) **Step 1: Invariance of interatomic distances** Interatomic distance before rotation: $$ r\_{ij} = \|\mathbf{r}\_i - \mathbf{r}\_j\| $$ Interatomic distance after rotation: $$ r'\_{ij} = \|R\mathbf{r}\_i - R\mathbf{r}\_j\| = \|R(\mathbf{r}\_i - \mathbf{r}\_j)\| $$ From the properties of the rotation matrix: $$ \|R\mathbf{v}\| = \|\mathbf{v}\| $$ Therefore: $$ r'\_{ij} = r\_{ij} $$ **Interatomic distances are invariant to rotation!** **Step 2: SchNet's message passing** SchNet's message is a function of the interatomic distance $r\_{ij}$: $$ m\_{ij} = W(r\_{ij}) \odot h\_j $$ Here, $W(r\_{ij})$ is a continuous filter (a linear combination of Gaussian basis functions): $$ W(r\_{ij}) = \sum\_{k=1}^{K} w\_k \exp\left(-\gamma (r\_{ij} - \mu\_k)^2\right) $$ **Step 3: Message after rotation** Since interatomic distances remain invariant after rotation: $$ m'\_{ij} = W(r'\_{ij}) \odot h'\_j = W(r\_{ij}) \odot h'\_j $$ **Step 4: Global representation** SchNet's final output aggregates the features of each atom: $$ E = \sum\_{i=1}^{N} f(h\_i) $$ Since each atom's features $h\_i$ depend only on interatomic distances before and after rotation, the aggregated result is also invariant: $$ E' = \sum\_{i=1}^{N} f(h'\_i) = E $$ **Conclusion**: Because SchNet takes only interatomic distances (rotational invariants) as input, the prediction does not change even when the entire molecule is rotated. This is the mathematical basis of **rotational invariance**. **Verification in code**:import torch
# Original coordinates
pos = torch.tensor([[0, 0, 0], [1, 0, 0], [0, 1, 0]], dtype=torch.float)
# Rotation matrix (90 degrees around the Z axis)
R = torch.tensor([[0, -1, 0], [1, 0, 0], [0, 0, 1]], dtype=torch.float)
pos_rotated = pos @ R.T
# Compute interatomic distance
dist_original = torch.norm(pos[0] - pos[1])
dist_rotated = torch.norm(pos_rotated[0] - pos_rotated[1])
print(f"Original distance: {dist_original.item():.6f}")
print(f"Distance after rotation: {dist_rotated.item():.6f}")
print(f"Difference: {abs(dist_original - dist_rotated).item():.10f}")
# Output: difference β 0 (within numerical error)
Problem 3 (Difficulty: hard)
Using GNNExplainer, write complete code that shows "the benzene ring contributes to toxicity" in a molecular toxicity prediction model.
Hint
Convert the molecule to a graph with RDKit and identify the important atoms with GNNExplainer.Sample Solution
import torch
import torch.nn.functional as F
from torch_geometric.nn import GCNConv, global_mean_pool
from torch_geometric.data import Data
from torch_geometric.explain import Explainer, GNNExplainer as GNNExplainerAlgo
from rdkit import Chem
from rdkit.Chem import Draw
import matplotlib.pyplot as plt
import numpy as np
# Step 1: Define the toxicity prediction model
class ToxicityGNN(torch.nn.Module):
def __init__(self, num_node_features, hidden_channels=64):
super().__init__()
self.conv1 = GCNConv(num_node_features, hidden_channels)
self.conv2 = GCNConv(hidden_channels, hidden_channels)
self.conv3 = GCNConv(hidden_channels, hidden_channels)
self.lin = torch.nn.Linear(hidden_channels, 1) # Toxicity score
def forward(self, x, edge_index, batch):
x = F.relu(self.conv1(x, edge_index))
x = F.relu(self.conv2(x, edge_index))
x = F.relu(self.conv3(x, edge_index))
x = global_mean_pool(x, batch)
x = self.lin(x)
return torch.sigmoid(x) # Score in 0-1
# Step 2: Convert SMILES to a graph
def smiles_to_graph(smiles):
mol = Chem.MolFromSmiles(smiles)
if mol is None:
return None, None
# Node features (one-hot of atomic number)
atom_features = []
for atom in mol.GetAtoms():
features = [0] * 10 # Top 10 elements
atomic_num = atom.GetAtomicNum()
if atomic_num < 10:
features[atomic_num] = 1
else:
features[9] = 1 # Other
atom_features.append(features)
x = torch.tensor(atom_features, dtype=torch.float)
# Edge indices
edge_indices = []
for bond in mol.GetBonds():
i = bond.GetBeginAtomIdx()
j = bond.GetEndAtomIdx()
edge_indices += [[i, j], [j, i]]
edge_index = torch.tensor(edge_indices, dtype=torch.long).t().contiguous()
return Data(x=x, edge_index=edge_index), mol
# Step 3: Train the model (simplified; in practice, train on training data)
model = ToxicityGNN(num_node_features=10)
model.eval() # Assume trained
# Step 4: Generate an explanation for a benzene-containing molecule
smiles = "c1ccccc1CC(=O)O" # Phenylacetic acid (benzene ring + acetic acid)
data, mol = smiles_to_graph(smiles)
batch = torch.zeros(data.num_nodes, dtype=torch.long)
# Toxicity prediction
with torch.no_grad():
toxicity_score = model(data.x, data.edge_index, batch)
print(f"SMILES: {smiles}")
print(f"Predicted toxicity score: {toxicity_score.item():.4f}")
# Step 5: Generate an explanation with GNNExplainer
explainer = Explainer(
model=model,
algorithm=GNNExplainerAlgo(epochs=200),
explanation_type='model',
node_mask_type='attributes',
edge_mask_type='object',
model_config=dict(
mode='binary_classification',
task_level='graph',
return_type='raw',
),
)
explanation = explainer(data.x, data.edge_index, batch=batch)
# Step 6: Identify important atoms
node_importance = explanation.node_mask.detach().numpy()
important_atoms = np.where(node_importance > node_importance.mean())[0]
print(f"\nImportant atoms (indices): {important_atoms.tolist()}")
# Check whether the benzene ring atoms (0-5) are important
benzene_ring = [0, 1, 2, 3, 4, 5]
benzene_importance = np.mean([node_importance[i] for i in benzene_ring])
other_importance = np.mean([node_importance[i] for i in range(6, data.num_nodes)])
print(f"\nMean importance of benzene ring: {benzene_importance:.4f}")
print(f"Mean importance of other atoms: {other_importance:.4f}")
if benzene_importance > other_importance:
print("β
The benzene ring strongly contributes to toxicity!")
else:
print("β The benzene ring's contribution is lower than other parts.")
# Step 7: Visualization
fig, axes = plt.subplots(1, 2, figsize=(14, 5))
# Molecular structure
img = Draw.MolToImage(mol, size=(400, 400))
axes[0].imshow(img)
axes[0].set_title(f'Molecular Structure\n{smiles}', fontsize=12)
axes[0].axis('off')
# Atom importance
axes[1].bar(range(data.num_nodes), node_importance, color='steelblue')
axes[1].axhline(y=node_importance.mean(), color='r', linestyle='--', label='Mean')
axes[1].set_xlabel('Atom Index', fontsize=12)
axes[1].set_ylabel('Importance', fontsize=12)
axes[1].set_title('GNNExplainer: Toxicity Contribution per Atom', fontsize=13)
axes[1].legend()
axes[1].grid(True, alpha=0.3, axis='y')
plt.tight_layout()
plt.show()
SMILES: c1ccccc1CC(=O)O
Predicted toxicity score: 0.7234
Important atoms (indices): [0, 1, 2, 3, 4, 5]
Mean importance of benzene ring: 0.8523
Mean importance of other atoms: 0.3241
β
The benzene ring strongly contributes to toxicity!
References
-
Ying, Z., et al. (2018). "Hierarchical Graph Representation Learning with Differentiable Pooling." NeurIPS 2018. URL: https://arxiv.org/abs/1806.08804 DiffPool paper. Pioneering research on differentiable graph pooling.
-
SchΓΌtt, K., et al. (2017). "SchNet: A continuous-filter convolutional neural network for modeling quantum interactions." NeurIPS 2017. DOI: 10.5555/3294771.3294866 SchNet paper. Foundation of GNNs that consider 3D information.
-
Klicpera, J., et al. (2020). "Directional Message Passing for Molecular Graphs." ICLR 2020. URL: https://arxiv.org/abs/2003.03123 DimeNet paper. High-accuracy GNN considering bond angles.
-
Batzner, S., et al. (2022). "E(3)-equivariant graph neural networks for data-efficient and accurate interatomic potentials." Nature Communications, 13, 2453. DOI: 10.1038/s41467-022-29939-5 NequIP paper. Latest research on equivariant GNNs.
-
VeliΔkoviΔ, P., et al. (2018). "Graph Attention Networks." ICLR 2018. URL: https://arxiv.org/abs/1710.10903 GAT paper. Pioneering research introducing attention mechanisms into GNNs.
-
Ying, R., et al. (2019). "GNNExplainer: Generating Explanations for Graph Neural Networks." NeurIPS 2019. URL: https://arxiv.org/abs/1903.03894 GNNExplainer paper. Realizes interpretability of GNNs.
Created: 2025-10-17 Version: 1.0 Template: chapter-template-v2.0 Author: GNN Introduction Series Project