SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\begin{array}{l} \nabla_ {\lambda_ {c t 1}} \mathcal {L} = \nabla_ {\lambda_ {c t 1}} \sum_ {k, i} r _ {i c t k} \cdot y _ {k} \left(\log \lambda_ {c t 1} - \log \Lambda_ {t 1}\right) (28) \\ = \frac {\sum_ {k , i} r _ {i c t k} \cdot y _ {k}}{\lambda_ {c t 1}} - \frac {\sum_ {k , i} \sum_ {c ^ {\prime} = 1} ^ {C} r _ {i c ^ {\prime} t k} \cdot y _ {k}}{\Lambda_ {t 1}} = 0 (29) \\ \Rightarrow \sum_ {k, i} r _ {i c t k} \cdot y _ {k} \left(\lambda_ {c t 1} + \sum_ {c ^ {\prime} \neq c} \lambda_ {c ^ {\prime} t 1}\right) = \sum_ {k, i} \sum_ {c ^ {\prime} = 1} ^ {C} r _ {i c ^ {\prime} t k} \cdot y _ {k} \cdot \lambda_ {c t 1} (30) \\ \Rightarrow \sum_ {k, i} y _ {k} \left(\sum_ {c ^ {\prime} \neq c} r _ {i c ^ {\prime} t k}\right) \lambda_ {c t 1} = \sum_ {k, i} y _ {k} \cdot r _ {i c t k} \left(\sum_ {c ^ {\prime} \neq c} \lambda_ {c ^ {\prime} t 1}\right) (31) \\ \Longrightarrow \lambda_ {c t 1} = \frac {\sum_ {k , i} y _ {k} \cdot r _ {i c t k} \left(\sum_ {c ^ {\prime} \neq c} \lambda_ {c ^ {\prime} t 1}\right)}{\sum_ {k , i} y _ {k} \left(\sum_ {c ^ {\prime} \neq c} r _ {i c ^ {\prime} t k}\right)} := \lambda_ {c t 1} ^ {*}. (32) \\ \end{array}
$$
$$
\nabla_ {\lambda_ {c t 0}} \mathcal {L} = 0 \Rightarrow \lambda_ {c t 0} = \frac {\sum_ {k , i} (1 - y _ {k}) r _ {i c t k} \left(\sum_ {c ^ {\prime} \neq c} \lambda_ {c ^ {\prime} t 0}\right)}{\sum_ {k , i} (1 - y _ {k}) \left(\sum_ {c ^ {\prime} \neq c} r _ {i c ^ {\prime} t k}\right)} := \lambda_ {c t 0} ^ {*}. \tag {33}
$$
Consider the gradient with respect to the  $\eta_c$  parameter,
$$
\begin{array}{l} \nabla_ {\eta_ {c}} \mathcal {L} = \sum_ {k, t, i} r _ {i c t k} \nabla_ {\eta_ {c}} \log \mathcal {N} \left(s _ {i t k}; \eta_ {c}\right) (34) \\ = \sum_ {k, t, i} r _ {i c t k} \nabla_ {\eta_ {c}} \frac {1}{2} \left(\log | \Sigma_ {c} ^ {- 1} | - \operatorname {T r} \left\{\Sigma_ {c} ^ {- 1} \left(s _ {i t k} - \mu_ {c}\right) \left(s _ {i t k} - \mu_ {c}\right) ^ {\top} \right\}\right). (35) \\ \end{array}
$$
The closed-form updates for  $\mu_c$  and  $\Sigma_c$  are
$$
\begin{array}{l} \nabla_ {\mu_ {c}} \mathcal {L} = \sum_ {k, t, i} r _ {i c t k} \Sigma_ {c} ^ {- 1} \left(s _ {i t k} - \mu_ {c}\right) = 0 (36) \\ \Rightarrow \mu_ {c} = \frac {1}{n _ {c}} \sum_ {k, t, i} r _ {i c t k} s _ {i t k} := \mu_ {c} ^ {*}, n _ {c} = \sum_ {k, t, i} r _ {i c t k}. (37) \\ \end{array}
$$
$$
\begin{array}{l} \nabla_ {\Sigma_ {c}} \mathcal {L} = \frac {1}{2} \sum_ {k, t, i} r _ {i c t k} \left(\Sigma_ {c} - \left(s _ {i t k} - \mu_ {c}\right) \left(s _ {i t k} - \mu_ {c}\right) ^ {\top}\right) = 0 (38) \\ \Rightarrow \Sigma_ {c} = \frac {1}{n _ {c}} \sum_ {k, t, i} r _ {i c t k} \left(s _ {i t k} - \mu_ {c}\right) \left(s _ {i t k} - \mu_ {c}\right) ^ {\top} := \Sigma_ {c} ^ {*}. (39) \\ \end{array}
$$
2.2 CAVI-based decoder
The CAVI-based decoder employs the same generative model as the CAVI-based encoder, with the exception that the behavior-dependent firing rates  $\lambda_{ct1}$  and  $\lambda_{ct0}$  are learned by the encoder and kept fixed, and the behavior  $y$  is treated as an unknown latent variable. We sample  $y$  from a prior distribution,  $y_{k} \sim \mathrm{Bernoulli}(\phi)$ , where  $\phi$  is a variational parameter that represents the probability that  $y_{k} = 1$ . The log-likelihood can be expressed as follows
$$
\begin{array}{l} \log p (s, z, y) = \sum_ {k, t, c, i} z _ {i c t k} \left\{\log \mathcal {N} \left(s _ {i t k}; \eta_ {c}\right) + y _ {k} \left(\log \lambda_ {c t 1} - \log \Lambda_ {t 1}\right) \right. \tag {40} \\ + (1 - y _ {k}) \left(\log \lambda_ {c t 0} - \log \Lambda_ {t 0}\right) \} + y _ {k} \log \phi + (1 - y _ {k}) \log (1 - \phi). \\ \end{array}
$$
We use the factorization  $q(z, y) = q(z)q(y) = \prod_{c,t} q(z_{ct})q(y)$  to approximate the posterior distribution  $p(z, y \mid s)$ . The ELBO of the decoder can be defined as
$$
\mathcal {L} _ {\mathrm {d e c}} ^ {\mathrm {C A V I}} := \mathbb {E} _ {q (z, y)} [ \log p (s, z, y) ] - \mathbb {E} _ {q (z, y)} [ \log q (z, y) ]. \tag {41}
$$
Algorithm 2 CAVI-based decoder
Input:  $\{s_{itk}\} ,\{\lambda_{ct0},\lambda_{ct1}\}$ $i = 1:n_{tk}$ $t = 1:T$ $k = 1:K$ $c = 1:C$  Initialize  $\{\mu_c\} ,\{\Sigma_c\}$  while ELBO not converged do for all  $k\in 1:K$  do Set  $q(y_{k} = 1)\propto y_{k1}^{*}$  .  $\triangleright$  eq.44) for all  $t\in 1:T$  do for all  $i\in 1:n_{tk}$  do Set  $q(z_{ictk} = 1)\propto \rho_{ictk}$  end for end for end for for all  $c\in 1:C$  do Set  $\mu_c = \mu_c^*$ $\Sigma_{c} = \Sigma_{c}^{*}$  end for Set  $\phi = \phi^{*}$  .  $\triangleright$  eq.46) Compute the ELBO  $\mathcal{L}_{\mathrm{dec}}^{\mathrm{CAVI}}$  end while Return  $q(z,y),\phi ,\mu ,\Sigma$
The exact updates for  $q(z)$  and  $q(y)$  that guarantee an increase in the ELBO are
$$
q (z) \propto \exp \left\{\mathbb {E} _ {q (y)} [ \log p (s, z, y) ] \right\}, \quad q (y) \propto \exp \left\{\mathbb {E} _ {q (z)} [ \log p (s, z, y) ] \right\}, \tag {42}
$$
$$
q \left(z _ {i c t k} = 1\right) \propto \exp \left\{\log \mathcal {N} \left(s _ {i t k}; \eta_ {c}\right) + \mathbb {E} \left[ y _ {k} \right] \log \left(\frac {\lambda_ {c t 1}}{\Lambda_ {t 1}}\right) + \left(1 - \mathbb {E} \left[ y _ {k} \right]\right) \log \left(\frac {\lambda_ {c t 0}}{\Lambda_ {t 0}}\right) \right\} := \rho_ {i c t k}, \tag {43}
$$
and  $\mathbb{E}[z_{ictk}] = \rho_{ictk} / \sum_{c'}\rho_{ic'tk} \coloneqq r_{ictk}$  are the unnormalized and normalized posterior probabilities of assigning spike  $i$  collected at time  $t$  in trial  $k$  to component  $c$ , respectively. The unnormalized posterior probabilities of  $y_{k} = 1$  and  $y_{k} = 0$  are
$$
q \left(y _ {k} = 1\right) \propto \exp \left\{\sum_ {t, c, i} \mathbb {E} \left[ z _ {i c t k} \right] \left(\log \lambda_ {c t 1} - \log \Lambda_ {t 1}\right) + \log \phi \right\} := y _ {k 1} ^ {*}, \tag {44}
$$
$$
q (y _ {k} = 0) \propto \exp \left\{\sum_ {t, c, i} \mathbb {E} [ z _ {i c t k} ] (\log \lambda_ {c t 0} - \log \Lambda_ {t 0}) + \log (1 - \phi) \right\} := y _ {k 0} ^ {*},
$$
and  $\mathbb{E}[y_k] = y_{k1}^* / (y_{k1}^* + y_{k0}^*) \coloneqq \nu_k$  represents the normalized posterior probability of  $y_k = 1$ .
After fixing  $q(z)$  and  $q(y)$ , the term in the ELBO which depends on  $\phi, \mu$  and  $\Sigma$  can be written as
$$
\begin{array}{l} \mathscr {L} ^ {\prime} := \sum_ {k, t, c, i} r _ {i c t k} \left\{\log \mathcal {N} \left(s _ {i t k}; \eta_ {c}\right) + \nu_ {k} \left(\log \lambda_ {c t 1} - \log \Lambda_ {t 1}\right) \right. \tag {45} \\ + (1 - \nu_ {k}) (\log \lambda_ {c t 0} - \log \Lambda_ {t 0}) \} + \nu_ {k} \log \phi + (1 - \nu_ {k}) \log (1 - \phi). \\ \end{array}
$$
Considering the gradient of the ELBO with respect to the  $\phi$  parameter, we obtain its update:
$$
\nabla_ {\phi} \mathcal {L} ^ {\prime} = 0 \Rightarrow \phi^ {*} = \frac {1}{K} \sum_ {k = 1} ^ {K} \nu_ {k}. \tag {46}
$$
The updates for  $\eta_c$  are computed in a similar manner as described in Equations 36-39.
3 MoG initialization
We employ the following procedure to initialize the MoG model used in both the ADVI-based and CAVI-based algorithms: 1) According to Figure 1 (a), the spike feature distribution is highly multimodal. To determine the appropriate number of modes, we utilize isosplit (Magland and Barnett, 2015) to cluster the spike features. This step helps in splitting the set of spike features into distinct clusters. 2) For each identified cluster, we compute the mean and variance of the spike features belonging to that cluster, which serve as the parameters for the corresponding Gaussian component. This automatic selection of the number of MoG components and the initialization of means and covariance matrices facilitate the initialization of the ADVI-based and CAVI-based algorithms.
4 Data preprocessing
We provide a brief overview of our data preprocessing pipeline, which involves the following steps.
Destriping. During the data collection process, we encounter line noise due to voltage leakage on the probe. This translates into large "stripes" of noise spanning the whole probe. To mitigate the impact of these noise artifacts, we apply a destriping procedure (Chapuis et al., 2022).
Subtraction-based spike detection and denoising. We employ the iterative subtraction-based procedure for spike detection and collision-correction described in Boussard et al. (2023).
Spike localization. We employ the method of Boussard et al. (2021) to estimate the location of each denoised spike.
Drift registration. Probe motion (or drift) in the electrophysiology data poses a challenge for downstream analyses. Decentralized registration (Windolf et al., 2022) is applied to track and correct for motion drift in the high-density probe recordings.
Figure 6: Motion drift in "good" and "bad" sorting recordings. (a) The motion-registered spike raster plot of a "good" sorting example that is less affected by drift. (b) The spike raster plot of a "bad" sorting example, which is still affected by drift even after registration.
5 Behavior decoder
To decode binary behaviors, such as the mouse's left or right choices, we utilize  $\mathcal{L}_2$ -penalized logistic regression. For decoding dynamic behaviors, such as wheel speed, we employ a sliding-window algorithm to aggregate the entries of the weight matrix,  $W_{kct}$ , over time. Within the time window  $[t - \delta, t + \delta]$ , where  $\delta$  is the window size, we stack  $2\delta$  weight matrix entries,  $\{W_{kct}\}_{c=1}^{C}$ , for time point  $t$  in trial  $k$ . This aggregated weight matrix is then used as input for ridge regression to predict the behavior  $y_{tk}$  at time  $t$ . The window size  $\delta$  and the regularization strength are model hyper-parameters, set through cross-validation to achieve the optimal decoding performance.
6 Model interpretation
In this section, we provide visualizations to gain insights into the effectiveness of our proposed decoder. We quantify the posterior entropy of each spike assignment in Figure 7 (a). Spike assignments with higher entropy correspond to a spread of posterior probabilities among multiple MoG components. In contrast, traditional spike sorting or thresholding methods result in deterministic spike assignments, leading to lower entropy and empirically reduced decoding performance. In Figure 7 (b), we compare the trial-averaged weight matrices  $(W)$  used for decoding between spike-sorted, spike-thresholded, and our proposed decoders.
Figure 7: Model interpretation. (a) The posterior entropy of the spike assignment is high, when the posterior probability of spike assignment,  $q(z_{ikct})$ , is spread out among several MoG components instead of being concentrated at a single component. The scatter plot shows that low-amplitude spikes that are difficult to assign have higher posterior entropy than high-amplitude spikes. (b) Visualizations of the averaged weight matrices  $W$ 's across trials in an example IBL session. For "W (spike-sorted and spike-thresholded)", the  $W$  matrix has one-hot rows and each entry  $W_{kct}$  is the spike count that belongs to KS unit (channel)  $c$  and time bin  $t$  in trial  $k$ . The purple crosses indicate the "good" KS units. For "W (ours)",  $W_{kct}$  is the sum of posterior probabilities of spike assignments, and the MoG mixing proportion  $\pi$  depends on the behavior  $y$  and changes over time. The arrangement of KS units (channels or MoG components) on the heat maps is based on the depth of the NP probe, ensuring comparability across the displayed  $W$  matrices.
7 Decoding across brain regions
In addition to the previously mentioned five brain regions (PO, LP, DG, CA1, VISa) depicted in Figure 3, we expanded our analysis to include two additional brain regions situated in the cerebellum: the arbor vitae (ARB) and the ansiform cruciform lobule (ANCR). We specifically include the cerebellum in our analysis due to the frequent occurrence of spike sorting issues in this area. In Figure 8 and 9, we present a comparison between our decoder and the spike-sorted decoder that utilizes all KS units across all the brain regions studied. According to Figure 8, our method consistently outperforms the decoder that relies on all KS units across all brain regions and the majority of IBL sessions. Furthermore, Figure 9 specifically demonstrates that our method consistently achieves superior decoding performance in the recorded regions of the cerebellum, where spike sorting quality issues are commonly encountered. This highlights the robustness and reliability of our method, particularly in challenging recording conditions.
Figure 8: Decoding comparisons across brain regions in the thalamus, hippocampus and visual cortex. We compare our decoder to the spike-sorted decoder using all KS units across various brain regions and behavioral tasks. Each point in the scatter plot represents one session from the 20 IBL session previously described in the experiments section. Sessions with "good" sorting quality are depicted in green, while sessions with "bad" sorting quality are shown in purple. The majority of the sessions lie above the gray diagonal line, indicating that our method consistently outperforms the decoder relying on all KS units.
Figure 9: Decoding comparisons across brain regions in the cerebellum. We evaluate our decoder against the spike-sorted decoder utilizing all KS units across various brain regions in the cerebellum. The scatter plot visualizes the results for each IBL session in the study, with all sessions depicted in purple indicating "bad" sorting quality. Note that the probes used in the cerebellum sessions were not uniformly implanted in the same set of brain regions. As a result, the number of available sessions for decoding in certain cerebellum regions is limited. Notably, the majority of sessions are positioned above the gray diagonal line, indicating that our method consistently achieves better decoding performance compared to the decoder relying on all KS units. This outcome highlights the robustness of our method, particularly in the context of the cerebellum where spike sorting quality issues are prevalent.
8 Ablation study
We conduct an ablation study to investigate the importance of various components in our decoding paradigm. Specifically, we examine how the integration of dynamic (behavior-dependent) mixing proportions in the MoG and higher-dimensional spike features influence decoding performance. Additionally, we analyze how the inclusion criteria for spike waveforms can affect decoding outcomes.
Effects of dynamic MoG mixing proportion. Table 3 presents a comparison between the ordinary MoG with a fixed mixing proportion (referred to as "fixed  $\pi$ ) and our proposed model with a dynamic mixing proportion (referred to as "dynamic  $\pi$ ). The results indicate that this approach leads to improved decoding performance compared to using a fixed mixing proportion in the MoG model.
Motion energy (R2)	Wheel speed (R2)	Choice (accuracy)
Fixed π	Dynamic π	Fixed π	Dynamic π	Fixed π	Dynamic π
All	0.664 (± 0.034)	0.742 (± 0.028)	0.470 (± 0.062)	0.564 (± 0.045)	0.948 (± 0.038)	0.957 (± 0.036)
PO	0.365 (± 0.031)	0.488 (± 0.046)	0.520 (± 0.019)	0.670 (± 0.015)	0.844 (± 0.016)	0.861 (± 0.035)
LP	0.145 (± 0.015)	0.464 (± 0.054)	0.114 (± 0.032)	0.342 (± 0.027)	0.917 (± 0.026)	0.931 (± 0.022)
DG	0.280 (± 0.033)	0.492 (± 0.042)	0.221 (± 0.040)	0.381 (± 0.035)	0.669 (± 0.084)	0.722 (± 0.050)
CA1	0.407 (± 0.021)	0.538 (± 0.038)	0.308 (± 0.051)	0.428 (± 0.030)	0.621 (± 0.062)	0.626 (± 0.064)
VISA	0.488 (± 0.046)	0.490 (± 0.047)	0.318 (± 0.056)	0.364 (± 0.074)	0.857 (± 0.065)	0.874 (± 0.066)
Effects of higher-dimensional spike features. In Table 4, we provide a comparison of decoding performance using two different sets of spike features. The first set includes spike location along the width and depth dimensions of the probe (denoted as  $x$  and  $z$ ) as well as the maximum peak-to-peak amplitude of the spike (denoted as  $a$ ). The second set includes the first and second principal components (PCs) of the spike waveforms (denoted as  $u_{1}$  and  $u_{2}$ ) in addition to  $x$ ,  $z$  and  $a$ . The spike features are visually represented using scatter plots in Figure 10. We report the mean and standard deviation of the decoding accuracy ( $R^2$ ) obtained from a 5-fold CV on a single session. Table 4 provides insights regarding the inclusion of additional waveform PC features for decoding. The findings suggest that the incorporation of these additional PC features does not contribute to significant improvements in decoding performance.
Table 3: Effects of dynamic mixing proportion of the MoG on decoding performance.
Motion energy (R2)	Wheel speed (R2)	Choice (accuracy)
(x,z,a)	(x,z,a,u1,u2)	(x,z,a)	(x,z,a,u1,u2)	(x,z,a)	(x,z,a,u1,u2)
All	0.531 (±0.026)	0.529 (±0.026)	0.484 (±0.042)	0.478 (±0.045)	0.917 (±0.019)	0.917 (±0.019)
PO	0.462 (±0.038)	0.457 (±0.039)	0.469 (±0.061)	0.464 (±0.051)	0.853 (±0.019)	0.853 (±0.025)
LP	0.490 (±0.026)	0.489 (±0.029)	0.479 (±0.035)	0.473 (±0.013)	0.864 (±0.022)	0.849 (±0.036)
DG	0.335 (±0.028)	0.321 (±0.037)	0.273 (±0.030)	0.264 (±0.025)	0.675 (±0.038)	0.679 (±0.049)
CA1	0.449 (±0.027)	0.440 (±0.046)	0.329 (±0.044)	0.328 (±0.040)	0.755 (±0.045)	0.758 (±0.053)
VISA	0.270 (±0.021)	0.237 (±0.017)	0.225 (±0.024)	0.206 (±0.013)	0.725 (±0.048)	0.732 (±0.054)
Table 4: Effects of incorporating higher-dimensional spike features on decoding performance.
Figure 10: Visualizations of spike features employed for decoding. Spike localization features,  $(x,z)$ , the locations of spikes along the width and depth of the NP1 probe, and waveform features,  $a$ , the maximum peak-to-peak (max ptp) amplitudes of spikes. Amplitude is measured in standard units (s.u.).  $u_{1}$  and  $u_{2}$  denote the first and second principal components (PCs) of the spike waveforms.
Effects of inclusion criteria for spike waveforms. To investigate whether the density-based decoder is performing better by using additional spikes that a spike sorter would miss, we conducted an experiment using different inclusion criteria for spike waveforms. We fitted our model using only spikes detected by Kilosort 2.5, and compared its performance to decoders using spike-sorted outputs and our subtraction-based spike detection on choice and motion energy decoding. The results are summarized in Table 5. As shown in the table, our decoder can achieve comparable or better decoding performance than the spike-sorted decoder when modeling the same spikes. This suggests that the gain in decoding performance can be attributed to the density-based approach, instead of the spike inclusion criteria.
Choice (accuracy)	Motion energy (R2)
Density-based (subtraction spikes)	0.876 (± 0.068)	0.589 (± 0.111)
Density-based (KS spikes)	0.876 (± 0.079)	0.579 (± 0.121)
Sorted (KS spikes)	0.887 (± 0.078)	0.503 (± 0.117)
Table 5: Effects of inclusion criteria for spike waveforms on decoding performance.
9 Comparison to a state-of-the-art clusterless decoder
In this section, we outline the specific experimental setup for evaluating our density-based decoder in comparison with the clusterless point process decoder (Denovellis et al., 2021) on both multiple tetrodes and high-density (HD) probes data.
Application to tetrodes. We utilized the code provided in the GitHub repository<sup>1</sup> of the clusterless point process decoder (Denovellis et al., 2021) to generate simulated neural and behavioral data. This synthetic dataset was designed to mimic recordings from 5 tetrodes, with each tetrode containing 4 channels. Spike amplitudes from each channel of these multiple tetrodes were selected as the spike features for both decoding methods. The objective of the decoding was to estimate the animal's position from these simulated spike features. As the original simulated position was too straightforward to decode, we intentionally distorted it by blending it with real position data sourced from the GitHub repository<sup>2</sup>. Moreover, we introduced random Gaussian noise to the simulated position for added complexity.
Application to HD probes. We evaluated both decoders on decoding wheel speed and motion energy from spike features extracted from NP1 probes across three IBL datasets. To preprocess the data, we followed the pipeline outlined above, extracting spike localization features and maximum peak-to-peak amplitudes as common features for both decoders.
We only focus on continuous behaviors for decoding since the clusterless point process decoder is designed exclusively for continuous behaviors. Notably, due to its continuous time nature, the clusterless point process decoder directly decoded behaviors without time binning. In contrast, our density-based model required time binning of both behavioral and spike feature data into equal-sized time intervals. Consequently, we used the time-binned behaviors for decoding with the density-based approach. We employed 5-fold cross-validation to assess the decoding performance of both decoders. We used a random walk as the state transition for the clusterless point process decoder, and used the estimated variance from the animal's behavior to set the variance of the random walk. The clusterless point process decoder uses a grid-based approximation of the inferred behavior, discretizing the behavior space into place bins; see Section "Decoding" in Denovellis et al. (2021). We determine place bin size based on the square root of the observed behavior variance. Denovellis et al. (2021) use kernel density estimation (KDE) to estimate the distributions of both the behavior variable and the spike features used for decoding. We set the KDE bandwidth that determines the amount of smoothing done for spike features to be 1.0, and the bandwidth for behavior to be the square root of the observed behavior variance; see Section "Encoding - clusterless" in Denovellis et al. (2021).
<sup>1</sup>https://github.com/Eden-Kramer-Lab/replay_trjectory_classification
$^{2}$ https://github.com/nelpy/example-analyses/blob/master/LinearTrackDemo.ipynb
10 Simulation for model validation
We conducted simulations to illustrate the principles of our method. The simulation aimed to show that our encoding model can learn the relationship between spike features and behaviors. We performed two tasks, decoding a binary variable,  $y_{k}$ , simulated from a Bernoulli distribution, and decoding a continuous variable,  $y_{k}$ , simulated from a Gaussian process. To mimic the data-generating process, we selected Gaussian components with "templates" extracted from a real dataset. The encoding model parameters,  $b$  and  $\beta$ , were also taken from learned parameters in the same dataset. Given  $b$ ,  $\beta$  and  $y_{k}$ , we simulated the "firing rates"  $\lambda$  for each Gaussian component in the mixture, as described in the Method section of our paper. Next, we generated spike features based on these simulated "firing rates," and applied the encoding model to infer the behavior-dependent  $\lambda$ . Figure 11 displays the learned  $\lambda$  for each component  $c$ , time  $t$ , and trial  $k$ . The learned "firing rates" closely resembled the simulated ones, indicating the model's ability to recover the primary associations between spike features and behaviors. With such associations, the decoding model can decode behaviors.
Figure 11: Our encoding model recovers the relationship between the simulated spiking activity and the simulated behavior correlate. Panel (a) shows a comparison of the simulated firing rates conditioned on the binary behavior variable with the learned firing rates by our encoding model. In Panel (b), we compare the simulated firing rates conditioned on the continuous behavior variable with the learned firing rates from our encoding model.
11 Relationship to spike sorting