vicariousinc / PGMax

Loopy belief propagation for factor graphs on discrete variables, in JAX!
https://pgmax.readthedocs.io
MIT License
64 stars 9 forks source link

Incomplete documentation for evidence #119

Open nathanielvirgo opened 2 years ago

nathanielvirgo commented 2 years ago

I would like to do the following, which is one step of the EM algorithm:

I'm currently stuck at the second bullet point. If I search the docs for "evidence" I find pgmax.fg.graph.Evidence and pgmax.fg.graph.update_evidence, but the documentation for these is very incomplete and doesn't specify the format in which the inputs should be provided or the procedure for using them.

The blog post doesn't use these but something else - it uses an argument evidence_updates to run_bp, but this is not mentioned in the documentation. I can't really understand what this argument is doing in the example, but I suspect it's not what I'm looking for.

So I'm a bit stuck on understanding the basic usage, and I'm wondering if I could either get a hint about it, and/or whether the documentation could be updated to make it a bit clearer.

StannisZhou commented 2 years ago

This is definitely another part of the documentation we need to improve.

You can update evidence using intuitive assignment syntax. https://github.com/vicariousinc/PGMax/blob/f7f3e3fe7fbb3a3c7582b7093c6ecd56eb04963a/examples/ising_model.py#L96 Here is an example for an Ising model. You index the evidence with variable names (tuples of the type (group_name, variable_name_in_group)) and do assignment. If you are working with variable groups (e.g. an NDVariableArray), you can also use intuitive assignment syntax with arrays (e.g. https://github.com/vicariousinc/PGMax/blob/f7f3e3fe7fbb3a3c7582b7093c6ecd56eb04963a/examples/ising_model.py#L103).

The evidence_updates argument to run_bp updates evidence on the fly for the given belief propagation run. You can use the elements in this dictionary in the above assignment syntax to achieve the same kind of evidence updates.

Some additional comments for this two different ways of doing evidence updates:

  1. Using the assignment syntax the updates are persistent. When you do run_bp later you will always work with the updated evidence.
  2. Using evidence_updates the updates are not persistent. It only updates the evidence for that particular run_bp run. But by structuring evidence_updates as a dictionary of jnp arrays and having run_bp as pure python functions, we obtain the ability to vectorize over/get gradients w.r.t. the evidence updates using jax.vmap/jax.grad. This is the main motivation for organizing things this way.
StannisZhou commented 2 years ago

Also you are more than welcome to contribute an example with directed models (since you obviously are working on one). Please refer to the developer installation instructions and contribution guide, and feel free to open a PR when it's ready! I can make more comments on the PR and make it part of our examples.