SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Fisher Information (FI) A quantity that plays an important role in our understanding of how neural noise properties affect behavioral discriminability and representation is FI defined as  $\mathbb{E}[\partial^2\log p(\pmb {x}) / \partial \pmb{x}^2 ]$  [3]. Unfortunately, estimating FI is challenging because it involves measuring infinitesimal changes in the neural response distribution with respect to the stimulus condition (Fig. 2F). Wishart processes enable us to exploit the continuous parametrization of experimental conditions to derive a natural estimator for FI. For constant covariances across conditions, the FI for multivariate normal distributions is given by  $\mu^{\prime}(\pmb {x})^{T}\pmb{\Sigma}^{-1}\pmb{\mu}^{\prime}(\pmb {x})$  where  $\mu^{\prime}(\pmb {x})$  is the gradient of the mean with respect to  $\pmb{x}$ . If covariances vary across conditions, this quantity is replaced with  $\pmb {\mu}'(\pmb {x})^T\pmb {\Sigma}^{-1}(\pmb {x})\pmb {\mu}'(\pmb {x}) + \frac{1}{2}\mathrm{tr}([\pmb{\Sigma}^{-1}(\pmb {x})\pmb{\Sigma}'(\pmb {x})]^2)$  where  $\pmb{\Sigma}'(\pmb {x})$  is the gradient of the covariance w.r.t. the conditions and  $\mathrm{tr}$  is the trace operator. Wishart process models enable the continuous estimation of FI as a function of  $\pmb{x}$ . This is achieved by sampling from the gradient of the inferred posterior distribution (see Supplement B.6).
We qualitatively and quantitatively compared Wishart processes to common covariance estimation methods used in neural data analysis: Ledoit-Wolf [19], empirical covariance (eq. 1), grand empirical covariance (eq. 2), and graphical LASSO [21]. Our primary measure of performance is the log
likelihood that the models assign to held-out data. Because the baseline methods are not Bayesian, we handicap ourselves by using a single sample from the posterior distribution to calculate held-out log-likelihoods. We also visualize and qualitatively evaluate contours of the predictive probability density ("covariance ellipses") in the top two principal component subspace ("PC space"). For synthetic data, we report the difference between inferred and true covariances measured in the operator norm (Fig. 2C1-C4). We observed that generally log-likelihood and operator norm results are aligned. However, since they measure the performance in two different ways it is possible for a model to outperform another in log probability but vice versa on the covariance operator norm.
3.1 Synthetic Data
We first validate our approach on simulated data over a periodic 1-dimensional space, with conditions sampled at  $C$  equispaced angles on the interval  $[0,2\pi)$ . The mean and covariance of neural responses were sampled according to the Wishart process generative model (eq. 4) and trials for each condition were conditionally sampled according to eq. (3) from a multivariate Gaussian distribution. An example of generated data is shown in Fig. 2A1. Due to our choice for the condition space, we use periodic kernels both for the Gaussian and Wishart process priors. For simplicity, we set  $\Lambda(x)$  to the identity matrix both in the data generative and inference models. To match experimental observations [47], we simulated low-dimensional noise covariance structure by setting scale matrix to  $\boldsymbol{L}\boldsymbol{L}^{\top} = \boldsymbol{U}\mathrm{diag}(s_{1:N})\boldsymbol{U}^{T}$  where  $\boldsymbol{U}$  is a uniformly sampled orthogonal matrix in  $N$  dimensions and  $s_1 > s_2 > \dots > s_N > 0$  are logarithmically spaced numbers between 1 and  $10^{-5}$ . Except for the number of dimensions, conditions, trials, and Wishart smoothness ( $\lambda_{\Sigma}$ ), which we vary over a range to assess the performance of different models, all other hyperparameters are kept constant ( $P = 2$ ,  $\gamma = 0.001$ ,  $\beta = 1$ , and  $\lambda_{\mu} = 1$ ). Supplement B.1 provides further details for these simulations, covering all specific hyperparameter settings for each panel of Fig. 2.
We fit Wishart process models to these simulated datasets using the procedures described in section 2.4. These models qualitatively capture the ground truth covariance structure (Fig. 2A2-A5) and quantitatively outperform alternative methods across a range of synthetic datasets. Wishart process inference performs particularly well in high-dimensional regimes relevant to neural data analysis with large  $N$ , small  $K$ , and intermediate levels of smoothness across conditions (Fig. 2B-C). We also investigated the ability of Wishart process models to improve decoding accuracy on this dataset via QDA, as summarized in section 2.5. The Wishart process model outperformed the baseline models across a range of data smoothness suggesting that it captures the boundaries between classes more accurately (Fig. 2E). Furthermore, in Fig. 2G we show that on synthetic data the Wishart process model accurately estimates the FI continuously across conditions.
3.2 Primate Reaching
To study how the brain plans upcoming actions, animals are often trained to prepare voluntary movements to cued locations during a "delay period" before executing those movements after a "go cue" is delivered [53-55]. During this delay period, neural firing rates in the primary motor cortex (M1) and dorsal premotor cortex (PMd) correlate with features of the upcoming movement (e.g. direction and distance) in nonhuman primates [56, 57]. Trial-to-trial variability in the delay period has been suggested to correlate with fluctuations in reaction time [58], but a quantitative account of this variability with dense target configurations is challenging to achieve [29].
We analyzed 200ms of spiking activity prior to the onset of the "go cue" across  $N = 192$  multiunit channels from both PMd and M1 in a nonhuman primate executing point-to-point hand reaches. The dataset contains  $C = 48$  conditions, where targets were organized in three concentric circles at radii of  $4\mathrm{cm}$ ,  $8\mathrm{cm}$ , and  $12\mathrm{cm}$ , with 16 equidistant targets around each ring (Fig. 3A). We used a squared exponential kernel (eq. 5) to model smoothness across reaches of different lengths and a  $2\pi$ -periodic kernel (eq. 6) to model smoothness across reach angles. Models were fitted on  $60\%$  of the available trials, radial and angular smoothness parameters for both GP and WP kernels were selected on a validation set of  $25\%$  of the data, and log-likelihood scores were generated using the remaining  $15\%$  of trials. This procedure was repeated over 8 randomized data partitions.
The Wishart process model outperforms competing estimators (Fig. 3B). Since  $N > K$ , the empirical covariance (eq. 1) is singular and therefore produces negative infinity log-likelihoods. Instead,
3Data was previously published and made available by [29].
Figure 3: Results on primate reaching task. (A) Task schematic (left) and set of experimental conditions (right). (B) Log-likelihood of held-out trials across 8 randomized folds. Horizontal lines indicate the maximum, minimum, and mean of the log-likelihood distribution. (C) Marginal variance of two example neurons across reaching angles, averaged over reach radii (green) and Wishart process model estimate (black). Shaded green denotes mean  $\pm$  SEM. (D) Wishart (left) and Ledoit-Wolf (right) estimates of mean and covariance, projected onto the first two principal components. To demonstrate the ability of the Wishart Process model to predict covariance in unseen conditions, the middle ring of targets was held out in training (red ellipses). (E) Predicted mean and (F) covariance in held-out conditions shown in panel D. Covariance ellipses, grand-empirical estimates (dashed lines), and samples (black dots) are visualized in the top-2 principal component subspace.
weighted averages of the grand empirical and empirical covariance were calculated according to  $\boldsymbol{\Sigma}^{\mathrm{WA}}(\boldsymbol{x}_c) = \alpha \boldsymbol{\Sigma}^{\mathrm{E}}(\boldsymbol{x}_c) + (1 - \alpha)\boldsymbol{\Sigma}^{\mathrm{GE}}$  (WA, E, and GE denote weighted average, empirical, and grand-empirical estimators respectively). This baseline is a simplified variant of [20] shrinking the per-condition covariances towards the grand-empirical estimate. Performance results produced by this baseline are shown in Fig. 3B at different values of  $\alpha$ .
These results suggest that WP is able to capture the smooth change in the covariance across conditions. Since smooth changes in the covariance imply smooth changes in the marginal variance (diagonal terms in the covariance), we plot the marginal variance of individual neurons as a function of the reaching angle to visualize the smooth change. We see that the variance indeed changes across conditions and that the Wishart process model provides a smooth fit to these noisy variance estimates (Fig. 3C). This smoothness is absent in the predictions of the Ledoit-Wolf and graphical LASSO estimators, while the grand covariance estimate is constant and, in essence, over-smoothed (Supp. Fig. 2). The out-of-sample log probability results suggest that the smoothness captured by WP is indeed a property that is present in the data and has physiological relevance (Fig. 3B).
We also investigated the Wishart model's ability to interpolate covariance estimates for unseen conditions. Here, the model was retrained on only trials from the inner and outermost rings (4cm and 12cm). The neural response means and covariances around the center ring were then estimated by sampling from the model posterior evaluated at these hold-out condition locations (8cm, all angles).
Training Conditions
Condition 1
Figure 4: Results on drifted gratings dataset. (A) Task schematic (left) and set of experimental conditions (right). (B) Cross-validated log probability as a function of Wishart smoothness  $(\lambda_{\Sigma})$  and the number of components  $(P)$  achieving their optimal values at  $\lambda_{\Sigma} \approx 1.5$  and  $P = 0$  and outperforming alternatives. (C) Top: mean responses and noise covariances of the full set of conditions plotted in the PC space for the Ledoit-Wolf (left) vs. Wishart (right) estimators (covariances are scaled down by a factor of 100 for visualization purposes). Colors are in the increasing order of condition number matching the colors in (A); Bottom: Train and test trials (black and red dots) in two example conditions plotted in the top two PC subspace. Covariance ellipses are color-coded by the model as in panel B. (D) Generalizing to unseen conditions: Top:  $80\%$  of the conditions are used for training (shown in colors) and the remaining conditions (shown in red) are withheld. Bottom: Inferred covariances by Wishart Process model for two held-out conditions (as in panel C, bottom). Notice that panels C, and D are from two separate analyses. In panel C, we train on all conditions and test on held out trials. In panel D, we train on a subset of conditions and test on held out conditions (interpolation over stimulus space). Condition 1 was held out only in the second analysis.
Condition 25
Unseen Test Conditions
Condition 1
Wishart model estimations perform comparably to empirical mean and grand-empirical covariance estimations calculated on these hold-out conditions (Fig. 3E-F).
3.3 Allen Brain Observatory (Neuropixels)
Perceptual discrimination is a longstanding area of research in psychology and neuroscience, dating back to work in the early 1800s by Weber and Fechner [59]. A classic paradigm requires subjects to discriminate between pairs of oriented visual gratings, which humans and nonhuman primates accomplish with high precision [60-62]. Discrimination thresholds on this task are thought to be related to the scale and orientation of noise in neural representations [7] (Fig. 1A), and recent works have examined this hypothesis by simultaneously recording large populations of neurons in rodent models [10, 25, 26]. The resulting analyses are delicate and require a large number of trials.
To investigate, we analyzed simultaneously recorded responses in mouse primary visual cortex to drifting visual gratings in the Allen Brain Observatory (Visual Coding: Neuropixels dataset). We analyzed a session (session-id = 756029989) with 81 neurons. Spike count responses in a 200 ms
<sup>4</sup>https://portal.brain-map.org/explore/circuits/visual-coding-neuropixels
window from the stimulus onset were recorded across 4 grating orientations, 2 directions, and 5 temporal frequencies as illustrated in Fig. 4A. We used periodic kernels for the orientation and direction axes (eq. 6), and a squared exponential kernel (eq. 5) for the temporal frequency axis of the condition space. For each of the resulting 40 conditions, 15 trials are collected, which we split into 13 train vs. 2 test trials. For choosing the number of components  $(P)$  and the smoothness of the Wishart kernel  $(\lambda_{\Sigma})$ , we quantified the cross-validated log probability of test trials over 10 randomized folds. As depicted in Fig. 4B the optimal performance of the Wishart model is achieved for  $\lambda_{\Sigma} \approx 1.5$  and  $P = 0$ .
This selection of the model parameters allows us to capture the structure that is shared among all conditions (note the similarity between Wishart and grand empirical covariance in Fig. 4C), while representing the condition-specific changes in the covariance structure (note the smooth change of matrix orientation across conditions in the top-right plot of Fig. 4C). Finally, we present covariance interpolation results across unseen conditions where we fit the model according to 32 randomly selected train conditions (colored ellipses in Fig 4D top-left) and leave out the complete set of trials for the remaining 8 conditions (red ellipses in Fig 4D top-right). The interpolated means and covariances for two example test conditions present a strong agreement with their corresponding held-out trials (Fig 4D bottom row).
We proposed a method to estimate covariance across large populations of co-recorded neurons, which is widely regarded as a challenging problem [8, 9] with important implications for theories of neural coding [3, 5, 7]. Specifically, we formulated probabilistic Gaussian and Wishart process models [42] to model smoothly changing covariance structure across continuously parameterized experimental conditions. We found this approach outperforms off-the-shelf covariance estimators that do not exploit this smoothness (e.g. Ledoit-Wolf) as well as the grand covariance estimator (eq. 2) which is infinitely smooth (i.e. does not allow for condition-dependent covariance). We investigated two task modalities (vision and motor) and model organisms (mouse and nonhuman primate). However, we expect these modeling principles will be even more broadly applicable, including to other sensory modalities with smoothly parameterized stimuli (e.g. acoustic tones or odorant mixtures) and in neural circuits supporting navigation [30, 31].
We believe our work also provides an important conceptual advance. Given the limited duration of experiments, it may seem that there is a strict tradeoff between accurately estimating neural responses across many experimental conditions and estimating trial-by-trial noise statistics in any fixed condition. Our results show that this tradeoff is not so simple. Indeed, if neural responses are sufficiently smooth across conditions, the Wishart process model could predict neural covariance accurately by densely sampling conditions with only a single trial per condition (e.g. visual gratings at random orientations, as done in [25]). Intuitively, the model works by pooling information across nearby conditions, and so can even predict covariance in entirely unseen conditions (as demonstrated in Figs. 3D and 4D). Thus, the limiting factor in estimating noise covariance is not the number of trials per condition, but the smoothness of the neural responses over the space of experimental conditions.
We conclude by discussing the current limitations of our model and opportunities for future work. First, the model requires us to explicitly specify kernel functions  $(k_{\mu}, k_{\Sigma})$  over the space of experimental conditions. We studied simple paradigms where a reasonable parametric form of this kernel could be specified and the hyperparameters could be fine-tuned through cross-validation. Extending the model to more complex stimulus sets (e.g. to naturalistic image stimuli) would be of great interest, and likely require more sophisticated approaches for learning kernels and feature representations of these stimulus sets. Second, we assumed that trial-to-trial noise followed a Gaussian distribution. A deeper investigation of noise models, including multivariate extensions of Poisson distributions [39, 47, 63], is warranted. Finally, we utilized a basic variational inference procedure to approximate the posterior distribution over neural covariance. While fully representing the posterior is intractable (as is typically the case), more sophisticated optimization techniques may further benefit performance [64]. Nonetheless, our results show that Wishart process models are performant even in the absence of the aforementioned extensions. These models outperform common practices in neuroscience and open the door to new analyses, such as the prediction of noise covariance in held-out conditions.
We thank Saurabh Vyas (Columbia) and Eric Trautmann (Columbia) for helpful discussions and for providing data from the nonhuman primate motor cortex. We also thank Sam Zheng (NYU) and Sarah Harvey (NYU) for their feedback on the manuscript.
[58] Afsheen Afshar, Gopal Santhanam, Byron M. Yu, Stephen I. Ryu, Maneesh Sahani, and Krishna V. Shenoy. "Single-Trial Neural Correlates of Arm Movement Preparation". In: Neuron 71.3 (2011), pp. 555-564. ISSN: 0896-6273.
[59] Gustav Theodor Fechner. Elemente der psychophysik. Vol. 2. Breitbart u. Härtel, 1860.
[60] Michael A. Webster, Karen K. De Valois, and Eugene Swithkes. "Orientation and spatial-frequency discrimination for luminance and chromatic gratings". In: J. Opt. Soc. Am. A 7.6 (June 1990), pp. 1034-1049.
[61] D C Burr and S A Wijesundra. "Orientation discrimination depends on spatial frequency". In: Vision Res. 31.7-8 (1991), pp. 1449-1452.
[62] R Vogels and G A Orban. "How well do response changes of striate neurons signal differences in orientation: a study in the discriminating monkey". In: J. Neurosci. 10.11 (Nov. 1990), pp. 3543-3558.
[63] Olivier J Hénaff, Yoon Bai, Julie A Charlton, Ian Nauhaus, Eero P Simoncelli, and Robbe L T Goris. "Primary visual cortex straightens natural video trajectories". In: Nat. Commun. 12.1 (Oct. 2021), p. 5982.
[64] Yuri Burda, Roger Grosse, and Ruslan Salakhutdinov. "Importance weighted autoencoders". In: 2015.