Neural Relational Inference
for Interacting Systems
Petar Veličković
Artificial Intelligence Group
Department of Computer Science and Technology, University of Cambridge, UK
MILA Graph Representation Reading Group Meeting

25 July 2018

Introduction

I In this talk, I will survey the recently published Neural

Relational Inference model (Kipf, Fetaya et al., ICML 2018).
I This model enables the discovery and exploitation oflatent interactions between objects, through the synergy ofgraph convolutional networks and variational autoencoders.
I Exciting results + avenues for further work!

Graphs are everywhere!

. . . but can we always see them?

Relational inference

I Virtually all graph convolutional techniques require a graph tobe provided as input!
I However, often we will only have access to node features. . .
I Approaches such as Relational Networks (Santoro et al.,2017), or VAIN (Hoshen, 2017) circumvent this by assuming acomplete graph (i.e. all-pairs interactions).
I But most interaction graphs have properties (such as sparsity)that we may wish to explicitly demand!
I Furthermore, we may wish to identify and decouple different

types of interaction.

Our task for today

Motivation: predicting trajectories
I Input: Trajectories (e.g. coordinates) ~xi≤t for each particle i.
I Output: Future trajectories ~xi>t for each particle i.

I The interaction graph between particles will be a byproduct!

Simple baseline #1: RNN
I Let ~
x t denote the coordinates of all particles at time t:

~x t = ~x1t , ~x2t , . . . , ~xnt
I We can now define a recurrent neural network (e.g. LSTM orGRU) to operate on this sequential input:~ht = RNN(~ht−1 , ~x t )I From its hidden states, we can predict the future timesteps:~x t+1 = f (~ht+1 )where f is an MLP.

Graph convolutional network

Graph convolutional network

In a nutshell, obtain higher-level representations of a node i byleveraging its neighbourhood, Ni !
~h`+1 = g ` (~h` , ~h` , ~h` , . . . ) (a, b, c, · · · ∈ Ni )a b ciwhere g ` is the `-th graph convolutional layer.

The MPNN framework
I The NRI model leverages a graph convolutional layer inspiredby message-passing neural networks (Gilmer et al., 2017).

The MPNN framework
I The NRI model leverages a graph convolutional layer inspiredby message-passing neural networks (Gilmer et al., 2017).
` , for each edge i → j in theI First, compute edge messages, ~hi→jgraph. Apply a simple MLP, fe` , over the features of i and j:~h` = f ` (~h` , ~h` )e ii→jj

The MPNN framework
I The NRI model leverages a graph convolutional layer inspiredby message-passing neural networks (Gilmer et al., 2017).
` , for each edge i → j in theI First, compute edge messages, ~hi→jgraph. Apply a simple MLP, fe` , over the features of i and j:~h` = f ` (~h` , ~h` )e ii→jj
I Then, aggregate all messages entering a node j to obtain thenext-level features, ~hj`+1 . Apply a simple MLP, fv` , over thesummed messages.
~h`+1 = f ` v
jX

j∈Ni~h` i→j

MPNN: initial setup

~h`
5~h`
3~h`
1~h`4~h`2~h`6MPNN, computing messages~h`5~h`3~h`1~h`3→4~h`4~h`2fe`~h`6~h`3~h`4

MPNN, aggregating messages~h`5~h`3~h`3→4~h`1~h`3→4~h`5→4fe`P~h`2~h`2→4~h`6→4~h`6~h`3~h`4

MPNN, computing node features~h`5~h`3~h`3→4~h`1P~h`2~h`2→4~h`3→4~h`5→4fv`~h`+14fe`~h`6→4~h`6~h`3~h`4

MPNN, next-level features~h`5~h`3~h`3→4~h`+14~h`1~h`2~h`2→4~h`3→4~h`5→4fv`~h`+14fe`~h`6→4~h`6~h`3~h`4

Simple baseline #2: Complete graph

I As another baseline approach, we may use this kind of layer topredict trajectories using a complete graph (assume all pairs ofnodes interact).
I The equations of the baseline become equivalent to:~ht = fe (~x t , ~x t )
i→jij


X
~ht ~xjt+1 = fv i→ji6=jwith a few kinks, specific to the trajectory predicting task. . .

Simple baseline #2: Complete graph

I First, to simplify the job of the network, have it only predictchanges in position:~ht = fe (~x t , ~x t )i→jij~xjt+1 = ~xjt + fv 
Xi6=j~ht i→j

Simple baseline #2: Complete graph

I Also, explicitly model uncertainty; will be useful for thevariational framework later on.

~ht = fe (~x t , ~x t )
i→ji
j

X~ht 
µ~ t+1 = ~xjt + fv i→j
ji6=j~xjt+1 ∼ N (~µt+1, σ 2 I)jSimple baseline #2: Complete graph, with GRU
I The model thus far assumed the Markov property (i.e. that ~x t+1
depends fully on ~x t ). This is OK for physics, but if necessary,

we can alleviate the constraint by using a recurrent update:~ht = fe (~x t , ~x t )i→jij~ht+1 = GRU ~x t ,jj


X
i6=j



t
~ht+1~µ~ t+1=x+fv
jj
j~xjt+1 ∼ N (~µjt+1 , σ 2 I)

~ht  , ~ht i→jj

Interaction graphI This baseline can be improved if we specify an explicit

interaction graph. Initially, assume there are K edge types (withone type reserved for “no edge”).
I Then, define a binary tensor z ∈ RV ×V ×K such that zijk denotes

whether the edge i → j is of the k -th type.
I Assume an edge cannot have more than one type, i.e., ~zij isone-hot.

Leveraging the interaction graphI Now this graph can be exploited—define a separate MLP fek foreach edge type. For the Markov decoder:~ht = fe (~x t , ~x t )i→jij

X~ht µ~ jt+1 = ~xjt + fv i→ji6=j~xjt+1 ∼ N (~µt+1, σ 2 I)
j

The NRI decoder

I Now this graph can be exploited—define a separate MLP fek foreach edge type. For the Markov decoder:~ht =i→jXzijk fek (~xit , ~xjt )kX~ht 
µ~ jt+1 = ~xjt + fv i→ji6=j~xjt+1 ∼ N (~µt+1, σ 2 I)
j

The NRI decoder~x5t~x3t~x1t~x4t~x2t~x6t

The NRI decoder, computing messages. . .

~x5t~x3t~ht
3→4~x1t~x4t~x2t~ht
6→4~x6t

The NRI decoder, computing messages. . .

~x5t~x3t~ht
3→4~x1t~ht
5→4~x4t~x2t~ht
2→4~ht6→4~x6t

The NRI decoder, computing messages. . .

~x5t~x3t~h6
3→4~x1t~ht
1→4~x2t~ht
5→4µ~ t+1
4~ht
2→4~ht6→4~x6t

Latent graph inference

I We are still tasked with discovering the entries of the tensor z.

I Idea: Use MPNNs over a complete graph once more—thenclassify edge types based on the edge messages ~hi→j .
I This time, stack two layers—so that edges can be derived

based on global interactions!
~h1 will only depend on ~xi and ~xj ;i→jI ~
h2 will depend on all the nodes in the graph.
Ii→jThe NRI encoder

I In equation form:~h1 = f (~xj )j~h1 = f 1 (~h1 , ~h1 )
ei→ji
j

X
~h1 
~h2 = f 1 v
i→j
j
i6=j

~h2 = f 2 (~h2 , ~h2 )ei
ji→j2
zij ∼ Categorical(softmax(~hi→j))where f is an embedding, and fe1 , fv1 and fe2 are MLPs.

The variational setupI The encoder gives us the probability distribution q(z|~x ), andthe decoder gives us the probability distribution p(~x |z).
I Combine learning the two in a VAE-style framework bymaximising the evidence lower bound (ELBO):
L = Ez∼q(z|~x ) [log p(~x |z)] − DKL (q(z|~x )kp(z)){z}|{z} |
Reconstruction accuracyRegularisation

I The prior p(z) can encode desirable properties of the latentgraph. Sparsity is enforced by setting the probability of “noedge” to be higher than the other types.

Backpropagating through the sampling

I The operation of selecting zij is a discrete decision—therefore,we cannot directly propagate gradients through it.
I Can use the Gumbel softmax trick to circumvent this:2~zij = softmax((~hi→j+ ~g )/τ )where gk ∼ Gumbel(0, 1) and τ is a temperature parameter(converges to one-hot when τ → 0).
I This is a continuous approximation to the discrete

distribution—and gradients can be propagated through it.

Avoiding degenerate decoders

I Optimising the ELBO directly would involve only single-step
predictions (predicting ~x t+1 from ~x t ). This can often be nicely

approximated by ignoring relational structure altogether!

I To enforce robust decoders, predict many steps at once! EveryM steps, feed back the ground-truth input.

Avoiding degenerate decoders, cont’dµ~ 2j = decode(~xj1 )µ~ 3j = decode(~µ2j )µ~ 4j = decode(~µ3j )..
.
µ~ jM+1 = decode(~µMj )µ~ M+2= decode(~xjM+1 )jµ~ M+3= decode(~µjM+2 )j..
.

Putting it all together

For a given training trajectory, ~x , of length T :1. Compute q(z|~x ) using the encoder.
2. Sample ~zij from q(z|~x ) using the Gumbel softmax trick.
3. Execute the decoder to obtain µ~ t for t ∈ {2, 3, . . . , T }.
4. Compute the reconstruction error (of µ~ t against ~x t ) andKL-divergence (of q(z|~x ) against the prior p(z)).
5. Optimise the ELBO using gradient descent.

The NRI architecture

Legend:

xe →v

v →e…

: Node emb.

: Edge emb.

v →ext…
ΣEncoder: MLP…
q φ(z|x): Concrete distribution

v →ee →v…
ΣDecoder: SamplingΔxt

Physics simulations: latent graph discovery

Physics simulations: trajectory prediction

It might seem as if the LSTM outperforms the NRI on Kuramoto!
Qualitative analysis may show otherwise. . .

Physics simulations: qualitative results

Physics simulations: qualitative results

Motion capture

The graph is now dynamic! Re-evaluate at every decoding step.

Motion capture: trajectory prediction

truthprediction

Motion capture: qualitative results

t=0t = 10t = 20t = 30

truthprediction

Motion capture: qualitative results

t=0t = 10t = 20t = 30

Concluding remarks

I The NRI is an extremely versatile model for inferring latentinteraction graphs from pointwise trajectories.
I Latent graph discovery is still in its early phases ofdevelopment—plentiful improvements possible!
I Limitation: does not scale to large graphs! O(V 2 ) memory

requirements, and computing edge messages makes
subsampling cumbersome.
I Should not be required—most real-world graphs are sparse!

But techniques we have thus far need to start with completegraph, and gradually discover sparsity. . .

Thank you!

Questions?
petar.velickovic@cst.cam.ac.uk
http://www.cst.cam.ac.uk/∼pv273/