Pathfinder 变分推断#

Pathfinder [Zhang 等人, 2021] 是一种变分推断算法，能够从贝叶斯模型的后验分布中生成样本。它相较于广泛使用的ADVI算法具有优势。在大规模问题上，它的扩展性应优于大多数MCMC算法，包括动态HMC（即NUTS），但代价是后验估计的偏差更大。有关该算法的详细信息，请参阅arxiv预印本。

该算法在 实现 于 BlackJAX，这是一个用于 JAX 的推理算法库。通过 PyMC 的 JAX 后端（通过 pytensor），我们可以使用一些简单的包装代码在任何 PyMC 模型上运行 BlackJAX 的路径查找器。

此包装代码在 pymc-experimental 中实现。本教程展示了如何在您的 PyMC 模型上运行 Pathfinder。

您首先需要安装 pymc-experimental：

pip install git+https://github.com/pymc-devs/pymc-experimental

安装其他包的说明：

import arviz as az
import matplotlib.pyplot as plt
import numpy as np
import pymc as pm
import pymc_experimental as pmx

print(f"Running on PyMC v{pm.__version__}")

Running on PyMC v5.15.1+17.g508a1341f

首先，定义你的 PyMC 模型。这里，我们使用 8-schools 模型。

# Data of the Eight Schools Model
J = 8
y = np.array([28.0, 8.0, -3.0, 7.0, -1.0, 1.0, 18.0, 12.0])
sigma = np.array([15.0, 10.0, 16.0, 11.0, 9.0, 11.0, 10.0, 18.0])

with pm.Model() as model:
    mu = pm.Normal("mu", mu=0.0, sigma=10.0)
    tau = pm.HalfCauchy("tau", 5.0)

    z = pm.Normal("z", mu=0, sigma=1, shape=J)
    theta = mu + tau * z
    obs = pm.Normal("obs", mu=theta, sigma=sigma, shape=J, observed=y)

接下来，我们调用 pmx.fit() 并传入我们希望它使用的算法。

with model:
    idata = pmx.fit(method="pathfinder", num_samples=1000)

Running pathfinder...

Transforming variables...

就像 pymc.sample() 一样，这会返回一个包含后验样本的 idata。请注意，由于这些样本并非来自 MCMC 链，因此无法以常规方式评估收敛性。

az.plot_trace(idata)
plt.tight_layout();
../_images/8ebdbd420e8fa5b19a46c2ec0a3778ddc55fa6dbe577bd6bf2ed0a94de6772cd.png

参考资料#

[1]

张璐, 鲍勃·卡彭特, 安德鲁·盖尔曼, 和 阿基·维塔里. Pathfinder: 并行拟牛顿变分推断. arXiv 预印本 arXiv:2108.03782, 2021.

作者#

  • 由Thomas Wiecki于2022年10月11日撰写（pymc-examples#429）

  • 由Reshama Shaikh于2023年2月5日重新执行的笔记本

  • Chris Fonnesbeck 于 2024 年 7 月 17 日修复了错误

水印#

%load_ext watermark
%watermark -n -u -v -iv -w -p xarray

Last updated: Wed Jul 17 2024

Python implementation: CPython
Python version       : 3.11.9
IPython version      : 8.25.0

xarray: 2024.6.0

numpy            : 1.26.4
matplotlib       : 3.8.4
arviz            : 0.18.0
pymc_experimental: 0.1.1
pymc             : 5.15.1+17.g508a1341f

Watermark: 2.4.3

许可证声明#

本示例库中的所有笔记本均在MIT许可证下提供，该许可证允许修改和重新分发，前提是保留版权和许可证声明。

引用 PyMC 示例#

要引用此笔记本，请使用Zenodo为pymc-examples仓库提供的DOI。

重要

许多笔记本是从其他来源改编的：博客、书籍……在这种情况下，您应该引用原始来源。

同时记得引用代码中使用的相关库。

这是一个BibTeX的引用模板：

@incollection{citekey,
  author    = "<notebook authors, see above>",
  title     = "<notebook title>",
  editor    = "PyMC Team",
  booktitle = "PyMC examples",
  doi       = "10.5281/zenodo.5654871"
}

渲染后可能看起来像：