For (3), you may want to try modeling only K-1 categories. Then each set of parameters is interpreted as log odds vs. that fixed category.
(2) sounds like a recipe for some pretty strong non-identifiability as you have moving components and states.
For the code (4), can’t you declare unA_ij as a matrix and define it by matrix multiplication? That’s more efficient than the loop. Every time I look at one of these complicated mixtures it makes me realize we need better vectorized mixtures built-in. But it’s a design challenge figuring out how to specify all the shared computations efficiently and on the right scale (log, log odds, linear). If you define accumulator in the forward algorithm as a vector, you can vectorize the innermost loop. Do you really need those alpha_tk or can you define them as generated quantities? As generated quantities, there’s no autodiff overhead. I didn’t follow the forwards-backwards bit. I found I only had to calculate the forward algorithm and then I could read the likelihood off the final row with a final log_sum_exp. I’ve been wondering if using more sufficient stats in forward-backward would be a win.
You want to vectorize that look over K in the model. It’s much more efficient. I don’t like variable names like “hyperparams”–it’s too much like naming a variable “variable”. They’re also too easy to get confused that way. The ony tricky vectorization is mu_kl, which will be
mu_kl ~ normal(hypermu_k, hyperparams[3]);
The others you just drop the loop and the j index.
The moveHMM package in R has a nice overview of the movement HMMs with covariates. The CRF paper by McCallum, Lafferty and Pereira also compares some similar alterntives and introduces CRFs—it’s in a natural languae setting, but there have been lots of other applications.
Do you really need to run Viterbi within the algorithm? What do you use it for? Viterbi only finds the first-best output (or the n-best in the natural generalization with backtracking).