Tree-Structured Recurrent Switching Linear Dynamical Systems for Multi-Scale Modeling

1 INTRODUCTION . Complex systems can often be described at multiple levels of abstraction . A computer program can be characterized by the list of functions it calls , the sequence of statements it executes , or the assembly instructions it sends to the microprocessor . As we zoom in , we gain an increasingly nuanced view of the system and its dynamics . The same is true of many natural systems . For example , brain activity can be described in terms of high-level psychological states or via detailed ion channel activations ; different tasks demand different levels of granularity . One of our principal aims as scientists is to identify appropriate levels of abstraction for complex natural phenomena and to discover the dynamics that govern how these systems behave at each level of resolution . Modern machine learning offers a powerful toolkit to aid in modeling the dynamics of complex systems . Bayesian state space models and inference algorithms enable posterior inference of the latent states of a system and the parameters that govern their dynamics ( Särkkä , 2013 ; Barber et al. , 2011 ; Doucet et al. , 2001 ) . In recent years , this toolkit has been expanded to incorporate increasingly flexible components like Gaussian processes ( Frigola et al. , 2014 ) and neural networks ( Chung et al. , 2015 ; Johnson et al. , 2016 ; Gao et al. , 2016 ; Krishnan et al. , 2017 ) into probabilistic time series models . In 1This work was supported by National Science Foundation ( NSF IIS-1734910 , CCF-1617986 , HRD-1612689 ) and National Institute of Health ( NIH R01EB026946 ) . SWL was supported by the Simons Collaboration on the Global Brain ( SCGB-418011 ) . neuroscience , sequential autoencoders offer highly accurate models of brain activity ( Pandarinath et al. , 2018 ) . However , while these methods offer state of the art predictive models , their dynamics are specified at only the most granular resolution , leaving the practitioner to tease out higher level structure post hoc . Here we propose a probabilistic generative model that provides a multi-scale view of the dynamics through a hierarchical architecture . We call it the tree-structured recurrent switching linear dynamical system , or TrSLDS . The model builds on the recurrent SLDS ( Linderman et al. , 2017 ) to approximate latent nonlinear dynamics through a hierarchy of locally linear dynamics . Once fit , the TrSLDS can be queried at different levels of the hierarchy to obtain dynamical descriptions at multiple levels of resolution . As we proceed down the tree , we obtain higher fidelity , yet increasingly complex , descriptions . Thus , depth offers a simple knob for trading off interpretability and flexibility . The key contributions are two-fold2 : first , we introduce a new form of tree-structured stick breaking for multinomial models that strictly generalizes the sequential stick breaking of the original rSLDS , while still permitting Pólya-gamma data augmentation ( Polson et al. , 2013 ) for efficient posterior inference ; second , we develop a hierarchical prior that links dynamics parameters across levels of the tree , thereby providing descriptions that vary smoothly with depth . The paper is organized as follows . Section 2 provides background material on switching linear dynamical systems and their recurrent variants . Section 3 presents our tree-structured model and Section 4 derives an efficient fullyBayesian inference algorithm for the latent states and dynamics parameters . Finally , in Section 5 we show how our model yields multi-scale dynamics descriptions for synthetic data from two standard nonlinear dynamical systems—the Lorenz attractor and the FitzHugh-Nagumo model of nonlinear oscillation—as well as for a real dataset of neural responses to visual stimuli in a macaque monkey . 2 BACKGROUND . Let xt ∈ Rdx and yt ∈ Rdy denote the latent state and the observation of the system at time t respectively . The system can be described using a state-space model : xt = f ( xt−1 , wt ; Θ ) , wt ∼ Fw ( state dynamics ) ( 1 ) yt = g ( xt , vt ; Ψ ) , vt ∼ Fv ( observation ) ( 2 ) where Θ denotes the dynamics parameters , Ψ denotes the emission ( observation ) parameters , and wt and vt are the state and observation noises respectively . For simplicity , we restrict ourselves to systems of the form : xt = f ( xt−1 ; Θ ) + wt , wt ∼ N ( 0 , Q ) , ( 3 ) If the state space model is completely specified then recursive Bayesian inference can be applied to obtain an estimate of the latent states using the posterior p ( x0 : T |y1 : T ) ( Doucet et al. , 2001 ) . However in many applications , the parametric form of the state space model is unknown . While there exist methods that perform smoothing to obtain an estimate of x0 : T ( Barber , 2006 ; Fox et al. , 2009 ; Djuric & Bugallo , 2006 ) , we are often interested in not only obtaining an estimate of the continuous latent states but also in learning the dynamics f ( · ; Θ ) that govern the dynamics of the system . In the simplest case , we can take a parametric approach to solving this joint state-parameter estimation problem . When f ( · ; Θ ) and g ( · ; Ψ ) are assumed to be linear functions , the posterior distribution over latent states is available in closed-form and the parameters can be learned via expectationmaximization . On the other hand , we have nonparametric methods that use Gaussian processes and neural networks to learn highly nonlinear dynamics and observations where the joint estimation is untractable and approximations are necessarily imployed ( Zhao & Park , 2016 ; 2018 ; Frigola et al. , 2014 ; Sussillo et al. , 2016 ) . Switching linear dynamical systems ( SLDS ) ( Ackerson & Fu , 1970 ; Chang & Athans , 1978 ; Hamilton , 1990 ; Ghahramani & Hinton , 1996 ; Murphy , 1998 ) balance between these two extremes , approximating the dynamics by stochastically transitioning between a small number of linear regimes . 2Source code is available at https : //github.com/catniplab/tree_structured_rslds 2.1 SWITCHING LINEAR DYNAMICAL SYSTEMS . SLDS approximate nonlinear dynamics by switching between a discrete set of linear regimes . An additional discrete latent state zt ∈ { 1 , . . . , K } determines the linear dynamics at time t , xt = xt−1 +Aztxt−1 + bzt + wt , wt ∼ N ( 0 , Qzt ) ( 4 ) where Ak , Qk ∈ Rdx×dx and bk ∈ Rdx for k = 1 , . . . , K. Typically , zt is endowed with Markovian dynamics , Pr ( zt|zt−1 = k ) = πk . The conditionally linear dynamics allow for fast and efficient learning of the model and can utilize the learning tools developed for linear systems ( Haykin , 2001 ) . While SLDS can estimate the continuous latent states x0 : T , the assumption of Markovian dynamics for the discrete latent states severely limits their generative capacity . 2.2 RECURRENT SWITCHING LINEAR DYNAMICAL SYSTEMS . Recurrent switching linear dynamical systems ( rSLDS ) ( Linderman et al. , 2017 ) , also known as augmented SLDS ( Barber , 2006 ) , are an extension of SLDS where the transition density of the discrete latent state depends on the previous location in the continuous latent space zt|xt−1 , { R , r } ∼ πSB ( νt ) , ( 5 ) νt = Rxt−1 + r , ( 6 ) where R ∈ RK−1×dx and r ∈ RK−1 represents hyperplanes . πSB : RK−1 → [ 0 , 1 ] K maps from the reals to the probability simplex via stick-breaking : πSB ( ν ) = ( π ( 1 ) SB ( ν ) , · · · , π ( K ) SB ( ν ) ) , π ( k ) SB = σ ( νk ) ∏ j < k σ ( −νj ) , ( 7 ) for k = 1 , . . . , K−1 and π ( K ) SB = ∏K−1 k=1 σ ( −νk ) where νk is the kth component of of ν and σ ( ν ) = ( 1 + e−ν ) −1 is the logistic function ( Fig . 1 ) . By including this recurrence in the transition density of zt , the rSLDS partitions the latent space into K sections , where each section follows its own linear dynamics . It is through this combination of locally linear dynamical systems that the rSLDS approximates eq . ( 3 ) ; the partitioning of the space allows for a more interpretable visualization of the underlying dynamics . Recurrent SLDS can be learned efficiently and in a fully Bayesian manner , and experiments empirically show that they are adept in modeling the underlying generative process in many cases . However , the stick breaking process used to partition the space poses problems for inference due to its dependence on the permutation of the discrete states { 1 , · · · , K } ( Linderman et al. , 2017 ) . 3 TREE-STRUCUTRED RECURRENT SWITCHING LINEAR DYNAMICAL SYSTEMS . Building upon the rSLDS , we propose the tree-structured recurrent switching linear dynamical system ( TrSLDS ) . Rather than sequentially partitioning the latent space using stick breaking , we use a treestructured stick breaking procedure ( Adams et al. , 2010 ) to partition the space . Let T denote a tree structure with a finite set of nodes { , 1 , · · · , N } . Each node n has a parent node denoted by par ( n ) with the exception of the root node , , which has no parent . For simplicity , we initially restrict our scope to balanced binary trees where every internal node n is the parent of two children , left ( n ) and right ( n ) . Let child ( n ) = { left ( n ) , right ( n ) } denote the set of children for internal node n. Let Z ⊆ T denote the set of leaf nodes , which have no children . Let depth ( n ) denote the depth of a node n in the tree , with depth ( ) = 0 . At time instant t , the discrete latent state zt is chosen by starting at the root node and traversing down the tree until one of theK leaf nodes are reached . The traversal is done through a sequence of left/right choices by the internal nodes . Unlike in standard regression trees where the choices are deterministic ( Lakshminarayanan , 2016 ) , we model the choices as random variables . The traversal through the tree can be described as a stick breaking process . We start at the root node with a unit-length stick π = 1 , which we divide between its two children . The left child receives a fraction πleft ( ) = σ ( ν ) and the right child receives the remainder πright ( ) = 1− σ ( ν ) such that ν ∈ R specifies the left/right balance . This process is repeated recursively , subdividing πn into two pieces at each internal node until we reach the leaves of the tree ( Fig . 1 ) . The stick assigned to each node is thus , πn = { σ ( νpar ( n ) ) I [ n=left ( par ( n ) ) ] ( 1− σ ( νpar ( n ) ) ) I [ n=right ( par ( n ) ) ] πpar ( n ) n 6= , 1 n = . ( 8 ) We incorporate this into the TrSLDS by allowing νn to be a function of the continuous latent state νn ( xt−1 , Rn , rn ) = R T nxt−1 + rn , ( 9 ) where the parametersRn and rn specify a linear hyperplane in the continuous latent state space . As the continuous latent state xt−1 evolves , the left/right choices become more or less probable . This in turn changes the probability distribution πk ( xt−1 , Γ , T ) over the K leaf nodes , where Γ = { Rn , rn } n∈T . In the TrSLDS , these leaf nodes correspond to the discrete latent states of the model , such that for each leaf node k , p ( zt = k | xt−1 , Γ , T ) = πk ( xt−1 , Γ , T ) . ( 10 ) In general , the tree-structured stick-breaking is not restricted to balanced binary trees . We can allow more than two children through an ordered sequential stick-breaking at each level . In this sense , tree-structured stick-breaking is a strict generalization of stick-breaking . We also note that similar to rSLDS , the model can be made more flexible by introducing a dependence on the previous discrete latent in eq . ( 9 ) but for the rest of the paper , we stick to eq . ( 8 ) .

Summary: The authors develop a tree structured extension to the recently proposed recurrent switching linear dynamical systems. Like switching linear dynamical systems (sLDS) the proposed models capture non-linear dynamics by switching between a collection of linear regimes. However, unlike SLDS, the transition between the regimes is a function of a latent tree as well as the preceding continuous latent state. Experiments on synthetic data as well as neural spike train data are presented to demonstrate the utility of the model.