Pyro中高斯HMM的发射均值全部收敛至同一值的问题

问题背景 #

手动生成含3个状态的一阶马尔可夫序列（真实均值[-10, 0, 10]，标准差[1,2,1]），使用Pyro构建的高斯HMM模型训练时ELBO持续下降，但三个状态的均值最终收敛至同一值；同一序列用hmmlearn可正确推断真实参数。

核心问题分析 #

1. 缺失初始观测建模：原模型仅从t=1开始处理观测，未对t=0的初始状态对应观测y_0建模，丢失了初始状态的关键观测信息，导致状态区分度不足。
2. 参数初始化不合理：原p_mu用torch.randn初始化，初始值集中在0附近，与真实均值跨度大的分布不符，易让优化陷入局部最优，导致所有均值收敛到同一值。
3. ELBO类型不匹配：使用Trace_ELBO处理带枚举离散变量的模型，梯度计算精度不足，影响参数推断效果。

修正后的模型与训练代码 #

修正后的Model函数

def model(observations, num_state):
    assert not torch._C._get_tracing_state()
    with poutine.mask(mask=True):
        p_transition = pyro.sample("p_transition",
                                   dist.Dirichlet((1 / num_state) * torch.ones(num_state, num_state)).to_event(1))
        p_init = pyro.sample("p_init",
                             dist.Dirichlet((1 / num_state) * torch.ones(num_state)))
    # 改用接近真实值的初始化，避免局部最优
    p_mu = pyro.param(name="p_mu",
                      init_tensor=torch.tensor([-8.0, 1.0, 9.0]),
                      constraint=constraints.real)
    p_tau = pyro.param(name="p_tau",
                       init_tensor=torch.tensor([1.0, 2.0, 1.0]),
                       constraint=constraints.positive)
    
    current_state = pyro.sample("x_0",
                                dist.Categorical(p_init),
                                infer={"enumerate": "parallel"})
    # 新增初始观测y_0的建模
    pyro.sample("y_0",
                dist.Normal(Vindex(p_mu)[current_state], Vindex(p_tau)[current_state]),
                obs=observations[0])
    
    for t in pyro.markov(range(1, len(observations))):
        current_state = pyro.sample("x_{}".format(t),
                                    dist.Categorical(Vindex(p_transition)[current_state, :]),
                                    infer={"enumerate": "parallel"})
        pyro.sample("y_{}".format(t),
                    dist.Normal(Vindex(p_mu)[current_state], Vindex(p_tau)[current_state]),
                    obs=observations[t])

修正后的训练代码

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
obs = torch.tensor(obs, dtype=torch.float32).to(device)
torch.set_default_tensor_type("torch.cuda.FloatTensor" if torch.cuda.is_available() else "torch.FloatTensor")

# 改用TraceEnum_ELBO适配枚举离散变量的模型
guide = AutoDelta(poutine.block(model, expose_fn=lambda msg: msg["name"].startswith("p_")))
elbo = TraceEnum_ELBO(max_plate_nesting=1)
optim = Adam({"lr": 0.01})  # 增大学习率加速参数分离
svi = SVI(model, guide, optim, elbo)

修正说明 #

• 补充初始观测：加入y_0的观测建模，让模型利用初始状态的观测信息，提升状态区分能力。
• 优化参数初始化：将均值和标准差的初始值设置为接近真实值，引导优化过程避开局部最优。
• 更换ELBO类型：TraceEnum_ELBO针对带枚举离散变量的场景优化了梯度计算，能更准确地更新参数。
• 调整学习率：从0.001提升至0.01，加快参数收敛速度，避免因学习率过小导致参数无法分离。

内容的提问来源于stack exchange，提问作者NNNNNNN