Efficient hierarchical Bayesian inference for spatio-temporal regression models in neuroimaging

1 Introduction . Probabilistic graphical models for regression problems , where both the model parameters and the observation noise can have complex spatio-temporal correlation structure , are prevalent in neuroscience , neuroimaging and beyond [ 1 ] . Example domains where these models are being applied include Gaussian process inference [ 2 ] , sparse signal recovery and multi-task learning [ 3–6 ] , array signal processing [ 7 ] , fMRI analysis [ 8 ] , climate science [ 9 ] , computer hardware monitoring [ 10 ] , and brain-computer interfaces ( e.g . [ 11 , 12 ] ) . The focus in this paper will be on brain source imaging ( BSI ) [ 13 ] , i.e . the reconstruction of brain source activity from non-invasive magneto- or electroetoencephalography ( M/EEG ) sensors [ 14 ] . Although a number of generative models that encompass spatio-temporal correlations have been proposed within the BSI regression problem , inference solutions have either imposed specific simplifications and model constraints [ 15–17 ] or have ignored temporal correlations overall [ 5 ] , and therefore not fully addressed the inherent spatio-temporal problem structure . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . We propose a novel and efficient algorithmic framework using Type-II Bayesian regression that explicitly considers the spatio-temporal covariance structure in both model coefficients and observation noise . To this end , we focus on the M/EEG inverse problem of brain source reconstruction , for which we adopt a multi-task regression approach and formulate the source reconstruction problem as a probabilistic generative model with Gaussian scale mixture priors for sources with spatio-temporal prior covariance matrices that are expressed as a Kronecker product of separable spatial and temporal covariances . Their solutions are constrained to the Riemannian manifold of positive definite ( P.D . ) temporal covariance matrices . Exploiting the concept of geodesic convexity on the Riemannian manifold of covariance matrices with Kronecker structure , we are then able to derive robust , fast and efficient majorization-minimization ( MM ) optimization algorithms for model inference with provable convergence guarantees . In addition to deriving update rules for full-structural P.D . temporal covariance matrices , we also show that assuming a Toeplitz structure for the temporal covariance matrix results in computationally efficient update rules within the proposed MM framework . 2 Spatio-temporal generative model . Let us consider a multi-task linear regression problem , mathematically formulated as Yg = LXg + Eg for g = 1 , . . . , G , ( 1 ) in which a forward matrix , L ∈ RM×N , maps a set of coefficients or source components Xg ∈ RN×T to measurements Yg ∈ RM×T , with independent Gaussian noise Eg ∈ RM×T . G denotes the number of sample blocks in the multi-task problem , whileN , M , and T denote the number of sources or coefficients , the number of sensors or observations , and the number of time instants , respectively . The problem of estimating { Xg } Gg=1 given L and { Yg } Gg=1 can represent an inverse problem in physics , a multi-task regression problem in machine learning , or a multiple measurement vector ( MMV ) recovery problem in signal processing [ 18 ] . In the context of BSI , { Yg } Gg=1 refers to M/EEG sensor measurement data , and { Xg } Gg=1 refers to brain source activity contributing to the observed sensor data . Here , G can be defined as the number of epochs , trials or experimental tasks . The goal of BSI is to infer the underlying brain activity from the M/EEG measurements given the lead field matrix L. In practice , L can be computed using discretization methods such as the finite element method for a given head geometry and known electrical conductivities using the quasi-static approximation of Maxwell ’ s equations [ 19 ] . As the number of locations of potential brain sources is dramatically larger than the number of sensors , the M/EEG inverse problem is highly ill-posed , which can be dealt with by incorporating prior assumptions . Adopting a Bayesian treatment is useful in this capacity because it allows these assumptions to be made explicitly by specifying prior distributions for the model parameters . Inference can be performed either through Maximum-a-Posteriori ( MAP ) estimation ( Type-I Bayesian learning ) [ 20–22 ] or , when the model has unknown hyperparameters , through Type-II MaximumLikelihood ( ML ) estimation ( Type-II Bayesian learning ) [ 23–25 ] . Here we focus on full spatio-temporal Type-II Bayesian learning , which assumes a family of prior distributions p ( Xg|Θ ) parameterized by a set of hyperparameters Θ . We further assume that the spatio-temporal correlation structure of the brain sources can be modeled by a Gaussian p ( xg|Γ , B ) ∼ N ( 0 , Σ0 ) , where xg = vec ( X > g ) ∈ RNT×1 and where the covariance matrix Σ0 = Γ ⊗ B has Kronecker structure , meaning that the temporal and spatial correlations of the sources are modeled independently through matrices Γ and B , where ⊗ stands for the Kronecker product . Note that the decoupling of spatial and temporal covariance is a neurophysiologically plausible assumption as neuronal sources originating from different parts of the brain can be assumed to share a similar autocorrelation spectrum . In this paper , we set Γ = diag ( γ ) , γ = [ γ1 , . . . , γN ] > , and B ∈ ST++ , where ST++ denotes the set of positive definite matrices with size T × T . Note that using a diagonal matrix Γ amounts to assuming that sources in different locations are independent . This assumption may be relaxed in our future work . On the other hand , modeling a full temporal covariance matrix B amounts to assuming non-stationary dynamics of the sources , which is appropriate in event-related experimental designs . To deal with the more stationary sources ( ongoing brain activity ) , we will also constrain B to have Toeplitz structure later on . Based on these specifications , the prior distribution of the i-th brain source is modeled as p ( ( Xg ) i.|γi , B ) ∼ N ( 0 , γiB ) , for i = 1 , . . . , N , where ( Xg ) i. denotes the i-th row of source matrix Xg . Analogous to the definition of the sources , we here also model the noise to have spatio-temporal structure . The matrix Eg ∈ RM×T represents T time instances of zero-mean Gaussian noise with full covariance , Σe = Λ⊗Υ , where eg = vec ( E > g ) ∈ RMT×1 ∼ N ( 0 , Σe ) , and where Λ and Υ denote spatial and temporal noise covariance matrices , respectively . Here , we assume that the noise and sources share the same temporal structure , i.e. , Υ = B . We later investigate violations of this assumption in the simulation section . Analogous to the sources , we consider spatially independent noise characterized by diagonal spatial covariance Λ = diag ( λ ) , λ = [ λ1 , . . . , λM ] > . We consider both the general , heteroscedastic , case in which the noise level may be different for each sensor as well as the homoscedastic case in which the noise level is uniform . For later use , we also define augmented versions of the source and noise covariance matrices as well as the lead field . Specifically , we set H and Φ so that H : = [ Γ,0 ; 0 , Λ ] , and Φ : = [ L , I ] . These definitions allow us to unify source and noise covariance parameters within the single variable H , which facilitates the concurrent estimation of both quantities . Φ now plays the same role as the lead field L , in the sense that it maps ηg : = [ x > g , e > g ] to the measurements ; i.e. , yg = Φηg . Figure 1- ( a ) illustrates a probabilistic graphical model summarizing the spatio-temporal generative model of our multi-task linear regression problem . The MMV model Eq . ( 1 ) can be formulated equivalently in single measurement vector ( SMV ) form by vectorizing the spatio-temporal data matrices and using the Kronecker product as follows : yg = Dxg+eg , where yg = vec ( Y > g ) ∈ RMT×1 , and D = L⊗IT . In a Type-II Bayesian learning framework , the hyper-parameters of the spatio-temporal source model are optimized jointly with the model parameters { Xg } Gg=1 . In our case , these hyper-parameters comprise the unknown source , noise and temporal covariance matrices , i.e. , Θ = { Γ , Λ , B } . 3 Proposed method — full Dugh . The unknown parameters Γ , Λ , and B can be optimized in an alternating iterative process . Given initial estimates , the posterior distribution of the sources is a Gaussian of the form p ( xg|yg , Γ , Λ , B ) ∼ N ( x̄g , Σx ) , whose mean and covariance are defined as follows : x̄g = Σ0D > ( Λ⊗B + DΣ0D > ) −1yg = Σ0D > Σ̃−1y yg , ( 2 ) Σx = Σ0 −Σ0D > Σ̃−1y DΣ0 , ( 3 ) where Σy = LΓL > + Λ , and where Σ̃y = Σy ⊗B denotes the spatio-temporal statistical model covariance matrix . The estimated posterior parameters x̄g and Σx are then in turn used to update Γ , Λ , and B as the minimizers of the negative log of the marginal likelihood p ( Y|Γ , Λ , B ) , which is given by : Lkron ( Γ , Λ , B ) = T log |Σy|+M log |B|+ 1 G G∑ g=1 tr ( Σ−1y YgB −1Y > g ) , ( 4 ) where | · | denotes the determinant of a matrix . A detailed derivation is provided in Appendix A . Given the solution of the hyperparameters Γ , Λ , and B , the posterior source distribution is obtained by plugging these estimates into Eqs . ( 2 ) – ( 3 ) . This process is repeated until convergence . The challenge in this high-dimensional inference problem is to find ( locally ) optimal solutions to Eq . ( 4 ) , which is a non-convex cost function , in adequate time . Here we propose a novel efficient algorithm , which is able to do so . Our algorithm thus learns the full spatio-temporal correlation structure of sources and noise . Hashemi et al . [ 26 ] have previously demonstrated that majorizationminimization ( MM ) [ 27 ] is a powerful non-linear optimization framework that can be leveraged to solve similar Bayesian Type-II inference problems . Here we extend this work to our spatiotemporal setting . Building on the idea of majorization-minimization , we construct convex surrogate functions that majorize Lkron ( Γ , Λ , B ) in each iteration of the proposed algorithm . Then , we show the minimization equivalence between the constructed majorizing functions and Eq . ( 4 ) . These results are presented in theorems 1 and 3 . Theorems 2 and 4 propose an efficient alternating optimization algorithm for solving Lkron ( Γ , Λ , B ) , which leads to update rules for the spatial and temporal covariance matrices Γ and B as well as the noise covariance matrix Λ . Starting with the estimation of the temporal covariance based on the current source estimate , we can state the following theorem : Theorem 1 . Optimizing the non-convex Type-II cost function in Eq . ( 4 ) , Lkron ( Γ , Λ , B ) , with respect to the temporal covariance matrix B is equivalent to optimizing the following convex surrogate function , which majorizes Eq . ( 4 ) : Ltimeconv ( Γk , Λk , B ) = tr ( ( Bk ) −1B ) + tr ( MktimeB −1 ) , ( 5 ) where Mktime is defined as Mktime : = 1 MG G∑ g=1 Y > g ( Σky ) −1 Yg , ( 6 ) and where Γk , Λk , and Σky denote the source , noise and statistical model covariance matrices at the k-th iteration , which are treated as constants when optimizing over B . Proof . A detailed proof is provided in Appendix B . The solution of Eq . ( 5 ) can be obtained by taking advantage of the Riemannian geometry of the search space . The following theorem holds . Theorem 2 . The cost function Ltimeconv ( Γk , Λk , B ) is strictly geodesically convex with respect to the P.D . manifold and its minimum with respect to B can be attained by iterating the following update rule until convergence : Bk+1 ← ( Bk ) 1/2 ( ( Bk ) −1/2Mktime ( B k ) −1/2 ) 1/2 ( Bk ) 1/2 . ( 7 ) Proof . A detailed proof is provided in Appendix C. Remark 1 . A geometric interpretation of the update rule ( 7 ) is that it finds the geometric mean between a spatially whitened version of the empirical temporal covariance matrix , Mktime , and the model temporal covariance matrix from the previous iteration , Bk . A geometric representation of the geodesic path between the pair of matrices { Bk , Mktime } on the P.D . manifold and the geometric mean between them , representing the update for Bk+1 , is provided in Figure 1- ( b ) . We now derive an update rule for the spatial source and noise covariance matrices using an analogous approach . To this end , we first construct a convex surrogate function that majorizes Lkron ( Γ , Λ , B ) in each iteration of the optimization algorithm by considering augmented variables comprising both the source and noise covariances , i.e. , H : = [ Γ,0 ; 0 , Λ ] , and Φ : = [ L , I ] . The following theorem holds . Theorem 3 . Minimizing the non-convex Type-II cost function in Eq . ( 4 ) , Lkron ( Γ , Λ , B ) , with respect to the spatial covariance matrix H is equivalent to minimizing the following convex surrogate function , which majorizes Eq . ( 4 ) : Lspaceconv ( Γ , Λ , Bk ) = Lspaceconv ( H , Bk ) = tr ( Φ > ( Σky ) −1ΦH ) + tr ( MkSNH −1 ) , ( 8 ) where MkSN is defined as MkSN : = H kΦ > ( Σky ) −1Mkspace ( Σ k y ) −1ΦHk , with Mkspace : = 1 TG G∑ g=1 Yg ( B k ) −1Y > g , ( 9 ) and where Bk denotes the temporal model covariance matrix estimated in the k-th iteration . Proof . A detailed proof is provided in Appendix D. While , in principle , update rules for full-structured spatial source and noise covariances may be conceived , in analogy to the ones presented for the temporal covariance , we here restrict ourselves to the discussion of diagonal spatial covariances . Note , though , that such a choice of prior does not prohibit the reconstruction of parameters with more complex correlation structure . We have H = diag ( h ) , h = [ γ1 , . . . , γN , σ21 , . . . , σ 2 M ] > , and Σy = ΦHΦ > . The update rule for H then takes a simple form , as stated in the following theorem : Theorem 4 . The cost function Lspaceconv ( H , Bk ) is convex in h , and its minimum with respect to h can be obtained according to the following closed-form update rule , which concurrently estimates the scalar source variances and heteroscedastic noise variances : Hk+1 = diag ( hk+1 ) , hk+1i ← √ gki zki for i = 1 , . . . , N +M , where ( 10 ) g : = diag ( MkSN ) ( 11 ) z : = diag ( Φ > ( Σky ) −1Φ ) ( 12 ) Proof . A detailed proof can be found in Appendix E. The final estimate of the posterior distribution ( 2 ) can be obtained by starting from any initial guess for the model parameters , H0 = { Γ0 , Λ0 } and B0 , and iterating between update rules ( 7 ) , ( 10 ) , and ( 2 ) until convergence . We call the resulting algorithm full Dugh . Convergence is shown in the following theorem : Theorem 5 . Optimizing the non-convex ML cost function Eq . ( 4 ) with alternating update rules for B and H = { Γ , Λ } in Eq . ( 7 ) and Eq . ( 10 ) defines an MM algorithm , which is guaranteed to converge to a stationary point . Proof . A detailed proof can be found in Appendix G .

Summary: This paper proposes an efficient method for source reconstruction from M/EEG data that estimates the spatiotemporal correlation structure in both the source space and the measurement error. The spatiotemporal pattern of the source and noise covariances is given by a Kronecker product structure with an independent (diagonal) spatial structure and either a full or Toeplitz temporal structure, where the temporal component is shared across the source and noise components. Since the full optimization problem is non-convex, the authors propose a majorization-minimization (MM) algorithm with provable convergence and derive simple, efficient iterative updates for each component. Two versions of the model are proposed: full Dugh, which uses a full positive definite matrix for the temporal covariance; and thin Dugh, which can be applied in the stationary data setting and uses a Toeplitz temporal covariance structure with an alternative set of efficient updates for the MM algorithm. On simulated data with known sources, thin Dugh achieves comparable or better performance as compared to existing source reconstruction methods (Champagne, eLORETA, and S-FLEX). When applied to simulated data that was generated while varying the true underlying temporal covariance structure, full Dugh is shown to outperform thin Dugh when the true matrix is full and vice versa when it has Toeplitz structure. Finally, experiments on real MEG data from subjects presented with visual or auditory stimuli show that full and thin Dugh can recover plausible and well-localized temporal and spatial patterns for the reconstructed sources, even with a very limited number of trials.