Probabilistic Tensor Decomposition of Neural Population Spiking Activity
Authors: Hugo Soulat
Sepiedeh Keshavarzi
Troy William Margrie
Maneesh Sahani
Venue: NeurIPS 2021 Spotlight (2021)
Keywords: Probabilistic
Tensor Decomposition
Neuroscience
Spike
Population Activity
Count

Abstract: The firing of neural populations is coordinated across cells, in time, and across experimental
conditions or repeated experimental trials; and so a full understanding of the computational
significance of neural responses must be based on a separation of these different contributions to
structured activity.

Tensor decomposition is an approach to untangling the influence of multiple factors in data that is
common in many fields.  However, despite some recent interest in neuroscience, wider applicability
of the approach is hampered by the lack of a full probabilistic treatment allowing principled
inference of a decomposition from non-Gaussian spike-count data.
Here, we extend the Pólya-Gamma (PG) augmentation, previously used in sampling-based Bayesian
inference, to implement scalable variational inference in non-conjugate spike-count models.

Using this new approach, we develop techniques related to automatic relevance determination to infer
the most appropriate tensor rank, as well as to incorporate priors based on known brain anatomy such
as the segregation of cell response properties by brain area.

We apply the model to neural recordings taken under conditions of visual-vestibular sensory
integration, revealing how the encoding of self- and visual-motion signals is modulated by the
sensory information available to the animal.

Probabilistic Tensor Decomposition of Neural Population Spiking Activity
Hugo Soulat
Sepiedeh Keshavarzi Sainsbury Wellcome Centre University College London London, WIT 4JG s keeshaaarziuucl ac uk
Gatsby Unit
University College London
London, WIT 4JG hugos&gatsby ucl ac
Troy W. Margrie Sainsbury Wellcome Centre University College London London, WIT 4JG t m.aaarieeucc... ac uk
Maneesh Sahani
Gatsby Unit
University College London
London, WIT 4JG maneesh@gatsby ucl. ac uk
Abstract
The firing of neural populations is coordinated across cells, in time, and across experimental conditions repeated experimental trials, and so full understanding of the computational significance of neural responses must be based on separation of these different contributions to structured activity. Tensor decomposition is an approach to untangling the influence of multiple factors data that common many fields. However, despite some recent interest in neuroscience, wider appli- cability of the approach is hampered by the lack of a full probabilistic treatment allowing principled inference of decomposition from non-Gaussian spike-count data. Here, we extend the Pólya-Gamma (PG) augmentation, previously used sampling-base based Bayesian inference, implement scalable variational inference in non-conjugate spike-coum count models. Using this new approach, we develop tech- niques related to automatic relevance determination to infer the most appropriate tensor rank, as well as to incorporate priors based on known brain anatomy such as the segregation of cell response properties by brain area. We apply the model neural recordings taken under conditions of visual-vestibular sensory integration, revealing how the encoding self- and visual-motion signals is modulated by the sensory information available to the animal.
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 increasingly common, the information they contain can only be put to full use once we are able to separate the underlying factors that contribute the neural activity in simple and interpretable way. Rich data sets often naturally within a multidimensional array or tensor, and variety of tensor decomposition techniques are available 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 BEG [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 trial variability [6, 7] or
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

A
B
shape
Va
signa,
V
V
Pa
XJ
A
A
Way
C
(3)
ARD
12
factors
W
k
latent
Ug
= (dr..., dp)
10
100
Figure 1: Probabilistic tensor decomposition model. (A) Graphical model. (B) The Offset tensor is constrained to vary along limited set of dimensions. (C) W is low rank tensor.
experimental condition [8] in neural spike-count data, even though the models employed have not yet provided full probabilistic treatment appropriate to counts.
Canonical Polyadic (CP) decomposition [9] is a widely used technique in which a tensor is decom- posed into a sum of rank-1 elements. When applied to noisy data, it is usually cast as a least-qquares problem and therefore implicitly assumes that noise is normally distributed. Spike counts violate this assumption but alternated- iiternated-gradient--e based optimization methods can be adapted to non- Gau-saas likelihood functions. Hong et al. [10] introduced 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 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 is 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ólya- Gamma (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-ling -inear 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- vations, which enables 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 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"
obser-
[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 decom- position and PG augmentation schemes. In section 3, we derive 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 D-dimensional tensor
X' E R //xx xxpx
(I)
with dimensions corresponding to neurons, time in trial, various (factorial) experimental manipula- tions, and possibly trial number. Our goal is to decompose X' into pair of simpler objects: a low rank tensor W and an "offset" tensor V which varies along a restricted set of dimensions. Here W 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 rank-R Canonical Poliadic (CP) [9] decomposition of W, often denoted W = [[A(0) "..,,......]... A(D!I), such that for d = (d1....,de E e .1).
(1) di 1
Wa
For A. is an In R matrix, termed factor, whose rows at (n) For tensor T, we represent by Tin) its n-th unfolding and recall that:
[A(1]
(m)
Ain
Wmn = A B(7T
(3)
where B(*) is the Khatri-Rao product O A(m pęn
In addition, assume that V varies only along subset of D. <. D dimensions, remaining fixed along the complementary dimensions. For convenience, in what follows we will write d= d d where d' (di...., d'p)) contains the dimesions that vary and d (df...., d' are fixed, regardless the dimension ordering. Thus, instead of a I x ... x ID tensor, in practice need only estimate the I, x... X ID. tensor and define V such that for all d' d=
(o-D) those that
Viuud Vij
(4)
We then take to generated by a negative-binomial noise process from the decomposition determined by W and V.
2.2 Negative-binomial distribution
The negative binomial (NB) model can be seen as doubly stochastic relaxation of the Poisson distribution to overdispersed (FF >1) data. We model each spike count X₂ as Poisson-distributed with random mean μa drawn from a Gamma distribution with shape and scale eWatVa;
(Kilua) Poisson((aa),
(5)
(ualVa, Wa, {) Gamma(ç, e Wa+
Marginalizing the Poisson rates μı then yields the NB distribution [13]
l'(C + X2) (1 (pa)
P(KalVa, Wa.
(XiVa, () ~ NB(G.pa).
X²)I(()
where pa = 1/(1 + (WWa+va))) is the success probability. We recall that the mean and variance respectively given by (W,
of the NB distribution W+ (1 + e eW+Y), giving a Fano factor FF 1 + e WatVa that is greater than one and approaches the Poisson case when ( increases infinity while keeping the mean constant.
W+V and V(X', W, V)
2.3 Pólya-Gamma augmentation
Unfortunately, the negative binomial observation probability does not admit a conjugate prior on W and V which makes direct Bayesian inference intractable. This challenge is resolved by adopting a PG augmentation [1 1]. We first recall key properties of PG distributions before deriving the augmented model likelihood.
3

---

2.3.1 Definition and properties
For & E R and w E R, random variable U is PG(G, w) distributed iff:
20
U _ 2T
9k
2 1 1/2)2 + w² (1 (4
where 9h Gamma(E, 1). We recall three major results from [11].
i For { > 0, ~ PG(S, 0), then for all 2
exc
2 ( 11 c= (UUE.0)dU.
(8)
(1 + en
(ii) The densities of PG((, w) and PG(5. 0) are linked through:
p(Uls, w) = cosh (w//)e 2 p(U]s.0)
(iii) Although the PG density can be expressed as an alternating-sip 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:
(V W 1/2/22+
E (e-")) = coshf(w/2 cosh &
(10)
2
Moments of the distribution are obtained from derivatives of this transform at 0, giving:
E(U) tanh(w/2) and V(U)= 3 (w/2) (sinh(w) - w)
(11)
2w
with similar closed-form expressions for higher-order moments.
These properties will, respectively. allow us (i) augment our model with a latent tensor and obtain Gaussian likelihood in W and V, (ii) derive the latent posterior distribution and (iii) introduce and validate moment- moment-matching approximation to cross entropies, thus facilitating variational Bayesian inference.
2.3.2 Model augmentation
Let x ~ NB p). with p = 1/(1 + t Introducing random variable combine (6) and (8) to obtain
PG(C + x, 0) we
P(x|ç, I(C + x) x! (1+c-))c+a
(e-
r(C r()+z), (c+=) (ct2) xT(())
( e -02 p(ul( + 0)du.
(12)
This expression allows us condition on to obtain the augmented likelihood:
P(rlç. u, v) = F(x, () ( (1-=-p))
(13)
(with F gathering terms independent of u and v) which conjugate to a Gaussian prior on v.
Applying this augmentation to the tensor model, we introduce a tensor of PG variates of the same size X, with each U3 PG(C + Xz 0) to obtain the log- ligliiell
X3 -5- )
log P(X'U, = Σ Ma
Wa + Va
(14)
214
3 Structured variational inference
PG-based augmentation was developed to facilitate Gibbs sampling in otherwise non-conjugate mod- cls. However, despite the availability of closed-form updates, sampling quickly becomes impractical

---

for large datasets [14]. Thus, in this section, we propose variational Bayes estimation procedure and report the associated updates (derived in Supplementary Materials A) to fit variational posteriors over the low rank tensor, the constrained offset tensor, and the tensor of PG latents. We incorporate prior knowledge about the neural recordings and perform Automatic Relevance Determination (ARD) using Gamma priors the factor precision matrices. Last, we treat the shape Ç as model parameter that we efficiently optimize during the variational M-step by approximating PG cross entropy terms with Gamma moment matching.
3.1 Priors
The prior on 24 is defined implicitly by the model augmentation through the conditional P(U X22) =
PG(C + X2. 0). We set the prior the offset matrix to be element-wise normal: P(va-) (( 1/Ta.), for d' (di....,,f,, Finally, for n = 1.... D, 1... In we use diagonal
precision matrices (()) to define the priors on the factor rows within W:
P(a 14 (())) = /O (A)
(15)
The parametrisation of (m) can encode various forms of structure expected in the data, for example:
Case 1: Rank ARD. We assume that precision matrices shared across set of modes and rows and use a gamma prior to perform automatic relevance determination. That is, we take:
A (m)
Diag(AI....,,A)))))
where P(),) Gamma(kux Fox).
(16)
During the optimization procedure, the shared diagonal elements of A! may diverge, effectively reducing the rank of the tensor.
Case 2: Neuron-group constraints. Alternatively, we can impose a mode-specific precision. For example we can group neurons based on their recording site and impose shared precision matrices across neurons in each group. If neuron factors are gathered in A we denote by g(i) E G the group of neuron and define:
A (m) = Diag(Agga ()g(.)............. ,g(0))))))))))))) where P()g(())) Gamma(kox,,0m 0ox).
(17)
In this case, divergence of a subset of precisions during optimization may lead components linked to one or more groups of neurons to shrink away. Thus, the model will favour explanations of the activity of each group that use as few components as possible.
3.2 Variational distribution
For the variational distribution on the latents Z [21, V, W,, we adopt the factorisation:
(()) = ((4))(((())))) II ((((a) II g(vaa) II ((a )g(A)) 1ix 0x),
(18)
and then iteratively maximize the VB free energy (or ELBO) F(S) (log P(X, 2(5)) + H(q]. If q-o corresponds to the variational distribution with variable marginalised out, we recall that the variational E-step update has the form q(x) x exp(log P(X, 21()) [19] (when obvious, we drop the explicit form of the distribution from the angle brackets).
3.3 Variational E-step
Here we describe the updates to infer the variational distributions of W.V and U. Details are provided in Supplementary A. We write (D'a) == (X, - <)/(2(ua)), retaining the expectation brackets the notation to emphasise the dependence on the complementary variational distributions.
Update of low rank tensor W. For n = 1... D, i 1= In:
q a (n)T = N (m." (1)) sz(11)
(19)
5

---

where the mean and variance are given by:
-!!! = -(( ((Mm).))) B(n)) + + (1
(20)
(n) = I(m) (Bin)T)Diag (M()),))) (ool))(mmmm))) (Dlm).i)) Manze))
(21)
Update of factor precisions A. By conjugacy, the precision update takes the form:
(Alki, 0) Gamma(k>, 0x)
(22)
where the parameters 1x and 0^ depend on the form of the prior (Rank ARD Neuron Neuron-group group constraint).
Update of offset tensor V. For d' (di..... dr we need to sum over constrained dimension update g(va=):
g(vā.) = =N (mā.: Eão),
(23)
with
Ex+ = (ux= + Tdx)) ) ma+ = Ede (UgeJ's+ + Taalaa))
(24)
and
Vá. l ((Yánuã)) - (Waluad))
úg. ((lanaa))
(25)
Eie (Hānudda
de
Update of PG latent U. For d (d1,..., do), using (12) and (9), we deduce that the PG varia- tional distribution takes the form
g(4a) = PG (5 + Xr. 22), where Ω2 a2= VWWa va)+
(26)
3.4 M-step
Finally, we fix the variational distribution and maximize the free energy with respect to the shape parameter. The main computational challenge in this step is to estimate the values of PG cross entropies. As we need one such estimate for each entry in X, both sampling methods and numerical approximations of PG densities prove too expensive. We therefore developed a moment-matching procedure (see Supplementary Materials B) obtain efficient approximation for the regime where the conditional Fano factor is near 1.
Denote by mΩg and m² the th moments respectively of PG(E,w) and of the Gamma distribution with matched mean and variance. As we show in Supplementary Materials B, the learnt model introduces coupling between the parameters & and w and the inferred conditional Fano Factor, with the result that both parameters diverge as FF -> 1 but the ratio & remains finite. Furthermore, provided that Ω O(W V/I (ie. the variational posterior variance of the estimated tensor is of the order of its squared mean), the leading order terms in the polynomial expansions of mya and m2 are both equal to
2 and so their higher-order moments are asymptotically equivalent in the limit FF + 1. This motivates an approximation of the PG KI. divergence by the divergence between corresponding moment-matched gamma distributions.
Empirically, we find that the approximation is most accurate for small values KL (see Supplemen- tary Fig. 7), but in practice posterior and prior PG distributions may be very different. Fortunately, we can exploit the analytic properties of the PG distribution (Eq. 9) to re-write the free energy in terms of KL divergences between PG distributions which differ only in their first parameter:
F(5) =C Σ
log I(X' + C()((( - (log2 + (Wa/2) + (Va/2) + log cosh (9q/2))
d
(27)
Σ KL (PG(X2 (.CCK)H+H.)))0PP((++++>+>

---

A
B
D
Ky (50, w)
LL000 1|fe. w)
log d(proopP)) pc)
Jog d(pre, pc)
60
40
10
c
100
True
120
100 FF
361
W
FF
Figure 2: Moment-Matchin Pólya-Gamma (PG) and Gamma (G) distribution in the unity Fano factor
(FF) limit. (A) Numerical estimates KL divergences between PG(50, 1) and PG(E,W). (B) Moment matched Gamma KL values. K (50, 1,),, indicates that estimate K (5000, 1ll5.w) using closed formed KL divergence between two Gamma distribution moment matched to PG 1) and (G) Color crosses point to the minimum of the KL divergences. (C) Normalized L2 norm between the densities PPG and pG as a function of an effective Fano factor FF + e
d(fra, fG) 2||pPG PG / + I|pe li )). The plain traces indicate relationship between & and FF for different number of observed spikes. (D) Values of d(fpa, fa) along such traces.
where we have omitted terms independent of Ç, and Ç is the current estimate of the shape parameter. The remaining KL divergence in (27) is comparatively small, and can be well approximated by the moment-matched gamma divergence.
We verified the validity of this approximation numerical experiments (Supplementary B.5). The gamma divergence replicated that between the PG distributions over wide range (Figure A-B) with accuracy increasing monotonically as conditional FF + (Figure 2 C-E).
is important to stress that the quality of the approximation depends on the conditional rather than marginal Fano factor; that is the FF derived from residual variance after shared trial-to-trial variability captured by common factors in the population has been discounted. While marginal Fano factors that include all sources of variance in neural data often found to exceed 1.5, population models tend to identify conditional Fano factors much closer to [20]. Thus the moment-matchen gamma approximation makes possible scale PG augmentation methods large tensor decompositions in a regime appropriate to neural data analysis.
3.5 Dealing with missing data
Finally, we note that neural data analysis often faces missing data (failing electrode, confounded experiments, etc.) or multi-subject experiments, where a given neuron is only recorded in some trials. Although alternating gradient type updates, where the likelihood maximized only over the observed entries, can accommodate such cases, they can be highly inefficient as they deal with full size tensor. Instead, our variational scheme reduces the effective size of the observed dataset as the variational updates previously described are adjusted simply by considering observed entries only.
4 Experiments
In this section, we first explore the accuracy of tensor factorisations obtained by our proposed algorithm on simulated datasets. We confirm the robustness of the method to missing data, and compare the accuracy of ARD-based estimates of tensor rank to alternatives. We then benchmark the method against standard CP [9] and GCP [10] decompositions on neural spike data. We report performance in terms of variance explained, deviance explained, and decomposition similarity metric. Finally, we look at the decomposition factors themselves, how they can be interpreted, and the consequences for neural data analysis.'
All experiment were performed using Intel(R) Xeon(R) CPU E5-2620 v3 @ 2.40GHz with 65GB of RAM. Benchmark analysis took approximately 12 hours.
7

---

A
VE
B
DE
C
Similarities
0.76
07.
06
0.72
5
10
Figure 3: Held out (trial based) performances function of the tensor rank R between probabilistic model (VB-GCP), CP and Poisson GCP. (A) Variance Explained. (B) Poisson Deviance explained. (C) Similarity metric. Error bars indicate 1 standard deviation.
4.1 Simulations
We simulated two 5 dimensional data tensors (size 100 70 x 4) using the generative NB model described in (6). In both cases, the generative shape parameter was fixed to 80 and the generative tensor rank was R Assuming that first dimension accounts for neuron loadings, we simulated "neuron groups" by restricting each CP-component to load only on subset of neurons.
The first simulation included offset tensor that was constrained vary across the first and third dimension, and incorporated a multi-experiment stitching setting in which only 1/4 of the full tensor was observed. Supplementary Figures and 9 show that the algorithm, initialized with R = 6 components, infers posteriors concentrated around the correct tensor decomposition, while identifying the correct shape parameter Ç. Automatic relevance determination eliminates two of the components thus effectively estimating the rank of the data tensor. For comparison, using the exact same initialization but without ARD, two spurious components are identified (See Supplementary Figure 9).
The second simulation designed compare the quality of ARD-based rank estimates other empirical rank selection metrics. To ensure standard CP and GCP methods [7, were applicable we did not include an offset tensor. Results in Supplementary Figure 10 illustrate the fragility of alternative approaches.
4.2 Spike recordings
We then applied our method to neural spiking data [18] recorded in mice during a multisensory integration paradigm that aimed elucidate the contribution of vestibular and visual signals to self-motion representation in the cortex. In brief, high-density single-unit recordings (Neuropixels and Neuronexus) were made from the retrosplenial cortex (RSP) of adult mice while they were presented with three types of motion stimuli: (i) passive horizontal rotation of the mouse in the dark (vestibular). (ii) passive horizontal rotation of the mouse in presence of a stationary surround vertical grating (vestibular + visual = both), and (iii) horizontal rotation of the surround vertical grating while the mouse was stationary, thus simulating optic flow alone (visual). (Figure 4-A). Single units were isolated in both granular and dysgranular divisions of the RSP (RSPg, RSPd) and across all cortical layers (Figure 4-C, left). The motion profile of the rotation is shown in Figure 4-B and was adapted such that the visual motion component under 'both' and "visual' conditions were identical. The dataset consists of 676 neurons recorded under all experimental conditions over trials. Each trial lasted 7 seconds (arranged in 0.1 second non overlapping time bins) during which the mouse was rotated back and forth from reference 0 degree position to an angle π. For each of the 24 cross validation folds, split the dataset in half across trials. The algorithms were trained on the summed spikes of one half, and performance evaluated on the other half as function of the tensor rank.
We benchmark our algorithm against both CP and Poisson GCP2 The latter is defined using an exponential link function for easier model comparison as well as more natural interpretation of the neural data. In plots, we refer to our method as VB-GCP (Variational Bayes GCP). All algorithms were trained with maximum of 10000 iterations. For VB-GCP, we used Kod = 100, lax
hhanss//ithubb.co//gggsuuutttt and https://www.tensortoollloos

---

A
C
Vestibular
Both 0
Visual 0
Neurons
Time Bins
Conditions
0
RS5/
Visd
Bolt
Vau
c2
40
0.2
π
.0.2
B
Velocity
Speed
Position
-0.2
[GSM
Rt5Pg
Visit
Bolth
Viex
622
20
1
19
50c
Sec
Sec
Layer
RSPU
Layer
RSPg
Layer
Sec
Bot
VIU
Figure 4: Neural data analysis. (A) Experimental conditions: the mouse is rotated in the dark (Vestibular), rotated in presence of a surround static visual scene (Both), or is stationary while the surrounding visual scene is being rotated (Visual). (B) Temporal profiles of the rotation stimuli under all conditions. (C) First three discovered factors using our decomposition. Grey patches account for standard deviation around the mean based on variational posteriors estimates.
(see (16)). Neuron groups (see (17)) were based the neuron recording sites (Layer and RSP division), and the offset tensor was only allowed to vary across the neuron and experimental condition dimensions.
We report testing variance explained (VE) and Poisson deviance explained (DE) as well as a similarity metric (all defined in Supplementary Section C). The latter is essentially normalised scalar product between different the CP factors (on a scale of 0 to 1) that accounts for possible permutations of the components. Similarity is measured for factors obtained with different initializations and from different cross validation folds. Thus it reflects both algorithmic stability and stability of the extracted structure across different trials under consistent experimental conditions.
find that VB-GCP outperforms CP and GCP, even though VE and DE explicitly assess per- formance under the Gaussian and Poisson likelihood losses assumed by those methods (Figure A-B). At low tensor rank, some of the improvement might come from the inclusion of the constrained offset tensor in the model. In particular, for fixed rank R, model includes more parameters than standard (G)CP. For a dataset of size N x (ncurons X times conditions/trialss) standard R CP decomposition has R x (N + T + K) parameters, to which the offset adds N x K. However, VB-GCP test performance continues to climb with increased R, while the maximum- likelihod approaches appear to overfit even though they applied to model with fewer parameters (both VE and DE on training data exceed that of the VB model, but test performance falls; Figure 11 Supplementary Materials). Thus, the robustness of VB estimation plays important role in facilitating the offset expansion of the model. While the expanded model appears valuable in the context of our experimental data, it may be omitted in other contexts without disrupting the remainder of the algorithm.
We also note that VB-GCP generates more reliable decompositions across cross-valle validation folds (Figure C). This reliability reinforces the interpretability of the specific decompositions found, making them of particular practical value in neuroscience context.
The Poisson distribution standard model of stochastic neuronal firing [21] and so we present Poisson-GCI relatively familiar baseline model. Our dataset displayed only moderate overdisper- sion (see Supplementary Materials), and so was close to the Poisson noise assumption. Furthermore, we compared model performance in terms of Poisson Deviance Explained (DE), which is the loss directly optimized by Poisson GCP. We also fit NB-GCP models but did not find notable performance improvements compared to the Poisson case (see Supplementary Section E).

---

Finally, we considered the interpretability of the factors found by the different models. In Supple- mentary Figures 12,13 and 14, for each method, we report the rank R decomposition which was the most similar to others across folds. Contrary to CP and GCP, our method is able to report posterior estimates of the factors, and, thanks to the offset tensor model, directly renders patterns of modulation around condition-dependent baseline values. The discovered temporal dynamics are smoother and centered, hence easier to compare between each other and to relate to experimental variables. In Figure 4-C we show the first three inferred components (sorted by amplitude). Strikingly, and even though they are extracted in a fully unsupervised manner, they are highly reminiscent of the experimental motion variables (namely the angular velocity, position, and absolute angular speed, 4-B). Particularly, one can notice that component r = 3, which correlates almost perfectly with mouse absolute rotation velocity, corroborates the decoding analysis performed in [18]. In the latter, combined vestibular and visual information led to higher decoding accuracy. Here, it led to bigger modulatory effects population activity as shown in the "conditions" panel. Both analyses thus suggest eestibulo-visual integration of rotation speed. In summary, our analysis revealed the encoding of motion variables while disentangling the contribution of different sensory modalities and brain region or layers.
Discussion
We have introduced variational Bayesian framework for probabilistic tensor decomposition of count data, amenable and scalable to neural datasets. Our algorithm provides a principled way to estimate the rank of the data and can incorporate prior knowledge about the neural recordings (such as the neuron shape or the recording site). In practice, and although this comes at a significant computational cost compared to finely optimized blackbox tensor toolboxes [9], our algorithm leads to improved reconstruction performance, and perhaps more importantly to greater robustness compared to baseline (G)CP decompositions. That is, our method was both able better capture variance in the dataset, and recovered decompositions that were more stable across experimental trials. Such stability is key in order draw accurate conclusions from the multi-way analysis of neural datasets. One additional feature of our method is the introduction of the constrained offset tensor to the decomposition. In the context of the multisensory integration experiment [18] presently analyzed, factors out the baseline firing rates of neurons and centers temporal factors, which facilitates the interpretation of the modulatory impact of experimental parameters on neural activity.
Our method builds upon previous work using PG augmentation schemes [1 [1]] and extends it to a Bayesian treatment of tensor decomposition. Although the PG approach was also adopted in [14], the authors used it to develop an EM algorithm for maximum a posterior point-estimates of the tensor factors, with no posterior variance. Moreover, they proposed a different way of inferring the tensor rank, relying on a multiplicative Gamma prior on the factor amplitudes. Although this scheme also allows probabilistic treatment of the rank, the focus on the amplitudes prevents the model from incorporating knowledge about specific modes of the dataset (such as neuron groups in the present paper). In contrast, we made variational treatment of the factors possible by approximating PG KL divergences with their Gamma moment- matched counterparts in regions where we expect them to be well behaved. Nevertheless, as with many mean-field or factored variational inference procedures, it is possible that our method underestimate the true uncertainty in the model variables. Possible extensions could therefore involve low-rank approximations the posterior covariance between factors, or the extension of linear-response methods [22] to PG-augmented models that lie outside the exponential family.
Finally, our model can cope with Poisson (FF = 1) and overdispersed datasets (FF > 1) but underdispersion (FF 1) has been reported in several neural circuits like the early visual system [23], the somatosensory cortex [24], the Dorsal Premotor Cortex (after presentation of a stimulus) [20] or the auditory cortex [25]. Interestingly, in this last example, WeWeese and Zador found that single-unit spiking activity was underdispersed to the point that the majority of neurons exhibited binary behavior with few multi-spike responses. In this particular case, the Pólya- Gamma augmentation based tensor decomposition could easily be adapted by switching from a negative binomial to binomial observation model using the original model from [11]. In the other cases, neuron responses seem more complex and would require more sophisticated extension of our model. One solution to low variability settings is to model spiking history and neuron coupling (see for example [26]), but its incorporation to the tensor factorisation framework would require further developments.
10

---

Acknowledgments and Disclosure of Funding
We would like thank Joaquin Rapela for early contributions this project and Mare Deisenroth and Céline Marié for helpful comments on the manuscript. This work was funded by the Simons Foundation and the Gatsby Charitable Foundation.
References
[1] C. J. Appellof and E. R. Davidson. Strategies for Analyzing Data from Video Fluorometric Monitoring of Liquid Chromatographic Effluents. Analytical Chemistry, 53(13):2053-22) 2056, 1981.
[2] M. Alex O. Vasilescu and Demetri Terzopoulos. Multilinear analysis of image ensembles: Tensorfaces. In Lecture Notes in Computer Science (including subseries Lecture Notes in Artificial Intelligence and Lecture Notes Bioinformatics), volume 2350, pages 447-460. Springer Verlag. 2002.
[3] Fumikazu Miwakeichi, Eduardo Martinez-Montes, Pedro A. Haldés-Sosa, Nobuaki Nishiyama, Hiroaki Mizuhara, and Yoko Yamaguchi. Decomposing EEG data into seeeee-iiee---euueeeee components using Parallel Factor Analysis. Neurolmage, 22(3):1035 jul 2004.
[4] C. F. Beckmann and M. Smith. Tensorial extensions of independent component analysis for multisubject FMRI analysis. Neurolmage, 25(1):294-311, mar 2005.
[5] Justen Geddes, Gaute Einevoll, Evrim Acar, and Alexander Stasik. Multi-Linear Population Analysis (MLPA) of LFP Data Using Tensor Decompositions. Frontiers in Applied Mathematics and Statistics, 6,41, sep 2020.
[6] Amoo Onken, Jian K. Liu, P.Chamanthi Karunasekara, Ioannis Delis, Tim Gollisch, and Stefano Panzeri. Using Matrix and Tensor Factorizations for the Single-Trial Analysis of Population Spike Trains. PLoS Computational Biology, 12(11), nov 2016.
[7] Alex H. Williams, Tony Hyun Kim, Forea Wang. Saurabh Vyas, Stephen I. Ryu, Krishna V. Shenoy, Mark Schnitzer, Tamara G. Kolda, and Surya Ganguli. Unsupervised Discovery of Demixed, Low-Dimensional Neural Dynamics across Multiple Timescales through Tensor Component Analysis. Neuron, 989)) jun 2018.
[8] Jeffrey S Seely, Matthew Kaufman, Stephen Ryu, Krishna V Shenoy, John Cunningham, and Mark M Churchland. Tensor Analysis Reveals Distinct Population Structure that Parallels the Different Computational Roles of Areas MI and V1. PLoS Comput Biol, 12(11):1005164 2016.
[9] Tamara G Kolda and Brett W Bader. Tensor Decompositions and Applications * Society for Industrial and Applied Mathematics, 51(3):455-500, 2009.
[10] David Hong. Tamara G. Kolda, and Jed A. Duersch. Generalized Canonical Polyadic Tensor Decomposition. SIAM Review, 62(1):133-163, 2020.
[11] Nicholas G. Polson, James G. Scott, and Jesse Windle. Bayesian inference for logistic mod- els using Polya-Gamma latent variables. Journal of the American Statistical Association, 133350)):1399-1344 may 2012.
[12] Arto Klami, Dinh Phung, and Hang Li. Polya-Gamma augmentations for factor models. PMI.R, 19:112-128, feb 2014.
[13] Jonathan W Pillow and James G Scott. Fully Bayesian inference for neural models with negative-binomial spiking. Advances in Neural Information Processing Systems, 25, 2012.
[14] Piyush Rai, Changwei Hu, Matthew Harding, and Lawrence Carin. Scalable probabilistic tensor factorization for binary and count data. IJCAI International Joint Conference on Artificial Intelligence, 201 3015-Jaaaa:3770037776 2015.
[15] Giedrius T. Buračas, Anthony M. Zador, Michael R. DeWeese, and Thomas D. Albright. Efficient discrimination of temporal patterns by motion-sensitive neurons in primate visual cortex. Neuron, 9055)995--96,, may 1998.
[16] Adam S Charles, Gregory D Horwitz, and Jonathan W Pillow. Dethroning the Fano Factor: Flexible, Model-Based Approach Partitioning Neural Variability, Neural Computation, 30:1012-1045, 2018.
11

---

[17] Srinivas C Turaga, Lars Buesing, Adam M Packer, Henry Dalgleish, Noah Pettit, Michael Häusser, and Jakob Macke. Inferring neural population dynamics from multiple partial recordings of the same neural circuit. Advances in Neural Information Processing Systems, 26, 2013.
[18) Sepiedeh Keshavarzi, Edward F Bracey, Richard A Faville, Dario Campagner, Adam I. Tyson, Stephen C Lenzi, Tiago Branco, and Troy W Margrie. The retrosplenial cortex combines internal and external cues encode head velocity during navigation. bioRriv, page 2021.01.22.427789, jan 2021.
[19] Michael I. Jordan, Zoubin Ghahramani, Tommi S. Jaakkola, and Lawrence K. Saul. Introduction variational methods for graphical models. Machine Learning, (112)):18-233, nov 1999.
[20] Mark M. Churchland, Byron M. Yu, John P. Cunningham, Leo P. Sugrue, Marlene R. Cohen, Greg S. Corrado, William T. Newsome, Andrew M. Clark, Paymon Hosseini, Benjamin B. Scott, David C. Bradley, Matthew A. Smith, Adam Kohn, J. Anthony Movshon, Katherine M. Armstrong. Tirin Moore, Steve W. Chang, Lawrence H. Snyder, Stephen G. Lisberger, Nicholas J. Priebe, Ian M. Finn, David Ferster, Stephen I. Ryu, Gopal Santhanam, Maneesh Sahani, and Krishna V. Shenoy. Stimulus onset quenches neural variability: A widespread cortical phenomenon. Nature Neuroscience, 13(3):369-33 -778, mar 2010.
[21] Peter Dayan, Laurence F Abbott, et al. Theoretical neuroscience: computational and math- ematical modeling of neural systems. Journal of Cognitive Neuroscience, 15(1):154-155 2003.
[22] Ryan Giordano, Tamara Broderick, and Michael Jordan. Linear Response Methods for Accurate Covariance Estimates from Mean Field Variational Bayes. Advances Neural Information Processing Systems, 2015-Janua: 1 444--149, jun 2015.
[23] Prakash Kara, Pamela Reinagel, and R Clay Reid. Low response variability in simultaneously recorded retinal, thalamic, and cortical neurons. Neuron, 27(3):635-646, 2000.
[24] Samuel Andrew Hires, Diego A Gutnisky, Jianing Yu, Daniel H O' O'Connor, and Karel Svoboda. Low-noise encoding of active touch by layer in the somatosensory cortex. Elife, 4:c06619, 2015.
[25] Michael R DeWeese, Anthony M Zador, et al. Binary coding in auditory cortex. Advances in neural information processing systems, pages 117-124, 2003.
[26] Wilson Truccolo, Uri T Eden, Matthew R Fellows, John P Donoghue, and Emery N Brown. A point process framework for relating neural spiking activity to spiking history, neural ensemble, and extrinsic covariate effects. Journal of neurophysiology. 93(2):1074-1088 2005.
[27] Marlene R. Cohen and Adam Kohn. Mcasuring and interpreting neuronal correlations. Nature Neuroscience, 14(7):81 1-819, jul 2011.
[28] Enes Makalic and Daniel F. Schmidt. High-Dimensiona. Bayesian Regularised Regression with the BayesReg Package. arXiv.rrg, nov 2016.
12


Reviews:
[{"review_id": "iN-Tq1fS48n", "reviewer": "Reviewer_kXjW", "summary": "This manuscript details the inference of a probabilistic dimensionality reduction technique that is applicable to non-Gaussian data, for example spike counts. Tensor decomposition methods have been applied to neuroscience data in the past, and this paper extends this formulation to frequently seen spike data. The authors introduce a variational Bayes approach towards a Polya-Gamma augmentation of the model with negative binomial observations. ", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The paper is very thoroughly written and presented. The model confers many advantages and the authors have presented convincing results. The model successfully recovers model parameters on simulated data while performing automatic relevance determination. On neural spiking data, the model class offers clear gains in variance explained and Poisson deviance explained in held-out data for R<4, when compared against canonical polyadic (CP) decomposition and generalized canonical polyadic (GCP) decomposition. \n\nIt would help to show the results using the CP and GCP comparison methods on simulation data to better motivate the proposed method.\n\nIt would be good to comment on the comparison between CP, GCP and the proposed method for R>=4 for the neural spiking data. Figure 3 contains the metrics for a subset of the values of R; it would be great to fill this out if possible, or provide an explanation of why these values of R were chosen.\n\nImportant: please include a detailed explanation of why GCP was chosen with the Poisson observation model as opposed to the negative binomial observation model.\n\nPlease detail out GCP before the first time it is mentioned in the manuscript, as well as include a brief description of the method.\n\nThe authors mention Polya-Gamma and its acronym PG multiple times – it would be helpful to define PG once the first time they mention it, and then stick to the acronym.\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "2"}, "extra_text": {}}, {"review_id": "DthFUQ4X7eM", "reviewer": "Reviewer_dzgw", "summary": "In the paper \"Probabilistic Tensor Decomposition of Neural Population Spiking Activity\" the authors propose a probabilistic model for canonical polyadic tensor decomposition and apply it to neural data. The model assumes negative binomial observations and is used with variational Bayesian inference. The method is evaluated on synthetic data and on multisensory integration activity recorded from mice.", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Originality:\nAs noted by the authors, Rai et al. IJCAI 2015 used the Polya-Gamma approach as well in the tensor decomposition setting but did not achieve posterior variance estimates. For all I know, the proposed method enabling to probabilistically determine the rank and estimate variance in the CP decomposition setting is novel. The method makes use of a negative binomial observation model which is appropriate for many neuroscience datasets.\n\nQuality:\nThe quality of the work is outstanding. The method is derived in a sensible and rigorous way. However, I did not check all equations.\nThe evaluation includes both synthetic data with known ground truth as well as data recorded from mice illustrating the utility of the approach on real data as well as yielding interesting results with regard to multisensory integration factors.\nI found it surprising that the proposed method outperformed the baseline (G)CP decompositions on variance explained and Poisson deviance explained even though these are what the loss functions of the baseline models optimize for. In the Discussion the authors speculate that this is due to the constrained offset tensor. Couldn't this be easily tested with a variant of the proposed method without the offset tensor?\n\nClarity:\nThe manuscript is very well written. The work is well motivated and presented with a useful illustration (Figure 1). All figures are informative and can be understood easily.\n\nSignificance:\nGiven datasets of growing sizes in neuroscience, tensor decompositions are promising tools to find easy to interpret patterns in these datasets. A probabilistic model that can be used to quantify uncertainty and to determine the rank automatically is very useful in this context.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The author discuss limitations with regard to uncertainty estimates in the Discussion. The limitations paragraph could have been more extensive as the authors did not use all of the available space. I would have welcomed a discussion on the negative binomial assumption given that sometimes neuronal data have a fano factor less than 1.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "5"}, "extra_text": {}}, {"review_id": "-CfKD_dgqU", "reviewer": "Reviewer_1JXG", "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.", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "This proposed model is a novel approach to tensor decomposition that has strong potential to be widely applicable in neuroscience (using spike train data) than standard tensor decomposition models, which use squared error instead of a likelihood on counts.\n-\tAdditionally, it allows for a full Bayesian treatment with good use of the Polya-Gamma augmentation in variational inference, rather than slower sampling. Advantages of this approach are demonstrated by, for example, performing rank selection using an ARD prior.\n\nThe presentation is generally clear, and appropriate comparisons to existing methods are given.\nHowever, a couple points could improve the clarity of the results and utility of the approach:\n-\tThe real data example did not include much discussion about what the tensor composition showed. Although comparisons to CP and GCP decompositions are provided in the supplementary, it’s not immediately clear to the reader what is gained in this case.\n-\tFor comparing the GCP and CP to the probabilistic model (e.g., fig 3), it would be helpful to include the number of parameters in each model. I think this would help mainly because of the offset tensor included in the probabilistic model.\n\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "2"}, "extra_text": {}}]

Author rebuttal:

Meta-review: The authors develop tensor decompositions for count data, like neural spike trains, as well as a variational inference approach based on the Pólya-gamma augmentation. The reviewers and I agree that the paper is novel, well-presented, and of broad interest to the computational and statistical neuroscience community at NeurIPS.
Decision: Accept (Spotlight)