SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Yizi Zhang $^{1*}$  Tianxiao He $^{1,2*}$  Julien Boussard $^{1}$  Charlie Windolf $^{1}$  Olivier Winter $^{3}$  Eric Trautmann $^{1}$  Noam Roth $^{4}$  Hailey Barrell $^{4}$  Mark Churchland $^{1}$  Nicholas A. Steinmetz $^{4}$  The International Brain Laboratory $^{3}$  Erdem Varol $^{2}$  Cole Hurwitz $^{1}$  Liam Paninski $^{1}$
$^{1}$ Columbia University  $^{2}$ New York University  $^{3}$ The International Brain Laboratory  $^{4}$ University of Washington in Seattle
Neural decoding and its applications to brain computer interfaces (BCI) are essential for understanding the association between neural activity and behavior. A prerequisite for many decoding approaches is spike sorting, the assignment of action potentials (spikes) to individual neurons. Current spike sorting algorithms, however, can be inaccurate and do not properly model uncertainty of spike assignments, therefore discarding information that could potentially improve decoding performance. Recent advances in high-density probes (e.g., Neuropixels) and computational methods now allow for extracting a rich set of spike features from unsorted data; these features can in turn be used to directly decode behavioral correlates. To this end, we propose a spike sorting-free decoding method that directly models the distribution of extracted spike features using a mixture of Gaussians (MoG) encoding the uncertainty of spike assignments, without aiming to solve the spike clustering problem explicitly. We allow the mixing proportion of the MoG to change over time in response to the behavior and develop variational inference methods to fit the resulting model and to perform decoding. We benchmark our method with an extensive suite of recordings from different animals and probe geometries, demonstrating that our proposed decoder can consistently outperform current methods based on thresholding (i.e. multi-unit activity) and spike sorting. Open source code is available at https://github.com/yzhang511/density_decoding.
Decoding methods for large-scale neural recordings are opening up new ways to understand the neural mechanisms underlying cognition and behavior in diverse species (Urai et al., 2022). The emergence of high-density multi-electrode array (HD-MEA) devices introduced a tremendous increase in the number of extracellular channels that can be recorded simultaneously (Jun et al., 2017; Steinmetz et al., 2021), leading to scalable and high-bandwidth brain computer interfaces (BCI) systems (Musk et al., 2019; Paulk et al., 2022).
Traditional neural decoding methods assume that spiking activity has already been correctly spike-sorted. As a result, these methods are not appropriate for situations where sorting cannot be performed with high precision. Despite intensive efforts towards automation, current spike sorting algorithms still require manual supervision to ensure sorting quality (Steinmetz et al., 2018). Even after careful curation, current spike sorters suffer from many sources of errors including erroneous spike assignment (Deng et al., 2015). The dense spatial resolution of HD probes makes some known issues of spike sorting even more evident. With the increased density of the recording channels, the
probability of visibly overlapping spikes (spike collisions) is higher (Buccino et al., 2022). Even for the same HD dataset, different spike sorters have low agreement on the isolated units and can find a significant number of poorly sorted and noisy units (Buccino et al., 2020). Consequently, only single units that pass quality control metrics are included in many neural coding studies (IBL et al., 2022).
Because the spike-sorting problem remains unresolved, alternative approaches that do not rely on sorted single-units for decoding have been proposed. A popular choice is multi-unit threshold crossing that uses spiking activity on each electrode for decoding (Fraser et al., 2009; Trautmann et al., 2019). However, this approach ignores the fact that the signal on each electrode is a combination of signals from different neurons, thus making inefficient use of the data (Todorova et al., 2014). Ventura (2008) proposed a spike-sorting free decoding paradigm that estimates neuronal tuning curves from electrode tuning curves and then infers the behavior of interest using the estimated tuning curves and newly observed electrode spike trains. More recently, Chen et al. (2012); Kloosterman et al. (2014); Deng et al. (2015); Rezaei et al. (2021) developed spike feature decoding methods that use marked point processes to characterize the relationship between the behavior variable and features of unsorted spike waveforms. However, these decoders based on state-space models make explicit assumptions about the underlying system dynamics which reduce their flexibility in capturing complex relationships in the data. Moreover, these methods mainly utilize simple waveform features for decoding such as the maximum amplitude on each electrode and do not take advantage of HD spike features such as the estimated spike location.
To leverage the spatial spread and density of HD probes, Hurwitz et al. (2019); Boussard et al. (2021) developed spike localization methods. These methods estimate the source location of a detected spike; this is a low-dimensional feature that is informative about the firing neuron's identity. We propose a probabilistic model-based decoding method that scales to HD-MEA devices and utilizes these novel localization features in conjunction with additional waveform features. We use a mixture of Gaussians (MoG) model to encode the uncertainty associated with spike assignments in the form of parametric distributions of the spike features. Unlike traditional MoG models with a fixed mixing proportion, our method allows the mixing proportion to depend on the behavior of interest and change over time. This is motivated by the theory that behavioral covariates that modulate neurons' firing rates also contain information about spike identities and that such tuning information should be incorporated into spike sorting and neural decoding in order to obtain unbiased and consistent tuning function estimates (Ventura, 2009). To infer the functional relationship between spike features and behavioral correlates, we employ automatic differentiation variational inference (ADVI) (Kucukelbir et al., 2017) and coordinate ascent variational inference (CAVI) (Blei et al., 2017), which enable us to perform efficient and accurate inference while considering the behavior-modulated MoG model.
We apply our method to a large number of HD recordings and decode various types of behavioral correlates. Experimental results show that our decoder consistently outperforms decoders based on multi-unit threshold crossings and single-units sorted by Kilosort 2.5 (Pachitariu et al., 2023). We further validate the robustness of our method by applying it to recordings with different levels of sorting quality, HD probes with varying geometry, and recordings from multiple animal species. Consistent with previous results, our findings indicate that relying solely on "good" units, as determined by sorting quality metrics, leads to information loss and suboptimal decoding performance. This observation motivates our transition to a spike sorting-free decoding framework which enables us to extract more information from the spiking activity and improve decoding performance.
Consider an electrophysiological recording comprised of  $K$  trials, where each trial is divided into  $T$  equally spaced time bins. Let  $\{s_{itk}\}_{i=1}^{n_{tk}}$ ,  $s_{itk} \in \mathbb{R}^D$  denote a set of spike features, where  $i$  indexes the  $i$ -th spike,  $n_{tk}$  represents the number of spikes collected in the  $t$ -th time bin of the  $k$ -th trial, and  $D$  is the dimension of spike features. For example, the spike feature  $s_{itk} = (x_{itk}, z_{itk}, a_{itk}) \in \mathbb{R}^3$  includes the spike location along the x- and z-axis of the probe, and its maximum peak-to-peak (max ptp) amplitude. Let  $y_k \in \mathbb{R}^T$  be the observed time-varying behavior in the trial  $k$ , e.g., the speed of a rotating wheel controlled by a mouse. When the behavior in the trial  $k$  does not vary over time, it can take on either a binary ( $y_k \in \{0, 1\}$ ) or scalar ( $y_k \in \mathbb{R}$ ) value, e.g., the mouse responds to a
scalar-valued stimulus by making a binary decision.
The proposed decoding method comprises an encoder and a decoder model. During the training of the model, the encoder captures the relationship between the observed spike feature distribution and the observed behavior. During the testing phase, the decoder utilizes the learned relationship from the encoder to predict the unobserved behavior based on newly observed spike features. In the following section, we present a general formulation of the encoding-decoding paradigm and provide more detailed implementations in the supplementary materials.
Figure 1: Decoding paradigm and graphical model. (a) Spike localization features,  $(x,z)$ , the locations of spikes along the width and depth of the NP1 probe, and waveform features,  $a$ , the maximum peak-to-peak (max ptp) amplitudes of spikes. Amplitude is measured in standard units (s.u.). Spike features from the entire probe are shown, and we focus on a specific segment of the probe. (b) During the training phase, the encoder takes the observed spike features  $s$  and behavior  $y$  from the train trials as inputs and then outputs the variational parameters  $\theta$  which control the dependence of the firing rate  $\lambda$  on the behavior  $y$ . At test time, the decoder utilizes the learned model parameters  $\theta$  obtained from the encoder and the observed spike features  $s$  from the test trials to predict the corresponding behavior in the test trials. To ensure reliable decoding of behaviors, we initially calculate the  $\lambda$  during training using the learned  $\theta$  and observed behaviors  $y$  from the train trials. Then, we compute the  $\lambda$  during test time using the learned  $\theta$  and the estimated behavior  $\hat{y}$  obtained through multi-unit thresholding from the test trials. Finally, we generate the weight matrix  $W$  for both the train and test trials as input to the final behavior decoder, e.g., linear regression or neural networks (Glaser et al., 2020; Livezey and Glaser, 2021). (c) In the encoder, the firing rates of each MoG component  $\lambda_{ctk}$  are modulated by the observed behavior  $y_{tk}$  in the train trials. This modulation affects the MoG mixing proportion  $\pi_{ctk}$ , which in turn determines the spike assignment  $z_{itk}$  that generates the observed spike features  $s_{itk}$  in the train trials. In the decoder, the behavior  $y_{tk}$  in the test trials is unknown and considered as a latent variable to be inferred. The decoder uses the observed spike features  $s_{itk}$  from the test trials along with the fixed model parameters  $\theta_c$  learned by the encoder to infer the latent behavior  $y_{tk}$ .
Encoder. The multivariate spike feature distribution is modeled using a mixture of Gaussian (MoG). The encoder generative model is as follows:
$$
\lambda_ {c t k} = \lambda \left(t, y _ {k}, \theta_ {c}\right), \theta_ {c} \sim p \left(\theta_ {c}\right), \tag {1}
$$
$$
z _ {i t k} \sim \text {C a t e g o r i c a l} \left(z _ {i t k}; \pi_ {t k}\right), \pi_ {t k} = \left\{\pi_ {c t k} \right\} _ {c = 1} ^ {C}, \pi_ {c t k} = \frac {\lambda_ {c t k}}{\sum_ {c ^ {\prime}} \lambda_ {c ^ {\prime} t k}}, \tag {2}
$$
$$
s _ {i t k} \sim \mathcal {N} \left(s _ {i t k}; \eta_ {z _ {i t k}}\right), \eta_ {c} = \left(\mu_ {c}, \Sigma_ {c}\right), \tag {3}
$$
where  $\lambda (\cdot)$  is a function that describes the firing rate's dependence on behaviors  $y_{k}$ , while  $p(\theta_c)$  represents a general prior on  $\theta_c$  encompassing the model parameters for the mixture component  $c$ . Intuitively, the behavior-dependent  $\lambda$  governs the mixing proportion  $\pi$  of the MoG, which determines the specific component  $c$  from which a spike feature  $s$  is generated. As  $\lambda$  varies over time in response to  $y$ , spikes that are spatially close and share similar waveform features may originate from different MoG components at different time points within a trial. In our implementation, we parameterize  $\lambda$  using a generalized linear model (GLM), but alternative models such as neural networks (NN) can also be used; see the supplementary material for the GLM configuration.
We employ variational inference (VI) to learn the unknown quantities. In the standard MoG setting (Blei et al., 2017), our goal is to infer the spike assignment  $z$  which indicates the latent component from which the observation  $s$  originates. However, unlike the standard MoG, our spike assignment  $z$  is influenced by the firing rates  $\lambda$  of the neurons which are modulated by the behavior  $y$ . Consequently, learning the association between  $\lambda$  and  $y$  necessitates the estimation of the unknown model parameters  $\theta$ . Our objective is to simultaneously learn both the latent variables  $z$  and model parameters  $\theta$  based on the observed spike features  $s$  and behavior  $y$ . To accomplish this, we posit a mean-field Gaussian variational approximation
$$
q (z, \theta) = \prod_ {c, t} q \left(z _ {c t}\right) q \left(\theta_ {c}\right) \tag {4}
$$
for the posterior  $p(z, \theta \mid s, y)$ . Subsequently, we employ the CAVI or ADVI methods to maximize the evidence lower bound (ELBO) and compute updates for  $z$  and  $\theta$ . Analogous to the standard Expectation-Maximization (EM) algorithm for MoG, the proposed CAVI and ADVI procedures consist of an E step and a M step. The E step, which updates  $z$ , closely resembles that of the ordinary MoG, while the M step, responsible for updating  $\theta$ , differs. CAVI utilizes coordinate ascent to find  $\theta$  that maximizes the ELBO, while ADVI employs stochastic gradient ascent for  $\theta$  updates. For detailed information on the CAVI and ADVI model specifications, refer to Supplement 1 and 2.
Decoder. The decoder adopts the same generative model in Equation 1-3 as the encoder with two distinctions: 1)  $y_{k}$  is unobserved and considered a latent variable we aim to estimate, i.e.,  $y_{k} \sim p(y_{k})$ , where  $p(y_{k})$  is a general prior, and 2) the model parameters  $\theta$  are obtained from the encoder and kept constant. In practice, the choice of prior relies on the nature of  $y_{k}$ . For instance, if  $y_{k}$  is binary, we can sample from a Bernoulli distribution while a Gaussian process prior can capture the temporal correlation between time steps if  $y_{k} \in \mathbb{R}^{T}$ . The posterior  $p(z,y\mid s)$  is approximated using a mean-field Gaussian variational approach
$$
q (z, y) = \prod_ {c, t} q \left(z _ {c t}\right) q (y). \tag {5}
$$
We employ standard CAVI or ADVI methods to infer  $z$  and decode  $y$ .
Robust behavior decoding. In practice, we found that direct decoding of  $y$  using the approximated posterior  $q(y)$  in Equation 5 was not robust, leading to decoding results of inconsistent quality across different datasets. Although the factorization described in Equation 5 may not fully exploit the available information for predicting  $y$ , it is useful for learning about the spike assignment  $z$ . To enhance decoding robustness, we compute a weight matrix from the MoG outputs as input to the final behavior decoder. The weight matrix, denoted as  $W$ , has dimensions  $K \times C \times T$  and entries  $W_{kct} \coloneqq \sum_{i=1}^{n_{tk}} q(z_{ikct})$ , which capture the posterior probability of assigning spike  $i$  collected at time  $t$  of the trial  $k$  into the component  $c$ . In scenarios such as spike sorting or multi-unit thresholding, spikes are assigned deterministically to one of  $C$  sorted single units or thresholded channels. In this case,  $W$  has one-hot rows, and each entry  $W_{kct}$  represents the number of spikes belonging to trial
$k$ , time  $t$  and unit (channel)  $c$ . To obtain  $q(z_{ikct})$ , we rely on estimating the posterior  $\pi_{ctk}$ , which requires either the observed  $y_{tk}$  or the estimated  $\hat{y}_{tk}$ . During model training, we can substitute the observed  $y_{tk}$  into Equation 1 to calculate the posterior  $\pi_{ctk}$  for the train trials. At test time, we use the estimated  $\hat{y}_{tk}$  obtained from multi-unit thresholding along with the learned encoder parameters  $\theta$  to calculate the posterior  $\pi_{ctk}$  for the test trials. With the posterior  $\pi_{ctk}$  in hand, we then estimate  $q(z)$  to compute the weight matrix  $W$  for both the train and test trials, which serves as input to the final behavior decoder. The choice of the behavior decoder depends on the user's preference. For instance, we can use logistic regression as the behavior decoder for binary  $y_{k}$ , and ridge regression for  $y_{k} \in \mathbb{R}^{T}$  that exhibit temporal variations. Additional information regarding the selection of the behavior decoder can be found in Supplement 5.
We conducted experiments using both electrophysiological and behavior data obtained from the International Brain Laboratory (IBL) (IBL et al., 2021). The electrophysiological recordings were acquired using Neuropixels (NP) probes that were implanted in mice performing a decision-making task. Each recording comprises multiple trials with several behavioral variables recorded during each trial including the choice, face motion energy, and wheel speed; see Figure 2 (a) for details. Each trial has a duration of 1.5 seconds and is divided into 30 time bins of 50 milliseconds length. The NP probe spans multiple brain regions; an example of the brain parcellation can be seen in Figure 3. To prepare the recordings for decoding, we first applied IBL's standard destriping procedure (Chapuis et al., 2022) to reduce artifacts. Then, we used a subtraction-based spike detection and denoising method described in Boussard et al. (2023). After preprocessing, we computed the spike locations (Boussard et al., 2021) to acquire spike features for decoding and then utilized registration techniques (Windolf et al., 2022) to correct for motion drift in the recorded data. Further details about data preprocessing can be found in Supplement 4. In all experiments, we used spike locations along the width and depth of the NP probe, and maximum peak-to-peak amplitudes of spikes for decoding. We selected this set of spike features based on empirical evidence from our experiments, which showed their good decoding performance. Furthermore, previous studies (Boussard et al., 2023; Hilgen et al., 2017) have also recognized that these features were highly informative about unit identity. Figure 1(a) illustrates the spike localization and waveform features that were utilized for decoding.
We evaluate the performance of our decoding method by comparing it to the following baselines: (1) Spike-thresholded decoders which utilize the spiking activity on each electrode after a voltage-based detection step. (2) Spike-sorted decoders that utilize all single-units found using Kilosort (KS) 2.5 (Pachitariu et al., 2016). (3) Spike-sorted decoders based on "good" units which consist of KS units that have passed IBL's quality control procedure (IBL et al., 2022). The parameters used for KS were tuned across multiple IBL datasets as described by IBL's spike sorting white paper (Chapuis et al., 2022).
To assess the quality of decoding, we perform 5-fold cross validation (CV) and compute relevant decoding metrics. The coefficient of determination  $(R^2)$  is used to evaluate continuous behavior decoding (e.g., motion energy and wheel speed) while accuracy is used for discrete behaviors (e.g., choice). To demonstrate the efficacy of our approach in a wide range of settings, we conduct the following experiments.
Varying levels of spike sorting quality. Our objective is to compare the proposed decoding method to spike-sorted decoders using datasets which have varying levels of spike sorting quality. We apply our method to two datasets with high sorting quality ("good" sorting) and two with low sorting quality ("bad" sorting). The quality of sorting is assessed using IBL's quality metrics (Chapuis et al., 2022). Although motion registration has been performed, we find that the recordings that have "bad" sortings are more affected by motion drift than the recordings which have "good" sortings; see supplementary materials.
Different brain regions from 20 datasets. To demonstrate the efficacy of our method across many different datasets, we decode 20 IBL datasets (IBL et al., 2022). In these datasets, mice perform a behavioral task while NP1 probes record activity from multiple brain regions. These brain regions
Figure 2: Density-based decoding is robust to varying levels of spike sorting quality. (a) We decode various behaviors including choice, motion energy and wheel speed. In the experimental setup, the mouse detects the presence of a visual stimulus to their left or right and indicates the perceived location (choice) by turning a steering wheel in the corresponding direction. Motion energy is calculated within a square region centered around the mouse's whiskers. The example behavior traces are distinguished by different colors for each trial. (b) We compare decoders using two experimental sessions with "good" sorting quality (represented by the color green) and two sessions with "bad" sorting quality (represented by the color purple) based on IBL's quality metrics. The box plots display various statistical measures including the minimum, maximum, first and third quartiles, median (indicated by a gray dashed line), mean (indicated by a red solid line), and outliers (represented by dots). These decoding metrics are obtained from a 5-fold CV and are averaged across both "good" and "bad" sorting example sessions. (c) We compare the traces decoded by spike-sorted decoders and our method on example sessions with "good" sorting quality (indicated by green) and "bad" sorting quality (indicated by purple). (d) The scatter plots depict the decoding quality of motion energy, measured by  $R^2$ , with respect to various spike-sorting quality metrics. Each point represents one of the 20 IBL sessions, and different colors and shapes are used to distinguish between the type of decoder and sorting quality. The sorting quality metrics include "contamination," which estimates the fraction of unit contamination (Hill et al., 2011), "drift," which measures the absolute value of the cumulative position change in micrometers per second (um/sec) of a given KS unit, and "missed spikes," which approximates the fraction of missing spikes from a given KS unit (Hill et al., 2011). These metrics are averaged across all KS units in a session. The scatter plots demonstrate that decoding quality tends to decrease when sorting quality is compromised. However, our method outperforms spike-sorted decoders even in the presence of these sorting issues.
are repeatedly targeted across all the datasets. To explore how behaviors are linked to specific brain regions, we use spikes that are confined to a particular area of the mouse brain for decoding. We collect decoding results from the posterior thalamic nucleus (PO), the lateral posterior nucleus (LP), the dentate gyrus (DG), the cornu ammonis (CA1) and the anterior visual area of the visual cortex (VISa).
Different probe geometry. Our method is capable of decoding electrophysiological data from a variety of HD probes. To demonstrate this, we apply our method on Neuropixels 2.4 (NP2.4) and Neuropixels 1.0 in nonhuman primates (NP1-NHP) datasets. The NP2.4 and NP1-NHP recordings are preprocessed using an identical pipeline as employed for NP1; see Supplement 4 for details. For spike sorting, the NP2.4 and NP1 adopt the same KS parameters as outlined in IBL's spike sorting procedure. Different KS parameters are utilized for NP1-NHP probes which are detailed in Trautmann et al. (2023).