jax_sgmc.alias.re_sgld

jax_sgmc.alias.re_sgld(potential_fn, data_loader, cache_size=512, batch_size=32, temperature=1000.0, first_step_size=0.05, last_step_size=0.001, burn_in=0, accepted_samples=100, save_to_numpy=True, progress_bar=True)[source]

Replica Exchange Stochastic Gradient Langevin Diffusion.

reSGLD simulates a tempered and a default chain in parallel, which exchange samples at random following a (biased) markov jump process [1].

[1] https://arxiv.org/abs/2008.05367v3

resgld_run = alias.re_sgld(...)

sample = {"w": jnp.zeros((N, 1)), "sigma": jnp.array(2.0)}
init_samples = [(sample, sample), (sample, sample), (sample, sample)]

results = resgld_run(
  *init_samples, 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

  • temperature (float) – Temperature at which the helper chain should run

  • 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

  • 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 start_chain_{idx}.