2.1 Preliminaries

A point process is one of the stochastic processes used to model a sequence of events. Marked point processes are a subclass of point processes that generate events with feature values called marks. When employing marked point processes for statistical modeling, the goal is to estimate the conditional intensity function that describes the probability structure of the process. This function determines the instantaneous probability of observing an event with a particular mark. The definition of the conditional intensity function is as follows [36]: Here, is a mark, N(⋅) is a counting measure defined on the product space of time and mark, is a state that influences the occurrence of events, is an open ball of radius centered at κ, and |B(κ, Δκ)| is the Lebesgue measure of B(κ, Δκ). Although the general definition of this function depends on the history of the process, we do not consider such dependency in this paper.

To ensure estimation efficiency, we focus on a specific class of marked point processes. We assume that there is a state that changes with time, x_t, and events are generated from a marked inhomogeneous Poisson process whose intensity is determined by x_t. Under this assumption, the conditional intensity function can be written as where λ(x, κ) is a nonnegative function defined in the space . As mentioned in Section 1, we refer to this function as the “joint mark intensity.” This function explains how event occurrences depend on the state. More specifically, the value λ(x, κ)Δt is the expected number of events with mark κ occurring in an interval of length Δt when the state takes the value x.

Suppose that events are generated from this marked point process within an interval (0, T]. Then, the log-likelihood is given by In this expression, we need the continuous-time path of x_t to evaluate the exact log-likelihood. In practice, the values of x_t are available only at discrete time points t_j ≔ jΔt, j = 1, …, m, where Δt = T/m. Let be a state value at a spike time and be a state value at discrete time points. Note that we do not know the values of x_i since x_t is only measured at discrete time points . A straightforward approach to obtaining x_i is to use nearest-neighbor interpolation from the available measurements . Under this assumption, the log-likelihood is approximated as (1) In this approximation, the integral in the log-likelihood is replaced by Riemann summation. By maximizing this log-likelihood within a set of functions of λ(x, κ), we can estimate the joint mark intensity given the observations. For notational simplicity, we will omit Δt in this log-likelihood without any loss of generality. For justification, see Section S1.1 in S1 Text.

2.2 Joint mark intensity model

In this section and the one that follows, we assume that there is an observable state x_t and that this state affects the occurrence of unsorted spikes. An example of x_t corresponds to external stimuli presented to the animal during the experiment. The goal is to estimate the joint mark intensity λ(x, κ) from the state values at spike times , at discrete time points , and the unsorted spikes .

A straightforward way to express the joint mark intensity λ(x, κ) in a high-dimensional space is to use a neural network that takes pairs of x and κ as inputs and outputs nonnegative values. However, expressing λ(x, κ) as a neural network makes integration with respect to κ in the log-likelihood (1) challenging.

To address this issue, we propose a new joint mark intensity model based on neural networks. We first decompose the joint mark intensity λ(x, κ) into a base intensity term and a joint probability density of the state and mark, and define the latter density function by a variational autoencoder: where is a scale parameter, z is a low-dimensional representation of the state and mark, p(x, κ ∣ z) is a conditional density referred to as a decoder, and p(z) is a prior. We define p(x, κ ∣ z) as a distribution family whose parameters are neural networks with z as input.

Substituting this definition into the log-likelihood (1), we obtain (2) where p(x ∣ z) is the marginal density of x given z:

This model formulation has several advantages over existing models. First, it can represent highly complex data structures due to its use of neural networks. Second, since the integral with respect to κ in the log-likelihood becomes a closed form with our model, the log-likelihood can be calculated without heavy numerical integration when dim(z) ≪ dim(κ).

This log-likelihood still contains integration terms with respect to the latent variable z, and a simple Monte Carlo approximation of these terms results in large variance in gradient estimates [23]. To reduce this variance, we introduce two encoders, q(z ∣ x, κ) and q(z ∣ x), and derive the lower bound of this log-likelihood as explained in the next section.

2.3 Lower bound for point process log-likelihood

To obtain the lower bound for the point process log-likelihood (2), we provide the lower bound for and the upper bound for The main idea throughout this section is to rewrite the above evidence as a function of the divergence term and the evidence bound term. Based on the rewritten expression, minimizing divergence has the effect of tightening the evidence bound to the evidence. For step-by-step derivations of the bounds, see Section S1.2 in S1 Text. The lower bound introduced in this section can be generalized under a subclass of the f-divergence family [28]. For the details, see Section S1.3 in S1 Text.

First, we derive the lower bound for the log evidence log p(x, κ). This type of lower bound has been widely used in variational inference, e.g., variational autoencoders [23]. Here, we present the Rényi lower bound [24] as an example. The Rényi lower bound is an extension of the evidence lower bound based on the Kullback-Leibler divergence. Although this lower bound was originally derived from the Rényi divergence in [24], we use α-divergence [37] instead to ensure consistency with subsequent discussions. The definition of the α-divergence with a parameter α ∈ (0, 1) is Using this divergence, we obtain the following inequality: Here, p(z ∣ x, κ) is the posterior of z given x and κ, and q(z ∣ x, κ) is an encoder. The inequality is derived from the nonnegativity of the α-divergence.

Maximizing this lower bound with respect to the encoder reduces to a minimization of Note that equality holds when q(z ∣ x, κ) = p(z ∣ x, κ). Taking the logarithm of both sides of this inequality yields the Rényi lower bound for the log evidence: The parameter α controls the tightness of the bound. When α = 1, the bound is tightest and is equal to the objective in the importance-weighted autoencoder [38]. When α → 0, the bound reduces to the Kullback-Leibler lower bound used in the original variational autoencoder [23].

Next, we derive the upper bound for the evidence p(x). In the context of model learning, the evidence upper bound [24, 25] is less frequently used than the evidence lower bound. This is because the sign of the evidence in well-known objectives is inconsistent with upper bound minimization. In other words, the evidence should be larger when the model is better, but upper bound minimization results in smaller evidence. To use upper bound minimization for VAE learning, χ-VAE [39] adopts a two-step procedure for evidence maximization. The first step minimizes the upper bound of the evidence with respect to the encoder while fixing the decoder. The second step maximizes the evidence with respect to the decoder while fixing the encoder. This alternating optimization aims to find saddle points and tends to be unstable in practice. In contrast, maximizing the evidence lower bound achieves a better decoder and encoder simultaneously.

An illustrative application of evidence upper bound minimization for model learning can be found in undirected graphical model learning [40]. To highlight the connection to our model, we review the content of this paper. An undirected graphical model, such as a Markov random field, is defined by its unnormalized density, known as the energy function. Let be this unnormalized density governed by parameter θ. Then, the log-likelihood given the observations is The integral in the second term is the partition function and is not analytically tractable in practice. To handle this integral, the authors of the paper introduced a proposal distribution q(x, ϕ) and derived the lower bound using the χ upper bound as The parameters θ and ϕ are simultaneously estimated by maximizing the Monte Carlo estimate of this bound. In this case, the sign of the χ upper bound term is negative in the lower bound. Thus, upper bound minimization is consistent with likelihood maximization.

Since the point process log-likelihood contains the negative evidence −p(x), the minimization of the evidence upper bound is consistent with the maximization of the overall log-likelihood. We adopt the χ upper bound [25] derived from the χ-divergence to bound this evidence. The definition of the χ-divergence with a parameter β > 1 is Based on this divergence, we obtain the following inequality: Here, p(z ∣ x) is the posterior of z given x, and q(z ∣ x) is another encoder. The inequality is derived from the nonnegativity of the χ-divergence. This inequality is also derived from Jensen’s inequality.

The minimization of this upper bound with respect to the encoder reduces to the minimization of Equality holds when q(z ∣ x) = p(z ∣ x). The parameter β controls how the encoder approximates the true posterior. When β = 2, the encoder that minimizes this upper bound also minimizes the variance of the Monte Carlo estimate of the evidence p(x) where are i.i.d. samples from q(z ∣ x).

Substituting these two bounds gives the lower bound of the log-likelihood (2): In this equation, we move λ₀ into the inner brackets. In the training step, the expectations with respect to the encoders are replaced by Monte Carlo estimates using the reparameterization trick: (3) where and are i.i.d. samples from q(z ∣ x_i, κ_i) and q(z ∣ x_j), respectively.

We define the decoder and encoders as distribution families whose parameters are neural networks. Specific choices are Gaussian distributions whose mean vectors and precision matrices are modeled by neural networks: Here, N(⋅ ∣ μ, Λ) is a multivariate Gaussian density function with mean μ and precision matrix Λ. The neural networks μ_θ(⋅), ν_ϕ(⋅) return mean vectors, and Λ_θ(⋅), Ξ_ϕ(⋅) return precision matrices. For notational clarity, we indicate the input and output domain of each neural network in the superscript of the variable. For example, the neural network receives z and returns the mean vector defined on the space of x. In these networks, θ denotes the parameters of the decoders and ϕ denotes the parameters of the encoders.

To achieve estimation efficiency, we define the joint density p(x, κ ∣ z) as the product of two decoders such that x and κ are conditionally independent given z. This definition enables us to easily calculate p(x ∣ z) in the χ upper bound. In the above Gaussian case, p(x ∣ z) simply reduces to .

The encoders q(z ∣ x, κ) and q(z ∣ x) approximate the posteriors p(z ∣ x, κ) and p(z ∣ x), respectively, and thus the optimal q(z ∣ x) is determined by the optimal q(z ∣ x, κ) as To simplify parameter estimation, we use different neural networks for these encoders. Even in this approach, we do not need the preceding constraint; if the encoders q(z ∣ x, κ) and q(z ∣ x) are optimal, then this constraint automatically holds. On the other hand, developing encoders that satisfy this constraint reduces the model candidates and may help estimation efficiency and accuracy. Finding appropriate encoders for this purpose is our future work.

To estimate the parameters, we maximize the Monte Carlo estimate of the lower bound (3) with respect to the decoder parameter θ, the encoder parameter ϕ, and the scale parameter λ₀ using a stochastic gradient ascent algorithm. However, since the 1/β-th power function in the χ upper bound term is concave, this Monte Carlo estimate is a biased estimator and its expectation is not guaranteed to be a lower bound for the log-likelihood (2). If L → ∞, this Monte Carlo estimate converges to the true value [24]. When L is not sufficiently large, this bias may make parameter estimation unstable: the Monte Carlo estimate sometimes diverges to zero.

To address this issue, we utilize jackknife resampling [41] and doubly reparameterized gradients [42]. The jackknife estimator is a leave-one-out resampling method that eliminates the O(L⁻¹) bias. This estimator improves the order of convergence from O(L⁻¹) to O(L⁻²); however, it increases the variance [41]. The doubly reparameterized gradients reduce this variance. For the details, see Sections S1.4 and S1.5 in S1 Text.

2.4 State space representation

In Sections 2.2 and 2.3, we discussed joint mark intensity estimation given the observable state and unsorted spikes. Hereafter, we assume that a hidden state with some dynamics exists behind this data, and that the observed covariate and unsorted spikes are generated from distributions determined by this hidden state. We denote the observed covariate as y_t and the hidden state as x_t. In neural population modeling, y_t includes all observable values other than unsorted spikes during experiments, and x_t includes the values not directly observed through the experiments.

With this assumption, the model we aim to estimate is divided into three parts. The first part is the dynamics model of the hidden state. The second part is the observation model of the observed covariate given the hidden state. The last part is the marked point process model for unsorted spikes determined by the hidden state.

To consider hidden state dynamics in the discrete-time domain, we split the time axis into R equal-length bins and let x_r, r = 1, …, R be the hidden states and y_r, r = 1, …, R be the observed covariate values for each bin. We put n_r spikes within the r-th bin together and denote their marks as κ_ri, i = 1, …, n_r. In the place cell example, x_r is the hidden representation reflecting detailed information about the rat, y_r is the coordinate of the rat’s location, and κ_ri is a spike feature of the i-th spike occurring in the r-th bin. Hereafter, we denote , , and . Some bins may not contain any spikes, and so we define κ_r = ∅ for such bins. Note that we only observe y and κ, and we do not have access to information about x. Thus, the dimension of x_r is an unknown hyperparameter and needs to be decided beforehand.

This state space representation is similar to the black-box state space model used in [32]. The main difference between our model and this work lies in the observation model. In [32], observations were the number of spikes, and the observation model was a Poisson distribution. In our model, the observations consist of the covariate and unsorted spikes, and the observation model for the latter is the marked point process.

We assume that x, y, κ follow the state space model: where p(x₁) and p(x_r ∣ x_r−1) are the initial distribution and transition distribution for x_r, and p(y_r ∣ x_r) and p(κ_r ∣ x_r) are the observation distributions for y_r and κ_r, respectively. Fig 2 shows an illustration of this state space model. For estimation simplicity, we define these distributions as Gaussian distributions with neural networks [34]: where a₁ is the initial mean, V₁ is the initial precision matrix, F is the transition matrix, a is the transition offset, V is the transition precision matrix, and are the mean vector and precision matrix expressed as neural networks. For κ_r, we define the observation distribution as the point process, as described in the preceding sections:

Download:

• PPT
  PowerPoint slide
• PNG
  larger image
• TIFF
  original image

Fig 2. Illustration of our state space representation.

Circles indicate the variables, and lines connecting these circles indicate the probabilistic dependencies. The variables surrounded by solid circles are observed variables in the training step and evaluation tasks. We assume that the hidden state x_t influences the observed covariate y_t and unsorted spikes {(t_i, κ_i)}. To consider the state dynamics in the discrete domain, we divide the time axis into R bins of equal length and define x_r, r = 1, …, R as the hidden state and y_r, r = 1, …, R as the observed covariate for each bin. We group n_r spikes within the r-th bin and denote their marks as . The hidden state x_r follows linear dynamics, and the observed covariate y_r depends on this state through a nonlinear embedding. The unsorted spikes κ_r follow the marked point process defined by our joint mark intensity model. (A) In the training step, we estimate the parameters of our model by maximizing the lower bound of the log-likelihood given y and κ. With the estimated parameters, we evaluate our model in prediction and decoding tasks. (B) In the prediction task, we evaluate the negative log-likelihood (NLL) of y and κ in the test part of the data. (C) In the decoding task, we reconstruct the covariate y from unsorted spikes κ and calculate the mean squared error (MSE) between the reconstructed covariate and the true one.

https://doi.org/10.1371/journal.pcbi.1012620.g002

Under this assumption, the marginal log-likelihood cannot be expressed in closed form. Instead, we maximize the evidence lower bound based on Kullback-Leibler divergence by utilizing an encoder q(x ∣ y, κ): (4) This encoder q(x ∣ y, κ) should be easy to sample from and should allow for straightforward calculation of probabilities for generated samples. For this purpose, we adopt a structured variational family for this encoder [32, 34, 35]: Here, q(x_r ∣ y_r) and q(x_r ∣ κ_r) are encoders that provide a probabilistic guess at each hidden state x_r inferred from observations y_r and κ_r. The distributions q(x₁) and q(x_r ∣ x_r−1) connect neighboring nodes in a Markov manner.

To achieve sampling efficiency, we define these densities as Gaussian distributions. First, we define the distributions q(x₁) and q(x_r ∣ x_r−1) as where b₁ is the initial mean, W₁ is the initial precision matrix, G is the transition matrix, b is the transition offset, and W is the transition precision matrix.

The encoder q(x_r ∣ y_r) approximates the likelihood of x_r given y_r. We define this part as a Gaussian distribution with neural networks [34]: where are the mean vector and precision matrix expressed as neural networks, respectively.

The encoder q(x_r ∣ κ_r) approximates the point process likelihood. The point process likelihood is the product of n_r joint mark intensity terms λ(x_r, κ_ri) and the exponential term exp(−λ(x_r)). To express this product property, we define q(x_r ∣ κ_r) as the product of n_r + 1 Gaussian experts with neural networks: where are the mean vector and precision matrix expressed as neural networks, and are also the mean vector and precision matrix. Here, approximates λ(x_r, κ_ri), and approximates exp(−λ(x_r)).

The posterior distribution of x under these assumptions reduces to a Gaussian distribution whose precision matrix is a block tridiagonal matrix. By using the block tridiagonal structure, the computational cost of generating samples and calculating their densities is proportional to the number of bins R [34]. For more details of the sampling procedure, see Section S1.6 in S1 Text.

Once we get a sample x from the encoder q_θ(x ∣ y, κ), we then evaluate the inner expectation terms in the lower bound (4). While the prior p(x₁), p(x_r ∣ x_r−1) and the observation distribution for the covariate p(y_r ∣ x_r) are straightforward to evaluate, the point process likelihood p(κ_r ∣ x_r) is not analytically tractable. This term is replaced by its lower bound derived in Section 2.3: (5) where and are i.i.d. samples from q(z ∣ x_r, κ_ri) and q(z ∣ x_r), respectively.

For training our model, we estimate all the parameters simultaneously by maximizing the lower bound (4). The Jackknife resampling and the doubly reparameterized gradients are also applicable in this maximization. For a summary of the parameter estimation, see Section S1.7 in S1 Text.

2.5 Decoding covariate from unsorted spikes

Next, we consider a decoding problem. Assume that only the unsorted spikes κ are observed and that we already know the model parameters. Under this assumption, the decoding aims to calculate the posterior of the hidden state x and the observed covariate y given the unsorted spikes κ. Using Bayes’ rule and our state space model definition, this posterior is expressed as follows: where p(x, κ) = p(x)p(κ ∣ x) is the joint density of x and κ. Since the posterior of y is determined by the posterior of x and the decoder p(y ∣ x), all we need to do is calculate the posterior p(x ∣ κ). Once we calculate p(x ∣ κ) and obtain samples from this posterior, the posterior of y can be approximated as

The rest of this section focuses on calculating the posterior p(x ∣ κ). Under our model’s definition, the point process likelihood p(κ ∣ x) is not Gaussian, making the posterior p(x ∣ κ) analytically intractable. To tackle similar situations, previous studies have proposed a nonlinear filter algorithm for point process observations [16, 43, 44]. They recursively approximated the point process likelihood at each time step using the Laplace approximation and constructed a forward filtering algorithm. In this paper, we mainly use the direct optimization method described in [45, 46]. In this approach, we approximate the point process likelihood for all steps simultaneously and obtain the approximated posterior of x directly.

The Laplace approximation provides the approximated posterior by fitting a Gaussian distribution with a mean equal to the mode of the posterior and precision equal to the Hessian matrix of the negative log density evaluated at the mode. Denote the gradient and the Hessian matrix of log p(x, κ) with respect to x as ∇_x log p(x, κ) and , respectively. Then, the Laplace approximation for the posterior is given by Thus, approximating the posterior reduces to the optimization problem of finding the MAP solution of the joint density.

Newton’s algorithm can be applied to find this . The Newton update rule at the k-th iteration is (6) After K iterations of this update, starting from the initial value x⁽⁰⁾, the mean and precision matrix of the Laplace approximation can be determined as However, a challenge arises due to the dimensionality of x. In our problem, the length of x is proportional to the number of bins, R. This makes multiplying the inverse Hessian with the gradient ∇_x log p(x, κ) computationally expensive. This issue can be addressed by using the block tridiagonal structure of the Hessian matrix [45].

A critical issue specific to our model is that the joint density log p(x, κ) is not always guaranteed to be concave. This leads to two problems. The first problem is that the Newton update does not always increase the joint density, leading to convergence to poor solutions. The second problem is that the precision matrix may not be positive definite, making the approximate Gaussian distribution invalid.

To address these issues, we approximate the log joint density as a concave function at each iteration. The log joint density is the summation of the log prior and the point process log-likelihood as Since the point process log-likelihood is not analytically tractable under our model, this term is replaced by the lower bound (5). As the log prior follows a Gaussian distribution, the first term is already concave. Consequently, our primary task is to approximate the latter lower bound term to be concave.

Hereafter, let us consider the concave approximation of the lower bound around the point . To ensure the concavity, we employ the technique used in Gauss-Newton optimization [47]. This optimization is applied to the objective function that can be represented as a composite of a smooth map and a concave function.

To apply this technique, we first rewrite the Rényi lower bound and the χ upper bound as composite forms: where are smooth maps, are concave functions, and are composite functions. A detailed definition of these maps and functions is provided in Section S1.8 in S1 Text. With these notations, the lower bound of the point process log-likelihood at the r-th bin is Now, we are ready to apply the Gauss-Newton optimization. To obtain the concave approximation of this lower bound, we approximate F_ri(x_r), G_r(x_r) by a first-order Taylor expansion around , and use them as inputs to π and ρ: where and are the Jacobian matrices of F_ri and G_r at , respectively. Since π and ρ are concave functions and their inputs are linear in terms of x, both and are concave functions. With these approximations, the approximated lower bound is always concave. Consequently, the log joint density replaced with this approximated lower bound is also concave.

This established concavity addresses two prior concerns. First, the concavity ensures that the Newton step consistently moves in an ascending direction throughout optimization. Second, the precision matrix derived from the Laplace approximation is guaranteed to be positive definite, ensuring the validity of the resulting Gaussian distribution.

We substitute this approximated objective into the Newton update rule (6). Owing to the block tridiagonal structure of the Hessian matrix, the computational cost for each Newton step is directly proportional only to the number of bins R [45]. For detailed information, see Section S1.8 in S1 Text.

3.1 Tasks and metrics

The prediction task aims to evaluate the model’s ability to predict the future covariate and unsorted spikes. In this task, we calculated the negative log-likelihood (NLL) of the test data: This metric shows how well the model predicts the covariate and unsorted spikes in a future interval. A smaller NLL indicates better predictive performance. Our model defines the joint mark intensity as the integral of the decoder, making the point process log-likelihood not analytically tractable. For such types of models, we use the lower bound of the log-likelihood instead to calculate NLL. Since this value becomes larger than the true NLL, our model is at a disadvantage in the comparison. Thus, if our model achieves superior performance compared to other models, the gap from the true NLL does not change the relative ranking of the models.

In the decoding task, we reconstructed the observed covariate from the unsorted spikes κ. Then, we calculated the mean squared error (MSE) between the reconstructed covariate and the true covariate as where ‖⋅‖₂ is the l₂ norm. A smaller MSE indicates better decoding performance.

Fig 2 summarizes the observed variables in the training step and evaluation tasks.

3.2 Compared models

To demonstrate the effectiveness of our model, we compared it with other marked point process models. In this section, we provide a brief overview of three of these models. These models define the joint mark intensity in different ways while sharing the same state space model described in Section 2.4. For details on definitions, hyperparameters, parameter estimation, and decoding procedures, see Section S2 in S1 Text.

We did not apply the KDE-based model [9] due to memory issues, which scale linearly with the number of spikes and the length of waveforms (e.g., 170,000 × 256 for Gatsby_08282013).

For notational simplicity, we denote our model as JVAE in the figures, short for the joint mark intensity model using VAE.

3.3 Synthetic data

Data generation.

We generated datasets consisting of covariate time series and unsorted spikes influenced by this covariate. In line with our state space definition, we assumed that a hidden state exists and influences both the covariate and unsorted spikes. To demonstrate our model’s ability to handle nonlinear dynamics, we employed the Lorenz attractor as this hidden state. The Lorenz attractor is a chaotic solution to the nonlinear differential equation: where α, β, and γ are parameters.

Given the hidden state, the two types of observed covariates were defined as follows:

1. (i) The observed covariate equals the hidden state, (8) The data with this covariate was referred to as Lorenz_3D.
2. (ii) The observed covariate was a two-dimensional projection of the hidden state, (9) The data with this covariate was referred to as Lorenz_2D. In this case, we could not access information for x_2t, which also had an influence on the generation of unsorted spikes.

We assumed that there were four probes measuring different neural population activities and that unsorted spikes were generated from the marked inhomogeneous Poisson process that depends on the hidden state at each probe. To imitate the activities of multiple neurons, the joint mark intensity of this process was defined as a mixture representation: (10) Here, λ_l is a scale parameter of the l-th component, and are the most favorable directions at t = 0 and t = T, respectively. The intensity of the l-th component reaches its maximum value when x_t aligns with the most favorable direction determined by these vectors. In this definition, this direction gradually changes from c_l to d_l, which expresses the plasticity of spiking activities. The mark of the l-th component is distributed according to a Gaussian decoder N(κ ∣ μ^z→κ(z), Λ^z→κ(z)) with a Gaussian prior whose mean vector and precision matrix are and , respectively. We selected the parameters of this decoder and prior so that the generated marks resemble the actual extracellular waveforms with 32 time steps in [53].

Under this definition, we generated ten trials of data. We first calculated ten paths of the Lorenz attractor for a duration of T = 100, each with different initial values for x_t at t = 0. Using these paths, we defined the observed covariates y_t by (8) or (9), and generated the unsorted spikes {(t_i, κ_i)} by (10). To use our state space representation, we split the time axis into bins with a length of Δt = 0.01 and prepared the discrete observations y and κ. Fig 3A shows the data generated by this procedure.

Download:

• PPT
  PowerPoint slide
• PNG
  larger image
• TIFF
  original image

Fig 3.

(A) Synthetic data. We assumed that the dataset consisted of four sets of unsorted spikes measured at different probes. Each set of unsorted spikes consists of four components that are tuned to different state values and generate spikes with varied waveforms. Each column indicates the spike times and waveforms measured at each probe. (1st row) The hidden state and spike times: A blue line represents the hidden state x_t. Colored dots on this line denote the state values at spike times for each component. (2nd to 5th row) The waveform distribution of each component. The solid line and colored band show the mean and two-sigma interval of the waveform distribution. (B) The prediction performance for synthetic data. The first row shows the results from Lorenz_3D and the second shows the results from Lorenz_2D. (Left) Violin plots of NLL. These values are subtracted by the baseline NLL of JVAE (our model). The white dots represent the median, the thick black bars indicate the interquartile range, and the thin black bars show the range from the minimum to the maximum values. (Right) The p values in the Wilcoxon test for assessing differences between two models. Each cell indicates the result for a pairing of two models. These p values are corrected by the Holm-Sidak method. The color indicates levels of the p values. (C) The decoding performance for synthetic data. (Left) Violin plots of MSE. (Right) The p values in the Wilcoxon test for assessing differences between two models. (D) Decoding results for the synthetic data. The first row shows the results from Lorenz_3D and the second shows the results from Lorenz_2D. (1st column) The true observed covariate. (2nd to 4th columns) Reconstructed covariates.

https://doi.org/10.1371/journal.pcbi.1012620.g003

3.4 Unsorted place cell spiking activities

Data description.

We applied our model to place cell multi-unit spiking activities in the hippocampus of rats freely foraging on various tracks [53, 54]. A place cell has a receptive field known as a “place field” and emits spikes when a rat passes through its place field [55]. Estimating the firing rate of each place cell as a function of spatial coordinates provides insight into how a neuron population codes spatial information. The spiking activities of place cells also offer an excellent example of neural decoding. When using place cell data for neural decoding, the goal is to reconstruct the rat’s trajectories from the spikes. The success of decoding serves as evidence that the place cell population activity contains sufficient information to recover the rat’s trajectory.

Let us summarize the correspondence between mathematical variables in our model and the components of this data. The rat trajectory is the observed covariate y_t, and the sequence of spikes attached with their waveforms on the probes is the unsorted spikes {(t_i, κ_i)}. The state space model p(x, y, κ) explains the hidden dynamics behind these data and how this hidden state affects the rat trajectories and its place cell spiking activities. The joint mark intensity λ(x, κ) describes the firing rate of spikes with the waveform κ when the rat’s hidden state is x. The decoding task aims to reconstruct the rat trajectory y from the unsorted spikes κ.

We selected four trials from the dataset: two different rats foraging on a linear track and a circular track [53, 54]. We refer to these trials by their respective directory names in the dataset: Gatsby_08282013, Gatsby_08022013, Achilles_11012013, and Achilles_10252013. In those trials, probes with four, eight, or ten recording sites were inserted into the rat hippocampus, and electrical potentials were measured simultaneously at each recording site. Spikes were detected using threshold processing, and spike waveforms with 32 time steps were extracted from the electrical potential signals at each site. The waveforms of all sites were concatenated into one vector and used as a spike feature κ. One trial contains around 200,000 spikes in a 400-second interval. Fig 4A shows the spike waveforms of clusters assigned by spike sorting.

Download:

• PPT
  PowerPoint slide
• PNG
  larger image
• TIFF
  original image

Fig 4.

(A) Waveforms of spikes. Each panel shows the collection of spike waveforms classified into the same cluster. (B) Prediction performance for the rat place cell spiking activities. Each row shows the result of each trial. (Left) Violin plots of NLL. These values are subtracted by the baseline NLL of JVAE (our model). The white dots represent the median, the thick black bars indicate the interquartile range, and the thin black bars show the range from the minimum to the maximum values. (Right) The p values in the Wilcoxon test for assessing differences between two models. Each cell indicates the result for a pairing of two models. These p values are corrected by the Holm-Sidak method. The color indicates levels of the p values. (C) Decoding performance for the rat place cell spiking activities. (Left) Violin plots of MSE. (Right) The p values in the Wilcoxon test for assessing differences between two models. (D) Decoding results for the rat place cell spiking activities. Each row shows the result of each trial. (1st column) 2D plots of the rat trajectories in the circular or linear track. (2nd to 4th columns) True and reconstructed rat trajectories. Solid blue and red lines show the coordinates y₁ and y₂ of the actual rat trajectories during the experiment. Intervals without solid lines indicate that the rat trajectories were missing during those intervals due to device issues. The blue and red bands represent the two-sigma predictive intervals.

https://doi.org/10.1371/journal.pcbi.1012620.g004

We concatenated the spike waveforms measured at multiple sites into one vector and used this as a mark κ. This κ is a 320-dimensional vector at one of the probes. The observed covariate y_t is a two-dimensional spatial coordinate. To use our state space representation, we split the time axis into bins with length Δt = 0.025 and prepared the discrete observations y and κ.

3.5 Comparison

In this section, we first discuss three issues with the compared models used in the experiments. Next, we explain how our model overcomes these issues.

The first issue is the overfitting of the ground intensity. For synthetic data, CVAE demonstrated performance comparable to GMM in the prediction task but failed in the decoding task. The primary cause of the large MSE values for CVAE was the deviation of the reconstructed path from the true path. This deviation might be attributable to overfitting during the training phase. CVAE and RMPP utilize a neural network to define the ground intensity λ(x), adjusting this network to maximize the log-likelihood (7). Since the log-likelihood contains the term −λ(x), this network tends to assign zero mass to points where no spikes are observed. In fact, the optimal solution that maximizes the log-likelihood is given by where is a Dirac delta function equal to 1 when x = x_r. Thus, if the neural network comprises too many layers and units, it fits the training data too well, thereby failing to generalize to the test data. This overfitting negatively impacts the decoding performance. During decoding, we approximate the log-likelihood with a concave function. If the network for the ground intensity overfits to the training data and approaches the optimal solution mentioned above, its concave approximation significantly deviates from the true ground intensity. Such an incorrect approximation might result in a large deviation from the true path. Therefore, addressing overfitting is crucial to achieving better decoding results.

The second issue is the deviation of simple parametric models from the true joint mark intensity, which may occur in GMM and RMPP. This deviation might have been caused by two factors: nonlinear embeddings and complex waveform distributions. Both GMM and RMPP showed poorer performance for Lorenz_2D than our model. Our state space model attempts to linearize the Lorenz attractor dynamics in the hidden state space and recover missing information x_2t through its nonlinear embedding. This embedding determines the hidden state space and, consequently, the joint mark intensity within this space. However, this embedding does not guarantee that the joint mark density will be a mixture of unimodal functions. GMM and RMPP, which assume a Gaussian mixture for the joint mark intensity, fail to capture the complicated shape of the joint mark intensity induced by the embedding. For the place cell data, GMM’s prediction performance was worse than that of CVAE. The spike waveform in the place cell dataset was a high-dimensional vector, and its distribution was more complex than that of the synthetic data. Therefore, simple parametric models like Gaussian mixtures are insufficient to represent actual waveforms.

The third issue, specific to sorted decoding, is the discontinuity of the observed spikes. For synthetic data, GRU demonstrated performance comparable to both CVAE and GMM. However, as shown in Fig C in S1 Text, the reconstructed paths were noisier compared to unsorted decoding and exhibited abrupt jumps. The primary cause of these jumps is that the paths reconstructed through sorted decoding reflect the discontinuity observed in the spikes. The regression-based decoders do not impose any assumptions on the continuity of the hidden states. Thus, the similarity between reconstructed covariate values across consecutive bins was derived from the similarities between spike count vectors across consecutive bins. When spikes occur, these vectors exhibit abrupt jumps, causing discontinuity in the inputs of the decoders and resulting in noisy reconstructed covariate values. For place cell data, the performance of sorted decoding was inferior to CVAE in one trial. This might have been due to the discontinuity induced by unreliable spike assignment. Neural population activities often include a hash cluster, which is poorly isolated from other clusters [18, 56]. Spike sorting assigns labels to spikes within this cluster by using hard decision boundaries, overlooking the uncertainty of these spikes. Such labeling causes discontinuity in the input vectors, leading to noisy reconstructed paths.

Our model addresses these issues. First, the use of a variational autoencoder avoids overfitting. Our model maps a low-dimensional latent variable z to the joint representation of the hidden state x and the mark κ, and imposes a prior over z. This model definition acts as regularization. Second, by utilizing neural networks, our model can capture the complex interdependencies between hidden states and marks in a data-driven manner. This works effectively even when the nonlinear embedding distorts the true joint mark intensity. Third, our model defines joint mark intensity to represent unsorted spikes, thus avoiding the discontinuities introduced by the assignment of unclusterable spikes. The information of spikes with ambiguous waveforms has a small influence on decoding. Moreover, our state space model defines the hidden state dynamics, enabling us to obtain smoother reconstructed paths.

3.6 Application to Neuropixels data

To demonstrate the applicability of our model to data collected with high-density probes, we applied it to the open dataset [57]. This dataset includes electrophysiological recordings from multiple brain regions of mice using Neuropixels probes and behavioral data during decision-making tasks. For more details on the dataset, see [57].

The Neuropixels probe has 384 sites and measures a wide range of neuron populations. We used recordings from 56 sites located in the CA1 region. This dataset also includes spike sorting results provided by Kilosort [58], consisting of spike times, spike amplitudes, and spike depths on the probe. Using the times of spikes with estimated depths in the CA1 region, we extracted the waveforms for 32 time steps around each spike from the raw recording (Fig 5A). Then, we concatenated these values into one vector and used it as a spike feature κ. The dimension of κ was 1792. The mouse’s face motion energy and wheel speed were used as covariates. The interval used for this experiment was 400 seconds long and included approximately 150,000 observed spikes.

Download:

• PPT
  PowerPoint slide
• PNG
  larger image
• TIFF
  original image

Fig 5.

(A) Waveforms of spikes. Spike waveforms measured at 56 sites in the CA1 region were used for the spike feature. (B) Decoding performance for the Neuropixels data. (Left) Violin plots of MSE. These values are subtracted from the baseline MSE of JVAE_H (our model). The white dots represent the median, the thick black bars indicate the interquartile range, and the thin black bars show the range from the minimum to the maximum values. (Right) The p values in the Wilcoxon test for assessing differences between two models. Each cell indicates the results for a pairing of two models. These p values are corrected by the Holm-Sidak method. The color indicates levels of the p values. (C) Decoding result for the Neuropixels data. The first row shows the result of the motion energy, and the second shows the wheel speed. Solid blue and red lines show the actual values during the experiment. The blue and red bands represent the one-sigma predictive intervals.

https://doi.org/10.1371/journal.pcbi.1012620.g005

We applied our model (denoted as JVAE_H) and the Gaussian mixture model (GMM_H) to this dataset. Additionally, we applied the Gaussian mixture model to the data with a two-dimensional spike feature, including spike amplitudes and spike depths on the probe (GMM_L).

Fig 5B and 5C show the decoding results. JVAE_H outperformed both GMM_H and GMM_L, indicating that our model effectively handles high-dimensional spike features and is useful for Neuropixels data. GMM_L performed better than GMM_H, consistent with the comparison of decoding using low-dimensional and high-dimensional spike features presented in Section S4 in S1 Text.