JOSS review - Remarks on documentation

[gdalle]

Hi and congrats on the package!

I'm one of the reviewers for the JOSS paper you submitted, so here I'll list my questions and concerns about the documentation. This issue will be updated as my reading progresses so maybe don't start answering right away.

---

Home page

You can also call the low-level inference code (e.g. hmm_smoother()) directly, without first having to construct the HMM object.

This sentence was unclear to me: I still have to pass some model parameters to the smoother, right?

HMMs

Casino HMM: Inference

Overall really good and clear!

Now that we’ve initialize the model parameters, we can use them to draw samples from the HMM.

Typo (initializeD).

Casino HMM: Learning

First we construct an HMM and sample data from it, just as in the preceding notebook.

How would you handle data with several trajectories of varying lengths? Do you have to pad them into a 3D tensor, and then apply some kind of mask?

Perhaps the simplest learning algorithm is to directly maximize the marginal probability with gradient ascent.

How do you perform gradient ascent on the stochastic matrix $A$? Some kind of projection step? It seems far from obvious.

Also the example errors:

OverflowError: Python int too large to convert to C long

Finally, let’s compare the learning curve of EM to those of SGD.

The notion of "epoch" doesn't seem standard for EM, do you take it to mean one E step + one M step?

Not only does EM converge much faster on this example (here, in only a handful of iterations), it also converges to a better estimate of the parameters.

It would be interesting to explain why the asymptotic GD estimate is worse. Theoretically the finite size of the training data should also affect the EM algorithm, both are doing ERM on the loglikelihood.

print_params(em_params)

This example errors.

Gaussian HMM: Cross-validation and model selection

As in the preceding notebooks, we start by sampling data from the model. Here, we add a slight wrinkle: we will sample training and test data, where the latter is only used for model selection.

This example errors:

AttributeError: `np.issctype` was removed in the NumPy 2.0 release. Use `issubclass(rep, np.generic)` instead.

Also, in the "True HMM emission distribution" plot, it may not be obvious what the black lines stand for (transitions I assume).

Now fit a model to all the training data using the chosen number of states

It would be nice to add a comment explaining why the EM log prob can end up above the one of the true model.

SharedCovarianceGaussianHMM assumes the covariance matrices are the same for all states

Is it spherical, diagonal or generic?

AutoRegressive HMM demo

Very nice plots but not enough explanations.

Plot dynamics functions

Can you clarify what these plots represent? Adding titles / color legends would also help.

Sample emissions from the ARHMM

This example errors and the plot below (while a very pretty starfish) should also be explained.

The dotted lines represent the stationary point of the the corresponding AR state while the solid lines are the actual observations sampled from the HMM.

Maybe just specify that the stationary point is the limit of the recursion $y_{t} = Ay_{t-1} + b$ (not necessarily trivial for every reader).

Find the most likely states

Please annotate the plots.

Linear Gaussian SSMs

Tracking an object using the Kalman filter

params, _ = lgssm.initialize(jr.PRNGKey(0),...)

Why do we need an RNG to initialize the model even though all parameters are already fixed? This may also have been a relevant question for previous notebooks but I just noticed it now.

Sample some data from the model

This example errors.

Online linear regression using Kalman filtering

We perform sequential (recursive) Bayesian inference for the parameters of a linear regression model using the Kalman filter. (This algorithm is also known as recursive least squares.)

Conceptually, do we pay a performance penalty by doing this with an SSM-inspired formulation?

This is a LG-SSM, where...

It would be useful to remind the reader of the two equations defining LG-SSM in matrix form.

Online inference

This example errors.

Plot results

What is the difference between $w_0$ batch and $w_0$?

Parallel filtering and smoothing in an LG-SSM

This notebook shows how can reduce the cost of inference from O(T) to O(log T) time, if we have a GPU device.

• What is the difference between using your built-in parallel_smoother and just using jax.vmap on the normal smoother?
• Does this option exist for other models too, like HMMs?
• Can we also parallelize the EM algorithm, or some parts of it?

Test parallel inference on a single sequence

This example errors.

Also, are we supposed to see a difference between serial and parallel filtering on the plot? Obviously we expect both curves to be superposed but it is still a bit weird to have both in the legend and only see one.

MAP parameter estimation for an LG-SSM using EM and SGD

Data

This example errors.

Plot results

I don't understand what is going on in this plot. In particular, how you predict emissions from smoothing. When you predict, by definition you don't have access to observations beyond $t$. Shouldn't you predict from filtering instead?

Bayesian parameter estimation for an LG-SSM using HMC

Generate synthetic training data

This example errors.

Also I still don't understand what is meant by "smoothed emissions" (same problem as with the previous notebook), to me the only thing you can smooth is a state.

Implement HMC wrapper

An introduction to what HMC is and what you use to implement it (apparently blackjax) would be very useful here.

Call HMC

X and Y labels are wrong in the plot. Also, are we supposed to observe that the log probability increases? If so, the blue curve is not very convincing.

Use HMC to infer posterior over a subset of the parameters

Same remarks about the blue curve.

---

Related issues:

• [REVIEW]: Dynamax: A Python package for probabilistic state space modeling with JAX openjournals/joss-reviews#7069
• JOSS review - Remarks on code #377

[gdalle]

Making a second pass on the documentation, now that the errors have been fixed (the remarks above still apply).

General feedback

Important

I feel like I'm missing some documentation that is neither tutorial nor the full API reference. A kind of basic map to the different components of the package, how they fit together, and where to find the tools that I need. The Divio guide is a good resource to understand this distinction: I'm looking for something more explanatory and less demonstrative or exhaustive.

In nearly all tutorials, plots need more descriptions and axis labeling, especially when none of the axes represents time or iterations.

HMMs

casino HMM: Learning

As you can see, stochastic gradient descent converges much more quickly that full-batch gradient descent in this example. Intuitively, that’s because SGD takes multiple steps per epoch (i.e. each complete sweep through the dataset), whereas full-batch gradient descent takes only one.
The algorithms appear to have converged, but have they learned the correct parameters? Let’s see…
...
Ok, but not perfect!

The parameters learned by gradient descent are actually... pretty bad?
The second row of the transition matrix is way off, and the loadedness of one of the dies goes basically undetected. Is it just a bad initialization?

Linear Gaussian SSMs

Tracking an object using the Kalman filter

Sample some data from the model

It would be good to indicate the flow of time in this plot, which is not the horizontal axis.

Perform online filtering / Perform offline smoothing

What are the black circles?
These plots need more explanations and not just code.

Online linear regression using Kalman filtering

The final plot needs more explanations with words.

Parallel filtering and smoothing in an LG-SSM

I only just noticed that non-colorblind people probably see both serial and parallel curves superposed due to the different colors? But as a colorblind person I can't differentiate between red and green here. So maybe you could shift one of the curves slightly, to highlight that they follow each other?

MAP parameter estimation for an LG-SSM using EM and SGD

Data

What does the plot represent?

Bayesian parameter estimation for an LG-SSM using HMC

Call HMC

This example errors. Perhaps an insufficiently tight version bound on blackjax?

TypeError: kernel() got an unexpected keyword argument 'num_steps'

Nonlinear Gaussian SSMs

Tracking a spiraling object using the extended / unscented Kalman filter

We now show how to do this using a simple nonlinear Gaussian SSM, combined with various extensions of the Kalman filter algorithm.

What is the difference between these variants (e.g. unscented vs extended Kalman filter)? How does the user choose between them?

$q_t \in R^2$ [...] $r_t \in N(0, R)$

It would be clearer with different fonts for the set of real numbers $\mathbb{R}$ and the covariance matrix $R$.

Extended Kalman filter

Making one of the two lines dotted would help with readability. As would making the larger circles red instead of black.

Unscented Kalman filter

What changed between the two plots? How come the curves are suddenly superposed? We're missing a few explanations here.

Tracking a 1d pendulum using Extended / Unscented Kalman filter/ smoother

Sample data and plot it

Why does it look like the measurements are thresholded instead of just randomly distributed around the true angle?

Online learning for an MLP using extended Kalman filtering

Plot results

Again, the plots need more explanations, especially when the code is so verbose. What are the axes? What are the transparent curves, and the solid blue one? What are the dots?

Generalized Gaussian SSMs

Skipped for now

[slinderman]

Hi @gdalle,

Thanks for your thorough review, and for your patience as we've worked to address your comments. There's a lot to unpack here, but I've done my best to address everything.

First, you mentioned a lot of errors in running the notebooks. Sorry about that headache. I suspect that many stemmed from the numpy 2.0 incompatibility that we unfortunately inherit from Tensorflow Probability. Hopefully these issues are resolved now that we've pinned numpy <2.0 in our package requirements, as well as the tighter bounds on Python versions that we've introduced. Please let me know if issues persist.

---

One of your major comments was

I feel like I'm missing some documentation that is neither tutorial nor the full API reference. A kind of basic map to the different components of the package, how they fit together, and where to find the tools that I need. The Divio guide is a good resource to understand this distinction: I'm looking for something more explanatory and less demonstrative or exhaustive.

I'm not sure what exactly this would look like. The README and the home page of the documentation provides a basic map, and the demo notebooks are organized into categories for the major types of models and algorithms we support. Perhaps you could point me to an example of what a more explanatory page would look like? Divio points to this page as an example, but that feels a bit redundant to me. I'm still open to this suggestion though.

---

You can also call the low-level inference code (e.g. hmm_smoother()) directly, without first having to construct the HMM object.

This sentence was unclear to me: I still have to pass some model parameters to the smoother, right?

Technically no! The arguments to the core HMM inference routine are the initial state probabilities, the transition matrix (or matrices), and the emission potentials (log likelihoods). The emission potentials will typically be log likelihoods under an HMM with an emission model, and those calculations are performed in the higher level modeling modules. However, the low level routines allow for more general use cases as well. E.g., you could use potentials for a conditional random field, if you're doing structured prediction.

---

Typo (initializeD).

Thanks, fixed!

---

How would you handle data with several trajectories of varying lengths? Do you have to pad them into a 3D tensor, and then apply some kind of mask?

Unfortunately, we haven't added support for variable length time series yet, but it's high on our list of next steps. See #99. The simplest solution is what you suggest: pad and mask. However, when we add that functionality, I'd like to plumb down support for masking more generally. It's pretty easy in the HMM — just zero out the log likelihoods for masked data. It's slightly more tedious for LGSSM and other nonlinear SSMs, but the concept is the same. The real question is about the API — do we plumb a mask keyword arg or try to indicate masks via nans? Unfortunately neither is perfect, as discussed in the issue above.

---

How do you perform gradient ascent on the stochastic matrix A? Some kind of projection step? It seems far from obvious.

We use bijectors to map between constrained and unconstrained parameterizations. For the transition matrix, these are specified here: https://github.com/probml/dynamax/blob/main/dynamax/hidden_markov_model/models/transitions.py#L79. The properties use the SoftmaxCentered bijector to map from the $(K-1)$ dimensional simplex to $R^{K-1}$. Alternatively, you could use a projected/proximal gradient method, but that would require a bespoke implementation. The bijectors allow for vanilla (S)GD.

---

The notion of "epoch" doesn't seem standard for EM, do you take it to mean one E step + one M step?

Yes, we compare an epoch of SGD to a one EM iteration (one E + one M step). I agree that "epoch" is a strange word for an EM iteration, but it seems like the natural comparison since both SGD and EM have seen the full dataset after

---

It would be interesting to explain why the asymptotic GD estimate is worse. Theoretically the finite size of the training data should also affect the EM algorithm, both are doing ERM on the loglikelihood.

For this example, gradient descent and SGD get stuck in a local optimum and never escape. I did some tuning of the learning rate and momentum for SGD but it still got stuck. However, I added Adam to the list of optimizers, and with well-tuned learning rates it matches EM! I think this is a good demonstration of the importance of hyperparameter tuning for gradient-based methods. I've added it to the notebook. Thanks for the prompt!

---

Gaussian HMM: Cross-validation and model selection
Also, in the "True HMM emission distribution" plot, it may not be obvious what the black lines stand for (transitions I assume).

We've added a note to explain this plot better.

---

Now fit a model to all the training data using the chosen number of states
It would be nice to add a comment explaining why the EM log prob can end up above the one of the true model.

We've added a comment to explain this better.

---

SharedCovarianceGaussianHMM assumes the covariance matrices are the same for all states
Is it spherical, diagonal or generic?

It's a generic full covariance matrix. I added a comment to clarify.

---

AutoRegressive HMM demo
Very nice plots but not enough explanations.

I have added comments throughout this notebook to explain what's going on!

---

Plot dynamics functions
Can you clarify what these plots represent? Adding titles / color legends would also help.

Added further explanations.

---

Sample emissions from the ARHMM
This example errors and the plot below (while a very pretty starfish) should also be explained.
The dotted lines represent the stationary point of the the corresponding AR state while the solid lines are the actual observations sampled from the HMM.

Added further explanations.

---

Linear Gaussian SSMs

Tracking an object using the Kalman filter
params, _ = lgssm.initialize(jr.PRNGKey(0),...)
Why do we need an RNG to initialize the model even though all parameters are already fixed? This may also have been a relevant question for previous notebooks but I just noticed it now.

The initialize functions will randomly initialize parameters that aren't explicitly passed as keyword arguments. In this case, all the required parameters are specified, so the key isn't used. Moreover, since the key is an optional argument with a default value of zero, it could be omitted. In general, however, we thought it was good practice to explicitly pass the key in these notebooks because we expect users will often specify only a subset of the required parameters, in which case the key would be necessary. Rather than having the model silently initialize with the default key, we want the user to see that they have control over than random initialization. I agree it's not a perfect solution.

---

Online linear regression using Kalman filtering
We perform sequential (recursive) Bayesian inference for the parameters of a linear regression model using the Kalman filter. (This algorithm is also known as recursive least squares.)
Conceptually, do we pay a performance penalty by doing this with an SSM-inspired formulation?

It's certainly more expensive than doing least squares once with the full batch of data since you are updating the filtering distribution with each new data point, and that is cubic in the number of parameters. So for online linear regression, the cost is $O(ND^2 + TD^3)$ where $D$ is the parameter dimension. It is more efficient than naively running least squares $T$ times though, since it recursively updates the posterior with each new data point. I.e., as a function of the number of data points $N$, the cost is only $O(ND^2)$ rather than $O(TND^2)$ if you naively ran least squares $T$ times. Not sure it's worth spelling these details out in the notebook, but we can if you think it's important.

---

This is a LG-SSM, where...
It would be useful to remind the reader of the two equations defining LG-SSM in matrix form.

Good suggestion, done.

---

What is the difference between $w_0$ batch and $w_0$

I've edited the notebook to use "$\hat{\theta}$ offline" and "$\hat{\theta}$ online" to clarify that these are offline/online estimates of the regression parameters, respectively.

---

Parallel filtering and smoothing in an LG-SSM
This notebook shows how can reduce the cost of inference from O(T) to O(log T) time, if we have a GPU device.

What is the difference between using your built-in parallel_smoother and just using jax.vmap on the normal smoother?

It's very different and super cool! The parallel filter/smoothers use an associative scan rather than a sequential scan. So rather than taking $O(T)$ time, they take only $O(\log T)$ time on a parallel machine with $O(T)$ processors. this is parallelizing over time rather than parallelizing over batches!

There are some caveats to keep in mind: even on a GPU, you can exhaust the parallel resources, so you see diminishing returns for large enough $T$. These algorithms also require more memory usage and compute per step (it involves matrix-matrix multiplication rather than matrix-vector). Altogether, the final wall-clock time ends up depending a lot on your hardware.

Does this option exist for other models too, like HMMs?

Yes! See dynamax.hidden_markov_model.parallel_inference. I added a comment to the notebook.

Can we also parallelize the EM algorithm, or some parts of it?

You can parallelize the E-step using these algorithms. The M-step already operates on sufficient statistics that are summed across time, so there's not much parallelization to be done there.

Also, are we supposed to see a difference between serial and parallel filtering on the plot? Obviously we expect both curves to be superposed but it is still a bit weird to have both in the legend and only see one.

I think it was just a poor choice to use red and green for overlapping lines. I've updated the plot so the parallel and serial estimates should both be visible, even though they are overlapping.

---

MAP parameter estimation for an LG-SSM using EM and SGD
Plot results
I don't understand what is going on in this plot. In particular, how you predict emissions from smoothing. When you predict, by definition you don't have access to observations beyond t. Shouldn't you predict from filtering instead?

This notebook was a bit confusing, sorry. We have updated the notebook with many more comments
and explanations. This notebook shows how you can use an LG-SSM to reconstruct or denoise the emissions.

Technically, the posterior predictive distribution as defined in Dynamax is,
$\int \mathrm{N}(y_t \mid Hx_t, R) , \mathrm{N}(x_t \mid \mu_{t|T}, \Sigma_{t|T}) , \mathrm{dif} x_t$,
where $\mu_{t|T}$ and $\Sigma_{t|T}$ are the smoothing means and covariances, respectively.

We've updated the notebook to show the true emissions and latent states as well as the
posterior distribution over latent states and the posterior predictive distribution of
the observations (i.e., the reconstructions). The latent states are only identifiable
up to orthogonal transformations, but even if the true and inferred latent states don't
look the same, the reconstructions can still be accurate.

---

Bayesian parameter estimation for an LG-SSM using HMC
Generate synthetic training data
This example errors.

We have fixed the notebook to comply with the latest Blackjax syntax, so it should no longer error.

---

Implement HMC wrapper
An introduction to what HMC is and what you use to implement it (apparently blackjax) would be very useful here.

Agreed. We've significantly improved the exposition in this notebook. Now we provide a short intro to HMC and
why you might want to perform Bayesian inference. We've also fixed the plots adjusted the HMC hyperparameters
so that the curves look more reasonable.

---

Nonlinear Gaussian SSMs
Tracking a spiraling object using the extended / unscented Kalman filter
We now show how to do this using a simple nonlinear Gaussian SSM, combined with various extensions of the Kalman filter algorithm.
What is the difference between these variants (e.g. unscented vs extended Kalman filter)? How does the user choose between them?

We have updated these notebooks to provide more background on the EKK and UKF, and we have provided textbook references where the user can learn more about these algorithms. We do assume that the user has some familiarity with these algorithms, so we don't provide complete derivations or explanations. However, per your suggestion, we have tried to provide some guidance as to when you might prefer one or the other.

---

It would be clearer with different fonts for the set of real numbers R and the covariance matrix R.

Fixed.

---

Extended Kalman filter
Making one of the two lines dotted would help with readability. As would making the larger circles red instead of black.

Good suggestion. We've updated these plots.

---

Unscented Kalman filter
What changed between the two plots? How come the curves are suddenly superposed? We're missing a few explanations here.

We have added more explanation to these plots. In this example, the EKF and UKF provide very similar state estimates, and they are both close to the true states.

---

Tracking a 1d pendulum using Extended / Unscented Kalman filter/ smoother
Sample data and plot it
Why does it look like the measurements are thresholded instead of just randomly distributed around the true angle?

That is because in this example, the observation model is assumed to be $h(z_t) = \sin (z_{t,1}) + r_t$ where $r_t \sim \mathrm{N}(0,R)$ is additive Gaussian noise. Sarkka motivates this model as only observing the horizontal position of the pendulum, and he notes that it is similar to radar measurements where you only get noisy observations the range and bearing to an object.

---

Online learning for an MLP using extended Kalman filtering
Plot results
Again, the plots need more explanations, especially when the code is so verbose. What are the axes? What are the transparent curves, and the solid blue one? What are the dots?

This is a good criticism. We have significantly improved the explanations and plots in this notebook.

---

Generalized Gaussian SSMs
Skipped for now

Ok. I made some nice improvements to the "CMGF for onlinear logistic regression" and "Inferring States of a Poisson LDS" notebooks. The "CMGF for online learning of an MLP" notebook could still use some work, but it's very similar to the online logistic regression notebook, so I'm not so worried about it. Hopefully the changes are sufficient.

---

As for the paper itself, it is very clear and does a good job of introducing dynamax. I have three minor suggestions:

Dynamax supports canonical SSMs and allows the user to construct bespoke models as needed.
Can you give more details on how this is implemented API-wise? For instance, how generic can observation distributions be?
This question was a major motivation for my own take on HMMs, coded in Julia (https://joss.theoj.org/papers/10.21105/joss.06436). Hopefully I didn't misrepresent dynamax in the state of the art there.

For HMM with a new emission model, it's pretty straightforward. At a minimum, you just inherit from the base class, create a NamedTuple to specify the parameters and their properties (e.g. a bijector that maps to/from an unconstrained parameterization), and then a function to compute the emission distribution as a function of the latent state (and optionally any exogenous inputs to the system). Then the resulting model will inherit the SGD and EM algorithms (with SGD for the M-step) from the base class.

---

Dynamax provides a unique combination of low-level inference algorithms and high-level modeling objects that can support a wide range of research applications in JAX.
Would it be possible to provide concrete examples of what is missing from the previous libraries? Obviously JAX support is a big aspect, since I know that some of them are coded in Numpy (hmmlearn) or PyTorch (pomegranate)

The unique features, in my opinion, are JAX and the breadth of algorithms it includes. Perhaps the simplest tweak here is to delete the word "unique" since, I agree, almost everything in this package can be found to some degree in another package, but mostly not in JAX.

---

Parallel message passing routines that leverage GPU or TPU acceleration to perform message passing in sublinear time.
What do you mean by sublinear time? Isn't it just (roughly speaking) total sequential time divided by amount of parallelism?

It's actually logarithmic in the length of the sequence on a parallel machine! This is one of the really cool features, and it's pretty straightforward in JAX. See comments above in the notebook section.

[gdalle]

Thank you for addressing my remarks!

I'm not sure what exactly this would look like. The README and the home page of the documentation provides a basic map, and the demo notebooks are organized into categories for the major types of models and algorithms we support. Perhaps you could point me to an example of what a more explanatory page would look like? Divio points to this page as an example, but that feels a bit redundant to me. I'm still open to this suggestion though.

Okay, that's fair, I guess an organization by HMM type works fine.

It's certainly more expensive than doing least squares once with the full batch of data since you are updating the filtering distribution with each new data point, and that is cubic in the number of parameters. So for online linear regression, the cost is O ( N D 2 + T D 3 ) where D is the parameter dimension. It is more efficient than naively running least squares T times though, since it recursively updates the posterior with each new data point. I.e., as a function of the number of data points N , the cost is only O ( N D 2 ) rather than O ( T N D 2 ) if you naively ran least squares T times. Not sure it's worth spelling these details out in the notebook, but we can if you think it's important.

No, I don't think it is necessary indeed.

We've also fixed the plots adjusted the HMC hyperparameters
so that the curves look more reasonable.

When you freeze a subset of the parameters, isn't it weird that the HMC loglikelihood curve starts above the likelihood of the true params?

Overall the documentation is now fully satisfactory, closing this issue.