Chapter 2: Fundamentals of GNN Theory

This chapter organizes the basic mechanism of message passing so you can grasp it intuitively, even without equations. It clarifies the differences between representative models and how to choose among them.

💡 Supplement: Understanding comes faster if you separately consider three points: how much is transmitted (weights), how many times (number of layers), and how it is received (aggregation).

From Message Passing to Materials-Science-Specific GNNs

Learning Objectives

By reading this chapter, you will acquire the following:

• ✅ Understand the mathematical definition and representation of graphs
• ✅ Explain the 3 steps of message passing (aggregation → update → output)
• ✅ Understand the principles and differences of GCN, GAT, and GraphSAGE
• ✅ Learn the characteristics of materials-science-specific GNNs (SchNet, DimeNet)
• ✅ Implement a simple GNN in PyTorch
• ✅ Understand the importance of equivariant GNNs

Reading time: 25-30 min Code examples: 10 Exercises: 3

2.1 Mathematical Definition of Graphs

Basic Elements of a Graph

Definition:

A graph $G = (V, E)$ consists of a vertex set $V$ and an edge set $E \subseteq V \times V$.

Notation: - $n = |V|$: number of vertices - $m = |E|$: number of edges - $\mathcal{N}(v)$: set of neighbors of vertex $v$

Adjacency Matrix

Definition: $$ A \in {0, 1}^{n \times n}, \quad A_{ij} = \begin{cases} 1 & \text{if } (v_i, v_j) \in E \ 0 & \text{otherwise} \end{cases} $$

Implementation in Python:

import numpy as np

# Example: triangle graph (3 vertices, 3 edges)
n = 3
A = np.array([
    [0, 1, 1],  # Vertex 0: connected to 1, 2
    [1, 0, 1],  # Vertex 1: connected to 0, 2
    [1, 1, 0]   # Vertex 2: connected to 0, 1
])

print("Adjacency matrix:")
print(A)
print(f"\nNumber of vertices: {n}")
print(f"Number of edges: {A.sum() // 2}")  # Divide by 2 for undirected graphs

Output:

Adjacency matrix:
[[0 1 1]
 [1 0 1]
 [1 1 0]]

Number of vertices: 3
Number of edges: 3

Degree Matrix

Definition: $$ D \in \mathbb{R}^{n \times n}, \quad D_{ii} = \sum_{j=1}^{n} A_{ij} $$

Physical meaning: the number of connections of each vertex (in chemistry, the number of bonds)

# Degree matrix
D = np.diag(A.sum(axis=1))
print("Degree matrix:")
print(D)
print(f"\nDegree of each vertex: {np.diag(D)}")

Output:

Degree matrix:
[[2 0 0]
 [0 2 0]
 [0 0 2]]

Degree of each vertex: [2 2 2]

Laplacian Matrix

Definition: $$ L = D - A $$

Normalized Laplacian (commonly used in GNNs): $$ \tilde{L} = D^{-1/2} L D^{-1/2} = I - D^{-1/2} A D^{-1/2} $$

# Laplacian matrix
L = D - A
print("Laplacian matrix:")
print(L)

# Normalized Laplacian
D_inv_sqrt = np.diag(1 / np.sqrt(np.diag(D)))
L_norm = np.eye(n) - D_inv_sqrt @ A @ D_inv_sqrt
print("\nNormalized Laplacian:")
print(L_norm)

Output:

Laplacian matrix:
[[ 2 -1 -1]
 [-1  2 -1]
 [-1 -1  2]]

Normalized Laplacian:
[[ 1.  -0.5 -0.5]
 [-0.5  1.  -0.5]
 [-0.5 -0.5  1. ]]

Applications: - Spectral graph theory - Graph Fourier transform - Graph signal processing

Node Features and Edge Features

Node feature matrix $X \in \mathbb{R}^{n \times d}$: - Each row $x_i \in \mathbb{R}^d$: feature vector of vertex $i$ - Materials science: atomic number, electronegativity, number of valence electrons, etc.

Edge feature matrix $E \in \mathbb{R}^{m \times d_e}$: - Each row $e_{ij} \in \mathbb{R}^{d_e}$: feature of edge $(i, j)$ - Materials science: bond length, bond order, bond angle, etc.

# Example: features of a water molecule (H₂O)
X = np.array([
    [8, 2.55, 6],   # O: atomic number 8, electronegativity 2.55, 6 valence electrons
    [1, 2.20, 1],   # H1
    [1, 2.20, 1]    # H2
])

print("Node feature matrix (3×3):")
print(X)
print(f"Shape: {X.shape}")

2.2 How Message Passing Works

Message Passing Neural Network (MPNN)

This is the unified framework for GNNs (Gilmer et al., 2017).

Algorithm:

Step 1: Message Generation (Message)

Definition: $$ m_{ij}^{(t)} = \text{Message}(h_i^{(t)}, h_j^{(t)}, e_{ij}) $$

• $h_i^{(t)}$: hidden state of vertex $i$ at layer $t$
• $h_j^{(t)}$: hidden state of neighboring vertex $j$
• $e_{ij}$: feature of edge $(i, j)$

The simplest form: $$ m_{ij}^{(t)} = W \cdot h_j^{(t)} $$

import torch
import torch.nn as nn

class MessageFunction(nn.Module):
    def __init__(self, in_dim, out_dim):
        super().__init__()
        self.W = nn.Linear(in_dim, out_dim)

    def forward(self, h_j):
        """
        Generate messages from neighboring vertices

        Parameters:
        -----------
        h_j : Tensor (num_neighbors, in_dim)
            Features of neighboring vertices

        Returns:
        --------
        messages : Tensor (num_neighbors, out_dim)
            Generated messages
        """
        return self.W(h_j)

# Example
in_dim, out_dim = 16, 32
msg_fn = MessageFunction(in_dim, out_dim)

# Features of neighboring vertices (3 neighbors)
h_neighbors = torch.randn(3, in_dim)
messages = msg_fn(h_neighbors)
print(f"Message shape: {messages.shape}")
# Output: torch.Size([3, 32])

Step 2: Aggregation

Definition: $$ m_i^{(t)} = \text{Aggregate}\left( {m_{ij}^{(t)} : j \in \mathcal{N}(i)} \right) $$

Representative aggregation functions:

class AggregationFunction:
    @staticmethod
    def sum_agg(messages):
        """Sum aggregation"""
        return torch.sum(messages, dim=0)

    @staticmethod
    def mean_agg(messages):
        """Mean aggregation"""
        return torch.mean(messages, dim=0)

    @staticmethod
    def max_agg(messages):
        """Max aggregation"""
        return torch.max(messages, dim=0)[0]

# Example
messages = torch.tensor([
    [1.0, 2.0, 3.0],
    [4.0, 5.0, 6.0],
    [7.0, 8.0, 9.0]
])

print("Sum:", AggregationFunction.sum_agg(messages))
# Output: tensor([12., 15., 18.])

print("Mean:", AggregationFunction.mean_agg(messages))
# Output: tensor([4., 5., 6.])

print("Max:", AggregationFunction.max_agg(messages))
# Output: tensor([7., 8., 9.])

Step 3: Update

Definition: $$ h_i^{(t+1)} = \text{Update}\left( h_i^{(t)}, m_i^{(t)} \right) $$

Typical update formula: $$ h_i^{(t+1)} = \sigma\left( W_1 h_i^{(t)} + W_2 m_i^{(t)} \right) $$

class UpdateFunction(nn.Module):
    def __init__(self, hidden_dim):
        super().__init__()
        self.W1 = nn.Linear(hidden_dim, hidden_dim)
        self.W2 = nn.Linear(hidden_dim, hidden_dim)
        self.activation = nn.ReLU()

    def forward(self, h_i, m_i):
        """
        Update node features

        Parameters:
        -----------
        h_i : Tensor (hidden_dim,)
            Current node features
        m_i : Tensor (hidden_dim,)
            Aggregated message

        Returns:
        --------
        h_new : Tensor (hidden_dim,)
            Updated node features
        """
        return self.activation(self.W1(h_i) + self.W2(m_i))

# Example
hidden_dim = 32
update_fn = UpdateFunction(hidden_dim)

h_current = torch.randn(hidden_dim)
m_aggregated = torch.randn(hidden_dim)
h_new = update_fn(h_current, m_aggregated)

print(f"Before update: {h_current[:5]}")
print(f"After update: {h_new[:5]}")

The Overall Picture of Message Passing

class SimpleGNN(nn.Module):
    def __init__(self, in_dim, hidden_dim, num_layers):
        super().__init__()
        self.num_layers = num_layers

        # Parameters for each layer
        self.message_fns = nn.ModuleList([
            MessageFunction(hidden_dim, hidden_dim)
            for _ in range(num_layers)
        ])
        self.update_fns = nn.ModuleList([
            UpdateFunction(hidden_dim)
            for _ in range(num_layers)
        ])

        # Input transformation
        self.input_proj = nn.Linear(in_dim, hidden_dim)

    def forward(self, x, edge_index):
        """
        Parameters:
        -----------
        x : Tensor (num_nodes, in_dim)
            Node feature matrix
        edge_index : Tensor (2, num_edges)
            Edge list [[src], [dst]]

        Returns:
        --------
        h : Tensor (num_nodes, hidden_dim)
            Updated node features
        """
        # Input transformation
        h = self.input_proj(x)

        # Message passing layers
        for layer in range(self.num_layers):
            h_new = []

            # Update each vertex
            for i in range(x.size(0)):
                # Get neighboring vertices
                neighbors = edge_index[1][edge_index[0] == i]

                if len(neighbors) > 0:
                    # Step 1: Message generation
                    messages = self.message_fns[layer](h[neighbors])

                    # Step 2: Aggregation
                    m_i = torch.mean(messages, dim=0)

                    # Step 3: Update
                    h_i_new = self.update_fns[layer](h[i], m_i)
                else:
                    # If there are no neighbors
                    h_i_new = h[i]

                h_new.append(h_i_new)

            h = torch.stack(h_new)

        return h

# Usage example
model = SimpleGNN(in_dim=16, hidden_dim=32, num_layers=3)

# Graph data (triangle)
x = torch.randn(3, 16)  # 3 vertices, 16-dim features
edge_index = torch.tensor([
    [0, 0, 1, 1, 2, 2],  # Source
    [1, 2, 0, 2, 0, 1]   # Target
])

# Forward pass
h_out = model(x, edge_index)
print(f"Output shape: {h_out.shape}")
# Output: torch.Size([3, 32])

2.3 Representative GNN Architectures

Graph Convolutional Network (GCN)

Paper: Kipf & Welling (2017), ICLR

Core idea: spectral convolution on graphs

Update formula: $$ H^{(l+1)} = \sigma\left( \tilde{D}^{-1/2} \tilde{A} \tilde{D}^{-1/2} H^{(l)} W^{(l)} \right) $$

• $\tilde{A} = A + I$: adjacency matrix with self-loops
• $\tilde{D}_{ii} = \sum_j \tilde{A}_{ij}$: degree matrix
• $H^{(l)} \in \mathbb{R}^{n \times d}$: features at layer $l$
• $W^{(l)} \in \mathbb{R}^{d \times d'}$: learnable weights

Implementation in Python:

import torch
import torch.nn as nn
import torch.nn.functional as F

class GCNLayer(nn.Module):
    def __init__(self, in_features, out_features):
        super().__init__()
        self.linear = nn.Linear(in_features, out_features)

    def forward(self, X, A):
        """
        Parameters:
        -----------
        X : Tensor (num_nodes, in_features)
            Node feature matrix
        A : Tensor (num_nodes, num_nodes)
            Adjacency matrix

        Returns:
        --------
        H : Tensor (num_nodes, out_features)
            Updated features
        """
        # Add self-loops
        A_tilde = A + torch.eye(A.size(0), device=A.device)

        # Degree matrix
        D_tilde = torch.diag(A_tilde.sum(dim=1))

        # Normalization: D^(-1/2) * A * D^(-1/2)
        D_inv_sqrt = torch.diag(1.0 / torch.sqrt(D_tilde.diagonal()))
        A_norm = D_inv_sqrt @ A_tilde @ D_inv_sqrt

        # Graph convolution
        H = A_norm @ X
        H = self.linear(H)
        return F.relu(H)

# Usage example
gcn = GCNLayer(in_features=16, out_features=32)

# Graph data
X = torch.randn(5, 16)  # 5 vertices, 16 dims
A = torch.tensor([
    [0, 1, 1, 0, 0],
    [1, 0, 1, 1, 0],
    [1, 1, 0, 1, 1],
    [0, 1, 1, 0, 1],
    [0, 0, 1, 1, 0]
], dtype=torch.float32)

H = gcn(X, A)
print(f"GCN output shape: {H.shape}")
# Output: torch.Size([5, 32])

Characteristics: - ✅ Simple and fast - ✅ Beware of over-smoothing - ✅ Fixed weights (all neighbors treated equally)

Graph Attention Network (GAT)

Paper: Veličković et al. (2018), ICLR

Core idea: use attention to emphasize important neighbors

Attention coefficients: $$ \alpha_{ij} = \frac{\exp\left( \text{LeakyReLU}(a^T [W h_i | W h_j]) \right)} {\sum_{k \in \mathcal{N}(i)} \exp\left( \text{LeakyReLU}(a^T [W h_i | W h_k]) \right)} $$

Update formula: $$ h_i^{(l+1)} = \sigma\left( \sum_{j \in \mathcal{N}(i)} \alpha_{ij} W^{(l)} h_j^{(l)} \right) $$

class GATLayer(nn.Module):
    def __init__(self, in_features, out_features, dropout=0.6,
                 alpha=0.2):
        super().__init__()
        self.W = nn.Linear(in_features, out_features, bias=False)
        self.a = nn.Parameter(torch.zeros(2 * out_features, 1))
        self.leakyrelu = nn.LeakyReLU(alpha)
        self.dropout = nn.Dropout(dropout)

        nn.init.xavier_uniform_(self.a.data, gain=1.414)

    def forward(self, X, A):
        """
        Parameters:
        -----------
        X : Tensor (num_nodes, in_features)
        A : Tensor (num_nodes, num_nodes)

        Returns:
        --------
        H : Tensor (num_nodes, out_features)
        """
        # Linear transformation
        Wh = self.W(X)  # (N, out_features)
        N = Wh.size(0)

        # Attention computation
        # [Wh_i || Wh_j] for all edges
        Wh_repeat_interleave = Wh.repeat_interleave(N, dim=0)
        Wh_repeat = Wh.repeat(N, 1)
        concat = torch.cat([Wh_repeat_interleave, Wh_repeat], dim=1)
        concat = concat.view(N, N, -1)

        # Attention score
        e = self.leakyrelu(concat @ self.a).squeeze(2)

        # Mask (-inf where there is no edge)
        zero_vec = -9e15 * torch.ones_like(e)
        attention = torch.where(A > 0, e, zero_vec)

        # Softmax
        attention = F.softmax(attention, dim=1)
        attention = self.dropout(attention)

        # Weighted sum
        H = torch.matmul(attention, Wh)
        return F.elu(H)

# Usage example
gat = GATLayer(in_features=16, out_features=32)
H_gat = gat(X, A)
print(f"GAT output shape: {H_gat.shape}")
# Output: torch.Size([5, 32])

Characteristics: - ✅ Dynamic weights (automatically learns important neighbors) - ✅ Interpretability (visualization of attention coefficients) - ❌ High computational cost (about 2× that of GCN)

GraphSAGE (SAmple and aggreGatE)

Paper: Hamilton et al. (2017), NeurIPS

Core idea: sampling for mini-batch training

Update formula: $$ h_i^{(l+1)} = \sigma\left( W \cdot \text{Concat}\left( h_i^{(l)}, \text{Aggregate}({h_j^{(l)} : j \in \mathcal{S}(i)}) \right) \right) $$

• $\mathcal{S}(i)$: sampled neighbors (not all of them)

class GraphSAGELayer(nn.Module):
    def __init__(self, in_features, out_features, num_samples=10):
        super().__init__()
        self.num_samples = num_samples
        # Concat version: input is in_features * 2
        self.linear = nn.Linear(in_features * 2, out_features)

    def forward(self, X, A):
        """
        Parameters:
        -----------
        X : Tensor (num_nodes, in_features)
        A : Tensor (num_nodes, num_nodes)

        Returns:
        --------
        H : Tensor (num_nodes, out_features)
        """
        N = X.size(0)
        H_new = []

        for i in range(N):
            # Sample neighboring vertices
            neighbors = torch.nonzero(A[i]).squeeze()
            if neighbors.numel() > self.num_samples:
                # Random sampling
                perm = torch.randperm(neighbors.numel())
                sampled = neighbors[perm[:self.num_samples]]
            else:
                sampled = neighbors

            # Aggregation (Mean)
            if sampled.numel() > 0:
                h_neighbors = X[sampled]
                h_agg = torch.mean(h_neighbors, dim=0)
            else:
                h_agg = torch.zeros_like(X[i])

            # Concat
            h_concat = torch.cat([X[i], h_agg], dim=0)

            # Linear transformation
            h_new = self.linear(h_concat)
            H_new.append(h_new)

        H = torch.stack(H_new)
        return F.relu(H)

# Usage example
sage = GraphSAGELayer(in_features=16, out_features=32,
                      num_samples=3)
H_sage = sage(X, A)
print(f"GraphSAGE output shape: {H_sage.shape}")
# Output: torch.Size([5, 32])

Characteristics: - ✅ Scalable (handles large graphs) - ✅ Mini-batch training is possible - ✅ Inductive learning (generalizes to new vertices)

Comparison of the Three GNNs

• $m$: number of edges
• $n$: number of vertices
• $d$: feature dimension
• $k$: number of layers
• $s$: number of samples

2.4 Materials-Science-Specific GNNs

SchNet (Continuous-filter Convolutional NN)

Paper: Schütt et al. (2017), NeurIPS

Target: prediction of the quantum-chemical properties of molecules and materials

Core idea: 1. Continuous filters: convolution in 3D space rather than on a discrete graph 2. Distance dependence: explicitly models interatomic distances

Architecture:

Formula: $$ h_i^{(l+1)} = h_i^{(l)} + \sum_{j \in \mathcal{N}(i)} h_j^{(l)} \odot \phi\left( |r_i - r_j| \right) $$

• $\phi(d)$: continuous filter function (depends on distance $d$)
• $r_i, r_j$: 3D coordinates of the atoms

Filter function: $$ \phi(d) = \sum_{k=1}^{K} w_k \exp\left( -\gamma (d - \mu_k)^2 \right) $$

• Gaussian basis expansion (RBF: Radial Basis Function)

import torch
import torch.nn as nn

class GaussianBasis(nn.Module):
    def __init__(self, start=0.0, stop=5.0, num_gaussians=50):
        super().__init__()
        self.mu = nn.Parameter(
            torch.linspace(start, stop, num_gaussians),
            requires_grad=False
        )
        self.gamma = nn.Parameter(
            torch.tensor(10.0),
            requires_grad=True
        )

    def forward(self, distances):
        """
        Parameters:
        -----------
        distances : Tensor (num_edges,)
            Interatomic distances

        Returns:
        --------
        rbf : Tensor (num_edges, num_gaussians)
            Gaussian basis expansion
        """
        # (num_edges, 1) - (1, num_gaussians)
        diff = distances.unsqueeze(-1) - self.mu.unsqueeze(0)
        rbf = torch.exp(-self.gamma * diff ** 2)
        return rbf

class SchNetInteraction(nn.Module):
    def __init__(self, hidden_dim, num_gaussians):
        super().__init__()
        self.rbf_layer = GaussianBasis(num_gaussians=num_gaussians)
        self.filter_net = nn.Sequential(
            nn.Linear(num_gaussians, hidden_dim),
            nn.Softplus(),
            nn.Linear(hidden_dim, hidden_dim)
        )
        self.linear = nn.Linear(hidden_dim, hidden_dim)

    def forward(self, h, edge_index, distances):
        """
        Parameters:
        -----------
        h : Tensor (num_atoms, hidden_dim)
            Atom features
        edge_index : Tensor (2, num_edges)
            Edge list
        distances : Tensor (num_edges,)
            Interatomic distances

        Returns:
        --------
        h_new : Tensor (num_atoms, hidden_dim)
            Updated features
        """
        # RBF expansion
        rbf = self.rbf_layer(distances)

        # Generate filters
        W = self.filter_net(rbf)

        # Message passing
        src, dst = edge_index
        messages = h[dst] * W  # Element-wise product

        # Aggregation
        h_agg = torch.zeros_like(h)
        h_agg.index_add_(0, src, messages)

        # Update
        h_new = h + self.linear(h_agg)
        return h_new

# Usage example
schnet_layer = SchNetInteraction(hidden_dim=128,
                                 num_gaussians=50)

# Data
num_atoms = 5
h = torch.randn(num_atoms, 128)
edge_index = torch.tensor([[0, 1, 2, 3], [1, 2, 3, 4]])
distances = torch.tensor([1.5, 1.8, 2.0, 1.6])

h_new = schnet_layer(h, edge_index, distances)
print(f"SchNet output shape: {h_new.shape}")
# Output: torch.Size([5, 128])

Application examples: - QM9 dataset (molecular property prediction) - MD17 (molecular dynamics) - OC20 (catalyst adsorption energy)

Performance:

QM9 HOMO-LUMO gap:
- DFT calculation: 24 hours/molecule
- SchNet: 0.01 sec/molecule (MAE=0.04 eV)

DimeNet (Directional Message Passing NN)

Paper: Klicpera et al. (2020), ICLR

Extension: also considers bond angles

Core idea: - Uses not only distance but also angle information - Three-body interactions (triplet interaction)

Update formula: $$ m_{ij} = \sum_{k \in \mathcal{N}(j) \setminus {i}} W\left( d_{ij}, d_{jk}, \theta_{ijk} \right) h_k $$

• $\theta_{ijk}$: angle $\angle i-j-k$

Computing the angle:

import torch

def compute_angle(pos_i, pos_j, pos_k):
    """
    Compute the angle among three atoms

    Parameters:
    -----------
    pos_i, pos_j, pos_k : Tensor (3,)
        3D coordinates of the atoms

    Returns:
    --------
    angle : Tensor (1,)
        Angle (radians)
    """
    # Vectors
    v_ij = pos_j - pos_i
    v_jk = pos_k - pos_j

    # Dot product
    cos_angle = torch.dot(v_ij, v_jk) / (
        torch.norm(v_ij) * torch.norm(v_jk) + 1e-8
    )

    # Angle
    angle = torch.acos(torch.clamp(cos_angle, -1.0, 1.0))
    return angle

# Example: bond angle of a water molecule (H-O-H)
pos_O = torch.tensor([0.0, 0.0, 0.0])
pos_H1 = torch.tensor([0.96, 0.0, 0.0])
pos_H2 = torch.tensor([0.24, 0.93, 0.0])

angle = compute_angle(pos_H1, pos_O, pos_H2)
print(f"H-O-H angle: {torch.rad2deg(angle):.1f}°")
# Output: 104.5° (nearly matches the experimental value)

Performance:

QM9 dataset:
- SchNet: MAE=0.041 eV
- DimeNet: MAE=0.033 eV (20% improvement)

Computation time:
- SchNet: 0.01 sec/molecule
- DimeNet: 0.05 sec/molecule (5× slower)

GemNet (Geometric Message Passing NN)

Paper: Gasteiger et al. (2021), NeurIPS

Further extension: four-body interactions (dihedral angles)

Target: crystal structures, complex molecules

Core idea: - Consideration of dihedral (torsion) angles - Higher-order geometric information

Performance:

OC20 dataset (catalysis):
- SchNet: MAE=0.61 eV
- DimeNet++: MAE=0.49 eV
- GemNet: MAE=0.43 eV (highest accuracy)

Comparison of Materials-Science GNNs

2.5 The Importance of Equivariance

What Is Equivariance?

Definition:

A function $f$ is equivariant with respect to a transformation $T$ if $$f(T(x)) = T(f(x))$$ holds.

Meaning in materials science: - Even when a molecule is rotated or translated, the prediction is the same (or transforms correspondingly)

E(3) Equivariance

The E(3) group: isometric transformations of 3D Euclidean space - Rotation - Translation - Inversion

Importance: - Physical laws do not depend on the coordinate system - GNNs should behave the same way

Examples of Equivariant GNNs: NequIP, MACE

NequIP (Batzner et al., 2022): - E(3)-equivariant message passing - Use of spherical harmonics

Update formula: $$ m_{ij} = \phi\left( |r_i - r_j| \right) \otimes Y_l(r_{ij}) $$

• $Y_l$: spherical harmonics (retain angular information)
• $\otimes$: tensor product

MACE (Batatia et al., 2022): - Higher-order equivariance - More accurate force-field prediction

Performance:

MD17 dataset (molecular dynamics):
- SchNet: MAE(force) = 0.21 kcal/mol/Å
- NequIP: MAE(force) = 0.05 kcal/mol/Å (76% improvement)

Testing Equivariance

import torch
import torch.nn as nn

def test_equivariance(model, pos, edge_index):
    """
    Test the model's equivariance
    """
    # Original prediction
    pred_original = model(pos, edge_index)

    # Rotation matrix (90-degree rotation)
    angle = torch.tensor(torch.pi / 2)
    rotation = torch.tensor([
        [torch.cos(angle), -torch.sin(angle), 0],
        [torch.sin(angle), torch.cos(angle), 0],
        [0, 0, 1]
    ])

    # Rotate coordinates
    pos_rotated = pos @ rotation.T

    # Prediction after rotation
    pred_rotated = model(pos_rotated, edge_index)

    # Rotate the prediction
    pred_original_rotated = pred_original @ rotation.T

    # Compute the error
    error = torch.abs(pred_rotated - pred_original_rotated).mean()
    print(f"Equivariance error: {error.item():.6f}")

    if error < 1e-5:
        print("✅ The model is equivariant")
    else:
        print("❌ The model is not equivariant")

# Usage example (simplified)
class SimpleEquivariantModel(nn.Module):
    def forward(self, pos, edge_index):
        # Simplified: compute coordinate differences (equivariant)
        src, dst = edge_index
        diff = pos[dst] - pos[src]
        return diff

model = SimpleEquivariantModel()
pos = torch.randn(5, 3)
edge_index = torch.tensor([[0, 1, 2], [1, 2, 3]])

test_equivariance(model, pos, edge_index)

2.6 Column: Why Deep GNNs Are Difficult

Over-smoothing

Problem: as layers get deeper, all vertices end up with the same features

Cause: repeated message passing diffuses information

# Over-smoothing demo
import torch
import torch.nn.functional as F

def demonstrate_oversmoothing(X, A, num_layers=10):
    """
    Visualize over-smoothing
    """
    H = X
    smoothness = []

    for layer in range(num_layers):
        # Simple GCN layer
        D = torch.diag(A.sum(dim=1))
        D_inv_sqrt = torch.diag(1.0 / torch.sqrt(D.diagonal()))
        A_norm = D_inv_sqrt @ A @ D_inv_sqrt

        H = A_norm @ H
        H = F.relu(H)

        # Smoothness (similarity between vertices)
        similarity = F.cosine_similarity(
            H.unsqueeze(1), H.unsqueeze(0), dim=2
        )
        avg_similarity = similarity[torch.triu_indices(
            H.size(0), H.size(0), offset=1
        )[0], torch.triu_indices(
            H.size(0), H.size(0), offset=1
        )[1]].mean()

        smoothness.append(avg_similarity.item())
        print(f"Layer {layer+1}: Average similarity = {avg_similarity:.4f}")

    return smoothness

# Run
X = torch.randn(5, 16)
A = torch.eye(5) + torch.rand(5, 5) > 0.7
smoothness = demonstrate_oversmoothing(X, A.float(), num_layers=10)

Example output:

Layer 1: Average similarity = 0.2341
Layer 2: Average similarity = 0.4523
Layer 3: Average similarity = 0.6789
...
Layer 10: Average similarity = 0.9876

→ As layers get deeper, all vertices become similar

Countermeasures

1. Residual Connection: $$h_i^{(l+1)} = h_i^{(l)} + \text{GNN}(h_i^{(l)})$$

2. Jumping Knowledge Network: - Concatenate the outputs of all layers

3. PairNorm: - Normalize the features

class GNNWithResidual(nn.Module):
    def \_\_init\_\_(self, hidden\_dim):
        super().\_\_init\_\_()
        self.conv = GCNLayer(hidden\_dim, hidden\_dim)

    def forward(self, X, A):
        # Residual connection
        H = self.conv(X, A)
        return X + H  # Shortcut

2.7 Chapter Summary

What We Learned

1. Mathematical definition of graphs - Adjacency, degree, and Laplacian matrices - Node and edge features

2. Message passing - 3 steps: message generation → aggregation → update - Aggregation functions: Sum, Mean, Max, Attention

3. Representative GNN architectures - GCN: simple, fast - GAT: attention, high accuracy - GraphSAGE: scalable, mini-batch

4. Materials-science-specific GNNs - SchNet: distance-dependent, continuous filters - DimeNet: also considers angle information - GemNet: considers even dihedral angles

5. Equivariance - The importance of E(3) equivariance - Latest methods such as NequIP and MACE

Key Points

• ✅ Message passing is the unified framework for GNNs
• ✅ The choice of aggregation function greatly affects performance
• ✅ In materials science, geometric information (distance, angle) is important
• ✅ Equivariance guarantees physical laws
• ✅ Beware of over-smoothing (counter it with residual connections)

To the Next Chapter

In Chapter 3, we will learn hands-on PyTorch Geometric: - Environment setup (PyG, RDKit, ASE) - Molecular property prediction with the QM9 dataset - Crystal property prediction with Materials Project data - Model evaluation and hyperparameter tuning - Hands-on project

Chapter 3: Hands-On PyTorch Geometric →

Exercises

Exercise 1 (Difficulty: Easy)

Determine whether the following statements are true or false.

1. Message passing is performed in the order aggregation → update → message generation
2. Because GAT uses attention, it treats all neighbors with the same weight
3. SchNet explicitly considers interatomic distances

# Correct order
for layer in range(num\_layers):
    # Step 1: Message generation
    messages = message\_function(h\_neighbors)

    # Step 2: Aggregation
    m\_i = aggregate(messages)

    # Step 3: Update
    h\_i = update\_function(h\_i, m\_i)

Exercise 2 (Difficulty: Medium)

For the following graph, compute one layer of GCN forward propagation by hand.

Graph:

Vertices: 3 (v0, v1, v2)
Edges: v0-v1, v1-v2 (linear graph)

Node features:
X = [[1, 0],
     [0, 1],
     [1, 1]]

Adjacency matrix:
A = [[0, 1, 0],
     [1, 0, 1],
     [0, 1, 0]]

Weight matrix (simplified):
W = [[1, 0],
     [0, 1]]  (identity matrix)

Requirements: 1. Compute $\tilde{A} = A + I$ 2. Compute the normalized adjacency matrix $\hat{A} = \tilde{D}^{-1/2} \tilde{A} \tilde{D}^{-1/2}$ 3. Compute the GCN output $H = \hat{A} X W$ (no activation function)

import numpy as np

A\_tilde = np.array([
    [1, 1, 0],
    [1, 1, 1],
    [0, 1, 1]
], dtype=float)

D\_tilde = np.diag([2, 3, 2])
D\_inv\_sqrt = np.diag([1/np.sqrt(2), 1/np.sqrt(3), 1/np.sqrt(2)])

A\_hat = D\_inv\_sqrt @ A\_tilde @ D\_inv\_sqrt
print("Normalized adjacency matrix:")
print(A\_hat)

X = np.array([
    [1, 0],
    [0, 1],
    [1, 1]
], dtype=float)

H = A\_hat @ X
print("GCN output:")
print(H)

# Complete code
A = np.array([[0, 1, 0], [1, 0, 1], [0, 1, 0]], dtype=float)
X = np.array([[1, 0], [0, 1], [1, 1]], dtype=float)

# GCN
A\_tilde = A + np.eye(3)
D\_tilde = np.diag(A\_tilde.sum(axis=1))
D\_inv\_sqrt = np.diag(1.0 / np.sqrt(D\_tilde.diagonal()))
A\_hat = D\_inv\_sqrt @ A\_tilde @ D\_inv\_sqrt

H = A\_hat @ X
print("Final output:")
print(H)

Exercise 3 (Difficulty: Hard)

Implement SchNet's continuous filter function and visualize the filter's response to different interatomic distances.

Requirements: 1. Implement a Gaussian basis (RBF) function 2. Compute the RBF response for distances from 0.5 Å to 5.0 Å 3. Visualize with a heatmap 4. Discuss the physical meaning of the filter

import torch
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns

# ===== Implementation =====
class GaussianBasisFunction:
    def __init__(self, start=0.0, stop=5.0, num_gaussians=50,
                 gamma=10.0):
        """
        Gaussian basis function (RBF)

        Parameters:
        -----------
        start, stop : float
            Distance range
        num_gaussians : int
            Number of Gaussians
        gamma : float
            Spread parameter
        """
        self.mu = torch.linspace(start, stop, num_gaussians)
        self.gamma = gamma

    def __call__(self, distances):
        """
        Compute the RBF response

        Parameters:
        -----------
        distances : Tensor (num_distances,)

        Returns:
        --------
        rbf : Tensor (num_distances, num_gaussians)
        """
        # (num_distances, 1) - (1, num_gaussians)
        diff = distances.unsqueeze(-1) - self.mu.unsqueeze(0)
        rbf = torch.exp(-self.gamma * diff ** 2)
        return rbf

# ===== Visualization =====
# Generate RBF
rbf_layer = GaussianBasisFunction(
    start=0.0, stop=5.0,
    num_gaussians=50, gamma=10.0
)

# Distance samples (0.5–5.0 Å)
distances = torch.linspace(0.5, 5.0, 100)

# RBF response
rbf_response = rbf_layer(distances)  # (100, 50)

# Heatmap
plt.figure(figsize=(12, 6))
sns.heatmap(
    rbf_response.T.numpy(),  # Transpose (RBF x distance)
    cmap='viridis',
    xticklabels=10,
    yticklabels=10,
    cbar_kws={'label': 'RBF Response'}
)
plt.xlabel('Distance (Å)')
plt.ylabel('RBF Index')
plt.title('SchNet Continuous Filter: RBF Response')

# Set X-axis labels to actual distances
xticks = np.linspace(0, len(distances)-1, 10).astype(int)
xticklabels = [f'{distances[i]:.1f}' for i in xticks]
plt.xticks(xticks, xticklabels)

plt.tight_layout()
plt.savefig('schnet_rbf_heatmap.png', dpi=150)
plt.show()

# ===== RBF response at specific distances =====
fig, axes = plt.subplots(2, 2, figsize=(12, 10))
example_distances = [1.0, 1.5, 2.0, 3.0]  # Å

for ax, d in zip(axes.flatten(), example_distances):
    d_tensor = torch.tensor([d])
    rbf = rbf_layer(d_tensor).squeeze()

    ax.plot(rbf_layer.mu.numpy(), rbf.numpy(),
            marker='o', linewidth=2)
    ax.axvline(d, color='red', linestyle='--',
               label=f'Distance = {d}Å')
    ax.set_xlabel('RBF Center μ (Å)')
    ax.set_ylabel('RBF Response')
    ax.set_title(f'RBF Response at d = {d}Å')
    ax.legend()
    ax.grid(True, alpha=0.3)

plt.tight_layout()
plt.savefig('schnet_rbf_profiles.png', dpi=150)
plt.show()

# ===== Discussion of physical meaning =====
print("\n===== Physical meaning =====")
print("1. Short range (0.5-2.0 Å): covalent bonding region")
print("   - C-C: 1.54 Å, C=C: 1.34 Å, C-H: 1.09 Å")
print("   - RBF responds sharply (distinguishes presence/absence of bonds)")

print("\n2. Medium range (2.0-3.5 Å): non-covalent interactions")
print("   - Hydrogen bonds: 2.8 Å, van der Waals forces")
print("   - RBF responds gradually")

print("\n3. Long range (3.5-5.0 Å): weak interactions")
print("   - Electrostatic interactions, dispersion forces")
print("   - RBF response is small")

print("\n4. Role of the Gaussian basis:")
print("   - Continuous distance representation (no discretization)")
print("   - Differentiable for any distance")
print("   - Optimizable via machine learning (γ parameter)")

References

1. Kipf, T. N. & Welling, M. (2017). "Semi-Supervised Classification with Graph Convolutional Networks." ICLR. DOI: https://arxiv.org/abs/1609.02907

2. Veličković, P. et al. (2018). "Graph Attention Networks." ICLR. DOI: https://arxiv.org/abs/1710.10903

3. Hamilton, W. L. et al. (2017). "Inductive Representation Learning on Large Graphs." NeurIPS. DOI: https://arxiv.org/abs/1706.02216

4. Schütt, K. T. et al. (2017). "SchNet: A continuous-filter convolutional neural network for modeling quantum interactions." NeurIPS. DOI: https://arxiv.org/abs/1706.08566

5. Klicpera, J. et al. (2020). "Directional Message Passing for Molecular Graphs." ICLR. DOI: https://arxiv.org/abs/2003.03123

6. Gasteiger, J. et al. (2021). "GemNet: Universal Directional Graph Neural Networks for Molecules." NeurIPS. DOI: https://arxiv.org/abs/2106.08903

7. Batzner, S. et al. (2022). "E(3)-equivariant graph neural networks for data-efficient and accurate interatomic potentials." Nature Communications, 13, 2453. DOI: https://doi.org/10.1038/s41467-022-29939-5

Navigation

Previous Chapter

Chapter 1: Why Materials Science Needs GNNs ←

Next Chapter

Chapter 3: Hands-On PyTorch Geometric →

Series Contents

← Back to Series Contents

Author Information

Author: AI Terakoya Content Team Created: 2025-10-17 Version: 1.0

Update history: - 2025-10-17: v1.0 initial release

Feedback: - GitHub Issues: [repository URL]/issues - Email: yusuke.hashimoto.b8@tohoku.ac.jp

License: Creative Commons BY 4.0

Let's actually run a GNN in Chapter 3!