blackjax.adaptation.mass_matrix

Algorithms to adapt the mass matrix used by algorithms in the Hamiltonian Monte Carlo family to the current geometry.

The Stan Manual [stab] is a very good reference on automatic tuning of parameters used in Hamiltonian Monte Carlo.

Module Contents

Classes

WelfordAlgorithmState	State carried through the Welford algorithm.
MassMatrixAdaptationState	State carried through the mass matrix adaptation.

Functions

mass_matrix_adaptation(→ tuple[Callable, Callable, ...)	Adapts the values in the mass matrix by computing the covariance
welford_algorithm(→ tuple[Callable, Callable, Callable])	Welford's online estimator of covariance.

class WelfordAlgorithmState

State carried through the Welford algorithm.

mean

The running sample mean.

m2

The running value of the sum of difference of squares. See documentation of the welford_algorithm function for an explanation.

sample_size

The number of successive states the previous values have been computed on; also the current number of iterations of the algorithm.

mean: blackjax.types.Array

m2: blackjax.types.Array

sample_size: int

class MassMatrixAdaptationState

State carried through the mass matrix adaptation.

inverse_mass_matrix

The curent value of the inverse mass matrix.

wc_state

The current state of the Welford Algorithm.

inverse_mass_matrix: blackjax.types.Array

wc_state: WelfordAlgorithmState

mass_matrix_adaptation(is_diagonal_matrix: bool = True) → tuple[Callable, Callable, Callable]

Adapts the values in the mass matrix by computing the covariance between parameters.

Parameters:

is_diagonal_matrix – When True the algorithm adapts and returns a diagonal mass matrix (default), otherwise adaps and returns a dense mass matrix.

Returns:

• init – A function that initializes the step of the mass matrix adaptation.

• update – A function that updates the state of the mass matrix.

• final – A function that computes the inverse mass matrix based on the current state.

welford_algorithm(is_diagonal_matrix: bool) → tuple[Callable, Callable, Callable]

Welford’s online estimator of covariance.

It is possible to compute the variance of a population of values in an on-line fashion to avoid storing intermediate results. The naive recurrence relations between the sample mean and variance at a step and the next are however not numerically stable.

Welford’s algorithm uses the sum of square of differences \(M_{2,n} = \sum_{i=1}^n \left(x_i-\overline{x_n}\right)^2\) for updating where \(x_n\) is the current mean and the following recurrence relationships

Parameters:

is_diagonal_matrix – When True the algorithm adapts and returns a diagonal mass matrix (default), otherwise adaps and returns a dense mass matrix.

Note

It might seem pedantic to separate the Welford algorithm from mass adaptation, but this covariance estimator is used in other parts of the library.