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: The paper proposes a Bayesian framework for source estimation in M/EEG. The Bayesian framework explicitly estimates spatial and temporal covariance matrices for both signal and noise, before estimating the posterior mean. A majorization-minimization algorithm is proposed to estimate the model, and detailed proofs are provided, including a convergence proof. The model is benchmarked against three source-estimation methods on simulated signals. Finally, the model is used in empirical recordings with a challenging low number of trials.