Large scale Lasso with windowed active set for convolutional spike sorting

TL;DR

This paper introduces a scalable, parallelizable active set algorithm for large-scale convolutional Lasso problems in spike sorting, enabling efficient processing of extensive neural datasets and potential online application.

Contribution

A novel active set algorithm for convolutional Lasso that is efficient, scalable, and suitable for parallel implementation, advancing large-scale spike sorting capabilities.

Findings

Linear complexity with respect to temporal data size.

Effective spike recovery demonstrated in experiments.

Robustness to regularization parameter variations.

Abstract

Spike sorting is a fundamental preprocessing step in neuroscience that is central to access simultaneous but distinct neuronal activities and therefore to better understand the animal or even human brain. But numerical complexity limits studies that require processing large scale datasets in terms of number of electrodes, neurons, spikes and length of the recorded signals. We propose in this work a novel active set algorithm aimed at solving the Lasso for a classical convolutional model. Our algorithm can be implemented efficiently on parallel architecture and has a linear complexity w.r.t. the temporal dimensionality which ensures scaling and will open the door to online spike sorting. We provide theoretical results about the complexity of the algorithm and illustrate it in numerical experiments along with results about the accuracy of the spike recovery and robustness to the regularization parameter.

Equations

$$S=\sum _{r=1}^k{W}_r\ast A_r+N,$$

$$y=Ha+\eta ,$$

$$\hat{a}=\underset{a\in {\mathbb{R}}^{kn}}{\mathrm{argmin}}\frac{1}{2}{\Vert y-Ha\Vert }_2^2+\lambda {\Vert a\Vert }_1,$$

$$\forall j,\text{ if }|H_j^T(y-H{a}^{\ast })|<\lambda ,\text{then }{a}_j^{\ast }=0.$$

$$\chi =\{j\text{mbox}suchthat{a}_j\ne 0\text{mbox}and\exists \ell \in \chi ,0<|j-\ell |\le t\}.$$

$$\forall 1\le r\le k,A_r\sim PPP({\mu }_r),$$

Full text

I Introduction

The problem of spike sorting consists in recovering the shape of action potentials and the time activations (also called spikes) of individual neurons from the recordings of a population via an extracellular set of electrodes. Since the shape of action potentials is usually characteristic of a given neuron and can be considered stationary, it is possible to detect when an action potential occurs and to attribute it to a given neuron. This reconstruction of distinct simultaneous spike trains is a key preprocessing step to evaluate phenomenons, such as firing rate coding or synchronizations between neurons, that are presumably a sensible part of the neural code (see for instance [1, 2, 3]).

As such it has been the focus of numerous works in recent years [4, 5, 6]. The estimation of the shapes is usually treated as a clustering problem, for shape estimation, followed by template matching to associate each action potential to a neuron having the closest shape [7]. While these methods have been used in practice for a while, especially in the case of tetrode recordings (meaning only 4 electrodes), they often require to perform manual pre and post-processing in addition to selecting several parameters, notably in presence of spike synchronization. As such the results of the procedure strongly depend on the person achieving the task [8, 9]. Moreover, these manual calibrations might not be possible in the near future due to the development of new acquisition methods and the availability of larger datasets and recordings [10]. Thus future spike sorting methods should depend on few parameters easy to select, provide nice statistical properties and finally scale to large volume of data that will become the norm in the next few years: large number of electrodes (up to 4000), large number of recorded neurons, large number of spikes and synchronizations especially when the recordings takes place during epileptic seizures [11].

Convolutional model

Apart from the classical clustering methods, one particular method relies on convex optimization [5] and uses a convolutional model for the recorded signals on $d$ electrodes. To describe it, let us use the following notation: upper case letters correspond to matrices and lower case letter to vectors, $*$ stands for the convolution operator along time. The model is written as follows:

[TABLE]

where $S\in\mathbb{R}^{d\times n}$ is the observed signal on $d$ electrodes and $n$ temporal samples, $W_{r}\in\mathbb{R}^{d\times t}$ is the matrix containing the shapes of the action potential of neuron $r$ on every electrodes (all the shapes have a temporal spread of $t$) and $A_{r}\in\mathbb{R}^{1\times n}$ is a sparse signal called activation in the sequel (more precisely, the time activations of neuron $r$ corresponds to the indices of the non zero coefficients of $A_{r}$) and $N\in\mathbb{R}^{d\times n}$ is a noise matrix. An illustration of the model is provided in Figure 1. This model has been introduced for modeling speech signals in [12] and has been used more recently for spike sorting in [5, 6]. The main advantage is that this model explicitly takes into account synchronization as an additive superposition of shapes, which cannot be done with simpler clustering based methods [4].

The model (1) is usually estimated by alternative optimization w.r.t. the activations $A_{r}$ and the shapes $W_{r}$. But the more computationally intensive update is clearly $A_{r}$ since the number of variables in $A_{r}$ is proportional to $n$. Indeed, to give an order of magnitude, records last typically 30 minutes to hours, which means with the classical time resolution that $n\simeq 10^{8}$, whereas the number of electrodes $d$ range typically from 4 to 4000, the number of neurons from 1 to 1000 and the temporal spread of a shape is of the order of the ms, hence $t$ lies between between 30 and 150 samples.

This is why we focus on the problem of estimating $A_{r}$ when the shapes $W_{r}$ are known. In this case the estimation of the activations $A_{r}$ can be done using the basis pursuit estimator as proposed in [5]. This estimator, that is equivalent to the Lasso estimator, has several nice statistical properties such as the ability to accurately recover the support, i.e. activation times and active neurons, under appropriate conditions [13, 14, 15]. The main problem with this estimator in this case is the computational complexity of estimating the $A_{r}$ on large signals (large $n$). Indeed the number of variables in the optimization problem grows linearly with the length of the recorded signal. Several works have investigated similar optimization problems in [16, 17] but to the best of our knowledge, our approach is the first to really exploit the structure of the problem to probably attain linear complexity in $n$.

Contributions

We propose a novel active set algorithm working on temporal sliding windows that allows to estimate the positions and magnitude of the sparse activations in model 1. Our method relies on simple operations such as convolution and can be easily implemented on parallel GPU architecture which ensures a nice scaling w.r.t. the dimensionality of the problem. We also prove for a simple yet realistic scenario on neurons activations that our active sets method will need to solve most of the time low-dimensional problems and even compute their average dimension. This last result ensures that the complexity of solving the optimization problem will stay linear in the number of temporal samples $n$. Finally we perform numerical experiments that confirm the linearity of the algorithm in practice and illustrate the robustness of the Lasso in this convolutional model for spike sorting.

II Optimization problem and large scale windowed active set algorithm

II-A Large scale optimization problem for spike sorting

Due to the linearity of the convolution operator, we can reformulate (1) as a linear model. Vectorization of $S$ (resp: $N$) as $y\in\mathbb{R}^{dn}$ (resp: $\eta\in\mathbb{R}^{dn}$), lead to the following form:

[TABLE]

where $a\in\mathbb{R}^{kn}$ is the activation vector to be estimated, and $H=M_{dn,kn}(\mathbb{R})$ is a block Toeplitz matrix, coding for the convolution between the shapes of the neurons and the activation vector. More precisely the blocks of the matrix $H_{p,r}$ are Toeplitz matrices of size $n\times n$, coding for the convolution between the shape of neuron $r$ on electrode $p$ and the activations of neuron $r$. Matrix $H$ is very sparse and structured with a sparsity ratio $\approx\frac{t}{n}$ with $t\ll n$.

Estimating $a$ when the number of neurons is bigger than the number of electrodes would be impossible without additional structural assumptions. But biological evidences show that neurons tend to spike rarely (typically the number of non zero coefficients in $A_{r}$ should be of the order $10^{5}$ when $n\simeq 10^{8}$) therefore the number of indices to be activated in $a$ should remain small relative to $kn$. Since $a$ is a sparse vector, we can use an estimator promoting sparsity, such as the well-known Lasso estimator equivalent to what is proposed in [5] :

[TABLE]

where $\lambda>0$ is the only tuning parameter of the method that depends on the signal to noise ratio.

Standard approaches to solve problem (3) include well known algorithms such as LARS [18], coordinate descent [19] and more recently proximal gradient descent [20, 21, 22]. When the solution is known to be very sparse, active set approaches that iteratively add variables and solve only low-dimensional Lasso problems have been known to provide fast solvers [23]. Other accelerations such as dual acceleration combined with screening and active set also work very well in practice [24, 25].

In this work we propose a novel active set algorithm that takes into account the structure of the problem. The main ideas behind the optimization procedure are to i) use this sparsity to solve several small scale independent optimization problems and ii) to use the convolutional structure of $H$ to speedup computations and limit memory use.

II-B Active set for estimating the activations

The essential idea behind the active set (AS) algorithm is to solve only small Lasso problems (3) by working only on a small subset of indices that are considered meaningful. The subset of indices updated at each iteration is the active set and is denoted as $J$ in the following. A crucial property of the solution of the Lasso problem comes from the following KKT conditions (see [26] for more detail).

Proposition II.1**.**

Let $a^{*}$ be a solution of (3) and write $H_{j}$ as the $j$-th column of $H$, where $1\leq j\leq kn$.

[TABLE]

The property above gives a good test to determine if a variable $a_{i}$ set at [math] should remain at [math]. Thus an idea would be to initialize the vector $a$ at [math] and add all the variables that violate this condition iteratively until all the remaining zero variables satisfy the condition. This leads to the AS algorithm reported in Alg. 1. Starting from a null vector $a$ and an empty active set $J$, we add to $J$ the index $j\in J^{c}$ corresponding to the largest violation of the KKT condition $|H_{j}^{T}(y-Ha)|$. Then we update our estimation of $a$ by computing the Lasso solution for the problem $(H_{J},y)$, where $H_{J}$ is the submatrix of $H$ obtained by keeping the columns $H_{j}$ of $H$, for $j\in J$. We repeat this step until (4) is satisfied for all $j\in\{1,\dots,nk\}\backslash{J}$. Overall the active set strategy will need to solve much smaller problems than the original one when the solution is sparse but it requires access to an efficient Lasso solver.

II-C Optimization of the active set algorithm

The active set algorithm described above is efficient but may be further optimized using the structure of the problem. In the following, we present optimization for the most expensive steps in Algorithm 1: the computation of the KKT condition (line 3-4) and the Lasso estimation on $J$ (line 6).

Memory and linear operator optimization

A first optimization that greatly reduces the time and space complexity is to use convolution to compute the KKT violation of line 3. Note that $Ha$ can be computed with $O(td|J|)$ using the sparsity of $a$ and the correlation $H^{T}(y-Ha)$ can be computed with $O(nkdt)$ operations that is the main computational bottleneck. The sub-problem in line 6 is of size $|J|$ but $H_{J}$ is actually very sparse and has only $td|J|$ non-zero lines.

Groups of activations

We can exploit the convolutional structure of our problem to reduce the cost of the Lasso step (line 1 in Algorithm 1). Action potentials are short time localized events, of length $t$ in our model. Then finding a new activation only affects a portion of size $t$ in the signal. Therefore the only coordinates in $a$ that are updated by the Lasso are the ones in a temporal window of width $t$ around the new activation. Instead of updating $a$ on the whole active set, we only need to consider the indices $i$ in $J$ that overlap with the new activated sample $j$, in the sense that $|i-j|\leq t$.

This means that the true complexity of the AS update actually depends only on how big is the set of indices that need to be updated each time. An order of magnitude of the size of these overlapping groups (for short, overlaps) is computed in the next section, under plausible biological assumptions. If this size does not depend on $n$, this allows us to greatly decrease the size of the Lasso estimation (line 1). Note that a similar approach was proposed in [23] for 2D convolution but required the use of a connected component algorithm [27] that is not necessary in 1D where groups can be updated more efficiently.

Sliding window active set

We now discuss the core of our approach. First note that when using the optimization above, the computational bottleneck of the AS is KKT violation computation and maximum in lines 3 and 4 since it requires to compute the $O(nkdt)$ correlation in line 3 for each AS iteration. Since the number of iteration in the AS is $O(kn)$ this leads to a quadratic complexity in $n$ which limit its application on very long signals.

But finding the maximum along the whole signal is actually not necessary for convergence of AS as long as the KKT conditions are verified at the end. This is why we propose to compute KKT violation (line 3) and find the maximum (line 4) only in a temporal window of size $w>t$ in practice. We start with a window at the beginning of the signal and perform the active set. When there is no violation of the KKT on the window, two possible things occur. If there is an overlap strictly inside the window, then we can move the window to the right since this group is independent from the rest of the signal and its KKT will never be violated in the future. If the overlap goes beyond the window, then the size of the window is extended until the group is fully contained in it and the window can be translated again to the right. What makes this new approach very interesting in practice is that now the KKT violation in line 3 requires $O(wkdt)$ operations, that is independent on $n$. This means that if the resolution of the sub-Lasso in line 6 does not depend on $n$, then the whole AS algorithm complexity becomes $O(n)$ instead of $O(n^{2})$.

II-D Size of the overlaps

As we have seen before, the complexity depends strongly on the size of an overlap, meaning the larger set of time activations $\chi$ such that

[TABLE]

Often Poisson processes are used to model time activations of neurons in neuroscience [28]. Of course, these processes are in continuous time, whereas our record is in discrete time. But the discretization is so thin with respect to the firing rate of the neurons that the difference is negligible. Moreover the discretization process usually only discards spikes [28] and can only diminishes the size of an overlap.

Theorem II.2**.**

Assume that the activations of the neurons in 1 are drawn from independent Poisson point processes, that is

[TABLE]

where $\mu_{r}>0$ is the spiking intensity of neuron $r$. Let $\mu=\mu_{1}+\dots+\mu_{k}$. Then the mean size of an overlap is smaller than $\mu te^{\mu t}$.

Sketch of proof: Consider $A$ the union of the $A_{r}$. $A$ is a Poisson point process of intensity $\mu$. The gaps between two activations obey an exponential distribution with mean $1/\mu$ and are independent, then we can compute the probability that the length of an overlap is greater that $mt$, with $m\in\mathbb{N}$. $\square$

II.2 gives an average size of overlap independent from $n$. Despite the exponential term, the small spiking intensity $\mu$ make this size reasonable in practice. For instance, for a sampling frequency $F=30$kHz, a spiking activity of $30$Hz ($\mu=10^{-3}$), and spikes of length approximately $5.10^{-}{3}$s ($t=150$), we have an average number of activation in the overlap smaller than $7$ for $k=10$ and $60$ for $k=20$. Those values are in practice very small compared to $n=10^{8}$ corresponding to roughly an hour of recording. Note that when the number of electrodes becomes large we can extend this result to spatio-temporal overlap, leading to possible decrease in their size.

III Numerical experiments

In this section we illustrate the performance and computational complexity of our solver. For solving the Lasso, for both the full model and the AS approaches we use the accelerated proximal gradient FISTA [22] that is implemented in parallel in the SPAMS toolbox [29]. All experiments were performed of a simple notebook having 8GB memory and an Intel(R) Core(TM) i7-4810MQ CPU @ 2.80GHz. The code will be made available online upon publication.

III-A Computational complexity

In order to evaluate the computational complexity of our method we create a dataset from simulated data. We simulate realistic action potential shapes using the well-known Hodgkin-Huxley model [30] with implementation from [31], and generate matrix of activation $A_{r}$ as discrete point Poisson processes of intensity $\mu=10Hz$. We want to see the impact of the different optimization procedures discussed in II-C. First we solve the full Lasso that requires to pre-compute matrix $H$ which limits the size to the available memory (Global Solver), we also report the performance of a naive AS that also use the $H$ matrix to see the effect of AS on sparse solutions. Finally we implement the activation groups speedup (AS with groups II-C) and the sliding window AS discussed in II-C. Using $5$ neurons and $4$ electrodes, we generate signals of various lengths (that is $n$), in the noiseless scenario.

On Figure 2, we can visualize the impact of the optimizations of the active set that we described in section II-C. The global solver and the naive AS can only handle small values of $n$, as they are both memory intensive but the AS is one order of magnitude more efficient. It shows that the execution time of the active set with sliding window growths linearly with $n$ whereas the AS on groups grows quadratically as discussed in the previous section.

III-B Influence of the noise and the regularization parameter

Proper calibration of the regularization parameter $\lambda$ is crucial for the success of the estimation. We want to visualize the influence of this choice, especially for various noise levels. Using real shapes of action potentials recorded in [32] and that have been already spike sorted by classical algorithms, we simulate signals of size $n=500$ for different noise levels, with $k=2$ neurons firing at $50Hz$ and recorded by $d=4$ electrodes. A first measure of performance that we consider is the classical F1-score, which estimates a balance between false positive and false negative rates. This measure tends to be pessimistic as it penalizes equally short and long time deviations. We introduce a softer measure of performance: $CP(x,y)=1-\left\|K*(x-y)\right\|_{1}/(\left\|x\right\|_{1}+\left\|y\right\|_{1})$, where $K$ is a normalized rectangular function. Depending on the size of the support of $K$, this measure allows us to tolerate small time deviations.

The regularization parameter $\lambda$ should be chosen carefully in order for the Lasso to achieve good recovery. Taking $\lambda$ too small leads to too many activation whereas taking $\lambda$ too large leads to zero activation detected. As we can see on fig. 3, when the SNR is too weak, the admissible range for $\lambda$ becomes too narrow for any automatic tuning method to work. However for realistic signal-to-noise ratios, recovery is possible for a decent range of the regularization parameter $\lambda$, which confirms that the implementation of this method is conceivable in this spike sorting application.

IV Conclusion

In this paper, we provide an efficient algorithm to estimate the activations of the neurons. The windowed active set exhibits good performances, both in terms of estimation quality and execution time (scales linearly).

For future works, we intend to prove that the Lasso estimator recovers the correct support of activations, especially when the number of electrodes grows. Applications on real datasets will also be carried out, allowing us in particular to take into account the real structure of the noise. Finally, our windowed lasso open the door to online adaptation of the spikes shapes and we will investigate this by extending the works of [33].