Title: Towards robust and generalizable representations of extracellular data using contrastive learning
Venue: NeurIPS 2023
Accepted: yes
Keywords: contrastive learning

representations

extracellular

high-density

spike sorting

cell-type classification

transformers

invariance
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.

1

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 ( C ontrastive E mbeddings for E xtracellular D ata), for highdensity 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 .

1 Introduction

High-density extracellular recordings now allow for simultaneous recordings of large populations 13 of neurons across multiple brain regions with high temporal and spatial resolution. [1-4]. These 14 large-scale recordings are essential for gaining insights into key biological processes such as vision, 15 decision-making, and behavior which are distributed across brain regions [5]. Along with gaining 16 insights into brain function, these technologies also promise to improve the scalability and accuracy 17 of brain-computer interfaces which can restore motor function to paralyzed individuals [6]. 18

19

20

21

22

23

24

25

26

27

A major bottleneck for interpreting neural population activity is processing the raw extracellular signal [7]. Although extracellular recordings contain a precise record of the coordinated neural activity, it must be extracted algorithmically through a processing step called spike sorting . A crucial assumption in spike sorting is that each recorded neuron has a unique spatiotemporal extracellular waveform based on its morphology and position relative to the recording device [8]. Using this unique identifier, it is possible to assign a detected extracellular action potential (spike) back to its putative neuron (unit). Along with spike sorting, another important task in extracellular analysis is cell-type classification. As neural circuits are diverse and heterogeneous, it is becoming increasingly important to profile detected units using their morphoelectrical features [9]. At a coarse level, it is

28

29

30

31

32

33

34

35

36

37

38

39

40

41

42

43

44

45

46

47

48

49

50

51

52

53

54

55

56

57

58

59

60

61

62

63

64

65

66

67

68

69

70

71

72

73

74

75

76

77

78

possible to classify units as inhibitory or excitatory based on their extracellular profile [10]; however, it may be possible to divide units into finer subgroups [11-13].

For both spike sorting and cell-type classification, it is important to extract low-dimensional, meaningful features from extracellular waveforms. Despite the importance of feature extraction, current approaches are often ad hoc and lack robustness to common nuisance variables in extracellular recordings. By far the most common featurization method for spike sorting is principal components analysis (PCA) [3, 14-22]. Although PCA is relatively effective and scalable, it suffers from a few key drawbacks including: (1) a lack of robustness to extracellular nuisance variables such as spatiotemporally overlapping spikes (collisions), (2) an inability to model non-linear data, and (3) an objective function that aims to find features that explain variance rather than features that discriminate different waveforms. To improve the robustness of PCA, [21] introduced a supervised waveform denoiser that is able to 'clean' the waveforms before featurization. For morphoelectric cell-type classification, most feature extraction methods rely heavily on manually extracted features including action potential width, peak-to-peak amplitude, and the ratio of pre-hyperpolarization peak to the post-hyperpolarization peak [13]. Again, while scalable and effective, these features are too simple and ad hoc to fully capture morpho-eletrical differences [13]. Recently, a non-linear approach to cell-type classification, WaveMap, was introduced, utilizing UMAP [23] and Louvain community detection [24] to automatically find cell-type clusters [13].

In this work, we introduce a robust and generalizable feature learning method, CEED, for extracellular datasets. Our main hypothesis is that embeddings of extracellular waveforms that are invariant to both common and task-specific nuisance variables will be more useful for spike sorting and morphoelectric cell-type classification than current specialized feature extraction methods. For example, in extracellular recordings, there are a number of nuisance variables including collisions, correlated background noise, or variability in the time at which a spike is detected [21]. Each of these confounds pose a problem for traditional representation learning methods like PCA or manually extracted features. Along with these common nuisance variables, there are also task-specific nuisance variables such as the spatial position of detected spikes. For spike sorting, this information is crucial as spike locations are highly informative of neuron identity [25, 26]. However, for cell-type classification, neurons with different locations may still share a cell-type. In order to extract representations that are invariant to these nuisance variables, we utilize contrastive learning, which has been shown to approximately induce invariance in the representation space to a set of transformations [27]. We utilize a recent contrastive learning framework [28] for our training and implement a stochastic view generation module for extracellular waveforms. We validate our approach on multiple high-density extracellular recordings. Surprisingly, for cell-type classification, our representations appear to be more informative than state-of-the-art specialized methods even when performing zero-shot learning on an unseen animal and probe geometry. Our contributions are as follows:

We introduce a novel framework, CEED, for analyzing extracellular recordings based on invariance learning.

We implement a stochastic view generation module for both single-channel and multichannel extracellular waveforms.

We demonstrate that CEED works well with multiple neural network architectures, including a novel transformer-based architecture with a spatiotemporal causal attention mask (SCAM).

We show that CEED outperforms specialized featurization methods for spike sorting and morphoelectric cell-type classification.

2 Background

Contrastive Learning. Contrastive representation learning [28] falls under a broad class of selfsupervised learning methods [29-31] whose goal is to learn robust and generalizable representations by encouraging invariances to a prior-known set of transformations or nuisance variables. While contrastive learning has mainly been used to extract effective and transferable representations from image-based data [27, 28], more recently, it has also become a powerful tool in the sciences to learn invariances for physical systems [32, 33].

In computational neuroscience, representation learning has mostly been performed using generative 79 models [34]. While this has led to many interesting insights into behavior [35] and decision-making 80 [36], this paradigm is sensitive to nuisance variables and will not capture subtle changes in the 81

Recording

Contrastive Learning

Representations

ViewGenerationModule

Amplitude

Jitter

TwoViews

Transformer / MLP

Encoder

MAKE

SIMILAR

Collision

Negative

MAKE

examples

DIFFERENT

Crop

Batch of samples

Temporal

m//

Jitter

Tasks

Waveforms

Embeddings

Noise

sorting,

e.g.spike

Unit1

Unit2

saI aUn

0

cell-type

ExtractedWaveform

classification

Unito

observation space that do not contribute much to the explained variance. While recent research 82 into label-guided generative models [37-39] has somewhat alleviated these problems, it is still 83 unclear if good generative performance is important for representation learning. Most recently, a 84 number of contrastive learning methods have been introduced for learning robust and generalizable 85 representations of neural population data [40-42]. To our knowledge, however, contrastive learning 86 has not been applied to key primary data analysis tasks such as spike sorting or cell-type classification. 87

88

89

90

91

92

93

94

95

96

97

98

99

Spike sorting. In spike sorting, the goal is to extract a precise record of which neurons spiked at which time steps based on the raw extracellular data. For each electrode, the contributions of an unknown number of neurons are mixed together, making spike sorting a challenging blind-source separation task akin to the cocktail party problem [43]. While a number of spike sorting methods have been developed for microelectrode arrays (MEAs), almost all of these methods currently utilize PCA for feature extraction For example, the popular algorithms Klusta, HerdingSpikes2, Mountainsort4, SpykingCircus, Trideclous, and Kilosort all use a form of SVD/PCA in their processing pipelines. Despite the ubiquity of PCA as a feature extraction method for spike sorting, it lacks robustness to common nuisance variables in extracellular recordings, cannot capture non-linear features of the observed data, and prefers features that explain variance rather than those that discriminate different waveforms. More recently, it has been shown that non-linear autoencoders can have higher performance than PCA despite suffering from similar drawbacks [44].

Cell-type classification. Morphoelectric cell-type classification has many similarities to spike sorting. 100 Similarly to spike sorting, you are also interested in grouping detected extracellular waveforms 101 together. However, the granularity of these clusters should be lower, reflecting shared morphologies 102 among the units in the recording. Classically, pre-defined morphoelectrical features (e.g., action 103 potential width) were extracted from the waveforms to group together similar units [10, 45]. More 104 recently, a non-linear method based on UMAP was introduced for cell-type classification [13]. In this 105 work, it was shown that automatically discovering relevant features from extracellular waveforms can 106 lead to more informative cell-type groupings than when simply clustering pre-defined features. It is 107 important to note that, along with morphoelectric cell-type classification, it is also possible to classify 108 cell-types by their function, morphology, physiology, molecular properties [46-48]. 109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

128

129

130

131

132

133

134

135

136

137

138

139

140

141

142

143

144

145

146

147

148

149

150

151

152

153

154

155

156

157

3 Contrastive Embeddings for Extracellular Data (CEED)

We introduce Contrastive Embeddings for Extracellular Data (CEED), a novel framework for extracting robust and generalizable representations of extracellular data via contrastive learning. CEED consists of three main components (see Figure 1) - (1) a stochastic view generation module that augments waveforms with both common and task-specific nuisance variables. (2) A neural network-based encoder that can extract low-dimensional representations from both single-channel and multi-channel waveforms. (3) A contrastive loss function that encourages the representations of views of the same waveform to be similar while forcing the representations of views of different waveforms to be dissimilar. We postulate that the representations extracted by CEED, which are approximately invariant to nuisance variables, can outperform specialized feature extraction methods for both spike sorting and cell-type classification.

3.1 Notation

To introduce our method, we must first define some notation for extracellular recordings. Let P M := { p m | m ∈ M } be the position of all M channels on the multi-electrode array (MEA), where each p m ∈ R 3 is the location of channel m . Now define P ˜ M to be the set of positions for a subset of channels, ˜ M ⊆ M , on the MEA.

Let S := { s n } N n =1 be the set of N spiking events that are detected in an extracellular recording. Now, let the extracellular waveform of a spiking event s n on a channel m be defined as W n,m := { r (0) n,m , r (1) n,m , ..., r ( T ) n,m } ∈ R T , where T contiguous samples from the recording define a waveform and r ( t ) n,m is the value of the waveform at sample t . This can be thought of as a window of samples of size T on channel m which includes the spike s n .

The amplitude of a spike s n on a channel m can be defined as A n,m := max t | W n,m | . This is also known as the absolute maximum voltage recorded on channel m . The timestep t at which the absolute maximum voltage occurs can vary for each channel. Therefore, for a spike, we center the waveform on each channel using the timestep at which the maximum A n,m occurred (i.e., the timestep of the maximum amplitude channel). For all spikes in the recording, waveforms are aligned in this manner so that the amplitude on the maximum amplitude channel occurs at the same timestep t .

Let W n, ˜ M ∈ R T ×| ˜ M | be the set waveforms for a spike on a subset of channels and let A n, ˜ M ∈ R > 0 | ˜ M | be the set amplitudes for a spike on a subset of channels, where | ˜ M | is the number of channels.

3.2 Stochastic view generation module

3.2.1 Common invariances

Let us define the representation of the waveform W n, ˜ M for spike s n as z n ∈ R D . There are a number of nuisance variables that change the observed waveform of a spike without changing the underlying signal. Therefore, we can define a set of invariances we wish to impose on our representations,

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 n, ˜ M for spike s n such that

˜ W n, ˜ M = W n, ˜ M × (1 ± ϵ ) , then the representation ˜ z n for ˜ W n, ˜ M should be such that ˜ z n = z n .

Correlated background noise - We want representations that are invariant to the correlated background noise found in real extracellular recordings. We assume extracellular noise can be modeled as a spatiotemporal matrix Gaussian MN (0 , U, V ) , where U ∈ R T × T models the temporal correlation of a waveform and V ∈ R | M |×| M | models the spatial correlation between channels. For a spike s n , we want a representation such that if we sample n ∼ MN (0 , U, V ) and form ˜ W n, ˜ M = W n, ˜ M + n ˜ M , where n ˜ M is the noise on the channels subset ˜ M , then ˜ W n, ˜ M should still yield a representation ˜ z n = z n .

Spike collisions - We want representations that are invariant to the voltage distortions caused when spikes "collide" spatiotemporally. A collision is when two or more spikes occur at similar times in the same location on the probe leading to observed waveforms that are distorted. To this end, consider two spikes s n 1 and s n 2 with waveforms W n 1 , ˜ M 1 and W n 2 , ˜ M 2 such that ˜ M 1

158

159

160

161

162

163

164

165

166

167

168

169

170

171

172

173

174

175

176

177

178

179

180

181

182

183

184

185

186

187

188

189

190

191

192

193

194

195

196

197

198

199

200

201

202

203

204

205

206

and ˜ M 2 share some number of channels. Now let us shift s n 2 by k time steps ( -T/ 2 < k < 0 or 0 < k < T/ 2 ) such that for any channel m ∈ ˜ M 2 , W n 2 ,m := { r (0+ k ) n 2 ,m , r (1+ k ) n 2 ,m , ..., r ( T + k ) n 2 ,m } . We refer to this temporally shifted waveform as W shift n , ˜ M . A new waveform that is a linear sum

2 2 of these two waveforms ˜ W n, ˜ M 1 = W n 1 , ˜ M 1 + W shift n 2 , ˜ M 2 should yield a representation ˜ z n = z n 1 .

Channel subsets - We want representations that are invariant to changes in the subset of channels ˜ M used to define W n, ˜ M . 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 n, ˜ M ′ should still yield a representation ˜ z n = z n .

3.2.2 Task-specific invariances

While all feature extraction methods suffer from common extracellular nuisance variables, there are also task-specific nuisance variables. For example, the position and orientation of neurons in the recording are essential information for spike sorting. However, for cell-type classification, this information can be a confound when trying to find shared morphoelectrical features. Therefore, we propose an additional set of invariances for cell-type classification,

Cell position - For cell-type classification, we want representations that are invariant to the channel positions P M at which a spike is detected. So for a waveform W n,M , any uniform changes to the channel positions P ˜ M (without changing the waveforms at each channel) should still yield a representation ˜ z n = z n .

Cell amplitude - As amplitude is mainly a function of cell position and orientation, we want representations that are invariant to any uniform changes in the amplitudes. To this end, for any positive value a , a change in the amplitudes for the spike s n , ˜ W n, ˜ M = W n, ˜ M × a , should still yield a representation ˜ z n = z n .

To achieve these task-specific invariances for cell-type classification, we directly transform the training data. First, we extract the max channel waveform for each spike. Then, we normalize each waveform to be between -1 and 1, thus removing any positional information [13]. A drawback of this approach is that we discard multi-channel information that may be useful for classifying different cell-types [12] (see Section 7 for a more detailed discussion). The full view generation pipeline is detailed in Supplementary Materials Section A.1.

3.3 Encoder Architecture

For the encoder of CEED, we explored two neural network architectures - a transformer-based network with a novel Spatiotemporal Causal Attention Mask (SCAM) and a simple multi-layer perceptron (MLP) which is a more computationally efficient alternative (i.e., can fit on a single GPU). For a runtime comparison of these two architectures, see Supplementary Materials Section E.

3.3.1 Transformer-based encoder with spatiotemporal causal attention mask (SCAM)

Our first proposed architecture utilizes transformers [49] which have been highly successful across a series of tasks in natural language processing [50, 51]. Transformers have a natural inductive bias towards time-series and sequence-based data and are thus highly suitable for extracellular waveforms. For this work, we designed a novel spatiotemporal causal attention mask (SCAM) to obey causality across time and channels. Specifically, we allow every recorded time step in a waveform to attend to time steps on other channels as long as those data points precede it in time. Full details and visualization of the transformer-based architecture, including implementation details are available in Supplementary Materials Section A.2.

3.3.2 Multi-layer perceptron (MLP)

While a transformer-based architecture provides a natural inductive bias towards time-series data appropriate for extracellular recordings, a downside to such an architecture is its high computational complexity and the requirement that it needs multiple GPUs to train. To demonstrate the generality of the CEED framework, we also propose a simpler MLP architecture that can be trained on a single GPU. The MLP encoder is a straightforward model that consists of three layers with sizes [768 , 512 , 256] and ReLU activations between them.

207

208

209

210

211

3.4 Objective function

The encoder is trained as follows. Let ˜ z i and z i be representations of the two views of the input waveform i , where the two views differ by the invariances listed in Sections 3.2.1 and 3.2.2. The optimization objective for a batch of B samples follows the contrastive loss utilized by SimCLR [28] and several other representation learning methods [52, 53];

̸

where ˆ z i and ˆ ˜ z i are the L2 normalized representations and τ is a temperature hyperparameter. 212 Following [28], we also include a 2-layer MLP projector network after the encoder and the loss 213 function operates on the outputs of this projector network (see Supplementary Materials Section 214 A.3 for details and ablations of the projector network and its architecture). 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

4 Datasets

To train and evaluate our model, we make use of two publicly available extracellular recordings published by the International Brain Laboratory (IBL): the DY016 and DY009 recordings [54]. These multi-region, Neuropixels 1.0 recordings are taken from a mouse performing a decision-making task (see Supplementary Materials Section C for more details).

4.1 Spike sorting

To evaluate how useful the features learned by CEED are for spike sorting, we constructed three datasets using units found by Kilosort 2.5, a full spike sorting pipeline manually tuned by IBL. The first dataset was extracted from the DY016 extracellular recording. It consisted of a 10 unit train and test dataset where all 10 units were classified as "good" by IBL's quality metrics [55]. We selected these units for their high waveform diversity (qualitatively) and because they had a relatively high amplitude, i.e., peak-to-peak (ptp). For this dataset, we constructed training sets of 200 or 1200 spikes per unit with a test set of 200 spikes per unit. For each spike, we extracted waveforms from 21 channels centered on the maximum amplitude channel. Although we extract 21 channels for our data augmentations (see Supplementary Materials Section A.1), we train and evaluate our model (and baselines) on either 5 or 11 channel subsets.

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 245 the same 10 IBL "good" units used for the spike sorting datasets. To remove positional information 246 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. 248

249

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

5

Experiments

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 Results

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 GMMeven when when more principal components are afforded for the analysis.

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

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 wellisolated 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.

7 Discussion

In this paper, we introduced a novel representation learning method, CEED, for extracellular record323 ings. Our main hypothesis was that by finding representations of extracellular waveforms that are 324 robust to both common and task-specfic nuisance variables, we can outperform specialized feature 325 extraction approaches on two key tasks: spike sorting and morphoelectrical cell-type classification. 326 We validate CEED on multiple high-density extracellular datasets. For spike sorting, we show 327 that CEED extracts features that far outperforms those of PCA and a non-linear autoencoder on a 328 waveform discrimination task. For cell-type classification, we show that CEED is able to extract 329 discriminative features of extracellular waveforms that allow for finding morphological subgroups 330 in an unsupervised manner. Surprisingly, we find that CEED even outperforms a recent non-linear 331 cell-type classification method WaveMap, on an animal and probe geometry unseen during training. 332

GaussianMixtureModel(10components)

HDBSCAN (10 min cluster size,10 min samples)

(A)

Adjusted Rand Index

90 00 000 50 40

★

★CEED

Adjusted Rand Index

80

￥

★CEED

PCA

dI1

→Denoised PCA

Denoised PCA

4020

3

5

7

9

3

5

7

6

#of PCs

#of PCs

(B)

CEEDembedding,GMM Pred Labels

Denoised PCA, GMM Pred Labels

CEED embedding,True Labels

ARI: 93.62

Denoised PCA,True Labels

ARI: 62.80

2.5

2.0

2.0

CEED PC2

1.5

1.5

1.0

Denoised PC 2

2

0°0

0.5

0.5

1.0

1.0

1.5

1.5

-3

-2

CEED PC 1

。

1

2

CEED PC 1

-1

Denoised PC1

。

1

2

-1

Denoised PC 1

。

2

3

(C)

CEEDembedding,HDBSCANPred

Denoised PCA,HDBSCAN Pred

CEED embedding,TrueLabels

Denoised PCA,True Labels

2.5

2.5

3

2.0

2.0

1.5

2

2

CEED PC 2

1.0

Denoised PC 2

1

00

1

-0.5

0.5

2

1.0

1.0

1.5

1.5

-2

-1

1

2

1.5 1.0 0.5 0.0 0.5 1.0

CEED PC 1

1.52.02.5

-2

1

Denoised PC 1

-2

-1

Denoised PC 1

CEED PC 1

Despite the significant performance improvements of CEED, there are few limitations that must 333 be addressed before it can become a plug-and-play method for extracellular analysis. Firstly, the 334 training sets we utilize in this work are still quite small and lack neuron diversity. For CEED to 335 be generalizable to multiple recordings and animals, more diverse datasets must be constructed. 336 Secondly, the best results of CEED are obtained when using the transformer-based encoder which 337 requires multiple GPUs and is currently quite slow to train. Recent progress in acceleration software 338 [59] offer promising solutions to speed up computation and incorporating these methods into CEED 339 could be a future direction. Thirdly, all spike sorting results in the paper are from re-sorting already 340 sorted datasets; CEED must be incorporated into a full spike sorting pipeline in order to be used 341 by many different research groups. Finally, our cell-type results do not include any functional 342 classification which could help validate the clusters found by CEED. 343

344

8 Broader Impact

Although CEED has the potential to improve key tasks in extracellular analysis, a drawback of 345 our approach compared to simple approaches like PCA is the additional computational resources 346 it requires to train and run. Contrastive learning requires large batch sizes to achieve state-of347 the-art performance which means, if training the SCAM transformer model, we often have to 348 run experiments on large-scale, multi-gpu clusters. Moreover, transformer-based architectures 349 have hundreds of thousands of parameters, which further increases CEED's compute requirements. 350 As highly parameterized deep neural networks produce large amounts of carbon emissions [60], 351 CEED could have a possible negative impact on the environment. We also propose an MLP-based 352

353

354

355

356

357

358

359

360

361

UMAP1

(A)

CEED-93.4% accuracy (zero-shot learning)

(B)

8

BIC

Q

spikes: 87

6-

spikes: 76

pooolo

0

6

81012

UMAP 2

#of Clusters

1.0

2

spikes: 66

08

True Class

1

100

89

2

5

4

0.8

96

4

0.6

spikes: 173

spikes: 119

8

92

0.4

11

89

0.2

-2

spikes: 104

6

3

3

93

0.0

0.0

2.5

5.0

7.5

10.0

12.5

15.0

17.5

UMAP 1

1

2

3

4

5

6

(C)

WaveMap-88.8%accuracy

PredictedClass

6

(D)

8

splkes: 37

1.0

spikes: 162

1

81

6

11

3

7

0.8

UMAP 2

6

spikes: 112

True Class

8

84

2

6

0.6

100

6

9

84

0.4

4

7

93

0.2

3

3

3

3

91

V0.0

2

?

1

2

3

4

5

6

spikes: 83

Predicted Class

1

spikes: 92

spikes: 139

-2

0

2

4

6

8

10

12

architecture that performs comparably to a transformer while only using a single GPU, which can lower CEED's impact on the environment.

References

James J Jun, Nicholas A Steinmetz, Joshua H Siegle, Daniel J Denman, Marius Bauza, Brian Barbarits, Albert K Lee, Costas A Anastassiou, Alexandru Andrei, Ça˘ gatay Aydın, et al. Fully integrated silicon probes for high-density recording of neural activity. Nature , 551(7679):232236, 2017.

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.

P (Predicted | True)

P (Predicted I True)

Nicholas A Steinmetz, Cagatay Aydin, Anna Lebedeva, Michael Okun, Marius Pachitariu, 362 Marius Bauza, Maxime Beau, Jai Bhagat, Claudia Böhm, Martijn Broux, et al. Neuropix363 els 2.0: A miniaturized high-density probe for stable, long-term brain recordings. Science , 364 372(6539):eabf4588, 2021. 365

366

367

368

369

International Brain Laboratory, Kush Banga, Julius Benson, Niccolò Bonacchi, Sebastian A Bruijns, Rob Campbell, Gaëlle A Chapuis, Anne K Churchland, M Felicia Davatolhagh, Hyun Dong Lee, et al. Reproducibility of in-vivo electrophysiological measurements in mice. bioRxiv , pages 2022-05, 2022.

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

Nicholas A Steinmetz, Peter Zatka-Haas, Matteo Carandini, and Kenneth D Harris. Distributed coding of choice, action and engagement across the mouse brain. Nature , 576(7786):266-273, 2019.

Elon Musk et al. An integrated brain-machine interface platform with thousands of channels. Journal of medical Internet research , 21(10):e16194, 2019.

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-Technological Applications , pages 171-184, 2019.

Carl Gold, Darrell A Henze, Christof Koch, and Gyorgy Buzsaki. On the origin of the extracellular action potential waveform: a modeling study. Journal of neurophysiology , 95(5):3113-3128, 2006.

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 , 31(40):14235-14242, 2011.

Peter Barthó, Hajime Hirase, Lenaïc 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.

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 biophysical modeling and deep learning for multielectrode array neuron localization and classification. Journal of neurophysiology , 120(3):1212-1232, 2018.

Xiaoxuan Jia, Joshua H Siegle, Corbett Bennett, Samuel D Gale, Daniel J Denman, Christof Koch, and Shawn R Olsen. High-density extracellular probes reveal dendritic backpropagation and facilitate neuron classification. Journal of neurophysiology , 121(5):1831-1847, 2019.

Eric Kenji Lee, Hymavathy Balasubramanian, Alexandra Tsolias, Stephanie Udochukwu Anakwe, Maria Medalla, Krishna V Shenoy, and Chandramouli Chandrasekaran. Non-linear dimensionality reduction on extracellular waveforms reveals cell type diversity in premotor cortex. Elife , 10:e67490, 2021.

Samuel Garcia and Christopher Pouzat. Tridesclous , 2015. https://github.com/ tridesclous/tridesclous .

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):634-641, 2016.

Marius Pachitariu, Nicholas A Steinmetz, Shabnam N Kadir, Matteo Carandini, and Kenneth D Harris. Fast and accurate spike sorting of high-channel count probes with kilosort. Advances in neural information processing systems , 29, 2016.

Jason E Chung, Jeremy F Magland, Alex H Barnett, Vanessa M Tolosa, Angela C Tooker, Kye Y 407 Lee, Kedar G Shah, Sarah H Felix, Loren M Frank, and Leslie F Greengard. A fully automated 408 approach to spike sorting. Neuron , 95(6):1381-1394, 2017. 409

Gerrit Hilgen, Martino Sorbaro, Sahar Pirmoradian, Jens-Oliver Muthmann, Ibolya Edit Kepiro, 410 Simona Ullo, Cesar Juarez Ramirez, Albert Puente Encinas, Alessandro Maccione, Luca 411 Berdondini, et al. Unsupervised spike sorting for large-scale, high-density multielectrode arrays. 412 Cell reports , 18(10):2521-2532, 2017. 413

414

415

416

417

Pierre Yger, Giulia LB Spampinato, Elric Esposito, Baptiste Lefebvre, Stéphane Deny, Christophe Gardella, Marcel Stimberg, Florian Jetter, Guenther Zeck, Serge Picaud, et al. A spike sorting toolbox for up to thousands of electrodes validated with ground truth recordings in vitro and in vivo. Elife , 7:e34518, 2018.

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

Roland Diggelmann, Michele Fiscella, Andreas Hierlemann, and Felix Franke. Automatic spike sorting for high-density microelectrode arrays. Journal of neurophysiology , 120(6):3155-3171, 2018.

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, 2020.

Keming Chen, Yangtao Jiang, Zhanxiong Wu, Nenggan Zheng, Haochuan Wang, and Hui Hong. Htsort: Enabling fast and accurate spike sorting on multi-electrode arrays. Frontiers in Computational Neuroscience , 15:657151, 2021.

Leland McInnes, John Healy, and James Melville. Umap: Uniform manifold approximation and projection for dimension reduction. arXiv preprint arXiv:1802.03426 , 2018.

Vincent D Blondel, Jean-Loup Guillaume, Renaud Lambiotte, and Etienne Lefebvre. Fast unfolding of communities in large networks. Journal of statistical mechanics: theory and experiment , 2008(10):P10008, 2008.

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.

Julien Boussard, Erdem Varol, Hyun Dong Lee, Nishchal Dethe, and Liam Paninski. Threedimensional spike localization and improved motion correction for neuropixels recordings. Advances in Neural Information Processing Systems , 34:22095-22105, 2021.

Rumen Dangovski, Li Jing, Charlotte Loh, Seungwook Han, Akash Srivastava, Brian Cheung, Pulkit Agrawal, and Marin Soljaˇ ci´ c. Equivariant contrastive learning. arXiv preprint arXiv:2111.00899 , 2021.

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.

Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick. Momentum contrast for unsupervised visual representation learning, 2020.

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: A new approach to self-supervised learning, 2020.

Jure Zbontar, Li Jing, Ishan Misra, Yann LeCun, and Stéphane Deny. Barlow twins: Selfsupervised learning via redundancy reduction, 2021.

Yuyang Wang, Jianren Wang, Zhonglin Cao, and Amir Barati Farimani. Molecular contrastive learning of representations via graph neural networks. Nature Machine Intelligence , 4(3):279287, mar 2022.

Charlotte Loh, Thomas Christensen, Rumen Dangovski, Samuel Kim, and Marin Soljaˇ ci´ c. Surrogate- and invariance-boosted contrastive learning for data-scarce applications in science. Nature Communications , 13(1), jul 2022.

Cole Hurwitz, Nina Kudryashova, Arno Onken, and Matthias H Hennig. Building population 459 models for large-scale neural recordings: Opportunities and pitfalls. Current opinion in 460 neurobiology , 70:64-73, 2021. 461

462

463

464

Juan A Gallego, Matthew G Perich, Raeed H Chowdhury, Sara A Solla, and Lee E Miller. Long-term stability of cortical population dynamics underlying consistent behavior. Nature neuroscience , 23(2):260-270, 2020.

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

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-212, 2022.

Ding Zhou and Xue-Xin Wei. Learning identifiable and interpretable latent models of highdimensional neural activity using pi-vae. Advances in Neural Information Processing Systems , 33:7234-7247, 2020.

Omid G Sani, Hamidreza Abbaspourazad, Yan T Wong, Bijan Pesaran, and Maryam M Shanechi. Modeling behaviorally relevant neural dynamics enabled by preferential subspace identification. Nature Neuroscience , 24(1):140-149, 2021.

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 , 34:29379-29392, 2021.

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:10587-10599, 2021.

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.

Steffen Schneider, Jin Hwa Lee, and Mackenzie Weygandt Mathis. Learnable latent embeddings for joint behavioural and neural analysis. Nature , pages 1-9, 2023.

Glen D Brown, Satoshi Yamada, and Terrence J Sejnowski. Independent component analysis at the neural cocktail party. Trends in neurosciences , 24(1):54-63, 2001.

Eugen-Richard Ardelean, Andreea Coporîie, Ana-Maria Ichim, Mihaela Dîns , oreanu, and Raul Cristian Mures , an. A study of autoencoders as a feature extraction technique for spike sorting. Plos one , 18(3):e0282810, 2023.

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.

Thomas Euler, Silke Haverkamp, Timm Schubert, and Tom Baden. Retinal bipolar cells: elementary building blocks of vision. Nature Reviews Neuroscience , 15(8):507-519, 2014.

Hongkui Zeng and Joshua R Sanes. Neuronal cell-type classification: challenges, opportunities and the path forward. Nature Reviews Neuroscience , 18(9):530-546, 2017.

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):935953, 2020.

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems , 30, 2017.

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

Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding, 2019.

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.

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.

Aäron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding. CoRR , abs/1807.03748, 2018.

The International Brain Laboratory, Valeria Aguillon-Rodriguez, 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 , 10, 2021.

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.

Leland McInnes, John Healy, and Steve Astels. hdbscan: Hierarchical density based clustering. J. Open Source Softw. , 2(11):205, 2017.

Lawrence Hubert and Phipps Arabie. Comparing partitions. Journal of classification , 2:193-218, 1985.

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.

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.

Payal Dhar. The carbon impact of artificial intelligence. Nat. Mach. Intell. , 2(8):423-425, 2020. 536

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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\nSTRENGTHS:\n+ A nice application of modern machine learning methods (contrastive learning + transformers) to a classical problem (spike sorting)\n+ Clear presentation\n\nWEAKNESSES:\n- 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.\n\nQUESTIONS:\n* 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?\n\nLIMITATIONS:\nI found the limitations section to be adequate and clearly written.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-06-27"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nRebuttal 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nThe 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.\n\nQUESTIONS:\nIt 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)?\n\nLIMITATIONS:\nI 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.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-30"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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\nSTRENGTHS:\n- 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\nWEAKNESSES:\n- 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\nQUESTIONS:\nCold 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\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}

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

Note 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.

STRENGTHS:
- The topic is fascinating
- The paper is well written
- The potential impact of the results are promising

WEAKNESSES:
- The experimental setup is not fully convincing

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.

- As far as I could see, no model selection on the hyper-parameters is performed.

LIMITATIONS:
I think the authors correctly indicated a major potential limitation of the approach, namely its huge computational costs compared to literature alternatives.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-21