Can't find initial parameters in HMM with deterministic emissions

[simonsteiger]

Minimal working example

using Turing, HiddenMarkovModels, LinearAlgebra

init = [1.0, 0.0]
trans = [0.8 0.2; 0.0 1.0]
emiss = [0.9 0.1; 0.0 1.0]

dists = [Categorical(emiss[i, :]) for i in axes(emiss, 1)]
hmm = HMM(init, trans, dists)
obs_seqs = [last(rand(hmm, 15)) for _ in 1:20]
seq_ends = cumsum(length.(obs_seqs))
obs_seqs_concat = reduce(vcat, obs_seqs)

# Calling `logdensityof` doesn't error
logdensityof(hmm, obs_seqs_concat; seq_ends)

# Model disallowing transitions from state 2 to state 1
@model function hmm_trans_fixed(K, obs_seqs, seq_ends)
    # Initialization probabilities
    init ~ Dirichlet(ones(K))

    # Emission probabilities
    Φ ~ filldist(Dirichlet(ones(K)), K)

    # Transition probabilities, never from state 2 to 1
    Θ₁ ~ Dirichlet(ones(K))
    Θ = Matrix(Diagonal(ones(eltype(Θ₁), K)))
    Θ[:, 1] = Θ₁

    # Set up HMM object
    dists = [Categorical(Φ[:, i]) for i in axes(Φ, 2)]
    hmm = HMM(init, permutedims(Θ), dists)

    # Forward algorithm
    Turing.@addlogprob! logdensityof(hmm, obs_seqs; seq_ends)
    return (; hmm)
end

# This model runs and infers parameters correctly
model1 = hmm_trans_fixed(2, obs_seqs_concat, seq_ends)
chn1 = sample(model1, NUTS(), 500)

# Model where state 2 always emits the same observation
@model function hmm_trans_emiss_fixed(K, obs_seqs, seq_ends)
    # Initialization probabilities
    init ~ Dirichlet(ones(K))
	
    # Emission probabilities, state 2 always emits 2
    Φ₁ ~ Dirichlet(ones(K))
    Φ = Matrix(Diagonal(ones(eltype(Φ₁), K)))
    Φ[:, 1] = Φ₁
	
    # Transition probabilities, never from state 2 to 1
    Θ₁ ~ Dirichlet(ones(K))
    Θ = Matrix(Diagonal(ones(eltype(Θ₁), K)))
    Θ[:, 1] = Θ₁

    # Set up HMM object
    dists = [Categorical(Φ[:, i]) for i in axes(Φ, 2)]
    hmm = HMM(init, permutedims(Θ), dists)

    # Forward algorithm
    Turing.@addlogprob! logdensityof(hmm, obs_seqs; seq_ends)

    return (; hmm)
end

# This model fails with "failed to find valid initial parameters in 1000 tries. [...]"
model2 = hmm_trans_emiss_fixed(2, obs_seqs_concat, seq_ends)
chn2 = sample(model2, NUTS(), 500)

Description

Hi Turing.jl team!

I'm sorry to be back with another HMM-related issue. 🥲

In the example above, I am trying to build an HMM where specific transitions are forbidden. The scenario in the real data involves a sequence of problem solving behaviours which always end with either a success or a failure. Transitions out of these terminal states are impossible, and if a participant is in this state, it will always emit the observation corresponding to the terminal state.

In my simplified example above, I first only restrict the transition matrix to disallow transitions out of the terminal state. The second row is the terminal state, so I've set this to [0.0, 1.0].

This model runs and also picks up that the emission probabilities from the terminal state are effectively [0.0, 1.0].

If I try to set the emission probabilities of the terminal state to [0.0, 1.0], the model fails with failed to find valid initial parameters in 1000 tries. [...].

Thanks a lot for any help or tips!

PS: @gdalle already encountered this issue in a discourse post I made a little while back and did not know where this error came from. So I thought it's best to post it here.

Julia version info

versioninfo()

Julia Version 1.11.4
Commit 8561cc3d68d (2025-03-10 11:36 UTC)
Build Info:
  Official https://julialang.org/ release
Platform Info:
  OS: macOS (arm64-apple-darwin24.0.0)
  CPU: 8 × Apple M1 Pro
  WORD_SIZE: 64
  LLVM: libLLVM-16.0.6 (ORCJIT, apple-m1)
Threads: 1 default, 0 interactive, 1 GC (on 6 virtual cores)

Manifest

]st --manifest

Ran in a Pluto notebook with
Turing v0.37.0
HiddenMarkovModels v0.6.2

[penelopeysm]

This error is kind of vague but usually it means that the autodiff backend is returning NaN's or infinite gradients. In this case there's something in the logdensityof call that makes ForwardDiff return NaN's.

Here's the same code without Turing that demonstrates what's going on under the hood:

using LinearAlgebra, Distributions, HiddenMarkovModels, ForwardDiff
K = 2
init = rand(Dirichlet(ones(K)))
phi1 = rand(Dirichlet(ones(K)))
phi = [phi1[1] phi1[2]; 0.0 1.0]
theta1 = rand(Dirichlet(ones(K)))
theta = [theta1[1] theta1[2]; 0.0 1.0]
dists = [Categorical(phi[1, :]), Categorical(phi[2, :])]
hmm = HMM(init, theta, dists)

obs_seqs = [last(rand(hmm, 15)) for _ in 1:20]
obs_seqs_concat = reduce(vcat, obs_seqs)
seq_ends = cumsum(length.(obs_seqs))

# The function to calculate logp given a set of parameters
function logp(params)
    init = params[1:2]
    theta = [params[3] params[4]; 0.0 1.0]
    phi1 = params[5:6]
    if length(params) == 6  # phi2 was fixed
        phi2 = [0.0, 1.0]
    else
        phi2 = params[7:8]  # phi2 was not fixed
    end
    phi = [phi1[1] phi1[2]; phi2[1] phi2[2]]
    dists = [Categorical(phi[1, :]), Categorical(phi[2, :])]
    hmm = HMM(init, theta, dists)
    return logdensityof(hmm, obs_seqs_concat; seq_ends)
end

# Generate some random parameters by sampling from the model priors
function generate_params(fix_phi2)
    init = rand(Dirichlet(ones(2)))
    phi1 = rand(Dirichlet(ones(2)))
    phi2 = if fix_phi2
        Float64[]  # no need to generate these parameters
    else
        rand(Dirichlet(ones(2)))
    end
    theta1 = rand(Dirichlet(ones(2)))
    return vcat(init, theta1, phi1, phi2)
end

# This runs fine
x = generate_params(false)
ForwardDiff.gradient(logp, x)

# This gives NaN gradients regardless of what parameters are generated
x = generate_params(true)
ForwardDiff.gradient(logp, x)

[penelopeysm]

My suspicion is that fundamentally it boils down to the same problem in #2513 - the fact that the observation distribution is a Categorical([0.0, 1.0]) but the observation is 1 which leads to -Inf probability and NaN gradient.

In #2513 (and in the working model above) if this happens, it proves fine to just reject the sample and try again with fresh parameters because not every sample has this observation distribution.

However if the distribution is fixed to a Categorical([0.0, 1.0]) then this will basically happen on every sample(??) Not fully sure if this analysis is correct, but that's where I'm at right now.

[penelopeysm]

Ooh, here's a discovery: if we just use Categorical([0.0, 1.0]) instead of taking it from phi, then ForwardDiff doesn't convert it to a Dual number, and it works:

# The function to calculate logp given a set of parameters
function f(params)
    init = params[1:2]
    theta = [params[3] params[4]; 0.0 1.0]
    phi1 = params[5:6]
    if length(params) == 6
        # This line changed!
        dist2 = Categorical([0.0, 1.0])
    else
        dist2 = Categorical(params[7:8])
    end
    dists = [Categorical(phi1), dist2]
    hmm = HMM(init, theta, dists)
    return logdensityof(hmm, obs_seqs_concat; seq_ends)
end

So I wonder if this might work? It samples on my computer and the results seem reasonable:

# Model where state 2 always emits the same observation
@model function hmm_trans_emiss_fixed(K, obs_seqs, seq_ends)
    # Initialization probabilities
    init ~ Dirichlet(ones(K))
	
    # Emission probabilities, state 2 always emits 2
    Φ₁ ~ Dirichlet(ones(K))
	
    # Transition probabilities, never from state 2 to 1
    Θ₁ ~ Dirichlet(ones(K))
    Θ = Matrix(Diagonal(ones(eltype(Θ₁), K)))
    Θ[:, 1] = Θ₁

    # This is hardcoded - you'd have to make this more generic; but make sure to keep
    # the fixed-emission distributions as Float64
    dists = [Categorical(Φ₁), Categorical([0.0, 1.0])]

    # Set up HMM object
    hmm = HMM(init, permutedims(Θ), dists)

    # Forward algorithm
    Turing.@addlogprob! logdensityof(hmm, obs_seqs; seq_ends)

    return (; hmm)
end

#=
model2 = hmm_trans_emiss_fixed(2, obs_seqs_concat, seq_ends)
chn2 = sample(model2, NUTS(), 500)

Summary Statistics
  parameters      mean       std      mcse   ess_bulk   ess_tail      rhat   ess_per_sec
      Symbol   Float64   Float64   Float64    Float64    Float64   Float64       Float64

     init[1]    0.9246    0.0609    0.0026   414.9502   271.6760    0.9988       64.2537
     init[2]    0.0754    0.0609    0.0026   414.9502   271.6760    0.9988       64.2537
       Φ₁[1]    0.9091    0.0312    0.0011   731.8442   406.4313    0.9987      113.3237
       Φ₁[2]    0.0909    0.0312    0.0011   731.8442   406.4313    0.9987      113.3237
       Θ₁[1]    0.8221    0.0356    0.0014   613.7334   375.4466    1.0004       95.0346
       Θ₁[2]    0.1779    0.0356    0.0014   613.7334   375.4466    1.0004       95.0346

Quantiles
  parameters      2.5%     25.0%     50.0%     75.0%     97.5%
      Symbol   Float64   Float64   Float64   Float64   Float64

     init[1]    0.7680    0.8994    0.9393    0.9684    0.9956
     init[2]    0.0044    0.0316    0.0607    0.1006    0.2320
       Φ₁[1]    0.8439    0.8898    0.9121    0.9321    0.9589
       Φ₁[2]    0.0411    0.0679    0.0879    0.1102    0.1561
       Θ₁[1]    0.7498    0.7985    0.8232    0.8444    0.8858
       Θ₁[2]    0.1142    0.1556    0.1768    0.2015    0.2502
=#

[simonsteiger]

Ohh, thanks a lot for all the work you've put into investigating this! 😊 I will certainly try the fixed-emission distributions containing Float64 tomorrow.

#2513 was solved in gdalle/HiddenMarkovModels.jl#127 by the way, but it might well be that this is still related to the same problem in a way that I don't understand.

[gdalle]

Fundamentally, all of these stability issues come from HiddenMarkovModels using rescaling instead of log-scale computations. The latter would be even more stable, but it would prevent us from using BLAS calls, forcing us to call expensive logsumexp reductions instead. I plan on implementing it at some point, you can track progress here:

• Log scale gdalle/HiddenMarkovModels.jl#130