Gaussian HMM: Cross-validation and Model Selection#

A Gaussian HMM has emissions of the form,

(1)#\[\begin{align} p(y_t \mid z_t, \theta) &= \mathcal{N}(y_t \mid \mu_{z_t}, \Sigma_{z_t}) \end{align}\]

where the emission parameters \(\theta = \{(\mu_k, \Sigma_k)\}_{k=1}^K\) include the means and covariances for each of the \(K\) discrete states.

Dynamax implements a variety of Gaussian HMMs with different constraints on the parameters (e.g. diagonal, spherical, and tied covariances). It also includes prior distributions on the parameters. For example, it uses a conjugate normal-inverse Wishart (NIW) prior for the standard case.

This notebook shows how to:

Fit such models using expectation-maximization (EM)
Use cross-validation to choose the number of discrete states
Use the log probability of held-out data to choose among different models

Setup#

from functools import partial

import jax.numpy as jnp
import jax.random as jr
import matplotlib.pyplot as plt
from jax import vmap

from dynamax.hidden_markov_model import GaussianHMM
from dynamax.hidden_markov_model import DiagonalGaussianHMM
from dynamax.hidden_markov_model import SphericalGaussianHMM
from dynamax.hidden_markov_model import SharedCovarianceGaussianHMM
from dynamax.utils.plotting import CMAP, COLORS, white_to_color_cmap

Helper functions for plotting#

Generate sample data#

As in the preceding notebooks, we start by sampling data from the model. Here, we add a slight wrinkle: we will sample training and test data, where the latter is only used for model selection.

num_train_batches = 3
num_test_batches = 1
num_timesteps = 100

# Make an HMM and sample data and true underlying states
true_num_states = 5
emission_dim = 2
hmm = GaussianHMM(true_num_states, emission_dim)

# Specify parameters of the HMM
initial_probs = jnp.ones(true_num_states) / true_num_states
transition_matrix = 0.80 * jnp.eye(true_num_states) \
    + 0.15 * jnp.roll(jnp.eye(true_num_states), 1, axis=1) \
    + 0.05 / true_num_states
emission_means = jnp.column_stack([
    jnp.cos(jnp.linspace(0, 2 * jnp.pi, true_num_states + 1))[:-1],
    jnp.sin(jnp.linspace(0, 2 * jnp.pi, true_num_states + 1))[:-1],
    jnp.zeros((true_num_states, emission_dim - 2)),
    ])
emission_covs = jnp.tile(0.1**2 * jnp.eye(emission_dim), (true_num_states, 1, 1))
        
true_params, _ = hmm.initialize(initial_probs=initial_probs,
                                transition_matrix=transition_matrix,
                                emission_means=emission_means,
                                emission_covariances=emission_covs)

# Sample train, validation, and test data
train_key, val_key, test_key = jr.split(jr.PRNGKey(0), 3)
f = vmap(partial(hmm.sample, true_params, num_timesteps=num_timesteps))
train_true_states, train_emissions = f(jr.split(train_key, num_train_batches))
test_true_states,  test_emissions  = f(jr.split(test_key, num_test_batches))

---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
Cell In[4], line 22
emission_means = jnp.column_stack([
   jnp.cos(jnp.linspace(0, 2 * jnp.pi, true_num_states + 1))[:-1],
   jnp.sin(jnp.linspace(0, 2 * jnp.pi, true_num_states + 1))[:-1],
   jnp.zeros((true_num_states, emission_dim - 2)),
   ])
emission_covs = jnp.tile(0.1**2 * jnp.eye(emission_dim), (true_num_states, 1, 1))
---> 22 true_params, _ = hmm.initialize(initial_probs=initial_probs,
                               transition_matrix=transition_matrix,
                               emission_means=emission_means,
                               emission_covariances=emission_covs)
# Sample train, validation, and test data
train_key, val_key, test_key = jr.split(jr.PRNGKey(0), 3)

File ~/work/dynamax/dynamax/dynamax/hidden_markov_model/models/gaussian_hmm.py:651, in GaussianHMM.initialize(self, key, method, initial_probs, transition_matrix, emission_means, emission_covariances, emissions)
params["initial"], props["initial"] = self.initial_component.initialize(key1, method=method, initial_probs=initial_probs)
params["transitions"], props["transitions"] = self.transition_component.initialize(key2, method=method, transition_matrix=transition_matrix)
--> 651 params["emissions"], props["emissions"] = self.emission_component.initialize(key3, method=method, emission_means=emission_means, emission_covariances=emission_covariances, emissions=emissions)
return ParamsGaussianHMM(**params), ParamsGaussianHMM(**props)

File ~/work/dynamax/dynamax/dynamax/hidden_markov_model/models/gaussian_hmm.py:83, in GaussianHMMEmissions.initialize(self, key, method, emission_means, emission_covariances, emissions)
elif method.lower() == "prior":
   this_key, key = jr.split(key)
---> 83     prior = NormalInverseWishart(self.emission_prior_mean, self.emission_prior_conc,
                                self.emission_prior_df, self.emission_prior_scale)
   (_emission_covs, _emission_means) = prior.sample(seed=this_key, sample_shape=(self.num_states,))
else:

File /opt/hostedtoolcache/Python/3.9.19/x64/lib/python3.9/site-packages/decorator.py:232, in decorate.<locals>.fun(*args, **kw)
if not kwsyntax:
   args, kw = fix(args, kw, sig)
--> 232 return caller(func, *(extras + args), **kw)

File /opt/hostedtoolcache/Python/3.9.19/x64/lib/python3.9/site-packages/tensorflow_probability/substrates/jax/distributions/distribution.py:342, in _DistributionMeta.__new__.<locals>.wrapped_init(***failed resolving arguments***)
# Note: if we ever want to have things set in `self` before `__init__` is
# called, here is the place to do it.
self_._parameters = None
--> 342 default_init(self_, *args, **kwargs)
# Note: if we ever want to override things set in `self` by subclass
# `__init__`, here is the place to do it.
if self_._parameters is None:
 # We prefer subclasses will set `parameters = dict(locals())` because
 # this has nearly zero overhead. However, failing to do this, we will
 # resolve the input arguments dynamically and only when needed.

File ~/work/dynamax/dynamax/dynamax/utils/distributions.py:133, in NormalInverseWishart.__init__(self, loc, mean_concentration, df, scale)
self._df = df
self._scale = scale
super(NormalInverseWishart, self).__init__([
--> 133     InverseWishart(df, scale),
   lambda Sigma: tfd.MultivariateNormalFullCovariance(loc, Sigma / mean_concentration)
])
self._parameters = dict(loc=loc, mean_concentration=mean_concentration, df=df, scale=scale)

File /opt/hostedtoolcache/Python/3.9.19/x64/lib/python3.9/site-packages/decorator.py:232, in decorate.<locals>.fun(*args, **kw)
if not kwsyntax:
   args, kw = fix(args, kw, sig)
--> 232 return caller(func, *(extras + args), **kw)

File /opt/hostedtoolcache/Python/3.9.19/x64/lib/python3.9/site-packages/tensorflow_probability/substrates/jax/distributions/distribution.py:342, in _DistributionMeta.__new__.<locals>.wrapped_init(***failed resolving arguments***)
# Note: if we ever want to have things set in `self` before `__init__` is
# called, here is the place to do it.
self_._parameters = None
--> 342 default_init(self_, *args, **kwargs)
# Note: if we ever want to override things set in `self` by subclass
# `__init__`, here is the place to do it.
if self_._parameters is None:
 # We prefer subclasses will set `parameters = dict(locals())` because
 # this has nearly zero overhead. However, failing to do this, we will
 # resolve the input arguments dynamically and only when needed.

File ~/work/dynamax/dynamax/dynamax/utils/distributions.py:51, in InverseWishart.__init__(self, df, scale)
cho_scale = jnp.linalg.cholesky(scale)
inv_scale_tril = solve_triangular(cho_scale, eye, lower=True)
---> 51 super().__init__(
   tfd.WishartTriL(df, scale_tril=inv_scale_tril),
   tfb.Chain([tfb.CholeskyOuterProduct(),
              tfb.CholeskyToInvCholesky(),
              tfb.Invert(tfb.CholeskyOuterProduct())]))
self._parameters = dict(df=df, scale=scale)

File /opt/hostedtoolcache/Python/3.9.19/x64/lib/python3.9/site-packages/decorator.py:232, in decorate.<locals>.fun(*args, **kw)
if not kwsyntax:
   args, kw = fix(args, kw, sig)
--> 232 return caller(func, *(extras + args), **kw)

File /opt/hostedtoolcache/Python/3.9.19/x64/lib/python3.9/site-packages/tensorflow_probability/substrates/jax/distributions/distribution.py:342, in _DistributionMeta.__new__.<locals>.wrapped_init(***failed resolving arguments***)
# Note: if we ever want to have things set in `self` before `__init__` is
# called, here is the place to do it.
self_._parameters = None
--> 342 default_init(self_, *args, **kwargs)
# Note: if we ever want to override things set in `self` by subclass
# `__init__`, here is the place to do it.
if self_._parameters is None:
 # We prefer subclasses will set `parameters = dict(locals())` because
 # this has nearly zero overhead. However, failing to do this, we will
 # resolve the input arguments dynamically and only when needed.

File /opt/hostedtoolcache/Python/3.9.19/x64/lib/python3.9/site-packages/tensorflow_probability/substrates/jax/distributions/transformed_distribution.py:244, in _TransformedDistribution.__init__(self, distribution, bijector, kwargs_split_fn, validate_args, parameters, name)
self._zero = tf.constant(0, dtype=tf.int32, name='zero')
# We don't just want to check isinstance(JointDistribution) because
# TransformedDistributions with multipart bijectors are effectively
# joint but don't inherit from JD. The 'duck-type' test is that
# JDs have a structured dtype.
--> 244 dtype = self.bijector.forward_dtype(self.distribution.dtype)
self._is_joint = tf.nest.is_nested(dtype)
super(_TransformedDistribution, self).__init__(
   dtype=dtype,
   reparameterization_type=self._distribution.reparameterization_type,
   (...)
   parameters=parameters,
   name=name)

File /opt/hostedtoolcache/Python/3.9.19/x64/lib/python3.9/site-packages/tensorflow_probability/substrates/jax/bijectors/bijector.py:1705, in Bijector.forward_dtype(self, dtype, name, **kwargs)
 input_dtype = nest_util.broadcast_structure(
     self.forward_min_event_ndims, self.dtype)
else:
 # Make sure inputs are compatible with statically-known dtype.
-> 1705   input_dtype = nest.map_structure_up_to(
     self.forward_min_event_ndims,
     lambda x: dtype_util.convert_to_dtype(x, dtype=self.dtype),
     nest_util.coerce_structure(self.forward_min_event_ndims, dtype),
     check_types=False)
output_dtype = self._forward_dtype(input_dtype, **kwargs)
try:
 # kwargs may alter dtypes themselves, but we currently require
 # structure to be statically known.

File /opt/hostedtoolcache/Python/3.9.19/x64/lib/python3.9/site-packages/tensorflow_probability/python/internal/backend/jax/nest.py:324, in map_structure_up_to(shallow_structure, func, *structures, **kwargs)
def map_structure_up_to(shallow_structure, func, *structures, **kwargs):
--> 324   return map_structure_with_tuple_paths_up_to(
     shallow_structure,
     lambda _, *args: func(*args),  # Discards path.
     *structures,
     **kwargs)

File /opt/hostedtoolcache/Python/3.9.19/x64/lib/python3.9/site-packages/tensorflow_probability/python/internal/backend/jax/nest.py:353, in map_structure_with_tuple_paths_up_to(shallow_structure, func, expand_composites, *structures, **kwargs)
for input_tree in structures:
 assert_shallow_structure(
     shallow_structure, input_tree, check_types=check_types)
--> 353 return dm_tree.map_structure_with_path_up_to(
   shallow_structure, func, *structures, **kwargs)

File /opt/hostedtoolcache/Python/3.9.19/x64/lib/python3.9/site-packages/tree/__init__.py:778, in map_structure_with_path_up_to(***failed resolving arguments***)
results = []
for path_and_values in _multiyield_flat_up_to(shallow_structure, *structures):
--> 778   results.append(func(*path_and_values))
return unflatten_as(shallow_structure, results)

File /opt/hostedtoolcache/Python/3.9.19/x64/lib/python3.9/site-packages/tensorflow_probability/python/internal/backend/jax/nest.py:326, in map_structure_up_to.<locals>.<lambda>(_, *args)
def map_structure_up_to(shallow_structure, func, *structures, **kwargs):
 return map_structure_with_tuple_paths_up_to(
     shallow_structure,
--> 326       lambda _, *args: func(*args),  # Discards path.
     *structures,
     **kwargs)

File /opt/hostedtoolcache/Python/3.9.19/x64/lib/python3.9/site-packages/tensorflow_probability/substrates/jax/bijectors/bijector.py:1707, in Bijector.forward_dtype.<locals>.<lambda>(x)
 input_dtype = nest_util.broadcast_structure(
     self.forward_min_event_ndims, self.dtype)
else:
 # Make sure inputs are compatible with statically-known dtype.
 input_dtype = nest.map_structure_up_to(
     self.forward_min_event_ndims,
-> 1707       lambda x: dtype_util.convert_to_dtype(x, dtype=self.dtype),
     nest_util.coerce_structure(self.forward_min_event_ndims, dtype),
     check_types=False)
output_dtype = self._forward_dtype(input_dtype, **kwargs)
try:
 # kwargs may alter dtypes themselves, but we currently require
 # structure to be statically known.

File /opt/hostedtoolcache/Python/3.9.19/x64/lib/python3.9/site-packages/tensorflow_probability/substrates/jax/internal/dtype_util.py:247, in convert_to_dtype(tensor_or_dtype, dtype, dtype_hint)
elif isinstance(tensor_or_dtype, np.ndarray):
 dt = base_dtype(dtype or dtype_hint or tensor_or_dtype.dtype)
--> 247 elif np.issctype(tensor_or_dtype):
 dt = base_dtype(dtype or dtype_hint or tensor_or_dtype)
else:
 # If this is a Python object, call `convert_to_tensor` and grab the dtype.
 # Note that this will add ops in graph-mode; we may want to consider
 # other ways to handle this case.

File /opt/hostedtoolcache/Python/3.9.19/x64/lib/python3.9/site-packages/numpy/__init__.py:397, in __getattr__(attr)
   raise AttributeError(__former_attrs__[attr])
if attr in __expired_attributes__:
--> 397     raise AttributeError(
       f"`np.{attr}` was removed in the NumPy 2.0 release. "
       f"{__expired_attributes__[attr]}"
   )
if attr == "chararray":
   warnings.warn(
       "`np.chararray` is deprecated and will be removed from "
       "the main namespace in the future. Use an array with a string "
       "or bytes dtype instead.", DeprecationWarning, stacklevel=2)

AttributeError: `np.issctype` was removed in the NumPy 2.0 release. Use `issubclass(rep, np.generic)` instead.

# Plot emissions and true_states in the emissions plane
plot_gaussian_hmm(hmm, true_params, train_emissions[0], train_true_states[0], 
                  title="True HMM emission distribution")

../../_images/263a661121f40d06c65fcf1ef0a76e8d3b5b40afbec1ef7cb2efdfdcdf0e4161.png

# Plot emissions vs. time with background colored by true state
plot_gaussian_hmm_data(hmm, true_params, train_emissions[0], train_true_states[0])

../../_images/57051aaa1c7df49cad433ee33fb43f455340ba93bcc287d6ab32cb54fcb16605.png

Write a helper function to perform leave-one-out cross-validation#

This function fits the data into folds where each fold consists of all but one of the training sequences. It fits the model to each fold in parallel, and then computes the log likelihood of the held-out sequence for each fold. The average held-out log likelihood is what we will use for determining the number of discrete states.

def cross_validate_model(model, key, num_iters=100):
    # Initialize the parameters using K-Means on the full training set
    params, props = model.initialize(key=key, method="kmeans", emissions=train_emissions)
    
    # Split the training data into folds.
    # Note: this is memory inefficient but it highlights the use of vmap.
    folds = jnp.stack([
        jnp.concatenate([train_emissions[:i], train_emissions[i+1:]])
        for i in range(num_train_batches)
    ])
    
    def _fit_fold(y_train, y_val):
        fit_params, train_lps = model.fit_em(params, props, y_train, 
                                             num_iters=num_iters, verbose=False)
        return model.marginal_log_prob(fit_params, y_val)

    val_lls = vmap(_fit_fold)(folds, train_emissions)
    return val_lls.mean(), val_lls

Now run the cross-validation function on a sequence of models with number of states ranging from 2 to 10.

# Make a range of Gaussian HMMs
all_num_states = list(range(2, 10))
test_hmms = [GaussianHMM(num_states, emission_dim, transition_matrix_stickiness=10.) 
          for num_states in all_num_states]
results = []
for test_hmm in test_hmms:
    print(f"fitting model with {test_hmm.num_states} states")
    results.append(cross_validate_model(test_hmm, jr.PRNGKey(0)))
    
avg_val_lls, all_val_lls = tuple(zip(*results))

fitting model with 2 states
fitting model with 3 states
fitting model with 4 states
fitting model with 5 states
fitting model with 6 states
fitting model with 7 states
fitting model with 8 states
fitting model with 9 states

Plot the individual and average validation log likelihoods as a function of number of states#

plt.plot(all_num_states, avg_val_lls, '-ko')
for k, per_fold_val_lls in zip(all_num_states, all_val_lls):
    plt.plot(k * jnp.ones_like(per_fold_val_lls), per_fold_val_lls, '.')
plt.xlabel("num states ($K$)")
plt.ylabel("avg. validation log prob.")

Text(0, 0.5, 'avg. validation log prob.')

../../_images/8e555628b871f5194fea5076cb0bbd6b5d5b93453e450b3e12afcfaa73f2422c.png

There’s no right answer for how to choose the number of states, but reasonable heuristics include:

picking \(K\) that has the highest average validation log prob
picking \(K\) where the average validation log prob stops increasing by a minimum amount
picking \(K\) with a hypothesis test for increasing mean

Here, we’ll just choose the number of states with the highest average.

best_num_states = all_num_states[jnp.argmax(jnp.stack(avg_val_lls))]
print("best number of states:", best_num_states)

best number of states: 5

Now fit a model to all the training data using the chosen number of states#

# Initialize the parameters using K-Means on the full training set
key = jr.PRNGKey(0)
test_hmm = GaussianHMM(best_num_states, emission_dim, transition_matrix_stickiness=10.)
params, props = test_hmm.initialize(key=key, method="kmeans", emissions=train_emissions)
params, lps = test_hmm.fit_em(params, props, train_emissions, num_iters=100)

100.00% [100/100 00:02<00:00]

Plot the log probabilities over the course of training, along with those of the model with the parameters that generated the data.

# Evaluate the log probability of the training data under the true parameters
true_lp = vmap(partial(hmm.marginal_log_prob, params))(train_emissions).sum()
true_lp += hmm.log_prior(params)

# Plot log probs vs num_iterations
offset = 0
plt.plot(jnp.arange(len(lps)-offset), lps[offset:], label='EM')
plt.axhline(true_lp, color='k', linestyle=':', label="True")
plt.xlabel('num epochs')
plt.ylabel('log prob')
plt.legend()

<matplotlib.legend.Legend at 0x7fb0e9c77520>

../../_images/c482e0bda2751c27a8abc67bfd9f28ac7fa6e6605d5e31d8971f42258cb2fb54.png

Visualize the fitted model#

We’ll make the same plots as above, but now using the fitted parameters and the chosen number of states.

most_likely_states = test_hmm.most_likely_states(params, train_emissions[0])
plot_gaussian_hmm(test_hmm, params, train_emissions[0], most_likely_states, 
                  title=f"Fitted Model with {best_num_states} states", alpha=0.25)

../../_images/8da692ad44858c2a0c6eea78c896162aa7eed591203525718e300919a0234562.png

plot_gaussian_hmm_data(test_hmm, params, train_emissions[0], most_likely_states, xlim=(0, 100))

../../_images/85997647c13046499e4f1b8746650e37544bdce468d7d44908b1503db3a65cd3.png

Important

Note that the marginal log probability is invariant to relabeling of the states. Compare these plots to the ones at the top from the true model. You’ll see that the states have been permuted.

Comparing Gaussian HMMs with different constraints on the covariance#

Dynamax implements Gaussian HMMs with different constraints on the covariance matrices:

DiagonalGaussianHMM assumes the covariance matrices are diagonal (i.e. \(\Sigma_k = \mathrm{diag}([\sigma_{k,1}^2, \ldots, \sigma_{k,D}^2])\))
SphericalGaussianHMM assumes the covariance matrices are “spherical” (i.e. \(\Sigma_k = \sigma_k^2 I\))
SharedCovarianceGaussianHMM assumes the covariance matrices are the same for all states (i.e. \(\Sigma_k = \Sigma\))

In this last section, we will compare these models based on the marginal probability they assign to the test data.

Warning

Usually, you would use cross-validation to choose the number of discrete states for each type of model, and then evaluate them on test data. For simplicity, we will assume that we know the true number of states.

def fit_model(model, key, num_iters=100):
    # Initialize the parameters using K-Means on the full training set
    params, props = model.initialize(key=key, method="kmeans", emissions=train_emissions)
    params, lps = model.fit_em(params, props, train_emissions, num_iters=num_iters, verbose=False)
    test_lp = vmap(partial(model.marginal_log_prob, params))(test_emissions).sum()
    return test_lp

models = [
    GaussianHMM(true_num_states, emission_dim, transition_matrix_stickiness=10.),
    DiagonalGaussianHMM(true_num_states, emission_dim, transition_matrix_stickiness=10.),
    SphericalGaussianHMM(true_num_states, emission_dim, transition_matrix_stickiness=10.),
    SharedCovarianceGaussianHMM(true_num_states, emission_dim, transition_matrix_stickiness=10.),
]

# Fit the models and collect the test log probabilities
key = jr.PRNGKey(0)
test_lps = [fit_model(model, key) for model in models]
# Compare to the log probability under the true model
true_test_lp = vmap(partial(hmm.marginal_log_prob, true_params))(test_emissions).sum()

print("Marginal log probabilities of test data:")
for model, test_lp in zip(models, test_lps):
    print(f"{model.__class__.__name__: >30}: {test_lp: .1f}")
print(f"{'True Model': >30}: {true_test_lp: .1f}")

Marginal log probabilities of test data:
                   GaussianHMM:  92.7
           DiagonalGaussianHMM:  109.0
          SphericalGaussianHMM:  102.2
   SharedCovarianceGaussianHMM:  103.7
                    True Model:  109.8

Here, the DiagonalGaussianHMM has the highest test log likelihood, and it’s almost as high as the log probability under the true parameters. The standard GaussianHMM, which allows for arbitrary covariance matrices is underperforming on test data, likely because it overfits the training data with its extra flexibility.

Here, the true parameters were in fact spherical and shared across all states, so it’s perhaps a bit surprising that the more general diagonal covariance model wins out on test data. Keep in mind, however, that there is some amount of randomness in both the sampled data and the initialization of these models, so the particular ordering might change from one random seed to the next. In actual applications, you should compute means and standard errors by running these analyses with multiple random seeds.

Conclusion#

Dynamax provides a family of Gaussian HMMs for multivariate continuous emissions. This notebook showed how to use cross-validation to choose the number of discrete states, and how to compare model families (e.g. GaussianHMM vs DiagonalGaussianHMM) on test data.

Now that you have seen HMMs with categorical and Gaussian emissions, you can pattern-match to use the other HMMs. See the documentation for a complete list.

The next notebook highlights one more generalization of the standard HMM to autoregressive emissions, where the emissions depend on preceding emissions as well as the current discrete state.

Gaussian HMM: Cross-validation and Model Selection

Contents

Gaussian HMM: Cross-validation and Model Selection#

Setup#

Helper functions for plotting#

Generate sample data#

Write a helper function to perform leave-one-out cross-validation#

Plot the individual and average validation log likelihoods as a function of number of states#

Now fit a model to all the training data using the chosen number of states#

Visualize the fitted model#

Comparing Gaussian HMMs with different constraints on the covariance#

Conclusion#