Probabilistic Tensor Decomposition of Neural Population Spiking Activity

1 Introduction . Large-scale neural population recording offers a unique window through which to study brain computations that involve the coordinated action of many neurons . Although such rich data sets are increasingly common , the information they contain can only be put to full use once we are able to separate the underlying factors that contribute to the neural activity in a simple and interpretable way . Rich data sets often fit naturally within a multidimensional array or tensor , and a variety of tensor decomposition techniques are available to identify factorial contributions that shape such data . But whereas tensor decompositions have long been used in fields such as chemometrics [ 1 ] or computer vision [ 2 ] , until recently their use in neuroscience had been limited to the study of continuous traces like EEG [ 3 ] , fMRI signals [ 4 ] or LFP [ 5 ] . Encouragingly , in the last few years , tensor methods have proven useful to segregate the influence of time dynamics , trial to trial variability [ 6 , 7 ] or 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . experimental condition [ 8 ] in neural spike-count data , even though the models employed have not yet provided a full probabilistic treatment appropriate to counts . Canonical Polyadic ( CP ) decomposition [ 9 ] is a widely used technique in which a tensor is decomposed into a sum of rank-1 elements . When applied to noisy data , it is usually cast as a least-squares problem and therefore implicitly assumes that noise is normally distributed . Spike counts violate this assumption but alternated-gradient-based optimization methods can be adapted to non-Gaussian likelihood functions . Hong et al . [ 10 ] introduced a general framework for Generalized CP ( GCP ) decomposition that can be applied to Poisson or negative binomial ( NB ) distributed datasets but it has not been used in neuroscience to our knowledge . Moreover , practical methods have largely focused on point estimates of the decomposition rather than a full probabilistic treatment . As such , they do not provide principled ways to incorporate prior knowledge about the data ( such as the ( co ) location of recorded neurons in the brain ) , to automatically determine the rank of the observed tensor ( although it is a generalization of the matrix rank , tensor rank is not as well behaved and understood [ 9 ] ) , nor to estimate posterior confidence in results . One approach to the Bayesian treatment of count models is based on augmentation using PólyaGamma ( PG ) variables [ 11 ] . This solution enables Gibbs sampling in models which combine Gaussian latent structure with logistic , Poisson or negative binomial observations , including factor models [ 12 ] and generalised-linear regression [ 11 ] , and it has been used successfully in neuroscience settings [ 13 ] . While PG augmentation has also been considered for tensor decomposition [ 14 ] , these authors recognised the computational challenges posed by Gibbs sampling in typical size problems , and so focused primarily on point estimation while incorporating a limited range of priors . Here , we introduce a variational approach to PG-augmented models with negative-binomial observations , which enables a relatively efficient Variational Bayesian ( VB ) treatment . Our approach is based on a principled approximation of the PG cross entropy , which is especially well suited for neural datasets in which conditional Fano factors ( FF ) ( once the contribution of population-wide shared influences is discounted ) are often close to one [ 15 , 16 ] . The fully probabilistic formulation is able to handle missing observations , extending to situations where data from multiple animals or recording sessions are to be combined ( sometimes referred to as “ stitching ” [ 17 ] ) . VB also makes it possible to combine Automatic Relevance Determination ( ARD ) with knowledge about expected group structures based on anatomical or histological information about the recorded neurons ( brain area , morphology , protein expression etc. ) . Finally , we also augment standard CP decomposition models using a constrained offset tensor , which allows us to study modulation of neural activity around baseline values , improving readability and interpretability of the CP-factors . The paper is organised as follows : In section 2 we review background material on tensor decomposition and PG augmentation schemes . In section 3 , we derive a VB algorithm for approximate inference of the tensor factors and model parameters . Last , in section 4 , we analyse synthetic data and neural recordings in mice performing passive multisensory integration [ 18 ] and show that our method can estimate the population-level effects of temporal dynamics and experimental condition in a fully probabilistic manner . We show improvements in variance explained , deviance and , last but not least , decomposition robustness when compared to standard CP and GCP baselines . 2 Background . 2.1 Tensor decomposition . Consider observed spike counts gathered in a D-dimensional tensor X ∈ RI1×···×ID , ( 1 ) with dimensions corresponding to neurons , time in trial , various ( factorial ) experimental manipulations , and possibly trial number . Our goal is to decompose X into a pair of simpler objects : a low rank tensorW and an `` offset '' tensor V which varies along a restricted set of dimensions . HereW is designed to capture low-dimensional structure in the time-course of population activity that varies systematically with experimental condition , while the offset V models potential changes in baseline firing rates with trial or experimental condition . More specifically , we seek a rank-R Canonical-Polyadic ( CP ) [ 9 ] decomposition ofW , often denoted W = [ |A ( 1 ) , . . . , A ( D ) | ] , such that for d̄ = ( d1 , . . . , dD ) ∈ ∏D d=1 { 1 . . . Id } : Wd̄ = R∑ r=1 A ( 1 ) d1r . . . A ( D ) dDr . ( 2 ) For n = 1 . . . D , A ( n ) is an In ×R matrix , termed a factor , whose rows are a ( n ) i = [ A ( n ) i1 , . . . , A ( n ) iR ] . For a tensor T , we represent by T ( n ) its n-th unfolding and recall that : W ( n ) = A ( n ) B ( n ) ᵀ , ( 3 ) where B ( n ) is the Khatri-Rao product ⊙ p 6=n A ( p ) . In addition , we assume that V varies only along a subset of D∗ ≤ D dimensions , remaining fixed along the complementary dimensions . For convenience , in what follows we will write d̄ = d̄∗ ∪ d̄• where d̄∗ = ( d∗1 , . . . , d ∗ D∗ ) contains the dimesions that vary and d̄• = ( d•1 , . . . , d • D−D∗ ) those that are fixed , regardless of the dimension ordering . Thus , instead of a I1 × · · · × ID tensor , in practice we need only estimate the I1 × · · · × ID∗ tensor ν and define V such that for all d̄∗ , d̄• : Vd̄∗∪d̄• = νd̄∗ . ( 4 ) We then take X to be generated by a negative-binomial noise process from the decomposition determined byW and V . 2.2 Negative-binomial distribution . The negative binomial ( NB ) model can be seen as a doubly stochastic relaxation of the Poisson distribution to overdispersed ( FF > 1 ) data . We model each spike count Xd̄ as Poisson-distributed with random mean µd̄ drawn from a Gamma distribution with shape ζ and scale e Wd̄+Vd̄ : ( Xd̄|µd̄ ) ∼ Poisson ( µd̄ ) , ( µd̄|Vd̄ , Wd̄ , ζ ) ∼ Gamma ( ζ , eWd̄+Vd̄ ) . ( 5 ) Marginalizing over the Poisson rates µd̄ then yields the NB distribution [ 13 ] : P ( Xd̄|Vd̄ , Wd̄ , ζ ) = Γ ( ζ + Xd̄ ) Xd̄ ! Γ ( ζ ) ( 1− pd̄ ) ζ ( pd̄ ) Xd̄ ⇒ ( Xd̄|Vd̄ , Wd̄ , ζ ) ∼ NB ( ζ , pd̄ ) , ( 6 ) where pd̄ = 1/ ( 1 + e − ( Wd̄+Vd̄ ) ) is the success probability . We recall that the mean and variance of the NB distribution are respectively given by E ( X|ζ , W , V ) = ζeW+V and V ( X|ζ , W , V ) = ζeW+V ( 1 + eW+V ) , giving a Fano factor FF = 1 + eWd̄+Vd̄ that is greater than one and approaches the Poisson case when ζ increases to infinity while keeping the mean constant . 2.3 Pólya-Gamma augmentation . Unfortunately , the negative binomial observation probability does not admit a conjugate prior onW and V which makes direct Bayesian inference intractable . This challenge is resolved by adopting a PG augmentation [ 11 ] . We first recall key properties of PG distributions before deriving the augmented model likelihood . 2.3.1 Definition and properties . For ξ ∈ R+ and ω ∈ R , a random variable U is PG ( ξ , ω ) distributed iff : U D = 1 2π2 ∞∑ k=1 gk ( k − 1/2 ) 2 + ω2/ ( 4π2 ) , ( 7 ) where gk ∼ Gamma ( ξ , 1 ) . We recall three major results from [ 11 ] . ( i ) For ξ > 0 , if U ∼ PG ( ξ , 0 ) , then for all ζ , v : evζ ( 1 + ev ) ξ = 2−ξe ( ζ−ξ/2 ) v ∫ ∞ 0 e− U 2 v 2 p ( U |ξ , 0 ) dU . ( 8 ) ( ii ) The densities of PG ( ξ , ω ) and PG ( ξ , 0 ) are linked through : p ( U |ξ , ω ) = coshξ ( ω/2 ) e−U2 ω 2 p ( U |ξ , 0 ) . ( 9 ) ( iii ) Although the PG density can be expressed as an alternating-sign sum of inverse gamma distributions , the terms of the sum diverge rapidly as ξ increases . However , the Laplace transform of the density ( the sign-reflected moment generating function ) has the closed form : E ( e−tU ) = coshξ ( ω/2 ) cosh−ξ ( √ ω2/2 + t 2 ) . ( 10 ) Moments of the distribution are obtained from derivatives of this transform at 0 , giving : E ( U ) = ξ 2ω tanh ( ω/2 ) and V ( U ) = ξ 4ω3cosh2 ( ω/2 ) ( sinh ( ω ) − ω ) , ( 11 ) with similar closed-form expressions for higher-order moments . These properties will , respectively , allow us to ( i ) augment our model with a latent tensor and obtain a Gaussian likelihood in W and V , ( ii ) derive the latent posterior distribution and ( iii ) introduce and validate a moment-matching approximation to PG cross entropies , thus facilitating variational Bayesian inference . 2.3.2 Model augmentation . Let x ∼ NB ( ζ , p ) , with p = 1/ ( 1 + e−v ) . Introducing a random variable u ∼ PG ( ζ + x , 0 ) we combine ( 6 ) and ( 8 ) to obtain P ( x|ζ , v ) = Γ ( ζ + x ) x ! Γ ( ζ ) ( e−v ) ζ ( 1 + e−v ) ζ+x = Γ ( ζ + x ) xΓ ( ζ ) 2− ( ζ+x ) e ( x−ζ 2 ) v ∫ ∞ 0 e− u 2 v 2 p ( u|ζ + x , 0 ) du . ( 12 ) This expression allows us to condition on u to obtain the augmented likelihood : P ( x|ζ , u , v ) = F ( x , ζ ) e−u2 ( v− x−ζ 2u ) 2 ( 13 ) ( with F gathering terms independent of u and v ) which is conjugate to a Gaussian prior on v. Applying this augmentation to the tensor model , we introduce a tensor of PG variates U of the same size as X , with each Ud̄ ∼ PG ( ζ + Xd̄ , 0 ) to obtain the log-likelihood logP ( X|U , V , W , ζ ) =+C − 1 2 ∑ d̄ Ud̄ ( Wd̄ + Vd̄ − Xd̄ − ζ 2Ud̄ ) 2 . ( 14 )

Summary: The paper presents a probabilistic tensor decomposition model for count data with applications to neural recordings. A variational inference method using Polya-Gamma augmentation is derived to allow feasible model fitting on large neural datasets, and some theoretical results of the asymptotic behavior of the variational approximation of the PG augmentation are shown. The utility of the tensor decomposition approach is demonstrated using recordings from rodents during a sensory integration experiment.