Fitting summary statistics of neural data with a differentiable spiking network simulator

1 Introduction . Modelling neural recordings has been a fundamental tool to advance our understanding of the brain . It is now possible to fit recurrent spiking neural networks ( RSNNs ) to recorded spiking activity [ 1 , 2 , 3 , 4 , 5 ] . The resulting network models are used to study neural properties [ 6 , 7 , 8 , 9 ] or to reconstruct the anatomical circuitry of biological neural networks [ 4 , 5 , 10 ] . Traditionally a biological RSNN is modelled using a specific Generalized Linear Model [ 11 ] ( GLM ) often referred to as the Spike Response Model ( SRM ) [ 12 , 11 , 13 ] . The parameters of this RSNN model are fitted to data with the maximum likelihood estimator ( MLE ) . The MLE is consistent , meaning that if the amount of recorded data becomes infinite it converges to the true network parameters when they exist . However , when fitting neural activity in practice the MLE solutions are often reported to generate unrealistic activity [ 14 , 15 , 16 ] showing that this method is not perfect despite it ’ s popularity in neuroscience . We also find that these unrealistic solution emerge more easily when hidden neurons outside of the recorded units have a substantial impact on the recorded neurons . This is particularly problematic because the likelihood is not tractable with hidden neurons and it raises the need for new methods to tackle the problem . To address this , we optimize sample-and-measure loss functions in addition to the likelihood : these loss functions require sampling spiking data from the model and measuring the dissimilarity between recorded and simulated data . To measure this dissimilarity we suggest to compare summary statistics ∗Equal contributions . Senior authors . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . popular in neuroscience like the peristimulus time histogram ( PSTH ) and the noise-correlation ( NC ) . Without hidden neurons , this method constrains the network to generate realistic neural activity but without biasing the MLE solution in the theoretical limit of infinite data . In practice , it leads to network models generating more realistic activity than the MLE . With hidden neurons , the sample-and-measure loss functions can be approximated efficiently whereas the likelihood function is intractable . Although recovering the exact network connectivity from the recorded spikes remains difficult [ 10 ] , we show on artificial data that modelling hidden neurons in this way is crucial to recover the connectivity parameters . In comparison , methods like MLE which ignore the hidden activity wrongly estimate the connectivity matrix . In practice the method is simple to optimize with automatic differentiation but there were theoretical and technical barriers which have prevented earlier attempts . The first necessary component is to design an efficient implementation of back-propagation in stochastic RSNN inspired by straightthrough gradient estimators [ 17 , 18 ] and numerical tricks from deterministic RSNNs [ 19 ] . Previous generative models of spiking activity relying on back-propagation used artificial neural networks [ 20 , 21 ] which are not interpretable model in terms of connectivity and neural dynamics . Previous attempts to include hidden neurons in RSNN models did not rely on back-prop but relied on expectation maximization [ 11 , 22 ] or reinforce-style gradients [ 23 , 24 , 25 , 26 ] which have a higher variance [ 27 ] . There exist other methods to fit neural data using back-propagation and deep learning frameworks but they do not back-propagate through the RSNN simulator itself , rather they require to engineer and train a separate deep network to estimate a posterior distribution [ 28 , 29 , 30 , 31 ] or as a GAN discriminator [ 32 , 33 , 21 ] . The absence of a discriminator in the sample-and-measure loss function connects it with other simple generative techniques used outside of the context of neural data [ 34 , 35 , 36 ] . Our implementation of the algorithm is published online openly 2 . 2 A recurrent spiking neural network ( RSNN ) model . We will compare different fitting techniques using datasets of spiking neural activity . We denote a tensor of KD recorded spike trains as zD ∈ { 0 , 1 } KD×T×nV where nV is the total number of visible neurons recorded simultaneously and T is the number of time steps . To model the biological network which produced that activity , we consider a simple model that can capture the recurrent interactions between neurons and the intrinsic dynamics of each neuron . This recurrent network contains nV+H neurons connected arbitrarily and split into a visible and a hidden population of sizes nV and nH . Similarly to [ 1 , 4 , 37 , 14 ] we use a GLM where each unit is modelled with a SRM neuron [ 12 ] with ut , j being the distance to the threshold of neuron j and its spike zt , j is sampled at time step t from a Bernoulli distribution B with mean σ ( ut , j ) where σ is the sigmoid function . The dynamics of the stochastic recurrent spiking neural network ( RSNN ) are described by : zt , j ∼ B ( σ ( ut , j ) ) with ut , j = vt , j − vthr vthr ( 1 ) vt , j = nV+H∑ i=1 dmax∑ d=1 W dj , izt−d , i + bj + Ct , j , ( 2 ) whereW defines the spike-history and coupling filters spanning dmax time-bins , b defines the biases , vthr = 0.4 is a constant , and C is a spatio-temporal stimulus filter processing a few movie frames and implemented here as a convolutional neural network ( CNN ) ( this improves the fit accuracy as seen in [ 14 , 38 ] and in Figure S1 from the appendix ) . Equations ( 1 ) and ( 2 ) define the probability Pφ ( z ) of simulating the spike trains z with this model and φ represents the concatenation of all the network parameters ( W , b and the CNN parameters ) . Traditionally the parameters which best explain the data are given by the MLE : argmax φPφ ( zD ) [ 1 , 7 , 37 , 5 , 4 ] . When all neurons are visible , the likelihood factorizes as ∏ t Pφ ( zDt |zD1 : t−1 ) , therefore the log-likelihood can be written as the negative cross-entropy ( CE ) between zDt and σ ( u D t ) = Pφ ( zDt = 1|zD1 : t−1 ) where uDt is computed as in equation ( 2 ) with z = zD . So when all neurons are recorded and zD are provided in the dataset the computation of the MLE never needs to simulate spikes from the model and it is sufficient to minimize the loss function : LMLE = − logPφ ( zD ) = CE ( zD , σ ( uD ) ) . ( 3 ) 2Code repository : https : //github.com/EPFL-LCN/pub-bellec-wang-2021-sample-and-measure 3 The sample-and-measure loss functions . In this section we describe the sample-and-measure loss functions which include simulated data in a differentiable fashion in the optimization objective , a direct benefit is to enable the consideration of hidden neurons . We define the sample-and-measure loss functions as those which require sampling spike trains z ∈ { 0 , 1 } K×T×nV+H from the model Pφ and measuring the dissimilarity between the recorded and simulated data . This dissimilarity is defined using some statistics T ( z ) and the generic form of the sample-and-measure loss functions is : LSM = d ( T ( zD ) , EPφ [ T ( z ) ] ) , ( 4 ) where d is a dissimilarity function , like the mean-squared error or the cross entropy . To compute the expectations EPφ we use Monte-Carlo estimates from a batch of simulated trials z . For example to match the PSTH between the simulated and recorded data , we consider the statistics T ( z ) t , i = 1 K ∑ k z k t , i and evaluate the expectation with the unbiased estimate σ̄t , i = 1 K ∑ k σ ( u k i , t ) . Denoting the PSTH of the data as z̄Dt , i = 1 K ∑ k z k , D t , i and choosing d to be the cross-entropy , we define the sample-and-measure loss function for the PSTH : LPSTH = CE ( z̄D , σ̄ ) . ( 5 ) When all neurons are visible , we minimize the loss functionLMLE+SM = µMLELMLE+µSMLSM where µMLE , µSM > 0 are constant scalars . When there are hidden neurons , the log-likelihood is intractable . Instead we minimize the negative of a lower bound of the log-likelihood ( see appendix D for a derivation inspired by [ 40 , 23 , 24 , 25 ] ) : LELBO = CE ( zD , σ ( uV ) ) , ( 6 ) with σ ( uV ) being the firing probability of the visible neurons , where the visible spikes zV are clamped to the recorded data zD and the hidden spikes zH are sampled according to the model dynamics . Hence the implementation of LELBO and LSM are very similar with the difference that the samples used in LSM are not clamped ( but all our results about LSM are also valid when they are clamped if we use the extended definition given in Appendix D ) . To compute the gradients with respect to these loss functions we use back-propagation which requires the propagation of gradients through the stochastic samples z . If they were continuous random variables , one could use the reparametrization trick [ 27 ] , but extending this to discrete distributions is harder [ 17 , 18 , 41 , 42 , 43 ] . One way to approximate these gradients is to relax the discrete dynamics into continuous ones [ 41 ] or to use relaxed control variates [ 42 ] , but we expect that the relaxed approximations become more distant from the true spiking dynamics as the network architecture gets very deep or if the network is recurrent . Instead , we choose to simulate the exact spiking activity in the forward pass and use straight-through gradient estimates [ 17 , 18 ] in the backward pass by defining a pseudo-derivative ∂zt , i ∂ut , i over the binary random variables zt , i ∼ B ( σ ( ut , i ) ) . We use here the same pseudo-derivative ∂zt , i ∂ut , i = γmax ( 0 , 1 − |ut , i| ) as in deterministic RNNs [ 19 ] because the dampening factor ( here γ = 0.3 ) can avoid the explosive accumulation of approximation errors through the recurrent dynamics [ 44 ] . Although the resulting gradients are biased , they work well in practice . A geometrical description of a sample-and-measure loss function In the remaining paragraphs of this section we provide a geometrical representation and a mathematical analysis of the loss function LSM . For this purpose , we consider that the recorded spike trains zD are sampled from an unknown distribution PD and we formalize that our goal is to bring the distribution Pφ as close as possible to PD . In this view , we re-write LSM = d ( EPD [ T ( z ) ] , EPφ [ T ( z ) ] ) and we re-interpret LMLE as the Kullback-Leibler divergence ( DKL ) from Pφ to PD . This is equivalent because the divergence DKL ( PD , Pφ ) is equal to −EPD [ logPφ ( zD ) ] up to a constant . In Figure 1 we represent the losses LMLE and LPSTH in the space of distributions and we can represent LMLE = DKL ( PD , Pφ ) as the distance between PD and Pφ . To represent the sampleand-measure loss function LSM ( or specifically LPSTH in Figure 1 ) , we project the two distributions onto the space of summary statistics T represented in light green . Hence , these projections represent the expected statistics EPφ [ T ( z ) ] and EPD [ T ( zD ) ] and LSM can be represented as the distance between the two projected statistics . Although minimizing LMLE should recover the true distributions ( i.e . the biological network ) if the recorded data is sufficient and the model is well specified , these ideal conditions do not seem to apply in practice because the MLE solution often generates unrealistic activity ( see Figure 1D-E ) . Panels B and C in Figure 1 use the geometrical representation of panel A to summarize the two main scenarios where minimizing LMLE+SM is better than LMLE alone . In panel B , we describe a first scenario in which the model is misspecified meaning that it is not possible to find φ∗ so that PD = Pφ∗ . In this misspecified setting , there is no guarantee that the MLE solution yields truthful activity statistics and it can explain why the MLE solution generates unrealistic activity ( Figure 1E ) . In this case , adding a sample-and-measure loss function can penalize unrealistic solutions to solve this problem ( Figure 1B and F ) . Another possible scenario is sketched in panel C. It describes the case where the model is well specified but LMLE is flat around PD for instance because too few trials are recorded or some neurons are not recorded at all . In that case we suggest to minimize LMLE+SM to nudge the solution towards another optimum where LMLE is similarly low but the statistics T match precisely . In this sense , LSM can act similarly as a Bayesian log-prior to prefer solutions producing truthful activity statistics . Theoretical analysis of the sample-and-measure loss function To describe formal properties of the sample-and-measure loss function LSM , we say that two distributions are indistinguishable according to the statistics T if the expectation E [ T ( z ) ] is the same for both distributions . We assume that the dissimilarity function d ( T , T ′ ) reaches a minimum if and only if T = T ′ ( this is true for the mean-squared error and the cross-entropy ) . Then for any statistics T and associated dissimilarity function d we have : Property 1 . If the RSNN model is expressive enough so that there exists parameters φ◦ for which Pφ and PD are indistinguishable according to the statistics T , then φ◦ is a global minimum of LSM . Reciprocally , if this minimum is reached then Pφ and PD are indistinguishable according to T . This property is a direct consequence of our assumption on the function d. If T measures the PSTH it means that the optimized simulator produces the same PSTH as measured in the data . This can be true even if the model is misspecified which is why we represented in Figure 1B that the RSNN minimizing LMLE+PSTH lands in the dark green region where the PSTH of the data is matched accurately . As one may also want to target other statistics like the noise correlation ( NC ) , it is tempting to consider different statistics T1 and T2 with corresponding dissimilarity functions d1 and d2 and to minimize the sum of the two losses LSM1+SM2 = µ1LSM1 + µ2LSM2 where µ1 , µ2 > 0 are constant scalars . Indeed if d1 and d2 follow the same assumption as previously , we have ( see Figure 2 for an illustration ) : Property 2 . If the RSNN model is expressive enough so that there exists φ◦ for which PD and Pφ are indistinguishable according to both statistics T1 and T2 , then φ◦ is a global minimum for LSM1+SM2 . Reciprocally , if this minimum is reached , PD and Pφ are indistinguishable according to T1 and T2 . This is again a direct consequence of the assumptions on d1 and d2 . Additionally Figure 1C conveys the idea that LSM and LMLE are complementary and LSM can be interpreted as a log-prior . This interpretation is justified by the following Property which is similar to an essential Property of Bayesian log-priors . It shows that when it is guaranteed to recover the true model by minimizing LMLE , minimizing the regularized likelihood LMLE+SM will also recover the true model . Property 3 . If the RSNN is well specified and identifiable so that PD = Pφ∗ and in the limit of infinite data , then the global minimum of LMLE+SM exists , it is unique and equal to φ∗ . To prove this , we first note that all the conditions are met for the consistency of MLE so φ∗ is the unique global minimum of LMLE . Also the assumption PD = Pφ∗ is stronger than the assumption required in Properties 1 and 2 ( previously φ◦ only needed to match summary statistics : with parameters φ∗ the model is indeed matching any statistics ) so it is also a global minimum of LSM . As a consequence it provides a global minimum for the summed loss LMLE+SM . This solution is also unique because it has to minimize LMLE which has a unique global minimum . It may seem that those properties are weaker than the classical properties of GLMs : in particular LSM+MLE is not convex anymore because of LSM and the optimization process is not guaranteed to find the global minimum . This could be a disadvantage for LMLE+SM but we never seemed to encounter this issue in practice . In fact , as we argue later when analyzing our simulation results , the optimum of LMLE+SM found empirically always seem to be closer to the biological network than the global minimum of LMLE . We think that it happens because the conditions for the consistency of the MLE and Property 3 ( identifiability and infinite data ) are not fully met in practice . On the contrary , the Properties 1 and 2 hold with a misspecified model or a limited amount of recorded data which may explain the success of the sample-and-measure loss functions in practice . See Figure 1 for a geometrical interpretation .

Summary: This paper introduces a novel approach to fitting spiking network models to neural data. The proposed method uses differentiable recurrent spiking networks, which are fit to the data by minimising a loss function with the help of recent techniques such as the straight-through gradient estimate. The key novelty is to generalise the loss function to be minimized, to include not only the likelihood of the data, but also a number of regularizing terms that depend on desired spike train statistics, such as the PSTH or the noise correlations. These domain-specific terms in the loss function help with breaking degeneracies in the model, and with improving fit quality in the face of model misspecification. After a theoretical analysis of the properties of the proposed loss function(s), the paper shows examples of applications of the technique to some simulated data. A comparison to a state-of-the-art generative model for population spike trains, and an example of usage of the proposed method on a large database of experimental recordings, are presented in the supplementary material.