Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add low-rank-modified metric #684

Draft
wants to merge 1 commit into
base: main
Choose a base branch
from

Conversation

aseyboldt
Copy link

Add a low-rank-modified metric

This adds a new type of metric for hmc-style samlers:
Instead of just allowing a diagonal or a full mass matrix, we combine the diagonal mass matrix with a couple of additional eigenvectors and eigenvalues of the remaining correlation.

This leads to the following form for the mass matrix:

$$ \begin{align} P &= V(\Sigma^{-1} - I)V^T + I \\ M &= D^{-\frac{1}{2}}PD^{-\frac{1}{2}} \end{align} $$

$P$ has exactly the eigenvalues specified in $\Sigma$ with the corresponding eigenvectors in $V \in R^{n, k}$. All remaining eigenvalues are one.

All operations needed in hmc/nuts can be done in $O(nk)$, so if the number of eigenvalues is small this has a cost much closer to diagonal mass matrix estimation.

This PR for now only contributes the metric, but corresponding adaptation should come a bit later...

Related issue: #683

A few important guidelines and requirements before we can merge your PR:

  • If I add a new sampler, there is an issue discussing it already;
  • We should be able to understand what the PR does from its title only;
  • There is a high-level description of the changes;
  • There are links to all the relevant issues, discussions and PRs;
  • The branch is rebased on the latest main commit;
  • Commit messages follow these guidelines;
  • The code respects the current naming conventions;
  • Docstrings follow the numpy style guide
  • pre-commit is installed and configured on your machine, and you ran it before opening the PR;
  • There are tests covering the changes;
  • The doc is up-to-date;
  • If I add a new sampler* I added/updated related examples

Consider opening a Draft PR if your work is still in progress but you would like some feedback from other contributors.

momentum_generator,
kinetic_energy,
is_turning,
data=(diagonal_scale_std, eigenvalues, eigenvectors),
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I guess you want to store these as properties so it is easier for tuning later on? In any case I suggest removing it here, and add it in the subsequent PR when tuning is introduced (so we can discuss whether it is necessary)

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, that was mostly for debugging and I forgot to take it out. Sorry :-)

r"""Hamiltonian dynamic on euclidean manifold with normally-distributed momentum
:cite:p:`betancourt2013general`.

The gaussian euclidean metric is a euclidean metric further characterized
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
The gaussian euclidean metric is a euclidean metric further characterized
The Gaussian Euclidean metric is a Euclidean metric further characterized


The gaussian euclidean metric is a euclidean metric further characterized
by setting the conditional probability density :math:`\pi(momentum|position)`
to follow a standard gaussian distribution. A Newtonian hamiltonian
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
to follow a standard gaussian distribution. A Newtonian hamiltonian
to follow a standard Gaussian distribution. A Newtonian Hamiltonian

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants