Using Sum-Product Networks to estimate neural population structure in the brain
Abstract
We present a computationally efficient framework to model a wide range of population structures with high order correlations and a large number of neurons.
Our method is based on a special type of Bayesian network that has linear inference time and is founded upon the concept of contextual independence.
Moreover, we use an efficient architecture learning method for network selection to model large neural populations even with a small amount of data.
Our framework is both fast and accurate in approximating neural population structures.
Furthermore, our approach enables us to reliably quantify higher order neural correlations.
We test our method on publicly available large-scale neural recordings from the Allen
Brain Observatory.
Our approach significantly outperforms other models both in terms of statistical measures and alignment with experimental evidence.
Introduction
With the rise and fast growth of simultaneous neural population recording, modeling population structures and measuring correlations has become a focus of computational neuroscience Abbott & Dayan (1999); Averbeck et al.
; Azeredo da Silveira & Rieke (2021);
Urai et al.
Theoretical and Experimental works have demonstrated the necessity of measuring population correlations to investigate information coding Moreno-Bote et al.
(2014); Averbeck et al.
(2006), functional connectivity Dunn et al.
, learning Ganmor et al.
and arousal Vinck et al.
; Doiron et al.
Despite significant progress in recent years, research on measurement and analysis of population correlations still faces significant challenges Kohn et al.
Exact measurement of population correlations is an NP-hard problem in the general case since it requires computing every form of dependency among spiking neurons.
As a result, researchers have tried to come up with computationally efficient ways of approximation or indirect measurement of neural correlations.
Existing approaches are energy-based models rooted in statistical mechanics where the energy function incorporates couplings between subsets of variables (here neurons) Roudi et al.
(2009c)
; Tkaˇcik et al.
(2006); Sohl-Dickstein et al.
; Aurell & Ekeberg (2012).
However, these methods often carry auxiliary (and even unrealistic) assumptions about the neural dynamics, do not scale up for large populations, or are limited to very specific situations Roudi et al.
(2009b);
Averbeck et al.
Notably, generative models commonly used in other domains such as latent variable methods are often not applicable to neural populations as spiking neural data is discrete and sparse Zhao et al.
Furthermore, various parameters such as behavioural and emotional state of the animal affect firing patterns of neurons even in sensory cortex
Urai et al.
As a result, a recording long enough to train these models contains many external variable changes and confounding factors that make drawing scientific conclusions difficult.
In this paper, we take a probabilistic approach by modeling the joint probability distribution of neural activity with Bayesian networks.
Inference is NP-hard in general/classic Bayesian networks
Cooper (1990), making them impractical to model the population structure.
Therefore, we utilize a special family of Bayesian networks with linear inference time, first introduced as “Arithmetic circuits” Darwiche (2003);
Shen et al.
This family of networks has been designed to take advantage of “context-specific independence” of variables mainly for the purpose of computational efficacy, which also makes it suitable to extract local structures in the data Boutilier et al.
Shen et al.
We use a modification (and equivalent Rooshenas & Lowd (2014)) of arithmetic circuits, Sum-Product Networks (SPNs).
Sum product networks are more known and used by the community due to their success in image completion tasks Poon & Domingos (2011); Sanchez-Cauce et al.
In particular, we design a method that learns the architecture of the sum-product network from the spiking neural data in order to capture a wide range of population correlation/structure from local to global in polynomial time.
Due to the efficiency of our architecture learning and inference in SPNs, population structure estimation is polynomial in the size of the population.
In addition, we suggest a measure of high order population correlations based on our framework.
Our results include fitting on large-scale neural recording in different brain regions on more than 20 mice.
Our framework outperforms both of the best energy-based and latent variable models for neural population structure estimation.
Problem Definition and Related Work
One of the critical problems in computational neuroscience is providing an accurate statistical description of spike trains in a population of neurons.
As the full representation of the data, i.e. raw spike times, is high dimensional, spike trains are binned into small time windows.
The time bin should be short enough so each neuron spikes at most once in each bin (with some amount of tolerance in potentially losing some spikes).
In addition, this time bin should be large enough that the assumption of temporal independence of spikes holds.
With this time binning strategy, each neuron’s activity is a binary variable ( S i for neuron i is equal to 1 if there is a spike in the corresponding bin, otherwise 0 ) and each time bin represents an i.i.d sample/instance.
Therefore, spike trains of N neurons for the duration of T would be represented as a binary matrix D K × N where K = T/ ∆ t in which ∆ t is the bin length (Figure 1, left plot).
Consequently, the population activity has a probabilistic representation P ( S 1 , . . .
, S N ) , and the problem turns into modelling this joint distribution, given the data.
More specifically, the problem is to find a model m ∗ from a family of models M , and optimize its free parameters Θ m so as to satisfy the following:
[Equation]
In the existing approaches, M is set to maximum entropy (Ising) models, which are energy-based methods rooted in statistical mechanics Roudi et al.
(2009c)
; Schneidman et al.
Since learning maximum entropy models is computationally very expensive, these models are restricted to estimate the statistical properties of the population up to a constant order.
However, going beyond second order is not computationally feasible.
In fact, building the exact generative model is an intractable problem in the general case even for the second order (pairwise) correlations.
Therefore, even a Pairwise Maximum Entropy (PME) model requires further estimation where more accurate approximation algorithms requires thousands of samples for each pair, making them impossible to be used for large populations of neurons Roudi et al.
(2009c)
; Tkaˇcik et al.
Sohl-Dickstein et al.
; Aurell & Ekeberg (2012).
Moreover, there exist plausible scenarios, such as a dichotomized common input to loosely coupled neurons, in which pairwise correlations
To improve the structure estimation (and also as further evidence for inaccuracy of PME models), some methods modelled higher order correlations through the energy based approaches, e.g. “kpairwise” correlationsTkaˇcik et al.
, Restricted Boltzmann Machines (RBMs) on top of pair-wised correlations (semi RBMS (sRBMs))
Köster et al.
(2014), and sparse low-order correlations
Ganmor et al.
Founded upon energy based models, these methods also needs computationally expensive sampling strategies to estimate the partition function.
Moreover, to approximate the partition function, these methods often rely on pseudo-likelihoods which needs extensive amount of data points to be a accurate.
Here we use a special type of Bayesian network with linear inference, and consequently fast learning time, to estimate distributions of joint neural activity.
Our network approximates these distributions by ignoring correlations based on their effect size, rather than statistical order.
Moreover, the normalization process to obtain real probability values is also linear in the size of the network.
As a result, it is significantly more successful than energy-based models both in terms of efficiency and accuracy.
Model
Our framework is based on Sum-Product Networks (SPNs).
An SPN is a rooted directed acyclic graph representing a joint probability distribution of given variables.
This distribution is the result of a hierarchical combination of alternating mixtures (sum nodes) and factorizations (product nodes), with given variables as the leaf nodes of the network Poon & Domingos (2011); Sanchez-Cauce et al.
Specifically, each leaf node (a node with no children) represents a univariate probability distribution.
When the variable is categorical, the leaf node is a variable indicator ( I ( x ) ).
The scope of each leaf node is a singleton where its element is the variable that the leaf represents.
Other nodes are either sum or product .
A product node represents the product of its children (connected by a link).
A sum node represents the (normalized) weighted sum of its children.
The weight of each child is shown as the label of the link from the sum node to that child.
The scope of sum and product nodes is the union of scopes of their children.
Starting from a leaf as a one-node SPN, sum product networks can be built bottom-up by combining smaller SPNs through sum or product nodes.
The root of the SPN represents a valid joint probability distribution if the SPN is complete and decomposable .
An SPN is complete iff all children of each sum node have the same scope.
A sum-product network is decomposable iff all children of each product node have disjoint scopes.
Parameter Learning :
To model the joint probability distribution of the given variables, the weights on the links of sum nodes should be learned with a learning algorithm such as gradient descent or Expectation-Maximization (EM), given the data
Sanchez-Cauce et al.
Similar to other
Bayesian networks, the structure of the graph needs to be determined before parameter learning.
One approach is to use a random dense graph and rely on the parameter learning algorithm to select important components.
This strategy requires data sets orders of magnitudes larger than that gathered in neural experiments.
Structure Learning : Network structure can also be learned from the data directly based on the general computational properties of sum and product nodes Gens & Pedro (2013).
Sum nodes represent the sum of probability distributions over the same set of variables (see completeness above).
This means that they act as a union of sets.
Therefore, they explain the data best if their child distributions represent disjoint sets of data (instances), or in other words different “contexts” in the data.
A Product node computes the joint distribution of its children with a product, explicitly carrying an assumption of independence among the factors contributing to the product.
This means that recursive splittings of the data based on instances (e.g. by a clustering method) producing a sum node, and based on variables (e.g. via performing independence tests) resulting in a product node would give us a suitable
SPN structure for the given data.
While there exist few algorithms for structure learning Gens &
Pedro (2013); Vergari et al. (2015), we developed our own method due to our specific aim which is learning correlations in neural population.
× + + + ....
I S
I ¯ S 1 I S
I ¯ S
I S
I ¯ S ..  ..  ..
3.1 Neural population analysis with nSPN s
We construct a family of SPNs that we call Neural SPNs ( nSPN ) to model the joint probability distribution of spiking in populations of neurons.
These models are represented by directed graphs where leaves of an nSPN representing spiking activity of N neurons correspond to I s i and I ¯
s i ( 1 ≤ i ≤ N ) indicating whether neuron i fired or not in a given time bin (see section 2).
We suggest a systematic way to build the overall architecture of an nSPN suitable to capture population correlations even with a limited number of data points.
First we describe a shallow nSPN with three layers and then a mechanism to construct deep architectures that impose regularization for better fits.
These deep architectures are obtained by clustering based on the mutual information between different neurons.
Finally, we describe our measure of high order correlations and describe a mixture of nSPN models that aims to capture correlations with deep structures.
Shallow architecture of nSPN s: The simplest structure of an nSPN (in terms of graph complexity) consists of a product node as a root, linked to N sum nodes.
Each sum node represents the spike probability of a neuron i , linked to I s i and I
¯ s i leaf nodes.
This model, shown as nSPN 1 is basically a Naive Bayes model, assuming complete independence among neurons.
The right plot in figureshows an example nSPN 1 trained on the joint activity of three neurons in the left plot.
As mentioned before, each node of an SPN represents the probability of the leaves of the (sub-)tree of which it is the root.
As a result, (normalized) link weights of each sum node are aligned with the mean spike probability of the related neuron (note the assumption of independence).
Furthermore, given an assignment of probabilities to the leaf nodes and a configuration of spiking activity, the root represents the probability of that configuration’s occurrence according to the model (shown by the red number near the root node).
Within the SPN framework, a simple extension of the Naive Bayes model involves the introduction of multiple contexts , with population spiking being described by a Naive Bayes assumption conditioned on the context .
If there are two such contexts with associated probability, the model consists of a sum node as the root linked to two product nodes, each of which a root of an nSPN 1 (sub)-tree.
Similarly, a model can include an arbitrary number of contexts by increasing the number of product nodes under the root, e.g. nSPN b for b product nodes/contexts.
In principle, one can construct a shallow graph sufficient to model any neural data set by adding enough product nodes.
One problem with this approach is that there exist simple networks with deeper structure for which a three-layer model like we have described will require an exponential number of nodes Delalleau & Bengio (2011), which in addition to tractability also raises the issue of data limitations.
Moreover, all existing parameter learning algorithms for SPNs find local optima and shallow networks get stuck in local optima more frequently.
Most importantly, when the data is limited deeper architectures of our framework avoid over-fitting as their structure implements regularization.
× + + × + + ...
I S
I ¯ S
I S j
I ¯ S j × + + ...
I S
I ¯ S
I S j
I ¯ S j × + + ...
I S
j +
I ¯ S j +
I S N
I ¯ S N × + + ...
I S
j +
I ¯ S j +
I S N
I ¯ S N
Deep architecture of nSPN s: As computing all correlations is intractable, we seek an approach for constructing graphs that trades some fidelity for efficiency.
Our approach is to identify correlated ensembles of neurons and build deeper graphs that reflect these ensembles.
Since product nodes in
SPNs describe a factorization into independent distributions with separate scopes, a natural approach, given a scope of neurons { S 1 , . . .
, S N } , is to assign neurons to a number of groups based on a clustering algorithm, and form a model that has a product node as the root, with children that form the models for each group in the clustering.
Consider that such a clustering says to divide the neurons { S 1 , . .
S N } into { S 1 , . .
S j } and { S j +1 , . . .
, S N } .
One can then choose, for example, to model each group via the graphs nSPN 2 ( S 1 , . . .
, S j ) and nSPN 2 ( S j +1 , . . .
, S N ) .
This model would give the graph shown in Figure 2.
Roughly speaking, this model describes j neurons with two common contexts and N − j neurons with two different contexts that vary independently of those for the j neurons.
We can iterate this process to produce a graph with a finer scale.
Instead of using nSPN 2 to describe the distribution of { S 1 , . . .
, S j } we can iterate the clustering approach and separate these into { S 1 , . . .
, S k } and { S k +1 , . . .
, S j } .
These in turn can either be described by an nSPN b model or we can continue separating into clusters.
For example, nSPN bd is an SPN in which d levels of clustering by product nodes is performed (the tree has 2d + 3 layers), and each sum node connected to intermediate nodes has b branches.
Note that, one can easily use different number of branches for each node, or another criterion for splitting the variables without loss of generality (e.g minimum scope size for the clustering).
Correlation-based clustering of neurons
: A product node assumes its children are independent.
Therefore, to lose the least amount of population correlation, we would ideally minimize the mutual information between the two groups Gens & Pedro (2013).
Formally, splitting a variable set S into nonempty subsets of V and S \ V is optimal
iff V = arg
min
X ⊂ S I
( X, S \ X )
where I ( X, Y )
=
H ( X ) + H ( Y ) − H ( X, Y ) and H ( X ) = − (cid:80) ni
=1 P ( x i ) log P ( x i ) .
Since mutual information is a sub-modular function, the optimal split can be achieved in polynomial time by Queyranne’s algorithm if querying H ( X, Y ) is in polynomial time, e.g. through an oracle Queyranne (1998);
Narasimhan et al.
; Gens & Pedro (2013); Hidaka & Oizumi (2018).
This is not the case for us.
In fact, such query has the exact nature of our main goal, i.e. estimating the joint probability of the set.
Therefore, we use the marginal mutual information between pairs, I ( s i , s j )(1 ≤ i, j ≤ N ) instead.
With such a pairwise marginal entropy estimate, approaches that are based on maximum/minimum linkage criterion such as minimum/maximum spanning trees are the most suitable.
Particularly, we use agglomerative hierarchical clustering with the distance metric of log(2)
− I ( s i , s j ) between variable i and j and linkage criterion of average distance.
Measure of Neural Correlation and mixture-of-
nSPNs : The simplest architecture, nSPN 1 , assumes independence between all neurons.
On the other hand, the more complicated structures in our framework only assume independence within each context where the number of contexts are always more than 1.
This means that in stationary periods where the firing patterns of neurons do not change significantly, nSPN 1 fits to the data as well as more complex nSPN structures if and only of all neurons are independent of each other.
Any improvement in goodness of fit for more complicated structures (relative to nSPN 1 ) is due to the existence of neural correlations.
As a result, we can use the difference between the average log-likelihood (equation 1) of the best fitting nSPN and nSPN 1 as a measure for neural correlations.
For example, if our best model is nSPN b ∗ , our measure for population neural correlation is:
[Equation]
Any shallow structure SPN, i.e. nSPN b , is a super-set of nSPN 1 .
This is not true for deep architectures though.
As a result, especially for the purpose of measuring population correlations, a mixture of deep and shallow architectures should be used, e.g. a model constructed with a sum node as the root connected to a shallow structure such as nSPN 1 and a deep one such as nSPN bd .
We call this structure mixture-of- nSPN s.
Results
To assess our method on experimental neural recordings we modelled the neural data from the Neuropixels Visual Coding data set of the Allen Brain Observatory (https://observatory.brain-map.org).
In this data set, different stimuli (Gabors, flashes, drifting gratings, etc) were shown to 26 mice while the neural activity of their visual cortical and sub-cortical regions were recorded simultaneously with multiple Neuropixels probes.
Mice were head-fixed and on a wheel, free to run.
Importantly, stimuli in this experiment were passively viewed, meaning that it did not involve any choice or reward.
We modelled the neural population structure during the viewing of drifting gratings in 4 directions ( [0 ◦ , 45 ◦ , 90 ◦ , 135 ◦ ] ) with the contrast level of . 9 .
Each of these 4 stimuli was shown to 26 micetimes.
We analyzed the neural activity (spikes) of six regions in the visual cortex: VISp, VISl, VISal,
VISpm, VISrl, VISam.
Following other studies analyzing neural correlations, the time bin was set to ∆ t = 20 ms, balancing for minimizing spike loss due to binning, and accounting for activity delays between neurons.
We analyzed each of these 6 regions separately.
Each stimulus duration was 2 seconds.
We analyze the response from 400 ms after stimulus onset to 800 ms after stimulus onset.
The delay is to make sure the corresponding signal has reached the neurons in all areas.
Moreover, individual neurons’ change upon exposure to the same sensory stimuli (known as adaptation), hence the short period.
Since 400 ms results in only 20 data points, we combined the data of 5 consequent presentations with the same direction.
The order of stimulus condition (direction) of the stimulus type/block (drifting gratings) is random in the experiment.
We chose only 5 consequent presentations to keep the state of the animal (such as running speed or arousal level) as unchanged as possible.
The described data processing results in 15 blocks of population activity per each tuple (mouse, direction, area), each of which we modelled with a mixture-of- nSPN s (mixture of nSPN 1 and nSPN 21 )
We also fit our simplest structure nSPN 1 (Naive Bayes) in order to have a correlation-free baseline and compute the difference, ∆ ll SPN .
We fit both models on a modified version of data in which we shuffle the spikes of each neuron in the 400 ms of single stimulus showing for the primary visual cortical population (as the most sensitive area to events).
We observed a large difference between ∆ ll SPN in the original and shuffled data.
This difference confirmed our assumption about ∆ SPN in the original data is the sole product of neural correlations (Figure 3 left).
For an extra measure of caution we removed (mouse, direction) pairs with ∆ ll SPN of more than .
1 in the primary visual area for the shuffled data.
We used Minimum Flow Probability (MFP) Sohl-Dickstein et al.
(2011) that approximates pairwise maximum entropy variables through a direct fit to the data.
To the best of our knowledge this is the most accurate methods of pairwise approximation.
First, to further demonstrate the computational efficiency of our method, we compared the training time of SPN and MFP approaches for 16 example data blocks of different sessions (area V ISp ) on the same machine.
As demonstrated in the middle plot of figure 3, training time for the MFP is two orders of magnitude longer.
Importantly, for MFP we used an efficient package that takes advantage of a C++ implementation in Python Lee & Daniels
(2019), while the SPN was fully implemented in Python Molina et al.
We also used Restricted
Boltzmann Machine (RBM) as another comparison model Köster et al. (2014).
Instead of the whole joint probability distribution we used the (total) population activity as the measure of goodness of fit as used by others Roudi et al.
(2009a)
; Leen & Shea-Brown (2015).
In other words, we checked which model predicts the total number of spikes in a single bin better.
The (total) population activity is not an ideal measure for performing model comparison.
However, given the number of samples, i.e. 100, more desirable measures are prohibitive.
We also fitted an
RBM to the data.
The number of hidden units in the RBM, and the depth added to the mixtureof nSPN s ( nSPN 1 and nsPN 21 ) were both determined through 10-fold cross-validation within the 100 samples.
Compared to other models, the mixture-of- nSPN s produced significantly lower normalized JS-div for all of the 6 regions (Figure 3, right plot).
We also performed a Leave One Out Cross Validation (LOOCV) test, in which the samples of each of the 15 blocks is generated based on the otherfor each animal and direction (the validation set is outside of the fitted 100 samples and therefore unseen by the models).
For the SPN, the JS-div was also significantly lower than the MPF, in allregions.
However, the cross validated JS-div for the RBM was extremely high, showing over-fitting to the data.
Importantly, the non-cross validated result shows that even when over-fitted the RBM performed worse than the SPN.
Using nSPN models as a measure of neural correlations
Many experiments have suggested that arousal has the effect of reducing neural correlations, or decreasing the level of synchrony Gandal et al.
; Uhlhaas et al.
These experiments are mostly based on very noisy data such as the Local Field Potential (LFP) and rough estimates of signal frequency Gandal et al.
; Pfeffer et al.
; Vinck et al.
Having a quantitatively reliable measure based on spiking data, i.e. ∆ ll SPN , we examined the effect of arousal in the Allen
Brain Observatory.
Specifically, in each animal we looked at the (ratio to minimum) pupil diameter (a common measure of arousal in mice) changes at each of the chunks described above for all directions that have passed the shuffle test combined.
The mean distribution of Pearson correlation coefficient between pupil diameter and ∆ ll SPN across all animals was significantly below zero for all regions as shown in Fig.
4 (left).
To the best of our knowledge, even with LFP data, there is no work demonstrating the effect of arousal on higher visual areas during an experiment (See Vinck et al.
(2015) for VISp).
We repeated the same process for MPF by using the difference between log-likelihood of the full (first and second order) model with the first-order only version (second order variables set to zero).
The result was significantly different (Fig. 4, right).
Importantly and as mentioned, previous literature results (especially on VISp) are strongly in favour of the SPN results.
Moreover, this result further demonstrates the inability of the PME in modelling the full joint population activity (as opposed to the total population activity).
Discussion
Correlations are a potentially important channel of information in neural activity, with strong consequences for coding properties.
Using Sum Product Networks, we have constructed a computationally efficient approach to modeling structure and correlations in populations of spiking neurons.
Tested on experimental data, our approach outperformed both energy-based and latent variable methods.
Due to the nature of our scientific question in this paper, i.e. population correlations, we focused on short time spans in which the external variables and animal state remain constant.
Our framework, however, has the capacity to be applied more broadly in computational neuroscience, for example to capture population dynamics in different time scales or as a de-mixing tool for coded features in individual or population of neurons.
We believe, in these promising future directions, deep and mixed architectures play a larger role.
Acknowledgments
We wish to thank the Allen Institute founder, Paul G. Allen, for their vision, encouragement and support.