Relational State-Space Model for Stochastic Multi-Object Systems
Authors: Fan Yang
Ling Chen
Fan Zhou
Yusong Gao
Wei Cao
Venue:  (2020)
Keywords: state-space model
time series
deep sequential model
graph neural network

Abstract: Real-world dynamical systems often consist of multiple stochastic subsystems that interact with each other. Modeling and forecasting the behavior of such dynamics are generally not easy, due to the inherent hardness in understanding the complicated interactions and evolutions of their constituents. This paper introduces the relational state-space model (R-SSM), a sequential hierarchical latent variable model that makes use of graph neural networks (GNNs) to simulate the joint state transitions of multiple correlated objects. By letting GNNs cooperate with SSM, R-SSM provides a flexible way to incorporate relational information into the modeling of multi-object dynamics. We further suggest augmenting the model with normalizing flows instantiated for vertex-indexed random variables and propose two auxiliary contrastive objectives to facilitate the learning. The utility of R-SSM is empirically evaluated on synthetic and real time series datasets.

Published as a conference paper at ICLR 2020
RELATIONAL STATE-SPACE MODEL FOR STOCHASTIC MULTI-OBJECT SYSTEMS
Fan Yang', Ling Chen*i, Fan Zhou', Yusong Gao Wei Cao
CColleee of Computer Science and Technology, Zhejiang University, Hangzhou, China $ AAlibaba Group, Hangzhou, China
[fanyang01, lingchen, fanzhou) J jz edu. j sianchuan. gys, mingsong. cw) ) @al ibaba inc.. com com
ABSTRACT
Real-world dynamical systems often consist of multiple stochastic subsystems that interact with each other. Modeling and forecasting the behavior of such dynamics are generally not easy, due the inherent hardness in understanding the com- plicated interactions and evolutions of their constituents. This paper introduces the relational state-space model (R-SSM), sequential hierarchical latent variable model that makes use of graph neural networks (GNNs) to simulate the joint state transitions of multiple correlated objects. By letting GNNs cooperate with SSM, R-SSM provides a flexible way to incorporate relational information into the mod- eling of multi-object dynamics. We further suggest augmenting the model with normalizing flows instantiated for vertex-indexed random variables and propose two auxiliary contrastive objectives to facilitate the learning. The utility of R-SSM is empirically evaluated on synthetic and real time series datasets.
INTRODUCTION
Many real-world dynamical systems can be decomposed into smaller interacting subsystems if we take a fine-grained view. For example, the trajectories of coupled particles are co-determined by per- particle physical properties (e.g., traffic flow can be viewed as the coevolution of large number of vehicle dynamics. Models that are able better capture the complex behavior of such multi-object systems are of wide interest to various communities, e.g.. physics, ecology, biology, geoscience, and finance.
and velocity) and their physical interactions (e.g., gravity);
sptle- models (SSMs) are wide class sequential latent variable models (LVMs) that serve as workhorses for the analysis of dynamical systems and sequence data. Although SSMs are tra- ditionally designed under the guidance of domain-specific knowledge or tractability consideration, recently introduced deep SSMs (Fraccard. 2018) use neural networks (NNs) to parameterize flex- develop deep SSMs
ible state transitions and emissions, achieving much higher expressivity. for multi-object systems, graph neural networks (GNNs) emerge to be promising choice, they have been shown to be fundamental NN building blocks that can impose relational inductive bias explicitly and model complex interactions effectively (Battaglia all 2018)
Recent works that advocate GNNs for modeling multi-object dynamics mostly make use of GNNs in autoregressive (AR) fashion. AR models based recurrent (G)NNs can be viewed special instantiations of SSMs in which the state transitions are restricted to being deterministic (Fraceard. 2018, Section 4.2). Despite their simplicity, it has been pointed out that their modeling capability is bottlenecked by the deterministic state transitions (Chung et all, 20T5: Fraccaro cf all 2016) and the oversimplified observation distributions (Yang et all, 2018).
In this study, we make the following contributions: (i) We propose the relational state-space model (R-SSM), novel hierarchical deep SSM that simulates the stochastic state transitions of interacting objects with GNNs, extending GNN-based dynamics modeling to challenging stochastic multi-ob. ject systems. (ii) We suggest using the graph normalizing flow (GNF) to construct expressive joint
"Corresooniigg author.

---

Published as a conference paper at ICLR 2020
state distributions for R-SSM, further enhancing its ability to capture the joint evolutions of corre- lated stochastic subsystems. (iii) We develop structured posterior approximation to learn R-SSM using variational inference and introduce two auxiliary training objectives to facilitate the learning.
Our experiments on synthetic and real-world time series datasets show that R-SSM achieves com- petitive test likelihood and good prediction performance in comparison to GNN-based AR models and other sequential LVMs. The remainder of this paper is organized as follows: Section briefly reviews neccesary preliminaries. Section 3 introduces R-SSM formally and presents the methods to learn R-SSM from observations. Related work is summarized in Section 4 and experimental evaluation presented in Section We conclude the paper in Section 6.
2 PRELIMINARIES
In this work, an attributed directed graph is given by a 4-tuple g = E V.E), where V= [N] = NN) ; is the set of vertices, & C [M] x [M] is the set of edges, E Nxdu is matrix static vertex attributes, and E R xd. is a sparse tensor storing the static edge attributes. The set of direct predecessors of vertex is notated as N₁ = (pl(p, i) E}. We use the notation x; to refer to the i-th row of matrix X and write Xij to indicate the (i.j)-th entry of tensor X (if the corresponding matrix or tensor appears in the context). For sequences, we write Xgr = X1 = (x1,..., Xx) and
switch to xj
(6) for referring the i-th row of matrix X..
2.1 GRAPH NEURAL NETWORKS
GNNs are a class of neural networks developed to process graph-structured data and support rela- tional reasoning. Here focus vertex-centric GNNs that iteratively update the vertex repre- sentations of a graph G while being equivariant (Maron et all 2019) under vertex relabeling. Let H Nxd xd be matrix of vertex representations, in which the i-th row h, E IR is the vectorized representation attached to vertex i. Conditioning on the static graph structure and attributes given by G. a GNN just takes the vertex representations H along with some graph-level context g E R& as input and returns new vertex representations H' Nxd xd' output, i.e., H' GNN g, H).
When updating the representation of vertex from h, to h/, a GNN takes the representations of other nearby vertices into consideration. Popular GNN variants achieve this through multi-round message passing paradigm, in which the vertices repeatedly send messages to their neighbors, ag- gregate the messeages they received, and update their own representations accordingly. Formally, the operations performed by basic block of message- message-passing GNN are defined as follows:
VÚ,ii) E
M MESSAGE (g. Vj, Vi. ejis hi)
V
A AGGREGATE ((Mp-s)peN)
(2)
VIEV:
h COMBINE (g. V,, A;)
Throughout this work, implement Equations (1) and (2) by adopting a multi-head attention mech- anism similar to Vaswani ct all (2017) and Velkovi et all (20T)). For Equation (L), we use either RNN cell or a residual block (He et all 2 depending on whether the inputs GNN are RNN states or not. We write such block as H' = MHA(G, g. H) and give its detailed implementation in the Appendix. A GNN simply stacks L
parateerrparmmeteiizeee MHA blocks and iteratively computes H = ( (()) H', in which H MHA g, H(-)) for I 1, I,.... L. We write this construction as H' = (NN H) and treat it as black box to avoid notational clutter.
2.2 TTATE-SPACE MODELS
State-space models are widely applied analyze dynamical systems whose true states are not directly observable. Formally, an SSM assumes the dynamical system follows a latent state pro- cess {2ch211 which possibly depends on exogenous inputs (urh21- Parameterized by some (unknown) static parameter 6, the latent state process is characterized by initial density 2i ~ To(-)u;) and a transition density Ze11 ~ fo (: ugz+1). Moreover, at each time step, some noisy measurements of the latent state observed through observation density:
2

---

Published as a conference paper at ICLR 2020
H
Hi+i
He-l
Heal
2
22
21
X.
X
(a) Generative model.
(b) Collapsed generative model.
(c) Inference model.
Figure I: Graphical structures of R-SSM. Diamonds represent deterministic states and circles repre- sent random variables. To be concise, the dependencies the graph & and exogenous inputs U are not shown. (b) is the result collapsing all deterministic states in (a) and writing 2, = (z). Ze). In (c), solid lines represent the computation shared with the generative model and dashed lines rep- resent additional computation for inference.
X ~ go (lIzge, us+) The joint density of x1.7 and 21:7 factors as: p(XITT ZieT ( Te(zi II fe(zi ugd) II ge (xilzsr, uge).
The superior expressiveness of SSMs can be seen from the fact that the marginal predictive dis- tribution p(x, X ugi)
p(x,, Zga Xxt, ugr) dzgr can be far more complex than unimodal distributions and their finite mixtures that common in AR models. Recently developed deep SSMs use RNNs to compress zg, (and uge)) into fixed-size vectors to achieve tractability. As shown next section, R-SSM can be viewed as enabling multiple individual deep SSMs to communicate.
2.3 NORMALIZING FLOWS
Normalizing flows (Rezende & Mohamed, 2015) are invertible transformations that have the capabil- ity to transform a simple probability density into complex one (or vice versa). Given two domains X' R and y C TR,,, let f X y be an invertible mapping with inverse f-1. Applying f to random variable E X' with density p(z). by the change of variables rule, the resulting random variable = f(z) E y will have a density:
0f-1 On'
p(z) det 0f
-1
p(z') - p(z) det
On
A series of invertible mappings with cheap-to-evaluae determinants can be chained together to achieve complex transformations while retaining efficient density calculation. This provides a pow- erful way construct expressive distributions.
RELATIONAL STATE-SPACE MODEL
Suppose there is a dynamical system that consists of multiple interacting objects, and observing this system specific time accomplished by acquiring measurements from every individual object simultaneously. We further assume these objects are homogeneous, i.e., they share the same mea- surement model, and leave systems whose constituents are nonhomogencous for future work. To generatively model time-ordered series of observations collected from this system, the straightfor- ward approach that builds an individual SSM for each object usually unsatisfactory. as it simply assumes the state of each object evolves independently and ignores the interactions between objects. To break such an independence assumption, our main idea is to let multiple individual SSMs interact through GNNs, which expected to capture the joint state transitions of correlated objects well.
3

---

Published as a conference paper at ICLR 2020
3.1 GENERATIVE MODEL
Given the observations for a multi-object dynamical system, our model further assumes its inter- action structure is known as prior knowledge. The interaction structure is provided as directed graph, in which each object corresponds to a vertex, and a directed edge indicates that the state of its head is likely be affected by its tail. In situations where such graph structure is not available, complete graph can be specified. However, to model dynamical systems comprising a large number of objects, it is often beneficial to explicitly specify sparse graph structures, because they impose stronger relational inductive bias and help save the computational cost.
relational saate-ppae model assumes set of correlated dynamical subsystems evolve jointly under the coordination of graph neural networks. Formally, given a graph g = (V.E, V, E)., in which an edge (i.j) E & indicates that the state of vertex j may be affected by vertex Let u&!) Rd. (8) dh be the input and observation for vertex at time step f, respectively. For T steps, we introduce a set of unobserved random variables { 1.V in which z.! E Rd represents the latent state of vertex at time step f. Furthermore, we introduce global latent variable z E Rdd for each time step to represent the global state shared by all vertices. Conditioning on the graph and exogenous inputs, an R-SSM factorizes the joint density of observations and latent states as follows:
and X
N {z (1) 1 N po (fx)T)) i= (irsii 1 zi 9,
(1) N (1rr)==)
T
II fo ((
Ji=1 Ji=1 2
( (2 (
T
N go
10 N St [ i- 9, (1
(4)
from now The joint transition density fo is further factorized product of global transition density fä and local transition density fě. i.e., fu (Z. = f (z.l...) .) fā (Zelz......
For notational simplicity, we switch to the matrix notation Z: = 2 2
(1)
(N) T
To instan-
tiate these conditional distributions, a GNN accompanied by RNN cells is adopted to recurrently compress the past dependencies at each time step into fixed-size context vectors. Specifically, the observations are assumed be generated from following process:
6, (9) = RNN (h 1--1 [2221 (T],
b READOUT H.)
h RNN: (h [2 h[]),
z (. hh9
H: = GNN, (g, 21 H.),
Z, ~ fo (: H.)
( go( 12 (2) (0)
h 2/. h2)
for i I,....N and t 1....,T, T, where h h and 2 z Here hÿ, z h& and z learnable initial states. The READOUT function aggregates the context vectors of all vertices into a global context vector in a permutation-invarian manner. The global transition density I is specified to be a diagonal Gaussian distribution whose mean and variance are parameterized by the output of multilayer perceptron (MLP), and the local transition density fā will be discussed later. The local observation distribution ge can be freely selected in line with the data, and in our experiments it is either Gaussian distribution or a mixture of logistic distributions parameterized by MLPs.
The graphical structure of two consecutive steps of the generating process is illustrated in Figure Da note that the N + 1 latent state processes
An intuitive way to think about this generative model interact through the GNN, which enables the new state of a vertex to depend on not only its own state trajectory but also the state trajectories of other vertices and the entire graph.
3.2 LEARNING AND INFERENCE
As illustrated in Figure IR writing Z, = (z], Ze) and suppressing the dependencies the graph g and exogenous inputs U UKT, an R-SSM can be interpreted as an ordinary SSM in which the entire graph evolves whole, i.e., the joint density of latent states and observations factors as:
4

---

Published as a conference paper at ICLR 2020
P(X1:7; Z1:r) IT'=1 Po Zze (x Given observations X1:1, we interested in learn- ing unknown parameters 8 and inferring unobserved states ZI.T. For the learning task we wish to maximize the marginal likelihood pa (Xir) J Po 21:T) dZ but in our case the integral is intractable. We adopt a recently developed variational inference (VT) approach called variational sequential Monte Carlo (VSMC) (Maddison which maximizes a variational lower bound on the log marginal likelihood instead and learns the proposal distributions for the inference task simultancously.
all 2017; Naesseth all, 2018; Le et all, 2018),
Given a sequence of proposal distributions {4 (2, Zct X parameterized by 6, running the sequential Monte Carlo (SMC) algorithm with K particles yields an unbiased marginal likelihood where w is the unnormalized importance weight of particle k at time The variational lower bound is obtained by applying the Jensen's 's inequality:
estimator pe...K (Xir) = II- [/KK]]]]]] Ek= wf].
LK (0, Φ) = [log oo (1.T) E log IE 6. (XiT) = log pa [X177 1:T))
(5)
Assuming the proposal distributions are reparameterizable (Kingma & Welling 2014),, we use the biased gradient estimator V (0, R V log (:)) to maximize LSMC
Proposal design. We make the proposal for Z, depend on the information up time L and share some parameters with the generative model. We also choose to factorize qa (2,l...) = r (z)) -.)l-..)r& (Z, ..... The proposal distributions for all time steps are structured as follows:
1 = RNN, (b [x,", u, I,,,),
by = READOUT a, (.)
B, GNN (g. bj. B.), 2 ~ ( bbg b,),
by READOUT (9. B₂)
Z. ~ Z.-r B.)
fori L...,,..,.. N and t .,..., T. where hy and HiT are computed using the relevant parts of the generative model. r is specified to be diagonal Gaussian parameterized by an MLP, and r3 will be discussed soon. Here B, can be interpreted as belief state (Gregor et all 2019), which summarizes past observations X<: (and inputs Us+) deterministically. The graphical structure of this proposal design is shown in Figure Id and the detailed VSMC implementation using this proposal is given in Appendix A.4
3.3 GRAPH NORMALIZING FLOW
The local transition density fă (Ze]...) and the local proposal density T3 (Z.....) may be param- eterized in several ways. One simple and efficient starting point is tribution: fø
Gaussian dis- which assumes that the object states are conditionally
= II,, ((2("),.. independent, i.c., the joint state distribution is completely factorized over objects. We believe that such an independence assumption is an oversimplification for situations where the joint state evolu- tion is multimodal and highly correlated. One possible way to introduce inter-object dependencies is modeling joint state distributions as Markov random fields (MRFs) (Naesseth et all, 201 but this will significantly complicate the learning.
Here we introduce the Graph Normalizing Flow (GNF) U. which adapts Glow (Kingma & Dhariwal, 2018) to graph settings and enables us build expressive joint distributions for correlated random variables indexed by graph nodes. As described carlier, the key ingredient for flow is a series invertible mappings that are iteratively applied the samples of a base distribution. Now we are interested in the case where the samples are vertex states Z,, and thus the invertible mappings should be further constrained to be equivariant under vertex relabeling- This rules out popular autoregressive flows, e.g., IAF (Kingma et al, 20IG) and MAF (Papamakarios et an, 2017).
Our GNF built upon the coupling layer introduced in Dinhetal (2017)), which provides flexible framework to construct efficient invertible mappings. A GNF coupling layer splits the input Z E I xDD into two parts, Z₁ E IR N xd and Z₁ E (DD--)) The output Z' = N xD is formed as:
Z Z. Z Za O exp (Z.)) - +(Z.), Z' IZ Zr
GNF has heen independently developed by al (01) for different purpose.
5

---

Published as a conference paper at ICLR 2020
where @ denotes the element-wise product, and the functions s(-) and t(-) are specified to be GNNs to enforce the equivariance property. A GNF combines a coupling layer with a trainable element- wise affine layer and an invertible x convolution layer (Hoogeboom et all, 2019), organizing them as: Input + Affine - Coupling - Convix1 + Output. A visual illustration of this architecture is provided in Appendix A3.
In order to obtain more expressive prior and variational posterior approximation, the local transition density and local proposal density can be constructed by stacking multiple GNFs on top of diagonal Gaussian distributions parameterized by MLPs. With the message passing inside the coupling layers, GNFs can transform independent noise into correlated noise and thus increase model expressivity. The 1 x 1 convolution layers free us from manually permuting the dimensions, and the element-wise affine layers enable us to tune their initial weights to stablize training.
3.4 AUXILIARY CONTRASTIVE PREDICTION TASKS
In our initial experiments, we found that learning R-SSM suffered from the posterior collpase phenomenon, which is a well known problem in the training of variational autoencoders (VAEs). means that the variational posterior approximation qs (a(2, 2ct; Xxx) degenerate into the prior fo (o(z))2 Zxx) in the early stage of optimization, making the training dynamics get stuck in undesir- able local optima. Besides, we also encountered more subtle problem inherent in likelihood-base training of deep sequential models. That is, for relatively smooth observations, the learned model tended to only capture short-term local correlations but not the interaction effects and long-ten term transition dynamics.
Motivated by recent advances in unsupervised representation learning based on mutual information maximization, in particular the Contrastive Predictive Coding (CPC) approach (Oord et all 2018), we alleviate these problems by forcing the latent states to perform two auxiliary contrastive predic- tion tasks, At each time step the future observations of each vertex i are summarized into a vector
using a backward RNN: c, = R (x) Then we define two auxiliary CPC objectives:
T-1
Av,1 (2 (0) (.(
T-I
N
C E
log
E
Σ log
Eeer A (2
Eeen, An.2 (h,
where 2(()
6100 MLP (E h("),, and Zi is a set that contains cf" and some negative samples. The expectation is over negative samples and the latent states sampled from the filtering distributions. The positive score functions Ap.1 and X are specified to be simple log-bilinear models.
[2]
22
Intuitively, T encourages the latent states to encode useful information that helps distinguish the future summaries from negative samples. La encourages the deterministic states reflect the interaction effects, as it contrastingly predicts the future summary of vertex i based on the states of i's neighbors only. The negative samples are selected from the future summaries of other vertices within the minibatch. The final objective to maximize is L LSMC 3 3,CT.x + B2 in which 31 Z and 32 Z are tunable hyperparameters. The procedure to estimate this objective is
described in Appendix A.4.
RELATED WORK
GNN-based dynamics modeling. GNNs (Scarselh
0009 Duvenaud
2015:
2016: Defferrard et all 2016: Gilmer et all 2017: Hamilion et all 2017 Velkkvi et all. 2018: Xu
etall 20TY Maron etall 2019) provide a promising framework to learn on graph-structured data and impose relational inductive bias in learning models. We refer the reader to Battaglia et all (2()) for a recent review. GNNs (or neural message passing modules) are the core components of re- cently developed neural physics simulators (Battagla et all, 2016: Watters el all 2017: Chang et all 2017: [lanner et all 2019; Sancheza Gonzalel et all 2018; Mrowca el all 2018: Li et all 2019) and spatiotemporal or multi-agent dynamics models (Alahi Saahettal all 2016 Hoshen, 2017: Lr all 201 Zhang et all 2018 Lagehetti et all 201% Chen et all 2020). these works, GNNs usually act au- toregressively be integrated into the sequence-to-sequence (seq2seq) framework (Sutskever

---

Published as a conference paper at ICLR 2020
20E.). Besides, recently they have been combined with generative adversarial networks (Goodtel
low et all 2014) and normalizing flows for multi-agent forecasting (Gupta et all, 2018: Kosaraju all 2012; Khinehart all 2019). R-SSM differs from all these works by introducing structured latent variables to represent the uncertainty on state transition and estimation.
GNNS in sequential LVMs. A few recent works have combined GNNs with a sequential latent variable model, including R-NEM (van Steenkiste et all, 2018), NRI (Kıpt et all 2018). SQAIR (Kosiorek 2018), VGRNN (Hajiramezanali et all 2019) MFP (Tang & Salakhutdinov, 2019). and Graph VRNN (Sun et all 2012, Yeh et all 2019). The latent variables in R-NEM and NRI are discrete and represent membership relations and types of edges, respectively. In contrast, the latent variables in our model are continuous and represent the states of objects. SQAIR is also deep SSM for multi-object dynamics, but the GNN is only used in its inference network. VGRNN focused on modeling the topological evolution of dynamical graphs. MFP employs a conditional VAE architecture, in which the per-agent discrete latent variables are shared by all time steps. The work most relevant to ours is Graph VRNN, in which the hidden states of per-agent VRNNS interact through GNNs. Our work mainly differs from it by introducing a global latent state process to make the model hierarchical and exploring the use of normalizing flows as well as the auxiliary contrastive objectives. More subtle differences are discussed in Section 5.22
Deep LVMs for sequential data. There has been growing interest in developing latent variable models for sequential data with neural networks as their building blocks, among which the works most relevant to ours stochastic RNNs and deep SSMs. Many works have proposed incorporating stochastic latent variables into vanilla RNNs to equip them with the ability to express more complex data distributions (Bayer & Osendorfer, 2014 Chung et all 2015; Fraccaro all, 2016; [oyaa atlll 2017: Ke et all 2019) or, from another perspective, developing deep SSMs by parameterizing flexible transition and emission distributions using neural networks (Krishnan et all, 2017: Eraeaar
et all. 2017: [Bessing et alj 2018: Zheng et all 2017: Hainer et all 2019). Approximate inference and parameter estimation methods for nonlinear SSMs have been extensively studied in the literature (Doucet & Johansen.
2009; Andricu ct 2010: Kantas et all 2015: Gu e0E;; Karl ct 2016
Marino et all 2018; Gregor et all, 201% Hirt & Dellaportas 2019). We choose VSMC (Maddison et all 2017: Naesseth et all 2018; Le et all, 20R) as it combines the powers of VI and SMC. The posterior collapse problem is commonly addressed by KL. annealing, which does not work with VSMC. The idea of using auxiliary costs to train deep SSMs has been explored in Z-forcing (Goyal et all 2017; iKe et all, 2019), which predicts the future summaries directly rather than contrastingly. As result, the backward RNN in Z-forcing may degenerate easily.
EXPERIMENTS
We implement R-SSM using the TensorFlow Probability library (Dillon ell 2012). The experi- ments are organized as follows: In Section 5.1. we sample toy dataset from simple stochastic multi-object model and validate that R-SSM can fit it well while AR models and non-relational models may struggle. In Section 5 R-SSM is compared with state-of-the-ard sequential LVMs for aulti-agent modeling on a basketball gameplay dataset, and the effectiveness of GNF is tested through ablation studies. Finally, in Section 5.3, the prediction performance of R-SSM is compared with strong GNN-based seq2seq baselines on a road traffic dataset. Due to the space constraint, the detailed model architecture and hyperparameter settings for each dataset are given in the Appendix. Below, all values reported with eтror bars are averaged over 3 or 5 runs.
5.1 SYNTHETIC TOY DATASET
First we construct simple toy dataset to illustrate the capability of R-SSM. Each example in this dataset is generated by the following procedure:
G ~ SBM(N,, K, po: P1): Vi ~ N0,1 = Normal 1)
=(1)
17 Tv: + ai
(1) E /M + 022
(6)
4 Normal (0)
(i) ~ Normal tanh(Ez) (0) )20)
7

---

Published as a conference paper at ICLR 2020
Table 2: Test log-likelihood and rollout quality comparisons on the basketball gameplay dataset (offensive players only).
Model
1000
ELBO
Speed
Distance
OOB
VRNN
2360
0.89
43.78
33.78
MI-VRNN
2362
0.79
38.92
15.52
R-SSM
2459.8
2372.3 8
0.83 =:01
40.75 15
1.84 16
+2
2463.3 1.4
2380.2 1:6
0,82 1:01
40.36 1:23
2.17 1:00
+NNF (4)
2483.2
2381.6 1:4
0.80 1:00
39.37 1:35
2.06 15:15
+GNF (8)
2501.6 1.2
2382.1 1.4
0.79 1:00
39.14 1.29
2.12 1.10
Ground Truth
0.77
37.78
2.21
for == =1...... and /. 1..... Here SBM short for the symmetric stochastic block model, in which each vertex i belongs to exact one of the K communities, and two vertices and are connected with probability Po they are in the same community, P1 otherwise. covariate vector V; E gd. is attached to each vertex i, and by Equation (6). the state of each vertex i can be affected by its neighbors N₁. Choosing the parameters dg = 4, N = 36, K 3. Pu = 1/3, P1 1/18, T 80, 01 5.0, a2 1.5, 7 [-1.5, 0.4, 2.0, 0.9] dr O= 0.05, and E = 2.5, we generate 10K examples for training, validation, and test, respectively. A typical example visualized in the Appendix.
vertex-speciffc
Despite the simple generating process, the resulting dataset is highly challenging for common mod- els to fit. To show this, compare R-SSM with several baselines, including (a) VAR: Fitting first-order vector autoregression model for each example; (b) VRNN: A variational RNN (Chung et all 2015) shared by all examples; (c) GNN-AR: A variant of the recurrent decoder of NRI (Kipl all 2018), which is exactly a baseb AR model when given the ground-truth graph. VAR {.1-T] only, while GNN-AR and R-SSM additionally given access to the graph structure (V,E) (but not the vertex covariates). GNF is not used in R-SSM because the true joint transition distribution is factorized over vertices.
and VRNN given access to the observations (-:(] L.
For each model, we calculate three metrics:
(1) LL: Average likeliool (or its lower bound) of test examples; (2) MSE: Average mean squared one-step prediction error given the first 75 time steps of each test example; (3) CP: Average coverage probability of 90% one-step prediction interval. For non- analytic models, point predictions and pre- diction intervals computed using 1000 Monte Carlo samples. The results are re- ported in Table MI
Table 1: Test log- likellioioood and prediction perfor-
mance comparisons on the synthetic toy dataset.
Model
LL
MSE
CP
VAR
-366
0.679
+.000
0.750
VRNN
2-2641
0.501
4.003
0.93 4: 4.002
GNN-AR
-94
0.286
00002
0.806 : 004
R-SSM
2588
0.029
1.001
0.883 :: 000
+
26647
0.024 1.001
0.897 1: 1.101
The generating process involves latent fac- tors and nonlinearities, so VAR performs poorly as expected. VRNN largely underfits the data and struggles to generalize, which may be caused by the different topologies under the examples. In contrast, GNN-AR and R-SSM gener- alize well as expected, while R-SSM achieves much higher test log-likelihood and produces good one-step probabilistic predictions. This toy case illustrates the generalization ability of GNNs and suggests the importance of latent variables for capturing the uncertainty in stochastic multi-object systems. We also observed that without L the training dynamics easily get stuck posterior collapse at the very early stage, and adding L help improve the test likelihood.
5.2 BASKETBALL GAMEPLAY
In basketball gameplay, the trajectories of players and the ball are highly correlated and demonstrate rich. dynamic interations. Here we compare R-SSM with a ofate-of-the-ar hierarchical sequential L.V for multi-agent trajectories (Zhan etl 201), in which the per-agent VRNNs coordinated
8

---

Published as a conference paper at ICLR 2020
Table 3: Test log-likelihood comparison on the basketball gameplay dataset (offensive players plus the ball).
Table 4: Forecast MAE comparison on the METR- LA dataset. h is the number of steps predicted into the future. The X baseline outputs X to predict X,,
Model
LL
1000
Model
h 3
h
12
et (2IT9))
GRNN
2264
X,-A DCRNN
3.97
4.99
6.65
VRNN
>>250
2.77
3.15
3.60
GVRNN
>-832
GaAN
2.71
3.12
3.64
R-SSM+C2 C2 +GNF (8)
>2761 11
2805 1:0
R-SSM
2.67 2.00
3.14 :.01
3.72 =.02
22783 11
2826 +:0
CP
0.896 1:000
0.891
1:001
0.883 -..00
by a global "macro intent" model. We note it as MI-VRNN. The dataser includes 107,146 training examples and 13.845 test examples, each of which contains the 2D trajectories of ten players and the ball recorded at 6Hz for 50 time steps. Following their settings, we use the trajectories of offensive team only and preprocess the data in exactly the same way make the results directly comparable. The complete graph of players is used as the input to R-SSM.
Several ablation studies are performed to verify the utility of the proposed ideas. In Table B. we report test likelihood bounds and the rollout quality evaluated with three heuristic statistics: average speed (feet/step). average distance traveled (feet), and the percentage of out-of-bound (OOB) time steps. The VRNN baseline developed by Zhan et all (2019) is also included for comparison. Note that the VSMC bound LSMC is tighter log-likelihond approximation than the ELBO (which is equivalent to L The rollout statistics of R-SSMs are calculated from 150K 50-step rollouts with 10 burn-in steps. Several selected rollouts are visualized in the Appendix.
As illustrated in Table 2. all R-SSMs outperform the baselines in terms of average test log- Again, we observed that adding LTT is necessary for training R-SSM successfully on this dataset. Training with the proposed auxiliary loss CTT and adding GNFs do improve the results. R-SSM with GNFs (4 in prior, 4 in proposal) achieves higher likelihood than R-SSM with GNFs, indicating that increasing the expressivity of joint state distributions helps fit the data better. As for the rollout quality, the OOB rate of the rollouts sampled from our model matches the ground-truth significantly better, while the other two statistics are comparable to the MI-VRNN baseline.
In Table B we also provide preliminary results for the setting that additionally includes the tra- jectory of the ball. This enables us to compare with the results reported by Yeh er all (20119) for Graph VRNN (GVRNN). The complete graph of ball and players served as input to R-SSM is an- notated with two node types (player or ball) and three edge types (player-to-bally toyyer-)-- R-SSM achieves competitive test likelihood, and adding GNFs helps improve the performance. We point out that several noticeable design choices of GVRNN may help outperform R-SSM: (i) GVRNN uses a GNN-based observation model, while R-SSM uses a simple factorized observation model. (ii) GVRNN encodes X 1it-1 into H, and thus enables the prior of Z, to depend on past observations, which is not the case in R-SSM. (iii) GVRNN uses several implementation tricks, e.g., predicting the changes in observations only (X, X,- + AX,) and passing raw ob- servations as additional input to GNNs. We would like to investigate the effect of these interesting differences in future work.
ball-to-player or
5.3 ROAD TRAFFIC
Traffic speed forecasting road networks is an important but challenging task, as the traffic dy- namics exhibit complex spatiotemporal interactions. In this subsection, we demonstrate that R-SSM is comparable to the state-of-the-art GNN-based seq2seq baselines on real-world traffic dataset. The METR-LA dataset (L cLl R contains 4 months of ID traffic speed measurements that were recorded via 207 sensors and aggregated into 5 minutes windows. For this dataset, all con- ditional inputs G = (V,E, E, V, E) and U1:1 are provided to R-SSM, in which & constructed by connecting two sensors if their road network distance is below threshold, V stores the geographic
Data Source: STATS, copyright 2019.

---

Published as a conference paper at ICLR 2020
positions and learnable embeddings of sensors, E stores the road network distances of edges, and UI:T provides the time information -f and day-of-weekk. We impute the missing values for training and exclude them from evaluation. GNF is not used because of GPU memory limitation. Following the settings in et all (2018), we train our model on small time windows spanning hours and use 7:1:2 split for training, validation, and test.
The comparison of mean absolute forecast errors (MAE) is reported in Table ! The three forecast horizons correspond to 15, 30, and 60 minutes. We give point predictions by taking the element- wise median of 2K Monte Carlo forecasts. Compared with DCRNN (Li et all 2018) and GaAN (Zhang al. 2018), R-SSM delivers comparable short-term forecasts but slightly worse long-term forecasts.
We argue that the results are admissible because: (i) By using MAE loss and scheduled sampling, the DCRNN and GaAN baselines are trained on the multi-step objective that they are later evaluated on, making them hard beat. (ii) Some stochastic systems are inherently unpredictable beyond a few steps due to the process noise, e.g., the toy model in Section 5. In such case, multi-step MAE may not a reasonable metric, and probabistic forecasts may prefered. The average coverage probabilities (CP) of 90% prediction intervals reported in Table E indicate that R-SSM provides good uncertainty estimates. (iii) Improving the multi-step prediction ability of deep SSMs is still an open problem with few recent attempts (Ke et aD, 2019; Hainer et all 2019). We would like to explore it in future work.
CONCLUSIONS
In this work, we present deep hierarchical state-space model in which the state transitions of correlated objects are coordinated by graph neural networks. To effectively learn the model from observation data, we develop a structured posterior approximation and propose two auxiliary con- trastive prediction tasks to help the learning. We further introduce the graph normalizing flow to enhance the expressiveness of the joint transition density and the posterior approximation. The ex- periments show that our model can outperform match the state-of-the-arts on several time series modeling tasks. Directions for future work include testing the model on high-dimensional observa- tions, extending the model to directly learn from visual data, and including discrete latent variables in the model.
ACKNOWLEDGMENTS
This work was supported by the National Key Research and Development Program of China (No. 2018 (0B8B005000000 and the Alibaba-Zhejiang University Joint Institute of Frontier Technologies. Fan Yang would like thank Qingchen Yu for her helpful feedback on early drafts of this paper.
REFERENCES
Alexandre Alahi, Kratarth Goel, Vignesh Ramanathan, Alexandre Robicquet, Li Fei-Fei, and Silvio Savarese. Social Istm: Human trajectory prediction in crowded spaces. In Proceedings of the IEEE Conference Computer Vision and Pattern Recognition, 2016.
Christophe Andrieu, Arnand Doucet, and Roman Holenstein. Particle markov chain monte carlo methods. Journal of the Royal Statistical Society: Series (Statistical Methodology). 72(3):269- 342.2010. 2010.
Peter Battaglia, Razvan Pascanu, Matthew Lai, Danilo Jimenez Rezende, and koray kavukcuoglu. Interaction networks for learning about objects, relations and physics. In Advances in Neural Information Processing Systems 29. 2016.
Peter Battaglia, Jessica B Hamrick. Victor Bapst, Alvaro Sanchez-Gonzaler Vinicius Zambaldi. Mateusz Malinowski, Andrea Tacchetti, David Raposo, Adam Santoro, Ryan Faulkner, et al. Re- lational inductive biases, deep learning, and graph networks. arXiv preprint arXiv: (00::/806.1111 0161, 2018.
Justin Bayer and Christian Osendorfer. Learning stochastic recurrent networks. arXiv preprint arXiv: 1.11,7610, 2014.
10

---

Published as a conference paper at ICLR 2020
Lars Buesing, Theophane Weber, Sebastien Racaniere, SM Eslami, Danilo Rezende, David P Re- ichert, Fabio Viola, Frederic Besse, Karol Gregor, Demis Hassabis, et al. Learning and querying fast generative models for reinforcement learning, arXiv preprint arXiv:8002 03006, 2018.
Michael Chang, Tomer Ullman, Antonio Torralba, and Joshua Tenenbaum. A compositional object- based approach to learning physical dynamics. International Conference on Learning Repre- sentations. 2017.
Weiqi Chen, Ling Chen, Yu Xie, Wei Cao, Yusong Gao, and Xiaojie Feng. Multi-range attentive bicomponent graph convolutional network for traffic forecasting, In AAAI Conference on Artificial Intelligence, 2020.
Junyoung Chung. Kyle Kastner, Laurent Dinh, Kratarth Goel, Aaron C Courville, and Yoshua Ben- gio. recurrent latent variable model for sequential data. In Advances Neural Information Processing Systems 28, 2015.
Michaël Defferrard, Xavier Bresson, and Pierre Vandergheynst. Convolutional neural networks on graphs with fast localized spectral filtering. Advances Neural Information Processing Systems 29. 2016.
Joshua V Dillon, Ian Langmore. Dustin Tran, Eugene Brevdo, Srinivas Vasudevan, Dave Moore, Brian Patton, Alex Alemi, Matt Hoffman, and Rif A Saurous. Tensorflow distributions. arXiv preprint arXiv:1711.1066 2017.
Laurent Dinh, Jascha Sohl-Dickstein, and Samy Bengio. Density estimation using real nvp. In International Conference on Learning Representations, 2017.
Arnaud Doucet and Adam M Johansen. tutorial on particle filtering and smoothing: Fifteen years later. Handbook of Nonlinear Filtering, 2009.
David K Duvenaud, Dougal Maclaurin, Jorge Iparraguirre, Rafael Bombarell, Timothy Hirzel, Alan Aspuru-Guzio and Ryan P Adams. Convolutional networks graphs for learning molecular fingerprints. In Advances in Neural Information Processing Systems 28. 2015.
Marco Fraccaro. Deep Latent Variable Models for Sequential Data. PhD thesis, 2018.
Marco Fraccaro, Søren Kaae Sø nderby, Ulrich Paquet, and Ole Winther. Sequential neural models with stochastic layers. In Advances in Neural Information Processing Systems 29, 2016.
Marco Fraccaro, Simon Kamronn, Ulrich Paquet, and Ole Winther. A disentangled recognition and nonlinear dynamics model for unsupervised learning. In Advances Neural Information Processing Systems 30, 2017.
Justin Gilmer, Samuel S. Schoenholz, Patrick F. Riley, Oriol Vinyals, and George E. Dahl. Neu- ral message passing for quantum chemistry. In Proceedings of the International Conference on Machine Learning, 2017.
Ian Goodfellow, Jean Douget-Abadic, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial nets. In Advances Neural Infor- mation Processing Systems 27. 2014.
Anirudhand Goyal, Alessandro Sordoni, Marc Mlee-Alexandre Côté, Nan Rosemary Ke, and Yoshua Bengio. Z-forcing: Training stochastic recurrent networks. In Advances in Neural Information Processing Systems 30, 2017.
Karol Gregor, George Papamakarios, Frederic Besse, Lars Buesing, and Theophane Weber. Tempo- ral difference variational auto-encoder. International Conference on Learning Representations, 2019.
Shixiang (Shane) Gu, Zoubin Ghahramani, and Richard E Turner. Neural adaptive sequential monte carlo. In Advances in Neural Information Processing Systems 28. 2015.
Agrim Gupta, Justin Johnson, Li Fei-Fei, Silvio Savarese, and Alexandre Alahi. Social gan: So- cially acceptable trajectories with generative adversarial networks. In Proceedings of the IEEE Conference Computer Vision and Pattern Recognition, 2018.

---

Published as a conference paper at ICLR 2020
Danijar Hafner, Timothy Lillicrap, lan Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. In Proceedings of the International Conference Machine Learning, 2019.
Ehsan Hajiramezanali, Arman Hasanzadeh, Nick Duffield, Krishna Narayanan, Mingyuan Zhou, and Xiaoning Qian. Variational graph recurrent neural networks. In Advances in Neural Informa- fion Processing Systems 32. 2019.
Will Hamilton, Zhitao Ying, and Jure Leskovec. Inductive representation learning large graphs. In Advances in Neural Information Processing Systems 30. 2017.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, 2016.
Marcel Hirt and Petros Dellaportas. Scalable bayesian learning for state space models using varia- tional inference with sme samplers. In Proceedings of the International Conference Artificial Intelligence and Statistics, 2019.
Emiel Hoogeboom, Rianne van den Berg, and Max Welling. Emerging convolutions for generative normalizing flows. In Proceedings of the International Conference on Machine Learning, 2019.
Yedid Hoshen. Vain: Attentional multi-agent predictive modeling. In Advances in Neural Informa- tion Processing Systems 30. 2017.
Michael Janner, Sergey Levine, William T. Freeman, Joshua B. Tenenbaum. Chelsea Finn, and Jiajun Wu. Reasoning about physical interactions with object-centric models. In International Conference on Learning Representations, 2019.
Nikolas Kantas, Arnaud Doucet, Sumeetpal S Singh, Jan Maciejowski, Nicolas Chopin, et al. On particle methods for parameter estimation in state-space models. Statistical Science, 30(3):328- 351, 2015.
Maximilian Karl, Maximilian Soelch, Justin Bayer, and Patrick van der Smagt. Deep variational bayes filters: Unsupervised learning of state space models from raw data. In International Con- ference on Learning Representations. 2016.
Nan Rosemary Ke, Amanpreet Singh, Ahmed Touati, Anirudh Goyal, Yoshua Bengio, Devi Parikh, and Dhruv Batra. Modeling the long term future in model-based reinforcement learning. In International Conference Learning Representations, 2019.
Diederik Kingma and Max Welling, Auto-encoding variational bayes. In International Conference on Learning Representations. 2014.
Durk P Kingma and Prafulla Dhariwal. Glow: Generative flow with invertible 1xl convolutions. In Advances in Neural Information Processing Systems 31. 2018.
Durk P Kingma, Tim Salimans, Rafal Jozefowicz, Xi Chen, Ilya Sutskever, and Max Welling. Im- proved variational inference with inverse autoregressive flow. In Advances in Neural Information Processing Systems 29, 2016.
Thomas Kipf, Ethan Fetaya, Kuan-Chieh Wang, Max Welling, and Richard Zemel. Neural relational inference for interacting systems. In Proceedings of the International Conference on Machine Learning, 2018.
Vineet Kosaraju, Amir Sadeghian, Roberto Martín-Martin, Ian Reid, S. Hamid Rezatofighi, and Silvio Savarese. Social-bigat: Multimodal trajectory forecasting using bicycle-gan and graph attention networks. In Advances in Neural Information Processing Systems 32. 2019.
Adam Kosiorek, Hyunjik Kim, Yee Whye Teh, and Ingmar Posner. Sequential attend, infer, repeat: Generative modelling of moving objects. In Advances in Neural Information Processing Systems 31. 2018.
12

---

Published as a conference paper at ICLR 2020
Rahul G Krishnan, Uri Shalit, and David Sontag. Structured inference networks for nonlinear state space models. In AAAI Conference on Artificial Intelligence, 2017.
Tuan Anh Le, Maximilian Igl. Tom Rainforth, Tom Jin, and Frank Wood. Auto-encoding sequential monte carlo. In International Conference on Learning Representations, 2018.
Yaguang Li, Rose Yu, Cyrus Shahabi, and Yan Liu. Diffusion convolutional recurrent neural net- work: Data-driven traffic forecasting. In International Conference on Learning Representations, 2018.
Yujia Li, Daniel Tarlow, Marc Brockschmidt, and Richard Zemel. Gated graph sequence neural networks. International Conference Learning Representations, 2016.
Yunzhu Li, Jiajun Wu, Russ Tedrake, Joshua B. Tenenbaum, and Antonio Torralba. Learning par- ticle dynamics for manipulating rigid bodies, deformable objects, and fluids. In International Conference on Learning Representations, 2019.
Jenny Liu, Aviral Kumar, Jimmy Ba, Jamie Kiros, and Kevin Swersky. Graph normalizing flows. In Advances in Neural Information Processing Systems 32. 2019.
Chris J Maddison, John Lawson, George Tucker, Nicolas Heess, Mohammad Norouzi, Andriy Mnih, Arnaud Doucet, and Yee Teh. Filtering variational objectives. In Advances in Neural Information Processing Systems 30. 2017.
Joseph Marino, Milan Cvitkovic, and Yisong Yue. A general method for amortizing variational filtering. In Advances in Neural Information Processing Systems 31. 2018.
Haggai Maron, Heli Ben-Hamu, Nadav Shamir, and Yaron Lipman. Invariant and equivariant graph networks. In International Conference on Learning Representations, 2019.
Damian Mrowca, Chengxu Zhuang, Elias Wang, Nick Haber, Li F Fei-Fei, Josh Tenenbaum, and Daniel L. Yamins. Flexible neural representation for physics prediction. In Advances Neural Information Processing Systems 31. 2018.
C. A. Naesseth, F. Lindsten, and T. B. Schön. High-dimensional filtering using nested sequential monte carlo. IEEE Transactions on Signal Processing, 67(16):4177-4188 2019.
Christian Naesseth, Scott Linderman, Rajesh Ranganath, and David Blei. Variational sequential monte carlo. In Proceedings of the International Conference on Artificial Intelligence and Statis- tics, 2018.
Aaron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predic- tive coding. arXiv preprint 1rXiv: 2018
George Papamakarios, Theo Pavlakou, and lain Murray. Masked autoregressive flow for density estimation. In Advances in Neural Information Processing Systems 30. 2017.
Danilo Rezende and Shakir Mohamed. Variational inference with normalizing flows. In Proceedings of the International Conference on Machine Learning, 2015.
Nicholas Rhinehart, Rowan McAllister, Kris Kitani, and Sergey Levine. Precog: tioned on goals in visual multi-agent settings. The IEEE International Conference on Computer Vision, 2019.
Prediction condi-
Alvaro Sanchez-Gonzaler Nicolas Heess, Jost Tobias Springenberg, Josh Merel, Martin Riedmiller, Raia Hadsell, and Peter Battaglia. Graph networks as learnable physics engines for inference and control. In Proceedings of the International Conference on Machine Learning, 2018.
Franco Scarselli, Marco Gori, Ah Chung Tsoi, Markus Hagenbuchner, and Gabriele Monfardini. The graph neural network model. IEEE Transactions Neural Networks, 20(1):61-80, 2009.
Chen Sun, Per Karlsson, Jiajun Wu, Joshua B Tenenbaum, and Kevin Murphy, Stochastic prediction of multi-agent interactions from partial observations. In International Conference on Learning Representations, 2019.
13

---

Published as a conference paper at ICLR 2020
Iya Sutskever, Oriol Vinyals, and Quoe V Le. Sequence to sequence learning with neural networks. In Advances in Neural Information Processing Systems 27. 2014.
Andrea Tacchetti, H. Francis Song, Pedro A. M. Mediano, Vinicius Zambaldi, János Kramár, Neil C. Rabinowitz, Thore Graepel, Matthew Botvinick, and Peter W. Battaglia. Relational forward mod- els for multi-agent learning. In International Conference on Learning Representations, 2019.
Yichuan Charlie Tang and Ruslan Salakhutdinov. Multiple futures prediction. In Advances in Neural Information Processing Systems 32. 2019.
Sjoerd van Steenkiste, Michael Chang, Klaus Greff, and Jürgen Schmidhuber. Relational neural expectation maximization: Unsupervised discovery of objects and their interactions. In Interna- tional Conference on Learning Representations, 2018.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, ukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Infor- mation Processing Systems 30. 2017.
Petar Velikovi, Guillem Cucurull, Arantxa Casanova, Adriana Romero, Pietro Liò, and Yoshua Ben- gio. Graph attention networks. In International Conference on Learning Representations, 2018.
Nicholas Watters, Daniel Zoran, Theophane Weber, Peter Battaglia, Razvan Pascanu, and Andrea Tacchetti. Visual interaction networks: Learning physics simulator from video. In Advances in Neural Information Processing Systems 30. 2017.
Keyulu Xu, Weihua Hu, Jure Leskovec, and Stefanie Jegelka. How powerful are graph neural networks? International Conference on Learning Representations, 2019.
Zhilin Yang. Zihang Dai, Ruslan Salakhutdinov, and William W. Cohen. Breaking the softmax bottleneck: a high-rank RNN language model. In International Conference on Learning Repre- sentations, 2018.
Raymond A. Yeh, Alexander G. Schwing, Jonathan Huang, and Kevin Murphy. Diverse generation for multi-agent sports games. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, 2019.
Eric Zhan, Stephan Zheng, Yisong Yue, Long Sha, and Patrick Lucey. Generating multi-agent trajectories using programmatic weak supervision. In International Conference on Learning Rep- resentations, 2019.
Jiani Zhang, Xingjian Shi, Junyuan Xie, Hao Ma. Irwin King, and Dit-Yat Yeung, Gaan: Gated attention networks for learning on large and spatiotemporal graphs. In Conference on Uncertainty in Artificial Intelligence, 2018
Xun Zheng, Manzil Zaheer, Amr Ahmed, Yuan Wang, Eric Xing, and Alexander J Smola. State space 1stm models with particle meme inference, arXiv preprint 1rX19:711.111199, 2017.
14

---

Published as a conference paper at ICLR 2020
A APPENDIX
A.1 MHA IMPLEMENTATION
Each attention head & E [K] is separately parameterized and operates follows. For V(j,i E, produces a message vector - together with an unnormalized scalar weight wy Then for vi V, it aggregates all messages sent i in permutation-invarr manner:
paN softmax (HUpm)pEN;
at = EpEN, 0p p-r
Putting all K heads together, it turns out that Mi-ri = {(w 1 Specifically, each attention head is parameterized in query-key-value style:
1k= and A, = {at Ik=1
E V: V: = MLP.(vi) E Rd.. h = h; V.
Q HWQ A = HWA, C- H-W = cj W/+i = a V dy + (e))(
for d = d d do dg. We E Rdxde, WA Rdxi, and We E Råxde
A.2 MODEL. DETAILS
In this work, the READOUT function is implemented by passing the concatenation of the outputs of a mean aggregator and an element-wise max aggregator through a gated activation unit. The transition densities in the generative model are specified to be:
(z h2) = Normal (-Img (h²). Eg (h?))
f H.) II Normal (z(")
h
ΣS h (10) )
(8)
where mg and Dg (similarly pá and Σn are 3-layer MLPs that share their first layer. Eg and Mn output diagonal covariance matrices using the softplus activation. The proposal densities r and r are specified in a similar way. Then GNFs can stacked on top of fō and r to make them more expressive.
A.3 EXPERIMENTS
We use the Adam optimizer with an initial learning rate of 0.001 and a gradient clipping of 1.0 for all experiments. The learning rate was annealed according to linear cosine decay. We set 31 = 32 = 1.0 for the auxiliary losses in all experiments.
Synthetic toy dataset. A typical example in the dataset is visualized in Figure A The architectures of the models are specified as follows.
(a) VRNN: Using 128-dimensional latent variables and a two-layer, 512-unit GRU.
(b) GNN-AR: Using two-layer GNN and one-layer, 128-unit GRU shared by all nodes.
(c) R-SSM: We let dg = d 8. All RNNs are specified to be two-layer, 32-unit LSTMs. All MLPs use 64 hidden units. The generative model and the proposal both use a 4-head MHA layer. 4 SMC samples and batch size of 16 are used in training.
Basketball player movement. We let d, = d, 32. All RNNs are specified to be two-layer, 64- unit LSTMs and all MLPs use 256 hidden units. The generative model uses one 8-head MHA layer and the proposal uses two 8-head MHA layers. Each GNF uses additional MHA layer shared by the functions s-)) and t(-). 4 SMC samples and a batch size of 64 are used in training. Fight selected rollouts from the trained model are visualized in Figure
Road traffic. We le dg = d = 8. A 32-dimensional embedding for each sensor is jointly learned a part of the vertex attribute. All RNNs specified to be two-layer, 32-umit LSTMs and all MLPs use 64 hidden units. The generative model and the proposal both use two 8-head MHA layers. 3 SMC samples and batch size of 16 are used in training.
15

---

Published as a conference paper at ICLR 2020
WW MALA
W mn
M M Mn MO MALA
Mamaaaan M
WWV WWW.AMAA (WIN
MMM M Ma SANOOII WO
M. Mar MAAN
M ww
LA MY MWA.A. ANN
Mard
TL AA M.M. MOW
MWM A.e MMMM
Mudia a. MWWMMA
MaM MLWL
WOW www mwm MMMM wWw.AMMMM NMM
UUMAM PLMM UM
W
MWWWMAM MM
Figure An example from the toy dataset (N 36, =
Figure 3: Selected rollouts from the trained model. Black dots represent the starting points.
16

---

Published as a conference paper at ICLR 2020
A.4 TRAINING
We optimize the VSMC bound estimated by the following SMC algorithm:
Algorithm 1 Estimate the VSMC bound L
Input: graph G, observations XI.T, exogenous inputs ULTT
Require:
generative model (fo fg 99 ) proposal (ra, roh. number of particles K
for = 1.... do Simulate zi, (,, X1, U1) Simulate Zi,6 ~ (- 9, X1, U1)
Set w 18 (Z. 9. 11 go(x{") le, (. (2 (l..)
end for Initialize LiNe log k==l wi/K
K
fort = 2 2...TT do
[z 7 Z11,b 1K = RESAMPLE ZEik: W 1K
for k 1...KKKK do
Simulate zi, ~r (-122 Zeik G,X,, Use) Simulate Ze,4 (-
Zeck 9 X Xst, Usi)
Set w = 18 (.. 122 Zec.s... Il (x, (8)
S0,4 Z.c.....))
rg (x )r. (l.)
Set 2²2,,-- (2 ki,k), Zgeu (Z Z,,,,
end for
Update LSMC LSMC log Ek=1 w²/KK
end for
Output: ASMC
LK
In our model, dependencies on 2 and Zet, are provided through the compact RNN states h and H₁ When GNFs are used in fā and re density calculation and backpropagation are automat- ically handled by the TensorFlow Probability library.
To estimate the auxiliary objectives L and LTDx, we reuse the resampled unweighted particles ( Z1 K T-1 generated by Algorithm to form Monte Carlo estimations for them:
T-1
N
A
(1) (2,,0)
T-1
N
(2) (1) 2 C
Faux
log
Ásux
log Ecco
K
Eceny Ap,1 (4 c)
K
( A 2
1=1
where 2!12
1 n, 6. = MLP (EjzinjeN; 12)) and 24 is a set that contains e and
some negative samples selected from the future summaries of other vertices within the minibatch.
17

---

Published as a conference paper at ICLR 2020
A.5
GRAPH NORMALIZING FLOW
element wise affine layer
Oy +
Z R
Zu E R Vxd
Zo E Nx(--D)
Message Passing
element multiplication
element addition
Za E Nxd
Z,, Bx(00-))
invertible convixt layer
W 0xD
Ž E RxD
Figure 4: Visual illustration of GNE, Multiple GNFs can be stacked together achieve more expressive transformation.
The element-wiss affine layer is proposed by Kingma & Dhariwal (2018) for normalizing the activa- tions. Its parameters 7 E RD and B E RR. are initialized such that the per-channel activations have roughly zero and unit variance at the beginning of training. The invertible linear transforma- tion W e I DxD is parameterized using a QR decomposition (Hoogeboom et al, 2019).
18


Reviews:
[{"review_id": "r1eQTVd4qB", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "* Summary:\nThe paper presents a relational state-space model that simulates the joint state transitions of correlated objects which are hierarchically coordinated in a graph structure. A structured posterior approximation is developed based on sequential graph neural networks. Two auxiliary contrastive predictive losses are proposed to help circumvent the posterior collapse problem. Graph normalizing flow is further incorporated into the framework to make the joint state transition density more expressive. The proposed R-SSM shows performance gains over state-of-the-arts in three benchmarks.\n\n* Comments:\nThe paper is generally well written and technically sound. The framework, including formulation of each of its components, is well defined. It is also helpful that the authors included preliminaries of the literature. The number of experiments are adequate. However, there are a few parts that require more extensive clarification and analysis.\n1. Different parts of Section 3 appear to be rather disconnected, the reader still has a hard time figuring out how the learning of the whole framework is carried out. It is desirable to include a sketch of learning algorithm.\n2. In the formulation GNF, what is the intuition or principle to decouple the state Z_t into two parts Z_a and Z_b? How does the mapping of Z_b into Z'_b help to make the state transition distribution more expressive?\n3. In Section 5.3, the authors mention that GNF was not used due to memory cost. Could it be discussed more thoroughly about the complexities of learning R-SSM and GNF?\n4. The model keeps track of a global state z^g, but it is not analyzed in experiments. It is strongly recommended that the authors discuss about the (global and individual) states and their transitions. It would provide great insights on how multiple objects interact with each other.\n\nMinor point:\n1. Please clarify what is X_{t-h} in Table 3?", "extra_scores": {"experience_assessment": "I have read many papers in this area.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory."}, "extra_text": {}}, {"review_id": "S1eg3T6ZoS", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This manuscript proposes a novel approach to dynamic modeling of time series data based on a state space model that incorporates a Graph Neural Network to model the relationship between the dynamics of different objects. The (complex) architecture is described in details and to some extent justified before its predictive performance is demonstrated on several simulated and real datasets.\nOverall, while the authors provide evidence of a better predictive performance with respect to several baselines, I am left a bit unconvinced by the study for the following reasons:\n1.\tLack of relevant baselines: it seems clear that the overall purpose of the approach and the nature of the dataset require fitting a state space model, however, baseline are overall focused on recurrent and autoregressive models which seem underequipped to address these problems. I wonder if the choice of more relevant state space model baseline would convincingly show a true benefit of the proposed approach. In particular, there is likely a large number of variations of the Kalman filter and particle filters that might be relevant. For example, the ensemble Kalman filter has proved accurate for weather forecasting.\n2.\tLack of interpretability: given the proposed approach relies on an intuitive representation of the model as representing the dynamics of several object tied by an interaction graph, the purely predictive results are not really matching the expectation of the reader to “see” how well these interactions are captured by the model. Can we check in some way that the latent graphical model is learnt properly, even in a toy dataset?\n3.\tWriting of the methods section: This is a more vague comment, but while reading the overall description of the approach, one is left wondering how critical are each part of the model and whether some complexity could be spared. Moreover, some descriptions are difficult to follow, perhaps for the reader less familiar with customary design choices in dynamic neural networks, e.g. the objective in section 3.4. In addition, some statements seem at least to lack justification, e.g. stating that AR approaches lead to unimodal distributions at the bottom of page 2.\n", "extra_scores": {"experience_assessment": "I do not know much about this area.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper thoroughly.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory."}, "extra_text": {}}, {"review_id": "ByxfCSnyqS", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary of what the paper claims and contributes\n---\nThis paper proposes a hierarchical latent variable model of sequential dynamic processes of multiple objects when each object exhibit significant stochasticity. The model leverages Graph Neural Networks to design the architecture of the model. Normalizing flows are used to construct \"Graph Normalizing Flows\" that model transition densities of the per-object latent random variables, although the full model has more latent variables than tabouthe per-object latents. The approach applies and combines a host of methods to design the objective, the model, and optimize the model under the objective -- for the objective: variational inference, Contrastive Predictive Coding; for the model: normalizing flows, graph neural networks (and their combination, termed by the authors as Graph Normalizing Flows); for the optimization: variational sequential Monte Carlo. From my understanding, the main claims are that the model is a flexible way to incorporate relational information, and that the resulting model outperforms other multi-object time-series forecasting approaches.\n\nExperiments are conducted on a synthetic toy dataset of 1d observations per-object, a real dataset of basketball player movement, corresponding to 2d observations of each object's position (the agent's and the ball's positions), and a real dataset of traffic speed forecasting (1d observations per-object). The experiments demonstrate that the model compares favorably to a baseline and a few other time-series forecasting models from the literature. The experiments show that the method achieves higher test log-likelihood (lower-bounds), and lower MSE than the compared methods. The paper discussed many related works, but it's not clear why the specific methods were chosen for comparisons. More motivation is needed for these comparisons.\n\nEvaluation\n---\n\n>Originality:\nAre the tasks or methods new?\nThe method is new.\n\nIs the work a novel combination of well-known techniques?\nYes, albeit a lot of techniques.\n\nIs it clear how this work differs from previous contributions?\nBesides the fact that the proposed method is a combination of many techniques, the main / key technical difference was actually unclear to me.\n\nIs related work adequately cited?\nFrom my understanding, yes, mostly; there's just a few places where the paper could be situated better:\n- Sec 1: \"as they have been shown to be fundamental NN building blocks ...\". Citations needed (e.g. include the Graph NN citations later here)\n- Sec 1: \"R-SSM achieves superior test likelihood and good prediction performance\". When compared to what methods (citations to VRNN and GNN-AR/NRI needed here too)?\n- The paper contrasted the work to other dynamics modeling works with the statement \"R-SSM differs from all these works by introducing structured latent variables to represent the uncertainty on state transition and estimation\". The paper also said VRNNs were the most related, because they employ \"hidden states of per-agent VRNNs interact through GNNs\". The paper (PRECOG, arXiv:1905.01296) appears related as it also constructs a latent variable graphical model of stochastically-transitioning multiple interacting objects with normalizing flows. Like the authors said of Graph VRNNs, it appears that one of the biggest differences to PRECOG looks like the proposed method includes more sources of stochasticity outside of the objects (the global uncertainty) to model the observations.\n\n>Quality:\nIs the submission technically sound?\nYes, although not much motivation is given for many of the specific design decisions adopted. The model is a complicated combination of many techniques from the literature, which makes it difficult to understand the importance of each specific piece of the model.\n- For instance, why was global stochasticity introduced? Because per-object stochasticity is modeled with normalizing flows, the model would have been trainable with maximum likelihood estimation of a distribution over observations that's efficiently analytically computable, but after the addition of global stochasticity, the joint distribution over the random variables becomes intractable, and training must resort to variational inference. This leads to, as the paper mentioned, RSSM suffering from posterior collapse, which, to my knowledge, does not occur when training direct maximum likelihood estimation procedures (e.g. GLOW, RealNVP).\n- As mentioned above, it's not clear what the key technical innovation of this approach is (Graph Normalizing Flows?) It would be good if the claimed innovation was made clear in both the abstract and the introduction.\n\nAre claims well supported by theoretical analysis or experimental results?\nThe experimental results provide evidence that 1) the model can flexibly incorporate relation information and some evidence that 2) the model outperforms some multi-object time-series forecasting approaches from the literature.\n\nIs this a complete piece of work or work in progress?\nThe submission appears complete.\n\nAre the authors careful and honest about evaluating both the strengths and weaknesses of their work?\nYes, and they included discussion of when the model is under-performing w.r.t. other methods (at longer time horizons).\n\n>Clarity:\nIs the submission clearly written?\nAs mentioned above, many of the numerous design decisions are not motivated well. Essentially, the paper is not very self-contained.\n\nIs it well organized?\nYes.\n\nDoes it adequately inform the reader?\nNo -- the paper assumes a lot of prior knowledge, which most potential readers would likely possess. This issue would be less significant if the main innovation was clearer.\n\n>Significance:\nAre the results important?\nIt's difficult to tell, because it's not very clear why, of the many cited works, the specific ones chosen for comparison were used. More motivation for the comparison is needed.\n\nAre others (researchers or practitioners) likely to use the ideas or build on them?\nPossibly, although the model is complicated, which makes it less likely.\n\nDoes the submission address a difficult task in a better way than previous work?\nIt's plausible that the model performs better, but it's unclear how widely-applicable the method is. The model appears to be restricted to lower dimensional settings.\n\nDoes it advance the state of the art in a demonstrable way?\nIt's plausible, but I am not sure, as mentioned above.\n\nDoes it provide unique data, unique conclusions about existing data, or a unique theoretical or experimental approach?\nNo\n\nAdditional feedback\n---\nSec 5.2 describe what \"OOB rate of the rollouts\" means.\n\nSec 3.3 Coupling layer equation is missing a final parenthesis.\n\n3.3 \"complicate the learning a lot.\" -> \"significantly complicate the learning.\" (original is subjectively too informal)\n\nSec 5.2 \"like to dig into it in future work\" -> \"like to explore it in future work\" (original is subjectively too informal)\n\nRecommendation\n---\nIt is unclear what the main key contributions are, outside of usually performing better than some related methods. This lack of clarity, along with a complicated design that is not very well-motivated, make me lean towards rejection. An improved version of this paper would make clear what the precise contributions are, explain how the related work did not achieve them, and why the specific comparisons were adopted.\n", "extra_scores": {"experience_assessment": "I have read many papers in this area.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper.", "review_assessment:_checking_correctness_of_derivations_and_theory": "N/A"}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper proposed what is termed Relational State Space Model (R-SSM) that can be used for modeling interacting time-series data. The model essentially consists of a set of (nonlinear) state space models whose states are jointly evolved in a way that take into account a known interaction structure between them (the relational part, even though technically it is just a coupling structure -- the term relational structure in the past has been used for models with objects and classes, for example see the difference between "coupled HMM" vs "relational HMM"). The authors also proposed a graph normalizing flow operation to model the joint state evolution. The main weakness of the paper is in the complexity of the model. However, from a modeling point of view, R-SSM seems suitable in situation when the interaction structure is known, and this is demonstrated in the experimental results when comparing against the baselines. 
Decision: Accept (Poster)