贡献的代码

嵌套采样

嵌套采样是一种非 MCMC 方法，适用于任意概率模型，特别适合复杂的后验

NestedSampler 提供了 jaxns 的封装。有关示例，请参阅 JAXNS 的 readthedocs；有关如何在 numpyro 模型上应用采样器的示例，请参阅 Nested Sampling for Gaussian Shells。可以处理任意模型，包括具有离散随机变量和不可逆变换的模型。

class NestedSampler(model, *, constructor_kwargs=None, termination_kwargs=None)[source]

基类: object

(实验性功能) jaxns 的封装，一个基于 JAX 的嵌套采样软件包。

有关每个参数的详细含义，请参阅参考文献 [1]。如果您在研究中使用嵌套采样器，请考虑引用此参考文献。

注意

要枚举离散隐变量，您可以将关键字 infer={“enumerate”: “parallel”} 添加到相应的 sample 语句。

注意

为了提高性能，请考虑在您的 NumPyro 程序开头启用 x64 模式 numpyro.enable_x64()。

参考文献

JAXNS: a high-performance nested sampling package based on JAX, Joshua G. Albert (https://arxiv.org/abs/2012.15286)

参数:

model (callable) – 使用 NumPyro 基本操作的调用
constructor_kwargs (dict) – 用于构造上游 jaxns.NestedSampler 实例的附加关键字参数。
termination_kwargs (dict) – 用于终止采样器的关键字参数。请参考上游 jaxns.NestedSampler.__call__() 方法。

示例

>>> 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]

运行嵌套采样器并收集加权样本。

参数:

rng_key (random.PRNGKey) – 用于采样的随机数生成器密钥。
args – model 所需的参数。
kwargs – model 所需的关键字参数。

get_samples(rng_key, num_samples, *, group_by_chain=False)[source]

从运行收集的加权样本中抽取样本。

参数:

rng_key (random.PRNGKey) – 用于抽取样本的随机数生成器密钥。
num_samples (int) – 样本数量。
group_by_chain (bool) – 如果为 True，输出数组将添加一个大小为 1 的链维度。

返回值:

后验样本的字典

get_weighted_samples()[source]: 获取加权样本及其对应的对数权重。

print_summary()[source]: 打印结果摘要。这是 jaxns.utils.summary() 的封装。

diagnostics()[source]: 绘制结果诊断图。这是 jaxns.plotting.plot_diagnostics() 和 jaxns.plotting.plot_cornerplot() 的封装。

Stein 变分推断

Stein 变分推断 (SteinVI) 是一系列基于 Stein 方法的用于近似贝叶斯推断的 VI 技术（参见参考文献 [1] 获取概述）。它正日益普及，因为它结合了传统 VI 的可扩展性以及非参数粒子方法的灵活性。

Stein 变分梯度下降 (SVGD) [2] 是一种最近的 SteinVI 技术，它通过迭代移动一组粒子 \(\{z_i\}_{i=1}^N\) 来近似分布 p(z)。作为一种基于粒子的方法，SVGD 非常适合捕捉隐变量之间的相关性。该技术保留了传统 VI 方法的可扩展性，同时提供了诸如马尔可夫链蒙特卡洛 (MCMC) 等方法的灵活性和建模范围。SVGD 擅长捕捉多峰性 [3]。

numpyro.contrib.einstein 是一个用于使用 Stein 混合推断算法 [4] 进行基于粒子推断的框架。该框架作用于 Stein 混合，即由 Stein 粒子参数化的一种受限引导程序混合。类似于 SVGD 的工作方式，Stein 混合可以通过根据 Stein 力移动 Stein 粒子来近似模型后验。因为 Stein 粒子参数化了一个引导，它们捕获的是一个邻域而不是单个点。

numpyro.contrib.einstein 模拟了 numpyro.infer.svi 的接口，因此对于使用 SVI 推断的现有模型的代码只需少量修改即可尝试 SteinVI。关于主要用法，请参阅贝叶斯神经网络示例。

该框架目前支持多种核函数，包括

RBFKernel
LinearKernel
RandomFeatureKernel
MixtureKernel
GraphicalKernel
RadialGaussNewtonKernel

基于 SteinVI 的示例包括

参考文献

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。
Stein Variational Gradient Descent: A General-Purpose Bayesian Inference Algorithm。 2016.
Qiang Liu, Dilin Wang. NeurIPS
Nonlinear Stein Variational Gradient Descent for Learning Diversified Mixture Models。 2019.
Dilin Wang, Qiang Liu. PMLR
ELBOing Stein: Variational Bayes with Stein Mixture Inference。 2024.
Ola Rønning, Eric Nalisnick, Christophe Ley, Padhraic Smyth, and Thomas Hamelryck. arXiv:2410.22948。

SteinVI 接口

class SteinVI(model, guide, optim, kernel_fn, num_stein_particles=10, num_elbo_particles=10, loss_temperature=1.0, repulsion_temperature=1.0, non_mixture_guide_params_fn=<function SteinVI.<lambda>>, **static_kwargs)[source]

使用 Stein 混合推断 [1] 进行变分推断。

示例

>>> from jax import random, numpy as jnp

>>> from numpyro import sample, param, plate
>>> from numpyro.distributions import Beta, Bernoulli
>>> from numpyro.distributions.constraints import positive

>>> from numpyro.optim import Adagrad
>>> from numpyro.contrib.einstein import MixtureGuidePredictive, SteinVI, RBFKernel

>>> def model(data):
...     f = sample("fairness", Beta(10, 10))
...     n = data.shape[0] if data is not None else 1
...     with plate("N", n):
...         sample("obs", Bernoulli(f), obs=data)

>>> def guide(data):
...     # Initialize all particles in the same point.
...     alpha_q = param("alpha_q", 15., constraint=positive)
...     # Initialize particles by sampling an Exponential distribution.
...     beta_q = param("beta_q",
...                     lambda rng_key: random.exponential(rng_key),
...                     constraint=positive)
...     sample("fairness", Beta(alpha_q, beta_q))

>>> data = jnp.concatenate([jnp.ones(6), jnp.zeros(4)])

>>> opt = Adagrad(step_size=0.05)
>>> k = RBFKernel()
>>> stein = SteinVI(model, guide, opt, k, num_stein_particles=2)

>>> stein_result = stein.run(random.PRNGKey(0), 200, data)
>>> params = stein_result.params

>>> # Use guide to make predictions.
>>> predictive = MixtureGuidePredictive(model, guide, params, num_samples=10, guide_sites=stein.guide_sites)
>>> samples = predictive(random.PRNGKey(1), data=None)

参数:

model (Callable) – 使用 NumPyro 基本操作的模型 Python 可调用对象。
guide (Callable) – 使用 NumPyro 基本操作的引导 Python 可调用对象。
optim (_NumPyroOptim) – _NumpyroOptim 的一个实例。相较于 Adam [1]，更推荐使用 Adagrad。
kernel_fn (SteinKernel) – 计算用于 Stein 混合推断的再生核的函数。我们目前推荐使用 RBFKernel。由于核函数选择标准尚未完全理解，这可能会改变。
num_stein_particles – 混合近似中的粒子数量（即混合分量）。默认值为 10。
num_elbo_particles – 用于近似吸引力梯度的蒙特卡洛抽样次数。更多粒子提供更好的梯度近似。默认值为 10。
loss_temperature (Float) – 吸引力的缩放因子。默认值为 1。
repulsion_temperature (Float) – 排斥力 [2] 的缩放因子。我们建议不对排斥力进行缩放。默认值为 1。
non_mixture_guide_param_fn (Callable) – 关于引导中应使用一个粒子进行优化的参数名称的谓词。这可能是大型正态网络或其他变换的参数。默认情况下，所有参数都排除在此选项之外。
static_kwargs – 用于模型和引导的静态关键字参数。这些参数在推断过程中不能改变。

参考文献

Rønning, Ola, et al. “ELBOing Stein: Variational Bayes with Stein Mixture Inference。”
arXiv preprint arXiv:2410.22948 (2024)。
Liu, Chang, et al. “Understanding and Accelerating Particle-Based Variational Inference。”
International Conference on Machine Learning. PMLR, 2019。
Wang, Dilin, and Qiang Liu. “Nonlinear Stein Variational Gradient Descent for Learning Diversified Mixture Models。”
International Conference on Machine Learning. PMLR, 2019。

class SVGD(model, optim, kernel_fn, num_stein_particles=10, guide_kwargs={}, **static_kwargs)[source]

Stein 变分梯度下降 [1]。

示例

>>> from jax import random, numpy as jnp

>>> from numpyro import sample, param, plate
>>> from numpyro.distributions import Beta, Bernoulli
>>> from numpyro.distributions.constraints import positive

>>> from numpyro.optim import Adagrad
>>> from numpyro.contrib.einstein import SVGD, RBFKernel
>>> from numpyro.infer import Predictive

>>> def model(data):
...     f = sample("fairness", Beta(10, 10))
...     n = data.shape[0] if data is not None else 1
...     with plate("N", n):
...         sample("obs", Bernoulli(f), obs=data)

>>> data = jnp.concatenate([jnp.ones(6), jnp.zeros(4)])

>>> opt = Adagrad(step_size=0.05)
>>> k = RBFKernel()
>>> svgd = SVGD(model, opt, k, num_stein_particles=2)

>>> svgd_result = svgd.run(random.PRNGKey(0), 200, data)

>>> params = svgd_result.params
>>> predictive = Predictive(model, guide=svgd.guide, params=params, num_samples=10, batch_ndims=1)
>>> samples = predictive(random.PRNGKey(1), data=None)

参数:

model (Callable) – 使用 NumPyro 基本操作的模型 Python 可调用对象。
guide (Callable) – 使用 NumPyro 基本操作的引导 Python 可调用对象。
optim (_NumPyroOptim) – _NumpyroOptim 的一个实例。相较于 Adam [1]，更推荐使用 Adagrad。
kernel_fn (SteinKernel) – 计算用于 SVGD 的再生核的函数。我们目前推荐使用 RBFKernel。由于核函数选择标准尚未完全理解，这可能会改变。
num_stein_particles – 混合近似中的粒子数量（即混合分量）。默认值为 10。

guide_kwargs (Dict) –

AutoDelta 的关键字参数。默认行为与 AutoDelta 的默认行为相同。

用法

opt = Adagrad(step_size=0.05)
k = RBFKernel()
svgd = SVGD(model, opt, k, guide_kwargs={'init_loc_fn': partial(init_to_uniform, radius=0.1)})

static_kwargs (Dict) – 用于模型和引导的静态关键字参数。这些参数在推断过程中不能改变。

参考文献

Liu, Qiang, and Dilin Wang. “Stein Variational Gradient Descent: A General Purpose Bayesian Inference Algorithm。”
Advances in neural information processing systems 29 (2016)。

class ASVGD(model, optim, kernel_fn, num_stein_particles=10, num_cycles=10, transition_speed=10, guide_kwargs={}, **static_kwargs)[source]

退火 Stein 变分梯度下降 [1]。

示例

>>> from jax import random, numpy as jnp

>>> from numpyro import sample, param, plate
>>> from numpyro.distributions import Beta, Bernoulli
>>> from numpyro.distributions.constraints import positive

>>> from numpyro.optim import Adagrad
>>> from numpyro.contrib.einstein import ASVGD, RBFKernel
>>> from numpyro.infer import Predictive

>>> def model(data):
...     f = sample("fairness", Beta(10, 10))
...     n = data.shape[0] if data is not None else 1
...     with plate("N", n):
...         sample("obs", Bernoulli(f), obs=data)

>>> data = jnp.concatenate([jnp.ones(6), jnp.zeros(4)])

>>> opt = Adagrad(step_size=0.05)
>>> k = RBFKernel()
>>> asvgd = ASVGD(model, opt, k, num_stein_particles=2)

>>> asvgd_result = asvgd.run(random.PRNGKey(0), 200, data)

>>> params = asvgd_result.params
>>> predictive = Predictive(model, guide=asvgd.guide, params=params, num_samples=10, batch_ndims=1)
>>> samples = predictive(random.PRNGKey(1), data=None)

参数:

model (Callable) – 使用 NumPyro 基本操作的模型 Python 可调用对象。
guide (Callable) – 使用 NumPyro 基本操作的引导 Python 可调用对象。
optim (_NumPyroOptim) – _NumpyroOptim 的一个实例。相较于 Adam [1]，更推荐使用 Adagrad。
kernel_fn (SteinKernel) – 计算用于 ASVGD 的再生核的函数。我们目前推荐使用 RBFKernel。由于核函数选择标准尚未完全理解，这可能会改变。
num_stein_particles – 混合近似中的粒子数量（即混合分量）。默认值为 10。
num_cycles – 推断期间的总周期数。这对应于参考文献 [1] 方程 4 中的 \(C\)。默认值为 10。
transition_speed – 推断期间两个阶段之间的过渡速度。这对应于参考文献 [1] 方程 4 中的 \(p\)。默认值为 10。

guide_kwargs (Dict) –

AutoDelta 的关键字参数。默认行为与 AutoDelta 的默认行为相同。

用法

opt = Adagrad(step_size=0.05)
k = RBFKernel()
asvgd = ASVGD(model, opt, k, guide_kwargs={'init_loc_fn': partial(init_to_uniform, radius=0.1)})

static_kwargs (Dict) – 用于模型和引导的静态关键字参数。这些参数在推断过程中不能改变。

参考文献

D’Angelo, Francesco, and Vincent Fortuin. “Annealed Stein Variational Gradient Descent。”
Third Symposium on Advances in Approximate Bayesian Inference, 2021。

SteinVI 核函数

class RBFKernel(mode='norm', matrix_mode='norm_diag', bandwidth_factor: ~collections.abc.Callable[[float], float] = <function RBFKernel.<lambda>>)[source]

计算在 [1] 中使用的高斯 RBF 核函数。该核函数由以下公式给出

\(k(x,y) = \exp(\frac{-1}{h} \|x-y\|^2)\),

其中带宽 \(h\) 使用中位数启发式计算

\(h = \frac{1}{\log(m)} \text{med}(\|x-y\|)\).

上述公式中，\(m\) 是粒子数量。

参数:

mode (str) – ‘norm’（默认），指定取每个粒子的范数；‘vector’，返回逐分量核；或 ‘matrix’，返回矩阵值核。
matrix_mode (str) – ‘norm_diag’（默认），对角线填充范数核；或 ‘vector_diag’，对角线为向量值核。
bandwidth_factor – 基于数据大小 n 的带宽乘数（默认 1/log(n)）

参考文献

Liu, Qiang, and Dilin Wang. “Stein Variational Gradient Descent: A General Purpose Bayesian Inference Algorithm。”
Advances in neural information processing systems 29 (2016)。

class LinearKernel(mode='norm')[source]

计算参考文献 [1] 中定理 3.3 的线性核函数。该核函数由以下公式给出

\(k(x,y) = x^T y + 1\).

参考文献

Liu, Qiang, and Dilin Wang. “Stein Variational Gradient Descent as Moment Matching。”
Advances in Neural Information Processing Systems 31 (2018)。

class RandomFeatureKernel(mode='norm', bandwidth_subset=None, bandwidth_factor: ~collections.abc.Callable[[float], float] = <function RandomFeatureKernel.<lambda>>)[source]

计算参考文献 [1] 的方程 5 和 6 中的随机核函数的高斯变量。该核函数由以下公式给出

\(k(x,y)= \frac{1}{m}\sum_{l=1}^{m}\phi(x,w_l)\phi(y,w_l)\),

其中 \(\phi(\cdot, w)\) 是方程 6 中的高斯随机特征映射。该映射由以下公式给出

\(\phi(z, w) = \sqrt{2}\left(h^{-1}w_1^Tz + w_0\right)\),

其中 \(h\) 是使用 RBFKernel 中的中位数技巧计算的带宽，\(w_0\sim\text{Uni}([0,2\pi])\)，\(w_1\sim\mathcal{N}(0,I)\)。

参数:

bandwidth_subset – 应使用多少粒子来计算带宽？（默认 None，表示所有粒子）
random_indices – 执行随机特征扩展的索引集（默认 None，表示所有索引）
bandwidth_factor – 基于数据大小 n 的带宽乘数（默认 1/log(n)）

参考文献

Liu, Qiang, and Dilin Wang. “Stein Variational Gradient Descent as Moment Matching。”
Advances in Neural Information Processing Systems 31 (2018)。

class MixtureKernel(ws: list[float], kernel_fns: list[SteinKernel], mode='norm')[source]

计算参考文献 [1] 方程 1 中的多个核函数的混合。该核函数由以下公式给出

\(k(x,y) = \sum_i w_ik_i(x,y)\),

其中 \(k_i\) 是一个再生核，且 \(w_i \in (0,\infty)\)。

参数:

ws – 混合中每个核的权重。
kernel_fns – 需要混合的不同核函数。

参考文献

Ai, Qingzhong, et al. “Stein variational gradient descent with multiple kernels。”
Cognitive Computation 15.2 (2023): 672-682。

class GraphicalKernel(mode='matrix', local_kernel_fns: dict[str, ~numpyro.contrib.einstein.stein_kernels.SteinKernel] | None = None, default_kernel_fn: ~numpyro.contrib.einstein.stein_kernels.SteinKernel = <numpyro.contrib.einstein.stein_kernels.RBFKernel object>)[source]

计算来自参考文献 [1] 中定理 1 的图核，也称为坐标核。该核函数由以下公式给出

\(k(x,y) = diag({k_l(x_l,y_l)})\),

对于坐标核 \(k_l\)。

参数:

local_kernel_fns – 参数与该参数核函数选择之间的映射（默认情况下，每个参数使用 default_kernel_fn）
default_kernel_fn – 当未为特定参数指定核函数时，默认的核函数选择。

参考文献

Wang, Dilin, Zhe Zeng, and Qiang Liu. “Stein variational message passing for continuous graphical models。”
International Conference on Machine Learning. PMLR, 2018。

class ProbabilityProductKernel(guide, scale=1.0, mode='norm')[source]

实验性功能 计算参考文献 [1] 方程 5 给出高斯的非归一化概率积核。

参考文献:

Jebara, Tony, Risi Kondor, and Andrew Howard. “Probability product kernels。”
The Journal of Machine Learning Research 5 (2004): 819-844。

class RadialGaussNewtonKernel[source]

径向高斯-牛顿核 [1,2]，也称为缩放 Hessian 核，是一个标量核，定义为

\[k(x,y) = \exp\left(-\frac{1}{2d}(x-y)^T M (x-y)\right),\]

其中 \(x,y \in R^d\) 是粒子，且 \(M\) 是一个度量矩阵。\(M\) 使用 Hessian 近似 \(A(x)\) 来近似对数后验的期望曲率。

矩阵 \(M\) 使用 \(m\) 个粒子计算如下 [2, Eq. 19, p.5]

\[M= \frac{1}{m} \sum_{i=1}^m A(x_i)\]

其中 Hessian 近似由以下公式给出

\[A(x) = J(x) J(x)^T,\]

其中 \(J(x)\) 是 ELBO 在 \(x\) 处的雅可比矩阵。

参考文献:

Maken, Fahira Afzal, Fabio Ramos, and Lionel Ott. “Stein Particle Filter for Nonlinear,
Non-Gaussian State Estimation。” IEEE Robotics and Automation Letters 7.2 (2022)。
Detommaso, Gianluca, et al. “A Stein variational Newton method。”
Advances in Neural Information Processing Systems 31 (2018)。

随机支持

class StochasticSupportInference(model: ~typing.Callable, num_slp_samples: int, max_slps: int)[source]

基类: ABC

在具有随机支持的程序中运行推断的基类。每个子类将输入模型分解为所谓的直线程序 (SLPs)，它们是模型中不同的控制流路径。然后，在每个 SLP 中单独运行推断，并将结果组合起来以产生整体后验。

注意

此实现假设所有随机分支都是基于使用 infer={"branching": True} 进行标注的离散采样站点的结果完成的。例如，

def model():
    model1 = numpyro.sample("model1", dist.Bernoulli(0.5), infer={"branching": True})
    if model1 == 0:
        mean = numpyro.sample("a1", dist.Normal(0.0, 1.0))
    else:
        mean = numpyro.sample("a2", dist.Normal(1.0, 1.0))
    numpyro.sample("obs", dist.Normal(mean, 1.0), obs=0.2)

参数:

model – 包含 Pyro 基本操作 primitives 的 Python 可调用对象。默认值为 NUTS。
num_slp_samples (int) – 从先验中抽取的样本数量，用于发现直线程序 (SLPs)。
max_slps (int) – 要发现的最大 SLP 数量。DCC 将不会在超过 max_slps 的 SLP 上运行推断。

在每个 SLP 上单独运行推断并合并结果。

参数:

rng_key (jax.random.PRNGKey) – 随机数生成器密钥。
args – 模型的参数。
kwargs – 模型的关键字参数。

class DCC(model: ~typing.Callable, mcmc_kwargs: dict[str, ~typing.Any], kernel_cls: type[~numpyro.infer.mcmc.MCMCKernel] = <class 'numpyro.infer.hmc.NUTS'>, num_slp_samples: int = 1000, max_slps: int = 124, proposal_scale: float = 1.0)[source]

基类: StochasticSupportInference

实现了来自 [1] 的用于具有随机支持的模型的分而治之和合并 (DCC) 算法。

参考文献

Divide, Conquer, and Combine: a New Inference Strategy for Probabilistic Programs with Stochastic Support, Yuan Zhou, Hongseok Yang, Yee Whye Teh, Tom Rainforth

示例

def model():
    model1 = numpyro.sample("model1", dist.Bernoulli(0.5), infer={"branching": True})
    if model1 == 0:
        mean = numpyro.sample("a1", dist.Normal(0.0, 1.0))
    else:
        mean = numpyro.sample("a2", dist.Normal(1.0, 1.0))
    numpyro.sample("obs", dist.Normal(mean, 1.0), obs=0.2)

mcmc_kwargs = dict(
    num_warmup=500, num_samples=1000
)
dcc = DCC(model, mcmc_kwargs=mcmc_kwargs)
dcc_result = dcc.run(random.PRNGKey(0))

参数:

model – 包含 Pyro 基本操作 primitives 的 Python 可调用对象。
mcmc_kwargs (dict) – 传递给 MCMC 的参数字典。
kernel_cls (numpyro.infer.mcmc.MCMCKernel) – 用于局部推断的 MCMC 核函数类。默认值为 NUTS。
num_slp_samples (int) – 从先验中抽取的样本数量，用于发现直线程序 (SLPs)。
max_slps (int) – 要发现的最大 SLP 数量。DCC 将不会在超过 max_slps 的 SLP 上运行推断。
proposal_scale (float) – 用于估计 SLP 归一化常数的提议分布尺度参数。

class SDVI(model: ~typing.Callable, optimizer, svi_num_steps: int = 1000, combine_elbo_particles: int = 1000, guide_init: ~typing.Callable = <class 'numpyro.infer.autoguide.AutoNormal'>, loss: ~numpyro.infer.elbo.ELBO = <numpyro.infer.elbo.Trace_ELBO object>, svi_progress_bar: bool = False, num_slp_samples: int = 1000, max_slps: int = 124)[source]

基类: StochasticSupportInference

实现了来自 [1] 的用于具有随机支持的模型支撑分解变分推断 (SDVI) 算法。此实现为每个 SLP 创建一个单独的引导，分别训练引导，然后通过根据它们的 ELBO 估计值按比例加权来合并引导。

参考文献

Rethinking Variational Inference for Probabilistic Programs with Stochastic Support, Tim Reichelt, Luke Ong, Tom Rainforth

示例

def model():
    model1 = numpyro.sample("model1", dist.Bernoulli(0.5), infer={"branching": True})
    if model1 == 0:
        mean = numpyro.sample("a1", dist.Normal(0.0, 1.0))
    else:
        mean = numpyro.sample("a2", dist.Normal(1.0, 1.0))
    numpyro.sample("obs", dist.Normal(mean, 1.0), obs=0.2)

sdvi = SDVI(model, numpyro.optim.Adam(step_size=0.001))
sdvi_result = sdvi.run(random.PRNGKey(0))

参数:

model – 包含 Pyro 基本操作 primitives 的 Python 可调用对象。
optimizer – _NumpyroOptim、jax.example_libraries.optimizers.Optimizer 或 Optax GradientTransformation 的一个实例。传递给 SVI。
svi_num_steps (int) – 为每个 SLP 运行 SVI 的步数。
combine_elbo_particles (int) – 估计 ELBO 的粒子数量，用于计算 SLP 权重。
guide_init – 引导的构造函数。这应该是一个可调用对象，返回一个 AutoGuide 实例。默认值为 AutoNormal。
loss – SVI 的 ELBO 损失函数。默认值为 Trace_ELBO。
svi_progress_bar (bool) – 是否为 SVI 使用进度条。
num_slp_samples (int) – 从先验中抽取的样本数量，用于发现直线程序 (SLPs)。
max_slps (int) – 要发现的最大 SLP 数量。DCC 将不会在超过 max_slps 的 SLP 上运行推断。

希尔伯特空间高斯过程近似

本模块包含用于在 [1] 和 [2] 中描述的希尔伯特空间高斯过程 (HSGP) 近似方法中的辅助函数。

警告

本模块为实验性功能。

为什么需要近似？

高斯过程的计算量随数据点数量增加而迅速增长。回想一下，我们需要对核矩阵求逆！高斯过程模型的计算复杂度为 \(\mathcal{O}(n^3)\)，其中 \(n\) 是数据点数量。HSGP 近似方法是一种将高斯过程模型的计算复杂度降低到 \(\mathcal{O}(mn + m)\) 的方法，其中 \(m\) 是近似中使用的基函数数量。

近似策略步骤

我们强烈建议阅读 [1] 和 [2] 以详细了解近似方法。在 [3] 中，您可以找到使用 NumPyro 和 PyMC 的实践方法。

这里我们提供近似方法的主要步骤和组成部分

每个平稳核 \(k\) 都有一个相关的谱密度 \(S(\omega)\)。对于最常见的核函数，有闭合公式。这些公式取决于核函数的超参数（例如，幅度（amplitudes）和长度尺度（length scales））。

我们可以将谱密度 \(S(\omega)\) 近似为 \(||\omega||\) 的多项式级数。我们称 \(\omega\) 为频率。

我们可以将这些多项式项解释为拉普拉斯算子的“幂”。关键的观察是，拉普拉斯算子的傅里叶变换是 \(||\omega||^2\)。

接下来，我们对拉普拉斯算子施加 Dirichlet 边界条件，这使得它成为自伴随算子并具有离散谱。

我们将 (2) 中的展开式与拉普拉斯算子在 (4) 的特征基中的幂之和相对应。

令 \(m^\star = \prod_{d=1}^D m_d\) 为近似的总项数，其中 \(m_d\) 是第 \(d\) 维近似中使用的基函数数量。那么，在非中心化参数化中，近似公式为

\[f(x) \approx \sum_{j = 1}^{m^\star} \overbrace{\color{red}{\left(S(\sqrt{\boldsymbol{\lambda}_j})\right)^{1/2}}}^{\text{所有超参数都在这里！}} \times \underbrace{\color{blue}{\phi_{j}(\boldsymbol{x})}}_{\text{易于计算！}} \times \overbrace{\color{green}{\beta_{j}}}^{\sim \: \text{Normal}(0,1)}\]

其中 \(\boldsymbol{x}\) 是一个 \(D\) 维输入向量，\(\boldsymbol{\lambda}_j\) 是拉普拉斯算子的特征值，\(\phi_{j}(\boldsymbol{x})\) 是拉普拉斯算子的特征函数，且 \(beta_{j}\) 是展开式的系数（参见参考文献 [2] 的方程 (8)）。我们期望对于级数中有限的 \(m^\star\) 项，这是一个很好的近似，只要输入值 \(x\) 不太接近边界 \(-L_d\) 和 \(L_d\)。

注意

即使周期核函数不是平稳的，仍然可以调整并找到类似的近似公式。然而，这些核函数不支持多维输入。有关更多详细信息，请参阅参考文献 [2] 的附录 B。

示例

这里有一个如何使用 NumPyro 的 HSGP 近似方法的示例。我们将使用平方指数核函数。其他核函数也可以类似地使用。

>>> from jax import random
>>> import jax.numpy as jnp

>>> import numpyro
>>> from numpyro.contrib.hsgp.approximation import hsgp_squared_exponential
>>> import numpyro.distributions as dist
>>> from numpyro.infer import MCMC, NUTS


>>> def generate_synthetic_data(rng_key, start, stop: float, num, scale):
...     """Generate synthetic data."""
...     x = jnp.linspace(start=start, stop=stop, num=num)
...     y = jnp.sin(4 * jnp.pi * x) + jnp.sin(7 * jnp.pi * x)
...     y_obs = y + scale * random.normal(rng_key, shape=(num,))
...     return x, y_obs


>>> rng_key = random.PRNGKey(seed=42)
>>> rng_key, rng_subkey = random.split(rng_key)
>>> x, y_obs = generate_synthetic_data(
...     rng_key=rng_subkey, start=0, stop=1, num=80, scale=0.3
>>> )


>>> def model(x, ell, m, non_centered, y=None):
...     # --- Priors ---
...     alpha = numpyro.sample("alpha", dist.InverseGamma(concentration=12, rate=10))
...     length = numpyro.sample("length", dist.InverseGamma(concentration=6, rate=1))
...     noise = numpyro.sample("noise", dist.InverseGamma(concentration=12, rate=10))
...     # --- Parametrization ---
...     f = hsgp_squared_exponential(
...         x=x, alpha=alpha, length=length, ell=ell, m=m, non_centered=non_centered
...     )
...     # --- Likelihood ---
...     with numpyro.plate("data", x.shape[0]):
...         numpyro.sample("likelihood", dist.Normal(loc=f, scale=noise), obs=y)


>>> sampler = NUTS(model)
>>> mcmc = MCMC(sampler=sampler, num_warmup=500, num_samples=1_000, num_chains=2)

>>> rng_key, rng_subkey = random.split(rng_key)

>>> ell = 1.3
>>> m = 20
>>> non_centered = True

>>> mcmc.run(rng_subkey, x, ell, m, non_centered, y_obs)

>>> mcmc.print_summary()

              mean       std    median      5.0%     95.0%     n_eff     r_hat
   alpha      1.24      0.34      1.18      0.72      1.74   1804.01      1.00
 beta[0]     -0.10      0.66     -0.10     -1.24      0.92   1819.91      1.00
 beta[1]      0.00      0.71     -0.01     -1.09      1.26   1872.82      1.00
 beta[2]     -0.05      0.69     -0.03     -1.09      1.16   2105.88      1.00
 beta[3]      0.25      0.74      0.26     -0.98      1.42   2281.30      1.00
 beta[4]     -0.17      0.69     -0.17     -1.21      1.00   2551.39      1.00
 beta[5]      0.09      0.75      0.10     -1.13      1.30   3232.13      1.00
 beta[6]     -0.49      0.75     -0.49     -1.65      0.82   3042.31      1.00
 beta[7]      0.42      0.75      0.44     -0.78      1.65   2885.42      1.00
 beta[8]      0.69      0.71      0.71     -0.48      1.82   2811.68      1.00
 beta[9]     -1.43      0.75     -1.40     -2.63     -0.21   2858.68      1.00
beta[10]      0.33      0.71      0.33     -0.77      1.51   2198.65      1.00
beta[11]      1.09      0.73      1.11     -0.23      2.18   2765.99      1.00
beta[12]     -0.91      0.72     -0.91     -2.06      0.31   2586.53      1.00
beta[13]      0.05      0.70      0.04     -1.16      1.12   2569.59      1.00
beta[14]     -0.44      0.71     -0.44     -1.58      0.73   2626.09      1.00
beta[15]      0.69      0.73      0.70     -0.45      1.88   2626.32      1.00
beta[16]      0.98      0.74      0.98     -0.15      2.28   2282.86      1.00
beta[17]     -2.54      0.77     -2.52     -3.82     -1.29   3347.56      1.00
beta[18]      1.35      0.66      1.35      0.30      2.46   2638.17      1.00
beta[19]      1.10      0.54      1.09      0.25      2.01   2428.37      1.00
  length      0.07      0.01      0.07      0.06      0.09   2321.67      1.00
   noise      0.33      0.03      0.33      0.29      0.38   2472.83      1.00

Number of divergences: 0

注意

可以在 [3]、[4] 和 [5] 中找到包含代码的更多示例。

参考文献

Solin, A., Särkkä, S. Hilbert space methods for reduced-rank Gaussian process regression. Stat Comput 30, 419-446 (2020)。

Riutort-Mayol, G., Bürkner, PC., Andersen, M.R. et al. Practical Hilbert space approximate Bayesian Gaussian processes for probabilistic programming. Stat Comput 33, 17 (2023)。

Orduz, J., A Conceptual and Practical Introduction to Hilbert Space GPs Approximation Methods.

Example: Hilbert space approximation for Gaussian processes.

Gelman, Vehtari, Simpson, et al., Bayesian workflow book - Birthdays.

注意

本模块的代码基于由 Omar Sosa Rodríguez 贡献的示例 Example: Hilbert space approximation for Gaussian processes 的代码。

eigenindices

eigenindices(m: list[int] | int, dim: int) → Array[source]

返回拉普拉斯算子前 \(D \times m^\star\) 个特征值的索引。

\[m^\star = \prod_{i=1}^D m_i\]

更多详细信息请参阅 [1] 中的方程 (10)。

参考文献

Riutort-Mayol, G., Bürkner, PC., Andersen, M.R. et al. Practical Hilbert space approximate Bayesian Gaussian processes for probabilistic programming. Stat Comput 33, 17 (2023)。

参数:

m (list[int] | int) – 每维所需的特征值索引数量。如果是整数，则在每一维计算相同数量的特征值。
dim (int) – 空间的维度。

返回值:

前 \(D \times m^\star\) 个特征值索引的数组。

返回类型:

Array

示例

>>> import jax.numpy as jnp

>>> from numpyro.contrib.hsgp.laplacian import eigenindices

>>> m = 10
>>> S = eigenindices(m, 1)
>>> assert S.shape == (1, m)
>>> S
Array([[ 1,  2,  3,  4,  5,  6,  7,  8,  9, 10]], dtype=int32)

>>> m = 10
>>> S = eigenindices(m, 2)
>>> assert S.shape == (2, 100)

>>> m = [2, 2, 3]  # Riutort-Mayol et al eq (10)
>>> S = eigenindices(m, 3)
>>> assert S.shape == (3, 12)
>>> S
Array([[1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2],
       [1, 1, 1, 2, 2, 2, 1, 1, 1, 2, 2, 2],
       [1, 2, 3, 1, 2, 3, 1, 2, 3, 1, 2, 3]], dtype=int32)

sqrt_eigenvalues

拉普拉斯算子在 \( guesswork: [-L_1, L_1] \times ... \times [-L_D, L_D] \) 空间中的前 \(m^\star \times D\) 个特征值的平方根。参见参考文献 [1] 中的方程 (56)。

参考文献

Solin, A., Särkkä, S. Hilbert space methods for reduced-rank Gaussian process regression. Stat Comput 30, 419-446 (2020)

参数:

ell (int | float | list[int | float]) – 每维区间长度除以 2。如果是浮点数，则在每一维使用相同的长度。
m (list[int] | int) – 每维要计算的特征值数量。如果是整数，则在每一维计算相同数量的特征值。
dim (int) – 空间的维度。

返回值:

前 \(m^\star \times D\) 个特征值的平方根数组。

返回类型:

Array

eigenfunctions

在\([-L_1, L_1] \times ... \times [-L_D, L_D]\)区间上，拉普拉斯算子的前\(m^\star\)个特征函数在x值处进行评估。参见[1]中的公式 (56)。如果x是一维的，则问题被假定为一维。否则，输入空间的维度根据x的最后一个维度的大小推断。其他维度被视为批处理维度。

示例

>>> import jax.numpy as jnp

>>> from numpyro.contrib.hsgp.laplacian import eigenfunctions

>>> n = 100
>>> m = 10

>>> x = jnp.linspace(-1, 1, n)

>>> basis = eigenfunctions(x=x, ell=1.2, m=m)

>>> assert basis.shape == (n, m)

>>> x = jnp.ones((n, 3))  # 2d input
>>> basis = eigenfunctions(x=x, ell=1.2, m=[2, 2, 3])
>>> assert basis.shape == (n, 12)

参考文献

Solin, A., Särkkä, S. Hilbert space methods for reduced-rank Gaussian process regression. Stat Comput 30, 419-446 (2020)

参数:

x (ArrayLike) – 要评估特征函数的点。如果x是一维的，则问题被假定为一维。否则，输入空间的维度根据x的最后一个维度推断。其他维度被视为批处理维度。
ell (float | list[float]) – 每个维度上区间长度的一半。如果为浮点数，则每个维度使用相同的长度。
m (int | list[int]) – 每个维度要计算的特征值的数量。如果为整数，则每个维度计算相同数量的特征值。

返回值:

在x处评估的前\(m^\star \times D\)个特征函数的数组。

返回类型:

Array

eigenfunctions_periodic

用于近似周期核的基函数。

参数:

x (ArrayLike) – 要评估特征函数的点。
w0 (float) – 周期核的频率。
m (int) – 要计算的特征函数的数量。

注意

如果你想用周期来参数化它，请使用 w0 = 2 * jnp.pi / period。

警告

不支持多维输入。

spectral_density_squared_exponential

平方指数核的谱密度。

参见[1]中的第 4.2 节和[2]中的第 2.1 节。

\[S(\boldsymbol{\omega}) = \alpha (\sqrt{2\pi})^D \ell^D \exp\left(-\frac{1}{2} \ell^2 \boldsymbol{\omega}^{T} \boldsymbol{\omega}\right)\]

参考文献

Rasmussen, C. E., & Williams, C. K. I. (2006). Gaussian Processes for Machine Learning.

Riutort-Mayol, G., Bürkner, PC., Andersen, M.R. et al. Practical Hilbert space approximate Bayesian Gaussian processes for probabilistic programming. Stat Comput 33, 17 (2023)。

参数:

dim (int) – 维度
w (ArrayLike) – 频率
alpha (float) – 幅度
length (float) – 长度尺度

返回值:

谱密度值

返回类型:

Array

spectral_density_matern

Matérn 核的谱密度。

参见[1]中的公式 (4.15) 和[2]中的第 2.1 节。

\[S(\boldsymbol{\omega}) = \alpha \frac{2^{D} \pi^{D/2} \Gamma(\nu + D/2) (2 \nu)^{\nu}}{\Gamma(\nu) \ell^{2 \nu}} \left(\frac{2 \nu}{\ell^2} + \boldsymbol{\omega}^{T} \boldsymbol{\omega}\right)^{-\nu - D/2}\]

参考文献

Rasmussen, C. E., & Williams, C. K. I. (2006). Gaussian Processes for Machine Learning.

Riutort-Mayol, G., Bürkner, PC., Andersen, M.R. et al. Practical Hilbert space approximate Bayesian Gaussian processes for probabilistic programming. Stat Comput 33, 17 (2023)。

参数:

dim (int) – 维度
nu (float) – 平滑度
w (ArrayLike) – 频率
alpha (float) – 幅度
length (float) – 长度尺度

返回值:

谱密度值

返回类型:

浮点数

diag_spectral_density_squared_exponential

在\([-L_1, L_1] \times ... \times [-L_D, L_D]\)区间上，评估平方指数核的谱密度在拉普拉斯算子的前\(D \times m^\star\)个平方根特征值处的值。

参数:

alpha (float) – 平方指数核的幅度
length (float) – 平方指数核的长度尺度
ell (float | int | list[float | int]) – 每个维度上区间长度的一半。如果为浮点数或整数，则每个维度使用相同的长度。
m (int | list[int]) – 每个维度要计算并包含在近似中的特征值的数量。如果为整数，则每个维度计算相同数量的特征值。
dim (int) – 空间的维度

返回值:

在拉普拉斯算子的前\(D \times m^\star\)个平方根特征值处评估的谱密度向量

返回类型:

Array

diag_spectral_density_matern

diag_spectral_density_matern(nu: float, alpha: float, length: float, ell: float | int | list[float | int], m: int | list[int], dim: int) → Array[source]

在\([-L_1, L_1] \times ... \times [-L_D, L_D]\)区间上，评估 Matérn 核的谱密度在拉普拉斯算子的前\(D \times m^\star\)个平方根特征值处的值。

参数:

nu (float) – 平滑度参数
alpha (float) – Matérn 核的幅度
length (float) – Matérn 核的长度尺度
ell (float | int | list[float | int]) – 每个维度上区间长度的一半。如果为浮点数或整数，则每个维度使用相同的长度。
m (int | list[int]) – 每个维度要计算并包含在近似中的特征值的数量。如果为整数，则每个维度计算相同数量的特征值。
dim (int) – 空间的维度

返回值:

在拉普拉斯算子的前\(D \times m^\star\)个平方根特征值处评估的谱密度向量

返回类型:

Array

diag_spectral_density_periodic

diag_spectral_density_periodic(alpha: float, length: float, m: int) → Array[source]

实际上不是谱密度，但它们的使用方式相同。这些只是周期核的低秩近似的前m个系数。参见[1]中的附录 B。

参考文献

Riutort-Mayol, G., Bürkner, PC., Andersen, M.R. et al. Practical Hilbert space approximate Bayesian Gaussian processes for probabilistic programming. Stat Comput 33, 17 (2023)。

参数:

alpha (float) – 幅度
length (float) – 长度尺度
m (int) – 特征值的数量

返回值:

“谱密度”向量

返回类型:

Array

hsgp_squared_exponential

使用平方指数核进行希尔伯特空间高斯过程近似。

该方法的主要思想是结合平方指数核的相关谱密度和 Dirichlet 拉普拉斯算子的谱，以获得 Gram 矩阵的低秩近似。更多详情请参见 [1, 2]。

参考文献

Solin, A., Särkkä, S. Hilbert space methods for reduced-rank Gaussian process regression. Stat Comput 30, 419-446 (2020)。

Riutort-Mayol, G., Bürkner, PC., Andersen, M.R. et al. Practical Hilbert space approximate Bayesian Gaussian processes for probabilistic programming. Stat Comput 33, 17 (2023)。

参数:

x (ArrayLike) – 输入数据
alpha (float) – 平方指数核的幅度
length (float) – 平方指数核的长度尺度
ell (float | int | list[float | int]) – 参数化 D 维盒长度的正值，使得输入数据位于区间 \([-L_1, L_1] \times ... \times [-L_D, L_E]\) 内。我们期望此近似在此区间内有效。
m (int | list[m]) – 每个维度（\(\left\{1, ..., D\right\}\)）要计算并包含在近似中的特征值的数量。如果为整数，则每个维度计算相同数量的特征值。
non_centered (bool) – 是否使用非中心化参数化。默认为 True。

返回值:

低秩近似线性模型

返回类型:

Array

hsgp_matern

使用 Matérn 核进行希尔伯特空间高斯过程近似。

该方法的主要思想是结合 Matérn 核的相关谱密度和 Dirichlet 拉普拉斯算子的谱，以获得 Gram 矩阵的低秩近似。更多详情请参见 [1, 2]。

参考文献

Solin, A., Särkkä, S. Hilbert space methods for reduced-rank Gaussian process regression. Stat Comput 30, 419-446 (2020)。

Riutort-Mayol, G., Bürkner, PC., Andersen, M.R. et al. Practical Hilbert space approximate Bayesian Gaussian processes for probabilistic programming. Stat Comput 33, 17 (2023)。

参数:

x (ArrayLike) – 输入数据
nu (float) – 平滑度参数
alpha (float) – 平方指数核的幅度
length (float) – 平方指数核的长度尺度
ell (float | int | list[float | int]) – 参数化 D 维盒长度的正值，使得输入数据位于区间 \([-L_1, L_1] \times ... \times [-L_D, L_D]\) 内。我们期望此近似在此区间内有效。
m (int | list[m]) – 每个维度（\(\left\{1, ..., D\right\}\)）要计算并包含在近似中的特征值的数量。如果为整数，则每个维度计算相同数量的特征值。
non_centered (bool) – 是否使用非中心化参数化。默认为 True。

返回值:

低秩近似线性模型

返回类型:

Array

hsgp_periodic_non_centered

在非中心化参数化中使用周期平方指数核的低秩近似。

参见[1]中的附录 B。

参考文献

Riutort-Mayol, G., Bürkner, PC., Andersen, M.R. et al. Practical Hilbert space approximate Bayesian Gaussian processes for probabilistic programming. Stat Comput 33, 17 (2023)。

参数:

x (ArrayLike) – 输入数据
alpha (float) – 幅度
length (float) – 长度尺度
w0 (float) – 周期核的频率
m (int) – 要计算并包含在近似中的特征值的数量

返回值:

低秩近似线性模型

返回类型:

Array