Neural Relational Inference for interacting systems
byKipf, Thomas, et al., ICML 2018

NRI – Backgrounds• Various kinds of time series data are governed by an underlying
network.

Lynn, Christopher W., and Danielle S. Bassett. "The physics of brain network structure,function and control." Nature Reviews Physics. 2019

Time series dataNetwork

NRI – Backgrounds• Various kinds of time series data are governed by an underlying
network.
𝑵Coupled oscillators𝜽ሶ 𝒏 𝒕 = 𝝎𝒏 + ෍ 𝑲𝒎𝒏 𝐬𝐢𝐧 𝜽𝒎 (𝒕) − 𝜽𝒏 𝒕𝒎=𝟏

Time series dataNetwork

NRI – Backgrounds• Besides examples in nature, human behaviors can also be governed by network
interactions.
• Motion of players in a basketball game• Traffic, social activities…

Time series dataNetworkCan we reversely inference the network interactions?
• Statistical methods exist, such as correlation and granger causality.
Here we present and reimplement a neural network method, Neural
Relational Inference.

NRI – Methods• NRI consists an encoder and a decoder network.
• In the implementation, NRI takes the advantage of Graph NeuralNetwork (GNN) and Variational Auto Encoder (VAE).
• We will introduce NRI step by step!

NRI – Methods
1. GNN refers to a set of neural networks that takes graph structured data asinputs.
2. GNN are neural models that captures the dependence of graphs via message
passing between the nodes of graphs (Zhou et al., 2020).

How GNN captures the dependenceof graphs?

Zhou et al., 2020

NRI – Methods• In neural network model, embedding/representation might be the key forsuccessful learning.
• In GNN, we design node embedding and edge embedding for training.

Zhou et al., 2020

NRI – Methods• In neural network model, embedding/representation might be the mostimportant key for successful learning.
• In GNN, we design node embedding and edge embedding for training.

The GNN Layer can be MLP, RNN or even CNN.

Zhou et al., 2020

NRI – Methods

Kipf et al., 2018

ℎ3𝑙 , ℎ4𝑙 are node embeddings. ℎ𝑙 4,3 is edge embedding.
𝑓𝑒𝑙 is a neural network that takes edge embeddings as inputs.
𝑓𝑣𝑙 is a neural network that takes node embeddings as inputs.

NRI – Methods• NRI consists of an encoder and a decoder. In the last stage, the edgeembeddings are used to perform edge classification.
• Although we use encoder as the example, ideas of node/edge embedding
are also in the design of decoder.

𝒛 represents edge types.
𝒒𝝓 𝐳 𝐱 is the probability ofedge types for edge 𝒛𝒊𝒋 .

Kipf et al., 2018

NRI – Methods• If our training dataset contains ground truth edge types, thenEncoder alone can inference the edge type for testing dataset.
This is true for simulation data.
• However, in some cases (e.g. online edge inference), what we ONLY
have is the time series.
That is, our training dataset does not contain ground truth edge types.

NRI – Methods• To overcome the barrier, NRI takes the advantage of VAE.
• VAE is suitable for data without labelling (unsupervisedlearning).

𝒙𝒙′

Encoder

Latent variable Decoder

The encoder compresses input 𝒙 tothe latent space.
The decoder produces 𝒙′ (from latentspace) as similar as 𝒙.

NRI – Methods• The decoder has a similar structure as the encoder.
• MLP layer• 𝑉 → 𝐸 and 𝐸 → 𝑉 operations

Encoder
Kipf et al., 2018

NRI – Methods• The decoder takes 𝑥𝑡 and edge type probabilities 𝑞𝜙 as input. Twotypes of inputs are combined.
• The decoder output Δ𝑥 𝑡 is the predicted increasement for the nexttime step.
j∗ 𝑞𝜙 (𝑧𝑖𝑗 |𝑥)

iEncoderKipf et al., 2018

NRI – Methods• To sum up, in case we ONLY have time series, we have to infer the edge(interaction) types through VAE. If the latent variable 𝑞𝜙 (𝒛|𝒙) represents edge
types correctly, the decoder should be able to produce matched increasement.

Encoder
Kipf et al., 2018

NRI – Implementations• First, let’s take a look at the encoder.

src/nri_encoder.py
scripts/train_enc.py

NRI – Implementations• let’s take a look at the decoder.

src/nri_decoder.py

scripts/train_dec.py

NRI – Implementations• Use `scripts/generate_dataset.py` to generate simulation data.
This file is provided by the paper author.
• You can use `run_encoder.py` and `run_decoder` to play with thetrained model. Use `traj_plot.ipynb` to generate prediction
trajectory.
• Let’s go to Github for details.

NRI – Results, Encoder• In the experiments, we use simulated 2D springs-coupled data intraining and testing.

Kipf et al., 2018

NRI – Results, Encoder• We train the encoder with 40 epochs, batch_size=5, each trial with 49 times
steps. Use `nn.CrossEntropyLoss` as the loss function. The hidden layer
dimension of MLP layer is 256.

Test loss: 0.20062
Test acc: 0.91605

NRI – Results, Encoder• We train the encoder with 40 epochs, batch_size=5, each trial with 49 times
steps. The hidden layer dimension of MLP layer is 256.

Any idea to better
present the result?

NRI – Results, Decoder-1
Use decoder alone, provided with edge types,
1 Step ahead prediction, then teacher force

Target trajectory

Prediction trajectory

Test error: 0.009824
Test baseline error: 0.030993

NRI – Results, Decoder-2
10 Steps prediction, then teacher force

Target trajectory

Prediction trajectory

NRI – Paper Results

NRI – Issues• Appreciate any helps for the followings issues.
• 1. Test and add support for GPU platform.
• 2. MLP building block has been implemented for the encoder and decoder.
Consider other building blocks (e.g. CNN, RNN).
• 3. For data without ground truth edge information, we need to combine theencoder and the decoder together (optional).
• 4. Visualization for encoder output result.

• Any minor adjustments are also welcomed! ☺

Reference
Kipf, Thomas, et al. "Neural relational inference for interacting systems." International
Conference on Machine Learning. PMLR, 2018.
Zhou, Jie, et al. "Graph neural networks: A review of methods and applications." AIOpen 1 (2020): 57-81.