Neural Relational Inference with Node-Specific Information 

1 INTRODUCTION . Our world includes many different types of systems that involve multiple entities interacting with each other , from biology to sports , from social media to driving situations . Modelling the behaviour of such dynamical systems is a challenging task , which requires uncovering different types of interactions among the entities and how they affect each other . Recent approaches in machine learning learn the interactions among the entities through graph-based models and attention-based models , in which the representation of an entity is updated through its relationship with other entities . Such relationships are usually either predefined or uncovered by learning . Kipf et al . ( 2018 ) proposed neural relational inference ( NRI ) for relationship uncovering in the framework of variational inference ( Kingma & Welling , 2013 ; Rezende et al. , 2014 ) . NRI has an encoder-decoder structure in which the latent codes represent different types of interactions among the entities . The distributions over the latent variables are inferred based on the input features of the entities and a graph structure is formed by sampling from these distributions . The decoder of the model is a GNN-based algorithm that runs over the uncovered graph to update the features of entities . The model shows prominent performance over many synthetic and real-world datasets . Uncovering the interaction among entities in NRI and other models in the GNN literature is studied in problems where the features of the entities are completely observable and the GNN-based algorithm is also run on the observable features . However , in many real-world problems there is a set of hidden features that affect the way entities interact with each other . For example consider the problem of predicting the future position of vehicles in a driving scenario . In order to uncover the relation of a target vehicle with other vehicles not only should we consider the features extracted from the observations from vehicles , e.g . their trajectories up to the current time , but also we should take into account the intention of the target vehicle . In fact , intention , which can be an immediate action or a longer term goal , forms a set of features that is only accessible by the target vehicle and other vehicle can not know it . In this paper we call such feature the individualized features or Node-Specific Information ( NSI ) in a graph representation of the system . Formally , NSI is a set of features that is ∗Work was done prior to joining Amazon . only accessible by one node in a graph structure but affects the interactions of that node with other nodes . Our goal is to efficiently exploit NSI to build more accurate graph structures and consequently achieve better performance in the downstream tasks . Towards this goal , the first step is to find a proper representation for NSI in our graph structure . We propose introducing a new set of nodes in the graph that carry NSI and call them private nodes . Observable features of the entities are then denoted by public nodes . Therefore each entity can be represented by a public and private node . We introduce our model in the framework of NRI , i.e . a variational inference model that uncovers different types of interaction among the entities . We carefully design the encoder and decoder part of our model to ensure that NSI remains an individual feature for one entity and does not affect the interaction modelling of the other entities . At the same time , through experiment we show that such modelling of NSI is very efficient and can result in significant improvement on the performance on downstream tasks . The main contributions of this work are the followings : • To the best of our knowledge the problem of having individualized features for each entity has not been previously studied in the framework of relational inference . We tackle this problem by introducing a new set of nodes in a graph structure that represents the individualized information . These nodes are used during the process of relational inference as well as performing the downstream task , i.e . trajectory prediction in our case . • We show that our proposed model can exploit NSI efficiently by introducing minimum additional computational complexity compared to the original NRI model and its variants . • The results of our experiments on real-world datasets show that our model can outperform the baselines and achieve the state-of-the-art results on the defined tasks . 2 RELATED WORK . Interaction modelling : Relational learning is a popular approach for the problems with dependency structure among the data points . Such dependency can be predefined in some domains . But in many domains it has to be learned . For example approaches like locally linear embedding ( LLE ) ( Roweis & Saul , 2000 ) and Isomap ( Tenenbaum et al. , 2000 ) use kNN for forming such relationships based on different measures of similarities among the data points . More recently , neural networks have become the dominant tools for learning these dependencies based on different architectures and paradigms ( Kipf & Welling , 2017 ; Hamilton et al. , 2017 ; Garcia Duran & Niepert , 2017 ; Monti et al. , 2017 ; Veličković et al. , 2017 ; Franceschi et al. , 2019 ) . The most relevant works to our proposed model are NRI ( Kipf et al. , 2018 ) and dynamic NRI ( dNRI ) ( Graber & Schwing , 2020 ) , which are discussed in more details in the next section . Recently , Li et al . ( 2020 ) introduced similar ideas for multi-modal trajectories prediction . Extensions of NRI in other directions than ours also appeared in the literature . For example , Li et al . ( 2019 ) tries to uncover interactions by imposing some structural constraints on the prior and Webb et al . ( 2019 ) introduces the idea of factorized graph for NRI . The other common approach for uncovering relationships among the entities is based on the idea of attention . This idea has been used in Narasimhan et al . ( 2018 ) ; Hoshen ( 2017 ) ; Van Steenkiste et al . ( 2017 ) ; Garcia & Bruna ( 2017 ) ; Monti et al . ( 2017 ) ; Veličković et al . ( 2017 ) , where the attention mechanism is the main tool for interaction uncovering , however , it is also used as a building block for GNNs . Future trajectory prediction as evaluation metric : We define our problem as uncovering the interactions of entities in a multi-agent dynamical system , where the evaluation is based on the accuracy of future trajectory prediction . Trajectory prediction is in fact an important problem in many multi-agent systems , including the ever growing area of autonomous driving . In fact , our approach falls into the category of multivariate time-series prediction ( Yu et al. , 2018 ; Wu et al. , 2019 ; Sen et al. , 2019 ; Salinas et al. , 2020 ; Rangapuram et al. , 2018 ; Li et al. , 2018 ; Bai et al. , 2018 ) , in which the prediction is based on the relationship among the series . Specifically we use the relationship of the entities in the GNNs framework . GNN has been widely used in the trajectory prediction , especially in the application of autonomous driving and significantly improved the performance in this area . For example , Salzmann et al . ( 2020 ) uses GNN to capture the relationship among different road users ( vehicles and pedestrian ) , Gao et al . ( 2020 ) uses graph attention networks ( GATs ) to learn the relationship among agents and different components of map data , and Liang et al . ( 2020 ) uses graph convolutional networks ( GCNs ) to learn the interaction among lanes and vehicles . In our experiment we consider scenarios in which the goal ( final ) position of the entities is given as the individualized information . In the context of goal-aware prediction , there have been some effort in the area of autonomous driving that are not based on explicit relational learning ( Rhinehart et al. , 2019 ; 2018 ) . In these papers , the prediction is based on an autoregressive flow-based model that considers a collective observation of all entities to make prediction for each of them . The goal is then added at the inference time and the latent codes are optimized in a way that the target entity reaches its goal . 3 BACKGROUND : NEURAL RELATIONAL INFERENCE ( NRI ) . NRI is an unsupervised model that learns to infer the interaction types among entities in a multi-agent systems in order to model the dynamics of the system . The model is defined as the problem of predicting the future trajectory of entities given the past trajectories . Formally , the trajectories of N entities are given for T time steps . Entity i is denoted by xi = ( x1i , x 2 i , ... , x T i ) . The set of all trajectories at time t is denoted by xt = { xt1 , xt2 , ... , xtN } and x = ( x1 , x2 , ... , xT ) denotes the whole trajectories for all agents . NRI tries to model the system by maximizing the log-likelihood of the observations , log p ( x ) , in the framework of variational inference , i.e . maximizes the evidence lower-bound ( ELBO ) : L ( θ , φ ) = Eqφ ( z|x ) [ log pθ ( x|z ) ] − KL [ qφ ( z|x ) ||p ( z ) ] , ( 1 ) where latent variable z has a categorical distribution and represents the interaction among entities . More specifically , zij is a K-dimensional vector that denotes the type of interaction between entities xi and xj . The entities in the NRI model are represented using nodes of a graph1 and therefore the interactions are directed edges on this graph . Parameters of p ( . ) and q ( . ) models are denoted by θ and φ , respectively . The three probability distributions in Eq . 1 are : • The variational posterior , qφ ( z|x ) , is implemented using amortized inference parameterized by a neural network , namely the encoder network . Given the input trajectories , the encoder network predicts the type of edges on the graph . The latent variable z is assumed to have a categorical distribution . Samples from this distribution form the edges of the graph . In order to backpropagate the error signals to the encoder layers , we need to make the sampling process differentiable . In NRI this is done by approximating the posterior distribution and reparamterezation of Gumbel distribution ( Jang et al. , 2017 ; Maddison et al. , 2017 ) : zij = softmax ( ( h2 ( i , j ) + g ) /τ ) ( 2 ) where h2 ( i , j ) is the last output of the encoder before the softmax layer and g ∈ R K shows i.i.d . samples drawn from Gumbel ( 0 , 1 ) and τ is a hyperparameter that controls the smoothness of the distribution . • The prior , p ( z ) = ∏ i 6=j p ( zij ) , is assumed to be a factorized uniform categorical distribu- tion over the edges . • The likelihood , pθ ( x|z ) , is implemented by the decoder network and predicts the future trajectories given the uncovered structure of the graph . The prediction in NRI is done in an autoregressive fashion . However , the ground truth trajectory is fed to the model for few steps during the training to improve the performance of the decoder ( teacher forcing ) . NRI in its original form has two main shortcomings : 1 . The latent variable z , which defines the edge types , is fixed for the whole prediction horizon . That is , the uncovered interactions among the agents are assumed to be fixed over the next time steps . This is not necessarily a valid assumption as the agents can dynamically change their interactions in the system . 1Throughout the paper , entities and nodes as well as interactions and edges are used interchangeably , based on the context . 2 . The prior distribution is assumed to be uniform and not conditioned on the previous observations . Both of these assumptions can degrade the performance of the model in longer prediction horizons since samples from the prior provides minimum information about the input . More recently , Graber & Schwing ( 2020 ) pointed out the above issues and addressed them in dynamic neural relation inference ( dNRI ) model . The conditional prior distribution in dNRI is defined as : p ( z|x ) : = T∏ t=1 p ( zt|x1 : t , z1 : t−1 ) , ( 3 ) which is implemented by another set of MLP and LSTM layers that form the encoder of the pθ ( . ) model . The experiments show that , by resolving those issues , dNRI achieves better prediction performance than NRI .

Summary: **Summary**: This paper introduces a neural relational inference model that makes use of the hidden features of each node in a variational inference framework. Specifically, the hidden/individual information is modeled as private node in the graph. Importantly, the task assumption made by the authors is that these individualized features cannot be observed by other entities.   **Contributions**:  1. The authors claim be to the first to study the use of individualized information for each entity in this direction. 2. The proposed approach achieve state-of-the-art results while only introducing minimum additional computational complexity.