Deep Markov Factor Analysis: Towards Concurrent Temporal and Spatial Analysis of fMRI Data

1 Introduction . Functional magnetic resonance imaging ( fMRI ) has been extensively used in cognitive neuroscience to study brain structures and their interactions at rest or during a cognitive task [ Glover , 2011 ] . fMRI also provides insights into how brain ’ s functional connectivity changes during different experimental conditions [ Preti et al. , 2017 , Azari et al. , 2020 ] . However , due to high dimensional nature of fMRI data ( tens of thousands of voxels in few seconds-long sessions ) , analyzing functional connectivity of brain could get very challenging [ Turk-Browne , 2013 ] . Most classical methods employ region of interest ( ROI ) -based approaches to reduce size of data for processing [ Poldrack , 2007 , Etzel et al. , 2009 , Farnoosh and Soltanian-Zadeh , 2017 , Gadgil et al. , 2020 ] . However , averaging across many voxels within each region could wash out signals from small number of task-relevant voxels with noise from non-activated voxels , therefore results in a loss of information [ Poldrack , 2007 ] . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) Recently , few approaches have been proposed , based on probabilistic generative models , for topographic factorization of fMRI data into a weighted summation of few localized activation sources ( i.e. , temporal weights and topographic spatial factors ) , among which topographic factor analysis ( TFA ) [ Manning et al. , 2014b ] and its multi-subject extension , hierarchical TFA ( HTFA ) [ Manning et al. , 2018 ] , and neural TFA ( NTFA ) [ Sennesh et al. , 2020 ] are the most noted ones . This factorization serves as a necessary preparation step for subsequent statistical analysis that can effectively characterize subject- and stimulus-level variations and reveal task- or cognitive state-related networks in brain . However , TFA approaches assume a prior in which temporal weights are conditionally independent as a function of time , which means they do not encode temporal dynamics . Given the non-linearity and complex time-dependencies inherent in fMRI , a model is required that can capture and represent these dependencies . In this paper , we propose deep Markov factor analysis ( DMFA ) 1 , a Bayesian model for factorization of fMRI data that learns a deep generative Markovian prior to reason about nonlinear temporal dynamics . This is realized by a chain of low dimensional temporal embeddings related through neural networks . This prior is further augmented by a discrete latent for multimodal dynamical estimation , and clustering subject- and task-level variations directly in its low dimensional temporal embedding . To accommodating high spatial dimensionality , DMFA generatively parameterize spatial factors from a low dimensional spatial latent through neural networks . We evaluate the performance of DMFA on a synthetic and two real large-scale fMRI datasets . Our experiments demonstrate that DMFA uncovers meaningful clusters in these data and achieves better predictive performance for unseen data relative to the state-of-the-art . 2 Related Work . Factor Analysis in fMRI : Factor analysis in neuroimaging includes a wide range of approaches for reducing data dimensionality to facilitate their interpretability and computational tractability . Principal component analysis ( PCA ) [ Pearson , 1901 ] and independent component analysis ( ICA ) [ Comon et al. , 1991 ] are among the most well-known classical factor analysis methods . To accommodate tensor data and mitigate scalability issues , multilinear versions of PCA and ICA have been proposed in Vasilescu and Terzopoulos [ 2005 ] , Beckmann and Smith [ 2005 ] , Cichocki [ 2011 ] , Richard and Montanari [ 2014 ] , Hopkins et al . [ 2015 ] . Specifically , for multi-subject fMRI study , Lee et al . [ 2008 ] proposed independent vector analysis ( IVA ) and Richard et al . [ 2020 ] developed MultiView ICA to model shared responses . Likewise , Chen et al . [ 2015 ] developed a shared response model ( SRM ) for aggregating multi-subject fMRI data and highlighting group differences , and Van Kesteren and Kievit [ 2021 ] incorporated structured residuals into the exploratory factor analysis ( EFA ) framework . Karahanoğlu and Van De Ville [ 2015 ] deconvolved hemodynamic response from rest fMRI time series and then performed temporal clustering on the resulting whole-brain innovation signals to recover the corresponding spatial patterns . However , spatial factors obtained by these methods are unstructured , often have the same size as the images in the original dataset , and may include many small and large voxel clusters across the brain , therefore are not directly interpretable [ Manning et al. , 2014b ] . More important , these methods are permutation invariant along temporal dimension ( i.e. , do not assume any relationships between temporal dimensions ) , therefore do not model temporal dynamics [ Yu et al. , 2016 ] . To enhance spatial interpretability , Manning et al . [ 2014b ] introduced topographic factor analysis ( TFA ) , a probabilistic technique that casts each brain image as a weighted sum of Gaussian spatial factors . In this way , each Gaussian blob ( a.k.a topographic factor ) can be thought of as a brain region for which activation levels are estimated over time , and is easily interpreted through its set of parameters . Later on , authors proposed a hierarchical version of TFA ( HTFA ) [ Manning et al. , 2018 ] to incorporate data from multiple subjects and enable hypothesis testing across subjects by applying hierarchical Gaussian priors for global and subject-specific parameters of temporal and topographic factors . These linearly-dependent hierarchical Gaussian priors , though , favor estimation of unimodal distributions and limit the model ’ s expressivity . Recently , Sennesh et al . [ 2020 ] proposed neural TFA ( NTFA ) for task fMRI , which extends TFA by incorporating neural networks onto its framework . NTFA assumes separate latent embeddings for participants and stimuli and map them into the temporal and spatial latents with neural networks . However , methods in this line of work 1The source code is available at https : //github.com/ostadabbas/DMFA essentially assume a prior in which temporal weights are conditionally independent as a function of time , which means that they do not encode temporal dynamics . Dynamical Factorization : A number of matrix/tensor factorization approaches have been proposed for modeling temporal dynamics in sequential data . These methods employ linear-Gaussian statespace models [ Sun et al. , 2014 ] , autoregressive temporal regularizer [ Bahadori et al. , 2014 , Yu et al. , 2016 , Takeuchi et al. , 2017 ] , and multilinear dynamical systems [ Rogers et al. , 2013 , Cai et al. , 2015 , Jing et al. , 2018 ] . In contrast to these methods which provide point estimates , Bayesian dynamical factorization has been proposed in Xiong et al . [ 2010 ] , Charlin et al . [ 2015 ] , Sun and Chen [ 2019 ] , in which linear temporal dynamics are applied to factor latents . While some of these methods have been successful in dynamical modeling of sequential data , they are less effective for high dimensional spatial data like fMRI as they do not explicitly adopt any structural constraints for spatial factors . Moreover , their linear dynamical assumptions lack the capacity to characterize nonlinear dependencies . Motivated by recent advances in deep learning , several studies have implemented neural networks into Gaussian state space models for nonlinear dynamical modelling [ Krishnan et al. , 2015 , Watter et al. , 2015 , Chung et al. , 2015 , Karl et al. , 2017 , Krishnan et al. , 2017 , Fraccaro et al. , 2017 , Becker et al. , 2019 , Farnoosh et al. , 2021 ] . A common practice is to learn a temporal latent model followed by a mapping to data space , a.k.a . a decoder , and an encoder for amortized variational inference . However , this encoding/decoding framework is not tractable in very high dimensional spatial data like fMRI . To be more specific , it is computationally intensive to feed high dimensional data for amortized estimation , or map to them directly from a latent space ( see Section 4.1 for a detailed discussion ) . Also , several studies have employed recurrent neural networks ( RNNs ) for temporal analysis of fMRI data [ Hjelm et al. , 2018 , Dvornek et al. , 2018 , Yan et al. , 2019 , Wang et al. , 2019 ] . RNN-based methods are not probabilistic and can only process reduced dimensional ROI data . Although convolutional RNN models can capture both temporal and spatial correlations , by construction they do not provide the familiar spatial correlation maps that we expect in neuroimaging analysis and that are necessary for neuroscientific interpretability . In addition , RNN model are harder to train on long sequences due to the vanishing/exploding gradient problem as they feed parameters sequentially . While this problem is much alleviated in LSTM , GRU , and Transformer models , these architectures still result in a huge computational graph and intense GPU memory consumption on long sequences , which leads to very long training times . In contrast , DMFA , as will be explained below , samples from a distinct posterior distribution at each time point and can fit arbitrarily long sequences without any issues ( by allowing parallel computations over time ) , and at the same time infers smooth spatial factors automatically . 3 Deep Markov Factor Analysis ( DMFA ) . Let ’ s assume a corpus of N fMRI data { Yn |n ∈ [ 1 : N ] } for Yn ∈ RT×V , where T and V are the number of time points and voxels , respectively . DMFA defines a hierarchical dynamical deep generative model over this corpus in a factorization framework : Yn ∼ Norm ( W > n Fn , σYI ) , Wn ∼ Norm ( µWθ ( ZWn ) , σWθ ( ZWn ) ) , Fn = [ fk , v : = RBFk ( v ; ρk , γk ) ] , ρ , γ ∼ Norm ( µFθ ( zFn ) , σFθ ( zFn ) ) , zFn ∼ Norm ( 0 , I ) , ZWn ∼ pθ ( ZWn |Cn ) , Cn ∼ Cat ( π ) . where Fn ∈ RK×V are K V spatial factors and Wn ∈ RK×T are their associated temporal weights , consistent with a factorization framework . pθ ( ZWn |Cn ) is a deep generative Markovian prior over a set of low dimensional temporal latents ZWn = { zWn , t | t ∈ [ 0 : T ] } conditioned on a discrete latent Cn for clustering . The temporal weights Wn are sampled from a Gaussian distribution whose mean and covariance are a function of ZWn and are parameterized by neural networks µWθ and σ W θ , respectively . The spatial factors Fn are determined by a set of K radial basis functions ( RBFs ) whose parameters { ρ , γ } are sampled from a Gaussian distribution parameterized by neural networks µFθ and σ F θ from a low dimensional spatial latent z F n. All networks have parameters , which are collectively denoted by θ . Finally , σY denotes observation noise . The graphical model for DMFA is depicted in Fig . 1 . Clustering Latent : We assume that each fMRI data Yn belongs to a specific cluster s ∈ [ 1 : S ] which is declared by its associated discrete latent Cn ∼ Cat ( π ) sampled from a categorical prior with assignment probability vector π . This discrete latent conditions the first temporal latent zWn,0 resulting in a Gaussian mixture distribution over this temporal latent : pθ ( z W n,0|Cn = s ) = Norm ( µs , Σs ) , where µs and Σs are determined by cluster assignment . Deep Generative Markovian prior : We assume that the temporal latents ZWn are related through a Markov chain , for which the transition probability pθ ( zWn , t|zWn , t−1 ) = Norm ( µZθ ( zWn , t−1 ) , σZθ ( zWn , t−1 ) ) is a Gaussian distribution whose mean and covariance are parameterized by neural networks from zWn , t−1 . We further blend the estimated mean from neural network with a linear transformation of zWn , t−1 to support both nonlinear and linear transitions ( we dropped n for brevity ) : µZθ ( z W t−1 ) = ( 1− g ) µ Z , L θ ( z W t−1 ) + g µ Z , NL θ ( z W t−1 ) where µZ , Lθ ( · ) represents a linear mapping , µ Z , NL θ ( · ) is the nonlinear mapping of neural network , and g ∈ [ 0 , 1 ] is a weighting vector estimated from zWt−1 with another neural network . Temporal Weights & Spatial Factors : As with the transition model , we assumed Gaussian distributions for temporal weights , such that their means and covariances are a function of their associated temporal latents : wn , t ∼ Norm ( µWθ ( zWn , t ) , σWθ ( zWn , t ) ) , where µWθ and σ W θ are neural network functions . The high spatial dimensionality in fMRI data encourages the need for a hierarchical analysis that summarizes spatial factors with fewer parameters . Therefore , consistent with Manning et al . [ 2014b , a , 2018 ] , we represent each spatial factor as a radial basis function ( i.e. , a Gaussian blob ) : fkv : = RBFk ( v ; ρk , γk ) = exp ( − ‖pv − ρk‖ 2 exp ( γk ) ) , where pv ∈ R3 denotes the position of voxel with index v , and ρk , γk ∈ R3 denote the center and extent of the Gaussian blob , respectively . The set of all RBF parameters { ρk , γk | k ∈ [ 1 : K ] } ∼ Norm ( µFθ ( z F n ) , σ F θ ( z F n ) ) are sampled from a Gaussian distribution whose mean and covariance are a function of the shared spatial latent zFn . Introducing z F n as a low dimensional spatial embedding encourages estimation of a multi-modal distribution among spatial factors .

Summary: This paper introduces Deep Markov Factor Analysis (DMFA). DMFA can model spatiotemporal BOLD fMRI time series by using time-varying latent factors that govern time-varying spatial maps consisting of spatial mixtures of RBFs. Latent time series are associated to a categorical clustering variable in order to enable grouping. Certain links in the Bayesian diagram are modeled using deep neural networks. Optimization is done by a variational approach. Experiments are performed on a simulated data set with known ground truth, the ABIDE autism data set, and a Depression study data set. Some clustering of relevant clinical outcomes is observed.