blackjax.sgmcmc.csgld

blackjax.sgmcmc.csgld#

Public API for the Contour Stochastic gradient Langevin Dynamics kernel [DLH+22, DLL20].

Classes#

ContourSGLDState

State of the Contour SgLD algorithm.

Functions#

init(→ ContourSGLDState)

build_kernel(→ Callable)

as_top_level_api(→ blackjax.base.SamplingAlgorithm)

Implements the (basic) user interface for the Contour SGLD kernel.

Module Contents#

class ContourSGLDState[source]#

State of the Contour SgLD algorithm.

Parameters:

  • position – Current position in the sample space.

  • energy_pdf – Vector with m non-negative values that sum to 1. The i-th value of the vector is equal to \(\int_{S_1} \pi(\mathrm{d}x)\) where \(S_i\) is the i-th energy partition.

  • energy_idx – Index i such that the current position belongs to \(S_i\).

position: blackjax.types.ArrayTree[source]#
energy_pdf: blackjax.types.Array[source]#
energy_idx: int[source]#
init(position: blackjax.types.ArrayLikeTree, num_partitions=512) → ContourSGLDState[source]#
build_kernel(num_partitions=512, energy_gap=10, min_energy=0) → Callable[source]#
Parameters:

  • num_partitions – The number of partitions we divide the energy landscape into.

  • energy_gap – The difference in energy \(\Delta u\) between the successive partitions. Can be determined by running e.g. an optimizer to determine the range of energies. num_partition * energy_gap should match this range.

  • min_energy – A rough estimate of the minimum energy in a dataset, which should be strictly smaller than the exact minimum energy! e.g. if the minimum energy of a dataset is 3456, we can set min_energy to be any value smaller than 3456. Set it to 0 is acceptable, but not efficient enough. the closer the gap between min_energy and 3456 is, the better.

as_top_level_api(logdensity_estimator: Callable, gradient_estimator: Callable, zeta: float = 1, num_partitions: int = 512, energy_gap: float = 100, min_energy: float = 0) → blackjax.base.SamplingAlgorithm[source]#

Implements the (basic) user interface for the Contour SGLD kernel.

Parameters:

  • logdensity_estimator – A function that returns an estimation of the model’s logdensity given a position and a batch of data.

  • gradient_estimator – A function that takes a position, a batch of data and returns an estimation of the gradient of the log-density at this position.

  • zeta – Hyperparameter that controls the geometric property of the flattened density. If zeta=0 the function reduces to the SGLD step function.

  • temperature – Temperature parameter.

  • num_partitions – The number of partitions we divide the energy landscape into.

  • energy_gap – The difference in energy \(\Delta u\) between the successive partitions. Can be determined by running e.g. an optimizer to determine the range of energies. num_partition * energy_gap should match this range.

  • min_energy – A rough estimate of the minimum energy in a dataset, which should be strictly smaller than the exact minimum energy! e.g. if the minimum energy of a dataset is 3456, we can set min_energy to be any value smaller than 3456. Set it to 0 is acceptable, but not efficient enough. the closer the gap between min_energy and 3456 is, the better.

Return type:

A SamplingAlgorithm.