GNN

Representation

• Node embedding: Each node embedded as a vector, and the entire graph represented as adjacency matrix $\mathbf{A}\in {\mathbb{R}}^{n\times n}$ and feature matrix (attribute vector) $\mathbf{X}\in {\mathbb{R}}^{n\times n}$
• Graph embedding: Entire graph represented as a vector.

Properties

• Permutation invariance (Graph embedding), $f(PA{P}^T,PX)=f(A,X)$: Permuting the node embeddings and attribute vector has no effect on the output. When we want to predict anything about the entire graph, for example classifying a molecule, then we want the model to permutation invariant.
• Permutation equivariance (Node embedding), $f(PA{P}^T,PX)=Pf(A,X)$: Permuting the node embeddings and attribute vector is equivalent to first applying the function, and then permuting. Or in other words, translation of input features result in an equivalent translation of outputs, or permuting the adjancency matrix means the output of f is permuted in a consistent way.
• We need to satisfy either of the two: i.e. invariance or equivariance.

Why do we not represent a graph as an MLP? Why do we not represent as a CNN? what are the similarities?

• Similarity: Locality, Weight Sharing, arbitrary input size
• Difference: Abstract shape

Message Passing

Take a graph: $\mathcal{G}=(V,E)$ along with a set of node features $\mathbf{X}\in {\mathbb{R}}^{d\times |V|}$, and generate node embeddings $z_u,\forall u\in V$.

what can node features be?

Depends on the problem we’re solving. For molecular graphs, this can be information for each atom in the molecule, for social graphs, information can be about each individual member in the graph. For no individual node features, input can still be statistics of the node in the graph, and can even contain some more information about the graph itself. To break permutation equivariance, we can assign positional encoding to each node, say as a one-hot vector encoding.

Aggregate: in each round k, each node aggregates the message (feature description) from its neighbours and update the weights. $m_{\mathcal{N}(v)}^{(k)}={\text{Agg}}^{(k)}\left(\{h_u^{(k-1)}:u\in \mathcal{N}(v)\}\right)$

• Initial embedding is set to be features of the node: $h_u^{(0)}={\mathbf{x}}_u,\forall u\in V$
• Differentiable, multiset function. Input being a set designs the graph as Permutation equivariant.
• Sum, Mean, Max/Min.
• $m_{\mathcal{N}(v)}^{(k)}={\text{MLP}}_2\left({\sum }_{u\in \mathcal{N}(v)}{\text{MLP}}_1(h_u,h_v)\right)$: Universal approximation of multiset functions.
• Receptive field of the graph increases with each iteration, as more information further away from the node is aggregated.
• Messages from the neighbours can encode structural information like degree of the neighbour node, useful in problems like analysing molecular graphs. Or can also encode feature-based information from local neighbourhood of the graph analogous to how CNNs aggregate feature information from spatially-defined patches.

Update: $h_v^{(k)}={\text{Update}}^{(k)}\left(h_v^{(k-1)},m_{\mathcal{N}(v)}^{(k)}\right)$.

• $h_v^{(k)}=\sigma \left(W_{\text{self}}{h}_v^{(k-1)}+W_{\text{neigh}}{m}_{\mathcal{N}(v)}^{(k)}+b^{(k)}\right)$, where $W_{\text{self}},W_{\text{neigh}}\in {\mathbb{R}}^{d^{(k)}\times d^{(k-1)}}$ and $b^{(k)}\in {\mathbb{R}}^{d^{(k)}}$

Readout: $h_{\mathcal{G}}=\text{READOUT}(\{h_v^{(K)}:v\in \mathcal{G}\})$: Outputs a final result after K final iteration. Like pooling in CNNs.

We can also define a graph-level equation for aggregate and update, and we can even batch aggregate and update in one equation using self-loops:

$$\begin{array}{c}{\mathbf{H}}^{(k)}=\sigma \left(\mathbf{A}{\mathbf{H}}^{(k-1)}{\mathbf{W}}_{\text{neigh}}^{(k)}+{\mathbf{H}}^{(k-1)}{\mathbf{W}}_{\text{self}}^{(k)}+{\mathbf{b}}^{(k)}\right)\\ {\mathbf{H}}^{(k)}=\sigma \left((\mathbf{A}+\mathbf{I}){\mathbf{H}}^{(k-1)}{\mathbf{W}}^{(k)}\right)\end{array}$$

Similarity to MLP with each node’s weight being a vector instead of a scalar, and aggregate and update together forms the linear and pointwise layer. So actually, a GNN can learn much more than a simple MLP.

Every iteration weight of single node gets updated using it's neighbour, but that d-dimensional vector will be saturated after many iterations? Will a single node embedding contain any meaningful information from its neighbourhood?

Over-smoothing: Representation of all the nodes in the graph can become very similar to one another. This makes it impossible to build deeper GNN models.

Formally, Define influence of a node’s input feature $h_u^{(0)}={\mathbf{x}}_u$ on the final layer embedding of all other nodes in the graph $h_v^{(K)},\forall v\in V$. For any pair of nodes u, v in the graph, influence of u on v is quantified using the Jacobian $I_K(u,v)=1^T\left(\frac{\partial h_v^{(K)}}{\partial h_u^{(0)}}\right)1$.

Building deeper models can hurt performance of GNN models as with each added layer, information loss about local neighbourhood increases and learned embeddings are over-smoothed.

More on the influence of self-update and deeper models can be found in GRL book and Xu et. al..

Generalisations

Aggregate

• Normalisation: Mean normalisation $m_{\mathcal{N}(u)}=\frac{{\sum }_{v\in \mathcal{N}(u)}{h}_v}{|\mathcal{N}(u)|}$ or symmetric normalisation $m_{\mathcal{N}(u)}=\frac{{\sum }_{v\in \mathcal{N}(u)}{h}_v}{\sqrt{|\mathcal{N}(u)\mathcal{N}(v)|}}$.

  • Why does symmetric normalisation work better than mean?

  • Normalisation leads to loss of structural information, as is provably, less powerful than sum aggregation. Normalisation is more useful when feature information is more useful than structural information.
• Set pooling: Use a universal set function approximator that can approximate any permutation-invariant aggregator function: $m_{\mathcal{N}(v)}^{(k)}={\text{MLP}}_2\left({\sum }_{u\in \mathcal{N}(v)}{\text{MLP}}_1(h_u,h_v)\right)$
• Janossy Pooling: Permutation sensitive function averaged over permutations
• Attention: $m_{\mathcal{N}(v)}={\sum }_{u\in \mathcal{N}(v)}{\alpha }_{vu}{h}_u$: weighted aggregation. Is used to increase the inductive bias of the model with prior information about importance of neighbours.
• Multiple attention heads: Compute K distinct attention weights using independent parametrised attention layers. Aggregate all message by projection and concatenation.
  • $m_{\mathcal{N}(u)}=[a_1\oplus a_2\oplus \cdots \oplus a_K]$
  • $a_k=W_k{\sum }_{v\in \mathcal{N}(u)}{\alpha }_{v,u,k}{h}_v$

Update

• Skip connections: Counter over-smoothing by directly preserving information from previous rounds of message passing.
  • ${\text{Update}}_{\text{concat}}(h_u,m_{\mathcal{N}(u)})=[{\text{Update}}_{\text{base}}(h_u,m_{\mathcal{N}(u)})\oplus h_u]$
  • ${\text{Update}}_{\text{interpolate}}(h_u,m_{\mathcal{N}(u)})=[{\alpha }_1{\text{Update}}_{\text{base}}(h_u,m_{\mathcal{N}(u)})\oplus {\alpha }_2{h}_u]$, where ${\alpha }_1,{\alpha }_2\in [0,1{]}^d$ and ${\alpha }_2=1-{\alpha }_1$, and ${\alpha }_1$ can be learned jointly with other representations.
  • Due to the analogous properties of CNNs, concatenation and skip connection as described in He et. al produces similar results.

Features and Relationships

• Edge attributes: $m_{\mathcal{N}(v)}={\sum }_{u\in \mathcal{N}(v)}{\text{MLP}}^{(k)}\left(h_u^{(k-1)},h_v^{(k-1)},{\mathbf{w}}_{uv}\right)$
• Multi-relational: aggregation can depend on the relationship between nodes.

Generalised Message Passing

$$\begin{array}{rl}{h}_{(u,v)}^{(k)}& ={\text{Update}}_{\text{edge}}\left(h_{(u,v)}^{(k-1)},h_u^{(k-1)},h_v^{(k-1)},h_{\mathcal{G}}^{(k-1)}\right)\\ m_{\mathcal{N}(u)}& ={\text{Aggregate}}_{\text{node}}\left(\{h_{(u,v)}^{(k)}\forall v\in \mathcal{N}(u)\}\right)\\ h_u^{(k)}& ={\text{Update}}_{\text{node}}\left(h_u^{(k-1)},m_{\mathcal{N}(u)},h_{\mathcal{G}}^{(k-1)}\right)\\ h_{\mathcal{G}}^{(k)}& ={\text{Update}}_{\text{graph}}\left(h_{\mathcal{G}}^{(k-1)},\{h_u^{(k)}\forall u\in V\},\{h_{(u,v)}^{(k)}\forall (u,v)\in E\}\right)\end{array}$$

Main improvement over baseline message passing is that during each iteration, the model generates a hidden edge embedding for all edges in the graph, and an overall graph embedding corresponding to the entire graph. This helps differentiate between edge and node level features and entire graph-level features. We can also define different loss functions for different type of embeddings, and tasks.

Approximation Theory

Graph Isomorphisms: Given two graphs ${\mathcal{G}}_1,{\mathcal{G}}_2$, declare whether two graphs are isomorphic. Formally, we say two graphs with adjacency matrix $A_1,A_2$ and feature matrix $X_1,X_2$ are isomorphic if and only if there exists a permutation matrix P such that $P{A}_1{P}^T=A_2$ and $P{X}_1=X_2$. Or informally, when they have same structure but differ in ordering of nodes in their adjacency matrices.

GL test and HOW POWERFUL ARE GRAPH NEURAL NETWORKS?

• Weisfeiler-Lehman isomorphism test
• Distinguishing capacity of GNNs

Problems I’m Seeing:

• A d-dimensional weight vector exists for each node in the graph, so the size scales with $\mathcal{O}(Nd)$. And number of edges will also mean that update function will be hard to compute, but is actually emabarrasingly parallel.
• How to choose depth K?
• How do you update the graph structure as you process more? Can we prune or connect more edges and nodes?

More Readings:

• GRL Book
• geometric-gnn-dojo/geometric_gnn_101.ipynb at main · chaitjo/geometric-gnn-dojo · GitHub
• A Gentle Introduction to Graph Neural Networks