Deep Markov Spatio-Temporal Factorization

We introduce deep Markov spatio-temporal factorization (DMSTF), a deep generative model for spatio-temporal data. Like other factor analysis methods, DMSTF approximates high-dimensional data by a product between time-dependent weights and spatially dependent factors. These weights and factors are in turn represented in terms of lower-dimensional latent variables that we infer using stochastic variational inference. The innovation in DMSTF is that we parameterize weights in terms of a deep Markovian prior, which is able to characterize nonlinear temporal dynamics. We parameterize the corresponding variational distribution using a bidirectional recurrent network. This results in a flexible family of hierarchical deep generative factor analysis models that can be extended to perform time series clustering, or perform factor analysis in the presence of a control signal. Our experiments, which consider simulated data, fMRI data, and traffic data, demonstrate that DMSTF outperforms related methods in terms of reconstruction accuracy and can perform forecasting in a variety domains with nonlinear temporal transitions.


Introduction
Analysis of large-scale spatio-temporal data is relevant to a wide range of applications in biology, marketing, traffic control, climatology, and neuroscience. Due to the high dimensionality of the data, methods for spatiotemporal analysis either exploit smoothness and correlations in the signal [4,26], or map high-dimensional data onto a lower-dimensional representation [37,44,39,2,6,45,8].
In this paper, we combine these two principles by building on recent approaches for hierarchical probabilistic factor analysis [25,24,35]. These models represent data Y ∈ R T ×D with T temporal and D spatial dimensions as a product Y W F between temporal weights W ∈ R T ×K and spatial factors arXiv:2003.09779v1 [cs.
LG] 22 Mar 2020 F ∈ R K×D . Hierarchical approaches to probabilistic factor analysis explicitly model the correlations between multiple instances {Y 1 , . . . , Y N } by representing W n and F n in terms of some set of latent variables Z n that are drawn from a shared prior. The result is a modeling framework that performs two levels of dimensionality reduction. At the instance level, the model represents highdimensional spatio-temporal data Y n in terms of lower-dimensional matrices W n and F n . At the level of the corpus, we further reduce dimensionality by learning a mapping from latent variables Z n to weights and factors.
The main contribution of this paper is that we develop methods for hierarchical probabilistic factor analysis that uncover common patterns of temporal variation across instances and are amenable to temporal forecasting. To this end, we introduce deep Markov spatio-temporal factorization (DMSTF), a model that learns a deep generative Markovian prior in order to reason about (potentially nonlinear) temporal dynamics forecasting.
DMSTF is related to recent studies that employ temporal smoothness assumptions [9] or more explicitly model dynamics [37,44,39,33,17,2,6] for matrix factorization of multidimensional times series. While most of these methods are not essentially probabilistic (providing only point estimates for imputation/prediction tasks), some Bayesian probabilistic matrix factorization methods have been proposed [34,45,8], and linear temporal dymanics on factor latents, W n , have been adapted to these methods in [43,7,38]. However, these methods do not explicitly adopt a priori assumptions about functional form of spatial factors, F n , when available. This makes them intractable for extremely high dimensional spatial data such as neuroimaging data. Moreover, the linear dynamical assumptions in these methods fall short of modeling potentially nonlinear transitions.
DMSTF builds on topographic factor analysis (TFA) [25] and neural topographic factor analysis (NTFA) [35], which were designed with neuroimaging applications in mind. However, the DMSTF generalizes from both these models, specifically by modeling temporal dynamics, hence, can be applied to spatiotemporal analysis more generally. We demonstrate this in our experiments, which evaluate model performance on simulated data, two large-scale fMRI datasets, and four traffic datasets. Our experiments demonstrate that DMSTF achieves better predictive performance for unseen data relative to existing baselines, both when evaluating the log-likelihood for new instances and when performing next time-point forecasting.
We summarize the contributions of this paper as follows: -DMSTF learns a deep Markovian prior that represents high-dimensional spatio-temporal data in terms of low-dimensional latent variables that can capture nonlinear dynamics of the data. -DMSTF learns a mapping from latent variables to spatial factors and temporal weights, which provide an intermediate representation for downstream regression or classification tasks. -DMSTF is able to perform clustering in the low-dimensional temporal latent space, which can provide visual insights about data.
-DMSTF introduces a control input that can be used for semi-or fullysupervised factorization.
The remainder of this paper is structured as follows. Section 2 presents an overview of related work. Section 3 described the DMSTF model architecture. Section 4 provides details on model training. Section 5 describes experiments on simulated data, neuroimaging datasets and traffic datasets.

Related Works
Factor analysis has been extensively used for reducing dimensionality in spatiotemporal data. Principal component analysis (PCA) [29] and independent component analysis (ICA) [11] are among the most well-known classical factor analysis methods. To accommodate tensor data, and mitigate scalability issues, multilinear versions have been proposed [15,32,10,41]. These methods do not naturally handle missing data. They are also permutation invariant along the batch dimension, and therefore cannot capture temporal dynamics [44]. Spatial factors obtained by these methods are also unstructured and difficult to interpret in many applications. In an early attempt to get temporally smooth structures, Chen and Cichocki [9] developed a non-negative matrix factorization model and applied temporal smoothness and spatial decorrelation regularizers to achieve physiologically meaningful components.
Since then, several matrix/tensor factorization approaches have been proposed for modeling temporal dynamics in multivariate/multidimensional time series. Sun et al. [37] presented a dynamic matrix factorization suited for collaborative filtering setting in recommendation systems using a linear-Gaussian dynamical state space model. Yu et al. [44] proposed to use an autoregressive temporal regularizer for matrix factorization to describe temporal dependencies in multivariate time series. [39] added an additional graph Laplacian regularizer to learn spatial autocorrelations and perform prediction on unknown locations. [33] applied multilinear dynamical systems to the latent core tensor obtained from Tucker decomposition of tensor time series data, similarly to [17]. Bahador et al. [2] enforced a low rank assumption on coefficient tensor of vector autoregressive models, and used a spatial Laplacian regularization for prediction in spatio-temporal data. [6] developed a probabilistic temporal tensor decomposition that models temporal dynamics in latent factor using a multilinear Gaussian distribution with a multilinear transition tensor (using multilinear dynamical system) with additional contextual constraints.
In contrast to the the methods above, which provide point estimates for imputation/prediction tasks, Bayesian probabilistic matrix/tensor factorization methods have been proposed (see [34,45,8]). In these approaches, latent factors have Gaussian priors, and Markov chain Monte Carlo (MCMC) methods are used for approximate inference in training and imputation. However, these models essentially focus on global matrix/tensor factorization without explicitly modeling the local temporal and spatial dependencies between factors. Hence, linear temporal dymanics on factor latents have been adapted to Bayesian Gaussian tensor factorization in [43,7,38]. While some of these methods have been effective in dynamical modeling of multidimensional time series, they do not explicitly adopt any functional form for the spatial domain, which makes them less effective for high dimensional spatial data such as fMRI. Moreover, the linear dynamical assumptions in these methods lack the capacity to characterize complex nonlinear dependencies.
Motivated by recent advances in deep learning, several papers have studied incorporation of neural networks into Gaussian state space models for nonlinear state space modelling [21,42,22,13,18,3]. A common practice is to learn a low-dimensional temporal generative model with an encoding scheme for performing amortized inference. However, these models are not applicable to high dimensional spatio-temporal data without factorization assumptions, and are not tractable in the case of missing data.
A number of fMRI-specific generative models have been proposed in which spatial factors are parameterized by Gaussian kernels to enhance interpretability [24,25,35]. The most closely related of these is NTFA [35], which represents the neuroimaging data for each trial, Y n , in terms of an embedding for the participant, Z p(n) , and an embedding for the stimulus, Z s(n) , which allows the model to capture commonalities in trials associated with a particular participant and stimulus. However, NTFA and other TFA variants assume a prior in which weights are conditionally independent as a function of time, which means that these models do not explicitly encode temporal dynamics. In this sense, DMSTF, which does reason about temporal dynamics, but does not explicitly model the interaction between participants and stimuli, is complementary to NTFA.

Deep Markov Spatio-Temporal Factorization
Model Structure and Variational Inference: DMSTF defines a hierarchical deep generative model for a corpus of instances {Y 1 , . . . Y N }, Here p θ (Z) is a deep generative prior over a low-dimensional set of local (instancelevel) variables Z n . The temporal weights W n are sampled from a Gaussian distribution that is parameterized by neural networks µ w θ and σ w θ . The spatial factors are defined as a deterministic transformation Φ(H n ) of a set of factor parameters H n , which are sampled from a distribution that is paremeterized by neural networks µ f θ and σ f θ . All networks have parameters, which we collectively denote by θ.
We train this model using stochastic variational methods [14,30,20,31]. These methods approximate the posterior p θ (W, H, Z | Y ) using a variational distribu- By maximizing this bound with respect to the parameters θ, we learn a deep generative model that defines a distribution over datasets p θ (Y ), which captures correlations between multiple instances Y n of the training data. By maximizing the bound with respect to the parameters φ, we perform Bayesian inference by approximating the distribution q φ (W, H, Z) p θ (W, H, Z | Y ) over latent variables for each instance.
Model Variants: We will develop three variants of DMSTF, which differ in the set of latent variables Z that they employ. The graphical models for these model variants are shown in Fig. 1. All three models incorporate a deep generative Markovian prior to represent temporal variation in W . To do so, they introduce a set of time-dependent weight embeddings z w t to define the distribution on each row w t of the weight matrix W . Moreover, the K rows of the factor matrix are sampled from a shared prior, which is conditioned on a factor embedding z f . The models differ in the following ways: - Fig. 1(a): This model simplifies the structure described above by assuming that the factor parameters H n and embeddings z f n do not vary across instances, but are instead shared at the corpus-level. Markovian Temporal Latent: We assume a Gaussian distribution for the latent transition probability. Given the latent z w t−1 we parameterize the mean and covariance of the diagonal Gaussian distribution p θ (z w t |z w t−1 ) using a neural network. In the model in Fig. 1(c) this distribution p θ (z w t |z w t−1 , u t−1 ) is additionally conditioned on a control variable u t−1 . Concretely, we pass z w t−1 (and u t−1 when applicable) to a multilayer perceptron (MLP) for estimating the Gaussian parameters. We combine a linear transformation of z w t−1 with the estimated mean from the neural network to support both linear and nonlinear dynamics [22]: where L θ (·) is a linear mapping, F θ (·) is the nonlinear mapping of neural network, and g ∈ [0, 1] is a scaling vector which itself is estimated from z w t−1 using a neural network. We show the neural network architectures we employ in Table 1.
Clustering Latent: In the models in Fig. 1(b) and Fig. 1(c), we assume that each sequence Y n belongs to a specific state out of S possible states, and is determined by the categorical variable c n in our temporal generative model. This is sampled from a categorical distribution c n ∼ Cat(π), where π = [π 1 , · · · , π S ] specifies cluster assignment probabilities. To this end, we assume that the the first temporal latent z w n,0 is distributed according to a Gaussian mixture: where the cluster assignment c n enforces µ s and diagonal covariance Σ s .

Temporal & Spatial Factors:
As with the transition model, we assume Gaussian distributions for temporal weights, and spatial factors. We parameterize the diagonal Gaussian distributions for temporal weights p θ (w t |z w t ) and factor parameters p θ (H | z f ) with neural networks. z f itself is sampled from a normal distribution: z f ∼ N (0, I). Introducing z f as a low dimensional spatial embedding in the model encourages estimation of a multi-modal distribution among spatial factors.
The form of the spatial factor parameters H depends on the task at hand. In the case of fMRI data, we employ the construction used in TFA [25,24,35], which represents each factor as a radial basis function with parameters H k = {ρ k , γ k }: This parameterization represents each factor as a Gaussian"blob". The vector r d ∈ R 3 denotes the position of voxel with index d. The parameter ρ k ∈ R 3 denotes the center of the Gaussian kernel, whereas γ k ∈ R controls its width.
In the case of traffic forecasting experiments, we learned spatial factors without any functional form constraints, hence, H directly parameterized {F kd } K,D k=1,d=1 by mean and covariance of a Gaussian distribution (i.e., Φ(·) is basically an identity mapping here).
Variational Distribution: We introduce trainable variational parameters (mean and sigma of a diagonal Gaussian) for each data point in our dataset to define a fully-factorized (i.e. mean-field) variational distributions on the variables We use a structured variational distribution for the variables q φ (z w 1:T | w 1:T ) in the form of a one-layer bidirectional recurrent neural network (BRNN) with ReLU activation, which is then combined with z w t−1 through another neural network to form distribution parameters of z w t : where f φ is a nonlinear mapping parameterized by an MLP, and Σ is diagonal. Although we can define variational parameters for categorical distributions q(c n ), we approximate it with posterior p(c n |z w n,0 ), as suggested in [16], to compensate information loss induced by the mean-field approximation: .
Kullback Leibler Divergence Terms: A significant advantage of our formulation is that we can analytically calculate the KL terms of evidence lower bound for both multivariate Gaussian, and categorical distributions, which leads to lower variance gradient estimates, and faster training as compared to e.g., noisy Monte Carlo estimates often used in literature.

Training Details
We describe the network architectures for the neural networks used in DMSTF in Table 1, where D z is the dimension of z w and z f , and D t and D e are the dimensions of hidden layers for Markovian and temporal latents respectively. For all the experiments in this paper, we assume D z = 2 (except for traffic dataset). We did all the programming in PyTorch v1.3 [28], and used the Adam optimizer [19] with learning rate of 1 × 10 −2 . We initialized all the parameters randomly except for spatial locations of fMRI data which we set the initial values to the local extrema in their averaged fMRI data. We clipped spatial locations and scales to the confines of the brain if needed. We used KL annealing [5] to suppress KL divergence terms in early stages of training, since these terms could be quite strong in the beginning, and we do not want them to dominate the log likelihood term (which controls reconstruction) in early stages. We used a linear annealing schedule to increase from 0.01 to 1 over the course of 100 epochs. The number of learnable parameters in our model is dominated by w t parameters, and therefore will be O(KNT). We run and learn all of the models on an Intel Core i7 CPU @3.7 GHz with 8 Gigabytes of RAM, which proves tractability of the learning process. Per-epoch training time varied from 0.1 in small datasets to 6.0 minutes in larger experiments.
Test/Prediction: We evaluate the performance of DMSTF for short-time prediction tasks (one-point-ahead forecasting) by adopting a rolling prediction scheme as in [8,44]. We predict observations for the next time point using the generative model and spatial factors:Ŷ t+1 =ŵ t+1 F , whereŵ t+1 ∼ p(ŵ t+1 |ẑ t+1 ), andẑ t+1 ∼ p(ẑ t+1 |z t ). Then, we fix the parameters of the generative model and variational parameters for spatial factors, and update variational parameters for z t+1 , w t+1 using Y t+1 , the actual observation at t + 1 (running model for a few epochs). We repeat these steps to make predictions in a rolling manner across a test dataset.

Toy Example
We generated N = 100 synthetic spatio-temporal data with T = 15, K = 2 using a nonlinear dynamical model (motivated by [22]): where ρ = 0.2, α = 0.5, β = −0.1. For spatial factors, F , we picked two Gaussian blobs centered at ±(7.5, 7.5, 7.5) with scales of 3, 4.5 respectively in a box of 30 × 30 × 30 at origin. And finally generated {Y n = W n F } N n=1 with additive noise. We train DMSTF with this synthetic dataset, and estimate the parameters of the model given the functional forms of the generative model. As depicted in Fig. 2 (a) DMSTF is able to recover the actual values of the parameters.

Synthetic Data
We generate synthetic fMRI data using a MATLAB package provided by [25], which is known to be useful for analysing fMRI models. The synthesized brain image for each trial (time point) is a weighted summation of a number of radial basis functions (spatial factors) randomly located in the brain. The synthesized fMRI data is then convolved with a Hemodynamic Response Function (HRF), and finally we add zero-mean Gaussian noise with a medium-level signal-tonoise ratio. Here, we considered 30 activation sources (spatial factors) randomly located in a standard MNI-152-3mm brain template with roughly 270, 000 voxels, and 150 trials. We randomly split these 30 activation sources into 3 groups, each having 10 of the Gaussian blobs. These three groups of sources are periodically activated in turn (according to some random weights) for 5 trials. We generated non-overlapping sequences of T = 5 time points from this synthetic fMRI data. This resulted in 10 data points for each activation group (N = 30). In order to train DMSTF, we set T = 5, K = 30, D t = 2, D e = 8, S = 3, and σ 0 = 1 × 10 −2 . As depicted in Fig. 2 (b), our model was able to successfully recover the 3 clusters of activation that were present in the dataset in a totally unsupervised manner. Predictions of learned generative model are visualized in Fig. 2 (d), which proves that we were able to capture nonlinear hemodynamic response function by using multi-layer perceptrons in our temporal generative model.

Neuroimaging Datasets
Autism Dataset: We used the publicly available preprocessed resting state fMRI (rs-fMRI) data from the Autism Brain Imaging Data Exchange (ABIDE) collected at 16 international imaging sites [12]. This dataset includes rs-fMRI imaging from 408 individuals suffering from Autism Spectrum Disorder (ASD), and 476 typical controls. Each scan has T = 145 ∼ 315 time points at TR = 2, and D = 271, 633 voxels. We split the signals into sequences of 75 time points. We take two approaches to evaluate the performance of our model in separating ASD from control: (1) Cluster data directly in the low dimensional latent, z w , using the clustering feature of DMSTF (called DMSTF), (2) Extract functional connectivity matrices from learned weights, W , followed by a 10fold SVM for classification (called DMSTF+SVM). As a baseline, we performed a 10-fold SVM classification on extracted connectivity matrices from averaged signals of 116 regions of interest (ROIs) in AAL atlas [40] (called ROI+SVM). Several studies have been done on this dataset to differentiate ASD group from control, all of them using supervised methods, and could achieve accuracies up to 69% (by carefully splitting data to be as homogeneous as possible, and reducing site-related variability) using the signals extracted from anatomically labeled regions in the brain [1,27,36]. We set T = 75, K = 100, D e = 15, D t = 5, S = 2, σ 0 = 1 × 10 −2 , and trained DMSTF for 200 epochs on the entire dataset (Full), and also datasets from 9 sites (with more balanced datasets) separately: Caltec, Leuven, MaxMun, NYU, SBL, Stanford, UM, USM, Yale. As shown in Fig. 3 (c) DMSTF and DMSTF+SVM outperform ROI+SVM in Caltec, MaxMun, SBL, Stanford, and Yale, while ROI+SVM only performs better in NYU dataset. DMSTF+SVM performs slightly better than ROI+SVM on the entire dataset (Please note that DMSTF is a clustering approach, hence, no error bars are provided in Fig. 3 (c)). Clustering results for Caltec, Maxmun, SBL, and Stanford are shown in Fig. 3 (a) in which ASD, and control seems to be partially separable (see Fig. S1 in Appendix B for more visualization results).
Depression Dataset: In this dataset [23], 19 individuals with major depressive disorder (MDD) and 20 never-depressed (ND) control participants listened to standardized positive and negative emotional musical and nonmusical stimuli during fMRI scanning. Each participant underwent 3 musical, and 2 nonmusical runs each for 105 time points at TR=3 with D = 353, 600 voxels. During each run, each stimulus type (positive, and negative) was presented for 33 seconds (∼ 11 time points) interleaved with instances of neutral tone of the same length. We discard instances of neutral tone, and split each run into non-overlapping sequences of T = 6 time points in agreement with stimuli design (each stimuli block is split into two sequences). In other words, each run has 4 sequences associated with "positive stimuli", and 4 with "negative stimuli" resulting in a total of 8 data points for each run. In the first experiment, we trained DMSTF on the entire musical runs (N = 39 × 3 × 8 = 936) by setting T = 6, K = 100, D e = 15, D t = 5, σ 0 = 1×10 −3 for 200 epochs. The results are shown in Fig. 4 (a,  Left). We observed that DMSTF fully separated data points associated with each subject into distinct clusters across the low-dimensional temporal latent space. In other words, DMSTF was able to re-unite pieces of signals associated with each subject without any kind of supervision. More importantly, DMSTF was able to partially separate data points associated with MDD group from control. As seen in Fig. 4 (a, Left), MDD group data points are fairly populated in the center of temporal latent while control group are dispersed across latent space. However, DMSTF was not able to meaningfully separate "negative" and "positive" music pieces in latent embedding from a subject-level perspective, since the variation between runs of a subject dominates stimulus-level variation. For this reason, in a second experiment, we focused on 5 subjects, and their first musical run from both MDD and control group and trained DMSTF respectively. Again, as expected, data points from each subject were distinctly clustered in latent space (see middle and right columns in Fig. 4 (a)). Additionally, DMSTF was able to fit two partially separating Gaussians to "positive", and "negative" stimuli per subject. However, since the number of data points for each subject and run is Table 2. Held-out Log-Likelihood. DMSTF results in models with higher held-out likelihood, and therefore better fit comparing to HTFA.

Dataset
HTFA DMSTF Autism (Caltech) −2.82 × 10 6 −2.33 × 10 6 Depression −6.64 × 10 5 −5.71 × 10 5 limited it is not clear how significant these clusters are. A dataset with longer runs could possibly answer that. In a third experiment, we incorporate control inputs u t , and evaluate prediction performance of DMSTF. We train DMSTF on 2 musical and 1 nonmusical runs for a subject with depression (8 × 3 = 24 sequences) using u t = 1, and u t = 0 respectively (i.e., u t is encoding musical, nonmusical stimuli). We predict the remaining musical and nonmusical runs (8 × 2 = 16 sequences) once with u t = 0, and another time with u t = 1.
As reported in Fig. 4 (d), nonmusical sequences are better predictable with u t = 0 than u t = 1 with p-value= 0.011 (vice versa for musical sequences with p-value=0.036). Sample predicted fMRI time series from both musical, and nonmusical runs are shown in Fig. 4 (c) (see Fig. S2 in Appendix B for more visualization results).
Comparison with HTFA [24]: We further evaluated DMSTF against HTFA, an established probabilistic model for multi-subject fMRI analysis, which uses unimodal Gaussian priors for both temporal weights, and spatial factor parameters, in terms of held-out log-likelihood (see Table 2). For autism, we used Caltec site dataset, and split each subject's fMRI into two half (each T = 70). We trained DMSTF on the first half, and test on the second half. For depression, we consider 4 sequences from each subject's run for training, and test on the remaining 4 sequences. To this end, after training DMSTF on each training set, we fixed all parameters of the generative model and only tuned variational parameters of the test set for temporal latents z w t , w t . And finally computed an importance sampling-based estimate of the log-likelihood [22]. The results are shown in Table 2, which proves that DMSTF results in models with higher likelihood on test set, hence a better fit as compared to HTFA.  . We picked the data from January 1 to January 28 (28 days) as in [8], and organized it into a tensor of 28 × 288 × 323. We compared DMSTF (variant (a)) with three state-of-the-art baselines on our short-time prediction task: TRMF [44], BayesTRMF developed in [8], and BTMF [8]. We picked the last seven days from the Birmingham dataset, and the last five days from the Guangzhou, Hangzhou, and Seattle datasets for prediction, then trained the models on the rest for each dataset with K=10, 30, 10, 30 respectively (consistent setup with [8]). For DMSTF, we additionally set {D z , D t , D e } = 5, σ 0 = 0 for all datasets, and learned spatial factors without any functional form constraints. We trained DMSTF for 500 epochs. We report root mean squared error (RMSE), and mean absolute percentage error (MAPE) for all models on the testsets in Table 3. DMSTF outperforms in short-time prediction on all datasets, doing significantly better particularly on the Birmingham dataset. Testset predictions for two sample locations from each dataset are shown in Fig. 5 (see Fig. S3 in Appendix B for a number of long-term prediction visualizations).

Conclusion and Future Work
We presented deep Markov spatio-temporal factorization, a new probabilistic model for robust factor analysis of high-dimensional spatio-temporal data. We employed a chain of low-dimensional Markovian latent variables connected by deep neural networks as a state-space embedding for temporal factors in order to model nonlinear dynamics in data, account better for noise and uncertainty, and enable generative prediction. We also employed a low-dimensional spatial embedding to generate a multi-modal distribution of spatial factors. We then demonstrated the tractability of DMSTF on fMRI data (with very high spatial dimensionality) by incorporating functional form assumptions, and on traffic data with high temporal dimensionality. DMSTF enables clustering in the lowdimensional temporal latent space to reveal structure in data (e.g., cognitive states in fMRI), providing informative visualizations about the data. We plan to extend DMSTF to accommodate higher order dynamics for long-term prediction tasks. We can readily achieve this in our setting by conditioning temporal latents on a time-lag set, such as conditioning z w t on z w t−1 , z w t−2 in a 2-Markov chain. A Derivation of Evidence Lower BOund (ELBO) for DMSTF We derive the lower bound on the log-likelihood of observations by incorporating conditional independences inferred from the graphical model of DMSTF in Fig. 1 (c) as follows: Y n ⊥ ⊥ Y ¬n |w n,1:T , H n , z w n,t ⊥ ⊥ z w n,¬(t,t−1) |z w n,t−1 , z w n,t ⊥ ⊥ z w ¬n |z w n,t−1 z w n,0 ⊥ ⊥ z w ¬n,0 |c n , w n,t ⊥ ⊥ w n,¬t |z w n,t , w n,t ⊥ ⊥ w ¬n |z w n,t , H n ⊥ ⊥ H ¬n |z f n Considering these conditional independencies, the joint distribution of observations and latent variables will be (denoting Z = {z w , z f }): We assume a fully factorized variational distribution on {c, z f , z w 0 , w}, and a structured variational distribution on q φ (z w 1:T |w 1:T ), hence: We then derive the ELBO by writing down the log-likelihood of observations, and plugging in p(·) and q(·) from Equation (1) and Equation (2) p(Y n |w n,1:T , H n )p(H n |z f n )p(z f n )p(z w n,0 |c n )p(c n ) q(H n )q(z f n )q(z w n,0 )q(c n ) · · · T t=1 p(w n,t |z w n,t )p(z w n,t |z w n,t−1 , u n,t−1 ) q(w n,t )q(z w n,t |z w n,t−1 , w n,1:T ) · · · z ={w n,1:T ,Hn} c =∅ q(w n,1:T , H n ) log p(Y n |w n,1: q(w n,1:T )q(z w n,t−1 |w n,1:T ) · · · q(z w n,t |z w n,t−1 , w n,1:T ) log p(z w n,t |z w n,t−1 , u n,t−1 ) q(z w n,t |z w n,t−1 , w n,1:T ) dz + T t=1 z ={z w n,t ,wn,t} c =∅ q(z w n,t )q(w n,t ) log p(w n,t |z w n,t ) q(w n,t ) dz where the integration over distributions of latent variables outside of the summation over n is c\c q(c\c ) Z\z q(Z\z )dZ\z = 1 for {c , z } of each term inside the summation (we are abusing notation here, for the sake of brevity). Considering that KL(q, p) = q log q p , and E q(x) [f (x)] = x q(x)f (x), we can rewrite each term of the summation in Equation (3) to summarize the ELBO: where, L rec n = E q φ (w n,1:T ,Hn) [log p θ (Y n |w n,1:T , H n )] L C n = −KL(q φ (c n )||p θ (c)) − cn q φ (c n )KL(q φ (z w n,0 )||p θ (z w n,0 |c n )) L z w n,t = −E q φ (w n,1:T ) E q φ (z w n,t−1 |w n,1:T ) KL q φ (z w n,t |z w n,t−1 , wn,1:T ) p θ z w n,t |z w n,t−1 , un,t−1) L W n,t = −E q φ (z w n,t ) [KL(q φ (w n,t )||p θ (w n,t |z w n,t ))] Kulback Leibler (KL) Divergence Terms: We can analytically calculate the KL terms of ELBO in Equation (4). For two multivariate (d-dimensional) Gaussian distributions q(·), and p(·), the KL divergence is: which would be further simplified here, in the case of diagonal covariances. The KL divergence between two S-dimensional categorical distributions q(·) and p(·) can be readily computed: where the subscript s denotes the probability of state s out of S possible states.

B More Visualizations for Autism, Depression, and Traffic Datasets
We have visualized real and reconstructed brain images from the nine subsets of autism dataset (Caltex, Leuven, MaxMun, NYU, SBL, Stanford, UM, USM, and Yale sites) along with z w 0 after training DMSTF on the full autism dataset in Fig. S1. DMSTF clustered together temporal latent variables associated with each acquisition site without supervision in z w 0 . As depicted, the variation among different acquisition sites dominates the cognitive differences between ASD group and control, hence, a downstream connectivity matrix classification (using learned temporal weights, W ) helps better in differentiating ASD group from control in multi-site analysis. In Fig. S2, we have visualized example predicted fMRI time-series from both musical and non-musical trials in the test-set of depression dataset using control variable u t = 1 for musical and u t = 0 for non-musical trials. In Fig. S3, we have visualized next-day (long-term) prediction results for Birmingham and Huangzhou subsets of traffic data for four sample locations. These predictions are purely obtained from the trained temporal generative model. Please note that the actual values for the predicted days are not available. showing the smoothing given by sparse factorization. Visualizing z w 0 after training DMSTF on the full autism dataset. DMSTF clustered together temporal latent variables associated with each acquisition site without supervision. As depicted, the variation among different acquisition sites dominates the variation in cognitive state of the brain (ASD group vs. control), hence, a downstream connectivity matrix classification helps better in differentiating ASD group from control in multi-site analysis.   S3. Visualizing next-day (long-term) prediction results for Birmingham and Huangzhou subsets of traffic data for four sample locations. Although DMSTF is wellsuited for short-time prediction, next-day forecasts (purely from the trained temporal generative model) show its capability in long-term predictions. Please note that the actual values for the predicted days are not available.