blackjax.vi.meanfield_vi

blackjax.vi.meanfield_vi#

Classes#

MFVIState

MFVIInfo

Functions#

step(, stl_estimator, MFVIInfo])

Approximate the target density using the mean-field approximation.

sample(rng_key, state[, num_samples])

Sample from the mean-field approximation.

as_top_level_api(logdensity_fn, optimizer[, ...])

High-level implementation of Mean-Field Variational Inference

generate_meanfield_logdensity(mu, rho)

Module Contents#

class MFVIState[source]#
mu: blackjax.types.ArrayTree[source]#
rho: blackjax.types.ArrayTree[source]#
opt_state: optax.OptState[source]#
class MFVIInfo[source]#
elbo: float[source]#
step(rng_key: blackjax.types.PRNGKey, state: MFVIState, logdensity_fn: Callable, optimizer: optax.GradientTransformation, num_samples: int = 5, objective: blackjax.vi._gaussian_vi.Objective = KL(), stl_estimator: bool = True) → tuple[MFVIState, MFVIInfo][source]#

Approximate the target density using the mean-field approximation.

Parameters:

  • rng_key – Key for JAX’s pseudo-random number generator.

  • state – Current state of the mean-field approximation.

  • logdensity_fn – Function that represents the target log-density to approximate.

  • optimizer – Optax GradientTransformation to be used for optimization.

  • num_samples – The number of samples that are taken from the approximation at each step to compute the Kullback-Leibler divergence between the approximation and the target log-density.

  • objective – The variational objective to minimize. KL() by default or RenyiAlpha(alpha). For alpha = 1, Renyi reduces to KL.

  • stl_estimator – Whether to use the stick-the-landing (STL) gradient estimator [RWD17]. The STL estimator has lower gradient variance by removing the score function term from the gradient. [ASD20] recommend keeping it enabled.

Return type:

Updated MFVIState and MFVIInfo containing the ELBO value.

sample(rng_key: blackjax.types.PRNGKey, state: MFVIState, num_samples: int = 1)[source]#

Sample from the mean-field approximation.

Parameters:

  • rng_key – Key for JAX’s pseudo-random number generator.

  • state – Current MFVIState containing the variational parameters.

  • num_samples – Number of samples to draw.

Return type:

A PyTree of samples with leading dimension num_samples

as_top_level_api(logdensity_fn: Callable, optimizer: optax.GradientTransformation, num_samples: int = 100, objective: blackjax.vi._gaussian_vi.Objective = KL(), stl_estimator: bool = True)[source]#

High-level implementation of Mean-Field Variational Inference

Parameters:

  • logdensity_fn – A function that represents the log-density function associated with the distribution we want to sample from.

  • optimizer – Optax optimizer to use to optimize the variational objective.

  • num_samples – Number of samples to take at each step to optimize the ELBO.

  • objective – The variational objective to minimize. KL() by default or RenyiAlpha(alpha). For a = 1, Renyi reduces to KL.

  • stl_estimator – Whether to use the STL gradient estimator. Only supported when objective is KL() or RenyiAlpha(alpha=1.0).

Return type:

A VIAlgorithm.

generate_meanfield_logdensity(mu, rho)[source]#