Constraining AR-HMM emission parameters

[JustFineNeuro]

Assuming an AR-HMM, I want to know if it's possible to constrain the AR model parameters. Specifically, is it possible to set the 'A' matrix to fixed parameters and only estimate parameters for the input matrix (labeled as 'V' in the ssm package)?

Thanks!

[slinderman]

Good question! I don’t think you can do that with the built in code. You could make a custom observation class and only include the updates for the other parameters, via a custom M step. It should end up looking pretty similar to the existing AR observations class. Does that make sense?

[JustFineNeuro]

Thanks for the rapid response! Yes, this makes sense.
I see the AR class in https://github.com/lindermanlab/ssm/blob/master/ssm/observations.py. Can you point me to where the M step is handled, I assume by HMM.fit(). But I couldn't find it easily.

[slinderman]

First let me just make sure I understand the model you'd like to fit. I think you're describing an autoregressive observation model,

$$ p(x_t \mid x_{t-1}, z_t, u_t) = \mathcal{N}(x_t \mid A x_{t-1} + B_{z_t} u_t, Q_{z_t}) $$

where the free parameters are ${B_k, Q_k}_{k=1}^K$. In your case, the $A$ matrix is fixed (and shared by all discrete states). Is that right?

If so, I think you could equivalently define an HMM on the residuals $r_t = x_t - A x_{t-1}$, then have a simple HMM with "linear regression" observations,

$$ p(r_t \mid z_t, u_t) = \mathcal{N}(r_t \mid B_{z_t} u_t, Q_{z_t}) $$

This is also called a "switching linear regression model." I don't think we have such an observation model built into SSM, but we do have one in dynamax (see here)

[JustFineNeuro]

Yes, that should work well. Namely, I know the feedforward matrix transition matrix of the system dynamics "A". What I am trying to do is parse behavioral states by assuming they only differ by projection into a control subspace "B". Side question, then: in dynamax, are there options for computing the HMM with multiple emission types? For example, a gaussian for $p(r_t|z_t,u_t)$ and Poisson for firing rates?