Dynamical Wasserstein Barycenters for Time-series Modeling

1 Introduction . We consider the problem of estimating the dynamically evolving state of a system from time-series data.1 The notion of “ state ” in such contexts typically is modeled in one of two ways . For many problems , the system state is a vector of continuous quantities ( Kalman , 1960 ; Krishnan et al. , 2016 ) , perhaps constrained in some manner . Alternatively , discrete-state models take on one of a countable number of options at each point in time , as exemplified by hidden Markov models ( HMMs ) ( Rabiner , 1989 ) or “ switching-state ” extensions ( Ghahramani and Hinton , 2000 ; Linderman et al. , 2017 ) . Many time-series characterization problems of current interest warrant a hybrid of continuous and discrete state representation approaches , where the system gradually transitions in a continuous manner among a finite collection of “ pure ” discrete states . For example , in human activity recognition using accelerometer sensor data ( Bi et al. , 2021 ) , some segments of data do correspond to distinct activities ( run , sit , walk , etc . ) , suggesting a discrete state representation . However , when using sensors with high-enough sampling rates , transition periods when the system is evolving from one state to another ( e.g . the individual accelerates from standing to running over a few seconds ) can be well 1Code available at https : //github.com/kevin-c-cheng/DynamicalWassBarycenters_Gaussian 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . resolved . Gradual evolution between pure states can also be observed in other domains of time series data , such as economics ( Chang et al. , 2016 ) or climate science ( Chang et al. , 2020 ) . Characterizing these systems requires a model with a continuous state space to capture the gradual evolution of the system among the discrete set of pure states . Motivated by this class of applications , we consider models for time series in which the system ’ s dynamical state is specified by a vector of convex combination weights for mixing a set of datagenerating distributions that define the individual pure states . Many approaches , such as mixture models , interpret such a simplex-constrained state vector ( Rudin , 1976 ) as assignment probabilities ; that is , the system is assumed to be in a pure state with uncertainty as to which . As a result , the data-generating distribution at moments of transition is a convex , linear combination of the pure-state emission distributions . While useful in many applications , such linear interpolation does not capture the gradual transitions among pure states in the time series of interest to us . To illustrate the shortcomings of linear interpolation , consider the toy data task in Fig . 1 , where a system gradually transitions between three pure states over time . During the transition periods ( e.g . at times 600 and 1400 ) , the linear interpolation method infers a data-generating distribution that is multi-modal , shown in Fig . 1 ( b ) . If we refer to our pure states as “ walk ” and ” run , ” this approach models the walk-to-run transition as sometimes walk and sometimes run . This does not intuitively capture the gradual nature of accelerating from walk to run in our intended applications . To overcome this limited representation , we consider another way to mix together pure-state distributions : displacement-interpolation ( McCann , 1997 ) , which is related to the Wasserstein distance ( Peyré and Cuturi , 2019 ) , a metric over the space of probability distributions ( Sriperumbudur et al. , 2010 ) . While the work of McCann ( 1997 ) is limited to combining two distributions , it is extended to multiple distributions using the notion of a Wasserstein barycenter ( Agueh and Carlier , 2011 ) . Fig . 1 ( c ) shows how a Wasserstein barycenter approach to time-series modeling infers data-generating distributions during transitions that are not multi-modal but instead place mass in between where the two pure-state distributions do . This intuitively captures gradual transition between two pure states . Inspired by this framework , in this work we develop a dynamical Wasserstein barycentric ( DWB ) model for time series intended to explain data arising as a system evolves between pure states . Our model uses a barycentric weight vector to represent the system state . Given an observed multivariate time-series and a desired number of states K , all parameters are estimated in an unsupervised way . Estimation simultaneously learns the data-generating distributions of K pure discrete states as well as the K-simplex valued barycentric weight vector state at each timestep . Given the nature of our model , we require that the state lie in the simplex at every timestep , a constraint not respected by the Gaussian noise that drives common continuous-state processes ( Welch , 1997 ) . Building on work by Nguyen and Volkov ( 2020 ) , we employ a random walk where the driving noise comes from independent , identically distributed ( IID ) draws from a mixture of two Beta distributions , representing stationary and transitional dynamics . By blending the current state and a mixture-of-Betas draw in a convex manner , we construct a new state that lies in the simplex . To specify the emission distributions of our model , we assume that each pure state generates data from a multivariate Gaussian . While a Gaussian model may not be suitable in all applications , this choice allows us to exploit useful properties of Gaussian densities under the Wasserstein distance ( Takatsu , 2011 ) . Specifically , a closed-form expression exists for the Wasserstein distance between Gaussians , the Wasserstein barycenter among Gaussians can be computed via a simple recursion , and the estimation of the Gaussian mean vectors and covariance matrices can be performed conveniently over a Riemannian product manifold . Empirically , we find our proposed DWB model with Gaussian pure-states performs well on human activity datasets , accurately characterizing both pure-states emission distributions and capturing the system state in pure states and transition periods . Contributions : We introduce a displacement-interpolation model for time series where the datagenerating distribution is given by the weighted Wasserstein barycenter of a set of pure-state emission distributions and a time-varying state vector . We propose a simplex-valued random walk with flexible dynamical structure to model the system state . We exploit the Riemannian structure of Gaussian distributions under the Wasserstein distance for parameter estimation for faster convergence speed . We evaluate on human activity data and demonstrate the ability of our method to capture stationary and transition dynamics , comparing with the linear interpolation mixture model and with a continuous state space model . Outline : Sec . 2 provides an overview of the Wasserstein distance , barycenter , and the associated geometry for Gaussian distributions . We then formalize our problem statement and estimation problem in Sec . 3 . Sec . 4 discusses the model parameters , covering the dynamical simplex state-space model in Sec . 4.1 and the pure-state parameters in Sec . 4.2 . Sec . 5 discusses the optimization of our model parameters leveraging geometric properties of the Wasserstein distance for Gaussians . Finally , Sec . 6 evaluates and discusses the advantages of our model in the context of human activity data . 2 Technical Background . A core component to our approach is to model the intermediate transition states of a time series using the Wasserstein barycenter ( Agueh and Carlier , 2011 ) of probability distributions , which generalizes the displacement-interpolation framework ( McCann , 1997 ) beyond two distributions . We refer the works of ( Peyré and Cuturi , 2019 ) and ( Villani , 2009 ) for a detailed discussion on these concepts . Consider the space of all Borel probability measures over Rd with finite second moment . The squared Wasserstein-2 distance for two distributions ρ1 , ρ2 with squared Euclidean ground cost is defined as , W22 ( ρ1 , ρ2 ) = inf M∈Π ( ρ1 , ρ2 ) ∫ Rd×Rd ‖α− β‖22M ( dα , dβ ) , ( 1 ) where Π ( ρ1 , ρ2 ) is the set of all joint distributions with marginals ρ1 , ρ2 , and M is the optimal transport plan , the element that minimizes the total transportation cost . When these measures are Gaussian , parameterized by their mean vectorsmi ∈ Rd , and symmetric positive-definite covariance matrices Si ∈ Symd+ , the squared Wasserstein-2 distance has a closed form solution ( Takatsu , 2011 ) , W22 ( ρ1 ( m1 , S1 ) , ρ2 ( m2 , S2 ) ) = ‖m1 −m2‖ 2 2︸ ︷︷ ︸ E2 ( m1 , m2 ) + tr ( S1 + S2 − 2 ( S 1 2 1 S2S 1 2 1 ) ) ︸ ︷︷ ︸ B2 ( S1 , S2 ) . ( 2 ) This distance decomposes into sum of the squared Euclidean distance between mean vectors , E2 ( m1 , m2 ) , and the squared Bures distance ( Bhatia et al. , 2017 ) between covariance matrices , B2 ( S1 , S2 ) . Thus , the contributions of the mean and covariance to the Wasserstein distance between Gaussians are decoupled , a property that is uncommon for Gaussian distribution distances ( Nagino and Shozakai , 2006 ) and has important implications for optimization and barycenter computation . The Wasserstein barycenter extends the notion of the weighted average of points in Rd using the Euclidean distance to the space of probability distributions with the Wasserstein distance . Given a set of K measures and the barycentric coordinate vector on the K-simplex , x ∈ ∆K , the Wasserstein barycenter is the measure that minimizes this weighted Wasserstein distance to the set of measures , ρB ( x , { ρk } Kk=1 ) = argmin ρ∈P2 ( Rd ) K∑ k=1 x [ k ] W22 ( ρk , ρ ) . ( 3 ) When ρk are Gaussian distributions , ρB defined in ( 3 ) is itself Gaussian ( Agueh and Carlier , 2011 ) with parametersmB , SB . Again , because of the decomposition of the Wasserstein distance in ( 2 ) , the Wasserstein barycentric problem in ( 3 ) can be solved separately for its components , mB = argmin m∈Rd K∑ k=1 x [ k ] E2 ( mk , m ) , SB = argmin S∈Symd+ K∑ k=1 x [ k ] B2 ( Sk , S ) . ( 4 ) The optimal mean can be computed in closed-form : mB = ∑ k x [ k ] mk . The optimal covariance matrix can be solved via the fixed-point iteration proposed in Álvarez Esteban et al . ( 2016 ) . 3 Problem Formulation . Model Definition . Our DWB model ’ s data-generating process is illustrated in Fig . 2 . First , the pure-state emission parameters Θ ≡ { ( mk , Sk ) } Kk=1 , define a Gaussian distribution ρk for each pure state k. Second , the state vector xt defines the system state at t and lies on the simplex . Given Θ and xt , we can form a time-varying Gaussian distribution ρBt ( xt , Θ ) ≡ N ( mBt , SBt ) for each t , which is a barycentric combination of the K pure Gaussians using weights xt via ( 4 ) . We can write the states for an entire sequence asX , comprised of an initial state and a sequence of simplex-valued state vectors , denoted X ≡ { x0 , { xt } Tt=1 } . Data Preprocessing . We are given a vector-valued time series of observations yτ ∈ Rd , τ = 1 , ... , T . Instead of modeling this data directly , to improve smoothness we model the empirical distribution of sliding windows of 2n+ 1 samples ( Aghabozorgi et al. , 2015 ) strided by δ samples . We retain only windows with complete data , with start times corresponding to τ = 1 , ( δ + 1 ) , ( 2δ + 1 ) , ... , b ( T − ( 2n+1 ) ) δ cδ + 1 , which we index sequentially as t ∈ { 1 , 2 , . . . T } . A window indexed at t corresponds to a window centered at τ = ( δ ( t− 1 ) + n+ 1 ) , which provides an estimates of the underlying distribution at yτ . For each window location t , we compute an unbiased empirical Gaussian distribution ρ̂t = N ( mt , St ) where , mt = 1 2n+ 1 2n+1∑ i=1 yδ ( t−1 ) +i , St = 1 2n 2n+1∑ i=1 ( yδ ( t−1 ) +i −mt ) ( yδ ( t−1 ) +i −mt ) T . ( 5 ) Minimizing the Wasserstein distance between this sequence of empirical distributions ρ̂t and the model-predicted distributions ρBt drives our model ’ s parameter learning . Estimation Objective . In practice , we are given an observed sequence of empirical distributions { ρ̂t } Tt=1 and a desired number of states K. We wish to estimate the state sequenceX and emission parameters Θ . We pose the estimation of X and Θ as the solution to an optimization problem seeking to balance fidelity to a prior model on the parameters of interest with the desire to minimize the time integrated Wasserstein distance between the predicted and observed distributions : X̂ , Θ̂ = argmin X , Θ − log ( p ( X ) p ( Θ ) ) + λ T∑ t=1 W22 ( ρ̂t , ρBt ( xt , Θ ) ) . ( 6 ) The scalar weight λ > 0 trades off the model ’ s fit to data ( measured by the Wasserstein distance ) with the probability of the state sequenceX and pure-state emission parameters Θ under assumed prior distributions . Our chosen priors p ( X ) and the p ( Θ ) are covered in the following section .

Summary: This paper deals with the problem of modelling continuous-space time series as sequences of high-level hidden states. In such models, transitions are usually represented as multinomial mixtures over hidden states. In this work, another approach is considered in which transition distributions are modelled as Wasserstein barycenters between pure distributions attached to hidden states. In the case of Gaussian distributions, such barycenters are also Gaussians, and when the match between empirical distributions and those model-predicted distributions is assessed using Wasserstein distance, closed-form formulations can be used. From those observations, a loss is derived that can be optimized by alternating between (i) gradient descent on the "innovation" parameters $H$ and $\Gamma$ and (ii) gradient descent on a Riemannian manifold.