Towards robust and generalizable representations of extracellular data using contrastive learning
Authors: Ankit Vishnubhotla
Charlotte Loh
Akash Srivastava
Liam Paninski
Cole Lincoln Hurwitz
Venue: NeurIPS 2023 poster (2023)
Keywords: contrastive learning
representations
extracellular
high-density
spike sorting
cell-type classification
transformers
invariance
TL;DR: A contrastive learning spike featurization method for high-density extracellular recordings.

Abstract: 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.

Towards robust and generalizable representations of extracellular data using contrastive learning
Ankit Vishnubhotla"
Charlotte Loh" MIT Massachusetts cloh@mit edu
Liam Paninski
Columbia University New York av3016@columbia. edu
Columbia University
New York
liam@stat. Lammstatt.coluuia..co columbia.edi
Akash Srivastava
Cole Hurwitz
MIT-IBM
Columbia University
Massachusetts A Srivastava@iibm com
New York ch3676@columbia. edu
Abstract
Contrastive learning is quickly becoming essential tool in neuroscience for ex- tracting robust and meaningful representations neural activity. Despite numerous applications to neuronal population data, there has been little exploration of how these methods can adapted to key primary data analysis tasks such as spike sort- ing cell-type classification. In this work, we propose 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 method across multiple high-density extracellular recordings. All code used to run CEED can be found at htthttss: //ggthubb com mikitvissnu233 CEED
6
10
Introduction
13 High-density extracellular recordings now allow for simultancous recordings of large populations
14 of neurons across multiple brain regions with high temporal and spatial resolution. [1-4]. These
15 large-scale recordings are essential for gaining insights into key biological processes such as vision, and behavior which are distributed across brain regions (5). Along with gaining 17 insights into brain function, these technologies also promise improve the scalability and accuracy
16 decision-maker
18 of proiutcomputer interfaces which can restore motor function to paralyzed individuals [6].
major bottleneck for interpreting neural population activity is processing the extracellular signal 77. Although extracellular recordings contain a precise record of the coordinated neural activity, it must be extracted algorithmically through processing step called spike sorting. A crucial assumption in spike sorting is that each recorded neuron has 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 detected extracellular action potential (spike) back to its putative neuron (unit). Along with spike sorting, another important task in extracellular analysis cell-type classification. As neural circuits are diverse and heterogeneous, it becoming increasingly important to profile detected units using their morphoelectrical features 9 At a coarse level.
20
22
23
24
25
26
27
Equal contribution
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

28
possible to classify units as inhibitory or excitatory based on their extracellular profile [T]; however, may be possible to divide units into finer subgroups [TT.
29
30
For both spike sorting and cell-type classification, it is important to extract low-dimensional, mean- ingful 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
31
32
33
analysis (PCA) 13 14H. Although PCA is relatively effective and scalable, it suffers from a few key drawbacks including: (1) 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 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-peat amplitude, and the ratio of pre-hyperoolarization peak to the post-hyperpolarizatioo peak [3]. Again, while scalable and effective, these features are too simple and ad hoc to fully capture morpho-eletrical differences [13). Recently, a non-linear approach to cell-type classification, WaveMap, was introduced, utilizing UMAP PX and Louvain community
34
35
36
37
38
39
40
41
43
44
45
detection [24] to automatically find cell-type clusters [3].
48
In this work, we introduce 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 morphoelectrie cell-type classification than current specialized feature extraction methods. For example, in extracellular recordings, there are 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 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 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 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:
47
48
49
50
51
53
54
55
57
58
59
61
82
63
64
1. We introduce a novel framework, CEED, for analyzing extracellular recordings based on invariance learning.
65
66
2. We implement a stochastic view generation module for both single-channel and multi- channel extracellular waveforms.
67
68
3. We demonstrate that CEED works well with multiple neural network architectures, including a novel transformer- architecture with spatiotemporal causal attention mask (SCAM).
70
4. We show that CEED outperforms specialized featurization methods for spike sorting and morphoelectric cell-type classification.
71
2
Background
73 Contrastive Learning. Contrastive representation learning (28 falls under a broad class of self- 74 supervised learning methods 31 whose goal is to learn robust and generalizable representations by encouraging invariances to prior- set of transformations or nuisance variables. While 76 contrastive learning has mainly been used to extract effective and transferable representations from
77 basee data [27 28), more recently, it has also become powerful tool in the sciences to learn
78 invariances for physical systems [3].
79 In computational neuroscience, representation learning has mostly been performed using generative
80 models (34. While this has led to many interesting insights into behavior (35] and decision-making
81
36), this paradigm is sensitive to nuisance variables and will not capture subtle changes in the
2

---

Recording
Contrastive Learning
View Generation Module
Representations
Amplitude Jitter
Two Views
WWW Hampeee
Transformer /MP Encoder
MAKE SIMILAR
Collision
Negative
MAKE DIFFERENT
examples
wwww
Ho sysstiin Vacte Aobb
Crop
Batch of samples
Temporal fitter
a M Apo Mhym
Tasks
Waveforms
Embeddings
Nosse
11 Pappoieee proiie /4W
e.g. spike
2
sorting.
cell-type
Extracted Waveform
dessiffication
Unit
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 pre-defined set of invariances (see Section 3.2.I. Using these views, neural-networf based encoder, which can take the form of a multi-layer perceptron (MLP) or 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 series of downstream neuroscience tasks, such as spike sorting or morphoelectrical cell-type classification.
82 observation space that do not contribute much the explained variance. While recent research
83 into label-guided generative models 37-39] has somewhat alleviated these problems, it is still unclear if good generative performance important for representation learning. Most recently, number of contrastive learning methods have been introduced for learning robust and generalizable representations of neural population data 420 To our knowledge, however, contrastive learning has not been applied to key primary data analysis tasks such as spike sorting or cell-type classification.
84
85
86
87
88
Spike sorting. In spike sorting, the goal is extract a precise record of which neurons spiked 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 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 form of SVD/PCA 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 non-linear features of the observed data, and prefers features that explain variance rather than those that discriminate different waveforms. More recently, has been shown that non-linear autoencoders can have higher performance than PCA despite suffering from similar drawbacks 44
89
90
91
92
93
94
95
96
97
96
99
100
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
101
102
103
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 [I0] 45 More
104
105
recently, a non-linear method based on UMAP was introduced for cell-type classification 13 In this
106
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 446/48).
108
3

---

110
3
Contrastive Embeddings for Extracellular Data (CEED)
111
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 D (1) stochastic view generation module that augments waveforms with both common and task-specific nuisance variables. (2) A neural encoder that can extract now-dimensional representations from both single-channe 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.
112
113
114
115
network-
116
117
118
119
120
121
3.1
Notation
To introduce our method, we must first define some notation for extracellular recordings. Let P' = {pm m E 3 be the position of all M channels on the multi-electrode array (MEA), where each Pm E I is the location of channel m. Now define Pa to be the set of positions for subset of channels, M S M,, on the MEA.
123
125
Let S = {sn } be the set of N spiking events that are detected in an extracellular recording. Now, let the extracellular waveform of a spiking event 8, channel m be defined as W == (ru,m, n.m,,....n
126
(()) r.m E TRT", where T contiguous samples from the recording define a waveform and ra.m is the value of the waveform at sample This can be thought of a window of samples of size T on channel m which includes the spike 8,-
128
(8)
129
130
The amplitude of spike sn channel can be defined as An.m max, [W.. This is also known as the absolute maximum voltage recorded on channel m. The timestep at which the absolute maximum voltage occurs can vary for each channel. Therefore, for spike, center the waveform on each channel using the timestep at which the maximum An,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.
132
134
136
[ be the set waveforms for spike on a subset of channels and let A n, E Let W n.M E
137
138
R
0
[M]
be the set amplitudes for a spike on subset of channels, where [M] is the number of channels.
3.2
Stochastic view generation module
140
3.2.1
Common invariances
Let us define the representation of the waveform W Ai for spike 8,, as 2n DRD. There are 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,
143
144
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 W ", for spike s, such that
145
146
W n.A x (1 + e), then the representation ž,, for Wa.i should be such that ž, z,.
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 spatiotemporal matrix Gaussian MA((0, U, V), where U E TT' models the temporal correlation of a waveform and V E R models the spatial correlation between channels. For a spike sn, we want a representation such that if we sample n ~ MN U, V) and form W n. = W n. + jr where ja is the noise on the channels subset M, then W Wa,a should still yield a representation 2,, Z
148
151
152
153
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 the probe leading to observed waveforms that are distorted. To
155
this end, consider two spikes Sn1 and Sn2 with waveforms W M and W My such that M4,

---

158
and M2 share some number of channels. Now let us shift Snz by k time steps (-T/2 < k <0
or 0 < T/2) such that for any channel m M2, Wnz.m freg.m Jng.. 3... Png.m }. We refer to this temporally shifted waveform as W ng. A new waveform that is a linear sum
(1+b) (T+R))
159
160
of these two waveforms W n,,Y, = W,,,, AZ, + Walgi Ns should yield a representation 2, = 2nn Channel subsets- We want representations that are invariant to changes in the subset of channels W..'.. So for a new subset of channels M', as long as the channel with the highest amplitude is still contained in M', then the waveforms W W. should still yield representation 2n Zn.
161
162
4.
163
M used to define W n.
166
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.
188
169
170
171
172
1.
Cell position - For cell-type classification, we want representations that are invariant to the channel positions PM at which a spike is detected. So for a waveform W WR.M, any uniform changes to the channel positions Pas (without changing the waveforms at each channel) should still yield representation ž., Zn--
174
175
176
2.
Cell amplitude- As amplitude is mainly a function of cell position and orientation, we want
177
representations that are invariant to any uniform changes in the amplitudes. To this end, for any
178
n.M x a, should still positive value a, a change in the amplitudes for the spike 8n, W n. W
179
yield a representation 2,, Zy.
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 that we discard multi-channel information that may be useful for classifying different cell-types [12] (see Section for more detailed discussion). The full view generation pipeline is detailed in Supplementary Materials Section A.I
181
183
185
186
3.3
Encoder Architecture
187
For the encoder of CEED, we explored two neural network architectures - a transformer- network with a novel Spatiotemporal Causal Attention Mask (SCAM) and simple multi-layer perceptron (MLP) which is a more computationally efficient alternative (i.e., can fit on a single GPU).
188
189
190
For runtime comparison of these two architectures, see Supplementary Materials Section
191
3.3.1
Transformer-based encoder with spatiotemporal causal attention mask (SCAM)
192
Our first proposed architecture utilizes transformers (49) which have been highly successful across a series of tasks in natural language processing (50] [1). Transformers have natural inductive bias towards time-series and sequence-bree basen data and are thus highly suitable for extracellular waveforms.
193
194
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 other channels long 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
196
198
200
3.3.2
Multi-layer perceptron (MLP)
While a transformer-baser architecture provides a natural inductive bias towards time-series data 201
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
202
203
205
single GPU. The MI.P encoder a straightforward model that consists of three layers with sizes 7788, 512, 256] and ReLU activations between them.
5

---

207
3,4
Objective function
208
The encoder is trained as follows. Let ž, and z; be representations of the two views of the input
209
waveform i, where the two views differ by the invariances listed in Sections 3.2.1 and 3.2.2 The
210
optimization objective for batch of B samples follows the contrastive loss utilized by SimCLR
28
and several other representation learning methods 52 53];
exp(2; 2, /.//T Ek exp(2; 24/T)
C
log
(I)
where 2; and é; are the L2 normalized representations and T is a temperature hyperparameter. Following [28] we also include 2-layer MLP projector network after the encoder and the loss
212
213
214
function operates on the outputs of this projector network (see Supplementary Materials Section
215
A.3 for details and ablations of the projector network and its architecture).
Datasets
216
To train and evaluate model, 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 taken from mouse performing ppcision-making task (see Supplementary Materials Section C for more details).
218
220
4.1
Spike sorting
To evaluate how useful the features learned by CEED for spike sorting, we constructed three datasets using units found by Kilosort 2.5, 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-peal, (ptp). For this dataset, we constructed training sets of 200 or 1200 spikes per unit with 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.I) we train and evaluate our model (and baselines) on either 5 or 11 channel subsets.
223
224
225
226
227
228
230
232
The second dataset was extracted from both the DY016 and DY009 extracellular recordings. This dataset consisted of a 400 training units and 10 test units (the same 10 units evaluated in the first dataset). We extracted 200 units from each recording
233
234
create the training set with 200 1200 training spikes per unit. The goal of this dataset was to test how well CEED could generalize across large of units with varying quality. To this end, we also evaluated the performance of CEED on 100 random test sets of varying sizes (see Supplementary Materials Section D)
235
236
237
238
The third and final dataset also extracted from both the DY016 and DY009 extracellular record- ings. Unlike the first and second dataset, where the training set contained spikes from the test units (in-distribution),
239
240
purposefully excluded all units in the test set from the training To this end, this dataset consisted of 390 training units and 10 test units (the same 10 units evaluated in the above datasets). For all units in the training set, we utilized either 200 or 1200 training spikes per unit. The goal of this dataset was to test how well CEED could generalize to out-of out-of-distribttion (OOD) units.
242
243
4.2
Cell-type
245
For our cell-type classification dataset, we utilized the DY016 extracellular recording. We extracted the same 10 IBI. "good" units used for the spike sorting datasets. To remove positional information from each spike, we only extract waveforms from the maximum amplitude channel and we normalize
247
each waveform as described in Section 3,2.2

---

5
Experiments
249
250
5.1
Spike sorting
251
Re-sorting units extracted using KS2.5 In this experiment, we aim to demonstrate that CEED
252
extracts more useful features for spike sorting than both PCA and non-linear autoencoder. We evaluate these methods on both an in-distribution (ID) and out-of-distributt
253
(OOD) waveform discrimination task. Specifically, we train CEED, PCA, and autoencoder on all three spike sorting datasets (see Section 5 and perform inference on the spikes from the test units. Then, we 're-sort' these test spikes back their putative units by performing clustering on the resulting embeddings. For clustering, we use a parametric clustering algorithm. the Gaussian Mixture Model (GMM), and a non-parametrie clustering algorithm, HDBSCAN (56). To evaluate how well each clustering corresponds to the ground-truth we use the Adjusted Rand Index (ARI) [57]. One can also compute the accuracy after optimal permutation (e.g.. Hungarian algorithm), however, this is will give similar results to the ARI. To strengthen our spike sorting baselines, we also compare CEED to PCA and an autoencoder trained and tested on denoised waveforms using the YASS waveform denoiser [2 58
254
255
256
257
258
259
260
261
262
The YASS denoiser was trained to denoise single channel waveforms
we apply it independently to each channel and it improves clustering of PCA across the board. For all baselines, we sweep across (3,5,7.9) principal components and 3-11 channel subset sizes.
264
266
5.2
Cell-type classification
Cell-type classification of IBL recordings In this analysis, we aim demonstrate that CEED can extract putative cell-type clusters from IBL extracellular recordings. For the IBL recordings, we extract average waveforms (templates) for 163 good units pooled over DY016 and DY009. We then run inference on the extracted templates using CEED and perform a GMM clustering on the resulting embeddings find putative cell-types. As there is ground-truth for the IBL recordings, we validate our results by visualizing the templates for each cell-type cluster.
268
270
271
272
Zero-shot cell-type classification. In this experiment, compare the inferred representations of CEED to a state-of-the-ar cell-type classification method, WaveMap, on an outtof-distribution (OOD) dataset. Not only is this dataset not seen during training, but both the animal and the probe geometry are completely different than the dataset used to train CEED. The dataset consists of 625 templates (average waveforms) extracted from units which recorded while a monkey performs
274
276
278
discrimination task [33. The probe used to record these waveforms is a Plexon U-probe. The goal of this experiment is to find putative cell-types that explain the waveform variability seen in the data.
280
We propose to run CEED on the 625 templates and then perform a GMM clustering of the resulting embeddings. We choose the number of clusters by again utilizing the Elbow method and BIC. Given
282
how OOD this dataset is compared to the recordings used to train CEED, success on this task would
283
demonstrate the robustness and generalizability of CEED,
Results
284
285
6.1
Spike Sorting
286
Re-sorting units extracted using KS2.5 For our re-sorting task, we find that CEED outperforms both PCA and the non-linear autoencoder using raw waveforms or denoised waveforms across all three datasets introduced in Section 4 The results for this analysis can be found in Table !] As can be seen, in both the ID and OOD regimes, CEED has much higher performance than both baseline models. For all methods we utilize the same number of latent dimensions for the analysis (5D). We also show in Supplementary Materials Section that the strong performance of CEED generalizes to many different sets of test units.
288
290
292
293
Along with this results table, we also quantitatively and qualitatively compare the performance of CEED to PCA on the 10 neuron train and test dataset in Figure Visually, it can be seen that CEED's
294
features are far more informative about each unit's 'itss identity than the representations found by PCA.
296
The performance of CEED also much higher when being clustered by either HDBSCAN or
GMM even when when more principal components are afforded for the analysis.
7

---

Method
10 neuron train set
400 neuron train set
390 neuron train set
10 newron ID (ARI)
10 neuron ID (ARI)
10 neuron 0OD (ARI)
CEED (1200 spikes, 11 channels)
.89 .04
.79 +: .09
.78 - .05
CEED (1200 spikes, 5 channels)
.8 .03
.76 + .07
.77 = .08
Denoised PCA (1200 spikes, 11 channels)
.39 f .05
.45 1: .04
.46 i 00
Denoised PCA (1200 spikes. 5 channels)
.46 007
.48 1: .04
.49 1 04
Autnencoder (1200 spikes, 11 channels)
.47 + .06
.35 + .03
.28 + .02
Autnencoder (1200 spikes, 5 channels)
.43 + .06
.37 + .03
.33 + .01
Table I: Benchmarking CEED, PCA, and an autoencoder on in-distribution (ID) and out-of- distribution (OOD) data. For evaluation, fit 50 GMMs the embeddings and compute the mean and std. of the adjusted rand index (ARI). First column: we train and test each method with spikes from 10 neurons. Second column: we train each method with spikes from 400 neurons and then test on the original 10 neurons which are included in the training set (ID). Third column: we train each method on spikes from 390 neurons and test on the original 10 neurons which are not included in the training set (OOD). This experiment demonstrates that CEED peforms well on OOD data and can outperform non-linear autoencoder. All CEED results are generated using the MLP encoder.
298
6.2
Cell-type Classification
Cell-type classification of IBL recordings The results for cell-type classification of the pooled IBL recordings are shown in Supplementary Materials Section B We choose the number of cell-type clusters by sweeping over 1-10 clusters and choosing the minimum BIC. We find that 4 clusters explain most of the waveform variability. Interestingly, we find good separation in CEED's ' embedding space between narrow-spiking and broad-spiking units which indicates that we may be able to discriminate between inhibitory and excitatory subtypes.
300
302
304
The results for cell-type classification of OOD single unit data are shown in Figure 3 Despite training on IBL. extracellular datasets from mouse brain recorded with Neuropixels 1.0, CEED is able to generalize to unseen data from completely different animal (monkey) and probe (Plexon U-probe). In 3A, we visualize the inferred representations from
Zero-shot cell-type classification.
306
307
308
CEED using 2D UMAP and by coloring the points with the output of GMM that is trained on the 5D contrastive representations. We choose the number of cell-type clusters by sweeping over 1-10 clusters and then using the Elbow Method on the BIC curve. With this method, we discover 6 putative cell-types (similarly to WaveMap). Upon visual inspection, our cell-types are more well- isolated from each other than those of WaveMap (Figure To quantitatively assess which cell-type classification method better reflects the 'real' differences in extracellular waveforms, we utilize a supervised classifier which is trained predict cell-type labels using input waveforms (introduced
310
312
314
315
316
in 13). In Figure 3] and Figure 3), we show the results of this data on 5 cross-validation folds. Despite never training on this dataset, CEED has much higher accuracy (93.4%) compared that
318
of MaveMap (88.8%) 8%). A small caveat is that we ran WaveMap using publically available code but were unable to precisely reproduce the original result in the paper which finds 8 cell-type clusters.
319
320
Despite this difference, the accuracy value reported in the WaveMap paper for the 8 cell-type clusters
321
(91% %) was still lower than that of CEED.
322
7
Discussion
323
In this paper, we introduced a novel representation learning method, CEED, for extracellular record- ings. Our main hypothesis was that by finding representations of extracellular waveforms that are robust both common and task-speefic nuisance variables, we can outperform specialized feature extraction approaches on two key tasks: spike sorting and morphoelectrical cell-type classification.
324
325
326
327
We validate CEED on multiple high-density extracellular datasets. For spike sorting, we show that CEED extracts features that far outperforms those of PCA and a non-linear autoencoder on waveform discrimination task. For cell-type classification, we show that CEED able to extract discriminative features of extracellular waveforms that allow for finding morphological subgroups in an unsupervised manner. Surprisingly, we find that CEED even outperforms a recent nin-linear
328
329
331
cell-type classification method WaveMap, on an animal and probe geometry unseen during training.
8

---

Gaussian Mixture Model (10 components)
HDBSCAN (10 min cluster size, 10 min samples)
(A)
Adjusted Rand
70
80
90
R PCA
CLLD
PCA
Adjusted fand Index
150
80
R
CEED
PCA
Denaised
60
40
50
20-
40
of PCs
PCs
(B)
CEED embedcing, GMM Fred Labels
Denoised Pred
Thue
Deno PCA. True Labes
ARI: 62.80
Compoed PC
C360
06-
-13
Donotvod PC
Denikeod PC
(C)
CEED ombedding. HD8SCAN Pred
Deneised PCA, HDBSCAN Fred Labels, 37.43
CEED embedding, Thue Labels
Labes. ARI: 90.83
Labels
Derosed
00 M
Figure 2: CEED outperforms PCA on spike sorting featurization. Here, we show results for CEED (using the SCAM architecture) and PCA when trained and evaluated on a 10 neuron dataset. (A) Clustering results for the two featurization methods using a parametric (GMM) and non- conpaarmettric (HDBSCAN) clustering algorithm. As can be seen, the featurization learned by CEED is much more discriminative than that of PCA for both clustering methods even when using supervised denoiser reduce noise the data. (B) Visualized results of a 10 component GMM clustering on the learned embeddings from CEED (left) and denoised PCA (right). (C) Visualized results of the HDBSCAN clustering applied the learned embeddings from CEED (left) and denoised PCA (right).
333
Despite the significant performance improvements of CEED, there are few limitations that must be addressed before it can become plug-and-play method for extracellular analysis. Firstly, the training sets we utilize in this work are still quite small and lack neuron diversity. For CEED to be generalizable to multiple recordings and animals, more diverse datasets must be constructed. Secondly, the best results of CEED are obtained when using the transformer-basee encoder which requires multiple GPUs and is currently quite slow to train. Recent progress in acceleration software (59] offer promising solutions to speed up computation and incorporating these methods into CEED could be a future direction. Thirdly, all spike sorting results in the paper are from re re-sorting already sorted datasets; CEED must be incorporated into a full spike sorting pipeline in order to be used by many different research groups. Finally, our cell-type results do not include any functional classification which could help validate the clusters found by CEED.
334
335
337
339
341
343
8 Broader Impact
344
345
Although CEED has the potential to improve key tasks in extracellular analysis, a drawback of our approach compared to simple approaches like PCA is the additional computational resources requires train and run. Contrastive learning requires large batch sizes to achieve of the-art performance which means, if training the SCAM transformer model, we often have to run experiments on large-scale, multi-gpu clusters. Moreover, transformer-based architectures have hundreds of thousands of parameters, which further increases CEED's compute requirements. As highly parameterized deep neural networks produce large amounts of carbon emissions (60].
346
347
348
349
351
CEED could have possible negative impact the environment. We also propose MLP-based

---

(A)
CEED 93.4% accuracy (zero-hhot learnin)
(B)
6
12
UMAP 2
Clasters
1.0
100
P (Predicted True)
89
Thue Casss
96
0.6
0.4
-2
0.0
2.5
5.0
7.5
10.0
12.5
15.0
17.5
UMAP
1
6
(C)
WaveMap 88.8% accuracy
Predicted Class
(D)
1.0
81
31
7-
0.8
P (Predicted True))
94
True Class
UMAP2
0.4
4-
0.2
2-
spoxt 6
Predicted Class
-2
0
2
6
8
10
12
UMAP
Figure 3: CEED outperforms WaveMap on morphoelectrical cell-type classification with zero- shot learning. (A) A 2D UMAP embedding of CEED's inferred representations for OOD dataset which consists of 625 single units extracted from a monkey recording. We cluster the representations using GMM and plot each cluster with distinct color. For each cluster, also plot the waveforms associated with the discovered cell-type. It can be seen that narrow-spiking waveforms (3,5,6) are well-separated from broad-spiking waveforms (1.2.4) which may indicate good separation of inhibitory and excitatory subtypes, respectively. Surprisingly, CEED is trained on different animal and probe geometry, but can still generalize this dataset, outperforming WaveMap
classification baseline introduced in [13]. (B) On the top, we demonstrate how we chose the number of clusters for the GMM, i.e., with the Elbow Method and BIC. On the bottom, we show the confusion matrix of gradient boosted decision tree classifier trained to map raw waveforms to the cell-type clusters extracted by CEED. The accuracy of each method is defined
the average of the diagonal. (C) 2D UMAP embedding the single unit dataset using WaveMap. The clusters are colored according to WaveMap's outputted labels. Narrow-spiking clusters (1,2,4) are more mixed with broad-spiking clusters (3,5.6). (D) The confusion matrix of gradient boosted decision tree classifier trained to map raw waveforms to the cell-type clusters extracted by WaveMap.
architecture that performs comparably transformer while only using single GPU, which can 354
lower CEED's 's impact on the environment.
355 References
356
[]] James Jun, Nicholas A Steinmetz, Joshua Siegle, Daniel Denman, Marius Bauza, Brian Barbarits, Albert K Lee, Costas A Anastassiou, Alexandru Andrei, Çağatay Aydın, al. Fully integrated silicon probes for high-density recording of neural activity. Nature, 551(7679):232- 236. 2017.
357
358
359
360
[2] Ashley I. Juavinett, George Bekheet, and Anne K Churchland. Chronically implanted neuropix- els probes enable high-yield recordings in freely moving mice. Elife, 8:c47188, 2019.
10

---

382
[3] Nicholas A Steinmetz, Cagatay Aydin, Anna Lebedeva, Michael Okun, Marius Pachitariu, Marius Bauza, Maxime Beau, Jai Bhagat, Claudia Böhm, Martijn Broux, et al. Neuropix- els 2.0: A miniaturized high-density probe for stable, long-term brain recordings. Science, 372(6539)) 2021.
363
364
365
388
[4]
International Brain Laboratory, Kush Banga, Julius Benson, Niccolò Bonacchi. Sebastian A Bruijns, Rob Camphell, Gaëlle A Chapuis, Anne K Churchland, M Felicia Davatolhagh, Hyun Dong Lee, et al. Reproducibility of in- electrophysiological measurements in mice. bioRviv, pages 2022-055 05, 2022.
367
368
369
370
[5] Nicholas A Steinmetz. Peter Zatka-Haas, Haas, Matteo Carandini, and Kenneth D Harris. Distributed coding of choice, action and engagement across the mouse brain. Nature, 576(7786):2662733 2019.
371
372
373
[6] Elon Musk et al. An integrated brain-machine interface platform with thousands of channels. Journal of medical Internet research, 21(10):e1619, 2019.
374
375
[7]
Matthias H Hennig, Cole Hurwitz, and Martino Sorbaro. Scaling spike detection and sorting for next-generation electrophysiology. In Vitro Neuronal Networks: From Culturing Methods to Neuro-Technologe
378
377
Applications, pages 171-184, 2019.
378
(8]
Carl Gold, Darrell A Henze, Christof Koch, and Gyorgy Buzsaki. On the origin of the extracellu- lar action potential waveform: a modeling study. Journal of neurophysiology. 95(5):3113-3128, 2006.
379
380
[9] Ganesh Vigneswaran, Alexander Kraskov, and Roger N Lemon. Large identified pyramidal cells in macaque motor and premotor cortex exhibit "thin spikes": implications for cell type classification. Journal of Neuroscience, 31440))44235114222 2011.
382
383
384
[10] Peter Barthó, Hajime Hirase, Lenaic Monconduit, Michael Zugaro, Kenneth D Harris, and Gyorgy Buzsaki. Characterization of neocortical principal cells and interneurons by network interactions and extracellular features. Journal of neurophysiology. 92(1):600-608 2004.
386
[11] Alessio P Buccino, Michael Kordovan, Torbjørn V Ness, Benjamin Merkt, Philipp D Häfliger, Marianne Fyhn, Gert Cauwenberghs, Stefan Rotter, and Gaute T Einevoll. Combining biophysi- cal modeling and deep learning for multielectrode array neuron localization and classification. Journal of neurophysiology. 121(3):122 2018.
388
380
[12] Xiaoxuan Jia, Joshua H Siegle, Corbett Bennett, Samuel D Gale, Daniel Denman, Christof Koch, and Shawn R Olsen. High-density extracellular probes reveal dendritic backpropagation and facilitate neuron classification. Journal of neurophysiology. (1(5)):831-1847) 2019.
392
394
[13] Eric Kenji Lee, Hymavathy Balasubramanian, Alexandra Tsolias, Stephanie Udochukwu Anakwe, Maria Medalla, Krishna V Shenoy, and Chandramouli Chandrasekaran. Non-linea dimensionality reduction on extracellular waveforms reveals cell type diversity in premotor cortex. Elife. 10:e67490, 2021.
396
398
[14]
Samuel Garcia and Christopher Pouzat. Tridesclous, 2015.
https://github.co com/
tri triescloou//trresccloous
400
[15] Cyrille Rossant, Shabnam N Kadir, Dan FM Goodman, John Schulman, Maximilian LD Hunter, Aman B Saleem, Andres Grosmark, Mariano Belluscio, George H Denfield, Alexander S Ecker, et al. Spike sorting for large, dense electrode arrays. Nature neuroscience, 19(4):63-6411 2016.
402
404
[16] Marius Pachitariu, Nicholas A Steinmetz, Shabnam N Kadir, Matteo Carandini, and Kenneth Harris. Fast and accurate spike sorting of high-channel count probes with kilosort. Advances in neural information processing systems, 29, 2016.
406
407
[17] Jason E Chung, Jeremy F Magland, Alex H Barnett, Vanessa M Tolosa, Angela C Tooker, Kye Y Lee, Kedar G Shah, Sarah Felix, Loren M Frank, and Leslie Greengard. fully automated approach to spike sorting. Neuron, 1381-194, 2017.
408
11

---

410
[18] Gerrit Hilgen, Martino Sorbaro, Sahar Pirmoradian, Jens-Oliver Muthmann, Ibolya Edit Kepiro, Simona Ullo, Cesar Juarez Ramirez, Albert Puente Encinas, Alessandro Maccione, Luca Berdondini, et al. Unsupervised spike sorting for large-scale, high-density multielectrode arrays. Cell reports, 18(10):2521-2555 2017.
411
413
414
[19] Pierre Yger, Giulia LB Spampinato, Elric Esposito, Baptiste Lefebvre, Stéphane Deny, Christophe Gardella, Marcel Stimberg, Florian Jetter, Guenther Zeck, Serge Picaud, et al. spike sorting toolbox for up to thousands of electrodes validated with ground truth recordings in vitro and in vivo. Elife, 7:e34518, 2018.
416
[20] Roland Diggelmann, Michele Fiscella, Andreas Hierlemann, and Felix Franke. Automatic spike sorting for high-density microelectrode arrays. Journal of neurophysiology, 1556))315 2018.
419
[21] JinHyung Lee, Catalin Mitelut, Hooshmand Shokri, Ian Kinsella, Nishchal Dethe, Shenghao Wu, Kevin Li, Eduardo Blancas Reyes, Denis Turcu, Eleanor Batty, et al. Yass: Yet another spike sorter applied to large-scale multi-electrode array recordings in primate retina. BioRxiv, pages 2020- 03,03333222 2020.
422
424
425
[22] Keming Chen, Yangtao Jiang, Zhanxiong Wu, Nenggan Zheng, Haochuan Wang, and Hui Hong. Htsort: Enabling fast and accurate spike sorting multi-electrode arrays. Frontiers in Computational Neuroscience, 15:657151, 2021.
427
428
[23]
Leland McInnes, John Healy, and James Melville. Umap: Uniform manifold approximation and projection for dimension reduction. arXiv preprint UrXiv:/802.03422 2018.
[24] Vincent D Blondel, Loan-Loup Guillaume, Renaud Lambiotte, and Etienne Lefebvre. Fast unfolding of communities in large networks. Journal of statistical mechanics: theory and experiment,
431
432
2008
[25] Cole Hurwitz, Kai Xu, Akash Srivastava, Alessio Buccino, and Matthias Hennig. Scalable spike source localization in extracellular recordings using amortized variational inference. Advances in Neural Information Processing Systems, 32, 2019.
434
435
436
[26] Julien Boussard, Erdem Varol, Hyun Dong Lee, Nishchal Dethe, and Liam Paninski. Three- dimensional spike localization and improved motion correction for neuropixels recordings. Advances in Neural Information Processing Systems, 34:22095-22105, 2021.
437
438
439
[27] Rumen Dangovski, Li Jing, Charlotte Loh, Seungwook Han, Akash Srivastava, Brian Che- ung, Pulkit Agrawal, and Marin Soljačić. Equivariant contrastive learning, arXiv preprint arXiv:2111.0099, 2021.
440
441
442
[28]
Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In International conference on machine learning, pages 1597-1607. PMLR, 2020.
443
444
445
[29] Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick. Momentum contrast for unsupervised visual representation learning, 2020.
448
447
[30] Jean-Bastien Grill, Florian Strub, Florent Altché, Corentin Tallec, Pierre H. Richemond, Elena Buchatskaya, Carl Doersch, Bernardo Avila Pires, Zhaohan Daniel Guo, Mohammad Gheshlaghi Azar, Bilal Piot, Koray Kavukcuoglu, Rémi Munos, and Michal Valko. Bootstrap your own latent: new approach to self-supervised learning, 2020.
448
449
450
451
[31] Jure Zbontar, Li Jing, Ishan Misra, Yann LeCun, and Stéphane Deny. Barlow twins: Self- supervised learning via redundancy reduction, 2021.
452
453
[32] Yuyang Wang, Jianren Wang, Zhonglin Cao, and Amir Barati Farimani. Molecular contrastive learning of representations via graph neural networks. Nature Machine Intelligence, 4(3):279- 287, mar 2022.
454
455
[33] Charlotte Loh, Thomas Christensen, Rumen Dangovski, Samuel Kim, and Marin Soljažić. Surrogate- and invariance- contrastive learning for data-searce applications science. Nature Communications, 13(1), jul 2022.
457
12

---

459
[34] Cole Hurwitz, Nina Kudryashova, Arno Onken, and Matthias H Hennig. Building population models for large-scale neural recordings: Opportunities and pitfalls. Current opinion in neurobiology, 70:64-73, 2021.
480
462
[35] Juan A Gallego, Matthew G Perich, Raeed H Chowdhury, Sara Solla, and Lee E Miller. I Long-term stability of cortical population dynamics underlying consistent behavior. Nature neuroscience, 23(():260-2)) 270, 2020.
464
485
[36] Zoe C Ashwood, Nicholas A Roy, Iris R Stone, International Brain Laboratory, Anne E Urai, Anne K Churchland, Alexandre Pouget, and Jonathan W Pillow. Mice alternate between discrete strategies during perceptual decision-making Nature Neuroscience, 25(2):201 2022.
466
467
468
[37] Ding Zhou and Xue-Xin Wei. Learning identifiable and interpretable latent models of high- dimensional neural activity using pi-vae. Advances in Neural Information Processing Systems, 73472 2020.
469
470
[38] Omid G Sani, Hamidreza Abhaspourazad, Yan T Wong, Bijan Pesaran, and Maryam M Shanechi. Modeling behaviorally relevant neural dynamics enabled by preferential subspace identification. Nature Neuroscience, 24(1): (11)):140-1999 149, 2021.
472
473
474
[39] Cole Hurwitz, Akash Srivastava, Kai Xu, Justin Jude, Matthew Perich, Lee Miller, and Matthias Hennig. Targeted neural dynamical modeling. Advances in Neural Information Processing Systems, 31:2:3499-9995995 2021.
475
476
477
[40]
Ran Liu, Mehdi Azabou, Max Dabagia, Chi-Heng Lin, Mohammad Gheshlaghi Azar, Keith Hengen, Michal Valko, and Eva Dyer. Drop, swap, and generate: A self-supervised approach for generating neural activity. Advances in neural information processing systems, 34: 34:10587-10599 2021.
478
480
481
[41] Mehdi Azabou, Mohammad Gheshlaghi Azar, Ran Liu, Chi-Heng Lin, Erik C Johnson, Kiran Bhaskaran-Nair, Max Dabagia, Bernardo Avila-Pires, Lindsey Kitchell, Keith B Hengen, et al. Mine your own view: Self-supervised learning through across-sample prediction. arXiv preprint arXiv:2102. 10106, 2021.
482
484
[42] Steffen Schneider, Jin Hwa Lee, and Mackenzie Weygandt Mathis. Learnable latent embeddings for joint behavioural and neural analysis. Nature, pages 1-9, 2023.
486
[43] Glen D Brown, Satoshi Yamada, and Terrence J Sejnowski. Independent component analysis at the neural cocktail party. Trends in neurosciences, 24(1):54 2001.
486
[44] Eugen-Richard Ardelean, Andreea Coporiie, Ana-Maria Ichim, Mihaela Dinşoreanu, and Raul Cristian Muresan. A study autoencoders as a feature extraction technique for spike sorting. Plos one, 18(3)x00222100 2023.
490
492
[45] David A McCormick, Barry W Connors, James W Lighthall, and David A Prince Comparative electrophysiology of pyramidal and sparsely spiny stellate neurons of the neocortex. Journal of neurophysiology. 54(4):782-806, 1985.
494
[46] Thomas Euler, Silke Haverkamp, Timm Schubert, and Tom Baden. Retinal bipolar cells: elementary building blocks of vision. Nature Reviews Neuroscience, 15(8)5007-519 2014.
496
[47] Hongkui Zeng and Joshua R Sanes. Neuronal cell-type classification: challenges, opportunities and the path forward. Nature Reviews Neuroscience, 180))530-544, 2017.
438
499
[48] Nathan W Gouwens, Staci A Sorensen, Fahimeh Baftizadeh, Agata Budzillo, Brian R Lee, Tim Jarsky, Lauren Alfiler, Katherine Baker, Eliza Barkan, Kyla Berry, et al. Integrated morphoelectric and transcriptomic classification of cortical gabaergic cells. Cell, 183(4)9935- 953, 2020.
500
502
[49] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, tukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural information processing systems, 30, 2017.
504
13

---

506
[50] Jacob Devlin, Ming-
Ming-Wei Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding, 2019.
507
508
[5
J
Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert- Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners, 2020.
510
512
[52] Kihyuk Sohn. Improved deep metric learning with multi-class n-pair loss objective. In Proceedings of the 30th International Conference on Neural Information Processing Systems, NIPS'16, 16, page 1857-1865, Red Hook, NY, USA, 2016. Curran Associates Inc.
515
517
[53] Aäron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding. CoRR, abs/ 007708378 2018.
518
[54] The International Brain Laboratory. Valeria
Dora Angelaki, Hannah Bayer, Niccolo Bonacchi, Matteo Carandini, Fanny Cazettes, Gaelle Chapuis, Anne K Churchland, Yang Dan, et al. Standardized and reproducible measurement of decision-making in mice. Elife,
520
622
10, 2021.
523
[55]
Mayo Faulkner Chapuis, Kenneth D Harris, Julia M Huntenburg. Cole Hurwitz, Hyun Dong Lee, Liam Paninski, Cyrille Rossant, Noam Roth, Nicholas A Steinmetz, Charlie Windolf, et al. Spike sorting pipeline for the international brain laboratory. channels, 10:6, 2022.
525
528
[56] Leland Mclnnes, John Healy, and Steve Astels. hdbscan: Hicrarchical density based clustering.
527
J. Open Source Softe., 2(11):205, 2017.
528
[57]
Lawrence Hubert and Phipps Arabie. Comparing partitions. Journal of classification, 2:193-218. 1985.
[58]
Julien Boussard, Charlie Windolf, Cole Hurwitz, Hyun Dong Lee, Han Yu, Olivier Winter, and Liam Paninski. Dartsort: A modular drift tracking spike sorter for high-density multi-electrode probes. bioRxiv, pages 2023-08, 2023.
531
533
[59] Jie Ren, Samyam Rajbhandari, Reza Yazdani Aminabadi, Olatunji Ruwase, Shuangyan Yang, Minjia Zhang, Dong Li, and Yuxiong He. Zero-offload: Democratizing billion-scale model training, 2021.
534
535
536
[60] Payal Dhar. The carbon impact of artificial intelligence, Nat. Mach. Intell., 2(8):423 425,34422222222 2020.
14


Reviews:
[{"review_id": "5WsvgnwepW", "reviewer": "Reviewer_xe16", "summary": "The paper proposes an approach based on transformer and contrastive learning for tackling problems from extra cellular recordings, including spike sorting and cell type classification. Experiments show the validity over standard PCA and traditional approaches in the field.\n\nNote after rebuttal: I have appreciated the work on the revised manuscript, and raised the score accordingly. A concern on the rigor of the model selection is still present, though.", "questions": "- One of the main issues that I see in this work is also pointed out in the limitations section: the computational complexity of the approach. While the performance seems interesting, the point on the much higher computational complexity compared to the baselines risks to reduce the relevance of the analysis. Anyway, I suggest the authors to clearly indicate in the paper the number of trainable parameters, the computational complexity and an analysis of the computational cost required for training the proposed method, in comparison to the literature approaches.\n\n- As far as I could see, no model selection on the hyper-parameters is performed.\n", "limitations": "I think the authors correctly indicated a major potential limitation of the approach, namely its huge computational costs compared to literature alternatives.", "rating": 5, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The topic is fascinating\n- The paper is well written\n- The potential impact of the results are promising", "weaknesses": "- The experimental setup is not fully convincing", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for acknowledging that the paper is well-written, interesting, and that the results are promising. Since the reviewer has highlighted these strengths, we hope that by addressing their concerns, we can improve the rating.\n- The reviewer indicated that the experimental setup is not fully convincing. We agree with the reviewer that our spike sorting baselines and analyses could be improved. To address this, we include two new experiments to quantify (1) the performance of our method on out-of-distribution (OOD) data, (2) the performance of our method in comparison to a non-linear autoencoder. Please see Table 1 of the attached PDF and the Global Response section for a summary of these results. As morphoelectrical cell-type classification is still very much an open problem, we utilize the experimental setup and metric introduced in the Eric Kenji Lee 2021 paper. We feel this is a fair way of comparing our method to WaveMap. We do agree that quantifying our method and other cell-type classification methods with functional measures (cell-type selectivity, brain region cell-type distributions, etc.) is an exciting direction and would make for interesting future work.\n- The reviewer indicated that they did not see a section about model selection of the hyper-parameters. In supplement A.2, we discuss hyperparameters and how the learning rate was tuned using a heldout validation set (separate from the test set). We also did some hyperparameter tuning to choose the optimizer (‘adam’ vs. ‘sgd’) and representation size. We will include these details in the supplement of the final paper. We agree that more rigorous hyperparameter tuning could improve the performance of the model, however, we believe that it is a strength of the method that we already have strong results without too much hyperparameter tuning. \n- The reviewer expressed concerns about the computational complexity of our method. We note that these concerns were shared across other reviewers so we provide a detailed response in the Global Response section. To briefly address the point here, we agree that our originally proposed methods had significant computational complexity which limited their applicability to new datasets. In order to address this we: (1) sped up the data augmentation significantly by only computing the augmentations on a small neighborhood of channels, (2) proposed an alternative MLP architecture that also has strong performance and is much faster to train (only requires 1 GPU), (3) quantified CEED on OOD extracellular spikes to illustrate that even without training on new data, CEED’s features are more discriminable than the baselines’ features. While we are still looking into ways to improve the computational complexity of SCAM, this is challenging to address as transformers are still an active area of research. Recent progress in quantization [1] and acceleration software [2] offer promising solutions to transformers’ runtime issues, but incorporating these methods into CEED would be a future direction.\n\n[1] Liu, Zhenhua, et al. \"Post-training quantization for vision transformer.\" NeurIPS 2021\n\n[2] Ren, Jie, et al. \"ZeRO-Offload: Democratizing Billion-Scale model training.\" 2021 USENIX Annual Technical Conference. 2021."}, {"review_id": "i3A89vw9Fn", "reviewer": "Reviewer_CNrX", "summary": "This paper presents a novel method for self-supervised learning of useful representations for data from extracellular, multielectrode recordings in electrophysiology. The method uses a transformer architecture with causal spatiotemporal attention masks, and contrastive learning based on a set of desirable and relevant invariances. The paper describes the proposed neural architecture and training method, and tests them on spike sorting and cell-type classification (two standard problems for this type of data), on synthetic as well as real data from publicly available databases.\n", "questions": "Cold you give some more information about the practical computational and data requirements for the method? In particular, I could not find exact details on the hardware used for training (and the time required). It would also be great to have an additional analysis showing for instance how the results in figure 2 change as the amount of training data is changed.\n", "limitations": "The limitations of the method are discussed adequately. The paper mentions potential environmental issues related to the heavy computational requirements for model training, as well as possible mitigation strategies.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- the design choices for the model (including the details of the invariances used for CL) are well explained.\n\n- the proposed method performs well, beating denoised PCA in spike sorting and WaveMap in cell-type classification.\n\n- the paper is overall well written.\n\n- the limitations of the method are clearly signposted.\n", "weaknesses": "- as reported by the authors, the method is currently slow to train and it is untested for spike sorting on more diverse data with waveforms that could have very different shapes than those in the training data (which would be the most useful case for general usage).\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their strong assessment of our work and for their useful questions/feedback. We hope to address their concerns and questions in the following response.\n- The reviewer correctly pointed out that CEED is untested for spike sorting of more diverse data with waveforms that have different shapes than the training data. We agree with the reviewer that our experiments, which focused on CEED’s in-distribution (ID) performance, were insufficient for showing the general usability of the method. To address this, we include a new experiment to quantify the performance of our method on out-of-distribution (OOD) data. To strengthen our baselines, we also benchmark a non-linear baseline (an autoencoder) to show that the increase in performance comes from more than just the non-linearity in CEED. Please see Table 1 in the attached PDF and Global Response section for a summary of these results. We thank the reviewer for proposing this suggested experiment as we believe it strengthens the submission significantly. \n- The reviewer asked about the computational complexity and hardware used for training our method. For our spike sorting experiments, we utilized 16 NVIDIA V100s in parallel. Our runtime was 35 seconds per epoch for the 10 neuron, 200 spike, 11 channel model. For the 400 neuron, 200 spike, 11 channel model, the runtime was 3.1 minutes per epoch. \n- The reviewer mentioned that the computational complexity of our method is a weakness of the model. We note that this concern was shared across many of the reviewers so we provide a detailed response in the Global Response section. To briefly address the point here, we agree that our originally proposed model had significant computational complexity which limited its applicability to new datasets. In order to address this we, (1) sped up the data augmentation significantly by only computing the augmentations on a small neighborhood of channels, (2) proposed an alternative MLP architecture that also has strong performance and is much faster to train (only requires 1 V100 GPU), (3) quantified CEED on OOD datasets to show that even without training, CEED can outperform the current baselines for spike sorting or cell-type classification. While we are still looking into ways to improve the computational complexity of SCAM, this is challenging to address as transformers are still an active area of research. Recent progress in quantization [1] and acceleration software [2] offer promising solutions to transformers’ runtime issues, but incorporating these methods into CEED would be a future direction.\n- The reviewer mentioned that it would be useful to show how the results in figure 2 change as the amount of training data is changed. This is a great suggestion and we would be happy to include this analysis in the final manuscript. As can be seen in Tables 3 and 4 of our attached PDF, CEED still outperforms all baselines when using only 200 spikes from each neuron (rather than the 1200 spikes we use in figure 2).\n\n[1] Liu, Zhenhua, et al. \"Post-training quantization for vision transformer.\" NeurIPS 2021\n\n[2] Ren, Jie, et al. \"ZeRO-Offload: Democratizing Billion-Scale model training.\" 2021 USENIX Annual Technical Conference. 2021."}, {"review_id": "J1aFbRdTlg", "reviewer": "Reviewer_MJau", "summary": "Rebuttal Update: I thank the authors for answering my questions and for the changes and additional experiments that they conducted. I have raised my score accordingly.\n\nThis paper proposes a contrastive framework to do spike sorting and cell type identification.", "questions": "It seems like the training data already requires knowledge of the spikes. Does this mean that the model cannot be trained unsupervised from MEA recordings alone?\n\nCan you check if your learned features are actually becoming invariant to your data augmentations? I remember work showing that, in practice, this is often not the case.\n\nPlease give a more extensive discussion of refs 41-43. How do they differ from your approach? Could one just use CEBRA to do the same?\n\nIs there no next token prediction in your transformer? If not, why do you use causal masks?\n\nLine 95, ref 45 is form 2015 – please look at more recent reviews on Spike Sorting. I am skeptical that PCA is still ubiquitous, that seems like a straw man.\n\nLines 99-106: Please comment of functional, anatomical and genetic cell type classification approaches. Electrical profiles are only way.\n\n116 -are. Also, this only gives approx. invariance (see above)\n\n121 Please define MEA\n\n145-147 please clarify\n\n161 as -> a\n\n177 is -> are\n\nAre there no tokens in your transformer? I.e., no discrete symbols?\n\nWhy did you pick K=5? This seems low? Please comment on hyper parameter search.\n\n190 SimCLR did not invent this loss function. Please cite the original references.\n\n192 Did you perform an ablation on the projection MLP? What happens?\n\n205 Only 10 neurons? Does your approach scale to big/relevant modern datasets (i.e., KILOsort...)?\n\n227 If you have ground truth labels, why not compute accuracy (after optimal permutation with, e.g., Hungarian algorithm)?", "limitations": "I really want to know how classic spike sort algorithms (including the whole BSS pipeline and source recovery) compare to your approach as you are moving more out of distribution. My intuition is that learned (deep) approaches break down faster in those regimes.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper is well written. It is easy to follow and the motivation is clear. The experiments are sound. The data augmentations encode very useful inductive biases for spike sorting.", "weaknesses": "The paper lacks more extensive benchmarking. There is a whole Zoo of spike sorting methods and dedicated benchmark (e.g., https://www.nature.com/articles/s41592-020-0902-0) – limiting the comparison to PCA and one recent method for clustering is not satisfactory.\n\nClassical spike sorting builds on a lot of theory. Proposing a learned method instead is interesting, but it is important to check how this performs against classic methods more extensively out of distribution.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their very detailed review and for their useful feedback. We hope to address their concerns point-by-point in the following response.\n- The reviewer pointed out that the manuscript lacks extensive benchmarking especially with the large number of spike sorting methods currently available for MEAs. This is a valid concern and highlights a clarity issue with our paper. We want to emphasize that *CEED is not a spike sorting method, but rather a feature extraction method for spike sorting and cell-type classification*. While there are a number of spike sorting methods for MEAs, almost all of these methods utilize PCA for feature extraction (Klusta, HerdingSpikes2, Mountainsort4, SpykingCircus, Trideclous, and even Kilosort uses a form of SVD/PCA). We argue that any method that uses PCA would benefit from switching to our more discriminable and robust features. Since directly incorporating CEED into these methods is a non-trivial coding challenge (especially given the short rebuttal window), we plan to add a new experiment to the final manuscript where CEED is incorporated into a simple spike sorting pipeline which includes detection, featurization, clustering, and template matching. We will use SpikeInterface for this analysis.\n- The reviewer correctly pointed out that the method needs more extensive benchmarking for out-of-distribution spike sorting datasets. We completely agree with this criticism and have added an experiment to quantify the performance of our method on out-of-distribution (OOD) data. We show that CEED outperforms all baselines (including a new non-linear autoencoder) on spikes from neurons outside the training set. Please see the Table 1 in the attached PDF and the General Response for a summary of this analysis. \n- The reviewer asked if CEED requires knowledge of the spikes and cannot be trained without supervision. CEED does require spike times and channel positions for all the spikes in the training set. These are easy to obtain experimentally. However, CEED does not require spike sorted data and can be trained after a simple voltage-based thresholding detection step. We agree it would be an interesting future work to find useful embeddings of extracellular data without spike detection.\n- The reviewer asked us to check if the proposed data augmentations actually induce approximate invariance in CEED’s representations. This is an important point and something we visualize in supplementary figure 6. We further extend this analysis in the attached PDF by adding a visualization of embeddings from 3 different neurons with all data augmentations (please see Figure 1).\n- The reviewer asked us to provide more discussion of refs 41-43 and how they differ from our method. References 41-43 all applied their method to neural response data (processed spike trains, obtained after spike sorting) and not to raw detections from extracellular data (before spike sorting). To our knowledge, we are the first contrastive learning method for extracellular data and had to develop specific data augmentations for this data modality. A large difference with CEBRA is that CEBRA explicitly mentions that they do not utilize data augmentations which could make the learned representations sensitive to the nuisance variables that are prevalent when working with MEAs (e.g. collisions). We would be happy to add this explanation to the final manuscript.\n- The reviewer asked about SCAM’s causal mask. As each next token in SCAM is conditioned on all the previous tokens and the last token is then used as the representation, we thought a causal mask would be appropriate.\n- The reviewer asked if there are no tokens in our transformer, i.e., no discrete symbols. For CEED, the signal at each time point of the data corresponds to a single \"token\" in our model. However, this value is continuous unlike the discrete inputs to a language transformer. Correspondingly, instead of a look-up table typically used in language transformers to map the discrete tokens into a higher dimensional embedding, we simply use a trainable linear layer to map the signal at each time point to a high dimensional embedding.\n- The reviewer asked us to add discussion about functional, anatomical, and genetic cell-type classification. We will add this discussion in the final manuscript.\n- The reviewer asked why we chose K=5 for the representation. We chose K=5 because the clustering performance and our heldout KNN metric saturates at this dimensionality. We would be happy to add this analysis to final version of the paper.\n- The reviewer asked if we performed an ablation of the projection MLP. We found that removing the MLP hurt the performance of our method. We would be happy to include this analysis in the final manuscript.\n- The reviewer asked if CEED can be scaled to big/relevant modern datasets which have potentially 100s of neurons. It is important to note that the clustering step in many spike sorting algorithms is performed on a local spatial neighborhood of channels which has anywhere from 1-10 neurons. Therefore, CEED can be used with this approach to cluster large extracellular datasets. We test this hypothesis in supplement figure 5 where we show that CEED can be trained and tested on a much larger dataset (spikes from 400 neurons) and still has large performance gains over the baseline feature extraction methods. In Table 1 of the attached PDF, we also show that when CEED is trained on spikes from 100s of neurons, it is able to generalize well to OOD data.\n- The reviewer asked why we did not evaluate our spike sorting experiments using the accuracy after optimal permutation with the Hungarian algorithm. We can add this metric to the final manuscript, however, we do not expect the conclusions to change as this metric is very correlated with the ARI.\n- The reviewer pointed out a few errors in the writing (missed citations, spelling, and abbreviations) that we will address in the final manuscript. \n"}, {"review_id": "qTDHmiF6MG", "reviewer": "Reviewer_5ReK", "summary": "The authors propose a contrastive learning method for obtaining representations of extracellular recordings which could be used for spike sorting and cell type classification. Transformer based encoder is used to generate low-dim representations of random views of spike waveforms which are then compared using a contrastive loss. The authors evaluate the proposed method on a simulated and real extracellular datasets, showing that it outperforms PCA embeddings for spike sorting and WaveMap on zero-shot cell types classification.\n\n", "questions": "* How does CEED compare to state-of-the-art spike sorting methods? The ARI scores above 90 suggest it is quite similar to Kilosort in terms of spike sorting performance, would it be fair to say that CEED performs on par with state-of-the-art spike sorting methods?\n\n* Have you tried comparing to non-linear embeddings rather than PCA (e.g. UMAP or autoencoders)? Would you believe CEED would outperform these methods as well?", "limitations": "I found the limitations section to be adequate and clearly written.", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "+ A nice application of modern machine learning methods (contrastive learning + transformers) to a classical problem (spike sorting)\n+ Clear presentation", "weaknesses": "- I found the notation in Section 3 to be a bit confusing. E.g. \\max over m and t in line 136 of A_{n, M} which doesn't have m and t in its indices. Superindices appear in line 149, even though W_{n, M} was defined without them in line 124. I think the notation could be simplified or at least made more consistent.\n- The experimentation evaluation is limited. The authors show the proposed method outperforms PCA on spike sorting, but typically PCA is not the only tool used for spike sorting. I am not sure I understand how close CEED gets to state-of-the-art spike sorting methods or manual spike sorting based on a variety of tools (PCA, autocorrelations, visualisations, etc.)\n- Ablation experiments could be useful as well. For example, it would be interesting to see if any of the proposed augmentations are more useful than the others or the effect of the transformer architecture choices.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their strong assessment of our work and for their detailed comments and questions. The review is very useful and led us to some changes that strengthened the paper significantly. In the following response, we address each point raised by the reviewer.\n- The reviewer correctly pointed out that the notation in section 3 was sometimes confusing or inconsistent. We agree with this assessment and will make the notation more clear and consistent in the final version of the paper. Specifically, we will make sure the superscripts and subscripts are used consistently and correctly. We will also try to remove unneeded subscripts.\n- The reviewer pointed out that the experimentation evaluation is limited and that the baselines (PCA and denoised PCA) are insufficient given the rich literature of spike sorting methods. This is a valid concern and highlights a clarity issue with our paper. We want to emphasize that *CEED is not a spike sorting method, but rather a feature extraction method for spike sorting and cell-type classification.* While there are a number of spike sorting methods for MEAs, almost all of these methods utilize PCA for feature extraction (Klusta, HerdingSpikes2, Mountainsort4, SpykingCircus, Trideclous, and even Kilosort uses a form of SVD/PCA). We argue that any method that uses PCA would benefit from switching to our more discriminable and robust features. Since directly incorporating CEED into these methods is a non-trivial coding challenge (especially given the short rebuttal window), we plan to add a new experiment to the final manuscript where CEED is incorporated into a simple spike sorting pipeline which includes detection, featurization, clustering, and template matching. We will use SpikeInterface for this analysis.\n- The reviewer asked how CEED compares to a state-of-the-art spike sorting method such as Kilosort. It is important to note that Kilosort has a number of additional processing steps (cluster splitting and merging, template matching, etc.) which make direct comparison challenging. However, as stated in the above paragraph, CEED is a feature extraction method which can be utilized by full spike sorting pipelines.\n- The reviewer asked if we have compared CEED to any non-linear methods (i.e. umap or autoencoders). This is a great question as without any comparisons to non-linear methods it is unclear if the performance gains come from the non-linearity or the proposed data augmentations and contrastive learning objective. For the cell-type classification experiments, we refer the reviewer to Figure 3 where we provide a direct comparison to a non-linear method, WaveMap, which utilizes UMAP. We show that CEED is slightly stronger than WaveMap even without training on the dataset they provide. For spike sorting, we did not compare to a non-linear baseline in the submitted manuscript. To correct for this, we now include results in Table 1 of the attached PDF and the General Response showing that CEED also significantly outperforms a non-linear baseline (an autoencoder) on spike sorting. We thank the reviewer for this suggested experiment as it strengthens the manuscript considerably.\n- The reviewer mentions that ablation experiments would be useful. We completely agree and have included an ablation study of CEED’s data augmentations (Tables 2 and 3) in the attached PDF. We also benchmark an MLP architecture for CEED which has comparable performance to our SCAM architecture (see Table 1, Column 1).\n"}]

Author rebuttal:
We thank all the reviewers for the detailed and useful reviews. Using this feedback, we have made improvements to CEED and run a number of new experiments which we detail below. We also provide some discussion of shared reviewer concerns below. Please see the attached PDF for the referenced tables and figure.
- **Computational complexity**. There were concerns about CEED’s high computational complexity and the requirement that it needs multiple GPUs to train. To improve the computational complexity of CEED, we propose a few modifications to the original method. (1) We realized our original data augmentations were slower than expected because we were computing them on additional channels that were not used during training. By only computing the augmentations on a small neighborhood of channels, we see a noticeable per epoch runtime boost. (2) While our SCAM architecture showed high performance across a number of datasets, we want to emphasize that the CEED framework is general and can be utilized with other architectures that require less computational resources. To illustrate this, we train a number of MLPs using our data augmentation and contrastive learning scheme. We show that the performance of the MLPs are comparable to the original SCAM architecture (see Table 1, Column 1) with a significant reduction in computational complexity (see Table 2). We use the MLP architecture for all results in the attached PDF. While we are still looking into ways to improve the computational complexity of SCAM, this is challenging as transformers are still an active area of research. Recent progress in quantization [1] and acceleration software [2] offer promising solutions, but incorporating these methods into CEED would be a future direction.
- **Out-of-distribution (OOD) training and a new non-linear baseline**. A few reviewers wanted to see non-linear baselines and OOD performance evaluation for our spike sorting experiments. We completely agree with both of these points and have added a new experiment to address this (see Table 1).  In this experiment, we benchmark the performance of CEED, our original baselines, and a new non-linear autoencoder on a new 10 neuron OOD dataset. For this OOD dataset, we train each method with spikes from a large number of neurons (390) and then test on spikes from 10 heldout neurons (Table 1, Column 3). We show that CEED significantly outperforms the original baselines and the new autoencoder baseline for both the in-distribution (ID) and OOD datasets. 
- **Ablations of the data augmentations**. Some of the reviewers wanted us to ablate our data augmentations to see which ones were the most important for CEED. Please see Table 3 of the attached PDF for an ablation of three of our data augmentations (all augmentations are detailed in supplement A.1). The most impactful data augmentation for CEED’s performance was the max channel shift augmentation (.38 drop in the ARI). To further test the importance of this max channel shift augmentation, we ran another experiment (see Table 4) where we created a new training and testing dataset. Instead of extracting each spike at its maximum amplitude channel; we instead extract each spike on *its template's maximum amplitude channel*. This means that spikes from the same neuron will be centered on the same channel. We designed this experiment to evaluate the performance of CEED and our baselines without having to account for max channel shifts due to noise. For fair comparison, we turn off  CEED’s max channel shift augmentation for this dataset. As can be seen, CEED still has the highest performance of all methods. We want to emphasize that this dataset utilizes ground-truth information (each neuron’s template) and is an ablation to better understand each method’s performance. We would be happy to add an ablation for all the data augmentations in the final version of the manuscript.
- **Visualizing approximate invariances of CEED's embeddings**. A few reviewers wanted to see if CEED's embeddings were approximately invariant to the proposed data augmentations. Please see Figure 1 in the attached PDF for a visualization of CEED embeddings for 3 neurons under all different data augmentations.
-  **CEED vs. full spike sorting pipelines**. A few reviewers had questions about how CEED compares to full spike sorting pipelines such as Kilosort. We want to emphasize that *CEED is not a spike sorting method, but rather a feature extraction method for spike sorting and cell-type classification*. While there are a number of spike sorting methods for MEAs, almost all of these methods currently utilize principal components analysis (PCA) for feature extraction (Klusta, HerdingSpikes2, Mountainsort4, SpykingCircus, Trideclous, and Kilosort use a form of SVD/PCA). We argue that any method that uses PCA would benefit from utilizing our more discriminable and robust features. Since directly incorporating CEED into these methods is a non-trivial coding challenge (especially given the short rebuttal window), we plan to add a new experiment to the final manuscript where CEED is incorporated into a simple spike sorting pipeline which includes detection, featurization, clustering, and template matching. We will use the SpikeInterface for this analysis.

[1] Liu, Zhenhua, et al. "Post-training quantization for vision transformer." NeurIPS 2021

[2] Ren, Jie, et al. "ZeRO-Offload: Democratizing Billion-Scale model training." 2021 USENIX Annual Technical Conference. 2021.

Meta-review: The paper proposes a self-supervised feature extraction technique for spike sorting from extracellular electrophysiological recordings. The reviewers agree that it’s a novel and technically solid approach. They criticize the simple PCA baseline and somewhat limited evaluation. The authors have addressed some of these concerns in the rebuttal. In addition, spike sorting pipelines are usually quite complex and spike sorting is inherently difficult to benchmark quantitatively. The authors’ work concerns only the feature extraction step. Hence, I consider the present work a potentially valuable contribution that should be accepted.
Decision: Accept (poster)