SNLDS

Switching non-linear dynamical system trained with collapsed amortized variational inference (CAVI-SNLDS).

Warning

Experimental. On a 54-series sample of TSB-SEG (TSSB, UTSA, SKAB, MoCap), with the default parameters, SNLDS scores below a random segmentation (bidirectional covering 0.074 against 0.377, ARI 0.066 against 0.208): its labels change about 200 times more often than the true regimes. Training five times longer did not help on the two series tried. It is kept for further study, not as a recommended detector.

Description

SNLDS models a time series with a discrete regime \(s_t \in \{1, \dots, K\}\), a continuous latent state \(z_t\) and an observation \(x_t\). Each regime has its own non-linear transition \(p(z_t \mid z_{t-1}, s_t)\), an MLP; the regime switches follow \(p(s_t \mid s_{t-1}, x_{t-1})\), and \(x_t\) is emitted from \(z_t\) by a shared network.

Training maximises an evidence lower bound. An RNN infers the continuous states \(q(z_{1:T} \mid x_{1:T})\), and the discrete states are integrated out exactly with the forward-backward algorithm.

After training, each time step receives the regime of highest posterior probability \(p(s_t \mid x_{1:T}, z_{1:T})\).

Two devices from the paper prevent the collapse onto a single regime: a cross-entropy regulariser between the regime posterior and a uniform prior, with weight \(\beta\), and a temperature \(\tau\) on the discrete transitions. Both start large and decay, \(\beta\) first, then \(\tau\) down to 1. tsseg follows the schedule of the paper’s Reacher experiment (\(\beta_0 = \tau_0 = 1000\), decay by 0.975 every 500 steps after steps 50,000 and 100,000, out of 300,000), scaled to n_train_steps: the decays start after n_train_steps / 6 and n_train_steps / 3 steps, with a factor 0.975 every n_train_steps / 600 steps.

Each gradient step uses batch_size windows of train_window points drawn at random positions of the series, which bounds the cost of a step; prediction runs on the whole series. The paper trains on minibatches of 32 short sequences (50 to 100 steps) from many trajectories; tsseg draws its windows from the single series it segments.

Type: state detection

Supervision: semi-supervised (n_states is the number of regimes)

Scope: univariate and multivariate

Requires: TensorFlow, TensorFlow Probability and tf-keras (pip install tsseg[snlds])

Note

Training is slow. A step costs \(O(BLK^2)\) for batch_size \(B\) windows of \(L\) points, independent of the length of the series: about 0.1 s on one CPU thread with the defaults, so about 3 minutes for the 2,000 default steps. The first call also spends a few seconds compiling the training graph.

Parameters

Name	Type	Default	Description
n_states	int	3	Number of discrete regimes \(K\).
hidden_dim	int	8	Dimension of the continuous latent state \(z_t\).
rnn_dim	int	4	Hidden size of the inference RNN.
rnn_type	str	"simplernn"	Inference RNN cell ("gru", "lstm", "simplernn").
n_train_steps	int	2000	Number of gradient steps; the schedules scale with it.
learning_rate	float	1e-4	Adam learning rate.
batch_size	int	8	Training windows per gradient step.
train_window	int / None	256	Length of the training windows (None: the whole series).
objective	str	"elbo"	Training objective ("elbo" or "iwae").
use_temperature_annealing	bool	True	Anneal the temperature of the discrete transitions down to 1.
temperature_init	float	1000.0	Initial temperature.
use_cross_entropy_reg	bool	True	Add the cross-entropy regulariser on regime usage.
xent_init	float	1000.0	Initial weight of the regulariser.
random_state	int / None	42	Seed of the training windows and of TensorFlow (a few labels can still differ between calls).
verbose	int	0	Log the objective every tenth of training when >= 1.
axis	int	0	Time axis.

Usage

from tsseg.algorithms import SNLDSDetector

detector = SNLDSDetector(n_states=3)
labels = detector.fit_predict(X)

Implementation: wrapper around the authors’ TensorFlow code (Google Research, Apache License 2.0), vendored in tsseg/algorithms/snlds/vendor with relative imports and a Keras 3 fix of the inference RNN.

Reference: Dong, Seybold, Murphy & Bui (2020), Collapsed Amortized Variational Inference for Switching Nonlinear Dynamical Systems, ICML.

API reference

SNLDS (Switching Non-Linear Dynamical System) detector for time series segmentation.

Wraps the CAVI-SNLDS model from:

Dong, Z., Seybold, B., Murphy, K., & Bui, H. (2020). Collapsed Amortized Variational Inference for Switching Nonlinear Dynamical Systems. ICML 2020.

The model learns discrete regime labels via variational inference with a forward-backward algorithm over discrete states, making it a natural state detector.

class tsseg.algorithms.snlds.detector.SNLDSDetector(*, n_states=3, hidden_dim=8, rnn_dim=4, rnn_type='simplernn', n_train_steps=2000, learning_rate=0.0001, batch_size=8, train_window=256, objective='elbo', use_temperature_annealing=True, temperature_init=1000.0, use_cross_entropy_reg=True, xent_init=1000.0, random_state=42, verbose=0, axis=0)

Bases: BaseSegmenter

State detector based on Switching Non-Linear Dynamical Systems (CAVI-SNLDS).

Uses collapsed amortized variational inference to learn discrete regimes in time series data. After training, the discrete state posterior p(s_t | x_{1:T}, z_{1:T}) is used to assign a regime label to each time-step.

Warning

Experimental. On a sample of TSB-SEG, with the default parameters, SNLDS scores below a random segmentation: its labels change far more often than the true regimes. See the documentation page.

Parameters:

• n_states (int) – Number of discrete regimes (K).

• hidden_dim (int) – Dimension of the continuous latent state z_t.

• rnn_dim (int) – Hidden dimension of the posterior inference RNN.

• rnn_type (str) – RNN cell type for inference network (“gru”, “lstm”, “simplernn”).

• n_train_steps (int) – Number of gradient steps for training.

• learning_rate (float) – Learning rate for Adam optimiser.

• batch_size (int) – Number of training windows per gradient step.

• train_window (int | None) – Length of the training windows, drawn at random positions of the series at each step. None (or a value at least the length of the series) trains on the whole series, with batch_size copies of it.

• objective (str) – Training objective (“elbo” or “iwae”).

• use_temperature_annealing (bool) – Anneal the temperature of the discrete transitions from temperature_init down to 1 (see Notes).

• temperature_init (float) – Initial temperature.

• use_cross_entropy_reg (bool) – Add the cross-entropy regulariser on regime usage, with a weight that decays from xent_init towards 0 (see Notes).

• xent_init (float) – Initial weight of the cross-entropy regulariser.

• random_state (int | None) – Seed of the training windows, and of TensorFlow (global seed, set at the start of fit and predict). TensorFlow’s sampling is not fully reproducible: a few labels (about 0.1 % on MoCap) can differ between two calls.

• verbose (int) – Log the objective ten times during training when >= 1.

• axis (int) – Time axis of the input array.

Notes

The schedules follow the paper’s Reacher experiment (appendix), scaled to n_train_steps: there, the regulariser weight starts at 1,000 and decays by a factor 0.975 every 500 steps after step 50,000, the temperature does the same after step 100,000 (down to 1), and training lasts 300,000 steps. Here the weight decays after n_train_steps / 6 steps and the temperature after n_train_steps / 3, by 0.975 every n_train_steps / 600 steps, so that both reach the same final values.

Training on random windows bounds the cost of a step, which otherwise runs a forward-backward pass over the whole series. The paper trains on minibatches of 32 short sequences (50 to 100 steps) from many trajectories; here the windows come from the single series to segment. Prediction uses the whole series.

The vendored model code comes from Google Research (Apache License 2.0), see tsseg/algorithms/snlds/vendor.

References

[1]

Dong, Z., Seybold, B., Murphy, K., & Bui, H. (2020). Collapsed Amortized Variational Inference for Switching Nonlinear Dynamical Systems. ICML.

set_fit_request(*, axis: bool | None | str = '$UNCHANGED$') → SNLDSDetector

Configure whether metadata should be requested to be passed to the fit method.

Note that this method is only relevant when this estimator is used as a sub-estimator within a meta-estimator and metadata routing is enabled with enable_metadata_routing=True (see sklearn.set_config()). Please check the User Guide on how the routing mechanism works.

The options for each parameter are:

• True: metadata is requested, and passed to fit if provided. The request is ignored if metadata is not provided.

• False: metadata is not requested and the meta-estimator will not pass it to fit.

• None: metadata is not requested, and the meta-estimator will raise an error if the user provides it.

• str: metadata should be passed to the meta-estimator with this given alias instead of the original name.

The default (sklearn.utils.metadata_routing.UNCHANGED) retains the existing request. This allows you to change the request for some parameters and not others.

Added in version 1.3.

Parameters:

axis (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for axis parameter in fit.

Returns:

self – The updated object.

Return type:

object

set_predict_request(*, axis: bool | None | str = '$UNCHANGED$') → SNLDSDetector

Configure whether metadata should be requested to be passed to the predict method.

Note that this method is only relevant when this estimator is used as a sub-estimator within a meta-estimator and metadata routing is enabled with enable_metadata_routing=True (see sklearn.set_config()). Please check the User Guide on how the routing mechanism works.

The options for each parameter are:

• True: metadata is requested, and passed to predict if provided. The request is ignored if metadata is not provided.

• False: metadata is not requested and the meta-estimator will not pass it to predict.

• None: metadata is not requested, and the meta-estimator will raise an error if the user provides it.

• str: metadata should be passed to the meta-estimator with this given alias instead of the original name.

The default (sklearn.utils.metadata_routing.UNCHANGED) retains the existing request. This allows you to change the request for some parameters and not others.

Added in version 1.3.

Parameters:

axis (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for axis parameter in predict.

Returns:

self – The updated object.

Return type:

object