Relational State-Space Model for Stochastic Multi-Object Systems

1 INTRODUCTION . Many real-world dynamical systems can be decomposed into smaller interacting subsystems if we take a fine-grained view . For example , the trajectories of coupled particles are co-determined by perparticle physical properties ( e.g. , mass and velocity ) and their physical interactions ( e.g. , gravity ) ; traffic flow can be viewed as the coevolution of a large number of vehicle dynamics . 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 finance . 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 . Although SSMs are traditionally designed under the guidance of domain-specific knowledge or tractability consideration , recently introduced deep SSMs ( Fraccaro , 2018 ) use neural networks ( NNs ) to parameterize flexible state transitions and emissions , achieving much higher expressivity . 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 ) . Recent works that advocate GNNs for modeling multi-object dynamics mostly make use of GNNs in an autoregressive ( AR ) fashion . 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 ) . 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 oversimplified observation distributions ( Yang et al. , 2018 ) . In this study , we make the following contributions : ( i ) 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 . ( ii ) We suggest using the graph normalizing flow ( GNF ) to construct expressive joint ∗Corresponding author . state distributions for R-SSM , further enhancing its ability to capture the joint evolutions of correlated stochastic subsystems . ( iii ) We develop structured posterior approximation to learn R-SSM using variational inference and introduce two auxiliary training objectives to facilitate the learning . Our experiments on synthetic and real-world time series datasets show that R-SSM achieves competitive test likelihood and good prediction performance in comparison to GNN-based AR models and other sequential LVMs . The remainder of this paper is organized as follows : Section 2 briefly reviews neccesary preliminaries . Section 3 introduces R-SSM formally and presents the methods to learn R-SSM from observations . Related work is summarized in Section 4 and experimental evaluation is presented in Section 5 . We conclude the paper in Section 6 . 2 PRELIMINARIES . In this work , an attributed directed graph is given by a 4-tuple : G = ( V , E , V , E ) , where V = [ N ] : = { 1 , . . . , N } is the set of vertices , E ⊆ [ N ] × [ N ] is the set of edges , V ∈ RN×dv is a matrix of static vertex attributes , and E ∈ RN×N×de is a sparse tensor storing the static edge attributes . The set of direct predecessors of vertex i is notated as N−i = { p| ( p , i ) ∈ E } . We use the notation xi to refer to the i-th row of matrix X and write xij to indicate the ( i , j ) -th entry of tensor X ( if the corresponding matrix or tensor appears in the context ) . For sequences , we write x≤t = x1 : t : = ( x1 , . . . , xt ) and switch to x ( i ) t for referring to the i-th row of matrix Xt . 2.1 GRAPH NEURAL NETWORKS . GNNs are a class of neural networks developed to process graph-structured data and support relational reasoning . 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 . Let H ∈ RN×d be a matrix of vertex representations , in which the i-th row hi ∈ Rd is the vectorized representation attached to vertex i . Conditioning on the static graph structure and attributes given by G , a GNN just takes the vertex representations H along with some graph-level context g ∈ Rdg as input and returns new vertex representations H′ ∈ RN×d′ as output , i.e. , H′ = GNN ( G , g , H ) . When updating the representation of vertex i from hi to h′i , a GNN takes the representations of other nearby vertices into consideration . 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 . Formally , the operations performed by a basic block of a message-passing GNN are defined as follows : ∀ ( j , i ) ∈ E : Mj→i = MESSAGE ( g , vj , vi , eji , hj , hi ) ( 1 ) ∀i ∈ V : Ai = AGGREGATE ( { Mp→i } p∈N−i ) ( 2 ) ∀i ∈ V : h′i = COMBINE ( g , vi , hi , Ai ) ( 3 ) Throughout this work , we implement Equations ( 1 ) and ( 2 ) by adopting a multi-head attention mechanism similar to Vaswani et al . ( 2017 ) and Velikovi et al . ( 2018 ) . For Equation ( 3 ) , we use either a RNN cell or a residual block ( He et al. , 2016 ) , depending on whether the inputs to GNN are RNN states or not . We write such a block as H′ = MHA ( G , g , H ) and give its detailed implementation in the Appendix . A GNN simply stacks L separately-parameterized MHA blocks and iteratively computes H = : H ( 0 ) , . . . , H ( L ) = : H′ , in which H ( l ) = MHA ( G , g , H ( l−1 ) ) for l = 1 , . . . , L. We write this construction as H′ = GNN ( G , g , H ) and treat it as a black box to avoid notational clutter . 2.2 STATE-SPACE MODELS . State-space models are widely applied to analyze dynamical systems whose true states are not directly observable . Formally , an SSM assumes the dynamical system follows a latent state process { zt } t≥1 , which possibly depends on exogenous inputs { ut } t≥1 . Parameterized by some ( unknown ) static parameter θ , the latent state process is characterized by an initial density z1 ∼ πθ ( ·|u1 ) and a transition density zt+1 ∼ fθ ( ·|z≤t , u≤t+1 ) . Moreover , at each time step , some noisy measurements of the latent state are observed through an observation density : xt ∼ gθ ( ·|z≤t , u≤t ) . The joint density of x1 : T and z1 : T factors as : p ( x1 : T , z1 : T |u1 : T ) = πθ ( z1|u1 ) ∏T t=2 fθ ( zt|z < t , u≤t ) ∏T t=1 gθ ( xt|z≤t , u≤t ) . The superior expressiveness of SSMs can be seen from the fact that the marginal predictive distribution p ( xt|x < t , u≤t ) = ∫ p ( xt , z≤t|x < t , u≤t ) dz≤t can be far more complex than unimodal distributions and their finite mixtures that are common in AR models . Recently developed deep SSMs use RNNs to compress z≤t ( and u≤t ) into fixed-size vectors to achieve tractability . As shown in next section , R-SSM can be viewed as enabling multiple individual deep SSMs to communicate . 2.3 NORMALIZING FLOWS . Normalizing flows ( Rezende & Mohamed , 2015 ) are invertible transformations that have the capability to transform a simple probability density into a complex one ( or vice versa ) . Given two domains X ⊆ RD and Y ⊆ RD , let f : X → Y be an invertible mapping with inverse f−1 . Applying f to a random variable z ∈ X with density p ( z ) , by the change of variables rule , the resulting random variable z′ = f ( z ) ∈ Y will have a density : p ( z′ ) = p ( z ) ∣∣∣∣det ∂f−1∂z′ ∣∣∣∣ = p ( z ) ∣∣∣∣det ∂f∂z ∣∣∣∣−1 A series of invertible mappings with cheap-to-evaluate determinants can be chained together to achieve complex transformations while retaining efficient density calculation . This provides a powerful way to construct expressive distributions . 3 RELATIONAL STATE-SPACE MODEL . Suppose there is a dynamical system that consists of multiple interacting objects , and observing this system at a specific time is accomplished by acquiring measurements from every individual object simultaneously . 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 . 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 . 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 . 3.1 GENERATIVE MODEL . Given the observations for a multi-object dynamical system , our model further assumes its interaction structure is known as prior knowledge . 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 . In situations where such graph structure is not available , a complete graph can be specified . However , to model dynamical systems comprising a large number of objects , it is often beneficial to explicitly specify sparse graph structures , because they impose stronger relational inductive bias and help save the computational cost . A relational state-space model assumes a set of correlated dynamical subsystems evolve jointly under the coordination of graph neural networks . Formally , given a graph G = ( V , E , V , E ) , in which an edge ( i , j ) ∈ E indicates that the state of vertex j may be affected by vertex i . Let u ( i ) t ∈ Rdu and x ( i ) t ∈ Rdx 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 ∈ Rdz represents the latent state of vertex i at time step t. Furthermore , we introduce a global latent variable zgt ∈ Rdg for each time step to represent the global state shared by all vertices . Conditioning on the graph and exogenous inputs , an R-SSM factorizes the joint density of observations and latent states as follows : pθ ( { x ( i ) 1 : T } N i=1 , { z ( i ) 1 : T } N i=1 , zg1 : T ∣∣∣G , { u ( i ) 1 : T } Ni=1 ) = T∏ t=1 fθ ( { z ( i ) t } N i=1 , zgt ∣∣∣ { z ( i ) < t } Ni=1 , zg < t , G , { u ( i ) ≤t } Ni=1 ) T∏ t=1 N∏ i=1 gθ ( x ( i ) t ∣∣∣z ( i ) t , zg≤t , { z ( i ) < t } Ni=1 , G , { u ( i ) ≤t } Ni=1 ) ( 4 ) For notational simplicity , we switch to the matrix notation Zt = [ z ( 1 ) t , z ( 2 ) t , . . . , z ( N ) t ] ⊤ from now on . The joint transition density fθ is further factorized as a product of global transition density fgθ and local transition density f ⋆ θ , i.e. , fθ ( Zt , z g t | . . . ) = f g θ ( z g t | . . . ) f⋆θ ( Zt|z g t , . . . ) . To instantiate these conditional distributions , a GNN accompanied by RNN cells is adopted to recurrently compress the past dependencies at each time step into fixed-size context vectors . Specifically , the observations are assumed to be generated from following process : h̃ ( i ) t = RNN v θ ( h ( i ) t−1 , [ z ( i ) t−1 , u ( i ) t ] ) , h̃gt = READOUTθ ( G , H̃t ) hgt = RNN g θ ( hgt−1 , [ zgt−1 , h̃ g t ] ) , zgt ∼ f g θ ( · ∣∣hgt ) Ht = GNNθ ( G , zgt , H̃t ) , Zt ∼ f⋆θ ( · ∣∣Ht ) x ( i ) t ∼ gθ ( · ∣∣z ( i ) t , h ( i ) t , zgt , hgt ) for i = 1 , . . . , N and t = 1 , . . . , T , where h ( i ) 0 = h ∗ 0 and z ( i ) 0 = z ∗ 0 . Here h g 0 , z g 0 , h ∗ 0 and z ∗ 0 are learnable initial states . The READOUT function aggregates the context vectors of all vertices into a global context vector in a permutation-invariant manner . The global transition density fgθ is specified 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 . 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 . The graphical structure of two consecutive steps of the generating process is illustrated in Figure 1a . 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 .

Summary: The paper presents a relational state-space model that simulates the joint state transitions of correlated objects which are hierarchically coordinated in a graph structure. A structured posterior approximation is developed based on sequential graph neural networks. Two auxiliary contrastive predictive losses are proposed to help circumvent the posterior collapse problem. Graph normalizing flow is further incorporated into the framework to make the joint state transition density more expressive. The proposed R-SSM shows performance gains over state-of-the-arts in three benchmarks.