Graph Isomorphism Networks (GIN)

A Bit of History

Keyulu Xu, Weihua Hu, Jure Leskovec, and Stefanie Jegelka published "How Powerful are Graph Neural Networks?" at ICLR 2019, directly answering a question left open by the GCN (2017) and GraphSAGE (2017) papers: is there a theoretical limit to what message-passing GNNs can distinguish? Their answer connected modern deep learning directly back to a much older piece of graph theory — the 1-Weisfeiler-Leman (1-WL) test, the 1968 Soviet color-refinement heuristic for graph isomorphism (covered in its own deep dive). Xu et al. proved that no message-passing GNN, regardless of architecture, can be more powerful at distinguishing non-isomorphic graphs than 1-WL — and then designed GIN specifically to achieve that theoretical ceiling.

The Expressiveness Question: 1-WL as the Ceiling

The 1-WL test iteratively refines a coloring of graph vertices: at each round, every vertex's new color is a hash of its own current color plus the multiset of its neighbors' colors. Two graphs are declared "possibly isomorphic" if this process converges to the same color distribution for both; if the distributions differ, the graphs are provably not isomorphic.

Xu et al.'s central theorem: a GNN's aggregation function is only as powerful as 1-WL if and only if both its neighbor-aggregation step and its node-update step are injective functions over multisets — meaning distinct multisets of neighbor features must always map to distinct aggregated outputs. This is precisely what most existing architectures at the time failed to guarantee.

Working Principle: Injective Sum Aggregation

GIN's layer-wise update rule is designed explicitly to satisfy the injectivity requirement, using a learnable scalar $\epsilon$ and a multi-layer perceptron (MLP) as the injective function (justified by the universal approximation theorem):

$$h_v^{(k)} = \text{MLP}^{(k)} \left( (1 + \epsilon^{(k)}) \cdot h_v^{(k-1)} + \sum_{u \in \mathcal{N}(v)} h_u^{(k-1)} \right)$$

• Sum (not mean or max) over neighbor representations — the only multiset aggregator that is provably injective for countable input spaces.
• $(1 + \epsilon^{(k)})$: a learnable weight distinguishing the node's own previous representation from its neighbors' contribution — analogous to GraphSAGE's self/neighbor concatenation, but additive.
• MLP: a multi-layer perceptron (not a single linear layer) — necessary because a single linear function cannot approximate arbitrary injective functions over multisets, but an MLP with sufficient width theoretically can.

For graph-level tasks (not just node-level), GIN additionally sums (not averages or maxes) node representations across every layer of the network to form the final graph representation — again preserving injectivity at the whole-graph level, a technique the paper calls "graph-level readout with jumping knowledge."

Worked Example: Why Sum Beats Mean/Max

Complexity Analysis

$$\text{Time per layer: } O(|E| \cdot d + |V| \cdot \text{MLP cost}) \qquad \text{Space: } O(|V| \cdot d)$$

Asymptotically identical to GCN and GraphSAGE — GIN's theoretical superiority in expressiveness comes at essentially no extra computational cost, only a modest increase in parameters from using an MLP rather than a single linear transformation per layer.

Implementation

import torch
import torch.nn as nn

class GINLayer(nn.Module):
    def __init__(self, in_features, hidden_features, out_features, eps=0.0, train_eps=True):
        super(GINLayer, self).__init__()
        self.mlp = nn.Sequential(
            nn.Linear(in_features, hidden_features),
            nn.ReLU(),
            nn.Linear(hidden_features, out_features)
        )
        if train_eps:
            self.eps = nn.Parameter(torch.tensor([eps], dtype=torch.float32))
        else:
            self.register_buffer('eps', torch.tensor([eps], dtype=torch.float32))

    def forward(self, x, adj):
        """
        x: node features (N, in_features)
        adj: adjacency matrix (N, N), WITHOUT self-loops (GIN adds them via eps term)
        """
        neighbor_sum = torch.matmul(adj, x)  # SUM aggregation (not mean!)
        out = (1 + self.eps) * x + neighbor_sum
        return self.mlp(out)


class GIN(nn.Module):
    """Multi-layer GIN with graph-level sum readout across all layers."""
    def __init__(self, in_features, hidden_features, num_layers=3):
        super(GIN, self).__init__()
        self.layers = nn.ModuleList()
        self.layers.append(GINLayer(in_features, hidden_features, hidden_features))
        for _ in range(num_layers - 1):
            self.layers.append(GINLayer(hidden_features, hidden_features, hidden_features))

    def forward(self, x, adj):
        layer_outputs = [x]
        h = x
        for layer in self.layers:
            h = layer(h, adj)
            layer_outputs.append(h)

        # Graph-level readout: SUM each layer's node features, then SUM across nodes
        graph_repr = sum(layer_out.sum(dim=0) for layer_out in layer_outputs)
        return graph_repr, h  # (graph-level embedding, final node embeddings)


# Test on a small 4-node graph (triangle + pendant)
x = torch.tensor([[1.0, 0.0], [0.0, 1.0], [1.0, 1.0], [0.5, 0.5]])
adj = torch.tensor([
    [0, 1, 1, 0],
    [1, 0, 1, 0],
    [1, 1, 0, 1],
    [0, 0, 1, 0]
], dtype=torch.float32)

gin = GIN(in_features=2, hidden_features=8, num_layers=3)
graph_embedding, node_embeddings = gin(x, adj)

print("Graph-level embedding:", graph_embedding)
print("\nFinal node embeddings:")
print(node_embeddings)

Real-World Applications

Exercises

1. Construct two small non-isomorphic graphs (e.g., a 6-cycle vs. two disjoint triangles) and show that mean-aggregation GCN produces identical node-feature statistics for both, while sum-aggregation GIN distinguishes them.
2. Explain why a single linear layer cannot serve as GIN's injective update function, but a sufficiently wide MLP theoretically can (hint: universal approximation theorem).
3. Prove that GIN's expressiveness cannot exceed 1-WL, even with an arbitrarily powerful MLP (hint: consider what information is available to the aggregation step in the first place).
4. Challenge: Implement a "regular graph" pair (two non-isomorphic k-regular graphs known to fool 1-WL) and verify that GIN also fails to distinguish them, confirming the theoretical ceiling in practice.

Limitations