SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Ankit Vishnubhotla*
av3016@columbia.edu
Charlotte Loh*
Liam Paninski
liam@stat.columbia.edu
MIT-IBM
Akash.Srivastava@ibm.com
Cole Hurwitz
ch3676@columbia.edu
Contrastive learning is quickly becoming an essential tool in neuroscience for extracting robust and meaningful representations of neural activity. Despite numerous applications to neuronal population data, there has been little exploration of how these methods can be adapted to key primary data analysis tasks such as spike sorting or cell-type classification. In this work, we propose a novel contrastive learning framework, CEED (Contrastive Embeddings for Extracellular Data), for high-density extracellular recordings. We demonstrate that through careful design of the network architecture and data augmentations, it is possible to generically extract representations that far outperform current specialized approaches. We validate our method across multiple high-density extracellular recordings. All code used to run CEED can be found at https://github.com/ankitvishnu23/CEED.
High-density extracellular recordings now allow for simultaneous recordings of large populations of neurons across multiple brain regions with high temporal and spatial resolution. [1-4]. These large-scale recordings are essential for gaining insights into key biological processes such as vision, decision-making, and behavior which are distributed across brain regions [5]. Along with gaining insights into brain function, these technologies also promise to improve the scalability and accuracy of brain-computer interfaces which can restore motor function to paralyzed individuals [6].
A major bottleneck for interpreting neural population activity is processing the raw extracellular signal [7]. Although extracellular recordings contain a precise record of the coordinated neural activity, it must be extracted algorithmically through a processing step called spike sorting. A crucial assumption in spike sorting is that each recorded neuron has a unique spatiotemporal extracellular waveform based on its morphology and position relative to the recording device [8]. Using this unique identifier, it is possible to assign a detected extracellular action potential (spike) back to its putative neuron (unit). Along with spike sorting, another important task in extracellular analysis is cell-type classification. As neural circuits are diverse and heterogeneous, it is becoming increasingly important to profile detected units using their morphoelectrical features [9]. At a coarse level, it is
1   
2   
3   
4   
5   
6   
7   
8   
9   
10   
11
possible to classify units as inhibitory or excitatory based on their extracellular profile [10]; however, it may be possible to divide units into finer subgroups [11-13].
For both spike sorting and cell-type classification, it is important to extract low-dimensional, meaningful features from extracellular waveforms. Despite the importance of feature extraction, current approaches are often ad hoc and lack robustness to common nuisance variables in extracellular recordings. By far the most common featurization method for spike sorting is principal components analysis (PCA) [3, 14-22]. Although PCA is relatively effective and scalable, it suffers from a few key drawbacks including: (1) a lack of robustness to extracellular nuisance variables such as spatiotemporally overlapping spikes (collisions), (2) an inability to model non-linear data, and (3) an objective function that aims to find features that explain variance rather than features that discriminate different waveforms. To improve the robustness of PCA, [21] introduced a supervised waveform denoiser that is able to 'clean' the waveforms before featurization. For morphoelectric cell-type classification, most feature extraction methods rely heavily on manually extracted features including action potential width, peak-to-peak amplitude, and the ratio of pre-hyperpolarization peak to the post-hyperpolarization peak [13]. Again, while scalable and effective, these features are too simple and ad hoc to fully capture morpho-electrical differences [13]. Recently, a non-linear approach to cell-type classification, WaveMap, was introduced, utilizing UMAP [23] and Louvain community detection [24] to automatically find cell-type clusters [13].
In this work, we introduce a robust and generalizable feature learning method, CEED, for extracellular datasets. Our main hypothesis is that embeddings of extracellular waveforms that are invariant to both common and task-specific nuisance variables will be more useful for spike sorting and morphoelectric cell-type classification than current specialized feature extraction methods. For example, in extracellular recordings, there are a number of nuisance variables including collisions, correlated background noise, or variability in the time at which a spike is detected [21]. Each of these confounds pose a problem for traditional representation learning methods like PCA or manually extracted features. Along with these common nuisance variables, there are also task-specific nuisance variables such as the spatial position of detected spikes. For spike sorting, this information is crucial as spike locations are highly informative of neuron identity [25, 26]. However, for cell-type classification, neurons with different locations may still share a cell-type. In order to extract representations that are invariant to these nuisance variables, we utilize contrastive learning, which has been shown to approximately induce invariance in the representation space to a set of transformations [27]. We utilize a recent contrastive learning framework [28] for our training and implement a stochastic view generation module for extracellular waveforms. We validate our approach on multiple high-density extracellular recordings. Surprisingly, for cell-type classification, our representations appear to be more informative than state-of-the-art specialized methods even when performing zero-shot learning on an unseen animal and probe geometry. Our contributions are as follows:
1. We introduce a novel framework, CEED, for analyzing extracellular recordings based on invariance learning.
2. We implement a stochastic view generation module for both single-channel and multichannel extracellular waveforms.
3. We demonstrate that CEED works well with multiple neural network architectures, including a novel transformer-based architecture with a spatiotemporal causal attention mask (SCAM).
4. We show that CEED outperforms specialized featurization methods for spike sorting and morphoelectric cell-type classification.
Contrastive Learning. Contrastive representation learning [28] falls under a broad class of self-supervised learning methods [29-31] whose goal is to learn robust and generalizable representations by encouraging invariances to a prior-known set of transformations or nuisance variables. While contrastive learning has mainly been used to extract effective and transferable representations from image-based data [27, 28], more recently, it has also become a powerful tool in the sciences to learn invariances for physical systems [32, 33].
In computational neuroscience, representation learning has mostly been performed using generative models [34]. While this has led to many interesting insights into behavior [35] and decision-making [36], this paradigm is sensitive to nuisance variables and will not capture subtle changes in the
28   
29   
30   
31   
32   
33   
34   
35   
36   
37   
38   
39   
40   
41   
42   
43   
44   
45   
46   
47   
48   
49   
50   
51   
52   
53   
54   
55   
56   
57   
58   
59   
60   
61   
62   
63   
64   
65   
66   
67   
68   
69   
70   
71   
72   
73   
74   
75   
76   
77   
78   
79
Figure 1: CEED framework. In CEED, we assume that waveforms are already extracted from an extracellular recording. Each waveform is then passed through our stochastic view generation module, where different views are obtained by applying transformations. These transformations induce a pre-defined set of invariances (see Section 3.2.1). Using these views, a neural-network based encoder, which can take the form of a multi-layer perceptron (MLP) or a transformer, is then trained to produce representations that respect the desired invariances. This is accomplished via contrastive learning, where representations from views of the same waveform are encouraged to be similar and views of different waveforms are encouraged to be dissimilar. When training completes, the learned representations can then be used for a series of downstream neuroscience tasks, such as spike sorting or morphoelectrical cell-type classification.
observation space that do not contribute much to the explained variance. While recent research into label-guided generative models [37-39] has somewhat alleviated these problems, it is still unclear if good generative performance is important for representation learning. Most recently, a number of contrastive learning methods have been introduced for learning robust and generalizable representations of neural population data [40-42]. To our knowledge, however, contrastive learning has not been applied to key primary data analysis tasks such as spike sorting or cell-type classification.
Spike sorting. In spike sorting, the goal is to extract a precise record of which neurons spiked at which time steps based on the raw extracellular data. For each electrode, the contributions of an unknown number of neurons are mixed together, making spike sorting a challenging blind-source separation task akin to the cocktail party problem [43]. While a number of spike sorting methods have been developed for microelectrode arrays (MEAs), almost all of these methods currently utilize PCA for feature extraction. For example, the popular algorithms Klusta, HerdingSpikes2, Mountainsort4, SpykingCircus, Trideclous, and Kilosort all use a form of SVD/PCA in their processing pipelines. Despite the ubiquity of PCA as a feature extraction method for spike sorting, it lacks robustness to common nuisance variables in extracellular recordings, cannot capture non-linear features of the observed data, and prefers features that explain variance rather than those that discriminate different waveforms. More recently, it has been shown that non-linear autoencoders can have higher performance than PCA despite suffering from similar drawbacks [44].
Cell-type classification. Morphoelectric cell-type classification has many similarities to spike sorting. Similarly to spike sorting, you are also interested in grouping detected extracellular waveforms together. However, the granularity of these clusters should be lower, reflecting shared morphologies among the units in the recording. Classically, pre-defined morphoelectrical features (e.g., action potential width) were extracted from the waveforms to group together similar units [10, 45]. More recently, a non-linear method based on UMAP was introduced for cell-type classification [13]. In this work, it was shown that automatically discovering relevant features from extracellular waveforms can lead to more informative cell-type groupings than when simply clustering pre-defined features. It is important to note that, along with morphoelectric cell-type classification, it is also possible to classify cell-types by their function, morphology, physiology, molecular properties [46-48].
82   
83   
84   
85   
86   
87   
88   
89   
90   
91   
92   
93   
94   
95   
96   
97   
98   
99   
100   
101   
102   
103   
104   
105   
106   
107   
108   
109
3 Contrastive Embeddings for Extracellular Data (CEED)
We introduce Contrastive Embeddings for Extracellular Data (CEED), a novel framework for extracting robust and generalizable representations of extracellular data via contrastive learning. CEED consists of three main components (see Figure 1) — (1) a stochastic view generation module that augments waveforms with both common and task-specific nuisance variables. (2) A neural network-based encoder that can extract low-dimensional representations from both single-channel and multi-channel waveforms. (3) A contrastive loss function that encourages the representations of views of the same waveform to be similar while forcing the representations of views of different waveforms to be dissimilar. We postulate that the representations extracted by CEED, which are approximately invariant to nuisance variables, can outperform specialized feature extraction methods for both spike sorting and cell-type classification.
To introduce our method, we must first define some notation for extracellular recordings. Let  $P_{M} \coloneqq \{p_{m} | m \in M\}$  be the position of all  $M$  channels on the multi-electrode array (MEA), where each  $p_{m} \in \mathbb{R}^{3}$  is the location of channel  $m$ . Now define  $P_{\tilde{M}}$  to be the set of positions for a subset of channels,  $\tilde{M} \subseteq M$ , on the MEA.
Let  $S \coloneqq \{s_n\}_{n=1}^N$  be the set of  $N$  spiking events that are detected in an extracellular recording. Now, let the extracellular waveform of a spiking event  $s_n$  on a channel  $m$  be defined as  $W_{n,m} \coloneqq \{r_{n,m}^{(0)}, r_{n,m}^{(1)}, \dots, r_{n,m}^{(T)}\} \in \mathbb{R}^T$ , where  $T$  contiguous samples from the recording define a waveform and  $r_{n,m}^{(t)}$  is the value of the waveform at sample  $t$ . This can be thought of as a window of samples of size  $T$  on channel  $m$  which includes the spike  $s_n$ .
The amplitude of a spike  $s_n$  on a channel  $m$  can be defined as  $A_{n,m} \coloneqq \max_t |W_{n,m}|$ . This is also known as the absolute maximum voltage recorded on channel  $m$ . The timestep  $t$  at which the absolute maximum voltage occurs can vary for each channel. Therefore, for a spike, we center the waveform on each channel using the timestep at which the maximum  $A_{n,m}$  occurred (i.e., the timestep of the maximum amplitude channel). For all spikes in the recording, waveforms are aligned in this manner so that the amplitude on the maximum amplitude channel occurs at the same timestep  $t$ .
Let  $\mathbf{W}_{n,\tilde{M}}\in \mathbb{R}^{T\times |\tilde{M}|}$  be the set waveforms for a spike on a subset of channels and let  $\mathbf{A}_{n,\tilde{M}}\in \mathbb{R}_{>0}^{|\tilde{M}|}$  be the set amplitudes for a spike on a subset of channels, where  $|\tilde{M} |$  is the number of channels.
3.2 Stochastic view generation module
3.2.1 Common invariances
Let us define the representation of the waveform  $\mathbf{W}_{n,\tilde{M}}$  for spike  $s_n$  as  $\mathbf{z}_n\in \mathbb{R}^D$ . There are a number of nuisance variables that change the observed waveform of a spike without changing the underlying signal. Therefore, we can define a set of invariances we wish to impose on our representations,
1. Amplitude voltage jitter - We want representations that are invariant to "small" amplitude variability. To this end, if we scale the amplitude of the waveform  $\mathbf{W}_{n,\tilde{M}}$  for spike  $s_n$  such that  $\widetilde{\mathbf{W}}_{n,\tilde{M}} = \mathbf{W}_{n,\tilde{M}}\times (1\pm \epsilon)$ , then the representation  $\tilde{\mathbf{z}}_n$  for  $\widetilde{\mathbf{W}}_{n,\tilde{M}}$  should be such that  $\tilde{\mathbf{z}}_n = \mathbf{z}_n$ .
2. Correlated background noise - We want representations that are invariant to the correlated background noise found in real extracellular recordings. We assume extracellular noise can be modeled as a spatiotemporal matrix Gaussian  $\mathcal{MN}(0,U,V)$ , where  $U\in \mathbb{R}^{T\times T}$  models the temporal correlation of a waveform and  $V\in \mathbb{R}^{|M|\times |M|}$  models the spatial correlation between channels. For a spike  $s_n$ , we want a representation such that if we sample  $n\sim \mathcal{MN}(0,U,V)$  and form  $\widetilde{\mathbf{W}}_{n,\tilde{M}} = \mathbf{W}_{n,\tilde{M}} + n_{\tilde{M}}$ , where  $n_{\tilde{M}}$  is the noise on the channels subset  $\tilde{M}$ , then  $\widetilde{\mathbf{W}}_{n,\tilde{M}}$  should still yield a representation  $\widetilde{\mathbf{z}}_n = \mathbf{z}_n$ .
3. Spike collisions - We want representations that are invariant to the voltage distortions caused when spikes "collide" spatiotemporally. A collision is when two or more spikes occur at similar times in the same location on the probe leading to observed waveforms that are distorted. To this end, consider two spikes  $s_{n_1}$  and  $s_{n_2}$  with waveforms  $\mathbf{W}_{n_1,\tilde{M}_1}$  and  $\mathbf{W}_{n_2,\tilde{M}_2}$  such that  $\tilde{M}_1$
110   
111   
112   
113   
114   
115   
116   
117   
118   
119   
120   
121   
122   
123   
124   
125   
126   
127   
128   
129   
130   
131   
132   
133   
134   
135   
136   
137   
138   
139   
140   
141   
142   
143   
144   
145   
146   
147   
148   
149   
150   
151   
152   
153   
154   
155   
156   
157
and  $\tilde{M}_2$  share some number of channels. Now let us shift  $s_{n_2}$  by  $k$  time steps  $(-T / 2 < k < 0$  or  $0 < k < T / 2)$  such that for any channel  $m\in \tilde{M}_2$ ,  $W_{n_2,m} := \{r_{n_2,m}^{(0 + k)},r_{n_2,m}^{(1 + k)},\dots,r_{n_2,m}^{(T + k)}\}$ . We refer to this temporally shifted waveform as  $\mathbf{W}_{n_2,\tilde{M}_2}^{\mathrm{shift}}$ . A new waveform that is a linear sum of these two waveforms  $\tilde{\mathbf{W}}_{n,\tilde{M}_1} = \mathbf{W}_{n_1,\tilde{M}_1} + \mathbf{W}_{n_2,\tilde{M}_2}^{\mathrm{shift}}$  should yield a representation  $\tilde{\mathbf{z}}_n = \mathbf{z}_{n_1}$ .
4. Channel subsets - We want representations that are invariant to changes in the subset of channels  $\tilde{M}$  used to define  $\mathbf{W}_{n,\tilde{M}}$ . So for a new subset of channels  $\tilde{M}'$ , as long as the channel with the highest amplitude is still contained in  $\tilde{M}'$ , then the waveforms  $\mathbf{W}_{n,\tilde{M}'}$  should still yield a representation  $\tilde{\mathbf{z}}_n = \mathbf{z}_n$ .
3.2.2 Task-specific invariances
While all feature extraction methods suffer from common extracellular nuisance variables, there are also task-specific nuisance variables. For example, the position and orientation of neurons in the recording are essential information for spike sorting. However, for cell-type classification, this information can be a confound when trying to find shared morphoelectrical features. Therefore, we propose an additional set of invariances for cell-type classification,
1. Cell position - For cell-type classification, we want representations that are invariant to the channel positions  $P_{M}$  at which a spike is detected. So for a waveform  $\mathbf{W}_{n,M}$ , any uniform changes to the channel positions  $P_{\tilde{M}}$  (without changing the waveforms at each channel) should still yield a representation  $\tilde{\mathbf{z}}_n = \mathbf{z}_n$ .
2. Cell amplitude - As amplitude is mainly a function of cell position and orientation, we want representations that are invariant to any uniform changes in the amplitudes. To this end, for any positive value  $a$ , a change in the amplitudes for the spike  $s_n$ ,  $\tilde{\mathbf{W}}_{n,\tilde{M}} = \mathbf{W}_{n,\tilde{M}} \times a$ , should still yield a representation  $\tilde{\mathbf{z}}_n = \mathbf{z}_n$ .
To achieve these task-specific invariances for cell-type classification, we directly transform the training data. First, we extract the max channel waveform for each spike. Then, we normalize each waveform to be between -1 and 1, thus removing any positional information [13]. A drawback of this approach is that we discard multi-channel information that may be useful for classifying different cell-types [12] (see Section 7 for a more detailed discussion). The full view generation pipeline is detailed in Supplementary Materials Section A.1.
3.3 Encoder Architecture
For the encoder of CEED, we explored two neural network architectures - a transformer-based network with a novel Spatiotemporal Causal Attention Mask (SCAM) and a simple multi-layer perceptron (MLP) which is a more computationally efficient alternative (i.e., can fit on a single GPU). For a runtime comparison of these two architectures, see Supplementary Materials Section E.
3.3.1 Transformer-based encoder with spatiotemporal causal attention mask (SCAM)
Our first proposed architecture utilizes transformers [49] which have been highly successful across a series of tasks in natural language processing [50, 51]. Transformers have a natural inductive bias towards time-series and sequence-based data and are thus highly suitable for extracellular waveforms. For this work, we designed a novel spatiotemporal causal attention mask (SCAM) to obey causality across time and channels. Specifically, we allow every recorded time step in a waveform to attend to time steps on other channels as long as those data points precede it in time. Full details and visualization of the transformer-based architecture, including implementation details are available in Supplementary Materials Section A.2.
3.3.2 Multi-layer perceptron (MLP)
While a transformer-based architecture provides a natural inductive bias towards time-series data appropriate for extracellular recordings, a downside to such an architecture is its high computational complexity and the requirement that it needs multiple GPUs to train. To demonstrate the generality of the CEED framework, we also propose a simpler MLP architecture that can be trained on a single GPU. The MLP encoder is a straightforward model that consists of three layers with sizes [768, 512, 256] and ReLU activations between them.
158   
159   
160   
161   
162   
163   
164   
165   
166   
167   
168   
169   
170   
171   
172   
173   
174   
175   
176   
177   
178   
179   
180   
181   
182   
183   
184   
185   
186   
187   
188   
189   
190   
191   
192   
193   
194   
195   
196   
197   
198   
199   
200   
201   
202   
203   
204   
205   
206
3.4 Objective function
The encoder is trained as follows. Let  $\tilde{\mathbf{z}}_i$  and  $\mathbf{z}_i$  be representations of the two views of the input waveform  $i$ , where the two views differ by the invariances listed in Sections 3.2.1 and 3.2.2. The optimization objective for a batch of  $B$  samples follows the contrastive loss utilized by SimCLR [28] and several other representation learning methods [52, 53];
$$
\mathcal {L} = \sum_ {i = 1} ^ {B} - \log \frac {\exp \left(\hat {\mathbf {z}} _ {i} \cdot \hat {\tilde {\mathbf {z}}} _ {i}\right) / \tau}{\sum_ {k \neq i} \exp \left(\hat {\mathbf {z}} _ {i} \cdot \hat {\mathbf {z}} _ {k} / \tau\right)} \tag {1}
$$
where  $\hat{\mathbf{z}}_i$  and  $\hat{\bar{\mathbf{z}}}_i$  are the L2 normalized representations and  $\tau$  is a temperature hyperparameter. Following [28], we also include a 2-layer MLP projector network after the encoder and the loss function operates on the outputs of this projector network (see Supplementary Materials Section A.3 for details and ablations of the projector network and its architecture).
To train and evaluate our model, we make use of two publicly available extracellular recordings published by the International Brain Laboratory (IBL): the DY016 and DY009 recordings [54]. These multi-region, Neuropixels 1.0 recordings are taken from a mouse performing a decision-making task (see Supplementary Materials Section C for more details).
4.1 Spike sorting
To evaluate how useful the features learned by CEED are for spike sorting, we constructed three datasets using units found by Kilosort 2.5, a full spike sorting pipeline manually tuned by IBL. The first dataset was extracted from the DY016 extracellular recording. It consisted of a 10 unit train and test dataset where all 10 units were classified as "good" by IBL's quality metrics [55]. We selected these units for their high waveform diversity (qualitatively) and because they had a relatively high amplitude, i.e., peak-to-peak (ptp). For this dataset, we constructed training sets of 200 or 1200 spikes per unit with a test set of 200 spikes per unit. For each spike, we extracted waveforms from 21 channels centered on the maximum amplitude channel. Although we extract 21 channels for our data augmentations (see Supplementary Materials Section A.1), we train and evaluate our model (and baselines) on either 5 or 11 channel subsets.