SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Table 4: Comparison of reconstruction loss (first row, the smaller the better) and top-1 accuracy of class and subclass prediction using MLP downstream classifier (second and third rows, the larger the better) achieved by different dynamics models  $f$  (linear, nonlinear, RNN, and transformer) in NeuPRINT. Two objectives are used to train each model (mean squared error (MSE) or negative log likelihood (NLL)).
Subclass	w/o population inputs	+ running speed	+ pupil size	+ frame state	+ population activity	+ center-surround activity
MLP	0.512	0.609	0.658	0.683	0.732	0.756
Table 5: Ablation studies of permutation-invariant inputs representing population activity, including running speed, pupil size, frame state, permutation-invariant population representation, permutation-invariant center-surround representation. Results are reported on the subclass prediction task with an MLP downstream classifier.
D Extensions to Other Visual Stimulus Conditions
We further test our model NeuPRINT and other baselines on recordings with 3 sets of visual stimuli (drifting gratings and two different natural scenes). The results in Table 6 suggest an increase in cell type identification accuracy from in-vivo activity as stimulus relevance and complexity increases.
E.1 NeuPRINT
Permutation-invariant summary of population dynamics: We use  $C$ ,  $S$  to represent  $D$ ,  $D$  dimensional permutation- and size-invariant (i.e.,  $N$ ) surrogates for brain activity at each time point.  $C$  and  $S$  compute identical statistics (mean and standard deviation) of population activity except that the statistics in  $C$  are calculated over the center partition (neurons whose distance to neuron  $i$  is at most  $\Delta$ ) and the statistics in  $S$  are calculated over the surround partition (neurons whose distance to neuron  $i$  is larger than  $\Delta$ ). Here,  $\Delta$  is a hyperparameter:
$$
\mu_ {C _ {t} ^ {(- i)}} = \operatorname {m e a n} (X _ {j} (t) \mid j: 0 <   d _ {j i} <   \Delta) \tag {4}
$$
$$
\sigma_ {C _ {t} ^ {(- i)}} = \operatorname {s t d} \left(X _ {j} (t) \mid j: 0 <   d _ {j i} <   \Delta\right) \tag {5}
$$
$$
\mu_ {S _ {t} ^ {(- i)}} = \operatorname {m e a n} \left(X _ {j} (t) \mid j: 0 <   \Delta \leq d _ {j i}\right) \tag {6}
$$
$$
\sigma_ {S _ {t} ^ {(- i)}} = \operatorname {s t d} \left(X _ {j} (t) \mid j: 0 <   \Delta \leq d _ {j i}\right) \tag {7}
$$
Batch sampling: For each batch of data we randomly sample 512 labeled and unlabeled neurons to be included in the batch. For each neuron we further randomly sample 2 sub-windows of 512 timesteps from its continuous calcium fluorescence traces. Each resulting sample for the  $i^{th}$  neuron is denoted  $X^{(i)}$ .
Multihead attention: We use the transformer encoder architecture [28] and the masking strategy as originally proposed in [47]. Random timesteps in  $X^{(i)}$  are masked (zero-out) with probability 0.25 and concatenated with the permutation-invariant population summary  $\bar{P}$  and time-invariant representation  $\phi^{(i)}$  along the feature dimension to form input  $\bar{X}^{(i)}$ . Note that the same learnable  $\phi^{(i)}$  is repeated at every timestep, enforcing time-invariance. Similar to [47], we embed input  $\bar{X}^{(i)}$  and
Task	Classifier	LOLCAT	Transformer+ISI	Transformer+Raw	NeuPRINT
Spontaneous Activity	Subclass	linear	0.474	0.491	0.474	0.683
mlp	—	0.561	0.439	0.756
Class	linear	0.608	0.680	0.669	0.793
mlp	—	0.640	0.664	0.807
Drifting Gratings	Subclass	linear	0.421	0.509	0.404	0.756
mlp	—	0.474	0.509	0.780
Class	linear	0.600	0.640	0.648	0.809
mlp	—	0.632	0.640	0.809
Natural Scenes 1	Subclass	linear	0.491	0.579	0.421	0.780
mlp	—	0.544	0.439	0.805
Class	linear	0.648	0.672	0.696	0.809
mlp	—	0.664	0.640	0.836
Natural Scenes 2	Subclass	linear	0.491	0.509	0.544	0.732
mlp	—	0.561	0.526	0.782
Class	linear	0.648	0.664	0.656	0.809
mlp	—	0.664	0.648	0.836
Table 6: Accuracy of transcriptomic subclass and class prediction of NeuPRINT and baselines on single-mouse recordings during spontaneous activity and visual stimuli-driven (drifting gratings, natural scenes) activity.
employ sinusoidal positional embedding to encode the temporal order in the input sequence, resulting in  $\tilde{X}^{(i)} = \mathrm{Emb}(\bar{X}^{(i)}) + \mathrm{E}$ .
For each input  $\tilde{X}^{(i)}$ , a set of weights  $W^{Q}\in \mathbb{R}^{T\times d_{q}}$ ,  $W^{K}\in \mathbb{R}^{T\times d_{k}}$ ,  $W^{V}\in \mathbb{R}^{T\times d_{v}}$  are learned to transform input  $\tilde{X}^{(i)}$  to a set of query, key, and value  $(Q,K,V)$ , where  $Q = \tilde{X}^{(i)}W^{Q}$ ,  $K = \tilde{X}^{(i)}W^{K}$ ,  $V = \tilde{X}^{(i)}W^{V}$ . Attention between temporal tokens for one attention head is computed as:
$$
\operatorname {A t t e n t i o n} (Q, K, V) = \operatorname {s o f t m a x} \left(\frac {Q K ^ {\top}}{\sqrt {d _ {k}}}\right) V \tag {8}
$$
Each head will find a different pattern in the data and produce an output of size  $d_v$ . The final attention output will be a concatenation of these single-head outputs. We use 2 heads in our model.
Feedforward layers and residual connections are subsequently applied to attention output:
$$
Z ^ {(i)} = \tilde {X} ^ {(i)} + \mathrm {M S A} \left(\tilde {X} ^ {(i)}\right) + \mathrm {F F} \left(\tilde {X} ^ {(i)} + \mathrm {M S A} \left(\tilde {X} ^ {(i)}\right)\right) \tag {9}
$$
where MSA represents the multihead attention operation, FF represents the feedforward layer with ReLU activation, and  $Z^{(i)}$  represents the reconstructed calcium trace with masked timesteps recovered.
Computational cost: NeuPRINT has 833K parameters and takes 2 hours in total for training and inference on a single NVIDIA Tesla V100 GPU. Details on the size of the datasets are mentioned in Section 4.1 of the main paper.
E.2 LOLCAT and Its Variants
We use the publicly available implementation of LOLCAT from [31]. We implement two additional variants of LOLCAT (Transformer+ISI and Transformer+Raw), following the philosophy in [31] but with the self-attention mechanism as used by NeuPRINT [28]. For LOLCAT and Transformer+ISI, continuous time series of calcium traces are divided into non-overlapping sub-windows of size 64, within which the Inter-Spike Interval (ISI) distribution is computed. We use suite2p Python package [48] to infer spikes and compute the ISI distribution with 16 bins, using a spike threshold of 0.2. To perform classification using the self-attention mechanism, we use a learnable special classification token [CLS] appended to the beginning of the input sequence to represent the classification output [47]. No positional embedding is added to the input sequence. Therefore the transformer output at the CLS token position will represent the pooling operation as in [31], and all sub-windows are treated as if they are independent trials, which enables Transformer+ISI to apply to any number of observed trials - an important design choice for LOLCAT.
We further implement Transformer+Raw, a variant of LOLCAT where the transformer operates directly on the raw calcium traces rather than the ISI distribution of calcium traces. Transformer+Raw
NeuPRINT	LOLCAT	Transformer +ISI	Transformer +Raw	random/ PCA/UMAP
Representation dim	64	—	—	—	64
Encoder hidden dim	70	[32, 16, 16]	[128] × 4	[128] × 4	—
MLP classifier hidden dim	2048	—	2048	2048	2048
Number of attention heads	2	4	4	4	—
Number of attention layers	1	1	1	1	—
Window size	200	2048	512	512	2048
Context window size	2	—	8	1	—
Batch size	1024	varies	varies	varies	—
Number of epochs	400	500	500	500	—
Number of downstream epochs	5000	—	—	—	1000
Learning rate	10-3	—	—	—	—
Downstream learning rate	10-4	10-4	10-4	10-4	10-3
Dropout	0.1	—	0.1	0.1	—
KNN neighbors	5	—	5	5	5
ISI sub-window size	—	64	64	—	—
Number of ISI bins	—	16	16	—	—
Table 7: Hyperparameters of our self-supervised representation framework NeuPRINT and other baselines including end-to-end supervised learning model LOLCAT, its variants Transformer+ISI and Transformer+Raw, unsupervised representation learning models PCA and UMAP, and chance-level prediction based on random features.
follows the same architecture as Transformer+ISI, except that now the positional embedding is added to the input sequence to denote the temporal relationship between timesteps in the trial window.
E.3 Random / PCA / UMAP
To probe the chance-level classification performance, we train and evaluate downstream classifiers using random vectors of size 64 as representations for individual neurons, equivalent to the 64-dim  $\phi^{(i)}$ . To compare NeuPRINT with unsupervised methods PCA and UMAP, we first project the data to a lower dimensional space using 64 components, then train and evaluate downstream classifiers on the low-dimensional representation.
E.4 Downstream Classifiers
We use 5 nearest neighbors for KNN downstream classifier. For the downstream MLP classifier, we use a multi-layer perceptron network with a single hidden layer of size 2048 and ReLU activation.
E.5 Hyperparameters
We include all of the important hyperparameters (representation dim, window size, number of epochs, learning rate, batch size, etc.) for our NeuPRINT and other models (supervised-learning baselines LOLCAT, Transformer+ISI and Transformer+Raw, unsupervised representation learning baselines UMAP and PCA, chance-level prediction based on random features) in Table 7.
F Pseudo Code
Our NeuPRINT framework includes three main components: an implicit dynamical system that uses the state-of-the-art transformer architecture to model neural dynamics; an optimization framework that fits the dynamical model and learns time-invariant representations for neurons; a supervised learning framework to train the downstream classifiers for subclass and class prediction, taking the learned time-invariant representations as inputs. The pseudo code for these three components is listed as follows:
Transformer(calciumfluoroscence_trace, permutation_invariant_populaton_summery, time_invariant Representation): input_embedding  $=$  linear(input_dim，hidden_dim) transformerEncoder  $=$  multihead attention( window_size , layer_dim , num_heads) outputDecoder  $=$  linear(hidden_dim，output_dim) masked_calciumfluoroscence_trace  $=$  mask輸puts(calciumfluoroscence_trace) input  $=$  concatenate( masked_calciumfluoroscence_trace, permutation_invariant_populaton_summery, time_invariantRepresentation) input  $=$  positional_encoding(input) input  $=$  input_embedding(input) context_mask  $=$  generate_context_mask( window_size , context_window_size) output  $=$  transformerEncoder(input，context_mask) output  $=$  outputDecoder(output) return output   
Time_Invariant_Self_Supervised_Representation_Learning( calciumfluoroscence_trace, permutation_invariant_populaton_summery): time_invariantRepresentation  $=$  zeros(neuron_dim，embedding_dim) recon_model  $\equiv$  Transformer( input_dim , hidden_dim , window_size , context_window_size , layer_dim , num_heads) optimizer  $=$  Adam( time_invariantRepresentation,recon_model, learning_rate) predicted_calciumfluoroscence_trace  $=$  recon_model( calciumfluoroscence_trace, permutation_invariant_populaton_summery, time_invariantRepresentation) recon_loss  $=$  mse( predicted_calciumfluoroscence_trace[masked_steps]， calciumfluoroscence_trace[masked_steps]) recon_loss .backward ()
optimizer step()
return time_invariant Representation
Downstream Classifier_Supervised_Learning(
    time_invariant Representation,
    groundtruth_class_label):
classifier = mlp(embeding_dim, hidden_dim, output_dim)
optimizer = Adam classifier, learning_rate)
predicted_class_label = classifier(time_invariant Representation)
classification_loss = cross_entropy(
    predicted_class_label,
    ground Truth class label)
classification_loss Backward()
optimizer step()
return predicted_class_label