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: The paper presents a deep Markov factor analysis (DMFA), which is a generative model that employs Markov property in a chain of low dimensional temporal embeddings as well as spatial inductive assumptions in order to capture temporal dynamics in functional magnetic resonance imaging (fMRI). The paper shows that DMFA has capability to cluster fMRI data in its low temporal embedding in regards to subject and cognitive state variability. The advantages of DMFA are demonstrated both through synthetic and application data.