Contributed Code¶
Nested Sampling¶
- class NestedSampler(model, *, constructor_kwargs=None, termination_kwargs=None)[source]¶
Bases:
object
(EXPERIMENTAL) A wrapper for jaxns , a nested sampling package based on JAX.
See reference [1] for details on the meaning of each parameter. Please consider citing this reference if you use the nested sampler in your research.
Note
To enumerate over a discrete latent variable, you can add the keyword infer={“enumerate”: “parallel”} to the corresponding sample statement.
Note
To improve the performance, please consider enabling x64 mode at the beginning of your NumPyro program
numpyro.enable_x64()
.References
JAXNS: a high-performance nested sampling package based on JAX, Joshua G. Albert (https://arxiv.org/abs/2012.15286)
- Parameters
model (callable) – a call with NumPyro primitives
constructor_kwargs (dict) – additional keyword arguments to construct an upstream
jaxns.NestedSampler
instance.termination_kwargs (dict) – keyword arguments to terminate the sampler. Please refer to the upstream
jaxns.NestedSampler.__call__()
method.
Example
>>> from jax import random >>> import jax.numpy as jnp >>> import numpyro >>> import numpyro.distributions as dist >>> from numpyro.contrib.nested_sampling import NestedSampler >>> true_coefs = jnp.array([1., 2., 3.]) >>> data = random.normal(random.PRNGKey(0), (2000, 3)) >>> labels = dist.Bernoulli(logits=(true_coefs * data).sum(-1)).sample(random.PRNGKey(1)) >>> >>> def model(data, labels): ... coefs = numpyro.sample('coefs', dist.Normal(0, 1).expand([3])) ... intercept = numpyro.sample('intercept', dist.Normal(0., 10.)) ... return numpyro.sample('y', dist.Bernoulli(logits=(coefs * data + intercept).sum(-1)), ... obs=labels) >>> >>> ns = NestedSampler(model) >>> ns.run(random.PRNGKey(2), data, labels) >>> samples = ns.get_samples(random.PRNGKey(3), num_samples=1000) >>> assert jnp.mean(jnp.abs(samples['intercept'])) < 0.05 >>> print(jnp.mean(samples['coefs'], axis=0)) [0.93661342 1.95034876 2.86123884]
- run(rng_key, *args, **kwargs)[source]¶
Run the nested samplers and collect weighted samples.
- Parameters
rng_key (random.PRNGKey) – Random number generator key to be used for the sampling.
args – The arguments needed by the model.
kwargs – The keyword arguments needed by the model.
Stein Variational Inference¶
Stein Variational Inference (SteinVI) is a family of VI techniques for approximate Bayesian inference based on Stein’s method (see [1] for an overview). It is gaining popularity as it combines the scalability of traditional VI with the flexibility of non-parametric particle-based methods.
Stein variational gradient descent (SVGD) [2] is a recent SteinVI technique which uses iteratively moves a set of particles \(\{z_i\}_{i=1}^N\) to approximate a distribution p(z). SVGD is well suited for capturing correlations between latent variables as a particle-based method. The technique preserves the scalability of traditional VI approaches while offering the flexibility and modeling scope of methods such as Markov chain Monte Carlo (MCMC). SVGD is good at capturing multi-modality [3][4].
numpyro.contrib.einstein
is a framework for particle-based inference using the Stein mixture algorithm.
The framework works on Stein mixtures, a restricted mixture of guide programs parameterized by Stein particles.
Similarly to how SVGD works, Stein mixtures can approximate model posteriors by moving the Stein particles according
to the Stein forces. Because the Stein particles parameterize a guide, they capture a neighborhood rather than a
single point.
numpyro.contrib.einstein
mimics the interface from numpyro.infer.svi
, so trying SteinVI requires minimal
change to the code for existing models inferred with SVI. For primary usage, see the
Bayesian neural network example.
The framework currently supports several kernels, including:
RBFKernel
LinearKernel
RandomFeatureKernel
MixtureKernel
GraphicalKernel
ProbabilityProductKernel
For example, usage see:
References
1. Stein’s Method Meets Statistics: A Review of Some Recent Developments (2021) Andreas Anastasiou, Alessandro Barp, François-Xavier Briol, Bruno Ebner, Robert E. Gaunt, Fatemeh Ghaderinezhad, Jackson Gorham, Arthur Gretton, Christophe Ley, Qiang Liu, Lester Mackey, Chris. J. Oates, Gesine Reinert, Yvik Swan. https://arxiv.org/abs/2105.03481
2. Stein variational gradient descent: A general-purpose Bayesian inference algorithm (2016) Qiang Liu, Dilin Wang. NeurIPS
3. Nonlinear Stein Variational Gradient Descent for Learning Diversified Mixture Models (2019) Dilin Wang, Qiang Liu. PMLR
SteinVI Interface¶
- class SteinVI(model: collections.abc.Callable, guide: collections.abc.Callable, optim: numpyro.optim._NumPyroOptim, kernel_fn: numpyro.contrib.einstein.stein_kernels.SteinKernel, num_stein_particles: int = 10, num_elbo_particles: int = 10, loss_temperature: float = 1.0, repulsion_temperature: float = 1.0, non_mixture_guide_params_fn: collections.abc.Callable[[str], bool] = <function SteinVI.<lambda>>, enum=True, **static_kwargs)[source]¶
Variational inference with Stein mixtures.
Example:
>>> from jax import random >>> import jax.numpy as jnp >>> import numpyro >>> import numpyro.distributions as dist >>> from numpyro.distributions import constraints >>> from numpyro.contrib.einstein import MixtureGuidePredictive, SteinVI, RBFKernel >>> def model(data): ... f = numpyro.sample("latent_fairness", dist.Beta(10, 10)) ... with numpyro.plate("N", data.shape[0] if data is not None else 10): ... numpyro.sample("obs", dist.Bernoulli(f), obs=data) >>> def guide(data): ... alpha_q = numpyro.param("alpha_q", 15., constraint=constraints.positive) ... beta_q = numpyro.param("beta_q", lambda rng_key: random.exponential(rng_key), ... constraint=constraints.positive) ... numpyro.sample("latent_fairness", dist.Beta(alpha_q, beta_q)) >>> data = jnp.concatenate([jnp.ones(6), jnp.zeros(4)]) >>> optimizer = numpyro.optim.Adam(step_size=0.0005) >>> stein = SteinVI(model, guide, optimizer, kernel_fn=RBFKernel()) >>> stein_result = stein.run(random.PRNGKey(0), 2000, data) >>> params = stein_result.params >>> # use guide to make predictive >>> predictive = MixtureGuidePredictive(model, guide, params, num_samples=1000, guide_sites=stein.guide_sites) >>> samples = predictive(random.PRNGKey(1), data=None)
- Parameters
model (Callable) – Python callable with Pyro primitives for the model.
guide – Python callable with Pyro primitives for the guide (recognition network).
optim (_NumPyroOptim) – An instance of
_NumpyroOptim
.kernel_fn (SteinKernel) – Function that produces a logarithm of the statistical kernel to use with Stein mixture inference.
num_stein_particles – Number of particles (i.e., mixture components) in the Stein mixture.
num_elbo_particles – Number of Monte Carlo draws used to approximate the attractive force gradient. (More particles give better gradient approximations)
loss_temperature (Float) – Scaling factor of the attractive force.
repulsion_temperature (Float) – Scaling factor of the repulsive force (Non-linear Stein)
non_mixture_guide_param_fn (Callable) – predicate on names of parameters in guide which should be optimized classically without Stein (E.g. parameters for large normal networks or other transformation)
static_kwargs – Static keyword arguments for the model / guide, i.e. arguments that remain constant during inference.
SteinVI Kernels¶
- class RBFKernel(mode='norm', matrix_mode='norm_diag', bandwidth_factor: collections.abc.Callable[[float], float] = <function RBFKernel.<lambda>>)[source]¶
Calculates the Gaussian RBF kernel function, from [1], \(k(x,y) = \exp(\frac{1}{h} \|x-y\|^2)\), where the bandwidth h is computed using the median heuristic \(h = \frac{1}{\log(n)} \text{med}(\|x-y\|)\).
References:
Stein Variational Gradient Descent by Liu and Wang
- Parameters
mode (str) – Either ‘norm’ (default) specifying to take the norm of each particle, ‘vector’ to return a component-wise kernel or ‘matrix’ to return a matrix-valued kernel
matrix_mode (str) – Either ‘norm_diag’ (default) for diagonal filled with the norm kernel or ‘vector_diag’ for diagonal of vector-valued kernel
bandwidth_factor – A multiplier to the bandwidth based on data size n (default 1/log(n))
- class LinearKernel(mode='norm')[source]¶
Calculates the linear kernel \(k(x,y) = x \cdot y + 1\) from [1].
References:
Stein Variational Gradient Descent as Moment Matching by Liu and Wang
- class RandomFeatureKernel(mode='norm', bandwidth_subset=None, bandwidth_factor: collections.abc.Callable[[float], float] = <function RandomFeatureKernel.<lambda>>)[source]¶
Calculates the random kernel \(k(x,y)= 1/m\sum_{l=1}^{m}\phi(x,w_l)\phi(y,w_l)\) from [1].
References:
Stein Variational Gradient Descent as Moment Matching by Liu and Wang
- Parameters
bandwidth_subset – How many particles should be used to calculate the bandwidth? (default None, meaning all particles)
random_indices – The set of indices which to do random feature expansion on. (default None, meaning all indices)
bandwidth_factor – A multiplier to the bandwidth based on data size n (default 1/log(n))
- class MixtureKernel(ws: list[float], kernel_fns: list[numpyro.contrib.einstein.stein_kernels.SteinKernel], mode='norm')[source]¶
Calculates a mixture of multiple kernels \(k(x,y) = \sum_i w_ik_i(x,y)\)
References:
Stein Variational Gradient Descent as Moment Matching by Liu and Wang
- Parameters
ws – Weight of each kernel in the mixture
kernel_fns – Different kernel functions to mix together
- class GraphicalKernel(mode='matrix', local_kernel_fns: typing.Optional[dict[str, numpyro.contrib.einstein.stein_kernels.SteinKernel]] = None, default_kernel_fn: numpyro.contrib.einstein.stein_kernels.SteinKernel = <numpyro.contrib.einstein.stein_kernels.RBFKernel object>)[source]¶
Calculates graphical kernel \(k(x,y) = diag({K_l(x_l,y_l)})\) for local kernels \(K_l\) from [1][2].
References:
Stein Variational Message Passing for Continuous Graphical Models by Wang, Zheng, and Liu
Stein Variational Gradient Descent with Matrix-Valued Kernels by Wang, Tang, Bajaj, and Liu
- Parameters
local_kernel_fns – A mapping between parameters and a choice of kernel function for that parameter (default to default_kernel_fn for each parameter)
default_kernel_fn – The default choice of kernel function when none is specified for a particular parameter