M ESSAGE P ASSING N EURAL P ROCESSES
A BSTRACT
Neural Processes (NPs) are powerful and ﬂexible models able to incorporate uncertainty when representing stochastic processes, while maintaining a linear time complexity.
However, NPs produce a latent description by aggregating independent representations of context points and lack the ability to exploit relational information present in many datasets.
This renders NPs ineffective in settings where the stochastic process is primarily governed by neighbourhood rules, such as cellular automata (CA), and limits performance for any task where relational information remains unused.
We address this shortcoming by introducing Message
Passing Neural Processes (MPNPs), the ﬁrst class of NPs that explicitly makes use of relational structure within the model.
Our evaluation shows that MPNPs thrive at lower sampling rates, on existing benchmarks and newly-proposed CA and Cora-Branched tasks.
We further report strong generalisation over densitybased CA rule-sets and signiﬁcant gains in challenging arbitrary-labelling and few-shot learning setups.
[Figure]
on a sample from the global latent variable and makes label predictions over the target set, ˆ y T . Inference: the predicted labels are added to the target examples, differentiated from the unlabelled targets by the label τ and purple nodes.
The dataset is again passed through the encoder, h , and aggregator, a , to produce the global latent variable as conditioned on the joint target and context set, as required in the ELBO objective (Equation 5) for training.
M ESSAGE P ASSING N EURAL P ROCESSES
We present Message Passing Neural Processes (MPNPs) as the synthesis of the MP and NP models.
Figure 1 illustrates the operation of an MPNP and Appendix B gives more details including pseudocode for the entire computation and derivations.
Problem Statement.
Given a partially-labelled set of nodes with features X and neighbours given by A , sampled from f : X , A → Y , with f ∼ D , the goal is to predict labels for a subset of the unlabelled nodes .
Dataset Sampling.
In the classiﬁcation setting, the context set for a dataset (a graph ) is deﬁned as a set C = (cid:8) ( x i , y i ) (cid:9) of nodes and their one-hot labels.
The information available to the encoder h is given by the set C ∪ (cid:8)
x j | j ∈ (cid:83) i ∈ context set N ( i ) (cid:9) , with |C| = m , which contains the context set and the k -hop neighbourhoods of all context nodes.
In this way, the MPNP uses the relational structure between the context set and other nodes to produce richer representations of the context nodes.
In turn, the global latent variable z is able to encode relational structure present in the underlying stochastic process.
The target set, T = { x i | i ∈ context set }∪{ x i | i / ∈ context set } , with |T |
= n , is a superset of the context set (though not necessarily containing the entire graph), without labels .
The decoder g also uses information from the k -hop neighbourhood when predicting target labels.
E XPERIMENTS
Recently introduced benchmark suites such as OGB (Hu et al., 2020) and those proposed by Dwivedi et al.
(2020) aim to improve the quality of evaluation of graph-based models.
However, we ﬁnd they do not host tasks that match the problem statement in Section 1.
To this end, in addition to standard variants of TUD (Kersting et al., 2020) and PPI (Zeng et al., 2019) tasks, we introduce new task formulations (Appendix C, D) of the existing ShapeNet (Chang et al., 2015; Yi et al., 2016) and full
Cora (McCallum et al., 2000; Bojchevski & G ¨ unnemann, 2018) datasets and entirely novel synthetic tasks based on cellular automata (Von Neumann et al., 1951; Wolfram, 1984; Turing, 1990).
2. F IXED L ABELLING T ASKS
We ﬁrst consider tasks where the same set of classes appear in every example in the same order (Tables 1,2, Figure 2).
Inductive GNNs are designed for this setting and provide a useful baseline performance.
Figures 3 and 6 (in the Appendix) show the superior uncertainty-modelling capabilities of the MPNP.
We then introduce tasks based on modelling Cellular Automata to show how
MPNPs extend available modelling capabilities, as baselines struggle in this setting (Figure 4).
The model is provided with the states of some cells over a generation and tasked with evolving others.
[Table] [Table]
2. A RBITRARY L ABELLING T ASKS
As the total number of classes could be very large and test examples may include unseen classes, using ﬁxed-classes is infeasible.
Instead, arbitrary labellings (1 , . . .
, k )
(Garnelo et al., 2018a) are assigned on a per-dataset basis, and models are required to adapt accordingly.
The Cora-ML task is a widely used community detection benchmark.
Papers are represented by bag-of-words vectors with edges indicating that one of the papers cited the other.
Our task, CoraBranched , is derived from a more complete dataset, with 70 classes over 11 computer science
[Figure]
disciplines (McCallum et al., 2000; Bojchevski & G ¨ unnemann, 2018).
There are 10 times as many classes and the bag-of-words feature vectors are tripled in length.
Given a partially labelled subgraph of the network, the task is to label the rest (Tables 3 and 4).
In the ShapeNet mixed-category setup (Table 5)
[Figure]
we model the process that produces n -part objects (say, n = 4 for chairs with arms, legs, seats and backs , as well as airplanes with engines, bodies, tails and wings .)
We thus provide an arbitrary permutation of class labels for each example
[Figure]
The results presented show that the richer context representations and structural bias of the MPNP are generally beneﬁcial, outperforming the NP on Cora-Branched, PPISP, one TUD task and
ShapeNet mixed (excluding 3-class @ 0.1%), while producing semantically-realistic uncertainties, as shown in Figure 3 (and Figure 6 in the Appendix
) Label propagation is more successful when
[Table] [Table]
more labels are available
[Table]
but MPNP greatly improves on it at low sampling rates, showing powerful capabilities in scarce data settings.
GNNs learn better when the generative process has little functional variation
[Table]
but perform poorly in the opposite case (mixed-class and few-shot), being entirely unsuitable in the arbitrary labelling setting.
The TUD biochemical datasets are the only ﬁxed-class setting where GNNs do consistently better than MPNPs, though we can attribute this to the lack of functional variation of the generative process in these narrow tasks.
On ShapeNet and PPISP ﬁxed-class tasks, MPNP surpasses the GNN in most cases.
[Table]
A B ACKGROUND AND R ELATED W ORK
We begin by reviewing the theoretical foundations of our building blocks (Neural Processes, Message Passing architectures) and related works.
The next section presents MPNPs as a combination of these ideas that operates on datasets with relational structure generated by stochastic processes
. A. N EURAL P ROCESSES
Problem Statement.
Given a set of points with features X , partially labelled by a function f :
X → Y sampled from a distribution over functions, D , the goal is to predict labels for a subset of the unlabelled points.
A Neural Process (NP) (Garnelo et al., 2018b) learns to represent a stochastic process with an underlying distribution D .
To achieve this, the NP is trained on a set of functions f : X → Y sampled from D and tested on a disjoint set.
For each function, f i , a dataset contains tuples ( x j , y j ) , where y j = f i ( x j ) .
Their joint probability distribution can be written as p ( y 1: n | x 1: n ) =
(cid:82) p ( f i ) p ( y 1: n | f i , x 1: n ) df i .
Assuming observation noise Y j ∼ N ( f i ( x j ) , σ 2 ) and a neural network γ modelling the stochastic process instance
f i (that is, γ ( x, z ) =
f i ( x ) , where z is a random vector that mimics the randomness of f i ), we obtain the generative model:
[Equation]
where p ( z ) is a multivariate normal distribution.
Learning the non-linear function γ requires amortised variational inference on the evidence lower bound (ELBO), using a neural-networkparameterised posterior q ( z | x 1: n , y 1: n ) .
Model generation starts with the NP receiving a set of m context points C i
= { ( x j , y j ) } mj =1 sampled from f i .
The model then predicts the values y j
= f
i ( x j ) for n target points T i = { x j } nj =1 ; namely, the m original context points and ( m − n ) previously unseen target points.
To match this setup, we isolate the context set x 1: m , y 1: m from the target set x m +1: n , y m +1: n in equation 1.
The ﬁnal ELBO is:
[Equation]
Framework Details.
We note the differences in NP processing with respect to typical machine learning setups.
An NP is trained over multiple datasets, or sets of samples S i from functions
f i ∼ D , with a given training episode consisting of samples from a single such function.
Sampling over the distribution of functions provides information about the variability of the stochastic process being modelled to the NP.
The context set C i described above is a (labelled) subset of S i , while the target set T i is an (unlabelled) superset of C i , with C i
⊆ T
i ⊆
S
i .
Section 1 explains how our model uses the context set to compute target set predictions, while expanding on each of the main components of the framework—encoder, aggregator, decoder—that are common across NP models.
A. M ESSAGE P ASSING AND G RAPH N EURAL N ETWORKS
Neural networks that operate on graph-structured data process node features X ∈ R n × d with relational information in the form of an adjacency matrix A ∈ { 0 , 1 } n × n .
The aim is to produce embeddings that are useful for downstream tasks such as node or graph classiﬁcation.
Graph neural networks typically use generalised convolutional layers to learn these embeddings.
We describe their operation via the universal Message Passing (MP) paradigm; the next section presents the speciﬁc
MP instance that our models use.
Let h ti ∈ R d (cid:48) be the features of the i -th node after t message passing steps, where d (cid:48) is the embedding dimensionality; optionally, we may have edge features e ij ∈ R k where A ij = 1 .
A message passing layer corresponds to a single message passing step, updating the node features as follows, where F and G are learnable functions, N ( i ) =
{ j | A ij = 1 } and (cid:3) is a permutation-invariant aggregation function:
[Equation]
A. N EURAL P ROCESS M ODELS
Garnelo et al.
(2018b) formulated the Neural Process as a favourable combination of neural networks and Gaussian Processes.
Conditional Neural Processes (CNPs) (Garnelo et al., 2018a) are NP instances without a global latent variable, which implies a deterministic dependence on the context set.
Attentive NPs (Kim et al., 2019), CNAPs (Requeima et al., 2019), Convolutional CNPs (Gordon et al., 2019) and Sequential NPs (Singh et al., 2019) make modiﬁcations to reduce underﬁtting, better adapt in the multi-task setting, and apply inductive biases for translation and temporal sequences, respectively.
Louizos et al.
(2019) propose the Functional NP that learns a graph of dependencies between latent representations of the points, without placing a prior over the latent global variable, though their tasks do not contain explicit relational information.
The Graph NP (Carr & Wingate, 2019) is most closely related to the MPNP, performing edge imputation using a CNP-based model and Laplacian-derived features for the context points.
However, despite the naming similarity, Graph
NPs and MPNPs address different tasks—the former was evaluated on link prediction, which is not in the scope of our work.
Moreover, our NP-based model is more ﬂexible, handles uncertainty and learns from neighbourhoods, rather than whole-graph features, for classifying individual dataset samples (nodes), while leveraging the structure between them (edges).
A. G RAPH L EARNING UNDER U NCERTAINTY
Graph Gaussian Processes (Ng et al., 2018) were designed as an extension to GPs, where the covariance function and prior exploit the existence of features in node neighbourhoods.
Graph GPs are the only Gaussian method for node classiﬁcation, but perform slightly worse than GCNs— a type of GNNs that we use as a baseline.
Moreover, the complexity is somewhat higher:
O ( max node degree 2 ∗ N ) vs. O ( N ) for (MP)NP, where N = set of observations/context nodes.
The Relational GP (Chu et al., 2007) models pairwise undirected links between data points, thus addressing a different task.
The Graph Convolutional GP (Walker & Glocker, 2019) is a translationinvariant model that operates similarly to convolutional layers, while generalising to non-Euclidean domains.
More recently, Opolka & Li`o (2020) have also proposed a Graph Convolutional GP model for link prediction, which uses a GP for node-level predictions, another GP that builds on the ﬁrst one for edge-level predictions, and a deep GP incorporating these building blocks to produce more expressive representations.
B MPNP D ETAILS
Generation and Inference.
Starting from Equation 1, with the function γ corresponding to the neural network g in Figure 1, and letting x N ( i ) denote features corresponding to an entire neighbourhood, we state the generative model for the MPNP (Appendix B.1 contains a complete derivation):
[Equation]
The decoder function g is a composition of learnable functions (linear projections, MP steps) and non-linearities, so it is trainable with amortised variational inference.
The variational posterior q ( z | x 1: n , y 1: n ) is also parameterised by a neural network ( h in Figure 1) that is permutationinvariant, as each of the functions in h satisﬁes this property (full proof in Appendix B.3).
Optimisation can be achieved using standard methods with the following variational approximation of the ELBO objective (fully derived in Appendix B.4), where D = x 1: n ∪ (cid:83) ni
=1 x N ( i ) ∪ y 1: n :
[Equation]
[Equation]
Aggregation in Challenging Settings.
The manner in which information is stored in the global latent variable z is crucial—at test time, the (context-conditioned) sample is processed together with the new target points, so it must reﬂect the behaviour of the new stochastic process in a way that is relevant to the task.
Despite a simple mean over r i being sufﬁcient for many tasks, it is often necessary to produce a class-aware representation.
Therefore, we adopt the alternative aggregation function used by Garnelo et al.
(2018a) for few-shot learning tasks,
[Equation]
where C is the set of classes in the current context , with | C | ﬁxed, as required.
This performs concatenation ( (cid:107) ) of per-class summaries aggregated with a .
Intuitively, different classes in the context set are clearly delimited in this scheme, which is especially helpful in few-shot learning settings, where novel classes are seen during testing.
Models using this scheme have the ‘-c’ sufﬁx.
B. G ENERATIVE M ODEL
Equation 1 lets us derive the MPNP generative model, where the function γ corresponds to the neural network g in Figure 1 and x N ( i ) denotes the features corresponding to the neighbourhood of node i :
[Equation]
In this derivation, line 2 assumes that p ( y i | x i , x N ( i ) , z ) takes the form of a normal distribution, with mean and variance being functions of x i , x N ( i ) , z .
Line 3 uses the fact that, in our model, γ =
ReLU ◦ L 2 ◦ MP T ◦
ReLU ◦ L 1 .
Let us ﬁrst consider the case for T = 1 .
The function G corresponds to a linear transformation L 1 = W MP applied to each of the (target node) neighbours’ feature vectors (here, we refer to the concatenated representations x j (cid:107) z ).
This is followed by leveraging the aggregation operator (cid:74) j ∈ N ( i ) within the neighbourhood of each target node.
Finally, F consists of applying the skip-connection (linear transformation) W skip to each of the target node feature vectors, followed by the ReLU activation of the MP step and ReLU ◦ L 2 .
The only difference for T =lies in the aggregator and linear transformations within the MP step being performed twice.
It is important to note that the variance σ 2 is output by the same network γ , as each prediction has its own associated uncertainty.
B. M ODEL P SEUDOCODE
Algorithm 1 summarises the MPNP label generation process described in the Message Passing
Neural Processes section.
B. E NCODER P ERMUTATION I NVARIANCE
We show that, for initial node representations
h
i , the transformation r i =
( L 2 ◦ MP T ◦ ReLU ◦
L 1 )( h i ) produced by the encoder is permutation-invariant:
[Equation]
Proof: Assume an arbitrary set of features X ∈ R n × d and an adjacency matrix A ∈ { 0 , 1 } n × n , where n is the number of nodes in the context set and d is the feature dimensionality.
We ﬁrst show that each of the operations within the encoder is permutation-invariant:
The linear projections L 1 , L 2 are applied to each of the node vectors
X i separately, so changing the order of input nodes will result in the same order in the output:
[Equation]
The same holds for the activation functions, which are applied element-wise:
[Equation]
The message passing operation is also permutation-invariant, since the transformation
A → P T AP preserves the structure of the graph, with node neighbourhoods undergoing the transformation
N ( i )
[Equation]
MP ( XΠ )
j
[Equation]
where j (cid:48)
=
Π j ,
[Equation]
i
[Equation]
Initialise node features 1 foreach i ∈ context set do h
i ← x i
y i foreach j ∈ (cid:83) i ∈ context set N ( i ) do h
j ← x j
Encoding 5 foreach i ∈ context set do h
i ← ReLU
) foreach j ∈ (cid:83) i ∈ context set N i do h
j ← ReLU ( L 1 ( j )) 9 foreach t ∈ 1 , ..., T do foreach i ∈ context set do h ti ← MP ( h t 1 ) 12 foreach i ∈ context set do r i ← L 2 ( h Ti ) //
Aggregation 14 r ← a ( { r i (cid:107) i ∈ context set } ) //
Decoding 15 Sample z (cid:48) ∼ N ( µ ( r ) , diag[ σ ( r )]) 16 foreach i set do h[-(cid:48) 
-]
i =
i
foreach j i ∈ target set N i ) do h
20 foreach i set do h
i ← ReLU ( L 1 ( h (cid:48)
i )) foreach j ∈ (cid:83) i ∈ target set N ( i ) do h (cid:48)
j ← ReLU ( L 1 ( h (cid:48) 0 j ))
24 foreach t ∈ 1 , ..., T do foreach i ∈ target set do h (cid:48)
ti ← MP ( h (cid:48)
t − 1 ) 27 foreach i ∈ target set do r (cid:48)
i ← ReLU ( L 2 ( h (cid:48)
Ti ))
ˆ y (cid:48)
i ∼ N (cid:16) softmax ( µ ( r (cid:48) i )) , diag[ (cid:0) 0 .
1 + 0 .
9 × softplus ( σ ( r (cid:48) i ))] (cid:1)(cid:17)
Each type of operation performed within the encoder is thus permutation-invariant.
Composing permutation-invariant functions yields a function which has this property itself, so it follows that the overall transformation is permutation-invariant.
B. ELBO
We derive the variational approximation to the ELBO objective stated under Generation and Inference .
In the derivation, we assume m context nodes and n target nodes (that is, n − m additional targets).
The aim is to maximise the log-likelihood of target labels y m +1: n , given the target node features x 1: n , context node features x 1: m , context labels y 1: m and neighbourhoods of context nodes.
We denote by x N ( i ) the features corresponding to an entire neighbourhood and let
[Equation]
(cid:0)
[Equation]
| x
[Equation]
mi
[Equation]
In the order given above, the (in)equalities use the following: rewriting the log-likelihood via the posterior distribution, substituting the ﬁrst term via the generative model, introducing a variational distribution q ( z | x 1: n , (cid:83) ni
=1 x N ( i ) , y 1: n ) (in our case, the encoder h and the aggregation a ) to approximate the posterior p ( z | x 1: n , y 1: n ) , multiplying by q (cid:0)
z | x 1: n , (cid:83) ni
=1 x N ( i ) , y 1: n (cid:1) and integrating over z , the result that ∀ p, q. KL ( p (cid:107) q ) ≥ 0 , separating terms, approximating p ( z | x 1: m , (cid:83) mj
=1 x N ( j ) , y 1: m ) with q ( z | x 1: m , (cid:83) mj
=1 x N ( j ) , y 1: m ) and applying the KL deﬁnition.
C T ASK D ESCRIPTIONS
Cellular Automata For the Life-like family of cellular automata we sample ∼ % of the possible 2 18 rule sets at random (Bernoulli p = 0 . 01 ).
For each selected rule set, we generate a random state on a 30 × 30 toroidal lattice (top connects to bottom, left connects to right) and check that every possible state is present (i.e. there are live cells with each of 0, 1, 2, ..., 8 neighbours and similarly a dead cell), then step forward one generation by applying the rule set to form the input-label pair.
[Table] [Table]
For density-based rules we use birth/survival functions with either the form of the top-hat function:
R 0 ( d, k 1 , k 2 ) =
  0 for d < k 1 , 1 for k 1 ≤ d ≤ k 2 , 0 for d > k 2 .
(12) or 1 − R 0 , i.e.:
R 1 ( d, k 1 , k 2 ) =
  1 for d < k 1 , 0 for k 1 ≤ d ≤ k 2 , 1 for d > k 2 .
The irregular graphs that the density-based rules operate on are generated using Scipy and NetworkX.
In each case we sample the number of nodes uniformly from the interval [100 , 200] .
For the planar Voronoi the nodes are positioned at uniformly at random in the unit square and the tessellation is generated using SciPy.
1 For spherical-Voronoi the nodes are positioned uniformly at random over the surface of the sphere and the tessellation is generated using SciPy.
For small-world the graphs are generated using the Watts-Strogatz model with p = 0 .
1 and k = 10 i.e. the network is initialised in a ring-lattice connected to its 10 nearest-neighbours on the ring and then edges are rerouted with probability 0 .
1 , using the NetworkX implementation.
For the scale-free case we use the Barabasi-Albert model with m = 3 , using the NetworkX implementation.
Cora We base our Cora tasks on the CitationFull dataset provided in PyTorch Geometricwhich is loading the data used by Bojchevski and G ¨ unnemann 6 , who in turn base their set on that originally gathered by Andrew McCallum of University of Massachussets Amherst.
7 Nodes are research papers with bag-of-word features (8710 words meet the threshold for inclusion by Bojchevski and G¨unnemann, which was given through correspondence with the authors as a minimum of appearing in 10 documents in the set) that use presence/absence rather than counts (multi-hot).
Edges indicate that one of the papers cited the other, though we do not distinguish between citing/cited graphs graphs albert graph and the graph is undirected.
The papers belong to one of 70 topics within 11 disciplines of Computer Science, and we present the relevant class-ID information for splitting by discipline in Table 7.
In the few-shot learning setup we separate out classes { 29 , 4 , 10 , 53 , 26 , 45 , 30 , 17 , 21 , 56 , 47 } for validation and { 59 , 1 , 42 , 48 , 15 , 62 , 16 , 67 , 61 , 49 , 38 } for testing, representing 14.94% and 10.79% of the total nodes, respectively.
Each of these splits contains a class from every branch (henceclasses) with an effort made to ensure the class-to-branch ratios were also approximately 15% and 10%, with preferential selection for the test set and an allowance for producing largely connected subgraphs.
For example, on the Encryption-branch there are three classes 15, 26, and 6, containing approximately one sixth, one third and one half of the nodes, respectively, with class 15 being selected for the test set and 26 for the validation set.
Practically the connectivity allowance means selecting class 48 (12.1% of AI) rather than 22 (9.0% of AI) for the test set and class 29 (14.7% of
ML) rather than 55 (14.9% of ML) for the validation set.
TUD Datasets Proteins and Enzymes are more commonly treated as graph-classiﬁcation tasks, but there is an intermediate labelling of secondary structural elements ( α -helices, β -sheets and βturns) that can be used in the node classiﬁcation setup.
DFHR, COX2 and BZR consist of small libraries of small molecule inhibitors against each respective protein target (Dihydrofolate Reductase, Cycloxygenase-2 and the Benzodiazapene Receptor).
In the typical graph-classiﬁcation task, molecules are deemed active or inactive on the basis of a thresholded half-maximal inhibitory concentration measure determined through in vitro biochemical assays.
The node-classiﬁcation task considered here requires the model to predict node labels representing encodings of atom-type.
Node features are xyz coordinates of the conformation provided in the datasets.
Protein-Protein Interaction Site Prediction
This node-classiﬁcation task utilises protein structural data collated in (Zeng et al., 2019), representing protein structures as graphs of interacting residues.
Nodes are featurised with low dimensional embeddings of physicochemical properties (Meiler et al., 2001), encodings of secondary structure, solvent accessibility metrics, and position-speciﬁc scoring matrices which capture evolutionary information as protein-protein interaction residues have been shown to be evolutionarily conserved.
Edge features represent one-hot encodings of intramolecular interaction types.
Node labels indicate whether or not that amino acid takes part in an experimentally determined protein-protein interaction.
Graphs are constructed using
Graphein (Jamasb et al., 2020).
D E XPERIMENTAL AND M ODEL D ETAILS D. E XPERIMENTS
Recently introduced benchmark suites such as OGB (Hu et al., 2020) and those proposed by Dwivedi et al.
(2020) aim to improve the quality of evaluation of graph-based models.
However, we ﬁnd they lack tasks that are appropriate for evaluating meta-learning frameworks of the kind we present, i.e. they do not host tasks that match the problem statement in Section 1.
To this end, in addition to standard formulations of TUD (Kersting et al., 2020) and PPI (Zeng et al., 2019) tasks, we present new task formulations of the existing ShapeNet (Chang et al., 2015; Yi et al., 2016) and full Cora (McCallum et al., 2000; Bojchevski & G¨unnemann, 2018) datasets and entirely novel synthetic tasks based on cellular automata (Von Neumann et al., 1951; Wolfram, 1984; Turing, 1990).
D.1. B ASELINES AND M ODEL D ETAILS
We evaluate against a variety of baselines that collectively leverage all sources of information present in the tasks (featural, relational & contextual).
This helps highlight where the advantages of the
MPNP lie in a given setting.
Table 8 from Appendix D indicates which baselines were run on each task and why.
The label propagation algorithm (LP) (Zhu & Ghahramani, 2002) makes direct use the context labels (nodes are labelled by their neighbours, who label their neighbours, and so on) and is best suited to segmentation-like tasks.
The strength of LP as a baseline is conﬁrmed in recent work showing simple models based on label propagation are competitive with state-of-the-art GNN models (Huang et al., 2021).
Where relevant, we include guessing the most common contextlabel (Mode), as this may signiﬁcantly outperform the uniform-prior ( 1 / | C | ) for some tasks.
Graph neural networks (GNNs) use training data in the inductive setup, but are not equipped to use the context points provided at test time.
GNNs are expected to perform well on tasks with ﬁxed classes and little variation in the generative process across the set of datasets being modelled.
We note that these models are not designed to handle arbitrary labelling tasks and their expected performance is bound by chance, i.e. E [acc . ]
= 1 / | C | : as predictions do not depend on class labellings, for any given task example we can construct a set of equivalent tasks by permuting the labels, and over the set of permutations the average performance will be chance (a formal derivation is provided in Appendix F).
As such, we do not include this baseline on these tasks.
In our setup, the GNN consists of GCN layers with skip-connections (Appendix D.4 contains a detailed description of the architecture).
Non-message-passing Neural Processes (NPs) are limited only by their inability to leverage relational information between points, though this is, of course, a serious limitation in the settings we consider.
We use the same Message Passing Neural Process and NP architectures for most Cora,
ShapeNet and biochemical tasks, with the addition of Maxout layers (Goodfellow et al., 2013) for
CA tasks.
Other modiﬁcations are described with the experiment in which they are used, and full model details for each scenario are provided in Appendix D.
D.1. F IXED L ABELLING T ASKS
We ﬁrst consider tasks where the same set of classes appear in every example and the class labelling is ‘ﬁxed’.
Inductive GNNs are designed for this setting and provide a useful baseline performance.
Two tasks are adapted from the TUD collection (Kersting et al., 2020): Enzymes and DHFR.
The
Enzymes dataset consists of proteins represented as networks of secondary-structural elements ( αhelices, β -sheets, β -turns; SSEs) with biochemical features describing these units and edges between connected elements.
DHFR is a library of small molecules that inhibit a particular protein, represented as graphs of atoms connected by bonds with spatial positions as features.
Table 1 shows the MPNP narrowly outperforms the NP at the Enzymes task and by a much greater margin for
DHFR, though in each case an inductive GNN is more successful.
This suggests that the relational information present in the Enzymes dataset is of secondary importance to the featural information of the SSEs, and that there is limited variation over both datasets, given that an inductive model can perform well without any context points.
Nevertheless, it is promising that MPNPs are able to use the relational information in DHFR to improve greatly on NPs.
The Protein-Protein Interaction (PPI) Site Prediction task involves predicting which nodes (amino acids) in an amino acid residue graph are involved in an experimentally-determined PPI (Zeng et al., 2019).
Solving this task is thought to depend strongly on being able to use relational information, and there is great variation between examples.
As expected following the TUD results, the MPNP excels in this setting, with SOTA-competitive results at plausible context rates presented in Table 2.
The preﬁx ‘R’ indicates that the message-passing scheme has an edge-type dependency, as in the R-GCN (Schlichtkrull et al., 2018).
Full details for this model are provided in
Appendix D.5.
The ShapeNet repository (Chang et al., 2015; Yi et al., 2016) is a collection of large-scale 3D shapes, represented as point clouds for our applications.
We embed the points as a nearest-neighbours graph ( A ) and use the ( x, y, z ) position as node features ( X ).
There are 16 object categories, each with a ﬁxed number of parts, ranging from two to six.
The labels have consistent meaning across datasets within a category.
For example, we model the process that produces chairs with arms, legs, seats and backs , which we can consistently label { 1 , 2 , 3 , 4 } .
Part labelling results are presented in Figure 2.
We use the mean-Intersection-over-Union (mIoU) metric, which is standard for segmentation tasks: the ratio of overlap (TP) to the union (TP+FP+FN) is found for each part, and averaged (higher is better, T/F P/N = true/false positives/negatives).
In 11 object categories, the MPNP outperforms the NP at more than 95% conﬁdence across the entire context sampling range, and is the top-performing model over some of the sampling range in 13 out of 15 categories.
At 30% sampling, label propagation dominates as expected.
We do not include results for guessing-the-mode as this strategy is unsuited to the mIoU metric (performance is bound by 1 /C for C -part objects as TP C = 0 for all but one part.)
Figure 3 shows the superior uncertainty-modelling capabilities of the MPNP (with extensive visualisations in the Appendix, Figure 6).
In the ﬁrst 3 plots, we visualise the uncertainty predictions for a table sample.
Though the models achieve similar mIoU, the MPNP is only signiﬁcantly uncertain at the borders between parts (a physically relevant uncertainty), whereas the NP is uncertain along the table-top edges, which are distant from any table-leg points in the internal geometry of the table.
On the right, we present the results of an active learning experiment similar to that described by Garnelo et al.
(2018a).
At each step, the target with the greatest uncertainty is added to the context set (i.e. labelled) and predictions are repeated.
This shows the power of useful uncertainty estimates in the MPNP.
We introduce tasks based on modelling Cellular Automata .
Irregular graph-CAs have been used to study trafﬁc networks (Malecki, 2017), social networks (Hunt et al., 2011), urban development (White, 1998) and logistics (Lopez et al., 2019), and cell dynamics (Bock et al., 2009).
Our aim is to show how MPNPs extend available modelling capabilities, as existing baselines are likely to struggle in this setting.
The model is provided with the states of some cells over a generation and tasked with evolving others.
To evaluate generalisation, we prevent rule-set overlap in the train, validation and test sets.
This contrasts with the existing work of Gilpin (2018), where the model learns a single rule-set, and that of Mordvintsev et al.
(2020), who train a CA to produce a desired pattern.
We provide an overview of these tasks, with full details given in Appendix C.
Conway’s Game of Life (Games, 1970) consists of cells in a 2D lattice governed by simple rules: cells become alive/are born (B) or stay alive/survive (S) depending on the number of living neighbours.
The Life-like family of CA are the generalisations of these rules over any number of neighbours 0–8, deﬁning 2 18 variants.
Neighbour counts can also be generalised to neighbourhood population-densities, and density-based rules can be adapted to irregular graphs and non-planar topologies.
We consider single-interval rules, such that cells live or die based on being inside or outside a continuous range of population-densities, on small-world, scale-free, and spherical Voronoi networks (an example of the latter is shown in Figure 5).
State evolution results are presented in Figure 4.
Here, the ‘Population/State Mode’ baselines are versions of ‘guessing-the-mode‘ that output the most common label over the whole context set or by initial state, respectively.
The NP is often able to match the state-mode strategy, but this is the ceiling to methods that do not take relational information into account.
The MPNP is able to learn effective representations that generalise well to the disjoint test set for density-based rules.
For density-based rules, MPNPs perform strongly across a variety of graph structures, while NPs are bound by simple strategies that guess the most common state change.
Neither model is able to perform well for the Life-like family, despite the existence of a solution to this problem for MPNPs, outlined in Appendix G.
D.1. A RBITRARY L ABELLING T ASKS
Garnelo et al.
(2018a) applied the CNP model in the arbitrary labelling setting, where each dataset includes samples drawn from a ﬁxed number k of class types, where the total number of types
K (cid:29) k .
As the total number of classes could be very large and test examples may include unseen classes, using ﬁxed-classes is infeasible.
Instead, arbitrary labellings (1 , . . .
, k ) are assigned on a per-dataset basis, and models are required to adapt accordingly.
The Cora-ML task is a widely used community detection benchmark.
Papers are represented by bag-of-words vectors with edges indicating that one of the papers cited the other.
Our task, CoraBranched , is derived from a more complete dataset, with 70 classes over 11 computer science disciplines (McCallum et al., 2000; Bojchevski & G¨unnemann, 2018).
There are 10 times as many classes and the bag-of-words feature vectors are tripled in length.
Given a partially labelled subgraph of the network, the task is to label the rest.
We consider the transductive setup (Yang et al., 2016;
Veliˇckovi´c et al., 2017) where every class is observed during training and a few-shot learning task where, at test time, the models are presented with classes unobserved during training, to evaluate how well the learned representations generalise.
Results for the transductive setting are presented in Table 3.
Both models perform well in the low-sampling rate regime, indicating a strong feature signal, though the MPNP-c signiﬁcantly outperforms the NP-c in every test, by up to 10% for 7-class.
LP performs best at higher sampling rates and for more classes, as expected.
Table 4 compares the quality of NP and MPNP representations in the few-shot learning context—the MPNP generalises better to unseen categories.
In the ShapeNet mixed-category setup, we model the process that produces n -part objects (say, n = 4 for chairs with arms, legs, seats and backs , as well as airplanes with engines, bodies, tails and wings .)
Here, labels have consistent meaning only within a given realisation, so using a ﬁxed ordering of labels implies a meaningless relationship between, say, chair-backs and airplane-wings .
We thus provide an arbitrary permutation of class labels for each example.
Table 5 shows results for the mixed-class part-grouped ShapeNet task.
The GNN struggles as expected, with performance below chance.
Label propagation is the strongest performer at high sampling rates, with the MPNP-c and NP-c at a relative advantage with fewer context points.
The NP-c performs best at 0.1% on 3-class, which may be due to category imbalances (80% of 3-part objects are tables) disrupting the MPNP-c. MPNP-c and label propagation otherwise divide the sampling range as top performers.
All models were trained on a
Titan
Xp
GPU or an
RTX
GPU, with torch.manual seed(0) across all experiments.
An 80/20 train/test split was used for
TUD datasets 9 and the ones provided by PyTorch Geometric 10 for all ShapeNet tasks.
The supplementary material includes code for all models and experiments described in this paper.
D. MPNP
The architecture of the MPNP can be summarised as follows:
encoder:
Linear( h ), ReLU, { MP( h ), ReLU }× T , Linear( r ); 2.
global latent variable encoder:
Linear( r ), [Linear( z ), Linear( z )] (mean & variance of z ); 3. decoder: Linear( h ), ReLU, { MP( h ), ReLU }× T , Linear( h ), ReLU, [Linear( C ), Linear( C )] (mean & variance of ˆy ).
Across all experiments, the Adam optimiser is used to maximise the ELBO (i.e. minimise the sum of the negative log-likelihood and KL-divergence in equation 5).
D.2. TUD
On Proteins and Enzymes, the MPNP hyperparameters are h = 64 , r = 128 , z = 256 ; for the
MPNP-c, h = 64 , r = 96 , z = 288 ; both have T = 2 .
On DHFR, COX2 and BZR, both MPNP and MPNP-c have h = 64 , r = 128 , z = 256 , T = 1 .
We trained both models for 400 epochs with learning rate 7 e × 10 − 5 on all datasets except for Enzymes, where we used 700 epochs and learning
[Table] [Table]
rate 1 × 10 − 4 .
For all datasets, we sample context and (additional) target points in the 10% – 25% range.
D.2. S HAPE N ET
Across all experiments, h = 64 , r = 128 , z = 256 , T = 2 .
The MPNP was trained for 400 epochs on ﬁxed-class and 500 epochs on mixed-class tasks, with 5% – 25% context and (additional) target points and a learning rate of 7 × 10 − 5 .
D.2. C ORA
In both the transductive and few-shot settings, h = 64 , r = 64 , T = 2 and z = N × 64 for Nclasses.
In the transductive setting the model is trained for 500 epochs where little if any overﬁtting is observed.
In the few-shot setting the model is trained for 400 epochs.
The model performs signiﬁcantly better on the training classes in the few-shot case, though this is expected.
A learning rate of 7 × 10 − 5 is used in both cases and we sample context and target points in the 10% – 50% range.
D.2. CA
The CA models use a modiﬁed architecture that includes Maxout layers (Goodfellow et al., 2013) that can be summarised as follows:
encoder:
MP( h ), ReLU, { Linear( h ), ReLU }× 3 , Maxout( h, 2 ), Linear( r ), ReLU; 2.
global latent variable encoder:
Linear( r ), [Linear( z ), Linear( z )] (mean & variance of z ); 3. decoder: MP( h ), ReLU, { Linear( h ), ReLU }× 3 , Maxout( h, 2 ), (concatenation with z ), { Linear( h ), ReLU }× 3 , [Linear( C ), Linear( C )] (mean & variance of ˆy ).
Maxout layers use a pool-size of 2 and the decoder delays concatenation with z until after the
Maxout layer (and the part before concatenation matches the encoder).
For both the life-like and density-based settings, h = 64 , r = 64 , z = 128 .
The models are trained for 200 epochs with a learning rate of 1 × 10 − 4 and we sample context and target points in the 30% – 50% range.
D.2. R-MPNP
The architecture of the R-MPNP can be summarised as follows:
encoder:
Linear( h ), ReLU, { R-MP( h ), ReLU }× T , Linear( r ); 2.
global latent variable encoder:
Linear( r ), [Linear( z ), Linear( z )] (mean & variance of z ); 3. decoder: Linear( h ), ReLU, { R-MP( h ), ReLU }× T , Linear( h ), ReLU, [Linear( C ), Linear( C )] (mean & variance of ˆy ).
Across all experiments, the Adam optimiser is used to maximise the ELBO (i.e. minimise the sum of the negative log-likelihood and KL-divergence in equation 5).
D.2. PPISP
The hyperparameters used are h = 64 , r = 64 , z = 256 .
Models were trained for 1000 epochs with a learning rate of 4 × 10 − 5 .
We sample context and target points in the 10% – 50% range.
D. NP BASELINE
The architecture of the NP consists of:
encoder:
Linear( h ), ReLU, Linear( h ), ReLU, Linear( r ); 2.
global latent variable encoder:
same as for the MPNP; 3. decoder: Linear( h ), ReLU, Linear( h ), ReLU, Linear( h ), ReLU, [Linear( C ), Linear( C )]
(mean & variance of ˆy ).
The Adam optimiser is also used here to maximise the ELBO.
D.3. TUD
On Enzymes, the NP and NP-c hyperparameters are h = 64 , r = 128 , z = 512 .
On Proteins, we used h = 64 , r = 64 , z = 512 for the NP and h = 64 , r = 96 , z = 288 for the NP-c.
On DHFR,
COX2 and BZR, both NP and NP-c have h = 64 , r = 64 , z = 512 .
We trained both models forepochs with learning rate 4 e − 5 on all datasets except for Enzymes, where we used 700 epochs.
For all datasets, we sample 10% – 25% context and (additional) target points.
D.3. S HAPE N ET
The same hyperparameters were used for all tasks: h = 64 , r = 64 , z = 512 .
The NP was trained for 400 epochs on ﬁxed-class and 500 epochs on mixed-class tasks, with 5% – 25% context and (additional) target points and a learning rate of 4 e − 5 .
D.3. C ORA
In both the transductive and few-shot settings, h = 64 , r = 64 and z = N × 64 for N -classes, matching the MPNP.
The model is trained for 500 epochs in the transductive setting and 400 in the few-shot setting.
A learning rate of 7 × 10 − 5 is used in both cases and we sample context and target points in the 10% – 50% range, matching the MPNP.
D.3. CA
Changes are made to the NP architecture for the CA tasks to match the changes made to the MPNP for this task, with MP layers replaced with linear layers with 2 h units to match the parameter count of the MP.
Otherwise the parameters match that of the MPNP:
h = 64 , r = 64 , z = 128 .
The models are trained for 200 epochs with a learning rate of 1 × 10 − 4 .
D.3. PPISP
The hyperparameters used are h = 64 , r = 64 , z = 256 .
Models were trained for 1000 epochs with a learning rate of 6 × 10 − 5 .
We sample context and target points in the 10%-50% range, matching the R-MPNP.
D. GNN BASELINE
This model consists of 3 GCN 11 layers with learnable skip-connections; the operation of a layer is:
[Equation]
We use h = 64 across all tasks and train the model for 500 epochs, with the Adam optimiser minimising the cross-entropy loss and a learning rate of 1 e − 4 .
The context and target ranges are as previously described, for each dataset.
Note that this model does not make use of the context labels.
D. R-GCN BASELINE
The model consists of 3 RGCN 12 layers; the operation of a layer is:
[Equation]
We use h = 64 and train the model for 400 epochs, with the Adam optimiser minimising the crossentropy loss and a learning rate of 7 × 10 − 5 .
This model leverages edge features in the messagepassing steps but does not make use of context labels.
E N UMERICAL R ESULTS AND U NCERTAINTY P LOTS
In this section, we present the numerical results used to generate the CA and ShapeNet plots in the main text.
Tables 9, 10 and 11 show the ShapeNet single-category performances, whereas Tableprovides the Cellular Automata results.
Figure 6 illustrates additional uncertainty visualisations for other classes in ShapeNet, reinforcing the ﬁnding that the estimates produced by MPNPs are semantically relevant.
[Table] [Table]
[Table] [Table]
F I NDUCTIVE GNN S WITH A RBITRARY L ABELLING
When introducing the baselines in the Experiments section, we noted that the expected performance of inductive GNNs in the arbitrary labelling setting is no better than chance.
This is because the predictions of such a model do not depend on the labelling scheme and for any particular labelling of a task we can produce a set of equivalent tasks by permuting the labels.
First consider the two class case: outputs are either 1 or 2 and labels are either A or B , giving the mutually exclusive, collectively exhaustive groups { 1A } , { 1B } , { 2A } , { 2B } , which we normalise to sum to 1 by dividing by the number of examples.
In the case that (A , B) =
(1 , 2) , the accuracy is:
[Equation]
and if the labels are permuted:
[Equation]
which average to:
[Equation]
Generalising, outputs are in 1 , ..., N and labels in { A , ..., Ω } , for the matrix of pairs:
[Equation]
 with the sum of all these elements being 1.
There are N!
permutations of the arbitrary labelling, and therefore N! equivalent tasks.
Each term in the matrix appears in (N − 1)!
accuracy sums (with that term ﬁxed, there are N − 1 free terms with (N − 1)!
permutations), so the mean accuracy is:
[Equation]
G A N MPNP S OLUTION TO THE L IFE - LIKE F AMILY
The Life-like rules can be viewed as 18 separate rules that act in parallel, one for each neighbourhood count (9) for each state (2) and hence the 2 18 variants noted in the main
A solution can be built using the concatenation encoder where the ﬁrst steps describe the situation being observed at a given node as a one-hot encoding in an 18-element vector, and then summarises these using a max aggregator (or sum or mean with corrections later) and then concatenating by whether the cell lives or dies (i.e. concatenate by class).
The max aggregation gives all the observed conditions that lead to a cell being alive in the next generation, and all those that lead to a cell being dead 13 .
This representation is then used without modiﬁcation as the latent variable.
The decoder ﬁrst extracts the observation to the format used by the encoder and then compares it with the latent variable, if it matches a condition found in the living-half of the latent variable, then the cell is alive in the next generation, if it matches a condition in the dead-half then the cell dies or stays dead.
The non-obvious parts are producing a one-hot encoding from a scalar (the neighbourhood count is produced simply by the MP) and checking the decoder observation against the latent variable.
Onehot encodings of length N can be produced using Maxout layers as follows.
First using a 2-pool
Maxout as:
[Equation]
setting W 1 = − 1 and W 2 = 1 , b 1 = ( − 1 , 0 , ..., N − 2) and b 2 = ( − 1 , − 2 , ..., − N ) .
The elements of this function take the form of the max of ( − x + j − 1) and ( x − j − 1) which is a ‘v’ with unit slopes centred at j with a minimum value of − 1 .
If we follow the Maxout with a linear layer ( − I ) and a ReLU activation, we can ﬁrst ﬂip the ‘v’ and then ﬂatten the edges to give a triangular hat centred at j with height 1.
Thus, if x = 2 the ﬁrst element (zeroth) is 0, the second element is 0, the third element is 1 and the rest are 0s.
In this way, the ﬁrst parts of the encoder and decoder can accurately represent the observed states.
To compare an observation against the latent variable, we can take the sum of the observation and latent and subtract 1 s (i.e. an AND).