jax_sgmc.alias.sgld

jax_sgmc.alias.sgld(potential_fn, data_loader, cache_size=512, batch_size=32, first_step_size=0.05, last_step_size=0.001, burn_in=0, accepted_samples=1000, rms_prop=False, alpha=0.9, lmbd=1e-05, save_to_numpy=True, progress_bar=True)[source]

Stochastic Gradient Langevin Dynamics.

SGLD with a polynomial step size schedule and optional speed up via RMS-prop adaption [1].

[1] https://arxiv.org/abs/1512.07666

rms_run = alias.sgld(...)

sample = {"w": jnp.zeros((N, 1)), "sigma": jnp.array(10.0)}
results = rms_run(sample, init_model_state=0, iterations=50000)[0]['samples']['variables']
Parameters:

  • potential_fn (minibatch_potential) – Stochastic potential over a minibatch of data

  • data_loader (DataLoader) – Data loader, e. g. numpy data loader

  • cache_size (int) – Number of mini_batches in device memory

  • batch_size (int) – Number of observations per batch

  • first_step_size (float) – First step size

  • last_step_size (float) – Final step size

  • burn_in (int) – Number of samples to skip before collecting samples

  • accepted_samples (int) – Total number of samples to collect, will be determined by random thinning if accepted samples < iterations - burn_in

  • rms_prop (bool) – Whether to adapt a manifold via the RMSprop strategy

  • alpha (float) – Decay speed of previous manifold approximations

  • lmbd (float) – Stabilization parameter

  • save_to_numpy (bool) – Save on host in numpy array instead of in device memory

  • progress_bar (bool) – Print the progress of the solver

Returns:

Returns a solver function which can be applied to multiple chains starting at init_sample. If the likelihood is stateful, an initial state must be provided.