Scalable Bayesian GPFA with automatic relevance determination and discrete noise models

1 Introduction . The adult human brain contains upwards of 100 billion neurons [ 3 ] . Yet many of our day-to-day behaviors such as navigation , motor control , and decision making can be described in much lower dimensional spaces . Accordingly , recent studies across a range of cognitive and motor tasks have shown that neural population activity can often be accurately summarised by the dynamics of a “ latent state ” evolving in a low-dimensional space [ 7 , 8 , 13 , 38 , 43 ] . Inferring and investigating these latent processes can therefore help us understand the underlying representations and computations implemented by the brain [ 20 ] . To this end , numerous latent variable models have been developed and used to analyze the activity of populations of simultaneously recorded neurons . These models range from simple linear projections such as PCA to sophisticated non-linear and temporally correlated models [ 10 , 16 , 22 , 43 , 50 ] . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . A popular latent variable model for neural data analysis is Gaussian process factor analysis ( GPFA ) , which has yielded insights into neural computations ranging from time tracking to movement preparation and execution [ 1 , 48 , 49 , 51 ] . However , fitting GPFA comes with a computational complexity of O ( T 3 ) and a memory footprint O ( T 2 ) for T time bins . This prohibits the application of GPFA to time series longer than a few hundred time bins without artificially chunking such data into “ pseudo-trials ” and treating these as independent samples . Additionally , canonical GPFA assumes a Gaussian noise model while recent work has suggested that non-Gaussian models often perform better on neural data [ 12 , 25 , 60 ] . Here , we address these challenges by formulating a scalable and fully Bayesian version of GPFA ( bGPFA ; Figure 1 ) with a computational complexity of O ( T log T ) and a memory cost of O ( T ) . To do this , we introduce an efficiently parameterized variational inference strategy that ensures scalability to long recordings while also supporting non-Gaussian noise models . Additionally , the Bayesian formulation provides a framework for principled model selection based on approximate marginal likelihoods [ 53 ] . This allows us to perform automatic relevance determination and thus fit a single model without prior assumptions about the underlying dimensionality , which is instead inferred from the data itself [ 5 , 41 ] . We validate our method on synthetic and biological data , where bGPFA exhibits superior performance to GPFA and Poisson GPFA with increased scalability and without requiring cross-validation to select the latent dimensionality . We then apply bGPFA to longitudinal , multi-area recordings from primary motor ( M1 ) and sensory ( S1 ) areas during a monkey self-paced reaching task spanning 30 minutes . bGPFA readily scales to such datasets , and the inferred latent trajectories improve decoding of kinematic variables compared to the raw data . This decoding improves further when taking into account the temporal offset between motor planning encoded by M1 and feedback encoded by S1 . We also show that the latent trajectories for M1 converge to consistent regions of state space for a given reach direction at the onset of each individual reach . Importantly , the distance in latent space to this preparatory state from the state at target onset is predictive of reaction times across reaches , similar to previous results in a task that includes an explicit ‘ motor preparation epoch ’ where the subject is not allowed to move [ 1 ] . This illustrates the functional relevance of such preparatory activity and suggests that motor preparation takes place even when the task lacks well-defined trial structure and externally imposed delay periods , consistent with findings by Lara et al . [ 32 ] and Zimnik and Churchland [ 62 ] . Finally , we analyze the task relevance of slow latent processes identified by bGPFA which evolve on timescales of seconds ; longer than the millisecond timescales that can be resolved by methods designed for trial-structured data . We find that some of these slow processes are also predictive of reaction time across reaches , and we hypothesize that they reflect task engagement which varies over the course of several reaches . 2 Method . In the following , we use the notationA to refer to the matrix with elements aij . We use ak to refer to the kth row or column ofA with an index running from 1 to K , represented as a column vector . 2.1 Generative model . Latent variable models for neural recordings typically model the neural activity Y ∈ RN×T of N neurons at times t ∈ RT as arising from shared fluctuations in D latent variables X ∈ RD×T . Specifically , the probability of a given recording can be written as p ( Y |t ) = ∫ p ( Y |F ) p ( F |X ) p ( X|t ) dF dX , ( 1 ) where F ∈ RN×T are intermediate , neuron-specific variables that can often be thought of as firing rates or a similar notion of noise-free activity . For example , GPFA [ 59 ] specifies p ( Y |F ) = ∏ n , t N ( ynt ; fnt , σ2n ) ( 2 ) p ( F |X ) = δ ( F −CX ) ( 3 ) p ( X|t ) = ∏ d N ( xd ; 0 , Kd ) withKd = kd ( t , t ) ) ( 4 ) That is , the prior over the dth latent function xd ( t ) is a Gaussian process [ 45 ] with covariance function kd ( · , · ) ( usually a radial basis function ) , and the observation model p ( Y |X ) is given by a parametric linear transformation with independent Gaussian noise . In this work , we additionally introduce a prior distribution over the mixing matrix C ∈ RN×D with hyperparameters specific to each latent dimension . This allows us to learn an appropriate latent dimensionality for a given dataset using automatic relevance determination ( ARD ) similar to previous work in Bayesian PCA ( Appendix I ; 5 ) rather than relying on cross-validation or ad-hoc thresholds of variance explained . Unlike in standard GPFA , the log marginal likelihood ( Equation 1 ) becomes intractable with this prior . We therefore develop a novel variational inference strategy [ 55 ] which also ( i ) provides a scalable implementation appropriate for long continuous neural recordings , and ( ii ) extends the model to general non-Gaussian likelihoods better suited for discrete spike counts . In this new framework , which we call Bayesian GPFA ( bGPFA ) , we use a Gaussian prior over C of the form cnd ∼ N ( 0 , s2d ) , where sd is a scale parameter associated with latent dimension d. Integrating C out in Equation 3 then yields the following observation model : p ( F |X ) = ∏ n N ( fn ; 0 , XTS2X ) , with S = diag ( s1 , . . . , sD ) . ( 5 ) Moreover , we use a general noise model p ( Y |F ) = ∏ n , t p ( ynt|fnt ) where p ( ynt|fnt ) is any distribution for which we can evaluate its density . 2.2 Variational inference and learning . To train the model and infer bothX and F from the data Y , we use a nested variational approach . It is intractable to compute log p ( Y |t ) ( Equation 1 ) analytically for bGPFA , and we therefore introduce a lower bound on log p ( Y |t ) at the outer level and another one on log p ( Y |X ) at the inner level . These lower bounds are constructed from approximations to the posterior distributions over latents ( X ) and noise-free activity ( F ) respectively . Distribution over latents At the outer level , we introduce a variational distribution q ( X ) over latents and construct an evidence lower bound ( ELBO ; 55 ) on the log marginal likelihood of Equation 1 : log p ( Y |t ) ≥ L : = Eq ( X ) [ log p ( Y |X ) ] − KL [ q ( X ) ||p ( X|t ) ] . ( 6 ) Conveniently , maximizing this lower bound is equivalent to minimizing KL [ q ( X ) ||p ( X|Y ) ] and thus also yields an approximation to the posterior over latents in the form of q ( X ) . We estimate the first term of the ELBO using Monte Carlo samples from q ( X ) and compute the KL term analytically . Here , we use a so-called whitened parameterization of q ( X ) [ 19 ] that is both expressive and scalable to large datasets : q ( X ) = D∏ d=1 N ( xd ; µd , Σd ) with µd = K 1 2 d νd and Σd = K 1 2 d ΛdΛ T dK 1 2 d T , ( 7 ) whereK 1 2 d is any square root of the prior covariance matrixKd , and νd ∈ RT is a vector of variational parameters to be optimized . Λd ∈ RT×T is a positive semi-definite variational matrix whose structure is chosen carefully so that its squared Frobenius norm , log determinant , and matrix-vector products can all be computed efficiently which facilitates the evaluation of Equations 8 and 9 . This whitened parameterization has several advantages . First , it does not place probability mass where the prior itself does not . In addition to stabilizing learning [ 39 ] , this also guarantees that the posterior is temporally smooth for a smooth prior . Second , the KL term in Equation 6 simplifies to KL [ q ( X ) ||p ( X|t ) ] = 1 2 ∑ d ( ‖Λd‖2F − 2 log |Λd|+ ||νd||2 − T ) . ( 8 ) Third , q ( X ) can be sampled efficiently via a differentiable transform ( i.e . the reparameterization trick ) provided that fast differentiableK 1 2 d v and Λdv products are available for any vector v : x ( m ) d = K 1 2 d ( νd + Λdηd ) with ηd ∼ N ( 0 , I ) , ( 9 ) where x ( m ) d ∼ q ( xd ) . This is important to form a Monte Carlo estimate of Eq ( X ) [ log p ( Y |X ) ] . To avoid the challenging computation ofK 1 2 d v for generalKd [ 2 ] , we directly parameterizeK 1 2 d , the positive definite square root ofK , which implicitly defines the prior covariance function kd ( · , · ) . In this work we use an RBF kernel forKd and give the expression forK 1 2 d in Appendix F. Additionally , we use Toeplitz acceleration methods to computeK 1 2 d v products in O ( T log T ) time and with O ( T ) memory cost [ 48 , 57 ] . We implement and compare different choices of Λd in Appendix F. For the experiments in this work , we use the parameterization Λd = ΨdCd , where Ψd is diagonal with positive entries and Cd is circulant , symmetric , and positive definite . This parameterization enables cheap computation of KL divergences and matrix-vector products while maintaining sufficient expressiveness ( Appendix F ) . All results are qualitatively similar when instead using a simple diagonal parameterization Λd = Ψd . Distribution over neural activity Evaluating log p ( Y |X ) = ∑ n log p ( yn|X ) for each sample drawn from q ( X ) is intractable for general noise models . Thus , we further lower-bound the ELBO of Equation 6 by introducing an approximation q ( fn|X ) to the posterior p ( fn|yn , X ) : log p ( yn|X ) ≥ Eq ( fn|X ) [ log p ( yn|fn ) ] − KL [ q ( fn|X ) ||p ( fn|X ) ] . ( 10 ) We repeat the whitened variational strategy described at the outer level by writing q ( fn|X ) = N ( fn ; µ̂n , Σ̂n ) with µ̂n = K̂ 1 2 ν̂n and Σ̂n = K̂ 1 2LnL T n ( K̂ 1 2 ) T , ( 11 ) where ν̂n ∈ RD is a neuron-specific vector of variational parameters to be optimized along with a lower-triangular matrix Ln ∈ RD×D ; and K̂ denotes the covariance matrix of p ( f |X ) , whose square root K̂ 1 2 = XTS follows from Equation 5 . The low-rank structure of K̂ enables cheap matrix-vector products and KL divergences : KL [ q ( fn|X ) ||p ( fn|X ) ] = 1 2 ( ‖Ln‖2F − 2 log |Ln|+ ||ν̂d||2 −D ) . ( 12 ) Note that the KL divergence does not depend onX in this whitened parameterization ( Appendix H ) . Moreover , q ( fn|X ) in Equation 11 has the form of the exact posterior when the noise model is Gaussian ( Appendix G ) , and it is equivalent to a stochastic variational inducing point approximation [ 18 ] for general noise models ( Appendix H ) . Finally , we need to compute the first term in Equation 10 : Eq ( fn|X ) [ log p ( yn|fn ) ] = ∑ t Eq ( fnt|X ) [ log p ( ynt|fnt ) ] . ( 13 ) Each term in this sum is simply a 1-dimensional Gaussian expectation which can be computed analytically in the case of Gaussian or Poisson noise ( with an exponential link function ) , and otherwise approximated efficiently using Gauss-Hermite quadrature ( Appendix K ; 18 ) .

Summary: The authors introduce a new variant of GPFA for neural data that includes an ARD prior over the mixing matrix to automatically learn dimensionality. The authors also develop specific inference methods for this work as well as exploit structure in the covariance matrix for stability/speed. The authors test their model on synthetic data to demonstrate that the bGPFA model correctly identifies true latent dimensionality and performs well in situations with Poisson and Negative Binomial noise.  The authors also show model performance on multi-region reaching data in monkey cortex and explore the presence of time delays in the model.