SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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 to create the training set with 200 or 1200 training spikes per unit. The goal of this dataset was to test how well CEED could generalize across a large set 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).
The third and final dataset was also extracted from both the DY016 and DY009 extracellular recordings. Unlike the first and second dataset, where the training set contained spikes from the test units (in-distribution), we purposefully excluded all units in the test set from the training set. 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-distribution (OOD) units.
4.2 Cell-type
For our cell-type classification dataset, we utilized the DY016 extracellular recording. We extracted the same 10 IBL "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 each waveform as described in Section 3.2.2.
207   
208   
209   
210   
211   
212   
213   
214   
215   
216   
217   
218   
219   
220   
221   
222   
223   
224   
225   
226   
227   
228   
229   
230   
231   
232   
233   
234   
235   
236   
237   
238   
239   
240   
241   
242   
243   
244   
245   
246   
247   
248
5.1 Spike sorting
Re-sorting units extracted using KS2.5 In this experiment, we aim to demonstrate that CEED extracts more useful features for spike sorting than both PCA and a non-linear autoencoder. We evaluate these methods on both an in-distribution (ID) and out-of-distribution (OOD) waveform discrimination task. Specifically, we train CEED, PCA, and an 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 to 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-parametric 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 [21, 58]. The YASS denoiser was trained to denoise single channel waveforms so 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.
5.2 Cell-type classification
Cell-type classification of IBL recordings In this analysis, we aim to 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 to find putative cell-types. As there is no ground-truth for the IBL recordings, we validate our results by visualizing the templates for each cell-type cluster.
Zero-shot cell-type classification. In this experiment, we compare the inferred representations of CEED to a state-of-the-art cell-type classification method, WaveMap, on an out-of-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 are recorded while a monkey performs a discrimination task [13]. 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. 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 how OOD this dataset is compared to the recordings used to train CEED, success on this task would demonstrate the robustness and generalizability of CEED.
6.1 Spike Sorting
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 1. 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 D that the strong performance of CEED generalizes to many different sets of test units.
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 2. Visually, it can be seen that CEED's features are far more informative about each unit's identity than the representations found by PCA. The performance of CEED is also much higher when being clustered by either HDBSCAN or a GMM even when when more principal components are afforded for the analysis.
250   
251   
252   
253   
254   
255   
256   
257   
258   
259   
260   
261   
262   
263   
264   
265   
266   
267   
268   
269   
270   
271   
272   
273   
274   
275   
276   
277   
278   
279   
280   
281   
282   
283   
284   
285   
286   
287   
288   
289   
290   
291   
292   
293   
294   
295   
296   
297
Method	10 neuron train set 10 neuron ID (ARI)	400 neuron train set 10 neuron ID (ARI)	390 neuron train set 10 neuron OOD (ARI)
CEED (1200 spikes, 11 channels)	.89 ± .04	.79 ± .09	.78 ± .05
CEED (1200 spikes, 5 channels)	.83 ± .03	.76 ± .07	.77 ± .08
Denoised PCA (1200 spikes, 11 channels)	.39 ± .05	.45 ± .04	.46 ± .04
Denoised PCA (1200 spikes, 5 channels)	.46 ± .07	.48 ± .04	.49 ± .04
Autoencoder (1200 spikes, 11 channels)	.47 ± .06	.35 ± .03	.28 ± .02
Autoencoder (1200 spikes, 5 channels)	.43 ± .06	.37 ± .03	.33 ± .01
Table 1: Benchmarking CEED, PCA, and an autoencoder on in-distribution (ID) and out-of-distribution (OOD) data. For evaluation, we fit 50 GMMs to 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 performs well on OOD data and can outperform a non-linear autoencoder. All CEED results are generated using the MLP encoder.
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.
Zero-shot cell-type classification. The results for cell-type classification of OOD single unit data are shown in Figure 3. Despite training on IBL extracellular datasets from a mouse brain recorded with Neuropixels 1.0, CEED is able to generalize to unseen data from a completely different animal (monkey) and probe (Plexon U-probe). In Figure 3A, we visualize the inferred representations from CEED using a 2D UMAP and by coloring the points with the output of a 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 3C). To quantitatively assess which cell-type classification method better reflects the 'real' differences in extracellular waveforms, we utilize a supervised classifier which is trained to predict cell-type labels using input waveforms (introduced in [13]). In Figure 3B and Figure 3D, we show the results of this data on 5 cross-validation folds. Despite never training on this dataset, CEED has a much higher accuracy  $(93.4\%)$  compared to that of WaveMap  $(88.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. Despite this difference, the accuracy value reported in the WaveMap paper for the 8 cell-type clusters  $(91\%)$  was still lower than that of CEED.
In this paper, we introduced a novel representation learning method, CEED, for extracellular recordings. Our main hypothesis was that by finding representations of extracellular waveforms that are robust to both common and task-specific nuisance variables, we can outperform specialized feature extraction approaches on two key tasks: spike sorting and morphoelectrical cell-type classification. 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 a waveform discrimination task. For cell-type classification, we show that CEED is 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 non-linear cell-type classification method WaveMap, on an animal and probe geometry unseen during training.
298   
299   
300   
301   
302   
303   
304   
305   
306   
307   
308   
309   
310   
311   
312   
313   
314   
315   
316   
317   
318   
319   
320   
321   
322   
323   
324   
325   
326   
327   
328   
329   
330   
331   
332
Denoised PCA, True Labels
(C)
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-parametric (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 a supervised denoiser to reduce noise in 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 to the learned embeddings from CEED (left) and denoised PCA (right).
Despite the significant performance improvements of CEED, there are few limitations that must be addressed before it can become a 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-based 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-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.
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 it requires to train and run. Contrastive learning requires large batch sizes to achieve state-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], CEED could have a possible negative impact on the environment. We also propose an MLP-based
333   
334   
335   
336   
337   
338   
339   
340   
341   
342   
343
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 a OOD dataset which consists of 625 single units extracted from a monkey recording. We cluster the representations using a GMM and plot each cluster with a distinct color. For each cluster, we 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 a different animal and probe geometry, but can still generalize to this dataset, outperforming WaveMap on a 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 a 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 as the average of the diagonal. (C) A 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 a gradient boosted decision tree classifier trained to map raw waveforms to the cell-type clusters extracted by WaveMap.
architecture that performs comparably to a transformer while only using a single GPU, which can lower CEED's impact on the environment.
[1] James J Jun, Nicholas A Steinmetz, Joshua H Siegle, Daniel J Denman, Marius Bauza, Brian Barbarits, Albert K Lee, Costas A Anastassiou, Alexandru Andrei, Căgatay Aydin, et al. Fully integrated silicon probes for high-density recording of neural activity. Nature, 551(7679):232-236, 2017.
[2] Ashley L Juavinett, George Bekheet, and Anne K Churchland. Chronically implanted neuropixels probes enable high-yield recordings in freely moving mice. *Elife*, 8:e47188, 2019.
353   
354   
355   
356   
357   
358   
359   
360   
361
362   
363   
364   
365   
366   
367   
368   
369   
370   
371   
372   
373   
374   
375   
376   
377   
378   
379   
380   
381   
382   
383   
384   
385   
386   
387   
388   
389   
390   
391   
392   
393   
394   
395   
396   
397   
398   
399   
400   
401   
402   
403   
404   
405   
406   
407   
408   
409
410   
411   
412   
413   
414   
415   
416   
417   
418   
419   
420   
421   
422   
423   
424   
425   
426   
427   
428   
429   
430   
431   
432   
433   
434   
435   
436   
437   
438   
439   
440   
441   
442   
443   
444   
445   
446   
447   
448   
449   
450   
451   
452   
453   
454   
455   
456   
457   
458
459   
460   
461   
462   
463   
464   
465   
466   
467   
468   
469   
470   
471   
472   
473   
474   
475   
476   
477   
478   
479   
480   
481   
482   
483   
484   
485   
486   
487   
488   
489   
490   
491   
492   
493   
494   
495   
496   
497   
498   
499   
500   
501   
502   
503   
504   
505
[50] Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding, 2019.
[51] 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.
[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, page 1857-1865, Red Hook, NY, USA, 2016. Curran Associates Inc.
[53] Aäron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding. CoRR, abs/1807.03748, 2018.
[54] The International Brain Laboratory, Valeria Aguillon-Rodriguez, Dora Angelaki, Hannah Bayer, Niccolo Bonacchi, Matteo Carandini, Fanny Cazettes, Gabelle Chapuis, Anne K Churchland, Yang Dan, et al. Standardized and reproducible measurement of decision-making in mice. *Elife*, 10, 2021.
[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.
[56] Leland McInnes, John Healy, and Steve Astels. hdbscan: Hierarchical density based clustering. J. Open Source Softw., 2(11):205, 2017.
[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.
[59] Jie Ren, Samyam Rajbhandari, Reza Yazdani Aminabadi, Olatunj Ruwase, Shuangyan Yang, Minjia Zhang, Dong Li, and Yuxiong He. Zero-offload: Democratizing billion-scale model training, 2021.
[60] Payal Dhar. The carbon impact of artificial intelligence. Nat. Mach. Intell., 2(8):423-425, 2020.
506   
507   
508   
509   
510   
511   
512   
513   
514   
515   
516   
517   
518   
519   
520   
521   
522   
523   
524   
525   
526   
527   
528   
529   
530   
531   
532   
533   
534   
535   
536