Basics about Graph Networks

Bryan
5 min readJun 4, 2024

--

Graph

A graph data structure consists of a set of nodes and a set of edges, where edges form connections between nodes. If graph is undirected, each edge is an unordered pair of two nodes. If graph is directed, each edge is an ordered pair of nodes.

There are several ways to represent graph: list of edges, adjacency matrix, adjacency list.

Using the above chart as example, its list of edges is:

[(0,1),(0,2),(1,3),(2,3)]

As adjacency matrix:

As adjacency list:

[{1,2},{3},{3},{}]

Jraph is a lightweight library for working with graph neural networks in jax. It provides a data structure for graph. In jraph, we can define the above graph as:

nodes = jnp.array([[[0.0, 0.0], [0.0, 0.0]], 
[[0.1, 0.1], [0.1, 0.1]],
[[0.2, 0.2], [0.2, 0.2]],
[[0.3, 0.3], [0.3, 0.3]]])
senders = jnp.array([0, 0, 1, 2])
receivers = jnp.array([1, 2, 3, 3])
edges = jnp.array([[0.5], [0.6], [0.7], [0.8]])
n_node = jnp.array([4])
n_edge = jnp.array([4])
global_context = jnp.array([[1]])
graph = jraph.GraphsTuple(
nodes=nodes,
edges=edges,
senders=senders,
receivers=receivers,
n_node=n_node,
n_edge=n_edge,
globals=global_context
)

Graph Networks

Overview

Graph Networks (GN) are a type of model designed to handle data structured as graphs. It is powerful for tasks where the relationships between entities are crucial, such as social network analysis, molecular modeling, and recommendation systems.

In general, people use GN to solve three types of problems on graphs:

  1. Node Classification: E.g. what is the topic of a paper given a citation network of papers?
  2. Link Prediction / Edge Classification: E.g. are two people in a social network friends?
  3. Graph Classification: E.g. is this protein molecule (represented as a graph) likely going to be effective?

Here are more examples:

Figure 2 (Battaglia, 2018, p8)

Note GN are not necessarily Graph Neural Networks (GNN). E.g, if the update functions use Neural Network layers, it categorizes as GNN. Otherwise, if the update functions use other types of functions, it’s not GNN. GN generalize various existing GNN models by providing a flexible, modular approach.

Fundamental Concepts

In GN, a graph usually has 3 components: nodes/vertices (V), edges/links (E), global context (u).

In a social network, nodes could represent individuals, edges could represent friendships or connections; in a transportation network, nodes could represent intersections or stations, edges could represent roads or routes between them.

Context refers to additional information or metadata about the graph as a whole. It can include global attributes that apply to the entire graph, such as the overall type of network (e.g., social, transportation) or global statistics (e.g., total number of nodes and edges, overall density).

In GN, aggregation functions play a crucial role in combining information from neighboring nodes. Examples of aggregation functions are sum, max, avg, etc.

Update functions are used to update the state of a node based on the aggregated information from its neighbors. They each take a set as input, and reduce it to a single element which represents the aggregated information. Typical update functions are MLP, GRU/LSTM, Concatenation, etc.

We will provide with more details about aggregation and update functions in the next section.

Message Passing GN

In general there are three mechanism to process data in GN: convolutional, attentional, message passing. Among them message passing is the most generic, and as a result can simulate the most complicated data structure. In message passing, senders and receivers work together to compute arbitrary vector based message, and then send across receiver nodes.

Veličković, 2021, slide 47

Let’s look at how a graph is updated in a message passing GN block:

1. Update-edge function (φe) is applied per edge, with information from edges, receiver nodes, sender nodes, context.

2. Aggregate-edge-for-nodes function (ρe→v) is applied to updated edges from step 1. The result will be used in the next step’s node update.

3. Update-node function (φv) is applied to each node, to compute an updated node attribute.

4. Aggregate-node-for-globals function (ρv→u) is applied to updated nodes from step 3. The result will then be used in the next step’s global update.

5. Update-global function (φu) is applied once per graph, and computes an update for the global attribute. (Battaglia, 2018, p12, 13)

Note GN doesn’t have to implement all of the steps listed above. E.g, if the GN only makes prediction on the nodes, it doesn’t have to compute global context.

Characteristics of Graph Networks

1.Graphs can express arbitrary relationships among entities, which means the GN’s input determines how representations interact and are isolated, rather than those choices being determined by the fixed architecture. For example, the assumption that two entities have a relationship — and thus should interact — is expressed by an edge between the entities’ corresponding nodes.

2. A single GN can operate on graphs of different sizes (numbers of edges and nodes) and shapes (edge connectivity).

3. The structure and functions within a GN block can be configured in different ways, which offers flexibility in what information is made available as inputs to its functions, as well as how output edge, node, and global updates are produced. E.g, update functions often uses MLP, meanwhile CNN or RNN are also suitable for different scenario. Aggregation functions usually use elementwise summation, but averages and max/min could also be used. (Battaglia, 2018, p 13, 15)

Summary

This article briefly introduces graph data structure and Graph Networks, particularly how information are updated across nodes and edges in message passing Graph Networks. If you want to see a cool real-world application of message passing GNN, read this article Dissecting DM Graphcast.

--

--