Introduction

State-space modeling is a promising and versatile analytic approach for characterizing neural signals that has received substantial attention over many decades [1]. State-space models can be highly effective for analyzing neural time series data, particularly if the model formulation and parameterization are able to accurately represent the dynamics of the data generating process [2–4]. Powerful neural data analysis methods have been developed using stationary linear models [1, 5–9], for which inference and learning algorithms can be readily derived using well-established methods. However, real neural signal dynamics are more often time-varying, due to either intrinsic fluctuations of physiological states or extrinsic influences from cognitive and environmental processes [10–15]. Accordingly, mismatches between chosen state-space models and true underlying dynamical processes can be a major source of error in such algorithms [16].

In practice, researchers can manage variations in temporal dynamics by allowing the model parameters to vary across time or fitting different model parameters to consecutive windows of data. Thresholds can also be used to exclude outliers and limit the range of time series data within some neighborhood where a locally time-invariant approximation remains valid. These approaches can be effective when the time-varying activity can be tracked by changing parameter values within a given model class, or when the temporal dynamics evolve slowly compared to the window length. However, they fall short when the neural signals of interest are by definition transient. Examples of such transient activity are abundant in neuroscience, including epileptic bursts, hippocampal ripples, and event-related potentials [17–21]. In other scenarios, neural dynamics can change abruptly at unpredictable times, for instance, during anesthesia-induced burst-suppression, different stages of sleep, and modulations from active learning and behavioral responses [22–30]. Time-varying parameters or windowed approximations would perform sub-optimally at best in these cases, and could miss rapid transitions in the dynamics if the scale of temporal variation or window length is mis-specified.

State-space models with switching have been proposed to represent signals that are composed of multiple segments with different dynamics, under the assumption that each segment is approximately linear and stationary. Two distinct approaches have been taken to develop corresponding algorithms for these switching state-space models. In one approach, approximate solutions were developed based on existing solutions [31] to stationary Gaussian state-space models, or Gaussian SSM for short. These early studies recognized the intractability of accounting for all possible switching transitions whose combinations grow exponentially with the length of the time series [32, 33]. Nevertheless, efficient inference algorithms have been developed by assuming that the transitions between multiple state-space dynamics follow a hidden Markov model (HMM) process and by approximating the conditional means and covariances of Gaussian hidden states [34–38]. These methods have focused on the filtered estimates of the switching model probabilities and therefore on the ability to forecast a transition in dynamics. While this inference is relevant in control systems and econometrics, data analysis questions in neuroscience are more often centered on segmenting a given time series into periods with distinct neural dynamics and subsequently characterizing the properties of those dynamics. Likely for this reason, these inference algorithms have not been extensively applied in neural signal processing.

In another more recent approach, advances in Bayesian methods have enabled the development of several new algorithms for switching state-space models [39–41]. These methods assume similar switching dynamics undergoing HMM transitions but leverage Markov chain Monte Carlo sampling techniques to perform state estimation and parameter learning. The adoption of blocked Gibbs sampling [41, 42] for switching state-space models is particularly powerful, as it allows extensions with non-parametric [39] and generalized linear models [42, 43] as well as with traditionally intractable structures such as recurrent dynamics [41]. These state-of-the-art methods have been successfully applied to the analysis of neural signals [43–47]. Empirical studies so far using these algorithms have employed less structured hidden state-spaces with dense transition matrices. A limitation of this general-purpose approach is that the candidate models may be less interpretable in relation to the neural mechanisms being studied compared to simpler and physiologically motivated models [48]. So far, the Bayesian sampling methods have been successfully applied to segmentation problems using high-dimensional embedding of neural signals [43–47]. Meanwhile, state-space models employing quasi-stationary sliding windows with closed-form solutions have been used to characterize time-varying signals [49–52]. The performance of such models in time-varying scenarios could benefit from extensions to incorporate switching dynamics.

In 2000, Ghahramani and Hinton proposed an insightful solution for switching state-space models inspired by variational approximations from graphical models [53]. Their generative models were similar to the hybrid models combining Gaussian SSMs and HMM, and they derived inference and learning computations in the recent Bayesian framework. But, instead of relying on Gibbs sampling, their approximate inference and learning solution utilized traditional exact inference algorithms as building blocks. The Ghahramani and Hinton algorithm is uniquely situated between the above two eras of research on switching state-space models, and it provides an opportunity to combine strengths from both approaches to study neural activity with transient dynamics using interpretable models and optimized approximations. This algorithm provides an accessible entry point to switching state-space models and facilitates construction of time-varying models of neural activity that could be further developed using more recent Bayesian methods. Despite the important and significant conceptual advance described by Ghahramani and Hinton, its applications in neuroscience have also been limited. Overall, the likelihood function for switching state-space models is non-convex and solutions are therefore sensitive to the initialization conditions. To address this issue, Ghahramani and Hinton used deterministic annealing, enabling the algorithm to perform comparably to past inference methods, but with little improvement. Moreover, the complexity of the algorithm and its computational requirements may have limited its adoption.

In this paper we introduce methods to significantly improve switching state-space model inference and learning under the framework initially proposed by Ghahramani and Hinton. First, we present a simple but rigorous derivation of the inference algorithm [53] under the variational Bayes approximation and expectation-maximization (EM) frameworks, complete with all closed-form equations involved along with practical considerations when applying this method to neuroscience problems. We then describe a novel initialization procedure that addresses the crucial challenge of navigating the non-concave log-likelihood function to find relevant optimal solutions. We show results on how this method significantly improves over deterministic annealing, enabling variational inference to outperform past inference algorithms. We also show that our approach improves parameter recovery, which makes it possible to more accurately characterize time-varying dynamics. We then extend the generative model structure to accommodate more complicated switching state-space models including nested models, which are frequently encountered in many applications including neuroscience. Finally, we apply the variational Bayesian learning algorithm to a long-standing problem of sleep spindle detection from electroencephalography (EEG) recordings during sleep in humans and show compelling results that switching state-space models can reliably identify transient neural activity.

Results

Throughout this work, we use regular and boldface lowercase letters to denote scalars and vectors, respectively. Matrices are denoted by boldface uppercase letters or by regular uppercase letters when the matrices are one-dimensional. The transpose of a matrix M is denoted by M^⊤, and M_ij indicates the element of the matrix at the i^th row and j^th column position. A variable indexed with another discrete variable s taking values in {1, ⋯, M}, e.g., z^(s), refers to the following:

We consider switching state-space models to consist of a set of M uncorrelated linear Gaussian SSMs encoding arbitrarily distinct dynamics and an HMM of a discrete variable s taking values in {1, ⋯, M}. The hidden states of Gaussian SSMs evolve in parallel and are allowed to be of different dimensions with appropriate mapping to observations. At every time point, the HMM selects one of the Gaussian SSMs to generate the observed data, giving rise to the switching behavior of this generative model. However, this flexible switching structure comes with high computational complexity: exact inference of the hidden states from observed data quickly becomes intractable, even for M = 2 with moderately long time series.

It has been noted that when the switching states are known, the hidden states of Gaussian SSMs can be efficiently estimated. Conversely, one can infer the hidden HMM states given the Gaussian SSM states [53]. Based on this insight, the intractability can be circumvented by using a surrogate distribution, q, which approximates the intractable posterior, p. Specifically, we introduce two auxiliary variables and : acts as the model evidence for the m^th Gaussian SSM to produce the observed data in the absence of known Gaussian SSM states, while represents the model responsibility for the m^th Gaussian SSM to explain the observed data when the switching states are unknown. Therefore, alternately updating and allows us to efficiently estimate posterior distributions of all hidden states in closed-form. The functional forms of these variables are obtained by maximizing a closeness metric between the two distributions p and q. This procedure is known as variational approximation [54, 55]. We can also use this approximate variational inference within an instance of generalized EM algorithm to learn (fine-tune) the Gaussian SSM and HMM parameters when the model parameters are unknown (unsatisfactory). Fig 1 outlines this variational Bayesian learning algorithm for switching state-space models.

Download:

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

Fig 1. Flowchart of the variational Bayesian learning as an instance of generalized EM algorithm.

M Gaussian SSMs, indexed by m ∈ {1, ⋯, M}, is parameterized by and augmented with an HMM parameterized by ρ and ϕ. The HMM determines the switching among the Gaussian SSMs to produce observed data. The observations are corrupted by observation noise with covariance R and indexed by t ∈ {1, ⋯, T}. We introduce and , two variational summary statistics, to approximate the true posterior distribution p with a different distribution q. The E-step requires inference of the hidden states, achieved through fixed-point iterations that improve the variational approximation incrementally. Once the E-step has stabilized, model parameters are updated in the M-step. The convergence of the inner (E-step) iterations and outer (EM) iterations, index by i, is checked using the negative variational free energy . This algorithm outputs 1) posterior estimates of model parameters, , 2) state inferences, i.e., the means and covariances of Gaussian SSM hidden states, and , and 3) estimates of M model probabilities of generating the observation at each time point, .

https://doi.org/10.1371/journal.pcbi.1011395.g001

The variational inference procedure requires a good initialization of or so that they can be iteratively updated to drive the surrogate posterior q closer to the intractable true posterior p as described in Materials and methods. Given the non-convex nature of the problem, a good initialization should lead to a good local minimum. In practice, having an informed initialization is often difficult since no prior information on the discrete variable s_t is available. One interesting quantity available from closed-form state inference is an interpolated density, defined as the conditional probability distribution of any particular observation, given all past and future observations. This interpolated density allows us to devise an informative initialization of the iterative variational inference procedure, instead of using deterministic annealing [53] (see details in section Initialization of fixed-point iterations of Materials and methods). Concretely, we use this density to compare between the Gaussian SSMs and establish the initial weights for the HMM, , which enables us to achieve superior performance in both segmentation and parameter estimation.

In the following results sections, we first show simulation studies to assess the performance of such variational inference and learning. As performance metrics, we evaluate segmentation accuracy of a time series by learned switching states against the ground truth, and parameter estimation errors where applicable, to compare the proposed algorithm with a few existing switching inference algorithms. In addition, we investigate the effects of varying data length on segmentation and parameter learning metrics. Lastly, we model real-world human sleep EEG using switching state-space models and apply the proposed algorithm to detect the occurrence of sleep spindles in an unsupervised manner.

Segmentation with posterior inference

We first focus on the variational inference part (E-step) of the algorithm that approximates the true posterior distribution p with a structured approximate posterior distribution q via fixed-point iterations (see the Variational approximation of hidden state posterior section of Materials and methods). To demonstrate improvements with our novel initialization procedure using the interpolated density, we repeated the simulations from [53] using the following switching autoregressive (AR) models of order 1:

The starting points for each AR model were drawn from the same distribution as their respective state noises. The switching state s_t followed a binary HMM process with initial priors ρ₁ = ρ₂ = 0.5 and a symmetric state-transition probability matrix with ϕ₁₁ = ϕ₂₂ = 0.95, ϕ₁₂ = ϕ₂₁ = 0.05. We generated 200 sequences of 200 time points from the above generative model and analyzed the segmentation accuracy of inference algorithms given the true parameters.

Four inference algorithms were compared: static switching [33], interacting multiple models (IMM) [35], variational inference with deterministic annealing (VI-A) [53], and our proposed variational inference with interpolated densities (VI-I). Static switching assumes the switching state s_t to be independent across time points and applies the Bayes rule directly for switching inference [33]. IMM utilizes Gaussian merging to approximate the posterior distribution at each time point to estimate the switching state [35]. VI-A was initialized with a temperature parameter that decreased to over 12 fixed-point iterations with . For comparison, VI-I also ran for 12 fixed-point iterations. Switching states were labelled with a 0.5 threshold on the posterior model probabilities.

An example simulated y is shown in Fig 2a along with the estimated switching state by random segmentation and the four inference algorithms. Histograms of percentage correct segmentation (Fig 2b) verify that the same results of VI-A are obtained as in [53] with mean accuracy 0.810. Surprisingly, both static switching and IMM were more accurate, with means at 0.827 and 0.864, respectively. VI-I achieved the best mean accuracy at 0.890, surpassing the other algorithms. Notably, in terms of the precision of correct segmentation, i.e., the width of the histograms in Fig 2b, all of static switching, IMM, and VI-I show superior performance over VI-A.

Download:

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

Fig 2. Simulation results: Segmentation performance when true parameters are known.

(a) An example simulated sequence switching between two AR1 models with different dynamics. The top panel shows the time trace with the two states marked in different colors. The bottom panel shows switching inference results on this example given true parameters of the underlying generative model. Time points estimated to be in the first model (s_t = 1) are marked in colored dots for each inference method, with accuracy shown in parentheses. (b) Histograms of segmentation accuracy across 200 repetitions. The mean segmentation accuracy for each method is displayed and marked by the dashed red line. True = ground truth; Random = random segmentation with a Bernoulli process; Static = static switching method; IMM = interacting multiple models method; VI-A = variational inference with deterministic annealing (orange color); VI-I = variational inference with interpolated densities (blue color).

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

Segmentation with parameter learning

An important advantage of the variational Bayesian framework over other approximate inference algorithms is the ability to learn model parameters using a generalized version of EM. When the underlying parameters are unknown, an accurate learning algorithm can improve segmentation performance over an arbitrary selection of parameters; conversely, if the learning algorithm is inaccurate, segmentation could become unreliable.

To investigate switching segmentation when the model parameters are unknown and must be learned from data, we conducted simulations similar to the study above using the following generative model:

Initial state values followed their model-specific state noise distributions. An HMM process identical as before was used for the switching state s_t, and we again generated 200 sequences of 200 time points. We initialized all algorithms with random parameters drawn from uniform distributions centered around the true values for the l^th sequence: where the notations for model parameters are described in the Real- and discrete-valued state-space models section of Materials and methods.

We compared the same four algorithms: for static switching and IMM inference algorithms, we estimated the segmentation of sequences using the randomly initialized parameters, since parameter learning is not part of these algorithms; for the VI-A EM and VI-I EM learning algorithms, we ran both the fixed-point iterations during the E-step and generalized EM iterations until convergence (see the Fixed-point iterations section of Materials and methods and S3 Appendix). The same temperature decaying in the variational inference analysis was used for the VI-A EM learning at every E-step. Switching states were labelled with a 0.5 threshold on the posterior model probabilities.

Percentage correct segmentation histograms show performance degradation as expected across all algorithms compared to inference with true parameters (Fig 3a). Despite the use of incorrect parameters, static switching and IMM achieved reasonable segmentation accuracy at 0.750 and 0.809, respectively, which decreased by 6–7% compared to the previous section where the true parameters were used. Notably, VI-A EM was much less accurate, with a mean accuracy of 0.705. In contrast, VI-I EM maintained excellent segmentation accuracy, outperforming the other algorithms, with a mean of 0.849 that only decreased 4.1% compared to the previous section. This pattern of better segmentation accuracy of VI-I EM over others was replicated when the same simulation study was repeated across a wide range of data length (see Fig 3c). Unlike static switching and IMM, both EM methods achieved greater segmentation accuracy as data length increased, followed by eventual plateauing. However, VI-I EM reached its plateau at a much higher accuracy than VI-A EM.

Download:

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

Fig 3. Simulation results: Segmentation performance when true parameters are unknown.

(a) Histograms of segmentation accuracy across 200 repetitions of sequences of 200 time points. Since the true parameters were not known, they were estimated using an instance of EM algorithm for VI-A and VI-I inference starting from a random initialization. Static switching and IMM treated those random initialization as true model parameters. The mean segmentation accuracy for each method is displayed and marked by the dashed red line. The mean accuracy previously obtained using true parameters is shown in gray. Random = random segmentation with a Bernoulli process; Static = static switching method; IMM = interacting multiple models method; VI-A EM = variational EM learning with deterministic annealing (orange color); VI-I EM = variational EM learning with interpolated densities (blue color). (b) Swarm plots showing the distributions of model parameters learned by the variational learning algorithms for sequences of 200 time points. Uniform distributions used to sample initial parameters are marked in bold fonts on the y-axes, as well as using solid black lines (true values) and dotted gray lines (upper and lower bounds of the ranges). (c) Changes of mean segmentation accuracy over sequences of varying data lengths. Shaded bounds denote the standard error of the mean around the average accuracy values. (d) Mean parameter estimation errors from the true values across 10 EM iterations for two different data lengths. For transition matrix F and state noise variance Q, normalized error is defined as abs(estimated—true)/true, and averaged across the two switching models. Absolute error is defined as abs(estimated—true). Shaded bounds denote the standard error of the mean around the average error values.

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

To characterize the robustness of variational learning and the ability to recover generative model parameters via EM, we analyzed the converged model parameters produced by the two variational learning algorithms relative to the true values (Fig 3b). Across model parameters, VI-A EM failed to identify the true values in most cases, suggesting the algorithm got trapped in local maxima of log-likelihood (see S4 Appendix for an analysis of negative free energy), which explains the poor segmentation accuracy in Fig 3a. VI-I EM estimated model parameters that have distributional modes at the true values, indicating more successful parameter recovery than VI-A EM.

Tracking updated model parameters through 10 EM iterations reveals that estimated parameters converged quickly (∼5 EM iterations) to their stationary values (Fig 3d). On average, the stationary values obtained by VI-I EM were closer to the true values than VI-A EM, which is consistent with Fig 3b. Additionally, some VI-A EM estimated parameters converged away from the true values. Increasing the data length reduced estimation errors in both methods while requiring slightly more EM iterations to reach desired convergence. The observation noise variance R did not vary substantially during EM iterations, since its magnitude was much smaller compared to the state noise variances Q^{1,2} (see S4 Appendix). Lastly, both methods tended to overestimate the probability of HMM to stay in the same model, a bias due to finite sampling that was attenuated for a longer data length (Fig 3d).

All of the simulations produced similar results when they were repeated with less informative distributions for model parameter initialization, albeit with slightly less accurate segmentation (see S4 Appendix). Overall, the variational learning method with interpolated densities showed robustness under uncertain model parameters.

Extensions of model structure and parameter estimation

The simulations studied so far employ observations that switch between two AR1 models. While the findings are useful as a proof of principle for variational Bayesian learning, they are not representative of neural signals in real applications that require switching inference, as such signals usually have more complex dynamics. In addition, a binary HMM switching between two parallel processes introduces large change-point deviations as can be seen in Fig 2a, which might be trivially detected by looking at the first derivative of the time series. To address these gaps, we conducted simulation analyses using a single bivariate AR1 model that switches between two state-transition matrices. Specifically, the generative model is as follows: where and . The switching state s_t followed a binary HMM process as before with initial priors ρ₁ = ρ₂ = 0.5 and a state-transition probability matrix with ϕ₁₁ = ϕ₂₂ = 0.95, ϕ₁₂ = ϕ₂₁ = 0.05. Therefore, this generative model consists of two identical AR1 models where the second model has time-varying influence on the first. An example simulated time series y is shown in Fig 4 along with the resultant segmentation by the four inference algorithms using true parameters. We generated 200 sequences of 200 time points and analyzed the segmentation accuracy of inference algorithms given the true parameters. The same temperature parameters as before were used for variational inference with deterministic annealing. Both VI-A and VI-I ran until convergence of the fixed-point iterations. Switching states were labelled with a 0.5 threshold on the posterior model probabilities.

Download:

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

Fig 4. Simulation results: Example segmentation on data generated from a different switching model class when true parameters are known.

The top two panels show the two sequences, y₁ and y₂, recorded as a bivariate observation y. Sequence y₂ has a non-zero influence on sequence y₁ as shown by the upward arrows, according to a switching state s_t. The time traces are also marked with different colors for the two switching states. The bottom panel shows the results of switching inference on the example data given true parameters. Time points estimated to be in the first model (s_t = 1) are marked in colored dots for each inference method, with accuracy shown in parentheses. True = ground truth; Random = random segmentation with a Bernoulli process; Static = static switching method; IMM = interacting multiple models method; VI-A = variational inference with deterministic annealing (orange color); VI-I = variational inference with interpolated densities (blue color).

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

It is evident from both the generative model equations and the example time series that this is a more difficult inference problem than previous studies using two AR1 models with uncoupled dynamics. Nevertheless, all inference algorithms provided informative estimates compared to random segmentation, with the variational inference algorithms obtaining better accuracy results (Fig 5a). VI-I produced segmentation with the best accuracy mean at 0.741, followed by VI-A at 0.719. Static switching and IMM had worse performance at 0.634 and 0.704 respectively. We note that the true generative model in this simulation is different from the assumed generative model employed during inference. Nonetheless, we see that VI-I still provides useful switching inference. This feature is discussed in more detail in the forthcoming section.

Download:

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

Fig 5. Simulation results: Segmentation performance on data generated from a different switching model class.

(a) Histograms of segmentation accuracy given true parameters across 200 repetitions. (b) Histograms of segmentation accuracy when model parameter were unknown across 200 repetitions. We followed the same strategy stated in Fig 3 when the true model parameters were not available. In both (a) and (b), the mean segmentation accuracy for each method is displayed and marked by the dashed red line. Random = random segmentation with a Bernoulli process; Static = static switching method; IMM = interacting multiple models method; VI-A = variational inference with deterministic annealing (orange color); VI-I = variational inference with interpolated densities (blue color). VI-A/VI-I EM denote the EM learning algorithms with the corresponding initialization procedure during E-steps.

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

Next, we proceed to variational learning when model parameters are unknown. We initialized all algorithms for the l^th sequence with parameters drawn from uniform distributions centered around the true values:

Multivariate static switching and IMM were applied using the randomly initialized parameters. For VI-A EM and VI-I EM learning algorithms, we ran both the fixed-point iterations (E-step) and generalized EM iterations until convergence. Identical annealing and switching state decision thresholds were used as before.

With the M-step of variational learning in this analysis, we demonstrate an important extension of the structured variational approximation to support jointly estimated parameters across models. Specifically, the structured variational algorithms assume the following generative model structure in this bivariate AR1 example: where the two candidate models have distinct transition matrices F⁽¹⁾ and F⁽²⁾ that could be updated separately during the M-step, likewise for Q⁽¹⁾ and Q⁽²⁾. However, the true generative model suggests that if two parallel models are used to approximate the switching dynamic, they should share all parameters except the F₁₂ element of the state-transition matrix that gates the effect of the second sequence on the first sequence. Simple manipulations of the update equations can exploit this shared structure to reduce the set size of estimated parameters and pool information across candidate models. This joint parameter learning across models is described in detail in S2 Appendix.

As expected, the segmentation results in Fig 5b show slightly inferior performance compared to when using true parameters. The VI-I EM algorithm achieved the best mean accuracy at 0.732. Distributions of converged parameters were comparable between VI-I EM and VI-A EM (see S4 Appendix), suggesting better segmentation could still be obtained in the absence of obvious improvement in parameter recovery. In addition, when compared against a few different variational algorithms that are able to assume the accurate generative model structure using a single hidden state [41], VI-I EM also showed improved segmentation consistently (see S4 Appendix for more details).

This simulation study suggests that our method could be effective in analyzing neural signals that exhibit subtle switching dynamics, and in time-varying Granger causality problems particularly when the causal structure is changing rapidly [56–59]. In the current setting, the switching state can provide a statistical inference on the improvement of temporal forecasting of the first sequence by the second sequence at every time point, without relying on log-likelihood ratio tests and windowing [60–63]. Our method also estimates parameters shared between the “full” and “reduced” models, while taking the observation noise into account [61, 64]. This application can be extended to higher-order models in a straightforward manner.

Switching state-space oscillator models

The performance of methods based on parametric state-space models, such as the switching inference algorithms studied here, is ultimately dependent on how well the chosen model class represents the signal of interest. Oscillations are ubiquitous in neural signals, and therefore they need to be modeled with high sensitivity and specificity for the methods to capture any transient oscillations. Here, we employ a recently developed state-space model that is uniquely suited for modeling oscillations [65]:

This model can be viewed as a phasor rotating around the origin with frequency ω (in radians) in the real and imaginary plane whose projection onto the real line produces the observed oscillation [66]. We refer to this model as an oscillator hereafter.

In many neuroscience applications, neural systems may exhibit different states defined by distinct combinations of oscillations, and the system may transition between these neural states over time. To test the applicability of our switching inference method for analyzing neural oscillations, we conducted a simulation study using hidden states composed of oscillators. We simulated up to 5 simultaneous oscillations at different frequencies with model parameters a = 0.98, σ² = 3 and R = 1 for 10 seconds under a sampling rate of 100 Hz. The frequencies were set at 1, 10, 20, 30, and 40 Hz respectively. With n total underlying oscillations, one can generate 2ⁿ − 1 possible states, each with a different combination of the oscillations. We used a multinomial switching variable s_t taking 2ⁿ − 1 values to select one of the states at each time point, to determine which oscillations were observed. The switching states followed an HMM process with uniform initial priors and a symmetric state-transition probability matrix with 0.98 on the diagonal. We repeated this data generation process 200 times for each of n = 2, ⋯, 5 with the above oscillators. Fig 6a shows one such simulation instance where n = 5.

Download:

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

Fig 6. Simulation results: Segmentation performance on switching state-space oscillator models when true model parameters are known.

(a) Power spectra of simulated oscillations at different frequencies; spectral estimation obtained via multitaper method that utilized 6 tapers corresponding to a time half-bandwidth product of 4. (b) Changes of mean segmentation accuracy over the number of switching states, as the number of underlying oscillations varies between 2 and 5. Shaded bounds denote the standard error of the mean around the average accuracy values across 200 repetitions. (c) An example switching state path with 5 underlying oscillations and 31 possible switching states. (d) Estimated switching states with variational inference. VI-A = variational inference with deterministic annealing (orange color); VI-I = variational inference with interpolated densities (blue color).

https://doi.org/10.1371/journal.pcbi.1011395.g006

For example, if there are two oscillations with one at 1 Hz and the other at 10 Hz, there are three possible switching states, and the observation equation takes the form: where the hidden states consist of only the 1 Hz oscillation, of only the 10 Hz oscillation, and of both oscillations. The switching state s_t therefore takes values in {1, 2, 3}. We note that candidate models with more oscillators will inherently be favored by log-likelihood measures, owing to their higher numbers of degrees of freedom. To account for this, we mean-centered the model evidence of each candidate model when initializing the E-step, as described in detail in S3 Appendix.

We compared the segmentation accuracy of VI-A and VI-I given true parameters as we increased the number of underlying oscillations from 2 to 5. For simplicity, we considered the scenario where the model parameters are known, but as in previous examples, variational learning of the model parameters via EM is certainly possible. In practice, target neural oscillations are often characterized with stationary methods first; these models could be used as-is for inference, or to initialize model parameters for variational learning. Fig 6b shows that as the number of switching states increased, the accuracy of VI-A dropped precipitously while VI-I maintained good segmentation accuracy. VI-I was able to segment complicated switching trajectories close to the ground truth, such as in the example displayed in Fig 6c with 31 possible switching states from 5 oscillations, while VI-A failed to do so as shown in Fig 6d. These results demonstrate that our switching inference method is capable of modeling transient oscillations in a manner that scales well to higher-dimensional problems.

Real-world application: Spindle detection

Sleep spindles are transient oscillatory bursts of neural activity observed in EEG recordings during non-rapid eye movement (NREM) sleep, with waxing and waning waveforms in the sigma frequency range (12 Hz–16 Hz) that last 0.5 s–3 s and that occur on a background of slow waves (0.5 Hz–2 Hz) [67]. In this section, we describe a novel spindle detection algorithm to illustrate how our developed switching method can be easily applied in the context of neural signal processing. As we will show, switching state-space models provide a rigorous and probabilistic framework for detection and extraction of spindles from sleep EEG recordings in an unsupervised way.

To model slow oscillations (δ) and spindles (ς) during sleep, we consider two independent state-space oscillators within a set of two candidate models: where and ; and where I₂ is an identity matrix, and for f ∈ {δ, ς}, . An HMM process determines the switching state s_t taking values in {1, 2} across time. This model represents transient sleep spindles by switching between two “neural states”: one with both slow oscillations and spindles, and the other with only slow oscillations. Furthermore, the slow (δ) oscillator parameters are shared between the two candidate models to reflect the stationary dynamics of slow oscillations in the background.

To perform spindle segmentation, we initialized oscillator parameters as described in the Initialization of Gaussian SSM parameters section of Materials and methods. Briefly, the entire time series was modeled with a stationary state-space model including both slow oscillation and spindle oscillators, and parameters were learned through a standard EM algorithm. For variational learning, we derived M-step equations that considered the slow oscillators jointly across both candidate models, to update a single set of parameters for slow waves. In other words, the dynamics of the slow oscillations were modeled by pooling together segments regardless of the presence of sleep spindles. As we describe in S2 Appendix, this inference model structure can be viewed as close approximations to a range of other possible generative models for spindles. Lastly, the HMM switching process was initialized with initial state priors ρ₁ = ρ₂ = 0.5 and a state-transition probability matrix with ϕ₁₁ = ϕ₂₂ = 0.99, ϕ₁₂ = ϕ₂₁ = 0.01 that were updated during the M-steps. As mentioned in the earlier section, models with more hidden oscillation components will be favored by log-likelihood measures during learning. Thus, we again matched the ensemble averages of model evidence when initializing each E-step (see S3 Appendix). For VI-A EM, we used the same temperature decay as before.

We first compared variational learning methods with other inference algorithms on a short segment of EEG data recorded during sleep. In addition to static switching and IMM, we also analyzed a special Gaussian merging algorithm derived for Gaussian SSMs where only the observation matrix is allowed to switch [34]. Fig 7b shows that these inference-only algorithms performed poorly, even in this case with strong spindle activity: their posterior model probabilities were biased to select spindles due to the nested structure of the candidate models. On the other hand, VI-A EM labelled spindles that appeared reasonable given the EEG waveform and spectrogram. In comparison, VI-I EM produced segmentation results that were more accurate with tighter margins around spindles (Fig 7a). We note that the posterior model probabilities for VI-A EM and VI-I EM were more polarized (i.e., closer to either 0 or 1) compared to the other inference methods. This is a feature of the fixed-point iterations as described in Materials and methods, pushing the model responsibility away from 0.5. We show in S4 Appendix that a softer segmentation can be obtained using interpolated densities, while a harder segmentation is also possible via the Viterbi algorithm [68].

Download:

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

Fig 7. Real-world example: Comparisons of switching inference algorithms for sleep spindle detection.

(a) In the top three panels, the spindle activity is visualized using a spectrogram, the original time trace, and after being bandpass filtered within 10 Hz–16 Hz. The fourth panel shows the posterior model probabilities of s_t = 1 estimated by variational EM learning with interpolated densities (VI-I EM, blue color). The margins of spindle events identified by VI-I EM are also marked with vertical dashed lines. The last two panels display the estimated real (blue) and imaginary (magenta) spindle waveforms with 95% confidence intervals from posterior covariances. The learned spindle center frequency is displayed in blue in parentheses. (b) Estimated posterior model probabilities of s_t = 1 by other algorithms in comparison. A model probability closer to 1 suggests the presence of spindles (S = Spindle), while being closer to 0 indicates no spindle (NS = No Spindle). The 0.5 model probability is marked with gray horizontal dashed lines. VI-A EM = variational EM learning with deterministic annealing (orange color); Static = static switching method; IMM = interacting multiple models method; S&S 1991 = the Gaussian merging method in [34].

https://doi.org/10.1371/journal.pcbi.1011395.g007

Besides estimating posterior model probabilities for spindle detection, this approach provides other parametric characterizations of spindle activity, such as the center frequency. For example, in Fig 7a, the algorithm learns that the spindles are centered around 12.7 Hz. In addition, the underlying spindle (and slow oscillation) waveform can be extracted via the inference procedure without any bandpass filtering (see Fig 7a). While the spindle waveform appears similar to the bandpass filtered signal, unlike bandpass filtering we can construct confidence intervals around the estimated spindle activity. We also emphasize that bandpass filtering requires a pre-defined frequency band, and this could introduce serious biases if the spindle frequency gets closer to the boundaries. Furthermore, the current method naturally limited spindle activity to the statistically relevant periods without arbitrary thresholding. Finally, these estimated hidden oscillator states (i.e., the real and imaginary parts) can be readily used to compute instantaneous amplitude and phase for downstream hypothesis testing [69, 70].

Next, we applied the best performing VI-I EM method to three randomly selected 30 s NREM stage 2 sleep recordings, robustly extracting sleep spindles across all time points (Fig 8). In addition, we compared the segmentation results to spindles identified by a wavelet-based automatic spindle detector [71], which has been shown to correspond well to human expert scoring [72]. Every spindle identified by the conventional method was also captured by VI-I EM. However, there were many other spindle events labelled by VI-I EM that were missed by the current “clinical standard”. These events nevertheless showed obvious sigma-frequency power, suggesting a potentially inadequate detection by expert scoring. Furthermore, wavelet-based detection relied on arbitrary thresholds that only selected the strongest portion of a spindle activity. In contrast, actual spindles likely lasted for a longer duration, as evident both from the spindle waveforms estimated by VI-I EM and from the original signal time traces (Fig 8).

Download:

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

Fig 8. Automatic segmentation of sleep spindles using the VI-I EM method.

Three 30 s EEG recordings of NREM-2 sleep were segmented with the variational EM learning method with interpolated densities (VI-I EM) to identify spindles in an unsupervised manner. In each of the (a)-(c) recordings, the three sub-panels visualize spindle activity using a spectrogram, the original time trace, and the estimated real part of spindle waveform with 95% confidence intervals from posterior covariances. The learned spindle center frequencies are displayed in blue in parentheses. The estimated posterior model probabilities for the candidate model with both slow oscillations and spindles are overlaid on the time traces in blue lines. Shaded pink bars indicate spindles identified by a wavelet-based method for comparison. S = Spindle; NS = No Spindle.

https://doi.org/10.1371/journal.pcbi.1011395.g008

These results demonstrate that switching oscillator models with variational learning could detect spindles in an unsupervised manner that may be more reliable than expert scoring. In each recording, VI-I EM learned individualized parameters for the oscillators to best capture the time-varying amplitudes, frequencies, and waveform shapes of the underlying spindles. The posterior model probabilities of the HMM process provide a probabilistic criterion for spindle detection, which is not achievable with a simple fixed threshold on bandpass filtered data as in the de facto practice in sleep research.

Discussion

In this paper we presented an algorithm for inference and learning with switching state-space models. This method holds promise in modeling time-varying dynamics of neural signals, in particular, neural time series data such as EEG. It takes a Variational Bayes approach to approximate the otherwise intractable posterior and enables state estimation and system identification via a generalized EM algorithm. We showed extensive simulation results on how the method can provide accurate segmentation of piece-wise linear dynamic regions, even when the true model parameters were unknown. We also extended the generative model structure to higher dimensions and to more subtle switching transitions instead of jumping between independent state-space models with distinct dynamics. Accordingly, we derived parameter learning rules that were modified to encode this structure and to provide more efficient parameter updates with joint estimation across switching models. Finally, we applied this learning algorithm to a real-world neural signal processing problem of sleep spindle detection, providing excellent detection and extraction of spindles with a rigorous statistical characterization. Taken together, our results add to a growing literature on how switching state-space models can provide powerful statistical inferences to study neural signals.

We were inspired by an earlier, influential paper by Ghahramani and Hinton that applied the variational approximation to switching state-space models [53], but we introduced several important innovations that significantly improved the performance of the variational EM algorithm for these models in practice. In addition, we described in detail a simpler derivation of variational algorithms for such models. Ghahramani and Hinton utilized the same basic form of generative models presented here, but postulated the functional forms of the marginal distributions by introducing additional parameters g_t and h_t and solved for their optimal values by minimizing the KL divergence between the exact and approximate Hamiltonian functions. In contrast, we derived the optimal functional forms of hidden state marginals to maximize the negative variational free energy, which is mathematically equivalent to minimizing the KL divergence. In this way, we obtained expressions for g_t and h_t as summary statistics of the hidden state marginals directly from the variational posterior (see the Variational approximation of hidden state posterior section of Materials and methods).

Given the non-convexity of the problem, robust initialization of the generalized EM iterations, in particular of the fixed-point iterations during the E-step, is crucial for achieving reliable segmentation and avoiding trivial local maxima in the negative free energy. Ghahramani and Hinton proposed a deterministic annealing approach [53], which can be viewed as a form of variational tempering [73]. As discussed in [53] and analyzed in this paper, the annealed learning does not improve performance over traditional inference-only algorithms. We proposed a novel initialization strategy for g_t, based on the insight that interpolated densities provide a local statistical comparison between the candidate models, agnostic of switching dynamics. The interpolated densities therefore provide natural starting values for g_t, which serves as the surrogate observation probabilities in the HMM switching process. Our results showed that this technique significantly improved the segmentation accuracy and system identification of our algorithm compared to the deterministic annealing approach. We found that this more informed initialization was essential in more complicated scenarios, such as time-varying bivariate AR processes and switching oscillator models, or in practical neural signal processing applications, such as sleep spindle segmentation.

This algorithm provides a unique take on inference and learning for switching state-space models. Compared to traditional inference-only algorithms, also known as assumed density filtering (ADF) [74], the variational learning method allows iterative parameter tuning that improves segmentation. These learned parameters also allow us to characterize the properties of the underlying physical system, which may be important in scientific applications such as neural signal processing. Recent algorithms have explored smoothed inference by extending ADF algorithms [75–77]. In contrast, our approach provides an elegant solution for smoothed inference using familiar message-passing algorithms such as the Kalman filter, RTS smoother, and forward-backward algorithm, which are applicable after the conditional dependencies between the real- and discrete-valued hidden states are ignored under the parallel-model variational approximation [53]. This approach is an instance of structured variational inference for time series data [55].

Other variational learning methods have been developed for switching linear dynamical systems [78–80] based on the same principle of approximating intractable posterior with factorized distributions. A notable difference is that these methods assume a single multi-modal hidden state that is modulated by the switching state instead of the multiple parallel state-space models in our generative process. As we explored in the bivariate AR1 simulation and spindle analyses, multiple parallel models can be an effective way to approximate multi-modality: the hidden states in parallel models are suppressed when they are not contributing to the observed data, since the Kalman filtering and smoothing within each model are weighed by the model responsibility h_t. In addition, when multiple models contain hidden states reflected in the observed data, the smoothed estimates across parallel models are closely linked since they are all conditioned on the same observed data. A future study could compare these two approaches and characterize their computational costs and learning performance. Nevertheless, our initialization strategy using interpolated densities could still be helpful in these multi-modal variational learning algorithms. Given that the observed data have a sufficiently high sampling rate relative to the underlying state-transition dynamics, interpolated densities can provide local comparisons among switching states to warm start the E-steps of variational learning algorithms close to the desired optimal solution.

More recent algorithms for switching state-space models embrace a Bayesian approach by using sampling techniques to obtain posterior distributions. Some of these methods can simultaneously learn the number of switching states [39] and more complicated recurrent structures [41]. Recurrent neural networks can also be adapted to support Kalman updates in a deep state-space model and to characterize switching dynamics after training [81]. These methods have varying computational cost and complexity, which require careful adaptation for neural signal processing or biomedical applications. A few studies have applied variational inference on more elaborate model structures and showed meaningful segmentation with neural signals [43–47].

In comparison to the existing applications of switching state-space models, we focus on interpretable linear Gaussian SSMs that can faithfully capture the distinct dynamics observed in neural signals. By applying an intuitive variational approximation on several candidate models constructed from clinical and neuroscience domain knowledge, the variational learning method we described here provides a general solution for segmenting a given time series into neural dynamic states of interest. This method, however, still has limitations encountered in other switching state-space modeling algorithms. First, many of the methodological considerations we described here, including the bias of log-likelihoods in nested structures, apply to other methods. In addition, a single phenomenon or physical process may have multiple switching state-space representations. For example, the process described in the univariate AR1 simulations can have an alternative generative process: two univariate hidden states may be concatenated into a single bivariate hidden state with fixed dynamics, while an observation matrix gets modulated by the switching state to choose particular indices of the bivariate hidden state. Parameter learning and interpretation can be difficult in this scenario if the parameters for the SSM dynamics and the corresponding observation matrices are both learned from the observed data. Thus, we opted to use fixed observation matrices as a constraint to facilitate recovering unique solutions. We nonetheless provide expressions to update the observation matrices in Eq 28, as in some neuroscience applications, switching observation matrices may best encode neural dynamics of interest. In this scenario, they can be learned in a data-driven way with our approach after sharing all SSM parameters across candidate models, analogous to constraining to a stationary SSM dynamic in a single multi-modal hidden state [75].

Our assumed generative structure, consisting of multiple independent dynamical models, permits efficient closed-form computations on the approximate posterior using existing state-space inference algorithms. While the fixed-point iterations and parameter updates can be rigorously derived and familiar methods can be easily plugged in, the non-concave negative free energy poses challenges for inference, including frequent local maxima with polarized model probability estimates that we report here. An alternative to the generalized EM algorithm is blocked Gibbs sampling, which has been successfully applied to switching state-space models [39–41]. A future direction for further development could be to apply blocked Gibbs sampling [82] to the class of neural oscillator models, in a manner that uses the interpolated density to improve computational efficiency.

In the same vein of simplifying the state estimation using using sampling techniques, black-box variational inference is also applicable to switching state-space models [83–85]. Black-box variational approximations are attractive because they provide a general-purpose tool requiring almost no distributional assumptions on the dynamical systems. However, they might not achieve as good performance in cases where we can identify candidate models and their switching structure from prior domain knowledge (see S4 Appendix for an example with the bivariate AR1 simulations). In addition, closed-form inferences tend to have greater computational efficiency than sampling techniques despite the requirement for fixed-point iterations, which converge quickly when the parallel models are reasonably initialized. One clear advantage of sampling techniques is that they are readily applicable in the case of non-linear observation models, such as observations generated by counting processes (e.g., point processes to model neural spiking activity with history dependence [86]). In this context, we note that our approach can also be extended to such non-linear observation models with relatively small effort. For example, point process observations can be handled by replacing Kalman filtering with extended Kalman filtering [87] or unscented Kalman filtering [88], with the rest of the switching inference procedure unaltered.

In general, the variational inference algorithm described here is best suited for scenarios where an informed guess of model parameters can be used for initialization. From a neural signal processing perspective, we are equally concerned about characterizing interpretable hidden dynamics of different neural states and achieving accurate segmentation of time-varying activity due to transitions among them. This intention is reflected in our use of multiple models in parallel, which compete with each other to explain observed data, rather than using a larger array of models with arbitrary structures and parameterizations that are intended to approximate neural dynamics of interest. In contrast, the parallel models we use favor a more parsimonious and informed set of candidate models to obtain interpretable characterization of dynamics and accurate segmentation. The use of the interpolated density is consistent with this premise, since it better exploits prior domain knowledge and the carefully constructed candidate models. In cases where there is an unknown number of switching states, the dynamics are less well-characterized, or the goal is to fit to data as well as possible, a number of sampling-based algorithms such as in [39, 41, 42, 44] could be highly effective.

The sleep spindles analyzed here serve as a good example of scenarios where switching state-space models can be directly applied to achieve much more powerful analyses than standard methods employed in the field. Spindles have been studied with visual scoring by trained human experts for decades [89], until recent efforts to develop reliable detection algorithms [90] in light of their integral roles in normal sleep physiology, memory consolidation, and psychiatric and neurodegenerative diseases [91]. While different in the exact implementation, almost all spindle detection algorithms are fundamentally based on setting an arbitrary threshold and picking strong transient signals in predefined frequency ranges [72], which is severely biased as shown in a recent study [92]. With informed prior knowledge of sleep oscillations, state-space oscillator models can be constructed. The exact parameters of these models should be learned from data since they vary across individuals and different stages of sleep. But one can decide on the number of switching models and their compositions given domain knowledge of sleep and the target spindle detection problem. In this study, we show that the variational learning method can robustly extract sleep spindles from NREM2 sleep in an unsupervised manner, which obviates the need for thresholding. At the same time, variational learning can better characterize spindle properties such as frequency and amplitude by focusing on periods when they are statistically detectable. As mentioned earlier, the use of state-space oscillator models allows one to directly estimate instantaneous amplitude, phase, and frequency without relying on traditional bandpass filtering and Hilbert transform [69], as well as to easily obtain derived measures such as phase–amplitude coupling [70]. In addition, the state-space formulation also allows for filtered estimates of posterior switching state probabilities for online detection of sleep spindles, making it potentially suitable for real-time applications (see S4 Appendix for an example). Thus, there are numerous analytical advantages to a switching state-space modeling approach to detect and extract spindles.

Similar applications can be widely found across neuroscience: burst-suppression induced by anesthetic drugs, epileptic bursts, event-related potentials during cognitive tasks, spike train recordings, or simply behavioral responses over a block of trials [1]. Switching state-space models provide a statistical framework to investigate temporal variations and extract distinct states that are ubiquitous in different datasets. Thus, we advocate for a broad adoption of switching methods in place of previous windowing techniques that limit the analysis of neural signals.

Conclusion

In this paper, we explored the use of switching state-space models for neural signal processing. We demonstrated the promising performance of a variational learning algorithm with interpolated densities, both in simulations and when applied to spindle detection in real sleep EEG data. Our method can be widely applied to neuroscience problems when the observed dynamics show transitions between periods that can be modeled by linear Gaussian state-space models.

Materials and methods

Switching state-space models

We employ one particular form of switching state-space models that allows time-varying dynamics among arbitrary real-valued state-space models with stationary parameters [53]. In addition to its more general formulation compared to modeling time-varying parameters within a single model, this construction offers an elegant solution under the variational Bayesian framework [54] as detailed later. The generative model consists of M linear Gaussian SSMs indexed by numbers from 1 to M which each contain continuous real-valued states, and one HMM whose states take on discrete integer values from 1 to M. Furthermore, the states within each of the Gaussian SSMs are assumed to evolve independently from other models. The discrete HMM too is independent of the Gaussian SSM and decides which one of the M state-space models is generating the current observation data point. The directed acyclic graph corresponding to this conditional independence relation is shown in Fig 9.

Download:

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

Fig 9. Generative structure with parallel switching state-space models.

A directed acyclic graph is shown to represent the conditional independence structure between the M real-valued Gaussian hidden state sequences for m ∈ {1, ⋯ M}, a discrete-valued hidden Markov chain {s_t}, and the observed data {y_t} up to time T. In this generative structure, the observation at a given time point depends only on the hidden states of the M Gaussian models at that point, with the discrete-valued state selecting one of the models to produce the observation, hence the name switching.

https://doi.org/10.1371/journal.pcbi.1011395.g009

The real-valued hidden states in the M Gaussian SSMs are labelled as with m ∈ {1, ⋯, M}, the discrete-valued hidden states of HMM as {s_t} where s_t ∈ {1, ⋯, M}, and the real-valued observations as {y_t}. The joint probability for the observations and hidden states therefore factors nicely as: (6) where all three parts have the familiar forms from classical Gaussian SSM and HMM literature. We provide each of these expressions below.

First of all, conditioned on the discrete hidden state, s_t = m, the observation is simply a multivariate Gaussian variable from the observation equation of the m^th Gaussian SSM. The probability distribution of the observation vector is given by: (7) where R is the observation noise covariance matrix, and G^(m) is the observation matrix of the linear Gaussian SSMs indexed by m.

Secondly, the real-valued states of the M Gaussian SSMs evolve independently and in parallel, with dynamics specified by the model-specific state-transition matrix, F^(m), and state noise covariance, Q^(m), starting from different initial states, : (8)

Lastly, the discrete-valued switching state evolves according to the HMM specified by the M × M state-transition matrix ϕ and initial state probabilities ρ, independent of the other M Gaussian SSMs with real-valued states: (9)

Intractable posterior of hidden states.

With the generative model defined, state estimation and system identification problems need to be solved. As reviewed above, both problems are encompassed in the EM algorithm, which alternates between 1) finding the posterior distribution of the hidden states given the current values of model parameters, and 2) optimizing the model parameters given expectations of hidden states under the posterior distribution from (1). However, the hidden state posterior here cannot be computed explicitly. To illustrate this difficulty, we start with the complete log-likelihood of observations and hidden states: (10) where , for m = 1, ⋯, M, is the set of parameters of the distributions that fully characterize the switching state-space generative model in Fig 9. Note that the HMM process does not have a separate observation probability parameter (e.g. ψ). Once the switching state is fixed to s_t = m, the observation equation dynamic of the corresponding m^th Gaussian SSM in the second term of Eq 10 specifies the likelihood of generating the observation at that time point.

State estimation, i.e., the E-step, is computationally difficult because of the products between the discrete switching states of the HMM and the real-valued states of the Gaussian SSMs. This conditional dependence between the hidden states (the ‘two’ causes) when conditioning on the observations (the ‘common’ effect) has been described in the causal inference literature [102, pp.40–45]. Graphically, conditioning on the children of v-structures results in a graph with parents connected with undirected edges for the new conditional (in)dependence relation (Fig 10a). The new edges introduced make it challenging to express the exact hidden state posterior in a factorized form. In particular, the individual Markov independence property within each Gaussian SSM is no longer present, since conditioning on alone does not make conditionally independent from the history due to the new viable paths through the other Gaussian SSMs. Then exact inferences can only be solved after considering the combinations of states across all Gaussian SSMs and HMM. For example, the optimal filtering problem alone requires a bank of elemental estimators with the size of the bank increasing exponentially with time [32], so naturally the smoothing problem, i.e., state estimation of , becomes intractable. Without the hidden state posterior distribution estimated, system identification, i.e., parameter learning, cannot be performed either under the classical EM algorithm [93].

Download:

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

Fig 10. Posteriors of switching state-space models.

(a) True posterior distribution. This is the resultant graph encoding the conditional independence relation after conditioning on the observed data {y_t} up to time T. The new edges between the hidden states make exact inference intractable. (b) Approximate posterior distribution. Compared to the true posterior, a structured approximation decouples the hidden Gaussian state-space models from each other and from the switching state. On this approximate distribution, efficient closed-form inference can be performed. The marginal distributions of the Gaussian hidden states {x_t} and the discrete-valued switching state {s_t} are now inter-dependent through variational summary statistics and .

https://doi.org/10.1371/journal.pcbi.1011395.g010

Variational approximation of hidden state posterior.

One possible solution is to use the Variational Bayes technique to approximate the true posterior of the hidden states, [53]. The idea is to work within a subspace of tractable posterior probability distributions defined over the hidden states, and choose the optimal approximation, , based on a lower bound on the marginal log-likelihood of observations: (11) known as the negative variational free energy [54, 103, 104]. Since the distribution q is approximating the true posterior, the conditioning on {y_t} in the expression is already implied, therefore omitted in all q notations.

The choice of maximizes the negative free energy so that it reaches the true log-likelihood. However, since exact E-step is intractable due to the structure of the posterior distribution, we consider the following family of distributions that factor over the HMM and Gaussian SSMs, i.e., (12)

In other words, we get rid of the dependence between the switching state and real-valued states by limiting to distributions within this functional subspace. The graphical model corresponding to this conditional independence relation is shown in Fig 10b. It only remains to find the optimal approximate posterior that is the closest to the true posterior . Since negative free energy is a lower bound on the marginal log-likelihood, we use the negative free energy as a measure of one-sided closeness to choose the optimal approximation as following: (13)

We proceed to a direct derivation of the optimal functional forms of the approximate posterior. Within the functional subspace , the negative free energy can be simplified: (14) where Following the development in [54], the optimal functional forms for the variational log-posteriors to maximize Eq 14 are given as: (15) (16) with z and z′ being normalizing constants so that the posteriors sum or integrate to 1.

Substituting the expressions from Eq 10 into Eq 15 reveals that the variational log-posterior of the switching state can be written as: (17) (18)

This equation is identical to the log-posterior density of the discrete-valued states in an HMM with observation probabilities proportional to . In other words, functions as the model log-evidence of the m^th state-space model in generating the data at time t in the HMM process. Thus, the desired posterior distribution, i.e., q(s_t = m) at individual time points, can be expressed in terms of the forward-backward variables (α_t(m), β_t(m)) from the forward-backward algorithm [94] and computed efficiently (S3 Appendix).

Similarly, using the expressions from Eqs 10 and 16, we obtain the following variational log-posterior of the real-valued states in the Gaussian SSMs: (19) (20) that factorizes over the M parallel state-space models, i.e., (21)

Each of the log-posteriors within the sum takes the familiar Gaussian density form: (22) that can be computed using a forward pass through Kalman filter, followed by a backward pass through RTS smoother to get the posterior mean and covariance matrices, , as shown in S3 Appendix.

Given the third term in Eq 22, Kalman smoothing here should be modified to use as the observation noise covariance. Scaling the observation noise leads to an intuitive interpretation: when , resulting in high effective observation noise, the hidden state update that follows one-step prediction during Kalman filtering is effectively skipped for that time point; when , the hidden state is updated as in recursions without switching. In this sense, represents the model responsibility of the m^th state-space model in generating the observed data at the time point t.

Fixed-point iterations.

The variables and are the bridges between the distributions p and q, replacing the stochastic dependencies with deterministic quantities. The variational log-posterior of the discrete switching states {s_t} in Eq 18 involves some summary statistics of the real-valued states w.r.t. the approximate distribution q, and vice-versa in Eq 20. This motivates an iterative procedure to update the forward-backward variables and posterior means and covariances in a cyclic fashion to find the optimal variational parameters :

1. Run forward-backward algorithm to update {α_t(m), β_t(m)}.
2. Compute from the updated {α_t(m), β_t(m)}.
3. Run Kalman filter and RTS smoother to update .
4. Compute from the updated .

During these iterations, the negative variational free energy can be computed using log-likelihood values from the Kalman filter and forward-backward algorithm as shown in S1 Appendix. We stop the fixed-point iterations when the model evidence and model responsibility, i.e., , have stabilized or when the negative free energy stops increasing (see S3 Appendix).

Under the assumed generative structure in Fig 9, the true posterior p is expected to contain polarized values (e.g., with model probabilities close to 1 or 0) due to a few factors. First, each transition between candidate models introduces a discontinuity in the observation sequence at the switch point, as the realization trajectories from different SSMs are distinct from each other. Second, since the SSM hidden state trajectories consist of multiple time points, they reside in a high dimensional space. The dimensionality further increases as SSM state dimension increases. Because of this high-dimensionality, the separation between any two candidate models becomes very substantial. Third, real recordings often have infrequent transitions between distinct dynamics. This empirical skewness leads to a high probability to stay within the same candidate model. Such probability propagates as a strong prior across time points amplifying the effect of the other factors in polarizing the values of posterior estimates.

Correspondingly, the fixed-point iterations could polarize the values of through the reciprocal interdependence between and . The observation noise covariance during Kalman smoothing is scaled by the model responsibility as to produce the model log-evidence following Eq 18. This log-evidence is then used by the forward-backward algorithm to compute the model responsibility for the next pass of fixed-point iterations. Thus, this cycle could amplify the for the best candidate model toward 1 while pushing the others close to 0, effectively assigning each time point to one of the candidate models instead of maintaining a mixture of models. This empirical behavior of the fixed-point iterations appears similar to automatic relevance determination [105], with treated as features weighted by at each time point. However, the log-likelihood expression and therefore the pruning mechanism of the respective solutions are different.

It should be emphasized that the negative variational free energy is not jointly concave w.r.t. the values of and . In fact, the free energy landscape contains abundant local maxima, and the fixed-point iterations stop updating whenever a local maximum is reached. Empirically, the fixed-point iterations indeed recover polarized even when the segmentation accuracy is low, i.e., when sub-optimal local maxima are reached. This behavior makes proper initialization the most important factor in the performance of the current variational inference approach. We address this technical challenge below in section Initialization of fixed-point iterations.

Parameter learning using generalized EM algorithm.

Once the approximation of the hidden state posterior can be computed efficiently, one can employ an EM-like algorithm to learn the model parameters of the Gaussian SSMs and HMM. To this end, starting from an initial guess of model parameters, we first repeat the fixed-point iterations until convergence to find the optimal variational parameters, , which parameterize the approximate hidden state posterior (E-step). Next, we update the model parameters of the individual Gaussian SSMs and HMM to maximize the negative variational free energy using the converged variational parameters (M-step). Iterating between these two steps constitutes an instance of generalized EM algorithm where in lieu of data log-likelihood, a lower bound on the log-likelihood is maximized over successive iterations [106]. It should be noted that while derivations differ, the overall form of this algorithm is the same as in [53]. A flowchart of this variational learning algorithm is shown in Fig 1.

In the following section, we derive closed-form update equations for the model parameters of the Gaussian SSMs and HMM, given the approximate hidden state posterior. We note that these M-step equations can be derived from the complete log-likelihood expression in Eq 10 under the expectation w.r.t. the variational approximate distribution q, because only the first term in Eq 14 depends on the parameters θ. It is therefore sufficient to maximize this term for an optimization problem analogous to Eq 13 but over θ.

Specifically, at the i^th EM iteration, with the converged variational parameters, , we compute the following summary statistics: S3 Appendix provides full recursions used to compute these summary statistics, which are needed to take the expectation of the complete log-likelihood in Eq 10: (23)

The optimal parameters, for m = 1, ⋯, M, which maximize Eq 23, can be obtained in closed-form by taking individual partial derivatives. The parameter subset can be updated for each Gaussian SSM through the usual equations described in [93]: (24) (25) (26) (27) with a slight exception for G^(m) due to the product with the switching state: (28)

We note that since the M Gaussian SSMs have independent dynamics, these updates can be completed efficiently in parallel.

Similarly, the usual update equations for an HMM can be used to update {ρ, ϕ} [101]: (29) (30)

As noted earlier, the observation probabilities of the HMM are not explicitly updated as a parameter, since the variational model evidence, , is converged through the fixed-point iterations and used as point estimates of the (log) observation probability to relate the hidden states of Gaussian SSMs to the observation at each time point.

Finally, the update equation for R pools the posterior estimates of hidden states across the M Gaussian SSMs: (31) where captures the contribution from the m^th model: (32)

This update equation for R is an instance of joint estimation of parameters shared among Gaussian SSMs. In a more general case where each state-space model has its individual observation noise covariance R^(m), the update equation takes the form: (33)

It is also possible for other Gaussian SSM parameters to be partially shared among different models. Closed-form update equations in these cases can be derived. This extension exploits the flexibility of the variational Bayesian method to accommodate different generative models. Two such examples are studied in this paper, and we provide the derivations of their update equations in S2 Appendix.

Initialization of fixed-point iterations.

As shown above, the M-step equations provide optimal updates of parameters in a convex optimization setting given the multivariate Gaussian distribution q and the posterior hidden state estimates. In contrast, the fixed-point iterations in the E-step are not guaranteed to achieve globally maximal negative free energy due to the non-concavity of the negative free energy w.r.t. the variational parameters. This makes the final approximate distribution at the end of fixed-point iterations very sensitive to initialization. Moreover, while the fixed-point iterations are defined explicitly step-by-step, it is not obvious how to initialize these iterations. Thus, here we describe two practical initialization approaches.

Deterministic annealing. This entropy-motivated initialization technique was proposed along with the variational framework to alleviate getting trapped in local maxima during the fixed-point iterations [53, 107]. Specifically, equal model responsibilities are used as the initial at the onset of the first E-step, followed by RTS smoothing and forward-backward algorithm to compute and with definitions modified by a temperature parameter : (34) (35)

The temperature cools over successive fixed-point iterations via the decay function . In essence, a larger allows the iterations to explore a bigger subspace with high negative free energy entropy, which gradually decreases in trying to identify an optimal approximate distribution. This initialization technique was necessary for the variational approximation to produce any reasonable results [53].

After the first E-step, is not re-initialized with equal model responsibilities to allow a warm start in subsequent iterations. However, this makes the algorithm severely limited by the results of the first round of fixed-point iterations. Under this initialization scheme, once the for the m^th model gets close to zero, that model cannot regain responsibility for that time point in subsequent iterations. One may choose to reset the to be equal across models at every E-step, but that tends to select only one of the Gaussian SSMs for the entire duration as a local maximum. Regardless, with the model parameters fixed during the E-step, potential local maxima are likely selected, since the annealing is equivalent across the M Gaussian SSMs. Also, initializing with equal responsibilities assumes all models to explain all time points equally well, which is certainly far from the true posterior p. Thus, a better initialization should discern among the Gaussian SSMs and try to initialize closer to the global maximum.

Interpolated densities. Here we propose a different initialization technique that statistically compares the M Gaussian SSMs for their likelihoods of generating the observation at each time point. Specifically, we initialize using the interpolated density [108, 109], i.e., the probability density of the observation at a given time point conditioned on all other time points under the m^th model. (36)

In other words, we attempt to initialize based on how well each Gaussian SSM predicts the current observation, y_t, based on all the past and future observations. We expect this informative initialization of to be close to the global maximum in general. Since conventional Kalman filtering and RTS smoothing cannot compute the interpolated density easily, we utilize a different implementation (see S3 Appendix) [109].

This new initialization technique is well-grounded in the view of as the log-evidence of generating the observation at each time point within the HMM of a discrete-valued state s_t = m, with m ∈ {1, ⋯, M}. In the absence of known Gaussian SSM states, the next best choice is to evaluate which model dynamic provides the closest interpolation from all other time points for the current observation. It can also be seen as a “smoothing” extension of using filtered densities in place of the HMM observation probabilities in the early switching state-space model inference literature [34, 35, 75, 110].

As extensively analyzed in the Segmentation with posterior inference and Segmentation with parameter learning sections in Results, variational inference and learning with interpolated densities substantially improve over deterministic annealing and offer greater segmentation accuracy compared to the other switching inference methods. However, unlike the definition of in Eq 18 that is comparable across arbitrary Gaussian SSMs due to the identical R, interpolated densities are bonafide normalized distribution density functions. Therefore, the model-specific parameters (e.g. Q^(m)), especially the hidden state dimensionality, could bias the interpolated densities during the initialization. A simple heuristic is suggested in S3 Appendix to address this bias if present and shows robust performance in the spindle detection problem.

Supporting information

Acknowledgments

We would like to thank Amanda M. Beck for useful discussions on modeling neural signals with oscillators.