A flow-based latent state generative model of neural population responses to natural images

1 Introduction . Characterizing the activity of sensory neurons is a major goal of neural system identification . While neural responses in the visual cortex vary with visual stimuli , they also exhibit variability to the repeated presentations of identical stimuli [ 1–4 ] . This stimulus-conditioned variability has significant and sophisticated correlations among neurons commonly referred to as noise correlations [ 4–6 ] and exhibits dependency on various factors such as the stimulus [ 7–9 ] , the behavioral task [ 10 , 11 ] , attention [ 12–14 ] , and the general brain state [ 15 , 16 ] . Understanding the nature of this correlated variability and its functional implication in the processing of sensory stimuli requires models that account for both stimulus-driven and shared stimulus-conditioned variability . The goal is thus to model the stimulus-conditioned response distribution p ( r|x ) of population activity r ∈ Rn over n neurons responding to an arbitrary sensory stimulus x . However , models that account for stimulusdriven and stimulus-conditioned correlated variability have been developed largely independently . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . In the recent decade , we have seen significant progress in modeling stimulus-driven activity , largely driven by the use of deep neural networks ( DNNs ) [ 17–22 ] . Typically , the expected response of the neurons conditioned on the stimulus is captured as a function of the stimulus via a deep network fθ ( x ) = E [ r|x ] with learnable parameters θ . These models can therefore predict how population responses depend on an arbitrary stimulus , and could even be used to derive stimuli that would yield desirable responses [ 23 , 24 ] . Typically , these networks are trained using Poisson-loss , assuming that the population activity r is distributed around the stimulus-conditioned mean fθ ( x ) with an independent Poisson distribution . Therefore , existing state-of-the-art networks commonly ignore stimulus-conditioned correlations among neural responses , and impose strong assumptions about the form of the marginal distribution ( i.e . Poisson ) for each neuron . As sensory populations are known to exhibit noise correlations and deviate from Poisson distributions [ 4 , 25 , 26 ] , this conditional independence assumption might limit the ability of these models to accurately capture p ( r|x ) . On the other hand , many of the existing models for stimulus-conditioned variability capture the variations in the population activity by specifically modeling the responses to repeated presentations of an identical stimulus . Many of these approaches employ statistical techniques such as maximumentropy or copula distributions to reduce the number of parameters needed to fit the target distribution [ 27–29 ] . A popular approach has been to describe the stimulus-conditioned variability in terms of a typically lower-dimensional shared latent state z : p ( r|x ) = ∫ p ( r|x , z ) p ( z|x ) dz [ 16 , 25 , 26 , 30–35 ] . Among these are hierarchical generative models that can capture more sophisticated relationships between the stimulus and noise correlations , as well as deviations from Poisson , such as over-dispersion [ 25 , 26 , 32 , 34 , 35 ] . While these approaches present powerful methods to capture stimulus-conditioned variability , they often fit p ( r|x ) separately for each unique stimulus and require responses to repeated presentations of the stimulus [ 16 , 25 , 26 , 29 , 35 ] . This limits their ability to yield predictions to a novel stimulus without requiring some stimulus-specific parameters to be learned . Furthermore , the increased complexity of the distribution usually requires a substantially more involved probabilistic machinery to make latent state inference and parameter fitting feasible . Consequently , most latent state models for neural data either ignore stimulus-driven variability altogether [ 30 , 31 , 34 ] , or employ a very simple model of stimulus-driven variations [ 16 , 25 , 26 , 32 ] . Here , we propose a new model that closes the gap between these two approaches by combining DNNbased models of stimulus-driven activity with a latent state model that accounts for shared stimulusconditioned variability . While DNNs can be trained effectively via gradient-based optimization , the challenge is to avoid the complex probabilistic machinery associated with existing latent state models , particularly those that require stimulus-specific parameters to be learned over repeated presentations of identical stimuli . To this end , we combine normalizing flows [ 36–41 ] with Gaussian Factor Analysis ( FA ) models [ 42 ] , where the stimulus-dependence occurs through a DNN that learns to shift the mean of the FA distribution based on the stimulus . FA models make use of multivariate Gaussian distributions with a particular low-rank structure of the covariance matrix . While the use of FA in capturing shared variability greatly simplifies inference and learning , it is not directly applicable to neural responses because neural responses are not Gaussian-distributed , particularly for low firing rates . To circumvent this problem , variance-stabilizing transformations , such as the square-root function , have been used in the past to make the responses more Gaussian-distributed [ 16 , 30 ] . However , there may be other transformations that capture the response distribution more accurately . Furthermore , since the transformation for one neuron may not be applicable to other neurons , ideally it would be learned for each neuron separately . To achieve this flexibility , we allow our model to learn neuron-specific transformations with a marginal normalizing flow . Normalizing flow models are density estimators that use a series of diffeomorphisms to transform the source density underlying the data into a simple distribution—typically an isotropic Gaussian of the same dimension . These transformations are usually chosen to have efficient-to-compute log-determinants , and typically act on the entire variable vector to capture any statistical dependencies between the dimensions . Here , we replace the isotropic Gaussian with an FA model to capture dependencies among dimensions and only use diffeomorphisms that act on each dimension separately , i.e . apply flow-based transformations on the marginals only . While this choice places certain restrictions on the complex dependencies between neurons that may be captured ( refer to section 4 Discussion for details ) , it has two important advantages : ( 1 ) The generative model is easy to train while combining state-of-the-art deep networks with flexible latent state models , and ( 2 ) the use of marginal flows allows for an easy mechanism to compute conditional distributions of one neuron given responses of other neurons that would not be easy to obtain with non-marginal flow models . image-com putable m odel transformed responses `` `` flow transform me an shared covariance latent m odel 1 2 k 1 2 3 n-1 n readout position network receptive field positions `` `` cortex positions `` `` neural responses `` `` la te nt s ta te s `` '' im ag e `` '' b a Figure 1 : Flow-based Factor Analysis model . a : Schematic of the flowbased model relating all relevant variables in the study . b : Schematic of the sub-network used by the image-computable model to map cortical positions into receptive field positions . Refer to section 2 Methods for the details . In summary , we make the following contributions . We ( 1 ) combine state-of-the-art DNN-based models with flow-based latent state models to jointly account for stimulus-driven and shared stimulusconditioned variability in neural population activity . Our model can predict the distribution of neural responses to unseen stimuli , without the need for repeated presentations to learn stimulus-conditioned variability . We ( 2 ) apply our method on the activity of thousands of neurons in response to natural images , recorded via two-photon Calcium imaging from multiple areas of the mouse visual cortex . We demonstrate that our model outperforms current state-of-the-art methods in capturing the distribution of responses . Finally , we ( 3 ) show that our model infers latent state structures with meaningful relations to behavioral variables such as pupil dilation as well as other functional and anatomical properties of visual sensory neurons . 2 Methods . 2.1 Models . Flow-based Factor Analysis model ( FlowFA ) For a given stimulus x and population response r ∈ Rn , where n is the number of neurons , we define our normalizing flow-based Factor Analysis ( FlowFA ) model of the stimulus-conditioned population activity p ( r|x ) as p ( r|x , θ , φ ) = N ( Tφ ( r ) ; fθ ( x ) , CC > + Ψ ) · |det∇rTφ ( r ) | . ( 1 ) FlowFA has two major parts : ( 1 ) A flow model Tφ with learnable parameters φ that transforms the population responses r such that the transformed responses v = Tφ ( r ) are well modelled by a ( 2 ) Gaussian Factor Analysis ( FA ) modelN ( v ; fθ ( x ) , CC > + Ψ ) ( Fig . 1a ) . Here , N ( v ; µ , Σ ) denotes a Gaussian distribution over v with mean µ and covariance Σ . According to the FA model , the random variable v is generated via v = fθ ( x ) + Cz + ε where z ∈ Rk is a low-dimensional latent state with k n and an isotropic Gaussian prior z ∼ N ( 0 , Ik ) whose samples map to v via the factor loading matrix C ∈ Rn×k . The effect of the stimulus x on the responses is captured by the mean of the FA distribution that depends on the stimulus , modeled as a deep network fθ ( x ) ∈ Rn with learnable parameters θ ( Fig . 1a , b ) . We further include neuron-specific , independent noise ε ∼ N ( 0 , Ψ ) where Ψ ∈ Rn×n is a diagonal covariance matrix . Since the flow model is a trainable change of variables , it introduces the absolute determinant |det∇rTφ ( r ) | of the Jacobian ∇ of Tφ with respect to r into Eq . ( 1 ) . The transform itself is a diffeomorphism , i.e . an invertible differentiable mapping Tφ : Rn 7→ Rn allowing us to evaluate the exact likelihood of each data point and easily draw samples from the model . Therefore , the model serves as a fully generative model from which samples of the stimulus-conditioned population responses can easily be generated for an arbitrary stimulus . In the model formulation presented here , we choose Tφ to act on each single dimension separately , i.e . Tφ ( r ) = [ Tφ1 ( r1 ) , ... , Tφn ( rn ) ] > . This choice results in a diagonal Jacobian which not only substantially simplifies the form of the determinant to det∇rTφ ( r ) = ∏n i=1 ∂Tφi ∂ri , but also allows us to easily compute conditionals and marginals ( see appendix A for the details ) . This would not generally be possible for diffeomorphisms with a non-diagonal Jacobian . Zero-Inflated Flow-based Factor Analysis model ( ZIFFA ) For two-photon Calcium imaging , a significant portion of inferred neural activity is zero , resulting in a sharp peak at zero in the response distribution ( i.e . zero-inflated distribution ) [ 43 ] . This zero-inflation is potentially a problem for the FlowFA model since the model would attempt to generate the peak at zero by mapping a large proportion of the Gaussian probability mass onto the “ zero ” responses , resulting in a poor fit to the response distribution . To avoid this , we extend FlowFA by modeling the zero responses with a separate peak ( similar to Wei et al . [ 43 ] ) and applying the FlowFA model to capture only the positive responses . We refer to this model as Zero-Inflated Flow-based Factor Analysis ( ZIFFA ) . More specifically , ZIFFA is a mixture model that models neural responses below and above a threshold value ρ with two separate , non-overlapping distributions . To capture the peak at zero , the responses below the threshold ( i.e . “ zero ” responses ) are modeled by a uniform distribution , while FlowFA is used to capture responses above the threshold : p ( r|x ) = ∏ { i : ri≤ρ } 1− qi ( x ) ρ · ∏ { i : ri > ρ } qi ( x ) ·N ( Tφ ( r+ ) ; fθ , + ( x ) , C+C > ++Ψ+ ) ·|∇Tφ ( r+ ) | , ( 2 ) where qi ( x ) is the probability of the response being above the threshold ρ modeled , jointly with the mean of the FA , as a function of the stimulus via a DNN fθ with learnable parameters θ. r+ and fθ , + ( x ) are the sub-vectors , and C+ and Ψ+ are the sub-matrices corresponding to responses above the threshold , and θ , C , Ψ are the same as defined in Eq . ( 1 ) . Refer to appendix B for the derivation . Control models We compare the FA-based models against two control models used for neural system identification that assume independence among neurons with specific forms of marginal distributions inspired by existing work : ( 1 ) Poisson [ 18 , 22 ] and ( 2 ) Zero-inflated Gamma ( ZIG ) [ 43 ] . To capture continuous neural responses measured with Calcium imaging , we relax the discrete Poisson distribution into a continuous distribution by assuming r = r̂ + where r̂ ∼ Poisson ( λ ) and ∼ Uniform [ 0 , 1 ) . This yields the likelihood function ppoiss ( r|x ) = n∏ i λi ( x ) brice−λi ( x ) bric ! , ( 3 ) where λ ( x ) = fθ ( x ) is the predicted firing rate of the neurons to input image x modeled as a DNN fθ with learnable parameters θ . The ZIG distribution is a mixture of a uniform and a gamma distribution separated at the value ρ with no overlap [ 43 ] : pZIG ( r|x ) = n∏ i ( 1− qi ( x ) ρ + qi ( x ) r κi−1 i Γ ( κi ) νi ( x ) κi exp ( − ri νi ( x ) ) ) , ( 4 ) where νi ( x ) is the scale parameter of the gamma distribution , and qi ( x ) is same as in Eq . ( 2 ) . To formulate ZIG as an image-computable model , νi ( x ) and qi ( x ) are jointly modeled using a DNN fθ with learnable parameters θ . Similar to Wei et al . [ 43 ] , we let the shape parameter κi be neuronspecific , but independent of the input . Importantly , we used the same value for ρ in both ZIG and ZIFFA models . Note that when the covariance matrix of the FA-based models is diagonal ( i.e . 0-dimensional latent state ) , these models assume independence among neurons and their performance is directly comparable to the control models .

Summary: This paper proposes a model of neural responses by capturing both stimulus dependent mean as well as variance. They transform the firing rate of each neuron using a reversible flow based model and model the transformed responses as a multi-variate gaussian. The mean of the multivariate gaussian is given by a neural network transformation of the input stimulus, and the low rank + diagonal covariance is independent of stimulus. They show the utility of this approach on synthetic as well as calcium imaging recordings from mouse visual cortex.