google-deepmind / google-deepmind/distrax

HMM.forward_backward returns zero posterior marginals when length is shorter than obs_seq

Open Beginner friendly
#337 1 comment 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
651
Forks
48
Avg merge
21h 4m
Merged PRs (30d)
3

Description

`HMM.forward_backward` documents `length` as the valid length of an observation sequence, useful for truncating padded sequences. However, when `length < len(obs_seq)`, the returned posterior marginals for valid timesteps can become all zeros.

Minimal repro:

```python
import distrax
import jax.numpy as jnp
from distrax._src.utils import hmm

model = hmm.HMM(
init_dist=distrax.Categorical(probs=jnp.array([0.6, 0.4])),
trans_dist=distrax.Categorical(
probs=jnp.array([[0.8, 0.2], [0.3, 0.7]])
),
obs_dist=distrax.Normal(
loc=jnp.array([0.0, 3.0]),
scale=jnp.array([0.5, 0.5]),
),
)

obs = jnp.array([0.05, 2.9, 0.1, 99.0, 99.0])

_, beta_pad, gamma_pad, ll_pad = model.forward_backward(obs, length=jnp.array(3))
_, beta_prefix, gamma_prefix, ll_prefix = model.forward_backward(obs[:3])

print(ll_pad, ll_prefix)
print(beta_pad[:3])
print(beta_prefix)
print(gamma_pad[:3])
print(gamma_prefix)
```

Observed:

```text
ll_pad == ll_prefix

beta_pad[:3]
[[0. 0.]
[0. 0.]
[0. 0.]]

gamma_pad[:3]
[[0. 0.]
[0. 0.]
[0. 0.]]
```

Expected:
`forward_backward(obs, length=3)` should match `forward_backward(obs[:3])` for the first three valid timesteps.

Likely cause:
In `HMM.backward`, padded suffix steps replace the backward carry with zeros when `t > length`. Those zeros then propagate backward into valid timesteps.

Possible fix:
For padded suffix steps, preserve `beta_prev` instead of replacing it with zeros, and add a regression test comparing `forward_backward(obs, length=L)` against `forward_backward(obs[:L])`.

Contributor guide

Open the contributing guide

Research direction

Start with HMM.backward and run the minimal reproduction from the issue. Trace how padded suffix steps update the backward carry, then add a regression test comparing forward_backward(obs, length=L) with forward_backward(obs[:L]); done means valid beta and gamma values match while the log likelihood remains equal.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
machine-learning
Issue type
Bug
Difficulty
2/5
Estimated time
1-3 hours
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
72/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.