Original: R ELATIONAL S TATE -S PACE M ODEL FOR S TOCHASTIC M ULTI -O BJECT S YSTEMS
Revised: R ELATIONAL S TATE -S PACE M ODEL FOR S TOCHASTIC M ULTI -O BJECT S YSTEMS

Original: Anonymous authors

Original: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: Real-world dynamical systems often consist of multiple stochastic subsystems that interact with each other.
Revised: Real-world dynamical systems often consist of multiple stochastic subsystems that interact with each other.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: By letting GNNs cooperate with SSM
Revised: By letting GNNs cooperate with SSM,

Original: R-SSM provides a ﬂexible way to incorporate relational information into the modeling of multi-object dynamics
Revised: R-SSM provides a ﬂexible way to incorporate relational information into the modeling of multi-object dynamics.

Original: We further suggest augmenting the model with normalizing ﬂows instantiated for vertex-indexed random variables and propose two auxiliary contrastive objectives to facilitate the learning
Revised: We further suggest augmenting the model with normalizing ﬂows instantiated for vertex-indexed random variables and propose two auxiliary contrastive objectives to facilitate the learning.

Original: The utility of R-SSM is empirically evaluated on synthetic and real time series datasets.
Revised: The utility of R-SSM is empirically evaluated on synthetic and real time series datasets.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Many real-world dynamical systems can be decomposed into smaller interacting subsystems if we take a ﬁne-grained view.
Revised: Many real-world dynamical systems can be decomposed into smaller interacting subsystems if we take a ﬁne-grained view.

Original: For example, the trajectories of coupled particles are co-determined by perparticle physical properties (e.g., mass and velocity) and their physical interactions (e.g., gravity); trafﬁc ﬂow can be viewed as the coevolution of a large number of vehicle dynamics.
Revised: For example, the trajectories of coupled particles are co-determined by perparticle physical properties (e.g., mass and velocity) and their physical interactions (e.g., gravity); trafﬁc ﬂow can be viewed as the coevolution of a large number of vehicle dynamics.

Original: Models that are able to better capture the complex behavior of such multi-object systems are of wide interest to various communities, e.g., physics, ecology, biology, geoscience, and ﬁnance.
Revised: Models that are able to better capture the complex behavior of such multi-object systems are of wide interest to various communities, e.g., physics, ecology, biology, geoscience, and ﬁnance.

Original: State-space models (SSMs) are a wide class of sequential latent variable models (LVMs) that serve as workhorses for the analysis of dynamical systems and sequence data.
Revised: State-space models (SSMs) are a wide class of sequential latent variable models (LVMs) that serve as workhorses for the analysis of dynamical systems and sequence data.

Original: Although SSMs are traditionally designed under the guidance of domain-speciﬁc knowledge or tractability consideration, recently introduced deep SSMs (Fraccaro, 2018) use neural networks (NNs) to parameterize ﬂexible state transitions and emissions, achieving much higher expressivity.
Revised: Although SSMs are traditionally designed under the guidance of domain-speciﬁc knowledge or tractability consideration, recently introduced deep SSMs (Fraccaro, 2018) use neural networks (NNs) to parameterize ﬂexible state transitions and emissions, achieving much higher expressivity.

Original: To develop deep SSMs for multi-object systems, graph neural networks (GNNs) emerge to be a promising choice, as they have been shown to be fundamental NN building blocks that can impose relational inductive bias explicitly and model complex interactions effectively (Battaglia et al., 2018).
Revised: To develop deep SSMs for multi-object systems, graph neural networks (GNNs) emerge to be a promising choice, as they have been shown to be fundamental NN building blocks that can impose relational inductive bias explicitly and model complex interactions effectively (Battaglia et al., 2018).

Original: Recent works that advocate GNNs for modeling multi-object dynamics mostly make use of GNNs in an autoregressive (AR) fashion.
Revised: Recent works that advocate GNNs for modeling multi-object dynamics mostly make use of GNNs in an autoregressive (AR) fashion.

Original: AR models based on recurrent (G)NNs can be viewed as special instantiations of SSMs in which the state transitions are restricted to being deterministic (Fraccaro, 2018, Section 4.2).
Revised: AR models based on recurrent (G)NNs can be viewed as special instantiations of SSMs in which the state transitions are restricted to being deterministic (Fraccaro, 2018, Section 4.2).

Original: Despite their simplicity, it has been pointed out that their modeling capability is bottlenecked by the deterministic state transitions (Chung et al., 2015; Fraccaro et al., 2016) and the oversimpliﬁed observation distributions (Yang et al., 2018).
Revised: Despite their simplicity, it has been pointed out that their modeling capability is bottlenecked by the deterministic state transitions (Chung et al., 2015; Fraccaro et al., 2016) and the oversimpliﬁed observation distributions (Yang et al., 2018).

Original: In this study, we make the following contributions: (i)
Revised: In this study, we make the following contributions: (i)

Original: We propose the relational state-space model (R-SSM), a novel hierarchical deep SSM that simulates the stochastic state transitions of interacting objects with GNNs, extending GNN-based dynamics modeling to challenging stochastic multi-object systems.
Revised: We propose the relational state-space model (R-SSM), a novel hierarchical deep SSM that simulates the stochastic state transitions of interacting objects with GNNs, extending GNN-based dynamics modeling to challenging stochastic multi-object systems.

Original: (ii) We suggest using the graph normalizing ﬂow (GNF) to construct expressive joint state distributions for R-SSM, further enhancing its ability to capture the joint evolutions of correlated stochastic subsystems.
Revised: (ii) We suggest using the graph normalizing ﬂow (GNF) to construct expressive joint state distributions for R-SSM, further enhancing its ability to capture the joint evolutions of correlated stochastic subsystems.

Original: (iii)
Revised: (iii)

Original: We develop structured posterior approximation to learn R-SSM using variational inference and introduce two auxiliary training objectives to facilitate the learning.
Revised: We develop structured posterior approximation to learn R-SSM using variational inference and introduce two auxiliary training objectives to facilitate the learning.

Original: Our experiments on synthetic and real-world time series datasets show that R-SSM achieves competitive test likelihood and good prediction performance in comparison to GNN-based AR models and other sequential LVMs.
Revised: Our experiments on synthetic and real-world time series datasets show that R-SSM achieves competitive test likelihood and good prediction performance in comparison to GNN-based AR models and other sequential LVMs.

Original: The remainder of this paper is organized as follows: Section 2 brieﬂy reviews neccesary preliminaries.
Revised: The remainder of this paper is organized as follows: Section 2 brieﬂy reviews neccesary preliminaries.

Original: Section 3 introduces R-SSM formally and presents the methods to learn R-SSM from observations.
Revised: Section 3 introduces R-SSM formally and presents the methods to learn R-SSM from observations.

Original: Related work is summarized in Section 4 and experimental evaluation is presented in Section 5.
Revised: Related work is summarized in Section 4 and experimental evaluation is presented in Section 5.

Original: We conclude the paper in Section 6.
Revised: We conclude the paper in Section 6.

Original:  P RELIMINARIES
Revised:  P RELIMINARIES

Original: In this work, an attributed directed graph is given by a 4-tuple: G = ( V ; E ; V ; E ) , where V =
Revised: In this work, an attributed directed graph is given by a 4-tuple: G = ( V , E , V , E ) , where V =

Original: [ N ] := f 1 ; : : : ; N g is the set of vertices, E (cid:18)
Revised: [ N ] := { 1 , . . , N } is the set of vertices, E ⊆

Original: [ N ] (cid:2) [ N ] is the set of edges, V 2 R N (cid:2)
Revised: [ N ] × [ N ] is the set of edges, V ∈ R N ×

Original: d v is a matrix of static vertex attributes, and E 2 R N (cid:2) N (cid:2) d e is a sparse tensor storing the static edge attributes.
Revised: d v is a matrix of static vertex attributes, and E ∈ R N × N × d e is a sparse tensor storing the static edge attributes.

Original: The set of direct predecessors of vertex i is notated as N (cid:0) i=
Revised: The set of direct predecessors of vertex i is notated as N − i

Original: f p j ( p; i ) 2 Eg .
Revised: = { p | ( p, i ) ∈ E} .

Original: We use the notation x i to refer to the i -th row of matrix X and write x ij to indicate the ( i; j ) -th entry of tensor X (if the corresponding matrix or tensor appears in the context).
Revised: We use the notation x i to refer to the i -th row of matrix X and write x ij to indicate the ( i, j ) -th entry of tensor X (if the corresponding matrix or tensor appears in the context).

Original: For sequences, we write x (cid:20) t = x 1: t :=
Revised: For sequences, we write x ≤ t = x 1: t :=

Original: ( x 1 ; : : : ; x t ) and switch to x ( i ) t for referring to the i -th row of matrix X t .
Revised: ( x 1 , . . . , x t ) and switch to x ( i ) t for referring to the i -th row of matrix X t .

Original: 2. G RAPH N EURAL N ETWORKS
Revised: 2. G RAPH N EURAL N ETWORKS

Original: GNNs are a class of neural networks developed to process graph-structured data and support relational reasoning.
Revised: GNNs are a class of neural networks developed to process graph-structured data and support relational reasoning.

Original: Here we focus on vertex-centric GNNs that iteratively update the vertex representations of a graph G while being equivariant (Maron et al., 2019) under vertex relabeling.
Revised: Here we focus on vertex-centric GNNs that iteratively update the vertex representations of a graph G while being equivariant (Maron et al., 2019) under vertex relabeling.

Original: Let
Revised: Let

Original: H 2 R N (cid:2)
Revised: H ∈ R N ×

Original: d be a matrix of vertex representations, in which the i -th row h
Revised: d be a matrix of vertex representations, in which the i -th row h

Original: i [-
-]
Revised: i [-
-]∈

Original: R d is the vectorized representation attached to vertex i .
Revised: R d is the vectorized representation attached to vertex i .

Original: 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 2 R d g as input and returns new vertex representations H ′2 R N (cid:2)
Revised: 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 ∈ R d g as input and returns new vertex representations H ′

Original:  d ′ as output, i.e., H ′ =
Revised: ∈ R N × d ′ as output, i.e., H ′ =

Original: GNN ( G ; g ; H ) .
Revised: GNN ( G , g , H ) .

Original: When updating the representation of vertex i from h i to h
Revised: When updating the representation of vertex i from h i to h

Original: ′ i , a GNN takes the representations of other nearby vertices into consideration.
Revised: ′ i , a GNN takes the representations of other nearby vertices into consideration.

Original: Popular GNN variants achieve this through a multi-round message passing paradigm, in which the vertices repeatedly send messages to their neighbors, aggregate the messeages they received, and update their own representations accordingly.
Revised: Popular GNN variants achieve this through a multi-round message passing paradigm, in which the vertices repeatedly send messages to their neighbors, aggregate the messeages they received, and update their own representations accordingly.

Original: Formally, the operations performed by a basic block of a message-passing GNN are deﬁned as follows:
Revised: Formally, the operations performed by a basic block of a message-passing GNN are deﬁned as follows:

Original: [Equation]
Revised: [Equation]

Original: Throughout this work, we implement Equations (1) and (2) by adopting a multi-head attention mechanism similar to Vaswani et al.
Revised: Throughout this work, we implement Equations (1) and (2) by adopting a multi-head attention mechanism similar to Vaswani et al.

Original: (2017) and Velikovi et al.
Revised: (2017) and Velikovi et al.

Original: For Equation (3), we use either a
Revised: For Equation (3), we use either a

Original: RNN cell or a residual block (He et al., 2016), depending on whether the inputs to GNN are RNN states or not.
Revised: RNN cell or a residual block (He et al., 2016), depending on whether the inputs to GNN are RNN states or not.

Original: We write such a block as H ′ =
Revised: We write such a block as H ′ =

Original: MHA ( G ; g ; H ) and give its detailed implementation in the Appendix.
Revised: MHA ( G , g , H ) and give its detailed implementation in the Appendix.

Original: A GNN simply stacks L separately-parameterized MHA blocks and iteratively computes H
Revised: A GNN simply stacks L separately-parameterized MHA blocks and iteratively computes H

Original: =: H (0) ; : : : ; H ( L ) =:
Revised: =: H (0) , . . . , H ( L ) =:

Original: H ′ , in which H ( l ) =
Revised: H ′ , in which H ( l ) =

Original: MHA ( G ; g ; H ( l (cid:0) 1) ) for l = 1 ; : : : ; L .
Revised: MHA ( G , g , H ( l − 1) ) for l = 1 , . . . , L .

Original: We write this construction as H ′ =
Revised: We write this construction as H ′ =

Original: GNN ( G ; g ; H ) and treat it as a black box to avoid notational clutter.
Revised: GNN ( G , g , H ) and treat it as a black box to avoid notational clutter.

Original: 2. S TATE -S PACE M ODELS
Revised: 2. S TATE -S PACE M ODELS

Original: State-space models are widely applied to analyze dynamical systems whose true states are not directly observable.
Revised: State-space models are widely applied to analyze dynamical systems whose true states are not directly observable.

Original: Formally, an SSM assumes the dynamical system follows a latent state process f z t g t (cid:21) 1 , which possibly depends on exogenous inputs f u t g t (cid:21) 1 .
Revised: Formally, an SSM assumes the dynamical system follows a latent state process { z t } t ≥ 1 , which possibly depends on exogenous inputs { u t } t ≥ 1 .

Original: Parameterized by some (unknown) static parameter (cid:18) , the latent state process is characterized by an initial density z 1 (cid:24) (cid:25) (cid:18) ( (cid:1)j u 1 ) and a transition density
Revised: Parameterized by some (unknown) static parameter θ , the latent state process is characterized by an initial density z 1 ∼ π θ ( ·| u 1 ) and a transition density

Original: z t +1 (cid:24) f (cid:18) ( (cid:1)j z (cid:20) t ; u (cid:20) t +1 ) .
Revised: z t +1 ∼ f θ ( ·| z ≤ t , u ≤ t +1 ) .

Original: Moreover, at each time step, some noisy measurements of the latent state are observed through an observation density:x t (cid:24) g (cid:18) ( (cid:1)j z (cid:20) t ; u (cid:20) t ) .
Revised: Moreover, at each time step, some noisy measurements of the latent state are observed through an observation density:

Original: The joint density of x 1: T and z 1: T factors as: p ( x 1: T ; z 1: T j u 1: T ) = (cid:25) (cid:18) ( z 1 j u 1 ) ∏ Tt =2 f (cid:18) ( z t j z <t ; u (cid:20) t ) ∏ Tt =1 g (cid:18) ( x t j z (cid:20) t ; u (cid:20) t ) .
Revised: H t ˜

Original: The superior expressiveness of SSMs can be seen from the fact that the marginal predictive distribution p ( x t j x <t ; u (cid:20) t ) = ∫ p ( x t ; z (cid:20) t j x <t ; u (cid:20) t ) d z (cid:20) t can be far more complex than unimodal distributions and their ﬁnite mixtures that are common in AR models. Recently developed deep H t ~ H t +
Revised: H t +

Original: H t + z gt z gt + h gt h gt
Revised: H t + z gt z gt + h gt h gt

Original: +
Revised: +

Original: Z t
Revised: Z t

Original: Z t +
Revised: Z t +

Original: X t
Revised: X t

Original: X t + (a) Generative model.
Revised: X t + (a) Generative model.

Original: Z
Revised: Z

Original: Z
Revised: Z

Original: Z(cid:1)(cid:1)(cid:1)
Revised: Z···

Original: Z T
Revised: Z T

Original: X
Revised: X

Original: X
Revised: X

Original: X(cid:1)(cid:1)(cid:1)
Revised: X···

Original: X T (b) Collapsed generative model.
Revised: X T (b) Collapsed generative model.

Original: H t ~ H t +
Revised: H t ˜ H t +

Original: H t + z gt z gt + h gt h gt
Revised: H t + z gt z gt + h gt h gt

Original: +
Revised: +

Original: Z t
Revised: Z t

Original: Z t +
Revised: Z t +

Original: B t
Revised: B t

Original: B t +1 X t
Revised: B t +1 X t

Original: X t + (c) Inference model. 
Revised: X t + (c) Inference model. x t

Revised: [Equation]

Original: SSMs use RNNs to compress z (cid:20) t (and u (cid:20) t ) into ﬁxed-size vectors to achieve tractability.
Revised: The joint density of x 1: T and z 1: T factors as: p [Equation] 1: T [Equation]

Revised: z 1:

Revised: [Equation]

Revised: [Equation]

Revised: The superior expressiveness of SSMs can be seen from the fact that the marginal predictive distribution p ( x t | x <t , u ≤ t ) =

Revised: R p ( x t , z ≤ t | x <t , u ≤ t ) d z ≤ t can be far more complex than unimodal distributions and their ﬁnite mixtures that are common in AR models.

Revised: Recently developed deep

Revised: SSMs use RNNs to compress z ≤ t (and u ≤ t ) into ﬁxed-size vectors to achieve tractability.

Original: As shown in next section, R-SSM can be viewed as enabling multiple individual deep SSMs to communicate.
Revised: As shown in next section, R-SSM can be viewed as enabling multiple individual deep SSMs to communicate.

Original: 2. N ORMALIZING F LOWS
Revised: 2. N ORMALIZING F LOWS

Original: Normalizing ﬂows (Rezende & Mohamed, 2015) are invertible transformations that have the capability to transform a simple probability density into a complex one (or vice versa).
Revised: Normalizing ﬂows (Rezende & Mohamed, 2015) are invertible transformations that have the capability to transform a simple probability density into a complex one (or vice versa).

Original: Given two domains X (cid:18) R D and Y (cid:18) R D , let f : X ! Y be an invertible mapping with inverse f (cid:0) 1 .
Revised: Given two domains X ⊆ R D and Y ⊆ R D , let f : X → Y be an invertible mapping with inverse f − 1 .

Original: Applying f to a random variable z 2 X with density p ( z ) , by the change of variables rule, the resulting random variable z ′ =
Revised: Applying f to a random variable z ∈ X with density p ( z ) , by the change of variables rule, the resulting random variable z ′ =

Original: f ( z ) 2 Y will have a density:
Revised: f ( z ) ∈ Y will have a density:

Original: [Equation]
Revised: [Equation]

Original: [Table]

Original: A series of invertible mappings with cheap-to-evaluate determinants can be chained together to achieve complex transformations while retaining efﬁcient density calculation.
Revised: A series of invertible mappings with cheap-to-evaluate determinants can be chained together to achieve complex transformations while retaining efﬁcient density calculation.

Original: This provides a powerful way to construct expressive distributions.
Revised: This provides a powerful way to construct expressive distributions.

Original:  R ELATIONAL S TATE -S PACE M ODEL
Revised:  R ELATIONAL S TATE -S PACE M ODEL

Original: Suppose there is a dynamical system that consists of multiple interacting objects, and observing this system at a speciﬁc time is accomplished by acquiring measurements from every individual object simultaneously.
Revised: Suppose there is a dynamical system that consists of multiple interacting objects, and observing this system at a speciﬁc time is accomplished by acquiring measurements from every individual object simultaneously.

Original: We further assume these objects are homogeneous, i.e., they share the same measurement model, and leave systems whose constituents are nonhomogeneous for future work.
Revised: We further assume these objects are homogeneous, i.e., they share the same measurement model, and leave systems whose constituents are nonhomogeneous for future work.

Original: To generatively model a time-ordered series of observations collected from this system, the straightforward approach that builds an individual SSM for each object is usually unsatisfactory, as it simply assumes the state of each object evolves independently and ignores the interactions between objects.
Revised: To generatively model a time-ordered series of observations collected from this system, the straightforward approach that builds an individual SSM for each object is usually unsatisfactory, as it simply assumes the state of each object evolves independently and ignores the interactions between objects.

Original: To break such an independence assumption, our main idea is to let multiple individual SSMs interact through GNNs, which are expected to capture the joint state transitions of correlated objects well.
Revised: To break such an independence assumption, our main idea is to let multiple individual SSMs interact through GNNs, which are expected to capture the joint state transitions of correlated objects well.

Original: 3. G ENERATIVE MODEL
Revised: 3. G ENERATIVE MODEL

Original: Given the observations for a multi-object dynamical system, our model further assumes its interaction structure is known as prior knowledge.
Revised: Given the observations for a multi-object dynamical system, our model further assumes its interaction structure is known as prior knowledge.

Original: The interaction structure is provided as a directed graph, in which each object corresponds to a vertex, and a directed edge indicates that the state of its head is likely to be affected by its tail.
Revised: The interaction structure is provided as a directed graph, in which each object corresponds to a vertex, and a directed edge indicates that the state of its head is likely to be affected by its tail.

Original: In situations where such graph structure is not available, a complete graph can be speciﬁed.
Revised: In situations where such graph structure is not available, a complete graph can be speciﬁed.

Original: However, to model dynamical systems comprising a large number of objects, it is often beneﬁcial to explicitly specify sparse graph structures, because they impose stronger relational inductive bias and help save the computational cost.
Revised: However, to model dynamical systems comprising a large number of objects, it is often beneﬁcial to explicitly specify sparse graph structures, because they impose stronger relational inductive bias and help save the computational cost.

Original: A relational state-space model assumes a set of correlated dynamical subsystems evolve jointly under the coordination of graph neural networks.
Revised: A relational state-space model assumes a set of correlated dynamical subsystems evolve jointly under the coordination of graph neural networks.

Original: Formally, given a graph G =
Revised: Formally, given a graph G =

Original: ( V ; E ; V ; E ) , in which an edge ( i; j ) 2 E indicates that the state of vertex j may be affected by vertex i .
Revised: ( V , E , V , E ) , in which an edge ( i, j ) ∈ E indicates that the state of vertex j may be affected by vertex i .

Original: Let u ( i ) t 2 R d u and x ( i ) t 2 Rd x be the input and observation for vertex i at time step t , respectively. For T steps, we introduce a set of unobserved random variables
Revised: Let u ( i ) t ∈ R d u and x ( i ) t ∈ R

Original: f
Revised: d

Original:  z ( i ) 1: T g Ni =1 , in which z ( i ) t [Equation] represents the latent state of vertex i at time step t .
Revised: x be the input and observation for vertex i at time step t , respectively. For T steps, we introduce a set of unobserved random variables { z ( i ) 1: T } Ni =1 , in which z ( i ) t ∈ R d z represents the latent state of vertex i at time step t .

Original: Furthermore, we introduce a global latent variable z gt 2 R d g for each time step to represent the global state shared by all vertices.
Revised: Furthermore, we introduce a global latent variable z gt ∈ R d g for each time step to represent the global state shared by all vertices.

Original: Conditioning on the graph and exogenous inputs, an R-SSM factorizes the joint density of observations and latent states as follows:
Revised: Conditioning on the graph and exogenous inputs, an R-SSM factorizes the joint density of observations and latent states as follows:

Original: [Equation]
Revised: [Equation]

Original: For notational simplicity, we switch to the matrix notation Z t =
Revised: For notational simplicity, we switch to the matrix notation Z t =

Original: [Equation]
Revised: [Equation]

Original: ) t ; z (2) t ; : : : ; z ( N ) t ] ⊤ from now on.
Revised: [Equation] ⊤ from now on.

Original: The joint transition density f (cid:18) is further factorized as a product of global transition density f g(cid:18) and local transition density f ⋆(cid:18) , i.e., f (cid:18) ( Z t ; z gt j : : : )
Revised: The joint transition density f θ is further factorized as a product of global transition density f gθ and local transition density f ⋆θ , i.e., f θ ( Z t , z gt | . . . )

Original: =
Revised: =

Original: f g(cid:18) ( z gt j : : : ) f ⋆(cid:18) ( Z t j z gt ; : : : ) .
Revised: f gθ ( z gt | . . . ) f ⋆θ ( Z t | z gt , . . . ) .

Original: To instantiate these conditional distributions, a GNN accompanied by RNN cells is adopted to recurrently compress the past dependencies at each time step into ﬁxed-size context vectors.
Revised: To instantiate these conditional distributions, a GNN accompanied by RNN cells is adopted to recurrently compress the past dependencies at each time step into ﬁxed-size context vectors.

Original: Speciﬁcally, the observations are assumed to be generated from following process:
Revised: Speciﬁcally, the observations are assumed to be generated from following process:

Original: [Equation]
Revised: [Equation]

Original: for i = 1 ; : : : ; N and t = 1 ; : : : ; T , [-where h ( i ) 
-]
Revised: · (cid:12)(cid:12) h gt (cid:1) H t = GNN θ (cid:0) G , [-where h ( i ) 
-]z gt , ˜

Original:  =  h (cid:3) 

Revised: H t (cid:1) , Z t ∼ f ⋆θ (cid:0) · (cid:12)(cid:12) H t (cid:1) x ( i ) t ∼ g θ (cid:0) · (cid:12)(cid:12) z ( i ) t , h ( i ) t , z gt , h gt (cid:1) for i = 1 , . . . , N and t = 1 , . . , T , where h ( i )

Original:  and z ( i ) 

Revised: = h ∗ 0 and z ( i )

Original: = z (cid:3) 0 .
Revised: = z ∗ 0 .

Original: Here h g 0 , z g 0 , h (cid:3) 0 and z (cid:3) 0 are learnable initial states.
Revised: Here h g 0 , z g 0 , h ∗ 0 and z ∗ 0 are learnable initial states.

Original: The READOUT function aggregates the context vectors of all vertices into a global context vector in a permutation-invariant manner.
Revised: The READOUT function aggregates the context vectors of all vertices into a global context vector in a permutation-invariant manner.

Original: The global transition density f g(cid:18) is speciﬁed to be a diagonal Gaussian distribution whose mean and variance are parameterized by the output of a multilayer perceptron (MLP), and the local transition density f ⋆(cid:18) will be discussed later.
Revised: The global transition density f gθ is speciﬁed to be a diagonal Gaussian distribution whose mean and variance are parameterized by the output of a multilayer perceptron (MLP), and the local transition density f ⋆θ will be discussed later.

Original: The local observation distribution g (cid:18) can be freely selected in line with the data, and in our experiments it is either a Gaussian distribution or a mixture of logistic distributions parameterized by MLPs.
Revised: The local observation distribution g θ can be freely selected in line with the data, and in our experiments it is either a Gaussian distribution or a mixture of logistic distributions parameterized by MLPs.

Original: The graphical structure of two consecutive steps of the generating process is illustrated in Figure 1a.
Revised: The graphical structure of two consecutive steps of the generating process is illustrated in Figure 1a.

Original: An intuitive way to think about this generative model is to note that the N + 1 latent state processes 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.
Revised: An intuitive way to think about this generative model is to note that the N + 1 latent state processes 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.

Original: 3. L EARNING AND INFERENCE
Revised: 3. L EARNING AND INFERENCE

Original: As illustrated in Figure 1b, writing Z t = ( z gt ; Z t ) and suppressing the dependencies on the graph G and exogenous inputs U 1: T , an R-SSM can be interpreted as an ordinary SSM in which the entire graph evolves as a whole, i.e., the joint density of latent states and observations factors as: p ( X 1: T ; Z 1: T ) = 
Revised: As illustrated in Figure 1b, writing Z t = ( z gt , Z t ) and suppressing the dependencies on the graph G and exogenous inputs U 1: T , an R-SSM can be interpreted as an ordinary SSM in which the entire graph evolves as a whole, i.e., the joint density of latent states and observations factors as: p ( X 1: T , Z 1: T ) = Q Tt = [Equation] p θ

Original: ∏ Tt =1 p (cid:18) ( Z t jZ <t ) p (cid:18) ( X t jZ t ) .
Revised: [Equation] | [Equation] <t ) p θ ( X t |Z t ) .

Original: Given observations X 1: T , we are interested in learning unknown parameters (cid:18) and inferring unobserved states Z 1: T .
Revised: Given observations X 1: T , we are interested in learning unknown parameters θ and inferring unobserved states Z 1: T .

Original: For the learning task we wish to maximize the marginal likelihood p (cid:18)
Revised: For the learning task we wish to maximize the marginal likelihood p θ ( X 1: T ) =

Original: ( X 1: T ) = ∫ p (cid:18) ( X 1: T ; Z 1: T )
Revised: R p θ ( X 1: T , Z 1: T )

Original: d Z 1: T , but in our case the integral is intractable.
Revised: d Z 1: T , but in our case the integral is intractable.

Original: We adopt a recently developed variational inference (VI) approach called variational sequential Monte Carlo (VSMC)
Revised: We adopt a recently developed variational inference (VI) approach called variational sequential Monte Carlo (VSMC)

Original: (Maddison et al., 2017; Naesseth et al., 2018; Le et al., 2018), which maximizes a variational lower bound on the log marginal likelihood instead and learns the proposal distributions for the inference task simultaneously.
Revised: (Maddison et al., 2017; Naesseth et al., 2018; Le et al., 2018), which maximizes a variational lower bound on the log marginal likelihood instead and learns the proposal distributions for the inference task simultaneously.

Original: Given a sequence of proposal distributions f q ϕ ( Z t jZ <t ; X (cid:20) t ) g Tt =1 parameterized by ϕ , running the sequential Monte Carlo (SMC) algorithm with K particles yields an unbiased marginal likelihood estimator ^ p (cid:18);ϕ;K ( X 1:T ) = ∏
Revised: Given a sequence of proposal distributions { q ϕ ( Z t |Z <t , X ≤ t ) } Tt =1 parameterized by ϕ , running the sequential Monte Carlo (SMC) algorithm with K particles yields an unbiased marginal likelihood estimator ˆ p θ,ϕ,K ( X 1:

Original: [Equation]
Revised: [Equation]

Original: [Equation] ∑ Kk =
Revised: ) =

Original: [Equation]
Revised: [Equation]

Original: kt ] , where w kt is the unnormalized importance weight of particle k at time t .
Revised: , where w kt is the unnormalized importance weight of particle k at time t .

Original: The variational lower bound is obtained by applying the Jensen’s inequality:
Revised: The variational lower bound is obtained by applying the Jensen’s inequality:

Original: [Equation]
Revised: [Equation]

Original: Assuming the proposal distributions are reparameterizable (Kingma & Welling, 2014), we use the biased gradient estimator ∇L SMC K ( (cid:18); ϕ ) (cid:25) E [ ∇ log ^ p (cid:18);ϕ;K ( X 1: T )] to maximize L SMC K .
Revised: Assuming the proposal distributions are reparameterizable (Kingma & Welling, 2014), we use the biased gradient estimator ∇L SMC K ( θ, ϕ ) ≈ E [ ∇ log ˆ p θ,ϕ,K ( X 1: T )] to maximize L SMC K .

Original: Proposal design.
Revised: Proposal design.

Original: We make the proposal for Z t depend on the information up to time t and share some parameters with the generative model.
Revised: We make the proposal for Z t depend on the information up to time t and share some parameters with the generative model.

Original: We also choose to factorize q ϕ ( Z t j : : : )
Revised: We also choose to factorize q ϕ ( Z t | . . . )

Original: =
Revised: =

Original: r gϕ ( z gt j : : : ) r ⋆ϕ ( Z t j z gt ; : : : ) .
Revised: r gϕ ( z gt | . . . ) r ⋆ϕ ( Z t | z gt , . . . ) .

Original: The proposal distributions for all time steps are structured as follows:
Revised: The proposal distributions for all time steps are structured as follows:

Original: [Equation]
Revised: [Equation]

Original: for i = 1 ; : : : ; N and t = 1 ; : : : ; T , where h gt and H 1: T are computed using the relevant parts of the generative model.
Revised: for i = 1 , . . . , N and t = 1 , . . , T , where h gt and H 1: T are computed using the relevant parts of the generative model.

Original: r gϕ is speciﬁed to be a diagonal Gaussian parameterized by an MLP, and r ⋆ϕ will be discussed soon.
Revised: r gϕ is speciﬁed to be a diagonal Gaussian parameterized by an MLP, and r ⋆ϕ will be discussed soon.

Original: Here B t can be interpreted as a belief state ( Gregor et al., 2019), which summarizes past observations X (cid:20) t (and inputs U (cid:20) t ) deterministically.
Revised: Here B t can be interpreted as a belief state (Gregor et al., 2019), which summarizes past observations X ≤ t (and inputs U ≤ t ) deterministically.

Original: The graphical structure of this proposal design is shown in Figure 1c, and the detailed VSMC implementation using this proposal is given in Appendix A.4.
Revised: The graphical structure of this proposal design is shown in Figure 1c, and the detailed VSMC implementation using this proposal is given in Appendix A.4.

Original: 3. G RAPH N ORMALIZING F LOW
Revised: 3. G RAPH N ORMALIZING F LOW

Original: The local transition density f ⋆(cid:18) (
Revised: The local transition density f ⋆θ (cid:0)

Original: Z t (cid:12)(cid:12) : : : ) and the local proposal density r⋆ϕ
Revised: Z t (cid:12)(cid:12) . . . (cid:1) and the local proposal density r

Original:  ( Z t j : : : ) may be parameterized in several ways.
Revised: ⋆ϕ ( Z t | . . . ) may be parameterized in several ways.

Original: One simple and efﬁcient starting point is (block-)diagonal Gaussian distribution: f ⋆(cid:18) (
Revised: One simple and efﬁcient starting point is (block-)diagonal Gaussian distribution: f ⋆θ (cid:0)

Original: Z t (cid:12)(cid:12) : : : )
Revised: Z t (cid:12)(cid:12) . . .

Original:  =
Revised: (cid:1) =

Original: ∏ Ni =1 N ( z ( i ) t j : : : ) , which assumes that the object states are conditionally independent, i.e., the joint state distribution is completely factorized over objects.
Revised: Q Ni =1 N ( z ( i ) t | . . . ) , which assumes that the object states are conditionally independent, i.e., the joint state distribution is completely factorized over objects.

Original: We believe that such an independence assumption is an oversimpliﬁcation for situations where the joint state evolution is multimodal and highly correlated.
Revised: We believe that such an independence assumption is an oversimpliﬁcation for situations where the joint state evolution is multimodal and highly correlated.

Original: One possible way to introduce inter-object dependencies is modeling joint state distributions as Markov random ﬁelds (MRFs) (Naesseth et al., 2019), but this will signiﬁcantly complicate the learning.
Revised: One possible way to introduce inter-object dependencies is modeling joint state distributions as Markov random ﬁelds (MRFs) (Naesseth et al., 2019), but this will signiﬁcantly complicate the learning.

Original: Here we introduce the Graph Normalizing Flow (GNF) 1 , which adapts Glow (Kingma & Dhariwal, 2018) to graph settings and enables us to build expressive joint distributions for correlated random variables indexed by graph nodes.
Revised: Here we introduce the Graph Normalizing Flow (GNF) 1 , which adapts Glow (Kingma & Dhariwal, 2018) to graph settings and enables us to build expressive joint distributions for correlated random variables indexed by graph nodes.

Original: As described earlier, the key ingredient for a ﬂow is a series invertible mappings that are iteratively applied to the samples of a base distribution.
Revised: As described earlier, the key ingredient for a ﬂow is a series invertible mappings that are iteratively applied to the samples of a base distribution.

Original: Now we are interested in the case where the samples are vertex states Z t , and thus the invertible mappings should be further constrained to be equivariant under vertex relabeling.
Revised: Now we are interested in the case where the samples are vertex states Z t , and thus the invertible mappings should be further constrained to be equivariant under vertex relabeling.

Original: This rules out popular autoregressive ﬂows, e.g., IAF (Kingma et al., 2016) and MAF (Papamakarios et al., 2017).
Revised: This rules out popular autoregressive ﬂows, e.g., IAF (Kingma et al., 2016) and MAF (Papamakarios et al., 2017).

Original: Our GNF is built upon the coupling layer introduced in Dinh et al.
Revised: Our GNF is built upon the coupling layer introduced in Dinh et al.

Original: (2017), which provides a ﬂexible framework to construct efﬁcient invertible mappings.
Revised: (2017), which provides a ﬂexible framework to construct efﬁcient invertible mappings.

Original: A GNF coupling layer splits the input Z 2 R N (cid:2) D into two parts, Z a 2 R N (cid:2) d and Z b 2 R N (cid:2)
Revised: A GNF coupling layer splits the input Z ∈ R N × D into two parts, Z a ∈ R N × d and Z b ∈ R N ×

Original: ( D (cid:0) d ) .
Revised: ( D − d ) .

Original: The output Z ′ 2 R N (cid:2) D is formed as:
Revised: The output Z ′ ∈ R N × D is formed as:

Original: [Equation]
Revised: [Equation]

Original: where ⊙ denotes the element-wise product, and the functions s ( (cid:1) ) and t ( (cid:1) ) are speciﬁed to be GNNs to enforce the equivariance property.
Revised: where ⊙ denotes the element-wise product, and the functions s ( · ) and t ( · ) are speciﬁed to be GNNs to enforce the equivariance property.

Original: A GNF combines a coupling layer with a trainable elementwise afﬁne layer and an invertible 1 (cid:2) 1 convolution layer (Hoogeboom et al., 2019), organizing them as: Input ! Afﬁne ! Coupling ! Conv 1 (cid:2) 1 !
Revised: A GNF combines a coupling layer with a trainable elementwise afﬁne layer and an invertible 1 × 1 convolution layer (Hoogeboom et al., 2019), organizing them as: Input → Afﬁne → Coupling → Conv 1 × 1 → Output.

Original: Output.

Original: A visual illustration of this architecture is provided in Appendix A.5.
Revised: A visual illustration of this architecture is provided in Appendix A.5.

Original: 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
Revised: 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

Original: Gaussian distributions parameterized by MLPs.
Revised: Gaussian distributions parameterized by MLPs.

Original: With the message passing inside the coupling layers,
Revised: With the message passing inside the coupling layers,

Original: GNFs can transform independent noise into correlated noise and thus increase model expressivity.
Revised: GNFs can transform independent noise into correlated noise and thus increase model expressivity.

Original: The 1 (cid:2) 1 convolution layers free us from manually permuting the dimensions, and the element-wise afﬁne layers enable us to tune their initial weights to stablize training.
Revised: The 1 × 1 convolution layers free us from manually permuting the dimensions, and the element-wise afﬁne layers enable us to tune their initial weights to stablize training.

Original: 3. A UXILIARY CONTRASTIVE PREDICTION TASKS
Revised: 3. A UXILIARY CONTRASTIVE PREDICTION TASKS

Original: 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).
Revised: 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).

Original: It means that the variational posterior approximation q ϕ ( Z t jZ <t ; X (cid:20) t ) degenerate into the prior f (cid:18) ( Z t jZ <t ) in the early stage of optimization, making the training dynamics get stuck in undesirable local optima.
Revised: It means that the variational posterior approximation q ϕ ( Z t |Z <t , X ≤ t ) degenerate into the prior f θ ( Z t |Z <t ) in the early stage of optimization, making the training dynamics get stuck in undesirable local optima.

Original: Besides, we also encountered a more subtle problem inherent in likelihood-based training of deep sequential models.
Revised: Besides, we also encountered a more subtle problem inherent in likelihood-based training of deep sequential models.

Original: That is, for relatively smooth observations, the learned model tended to only capture short-term local correlations but not the interaction effects and long-term transition dynamics.
Revised: That is, for relatively smooth observations, the learned model tended to only capture short-term local correlations but not the interaction effects and long-term transition dynamics.

Original: Motivated by recent advances in unsupervised representation learning based on mutual information maximization, in particular the Contrastive Predictive Coding (CPC) approach (Oord et al., 2018), we alleviate these problems by forcing the latent states to perform two auxiliary contrastive prediction tasks.
Revised: Motivated by recent advances in unsupervised representation learning based on mutual information maximization, in particular the Contrastive Predictive Coding (CPC) approach (Oord et al., 2018), we alleviate these problems by forcing the latent states to perform two auxiliary contrastive prediction tasks.

Original: At each time step t , the future observations of each vertex i are summarized into a vector using a backward RNN: c ( i ) t = RNN  ( x ( i ) >t ) .
Revised: At each time step t , the future observations of each vertex i are summarized into a vector using a backward RNN: c ( i ) t = RNN ψ ( x ( i ) >t ) .

Original: Then we deﬁne two auxiliary CPC objectives:
Revised: Then we deﬁne two auxiliary CPC objectives:

Revised: L aux

Original: [Equation]
Revised: [Equation]

Revised: =

Revised: [Equation]

Original: where ^ z ( i ) t =
Revised: where ˆ z ( i ) t =

Original: [ z gt ; z ( i ) t ] , ^ h ( i ) t =
Revised: [ z gt , z ( i ) t ] , ˆ h ( i ) t =

Original: MLP ( ∑ j ̸ = i ^ j 2N (cid:0) i h ( j ) t ) , and Ω t;i is a set that contains c ( i ) t and some negative samples.
Revised: [Equation] = i ∧ j ∈N − i h ( j ) t ) , and Ω t,i is a set that contains c ( i ) t and some negative samples.

Original: The expectation is over negative samples and the latent states sampled from the ﬁltering distributions.
Revised: The expectation is over negative samples and the latent states sampled from the ﬁltering distributions.

Original: The positive score functions (cid:21) ; 1 and (cid:21) ; 2 are speciﬁed to be simple log-bilinear models.
Revised: The positive score functions λ ψ, 1 and λ ψ, 2 are speciﬁed to be simple log-bilinear models.

Original: Intuitively, L aux 1 encourages the latent states to encode useful information that helps distinguish the future summaries from negative samples.
Revised: Intuitively, L aux 1 encourages the latent states to encode useful information that helps distinguish the future summaries from negative samples.

Original: L auxencourages the deterministic states to reﬂect the interaction effects, as it contrastingly predicts the future summary of vertex i based on the states of i ’s neighbors only.
Revised: L auxencourages the deterministic states to reﬂect the interaction effects, as it contrastingly predicts the future summary of vertex i based on the states of i ’s neighbors only.

Original: The negative samples are selected from the future summaries of other vertices within the minibatch.
Revised: The negative samples are selected from the future summaries of other vertices within the minibatch.

Original: The ﬁnal objective to maximize is L =
Revised: The ﬁnal objective to maximize is L =

Original: L SMC K 
Revised: L SMC K + β 1 L aux 1 + β 2 L aux 2 , in which β 1 ≥ 0 and β 2 ≥ 0 are tunable hyperparameters.

Original: [Equation]
Revised: The procedure to estimate this objective is described in Appendix A.4.

Original:  R ELATED W ORK
Revised:  R ELATED W ORK

Original: GNN-based dynamics modeling.
Revised: GNN-based dynamics modeling.

Original: GNNs (Scarselli et al., 2009; Duvenaud et al., 2015; Li et al., 2016; Defferrard et al., 2016; Gilmer et al., 2017; Hamilton et al., 2017; Velikovi et al., 2018; Xu et al., 2019; Maron et al., 2019) provide a promising framework to learn on graph-structured data and impose relational inductive bias in learning models.
Revised: GNNs (Scarselli et al., 2009; Duvenaud et al., 2015; Li et al., 2016; Defferrard et al., 2016; Gilmer et al., 2017; Hamilton et al., 2017; Velikovi et al., 2018; Xu et al., 2019; Maron et al., 2019) provide a promising framework to learn on graph-structured data and impose relational inductive bias in learning models.

Original: We refer the reader to Battaglia et al.
Revised: We refer the reader to Battaglia et al.

Original: (2018) for a recent review.
Revised: (2018) for a recent review.

Original: GNNs (or neural message passing modules) are the core components of recently developed neural physics simulators (Battaglia et al., 2016; Watters et al., 2017;
Revised: GNNs (or neural message passing modules) are the core components of recently developed neural physics simulators (Battaglia et al., 2016; Watters et al., 2017;

Original: Chang et al., 2017; Janner et al., 2019; Sanchez-Gonzalez et al., 2018; Mrowca et al., 2018; Li et al., 2019) and spatiotemporal or multi-agent dynamics models (Alahi et al., 2016; 
Revised: Chang et al., 2017; Janner et al., 2019; Sanchez-Gonzalez et al., 2018; Mrowca et al., 2018; Li et al., 2019) and spatiotemporal or multi-agent dynamics models (Alahi et al., 2016; Hoshen, 2017;

Original: Li et al., 2018; Zhang et al., 2018;Hoshen, 2017;
Revised: Li et al., 2018; Zhang et al., 2018;

Original: Tacchetti et al., 2019).
Revised: Tacchetti et al., 2019; Chen et al., 2020).

Original: In these works, GNNs usually act autoregressively or be integrated into the sequence-to-sequence (seq2seq) framework (Sutskever et al.,2014).
Revised: In these works, GNNs usually act autoregressively or be integrated into the sequence-to-sequence (seq2seq) framework (Sutskever et al.,

Original: Besides, recently they have been combined with generative adversarial networks (Goodfellow et al., 2014) and normalizing ﬂows for multi-agent forecasting (Gupta et al., 2018;
Revised: Besides, recently they have been combined with generative adversarial networks (Goodfellow et al., 2014) and normalizing ﬂows for multi-agent forecasting (Gupta et al., 2018;

Original: Kosaraju et al., 2019;
Revised: Kosaraju et al., 2019;

Original: Rhinehart et al., 2019).
Revised: Rhinehart et al., 2019).

Original: R-SSM differs from all these works by introducing structured latent variables to represent the uncertainty on state transition and estimation.
Revised: R-SSM differs from all these works by introducing structured latent variables to represent the uncertainty on state transition and estimation.

Original: GNNs in sequential LVMs.
Revised: GNNs in sequential LVMs.

Original: A few recent works have combined GNNs with a sequential latent variable model, including R-NEM (van Steenkiste et al., 2018), NRI (Kipf et al., 2018), SQAIR (Kosiorek et al., 2018), VGRNN (Hajiramezanali et al., 2019), MFP (Tang & Salakhutdinov, 2019), and Graph VRNN (Sun et al., 2019; Yeh et al., 2019).
Revised: A few recent works have combined GNNs with a sequential latent variable model, including R-NEM (van Steenkiste et al., 2018), NRI (Kipf et al., 2018), SQAIR (Kosiorek et al., 2018), VGRNN (Hajiramezanali et al., 2019), MFP (Tang & Salakhutdinov, 2019), and Graph VRNN (Sun et al., 2019; Yeh et al., 2019).

Original: The latent variables in R-NEM and NRI are discrete and represent membership relations and types of edges, respectively.
Revised: The latent variables in R-NEM and NRI are discrete and represent membership relations and types of edges, respectively.

Original: In contrast, the latent variables in our model are continuous and represent the states of objects.
Revised: In contrast, the latent variables in our model are continuous and represent the states of objects.

Original: SQAIR is also a deep SSM for multi-object dynamics, but the GNN is only used in its inference network.
Revised: SQAIR is also a deep SSM for multi-object dynamics, but the GNN is only used in its inference network.

Original: VGRNN is focused on modeling the topological evolution of dynamical graphs.
Revised: VGRNN is focused on modeling the topological evolution of dynamical graphs.

Original: MFP employs a conditional
Revised: MFP employs a conditional

Original: VAE architecture, in which the per-agent discrete latent variables are shared by all time steps.
Revised: VAE architecture, in which the per-agent discrete latent variables are shared by all time steps.

Original: The work most relevant to ours is Graph VRNN, in which the hidden states of per-agent VRNNs interact through GNNs.
Revised: The work most relevant to ours is Graph VRNN, in which the hidden states of per-agent VRNNs interact through GNNs.

Original: Our work mainly differs from it by introducing a global latent state process to make the model hierarchical and exploring the use of normalizing ﬂows as well as the auxiliary contrastive objectives.
Revised: Our work mainly differs from it by introducing a global latent state process to make the model hierarchical and exploring the use of normalizing ﬂows as well as the auxiliary contrastive objectives.

Original: More subtle differences are discussed in Section 5.2.
Revised: More subtle differences are discussed in Section 5.2.

Original: Deep LVMs for sequential data.
Revised: Deep LVMs for sequential data.

Original: 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 are stochastic RNNs and deep SSMs.
Revised: 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 are stochastic RNNs and deep SSMs.

Original: 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 al., 2015;
Revised: 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 al., 2015;

Original: Fraccaro et al., 2016; Goyal et al., 2017; Ke et al., 2019) or, from another perspective, developing deep SSMs by parameterizing ﬂexible transition and emission distributions using neural networks (Krishnan et al., 2017; Fraccaro et al., 2017;
Revised: Fraccaro et al., 2016; Goyal et al., 2017; Ke et al., 2019) or, from another perspective, developing deep SSMs by parameterizing ﬂexible transition and emission distributions using neural networks (Krishnan et al., 2017; Fraccaro et al., 2017;

Original: Buesing et al., 2018; Zheng et al., 2017; Hafner et al., 2019).
Revised: Buesing et al., 2018; Zheng et al., 2017; Hafner et al., 2019).

Original: Approximate inference and parameter estimation methods for nonlinear SSMs have been extensively studied in the literature (Doucet & Johansen, 2009;
Revised: Approximate inference and parameter estimation methods for nonlinear SSMs have been extensively studied in the literature (Doucet & Johansen, 2009;

Original: Andrieu et al., 2010; Kantas et al., 2015; Gu et al., 2015; Karl et al., 2016;
Revised: Andrieu et al., 2010; Kantas et al., 2015; Gu et al., 2015; Karl et al., 2016;

Original: Marino et al., 2018; Gregor et al., 2019; Hirt & Dellaportas, 2019).
Revised: Marino et al., 2018; Gregor et al., 2019; Hirt & Dellaportas, 2019).

Original: We choose VSMC (Maddison et al., 2017; Naesseth et al., 2018; Le et al., 2018) as it combines the powers of VI and SMC.
Revised: We choose VSMC (Maddison et al., 2017; Naesseth et al., 2018; Le et al., 2018) as it combines the powers of VI and SMC.

Original: The posterior collapse problem is commonly addressed by KL annealing, which does not work with
Revised: The posterior collapse problem is commonly addressed by KL annealing, which does not work with

Original: VSMC.
Revised: VSMC.

Original: The idea of using auxiliary costs to train deep SSMs has been explored in Z-forcing (Goyal et al., 2017; Ke et al., 2019), which predicts the future summaries directly rather than contrastingly.
Revised: The idea of using auxiliary costs to train deep SSMs has been explored in Z-forcing (Goyal et al., 2017; Ke et al., 2019), which predicts the future summaries directly rather than contrastingly.

Original: As a result, the backward RNN in Z-forcing may degenerate easily.
Revised: As a result, the backward RNN in Z-forcing may degenerate easily.

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: We implement R-SSM using the TensorFlow Probability library (Dillon et al., 2017).
Revised: We implement R-SSM using the TensorFlow Probability library (Dillon et al., 2017).

Original: The experiments are organized as follows:
Revised: The experiments are organized as follows:

Original: In Section 5.1, we sample a toy dataset from a simple stochastic multi-object model and validate that R-SSM can ﬁt it well while AR models and non-relational models may struggle.
Revised: In Section 5.1, we sample a toy dataset from a simple stochastic multi-object model and validate that R-SSM can ﬁt it well while AR models and non-relational models may struggle.

Original: In Section 5.2, R-SSM is compared with state-of-the-art sequential LVMs for multi-agent modeling on a basketball gameplay dataset, and the effectiveness of GNF is tested through ablation studies.
Revised: In Section 5.2, R-SSM is compared with state-of-the-art sequential LVMs for multi-agent modeling on a basketball gameplay dataset, and the effectiveness of GNF is tested through ablation studies.

Original: Finally, in Section 5.3, the prediction performance of R-SSM is compared with strong GNN-based seq2seq baselines on a road trafﬁc dataset.
Revised: Finally, in Section 5.3, the prediction performance of R-SSM is compared with strong GNN-based seq2seq baselines on a road trafﬁc dataset.

Original: Due to the space constraint, the detailed model architecture and hyperparameter settings for each dataset are given in the Appendix.
Revised: Due to the space constraint, the detailed model architecture and hyperparameter settings for each dataset are given in the Appendix.

Original: Below, all values reported with error bars are averaged over 3 or 5 runs.
Revised: Below, all values reported with error bars are averaged over 3 or 5 runs.

Original: 5. S YNTHETIC TOY DATASET
Revised: 5. S YNTHETIC TOY DATASET

Original: First we construct a simple toy dataset to illustrate the capability of R-SSM.
Revised: First we construct a simple toy dataset to illustrate the capability of R-SSM.

Original: Each example in this dataset is generated by the following procedure:
Revised: Each example in this dataset is generated by the following procedure:

Original: [Equation]
Revised: [Equation]

Revised: [Table]

Original: for i = 1 ; : : : ; N and t = 1 ; : : : ; T .
Revised: for i = 1 , . . . , N and t = 1 , . . , T .

Original:  symmetric stochastic  model, in which each vertex i belongs to exact one of the K communities, and two vertices i and j are connected with probability p 

Revised: Here SBM is short for the symmetric stochastic block model, in which each vertex i belongs to exact one of the K communities, and two vertices i and j are connected with probability p

Original: if they are in the same community, p 1 otherwise.
Revised: if they are in the same community, p 1 otherwise.

Original: A vertex-speciﬁc covariate vector v 
Revised: A vertex-speciﬁc covariate vector v i ∈

Original: i 2 R d v is attached to each vertex i , and by Equation (6), the state of each vertex i can be affected by its neighbors N i .
Revised: R d v is attached to each vertex i , and by Equation ( 6), the state of each vertex i can be affected by its neighbors N i .

Original: Choosing the parameters d v = 4 , N = 36 , K = 3 , p 0 = 1 = 3 , p 1 = 1 = 18 , T = 80 , (cid:11) 1 = 5 : 0 , (cid:11) 2 = (cid:0) 1 : 5 , (cid:17) = 
Revised: Choosing the parameters d v = 4 , N = 36 , K = 3 , p 0 = 1 / 3 , p 1 = 1 / 18 , T = 80 , α 1 = 5 . 0 , α 2 = − 1 . 5 , η = [ − 1 . 5 , 0 . 4 , 2 . 0 , − 0 .

Original: [ (cid:0) 1 : 5 ; 0 : 4 ; 2 : 0 ; (cid:0) 0 : 9] ⊤ , (cid:27) x = (cid:27) z 
Revised: 9] ⊤ , σ x = σ z = 0 .

Original: = 0 : 05 , and " = 2 : 5 , we generate 10K examples for training, validation, and test, respectively.
Revised: 05 , and ε = 2 . 5 , we generate 10K examples for training, validation, and test, respectively.

Original: A typical example is visualized in the Appendix.
Revised: A typical example is visualized in the Appendix.

Original: Despite the simple generating process, the resulting dataset is highly challenging for common models to ﬁt.
Revised: Despite the simple generating process, the resulting dataset is highly challenging for common models to ﬁt.

Original: To show this, we compare R-SSM with several baselines, including (a) VAR:
Revised: To show this, we compare R-SSM with several baselines, including (a) VAR:

Original: Fitting a ﬁrst-order vector autoregression model for each example; (b) VRNN:
Revised: Fitting a ﬁrst-order vector autoregression model for each example; (b) VRNN:

Original: A variational RNN (Chung et al., 2015) shared by all examples; (c) GNN-AR:
Revised: A variational RNN (Chung et al., 2015) shared by all examples; (c) GNN-AR:

Original: A variant of the recurrent decoder of NRI (Kipf et al., 2018), which is exactly a GNN-based AR model when given the ground-truth graph.
Revised: A variant of the recurrent decoder of NRI (Kipf et al., 2018), which is exactly a GNN-based AR model when given the ground-truth graph.

Original: VAR and VRNN are given access to the observations f x ( i ) 1: T g N i =1 only, while GNN-AR and R-SSM are [Table] + GNF ( [Table]
Revised: VAR and VRNN are given access to the observations { x ( i ) 1: T } Ni =1 only, while GNN-AR and R-SSM are additionally given access to the graph structure ( V , E ) (but not the vertex covariates).

Original: GNF

Original: [Table]

Original: [Table]

Original: additionally given access to the graph structure ( V ; E ) (but not the vertex covariates).

Original: GNF is not used in R-SSM because the true joint transition distribution is factorized over vertices.[Table] -AR
Revised: GNF is not used in R-SSM because the true joint transition distribution is factorized over vertices.

Original: [Table]
Revised: [Table]

Original: For each model, we calculate three metrics: 
Revised: For each model, we calculate three metrics: (

Original: (1) LL: Average log-likelihood (or its lowerbound) of
Revised: ) LL: Average log-likelihood (or its lower

Original:  test examples ; (2) MSE: Average
Revised: bound) of test examples; (2) MSE

Original: mean squared

Original:  one-step prediction error given the ﬁrst 75 time steps of 
Revised: Average mean squared one-step prediction error given the ﬁrst 75 time steps of each test example; (3)

Original: each
Revised: CP:

Original: test example

Original: ; (3) CP: Average coverage probability of a % one step prediction 
Revised: Average coverage probability of a 90% one-step prediction interval.

Original: interval

Original: For nonanalytic models, point predictions and prediction intervals are computed using
Revised: For nonanalytic models, point predictions and prediction intervals are computed using

Original: Monte Carlo samples.
Revised: Monte Carlo samples.

Original: The results are reported in Table 1.
Revised: The results are reported in Table 1.

Original: The generating process involves latent factors and nonlinearities, so VAR performs poorly as expected.
Revised: The generating process involves latent factors and nonlinearities, so VAR performs poorly as expected.

Original: VRNN largely underﬁts the data and struggles to generalize, which may be caused by the different topologies under the examples.
Revised: VRNN largely underﬁts the data and struggles to generalize, which may be caused by the different topologies under the examples.

Original: In contrast, GNN-AR and R-SSM generalize well as expected, while R-SSM achieves much higher test log-likelihood and produces good one-step probabilistic predictions.
Revised: In contrast, GNN-AR and R-SSM generalize well as expected, while R-SSM achieves much higher test log-likelihood and produces good one-step probabilistic predictions.

Original: 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.
Revised: 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.

Original: We also observed that without L auxthe training dynamics easily get stuck in posterior collapse at the very early stage, and adding L aux 2 help improve the test likelihood.
Revised: We also observed that without L auxthe training dynamics easily get stuck in posterior collapse at the very early stage, and adding L aux 2 help improve the test likelihood.

Original: 5. B ASKETBALL GAMEPLAY
Revised: 5. B ASKETBALL GAMEPLAY

Original: In basketball gameplay, the trajectories of players and the ball are highly correlated and demonstrate rich, dynamic interations.
Revised: In basketball gameplay, the trajectories of players and the ball are highly correlated and demonstrate rich, dynamic interations.

Original: Here we compare R-SSM with a state-of-the-art hierarchical sequential
Revised: Here we compare R-SSM with a state-of-the-art hierarchical sequential

Original: LVM for multi-agent trajectories (Zhan et al., 2019), in which the per-agent VRNNs are coordinated by a global "macro intent" model.
Revised: LVM for multi-agent trajectories (Zhan et al., 2019), in which the per-agent VRNNs are coordinated [Table] [Table] [Table]

Revised: by a global "macro intent" model.

Original: We note it as MI-VRNN.
Revised: We note it as MI-VRNN.

Original: The dataset  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.
Revised: The dataset 2 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.

Original: Following their settings, we use the trajectories of offensive team only and preprocess the data in exactly the same way to make the results directly comparable.
Revised: Following their settings, we use the trajectories of offensive team only and preprocess the data in exactly the same way to make the results directly comparable.

Original: The complete graph of players is used as the input to R-SSM.
Revised: The complete graph of players is used as the input to R-SSM.

Original: Several ablation studies are performed to verify the utility of the proposed ideas.
Revised: Several ablation studies are performed to verify the utility of the proposed ideas.

Original: In Table 3, 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.
Revised: In Table 3, 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.

Original: The VRNN baseline developed by Zhan et al.
Revised: The VRNN baseline developed by Zhan et al.

Original: (2019) is also included for comparison.
Revised: (2019) is also included for comparison.

Original: Note that the VSMC bound L SMC 1000 is a tighter log-likelihood approximation than the ELBO (which is equivalent to L SMC).
Revised: Note that the VSMC bound L SMC 1000 is a tighter log-likelihood approximation than the ELBO (which is equivalent to L SMC).

Original: The rollout statistics of R-SSMs are calculated from 150K 50-step rollouts with 10 burn-in steps.
Revised: The rollout statistics of R-SSMs are calculated from 150K 50-step rollouts with 10 burn-in steps.

Original: Several selected rollouts are visualized in the Appendix.
Revised: Several selected rollouts are visualized in the Appendix.

Original: As illustrated in Table 2, all R-SSMs outperform the baselines in terms of average test log-likelihood.
Revised: As illustrated in Table 2, all R-SSMs outperform the baselines in terms of average test log-likelihood.

Original: Again, we observed that adding L aux 1 is necessary for training R-SSM successfully on this dataset.[Table] + GNF [Table]
Revised: Again, we observed that adding L aux 1 is necessary for training R-SSM successfully on this dataset.

Original: [Table]

Original: Training with the proposed auxiliary loss L aux 2 and adding GNFs do improve the results.
Revised: Training with the proposed auxiliary loss L aux 2 and adding GNFs do improve the results.

Original: R-SSM with 8 GNFs (4 in prior, 4 in proposal) achieves higher likelihood than R-SSM with 4 GNFs, indicating that increasing the expressivity of joint state distributions helps ﬁt the data better.
Revised: R-SSM with 8 GNFs (4 in prior, 4 in proposal) achieves higher likelihood than R-SSM with 4 GNFs, indicating that increasing the expressivity of joint state distributions helps ﬁt the data better.

Original: As for the rollout quality, the OOB rate of the rollouts sampled from our model matches the ground-truth signiﬁcantly better, while the other two statistics are comparable to the MI-VRNN baseline
Revised: As for the rollout quality, the OOB rate of the rollouts sampled from our model matches the ground-truth signiﬁcantly better, while the other two statistics are comparable to the MI-VRNN baseline.

Original: In Table [Table] we also provide preliminary results for the setting that additionally includes the trajectory of the ball.
Revised: In Table 3, we also provide preliminary results for the setting that additionally includes the trajectory of the ball.

Original: This enables us to compare with the results reported by Yeh et al.
Revised: This enables us to compare with the results reported by Yeh et al.

Original: (2019) for
Revised: (2019) for

Original: Graph VRNN (GVRNN).
Revised: Graph VRNN (GVRNN).

Original: The complete graph of ball and players served as input to R-SSM is annotated with two node types (player or ball) and three edge types (player-to-ball, ball-to-player or player-to-player).
Revised: The complete graph of ball and players served as input to R-SSM is annotated with two node types (player or ball) and three edge types (player-to-ball, ball-to-player or player-to-player).

Original: R-SSM achieves competitive test likelihood, and adding GNFs helps improve the performance.
Revised: R-SSM achieves competitive test likelihood, and adding GNFs helps improve the performance.

Original: We point out that several noticeable design choices of GVRNN may help it outperform
Revised: We point out that several noticeable design choices of GVRNN may help it outperform

Original: R-SSM: (i) GVRNN uses a GNN-based observation model, while R-SSM uses a simple factorized observation model.
Revised: R-SSM: (i) GVRNN uses a GNN-based observation model, while R-SSM uses a simple factorized observation model.

Original: (ii) GVRNN encodes X 1: t (cid:0) 1 into H t and thus enables the prior of Z t to depend on past observations, which is not the case in R-SSM.
Revised: (ii) GVRNN encodes X 1: t − 1 into H t and thus enables the prior of Z t to depend on past observations, which is not the case in R-SSM.

Original: (iii) GVRNN uses several implementation tricks, e.g., predicting the changes in observations only ( X t = X t (cid:0) 1 + ∆ X t ) and passing raw observations as additional input to GNNs.
Revised: (iii) GVRNN uses several implementation tricks, e.g., predicting the changes in observations only ( X t = X t − 1 + ∆ X t ) and passing raw observations as additional input to GNNs.

Original: We would like to investigate the effect of these interesting differences in future work.
Revised: We would like to investigate the effect of these interesting differences in future work.

Original: 5. R OAD TRAFFIC
Revised: 5. R OAD TRAFFIC

Original: Trafﬁc speed forecasting on road networks is an important but challenging task, as the trafﬁc dynamics exhibit complex spatiotemporal interactions.
Revised: Trafﬁc speed forecasting on road networks is an important but challenging task, as the trafﬁc dynamics exhibit complex spatiotemporal interactions.

Original: In this subsection, we demonstrate that R-SSM is comparable to the state-of-the-art GNN-based seq2seq baselines on a real-world trafﬁc dataset.
Revised: In this subsection, we demonstrate that R-SSM is comparable to the state-of-the-art GNN-based seq2seq baselines on a real-world trafﬁc dataset.

Original: The METR-LA dataset (Li et al., 2018) contains 4 months of 1D trafﬁc speed measurements that were recorded via 207 sensors and aggregated into 5 minutes windows.
Revised: The METR-LA dataset (Li et al., 2018) contains 4 months of 1D trafﬁc speed measurements that were recorded via 207 sensors and aggregated into 5 minutes windows.

Original: For this dataset, all conditional inputs G =
Revised: For this dataset, all conditional inputs G =

Original: ( V ; E ; V ; E ) and U 1: T are provided to R-SSM, in which E is constructed by connecting two sensors if their road network distance is below a threshold, V stores the geographic positions and learnable embeddings of sensors, E stores the road network distances of edges, and
Revised: ( V , E , V , E ) and U 1: T are provided to R-SSM, in which E is constructed by connecting two sensors if their road network distance is below a threshold, V stores the geographic positions and learnable embeddings of sensors, E stores the road network distances of edges, and

Original: U 1: T provides the time information (hour-of-day and day-of-week).
Revised: U 1: T provides the time information (hour-of-day and day-of-week).

Original: We impute the missing values for training and exclude them from evaluation.
Revised: We impute the missing values for training and exclude them from evaluation.

Original: GNF is not used because of GPU memory limitation.
Revised: GNF is not used because of GPU memory limitation.

Original: Following the settings in Li et al.
Revised: Following the settings in Li et al.

Original: (2018), we train our model on small time windows spanninghours and use a 7:2:1 split for training, validation, and test.
Revised: (2018), we train our model on small time windows spanninghours and use a 7:1:2 split for training, validation, and test.

Original: The comparison of mean absolute forecast errors (MAE) is reported in Table 4.
Revised: The comparison of mean absolute forecast errors (MAE) is reported in Table 4.

Original: The three forecast horizons correspond to 15, 30, and 60 minutes.
Revised: The three forecast horizons correspond to 15, 30, and 60 minutes.

Original: We give point predictions by taking the elementwise median of 2K Monte Carlo forecasts.
Revised: We give point predictions by taking the elementwise median of 2K Monte Carlo forecasts.

Original: Compared with DCRNN (Li et al., 2018) and GaAN (Zhang et al., 2018), R-SSM delivers comparable short-term forecasts but slightly worse long-term forecasts.
Revised: Compared with DCRNN (Li et al., 2018) and GaAN (Zhang et al., 2018), R-SSM delivers comparable short-term forecasts but slightly worse long-term forecasts.

Original: 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 to beat.
Revised: 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 to beat.

Original: (ii) Some stochastic systems are inherently unpredictable beyond a few steps due to the process noise, e.g., the toy model in Section 5.1.
Revised: (ii) Some stochastic systems are inherently unpredictable beyond a few steps due to the process noise, e.g., the toy model in Section 5.1.

Original: In such case, multi-step MAE may not be a reasonable metric, and probabistic forecasts may be prefered.
Revised: In such case, multi-step MAE may not be a reasonable metric, and probabistic forecasts may be prefered.

Original: The average coverage probabilities (CP) of 90% prediction intervals reported in Table 4 indicate that R-SSM provides good uncertainty estimates.
Revised: The average coverage probabilities (CP) of 90% prediction intervals reported in Table 4 indicate that R-SSM provides good uncertainty estimates.

Original: (iii) Improving the multi-step prediction ability of deep SSMs is still an open problem with a few recent attempts (Ke et al., 2019; Hafner et al., 2019).
Revised: (iii) Improving the multi-step prediction ability of deep SSMs is still an open problem with a few recent attempts (Ke et al., 2019; Hafner et al., 2019).

Original: We would like to explore it in future work.
Revised: We would like to explore it in future work.

Original:  C ONCLUSIONS
Revised:  C ONCLUSIONS

Original: In this work we present a deep hierarchical state-space model in which the state transitions of correlated objects are coordinated by graph neural networks.
Revised: In this work, we present a deep hierarchical state-space model in which the state transitions of correlated objects are coordinated by graph neural networks.

Original: To effectively learn the model from observation data, we develop a structured posterior approximation and propose two auxiliary contrastive prediction tasks to help the learning.
Revised: To effectively learn the model from observation data, we develop a structured posterior approximation and propose two auxiliary contrastive prediction tasks to help the learning.

Original: We further introduce the graph normalizing ﬂow to enhance the expressiveness of the joint transition density and the posterior approximation.
Revised: We further introduce the graph normalizing ﬂow to enhance the expressiveness of the joint transition density and the posterior approximation.

Original: The experiments show that our model can outperform or match the state-of-the-arts on several time series modeling tasks.
Revised: The experiments show that our model can outperform or match the state-of-the-arts on several time series modeling tasks.

Original: Directions for future work include testing the model on high-dimensional observations, extending the model to directly learn from visual data and including discrete latent variables in the model
Revised: Directions for future work include testing the model on high-dimensional observations, extending the model to directly learn from visual data, and including discrete latent variables in the model. A CKNOWLEDGMENTS

Revised: This work was supported by the National Key Research and Development Program of China (No.

Revised: 2018YFB0505000) and the Alibaba-Zhejiang University Joint Institute of Frontier Technologies.

Revised: Fan Yang would like to thank Qingchen Yu for her helpful feedback on early drafts of this paper

Original: A A PPENDIX A. MHA IMPLEMENTATION
Revised: A A PPENDIX A. MHA IMPLEMENTATION

Original: Each attention head k 2 [ K ] is separately parameterized and operates as follows.
Revised: Each attention head k ∈ [ K ] is separately parameterized and operates as follows.

Original: For 8 ( j; i ) 2 E , it produces a message vector (cid:12) kj!
Revised: For ∀ ( j, i ) ∈ E , it produces a message vector β kj

Original:  i together with an unnormalized scalar weight ! kj ! i .
Revised: → i together with an unnormalized scalar weight ω kj → i .

Original: Then for 8 i 2 V , it aggregates all messages sent to i in a permutation-invariant manner: { [Equation] [Equation] softmax
Revised: Then for ∀ i ∈ V , it aggregates all messages sent to i in a permutation-invariant manner: (cid:8)

Original: [Equation]
Revised: [Equation]

Original: Putting all K heads together, it turns out that M j ! i = f ( ! kj ! i ; (cid:12) kj !
Revised: Putting all K heads together, it turns out that M j → i = { ( ω kj → i , β kj → i ) }

Original: i ) g

Original: Kk =1 and A i = f a ki g Kk =1 .
Revised: Kk =1 and A i = { a ki } Kk =1 .

Original: Speciﬁcally, each attention head is parameterized in a query-key-value style: 
Revised: Speciﬁcally, each attention head is parameterized in a query-key-value style: ∀ i ∈ V : [Equation] = MLP v ( v i

Original: [Equation]
Revised: [Equation]

Revised: ˜ d v ,

Revised: [Equation]

Original: =
Revised: =

Revised: [Equation]

Revised: , ˜ v i , g ]

Revised: Q =

Revised: [Equation]

Revised: j → i = c

Revised: j

Revised: [Equation]

Revised: =

Original: [Equation]
Revised: [Equation]

Original:  for ~ d = d + ~ d v + d g , W Q 2 R ~ d (cid:2)
Revised: p d [Equation] for ˜ d = d + ˜ d v + d g , W Q ∈ R ˜

Original: d  q , W K 2 R ~ d (cid:2) d q , and W V 2 R ~ d (cid:2) d c .
Revised: d × d q , W A ∈ R ˜ d × d q , and W C ∈ R ˜ d × d c .

Original: A. M ODEL DETAILS
Revised: A. M ODEL DETAILS

Original: 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.
Revised: 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.

Original: The transition densities in the generative model are speciﬁed to be:
Revised: The transition densities in the generative model are speciﬁed to be:

Original: [Equation]
Revised: [Equation]

Original: [Equation]
Revised: Σ gθ ( h gt )) , (7) f ⋆θ (cid:0) Z t (cid:12)(cid:12) H t (cid:1) =

Original: [Equation]
Revised: N Y i =

Original: (8) where (cid:22) g(cid:18) and (cid:6) g(cid:18)
Revised: Normal (cid:16) z ( i ) t (cid:12)(cid:12)(cid:12) µ ⋆θ (cid:16) h

Original:  (similarly (cid:22) ⋆(cid:18) and (cid:6) ⋆(cid:18) ) are 3-layer MLPs that share their ﬁrst layer.
Revised: ( i ) t (cid:17) , Σ ⋆θ (cid:16) h ( i ) t (cid:17)(cid:17) , (8) where µ gθ and Σ gθ (similarly µ ⋆θ and Σ ⋆θ ) are 3-layer MLPs that share their ﬁrst layer.

Original: (cid:6) g(cid:18) and (cid:6) ⋆(cid:18) output diagonal covariance matrices using the softplus activation.
Revised: Σ gθ and Σ ⋆θ output diagonal covariance matrices using the softplus activation.

Original: The proposal densities r gϕ and r ⋆ϕ are speciﬁed in a similar way.
Revised: The proposal densities r gϕ and r ⋆ϕ are speciﬁed in a similar way.

Original: Then GNFs can be stacked on top of f ⋆(cid:18) and r ⋆ϕ to make them more expressive.
Revised: Then GNFs can be stacked on top of f ⋆θ and r ⋆ϕ to make them more expressive.

Original: A. E XPERIMENTS
Revised: A. E XPERIMENTS

Original: We use the Adam optimizer with an initial learning rate of 0.001 and a gradient clipping of 1. for all experiments.
Revised: We use the Adam optimizer with an initial learning rate of 0.001 and a gradient clipping of 1. for all experiments.

Original: The learning rate was annealed according to a linear cosine decay.
Revised: The learning rate was annealed according to a linear cosine decay.

Original: We set (cid:12) 1 = (cid:12) 2 =0 : 1  for the auxiliary losses in all experiments.
Revised: We set β 1 = β 2 = 1 . 0 for the auxiliary losses in all experiments.

Original: Synthetic toy dataset.
Revised: Synthetic toy dataset.

Original: A typical example in the dataset is visualized in Figure 2.
Revised: A typical example in the dataset is visualized in Figure 2.

Original: The architectures of the models are speciﬁed as follows.
Revised: The architectures of the models are speciﬁed as follows

Original: (a) VRNN:
Revised: (a) VRNN:

Original: Using 128-dimensional latent variables and a two-layer, 512-unit GRU.
Revised: Using 128-dimensional latent variables and a two-layer, 512-unit GRU.

Original: (b) GNN-AR: Using a two-layer GNN and an one-layer, 128-unit GRU shared by all nodes.
Revised: (b) GNN-AR: Using a two-layer GNN and an one-layer, 128-unit GRU shared by all nodes.

Original: (c) R-SSM: We let d g =
Revised: (c) R-SSM: We let d g =

Original: d z = 8 .
Revised: d z = 8 .

Original: All RNNs are speciﬁed to be two-layer, 32-unit LSTMs.
Revised: All RNNs are speciﬁed to be two-layer, 32-unit LSTMs.

Original: All MLPs use 64 hidden units.
Revised: All MLPs use 64 hidden units.

Original: The generative model and the proposal both use a 4-head MHA layer.
Revised: The generative model and the proposal both use a 4-head MHA layer.

Original: 4 SMC samples and a batch size of 16 are used in training.
Revised: 4 SMC samples and a batch size of 16 are used in training.

Original: Basketball player movement.
Revised: Basketball player movement.

Original: We let d g =
Revised: We let d g =

Original: d
Revised: d

Original: z = 32 .
Revised: z = 32 .

Original: All RNNs are speciﬁed to be two-layer, 64unit LSTMs and all MLPs use 256 hidden units.
Revised: All RNNs are speciﬁed to be two-layer, 64unit LSTMs and all MLPs use 256 hidden units.

Original: The generative model uses one 8-head MHA layer and the proposal uses two 8-head MHA layers.
Revised: The generative model uses one 8-head MHA layer and the proposal uses two 8-head MHA layers.

Original: Each GNF uses an additional MHA layer shared by the functions s ( (cid:1) ) and t ( (cid:1) ) .
Revised: Each GNF uses an additional MHA layer shared by the functions s ( · ) and t ( · ) .

Original: 4 SMC samples and a batch size of 64 are used in training.
Revised: 4 SMC samples and a batch size of 64 are used in training.

Original: Eight selected rollouts from the trained model are visualized in Figure 3.
Revised: Eight selected rollouts from the trained model are visualized in Figure 3.

Original: Road trafﬁc.
Revised: Road trafﬁc.

Original: We let d g =
Revised: We let d g =

Original: d
Revised: d

Original: z = 8 .
Revised: z = 8 .

Original: A 32-dimensional embedding for each sensor is jointly learned as a part of the vertex attribute.
Revised: A 32-dimensional embedding for each sensor is jointly learned as a part of the vertex attribute.

Original: All RNNs are speciﬁed to be two-layer, 32-unit LSTMs and all
Revised: All RNNs are speciﬁed to be two-layer, 32-unit LSTMs and all

Original: MLPs use 64 hidden units.
Revised: MLPs use 64 hidden units.

Original: The generative model and the proposal both use two 8-head MHA layers.
Revised: The generative model and the proposal both use two 8-head MHA layers.

Original: 3 SMC samples and a batch size of 16 are used in training.
Revised: 3 SMC samples and a batch size of 16 are used in training.

Original: A. T RAINING
Revised: A. T RAINING

Original: We optimize the VSMC bound estimated by the following SMC algorithm:
Revised: We optimize the VSMC bound estimated by the following SMC algorithm:

Revised: Algorithm 1 Estimate the VSMC bound L SMC K

Original: Input: graph G , observations X 1: T , exogenous inputs U 1: T
Revised: Input: graph G , observations X 1: T , exogenous inputs U 1: T

Original: Require: generative model  f  f g(cid:18) ; f ⋆(cid:18) ; g (cid:18) g  , proposal f r gϕ ; r ⋆ϕ g , number of particles K for k = 1 : : : K do
Revised: Require: generative model { f gθ , f ⋆θ , g θ } , proposal { r gϕ , r ⋆ϕ } , number of particles K for k = 1 . . . K do

Original: Simulate z g 1 ;k (cid:24) r gϕ ( (cid:1)jG ; X 1 ; U 1 )
Revised: Simulate z g 1 ,k ∼ r gϕ ( ·|G , X 1 , U 1 )

Original: Simulate Z 1 ;k (cid:24) r ⋆ϕ ( (cid:1)j z g 1 ;k ; G ; X 1 ; U 1 )
Revised: Simulate Z 1 ,k ∼ r ⋆ϕ ( ·| z g 1 ,k , G , X 1 , U 1 )

Original: Set w k 1 = f g(cid:18)
Revised: Set w k 1 = f gθ ( z g 1 ,k )

Original: ( z g 1 ;k ) f⋆(cid:18)
Revised: f

Original:  ( Z 1 ;k j z g 1 ;k ; G ; U 1 )
Revised: ⋆θ ( Z 1 ,k | z g 1 ,k , G , U 1 )

Original: ∏ Ni =1 g (cid:18) ( x ( i ) 1 j z ( i ) 1 ;k ; z g 1 ;k ;::: ) r gϕ ( z g 1 ;k j ::: )
Revised: ∏ Ni =1 g θ ( x ( i ) 1 | z ( i ) 1 ,k , z g 1 ,k ,... ) r gϕ ( z g 1 ,k | ... )

Original: r ⋆ϕ ( Z 1 ;k j ::: ) end for
Revised: r ⋆ϕ ( Z 1 ,k | ... ) end for

Original: Initialize ^ L SMC K = log 
Revised: Initialize ˆ L SMC K = log P

Original: ∑ Kk =1 w k 1 = K for t = 2 : : : T do f
Revised: Kk =1 w k 1 / K for t = 2 . . .

Original:  z g< t;k ; Z < t;k g Kk =1 = RESAMPLE
Revised: T do { z g<t,k , Z <t,k } Kk =1 = RESAMPLE

Original: ( f z g<t;k ; Z <t;k ; w kt [-(cid:0) 
 g-] Kk =1 ) for k = 1 : : : K do
Revised: { z g<t,k , Z <t,k , w kt [-(cid:0) 
 g-]− 1 } Kk =1 ) for k = 1 . . . K do

Original: Simulate z gt;k (cid:24) r gϕ ( (cid:1)j
Revised: Simulate z gt,k ∼ r gϕ ( ·| z g<t,k , Z <t,k , G , X ≤ t , U ≤ t )

Original:  z g<t;k ; Z <t;k ; G ; X (cid:20) t ; U (cid:20) t )Simulate Z t;k (cid:24) r ⋆ϕ
Revised: Simulate Z t,k ∼ r ⋆ϕ ( ·| z g ≤ t,k , Z <t,k , G , X ≤ t , U ≤ t )

Original: ( (cid:1)j z g (cid:20) t;k ; Z <t;k ; G ; X (cid:20) t ; U (cid:20) t ) Set w kt =f g(cid:18)
Revised: Set w kt =

Original:  ( z gt;k j ::: ) f ⋆(cid:18) ( Z t;k j z g (cid:20) t;k ;
Revised: f gθ ( z gt,k | ... ) f ⋆θ ( Z t,k | z g ≤ t,k , Z <t,k ,... )

Original: Z <t;k ;::: ) ∏ Ni =1 g (cid:18) ( x ( i ) t j z ( i ) t;k ; z g (cid:20)
Revised: ∏ Ni =1 g θ ( x ( i ) t | z ( i ) t,k , z g ≤ t,k , Z <t,k ,... )

Original: t;k ; Z <t;k ;::: ) r gϕ ( z gt ;k j ::: )
Revised: r gϕ ( z gt ,k | ... )

Original: r ⋆ϕ ( Z t;k j ::: )
Revised: r ⋆ϕ ( Z t,k | ... )

Original: Set z g (cid:20) t;k = ( z g<
Revised: Set z g ≤ t,k = ( z g<t,k , z gt,k ) , Z

Original: t;k ; z gt;k ) , Z (cid:20) t;k
Revised: [Equation] t, [Equation]

Original: =
Revised: =

Original: ( Z <t;k ; Z t;k ) end for
Revised: ( Z <t,k , Z t,k [Equation] end for

Original: Update^
Revised: Update

Original: L SMC K

Original: [Equation]
Revised: [Equation]

Original: end for
Revised: for

Original: Output: ^
Revised: Output

Original: L SMC K
Revised: K

Original: In our model, dependencies on z g<t;k and Z <t;k are provided through the compact RNN states h gt;k and H t;k .
Revised: In our model, dependencies on z g<t,k and Z <t,k are provided through the compact RNN states h gt,k and H t,k .

Original: When GNFs are used in f ⋆(cid:18) and r ⋆ϕ , density calculation and backpropagation are automatically handled by the TensorFlow Probability library.
Revised: When GNFs are used in f ⋆θ and r ⋆ϕ , density calculation and backpropagation are automatically handled by the TensorFlow Probability library.

Original: To estimate the auxiliary objectives L auxand L aux 2 , we reuse the resampled unweighted particles ff z g 1: t;k ; Z 1: t;k g Kk =1 g T (cid:0) 1 t =1 generated by Algorithm 1 to form Monte Carlo estimations for them:^ L aux 1 = T (cid:0) 1 ∑ t =
Revised: To estimate the auxiliary objectives L auxand L aux 2 , we reuse the resampled unweighted particles {{ z g 1: t,k , Z 1: t,k } Kk =1 } T − 1 t =1 generated by Algorithm 1 to form Monte Carlo estimations for them:

Original: N ∑ i = K log (cid:21) ; 1 (^ z ( i ) t;k ; c ( i ) t ) ∑ c 2 Ω t;i (cid:21) ; 1 (^ z ( i ) t;k ; c ) ; ^ L aux 2 = T (cid:0) 1 ∑ t =
Revised: [Equation]

Original: N ∑ i = K log (cid:21) ; 2 (^ h ( i ) t;k ; c ( i ) t ) ∑ c 2 Ω t;i (cid:21) ; 2 (^ h ( i ) t;k ; c ) where ^ z ( i ) t;k
Revised: where ˆ z ( i ) t,k = [ z gt,k , z ( i ) t,k ] , ˆ h

Original: = [ z gt;k ; z ( i ) t;k ] , ^ h ( i ) t;k
Revised: ( i ) t,k

Original: = MLP  ( ∑ j ̸ = i ^
Revised: = MLP ψ ( P j ̸ = i ∧

Original: j 2N (cid:0) i h ( j ) t;k ) , and Ω t;i is a set that contains c ( i ) t and some negative samples selected from the future summaries of other vertices within the minibatch.
Revised: j ∈N − i h ( j ) t,k ) , and Ω t,i is a set that contains c ( i ) t and some negative samples selected from the future summaries of other vertices within the minibatch.

Original: A.
Revised: A.

Original: G RAPH N ORMALIZING F LOW
Revised: G RAPH N ORMALIZING F LOW

Original: [Figure]
Revised: [Figure]

Revised: B

Original: C D
Revised: C D

Original: E F
Revised: E F

Original: [Figure]
Revised: [Figure]

Original: 𝒁 𝑏′ ∈ ℝ 𝑁 ×
Revised: 𝑎′

Original: [Figure]
Revised: [Figure]

Revised: 𝒁 𝑏′ ∈ ℝ 𝑁 ×

Revised: 𝐷−𝑑

Original: 𝒁 ′ ∈
Revised: 𝒁 ′ ∈

Original: ℝ 𝑁 × 𝐷
Revised: ℝ 𝑁 × 𝐷

Original: A B
Revised: A B

Original: C D
Revised: C D

Original: E F * +
Revised: E F * +

Original: A B
Revised: A B

Original: C D
Revised: C D

Original: E F
Revised: E F

Original: A B
Revised: A B

Original: C D
Revised: C D

Original: E F =
Revised: E F =

Original: A B
Revised: A B

Original: C D
Revised: C D

Original: E F
Revised: E F

Original: A B
Revised: A B

Original: C D
Revised: C D

Original: E F
Revised: E F

Original: Message Passing
Revised: Message Passing

Original: A
Revised: A

Original: B
Revised: B

Original: C D
Revised: C D

Original: E
Revised: E

Original: F
Revised: F

Original: A
Revised: A

Original: B
Revised: B

Original: C D
Revised: C D

Original: E
Revised: E

Original: F
Revised: F

Original: A
Revised: A

Original: B
Revised: B

Original: C D
Revised: C D

Original: E
Revised: E

Original: F
Revised: F

Original: A
Revised: A

Original: B
Revised: B

Original: C D
Revised: C D

Original: E
Revised: E

Original: F
Revised: F

Original: [Figure]
Revised: [Figure]

Revised: A B

Original: C D
Revised: C D

Original: E F
Revised: E F

Original: [Figure]
Revised: [Figure]

Original: The element-wise afﬁne layer is proposed by Kingma & Dhariwal (2018) for normalizing the activations.
Revised: The element-wise afﬁne layer is proposed by Kingma & Dhariwal (2018) for normalizing the activations.

Original: Its parameters (cid:13) 2 R D and (cid:12) 2 R D are initialized such that the per-channel activations have roughly zero mean and unit variance at the beginning of training.
Revised: Its parameters γ ∈ R D and β ∈ R D are initialized such that the per-channel activations have roughly zero mean and unit variance at the beginning of training.

Original: The invertible linear transformation W 2 R D (cid:2) D is parameterized using a QR decomposition (Hoogeboom et al., 2019).
Revised: The invertible linear transformation W ∈ R D × D is parameterized using a QR decomposition (Hoogeboom et al., 2019).