SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Amin Nejatbakhsh
Isabel Garon
Alex H Williams
Center for Neural Science, New York University, New York, NY
Center for Computational Neuroscience, Flatiron Institute, New York, NY
{anejatbakhsh,igaron,awilliams}@flatironinstitute.org
The signaling capacity of a neural population depends on the scale and orientation of its covariance across trials. Estimating this "noise" covariance is challenging and is thought to require a large number of stereotyped trials. New approaches are therefore needed to interrogate the structure of neural noise across rich, naturalistic behaviors and sensory experiences, with few trials per condition. Here, we exploit the fact that conditions are smoothly parameterized in many experiments and leverage Wishart process models to pool statistical power from trials in neighboring conditions. We demonstrate that these models perform favorably on experimental data from the mouse visual cortex and monkey motor cortex relative to standard covariance estimators. Moreover, they produce smooth estimates of covariance as a function of stimulus parameters, enabling estimates of noise correlations in entirely unseen conditions as well as continuous estimates of Fisher information—a commonly used measure of signal fidelity. Together, our results suggest that Wishart processes are broadly applicable tools for quantification and uncertainty estimation of noise correlations in trial-limited regimes, paving the way toward understanding the role of noise in complex neural computations and behavior.
Nominally identical repetitions of a sensory stimulus or a behavioral action often elicit variable responses in neurons. Intuitively, this "noise" degrades the information capacity of neural representations. However, the precise nature of this effect is complex since neural responses are not independent random events. Indeed, pairs of co-recorded neurons exhibit significant "noise correlations" from trial-to-trial, with well-studied and intriguing consequences [1-5]. In particular, introducing noise correlations into a circuit can theoretically degrade or enhance its signaling capacity depending on the geometric configuration of neural representations in firing rate space (Fig. 1A; [6, 7]).
Estimating noise covariance is therefore an important problem that arises frequently in neural data analysis. However, this problem is widely regarded as challenging [8-10] because the number of estimated parameters (i.e., number of unique neuron pairs) grows quadratically with the number of co-recorded neurons, while the number of measurements per trial (i.e., number of neurons with a response) grows only linearly (for more discussion see, e.g., [11]). As a result, experimental investigations into noise correlation have been limited to extremely simple tasks where a very large number of trials can be collected over a small number of conditions. For example, Rumyantsev et al. [10] recorded neural responses over just two stimulus conditions (oriented visual gratings), repeated across hundreds of trials. We therefore have a rudimentary characterization of how noise correlations
<sup>1</sup>Throughout, we focus on the problem of estimating noise covariance instead of noise correlation. By accurately estimating the former, we can also estimate the latter since the correlation between two neurons is given by their covariance normalized by each neuron's marginal variance.
Figure 1: (A) Illustration of information limiting noise correlations. (B) Experimental datasets with smoothly parametrized conditions (see [7]). Left, a nonhuman primate makes point-to-point reaches to radial targets. Right, a rodent views drifting visual grating stimuli of different orientations and temporal frequencies (i.e. speed). (C) Parameterized condition spaces for the datasets in B. (D) Samples from Gaussian and Wishart process prior distributions for with  $C = 6$  conditions over 1D periodic variable (e.g. angle) and  $N = 2$  neurons at varying kernel smoothness parameters ( $\lambda$  in eq. 6). Dots and ellipses represent the mean and covariance of neural responses. Colors appear in the increasing order of condition value (e.g. angle) and dashed lines connect neighboring conditions. Increasing  $\lambda_{\mu}$  (horizontal axis) encourages smoothness in the means while increasing  $\lambda_{\Sigma}$  (vertical axis) encourages the ellipses to change smoothly in scale and orientation. (E-F) Graphical model of the Wishart model with Normal (E) and Poisson (F) observations (see Supplement B.5 for details of the Poisson model).
change across rich or naturalistic experimental conditions, despite interest among the theoretical community [5, 12, 13], and experimental evidence that noise statistics are context-dependent [14-16]. New statistical approaches that are capable of inferring noise covariance with fewer trials per condition are needed to make progress on this subject.
To address this need, we develop a probabilistic model based on Wishart processes [17]. The core insight we exploit is that similar experimental conditions—e.g. visual gratings at similar orientations or cued arm reaches to similar locations—ought to exhibit similar noise statistics. Thus, by building this insight into the prior distribution of a probabilistic model, we can pool statistical power across neighboring conditions and accurately infer noise correlation structure with fewer trials per condition. In fact, this principle also enables us to predict covariance in entirely unseen conditions by interpolating between (and potentially extrapolating from) observed conditions. We demonstrate the scientific potential and generality of this approach through two experimental applications spanning different modalities (vision and motor tasks) and species (mouse and nonhuman primate).
Consider an experiment with  $N$  neurons recorded over  $C$  distinct stimulus conditions, each with  $K$  repeated trials. Experimental advances continually enable us to record larger neural populations (large  $N$ ) and neuroscientists are increasingly interested in performing rich and naturalistic experiments (large  $C$ ). Both of these trends put us in a regime where  $K < N$ , meaning there are few trials per condition relative to the number of neurons [9].
To simplify exposition, we assume that there are an equal number of trials in each condition. However, the model easily applies to the more general case with unbalanced trial allocations.
Empirical Estimator: Our goal is to estimate the  $N \times N$  covariance matrix,  $\Sigma_{c}$ , describing trial-by-trial variability in each condition  $c \in \{1, \dots, C\}$ . If  $\mathbf{y}_{ck} \in \mathbb{R}^{N}$  is the measured response on trial  $k$  of condition  $c$ , then a straightforward estimator is the empirical covariance:
$$
\overline {{\boldsymbol {\Sigma}}} _ {c} = (1 / K) \sum_ {k} \left(\boldsymbol {y} _ {c k} - \overline {{\boldsymbol {y}}} _ {c}\right) \left(\boldsymbol {y} _ {c k} - \overline {{\boldsymbol {y}}} _ {c}\right) ^ {\top}, \quad \text {w h e r e} \quad \overline {{\boldsymbol {y}}} _ {c} = (1 / K) \sum_ {k} \boldsymbol {y} _ {c k}. \tag {1}
$$
This estimator behaves poorly in our regime of interest when  $K < N$ . Indeed, while the true covariance must be full rank, the empirical covariance is singular with  $N - K$  zero eigenvalues.
Regularized Estimators: Statisticians often refer to problems where  $K < N$  as the high-dimensional regime. Estimating covariance in this regime is actively studied and understood to be challenging [18]. The Ledoit-Wolf estimator [19] is a popular approach, which regularizes ("shrinks") the empirical covariance towards a diagonal estimate. Other estimators shrink the individual covariance estimates towards the grand covariance (see eq. 2 below) instead of a diagonal target [20]. Graphical LASSO [21] is another popular approach, which incorporates sparsity-promoting penalty on the inverse covariance matrix. All of these methods can be understood as forms of regularization, which reduce the variance of the covariance estimate at the expense of introducing bias.
Grand Empirical Estimator: The methods summarized above (or variants thereof) are common within the neuroscience literature [8, 22, 23]. It is also common to simply estimate the covariance as being constant across conditions [10, 24], in which case an estimate can be found by pooling trials across all conditions to achieve an empirical estimate of the grand covariance:
$$
\overline {{\boldsymbol {\Sigma}}} = \frac {1}{C} \sum_ {c} \overline {{\boldsymbol {\Sigma}}} _ {c} = \frac {1}{K C} \sum_ {k, c} \left(\boldsymbol {y} _ {c k} - \overline {{\boldsymbol {y}}} _ {c}\right) \left(\boldsymbol {y} _ {c k} - \overline {{\boldsymbol {y}}} _ {c}\right) ^ {\top}. \tag {2}
$$
This estimator can be thought of as having high bias (since covariance is often thought to be stimulus-dependent [14-16]), but low variance since many more trials can be averaged over after pooling across conditions. We compare our model to these baselines in section 3.
2.2 Wishart Process Model of Noise Correlations
Our insight is to incorporate an alternative form of regularization whenever the experimental conditions are smoothly parameterized. For instance, a classical stimulus set used by visual neuroscientists is a grating pattern varied over  $C$  orientations. In fact, such stimuli are commonly used for noise correlation estimation [8, 10, 25, 26]. Intuitively, two gratings with nearly identical orientations should evoke similar probabilistic responses. Thus, as stimulus orientation is incremented, the mean and covariance of the response should smoothly change along a manifold (Fig. 1A).
This idea readily extends to other experimental settings. For example, pure acoustic tones with smoothly varied frequencies are often used in auditory neuroscience [27]. In studies of olfactory coding, isolated odorants or mixtures of odorants can be presented at smoothly varied concentrations [28]. In motor neuroscience, voluntary reach movements can be cued to continuously varied target locations [29]. Finally, in the context of navigation, noise correlations have also been studied in neural ensembles coding for continuous variables like position [30] and head direction [31]. In principle, the modeling approach we develop is applicable to any of these experimental settings. This is in contrast to experimental settings that employ a discrete set of stimuli with no underlying continuous parameterization (such as olfactory responses to distinct classes of odors).
To formalize these concepts, let  $c \in \{1, \dots, C\}$  index the set of conditions, and let  $\boldsymbol{x}_c \in \mathcal{X}$  denote the stimulus or animal behavior in condition  $c$ . For example,  $\boldsymbol{x}_c$  could be a scalar denoting the orientation of a voluntary reach or a visual grating (Fig. 1B). However, the model we will describe is more general than this:  $\boldsymbol{x}_c$  could be a vector describing a multi-dimensional stimulus space. For example, we will later study motor preparatory activity to reach targets at different radii and angles (sec. 3.2). In this setting,  $\boldsymbol{x}_c$  is a 2-dimensional vector specifying these two parameters of the reach target (Fig. 1C, left). We will also study neural responses to drifting visual gratings (sec. 3.3). In this setting,  $\boldsymbol{x}_c$  is a vector specifying the speed and orientation of the grating stimulus (Fig. 1C, right).
Recall that  $\mathbf{y}_{ck} \in \mathbb{R}^N$  denotes the simultaneously measured response of  $N$  neurons to condition  $c \in \{1, \dots, C\}$  on trial  $k \in \{1, \dots, K\}$ . Our goal is to model these measured responses as samples
from some distribution  $F$  with mean  $\pmb{\mu}$  and covariance  $\pmb{\Sigma}$ :
$$
\boldsymbol {y} _ {c k} \sim F (\boldsymbol {\mu} (\boldsymbol {x} _ {c}), \boldsymbol {\Sigma} (\boldsymbol {x} _ {c})) \quad \text {i n d e p e n d e n t l y f o r} c, k \in \{1, \dots C \} \times \{1, \dots K \}. \tag {3}
$$
Importantly, the functions  $\pmb{\mu}:\mathcal{X}\mapsto \mathbb{R}^N$  and  $\Sigma :\mathcal{X}\mapsto \mathbb{R}_{++}^{N\times N}$  should be smooth so that the response distribution changes gradually as the experimental condition is altered. To accomplish this, we leverage the machinery of Gaussian processes, which can be used to specify a prior distribution over smooth functions [32]. While Gaussian process models have been extensively used within statistical neuroscience [33-37], they have typically been used to model smoothness in temporal dynamics rather than smoothness across experimental conditions (but see [38]).
In this paper, we assume that  $F(\boldsymbol{\mu}, \boldsymbol{\Sigma})$  is a multivariate Gaussian distribution, which enables tractable and convenient probabilistic inference (see sec. 2.4). Since many noise correlation analyses are applied to spike count data, it is of great interest to extend our model so that  $F$  is better suited to this case. Multivariate extensions of the Poisson distribution are an attractive option [39]. While the focus of this paper is on multivariate Gaussian distributions, we develop inference tools for the Poisson model and we leave further investigation of this model for future work (see 1F for the corresponding graphical model and Supplement B.5 for further details). We note that it is common to pre-process spike count data and calcium fluorescence traces (e.g. through square root or log transforms) so that they are better approximated by Gaussian models [33, 40, 41].
To model  $\pmb{\mu}(\cdot)$ , we stack independent Gaussian processes into an  $N$ -vector. To model  $\pmb{\Sigma}(\cdot)$ , we use a Wishart process [42]. In essence, the Wishart process (denoted by  $\mathcal{WP}$ ) stacks Gaussian processes into an  $N \times P$  matrix and uses this as an overcomplete basis for the covariance. Formally,
$$
\boldsymbol {\mu} (\cdot) \sim \mathcal {G P} ^ {N} \left(\boldsymbol {0}, k _ {\mu}\right), \quad \boldsymbol {U} (\cdot) \sim \mathcal {G P} ^ {N \times P} \left(\boldsymbol {0}, k _ {\Sigma}\right), \quad \boldsymbol {\Sigma} (\boldsymbol {x}) = \boldsymbol {U} (\boldsymbol {x}) \boldsymbol {U} (\boldsymbol {x}) ^ {\top}, \tag {4}
$$
where  $P \geq N$  is a hyperparameter, and  $\mathcal{GP}^N(m, k)$  denotes an  $N$ -dimensional Gaussian process with (constant) mean  $\boldsymbol{m} \in \mathbb{R}^N$  and a positive definite kernel function  $k: \mathcal{X} \times \mathcal{X} \mapsto \mathbb{R}$ , described below. For any fixed value of  $\boldsymbol{x} \in \mathcal{X}$ , the product  $\boldsymbol{U}(\boldsymbol{x})\boldsymbol{U}(\boldsymbol{x})^\top$  follows a Wishart distribution, which is a standard probabilistic model for covariance and inverse covariance (see sec. 3.6 of [43]). To our knowledge, Wishart processes have not previously been leveraged for neural data analysis.
The kernel functions  $k_{\mu}$  and  $k_{\Sigma}$  play important roles by respectively specifying the degree of smoothness in the mean and covariance of the neural response across conditions (Fig. 1D). The kernel functions must satisfy certain technical conditions, as reviewed in [44, 45]. We satisfy these conditions by defining a kernel for each coordinate of  $\boldsymbol{x} \in \mathcal{X}$  and evaluating the product,  $k(\boldsymbol{x}, \boldsymbol{x}') = \prod_{i} k_{i}(x_{i}, x_{i}')$ . This is guaranteed to produce a valid kernel ([44], Proposition 13.2). For non-periodic variables, we use a standard squared exponential kernel
$$
k \left(x, x ^ {\prime}\right) = \gamma \delta_ {x x ^ {\prime}} + \beta \exp \left[ - \left(x - x ^ {\prime}\right) ^ {2} / \lambda \right] \tag {5}
$$
where  $\{\gamma, \beta, \lambda\}$  are hyperparameters and  $\delta_{xx'} = 1$  when  $x = x'$  and otherwise evaluates to zero. For periodic conditions (e.g. orientation of visual gratings or radial reach directions) we use [46]:
$$
k \left(x, x ^ {\prime}\right) = \gamma \delta_ {x x ^ {\prime}} + \beta \exp \left[ - \sin^ {2} \left(\pi | x - x ^ {\prime} | / T\right) / \lambda \right] \tag {6}
$$
where  $T$  specifies the period of the kernel. For both kernels, the most critical hyperparameter is  $\lambda$  (often called the "bandwidth"), which specifies the degree of smoothness across conditions. The remaining hyperparameters,  $\gamma$  and  $\beta$ , respectively ensure that the Gaussian process is well-conditioned and scaled. We tune all of these hyperparameters to maximize heldout log-likelihood by cross-validation. Throughout, we use  $\{\gamma_{\mu}, \beta_{\mu}, \lambda_{\mu}\}$  and  $\{\gamma_{\Sigma}, \beta_{\Sigma}, \lambda_{\Sigma}\}$  to respectively represent the kernel hyperparameters for  $k_{\mu}$  and  $k_{\Sigma}$ .
2.3 Low-Rank Plus Diagonal Covariance and Other Model Variants
The vanilla Wishart process model in equations (3) and (4) can be extended in several ways, which we found to be useful in practical applications. First, we can augment the covariance model by adding a diagonal term, resulting in a model that is reminiscent of factor analysis:
$$
\boldsymbol {\Sigma} (\boldsymbol {x}) = \boldsymbol {U} (\boldsymbol {x}) \boldsymbol {U} (\boldsymbol {x}) ^ {\top} + \boldsymbol {\Lambda} (\boldsymbol {x}), \tag {7}
$$
where  $\Lambda(\pmb{x}) \in \mathbb{R}^{N \times N}$  is a condition-dependent nonnegative diagonal matrix where the diagonal terms are independent samples from a  $\mathcal{GP}^N$  that are transformed using a softmax function,
$f(x) = \log (1 + e^{x})$ , to ensure positivity. We term this extension of  $\mathcal{WP}$  as  $\mathcal{WP}^{lrd}$  (for low-rank plus diagonal). Intuitively, this model extension has fewer parameters than a vanilla  $\mathcal{WP}$  and hence discourages overfitting. Additionally, the diagonal term,  $\Lambda (\pmb {x})$ , helps ensure that the covariance matrices remain well-conditioned during inference. Indeed, we can choose values of  $P < N$  (or even  $P = 0$ ), which enables us to specify prior distributions over the covariance that are "low-dimensional" in the sense that a large fraction of variance can be explained by a small number of dimensions. This comports with numerous experimental observations [10, 47]. As can be seen in Supplement Fig. 5, the dimensionality of matrices sampled from a vanilla Wishart distribution scales linearly with  $N$ , while matrices sampled according to equation (7) saturate at lower values.
Next, we additionally incorporate a lower-triangular matrix  $\pmb{L} \in \mathbb{R}^{N \times N}$  into the covariance model to capture aspects of covariance that are condition-independent (i.e. do not depend on  $\pmb{x}$ ):
$$
\boldsymbol {\Sigma} (\boldsymbol {x}) = \boldsymbol {L} (\boldsymbol {U} (\boldsymbol {x}) \boldsymbol {U} (\boldsymbol {x}) ^ {\top} + \boldsymbol {\Lambda} (\boldsymbol {x})) \boldsymbol {L} ^ {\top}, \tag {8}
$$
In the literature  $\boldsymbol{L}\boldsymbol{L}^T$  is referred to as the scale matrix. Some intuition can be gained by assuming  $\boldsymbol{\Lambda}(\boldsymbol{x}) = \mathbf{0}$  and  $P > N$ , in which case  $\boldsymbol{U}(\boldsymbol{x})\boldsymbol{U}(\boldsymbol{x})^\top \propto \boldsymbol{I}$  in expectation under the marginal prior distribution. Thus, a good choice is to set  $\boldsymbol{L}$  to be the Cholesky factor of the grand empirical covariance as defined in eq. (2). In our experiments, we use this as an initialization for  $\boldsymbol{L}$  and we optimize it to maximize the evidence lower bound (ELBO) as described in the Supplement B.3.
It is also possible, and potentially beneficial, to instead model smoothness in the inverse covariance matrix (also known as the precision matrix). This results in an inverse Wishart process. We refer the reader to Supplement B.2 for additional details of the inverse Wishart model.
2.4 Latent Variable and Parameter Inference
Each of the probabilistic models above specifies a posterior distribution of the noise covariance condition on observed neural activity,  $p\left(\left\{\boldsymbol{\Sigma}(\boldsymbol{x}_c)\right\}_{c=1}^C \mid \{\boldsymbol{y}_{ck}\}_{c=1,k=1}^{C,K}\right)$ . More precisely, each model specifies a posterior on the latent covariance factor matrices,  $p\left(\left\{\boldsymbol{U}(\boldsymbol{x}_c)\right\}_{c=1}^C \mid \{\boldsymbol{y}_{ck}\}_{c=1,k=1}^{C,K}\right)$ , which defines the desired posterior over covariance through a measurable transformation (eq. 8). This distribution over latent covariance factors has  $NPC$  dimensions, precluding us from representing the posterior exactly (doing so would require us to tile  $\mathbb{R}^{NPC}$  with a lattice of points). We instead use mean field variational inference [48] to approximate the desired posterior. Since this approach has already been described for similar Wishart process models [49], we relegate a description to Supplement B.3. All codes are implemented in numpyro [50] and available at https://github.com/neurostatslab/wishart-process.
2.5 Connections to Population Decoding
As we discussed in section 1, the estimation of noise covariance matrices is an important step for a variety of downstream neural analyses. Here, we briefly highlight two applications of Wishart process models in the context of neural decoding and perceptual discrimination [3].
Quadratic Discriminant Analysis (QDA) A simple decoding task is to infer which condition  $c \in \{1, \dots, C\}$  gave rise to an observed neural response,  $\boldsymbol{y} \in \mathbb{R}^{N}$ , on a single trial. A natural approach is to use a maximum likelihood classifier,  $\text{class}(\boldsymbol{y}) = \arg \max_{c \in \{1, \dots, C\}} \log p(\boldsymbol{y} | \hat{\boldsymbol{\mu}}_c, \hat{\boldsymbol{\Sigma}}_c)$ . Doing so requires that we construct estimates for the mean responses,  $\hat{\boldsymbol{\mu}}_1, \dots, \hat{\boldsymbol{\mu}}_C$ , and covariances,  $\hat{\Sigma}_1, \dots, \hat{\Sigma}_C$ . Since covariance estimation is challenging, a common practice (see e.g. [10]) is to use the grand empirical covariance (eq. 2) and therefore set  $\hat{\Sigma}_1 = \dots = \hat{\Sigma}_C$ . Under this model, known as linear discriminant analysis (LDA; see e.g. [51] Ch. 4), it can be shown that the decision boundaries of the maximum likelihood classifier become linear.
Intuitively, if we can accurately estimate distinct covariances for each condition, we may unlock more powerful decoders by relaxing the assumption of equal covariance across conditions. Under this model, known as quadratic discriminant analysis (QDA), the decision boundaries are described by hyperbolas (A schematic is shown in Fig. 2D). Furthermore, Pagan, Simoncelli, and Rust [52] described a circuit implementation of QDA using a cascade of two linear-nonlinear models, suggesting that such classifiers are a simple and biologically plausible neural computation. We demonstrate the ability of Wishart process models to enable accurate decoding via QDA in Fig. 2E.
Figure 2: Results on synthetic data. See Supplement B.1 for all simulation details. (A) Simulated (A1) and inferred (A2-A5) response distributions using different methods. Means  $\pmb{\mu}(\pmb{x}_c)$  and covariances  $\pmb{\Sigma}(\pmb{x}_c)$  are shown in dashed lines and colored ellipses in PC space. Train and test trials  $\pmb{y}_{ck}$  are shown in black and red dots. (B1-B4) Mean log-likelihood of held-out trials as a function of the number of dimensions, conditions, trials, and covariance smoothness. Error bars represent median  $\pm$  SEM across 5 independent runs. (C1-C4) Same as in B, the operator norm between the true and inferred covariances. (D) Schematic of QDA classifier. (E) Using different estimates of covariance, we use QDA to infer the condition label of held-out trials. Error bars represent the standard deviation over 6 independent runs. (F) Schematic of Fisher Information, FI. (G) Estimation of FI using grand empirical covariances and Wishart model.