Pyro SVI 第一部分：随机变分推断简介

Posted on

支持随机变量推断作为一个通用的推断算法被重点设计在Pyro之中。现在我们就来看一下Pyro中我们如何使用变分推断。

Setup

我们将假设，我们已经用Pyro定义了模型（细节可见Pyro模型简介）Pyro的model是一个接受参数的随机函数model (*args, **kwargs)。model()的不同部分的映射关系如下：

  1. 观察Observation $\Leftrightarrow$ 带obs参数字段的pyro.sample
  2. 隐随机变量 $\Leftrightarrow$ pyro.sample()
  3. 参数 $\Leftrightarrow$ pyro.param

现在让我们定义一些符号。模型中有observation x，隐随机变量z和参数 $\theta$。他们的联合概率密度(joint probability density)如下：

我们假设 $p_{theta}(x,z)$ 由多个随机分布 $p_i$ 组成，并有以下的性质：

  1. 我们可以对每个 $p_i$ 取样
  2. 我们可以逐点计算$p_i$的对数概率密度函数
  3. $p_i$ 对参数 $\theta$ 可求导

模型学习

在这章中，我们判断学习到一个好的模型的根据是最大化对数证据（maximize log evidence），比如，我们想要找到 $\theta$ 使得：

那么其对数证据就是

在通常情况下，这是一个更困难的问题。因为对隐随机变量z积分通常是没法解的。进一步的，即使我们直到如何计算所有 $\theta$ 的对数证据，要计算以 $\theta$ 为变量的函数来最大化对数证据也通常是是一个困难的非凸优化问题。

除了要找到 $\theta_{max}$，我们想要计算隐变量z的后验：

注意到，分母的表达式是（通常不可计算的）证据。变分推断提供了一个找到 $\theta_{max}$ 并计算后验的 $p_{\theta_{max]}}(z|x)$ 估计值的方案。让我们看一下是如何实现的。

Guide

变分推断的基础思想就是引入一个带参数的分布 $q_{\phi}(z)$， $\phi$被称作变分参数。在很多地方，这个分布被叫做变分分布；而在Pyro中我们叫他guide（毕竟名字简短点！）。guide将作为后验的一个近似存在。

就像模型一样，guide被编码成一个随机函数guide()，包含pyro.sample和pyro.param。它不包含观测到的数据，因为guide需要是一个正则的分布。注意到，Pyro强制model()和guide()有同样的签名，比如，他们都应该接受相同的参数。

因为guide是后验$p_{\theta_{max}}(z|x)$的一个近似，所以guide需要对模型中所有隐变量提供合格的联合概率密度。之前，在Pyro中，用基础表达式pyro.sample()申明一个隐随机变量时，第一个参数是这个随机变量的名字。这些名字将会被用于对齐model和guide中的随机变量。举个例子，比如model中有一个随机变量z_1：

1
2
def model():
  pyro.sample("z_1", ...)

那么guide需要有一个对应的sample表达式：
1
2
def guide():
  pyro.sample("z_1", ...)

这两个例子的分布可以是不同的，但是名字一定要是一对一的。

一旦我们申明了一个guide（在之后的章节中，我们会给一些明确的例子），我们就已经准备好去推断了。那么学习的过程就被设定成了一个优化问题，通过在 $\theta - \phi$ 空间中不断迭代，使得guide不断逼近真正后验。要完成这事，我们需要定义一个合适的客观的函数

ELBO

通过一个简单的推导（见参考）我们可以得出需要求什么：证据下界Evidence Lower Bound(ELBO)。ELBO，是关于 $\theta$ 和 $\phi$ 的函数，它的定义是对于guide取样结果的期望：

通过假设，我们可以计算这个期望E中的对数概率。同时因为guide被设定是一个我们可以取样的参数分布，所以我们可以计算ELBO的蒙特卡洛估计（Monte Carlo estimates）。重要的是，ELBO是对于我们能选择所有的 $\theta$ 和 $\phi$ 的对数证据的下界：

所以如果我们用随机梯度来不断最大化ELBO，我们同时也是在推动（期望E中的）对数证据变大。进一步来说，ELBO和对数证据之间的间隔正是由guide和posterior后验之间的KL散度决定的：

KL散度是一种用来表示两个分布相近程度的方法（根据定义，其值是非负的）。所以，对于一个固定的 $\theta$ ，当我们一步一步在 $\phi$ 的空间里增大ELBO时，我们就相当于在减少guide和后验之间的KL散度，就好像是我们在把guide移向后验证，如图。通常情况下，我们同时在 $\theta$ 和 $\phi$ 上进行梯度下降，这样guide和model就可以竞逐，即guide跟随着后验 $log\:p_{\theta}(z|x)$ 移动。或许有些惊人，尽管目标是移动的，对于不同的情况，这个优化问题依旧可以（在一定合适的近似程度上）有解。

总体上来说，变分是简单的：我们所要做的就是定一个guide，然后用计算ELBO的梯度。实际上，对任意一堆model和guide求梯度会导致一些复杂的问题（会在后续教程SVI Part III中讨论）。在本章中，我们就先把这个问题当作已经解决了，现在来看一下Pyro中如何做变分推断。

SVI类

在Pyro中，做变分推断的工具被封装在了SVI类中了。（现在SVI只支持ELBO目标，但在未来Pyro会提供更多的变分目标的支持。）

使用者需要提供：model模型，guide，和optimizer优化器。在上文中我们已经讨论过模型和guide了，在下文中我们也会讨论优化器，现在让我们看一下这三个东西。要创建能对ELBO目标做优化的SVI，使用者需要如下：

1
2
3
import pyro
from pyro.infer import SVI, Trace_ELBO
svi = SVI(model, guide, optimizer, loss=Trace_ELBO())

SVI对象提供了两个函数，step()和evaluate_loss()，他们封装了变分学习和评估的逻辑：

  1. step()接受一个单独的梯度，并返回一个损失的估计值（比如，减去ELBO）。如果提供了参数，给step()的参数也会送到pipe tomodel()和guide()中。
  2. evaluate_loss()不需要提供梯度就可以返回损失的估计。就像是step()，如果提供了参数，这些参数也会被送到guide()和model()中。

对于损失就是ELBO的情况，两个函数也会接受可选的参数num_particles，它决定了计算损失的采样次数（对应evaluate_loss）和 损失和梯度(对应step)

优化器Optimizers

在Pyro中，model和guide允许是任意随机函数，只要他们满足：

  1. guide不包含带有obs参数字段的pyro.sample的表达式
  2. model和guide有相同的调用签名

这带来一些挑战，因为这就意味着model()和guide()不同的执行会因为例如某些只出现过几次的隐随机变量和参数，而有相当不同的表现。事实上，参数是在推断过程中动态创造的。换言之，我们试图其中寻找最优解的 $\theta$ 和 $\phi$ 的空间可能增长并改变。

为了支持这样的情况，Pyro需要在学习的过程中，每一个参数第一次出现时，动态地生成优化器。幸运的是，PyTorch正好有一个轻量的优化库（见torch.optim）让我们可以简单地动态地重新调整用途。

所有的这些都由optim.PyroOptim类控制，它是一个对于PyTorch优化器的轻量封装。PyroOptim接受两个参数：一个constructor给PyTorch优化器optim_constructor和一个特定的优化器参数optim_args。总体上来说，在优化的过程中，无论一个新的参数是否被见到，只要用optim_args给了参数，optim_constructor就会来初始化一个新的优化器。

大多数使用者或许将不会直接用到PyroOptim，相对的他们会直接和optim/__init__.py中定义的别名打交道。让我们看一下是如何发生的。我们有两种方法来申明优化器的实参argument。简单的方法是，optim_args是一个固定的字典（dictionary），申明用来初始化PyTorch优化器的所有形参parameter的实参arguments。

1
2
3
4
from pyro.optim import Adam

adam_params = {"lr": 0.005, "betas": (0.95, 0.999)}
optimizer = Adam(adam_params)

第二种方法允许一定程度上的控制。在此，使用者必须申明一个会在Pyro看到新的形参，然后创建优化器时调用到的callable函数。这个callable必须有如下的签名：

  1. module_name: 如果有包含参数的module的名字
  2. param_name: 参数的名字

这给了使用者，比方说自定义学习率（learning rate）等，不同参数的能力。要看这种程度的控制是如何有用的，可以移步disccusion of baselines。这里看一个简单的如何调用API的例子：

1
2
3
4
5
6
7
8
9
from pyro.optim import Adam

def per_param_callable(module_name, param_name):
    if param_name == 'my_special_parameter':
        return {"lr": 0.010}
    else:
        return {"lr": 0.001}

optimizer = Adam(per_param_callable)

这里简单地告诉了Pyro，参数my_special_parameter要使用0.010的学习率，而其他的参数使用0.001的学习率。

一个例子

我们用一个例子来结尾。你有一枚两面的硬币。你想要看看这枚硬币是否公平，是否正反两面出现的频率相同。你关于硬币公平性的先验信念基于两个观察：

  • 这是个美国铸币局发行的标准的25美分
  • 因为多年的使用，它有点损伤
    那么，当你期望这枚硬币在刚被造出来时是公平的，你也就允许了他的公平性不是1:1，而偏移了。那么，但出现正反比率为11:10的时候，你不会惊讶。相对的，当它出现5:1的比率时，你会被惊讶道——毕竟它没有损伤的那么严重。

为了把这个添加到概率模型中去，我们把正面作为1，反面作为0。我们将硬币的公平性视作实数 $f$ ，$f$满足 $f\in[0.0,1.0]$ 且 $f=0.50$ 对应着这是一枚公平的硬币。我们对 $f$ 的先验信念将是一个Beta分布，$Beta(10,10)$，一个在[0.0,1.0]区间上对称分布的概率，其峰值在 $f=0.5$。

要学习到比我们这种瞎猜而来的先验更准确的公平性，我们需要做实验并收集数据。比方说我们抛硬币10次，然后记录结果。当然实际上我们肯定要做更多次，不过现在先这样。

假设我们将收集到的数据叫做data，对应的模型是：

1
2
3
4
5
6
7
8
9
10
11
12
13
import pyro.distributions as dist

def model(data):
    # define the hyperparameters that control the beta prior
    alpha0 = torch.tensor(10.0)
    beta0 = torch.tensor(10.0)
    # sample f from the beta prior
    f = pyro.sample("latent_fairness", dist.Beta(alpha0, beta0))
    # loop over the observed data
    for i in range(len(data)):
        # observe datapoint i using the bernoulli
        # likelihood Bernoulli(f)
        pyro.sample("obs_{}".format(i), dist.Bernoulli(f), obs=data[i])

这里我们有一个隐随机变量('latent_fairness')，服从 $Beta(10,10)$ 分布。对它条件判断，我们观测到每个数据点都是伯努利的似然。注意到每次观测都有自己唯一的名字。

我们下一步的任务是定义一个对应的guide，一个合适的变分分布对应隐随机变量 $f$。唯一的要求是 $q(f)$ 应该是一个在[0.0,1.0]区间伤的概率分布，因为 $f$ 不应该超出这个范围。一个简单的选择是使用另一个由 $\alpha_q$ 和 $\beta_q$ 控制的beta分布。实际上，这个例子中，这正好是正确的选择，因为Beta分布和Bernoulli分布的共轭关系意味着，后验就是beta分布。在Pyro中，我们写如下代码：

1
2
3
4
5
6
7
8
def guide(data):
    # register the two variational parameters with Pyro.
    alpha_q = pyro.param("alpha_q", torch.tensor(15.0),
                         constraint=constraints.positive)
    beta_q = pyro.param("beta_q", torch.tensor(15.0),
                        constraint=constraints.positive)
    # sample latent_fairness from the distribution Beta(alpha_q, beta_q)
    pyro.sample("latent_fairness", dist.Beta(alpha_q, beta_q))

这里有几件值得注意的事：

  • 我们让随机变量的名字正好与model和guide中的变量对应了
  • model(data)和guide(data)接受同样的参数
  • 变分参数都是torch.tensor。字段requires_grad被pyro.param自动设置成True
  • 我们用constraint=constraints.positive来保证alpha_q和beta_q在优化过程中保持非负。

现在我们可以来做随机变分推断了

1
2
3
4
5
6
7
8
9
10
11
# set up the optimizer
adam_params = {"lr": 0.0005, "betas": (0.90, 0.999)}
optimizer = Adam(adam_params)

# setup the inference algorithm
svi = SVI(model, guide, optimizer, loss=Trace_ELBO())

n_steps = 5000
# do gradient steps
for step in range(n_steps):
    svi.step(data)

注意到，在step()中，我们传输进了数据data，然后这些数据也会被传递给model和guide

唯一一件我们忽略的时就是数据。那么我们来创建一些数据，整合所有代码到一起吧：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
import math
import os
import torch
import torch.distributions.constraints as constraints
import pyro
from pyro.optim import Adam
from pyro.infer import SVI, Trace_ELBO
import pyro.distributions as dist

# this is for running the notebook in our testing framework
smoke_test = ('CI' in os.environ)
n_steps = 2 if smoke_test else 2000

# enable validation (e.g. validate parameters of distributions)
assert pyro.__version__.startswith('1.3.0')
pyro.enable_validation(True)

# clear the param store in case we're in a REPL
pyro.clear_param_store()

# create some data with 6 observed heads and 4 observed tails
data = []
for _ in range(6):
    data.append(torch.tensor(1.0))
for _ in range(4):
    data.append(torch.tensor(0.0))

def model(data):
    # define the hyperparameters that control the beta prior
    alpha0 = torch.tensor(10.0)
    beta0 = torch.tensor(10.0)
    # sample f from the beta prior
    f = pyro.sample("latent_fairness", dist.Beta(alpha0, beta0))
    # loop over the observed data
    for i in range(len(data)):
        # observe datapoint i using the bernoulli likelihood
        pyro.sample("obs_{}".format(i), dist.Bernoulli(f), obs=data[i])

def guide(data):
    # register the two variational parameters with Pyro
    # - both parameters will have initial value 15.0.
    # - because we invoke constraints.positive, the optimizer
    # will take gradients on the unconstrained parameters
    # (which are related to the constrained parameters by a log)
    alpha_q = pyro.param("alpha_q", torch.tensor(15.0),
                         constraint=constraints.positive)
    beta_q = pyro.param("beta_q", torch.tensor(15.0),
                        constraint=constraints.positive)
    # sample latent_fairness from the distribution Beta(alpha_q, beta_q)
    pyro.sample("latent_fairness", dist.Beta(alpha_q, beta_q))

# setup the optimizer
adam_params = {"lr": 0.0005, "betas": (0.90, 0.999)}
optimizer = Adam(adam_params)

# setup the inference algorithm
svi = SVI(model, guide, optimizer, loss=Trace_ELBO())

# do gradient steps
for step in range(n_steps):
    svi.step(data)
    if step % 100 == 0:
        print('.', end='')

# grab the learned variational parameters
alpha_q = pyro.param("alpha_q").item()
beta_q = pyro.param("beta_q").item()

# here we use some facts about the beta distribution
# compute the inferred mean of the coin's fairness
inferred_mean = alpha_q / (alpha_q + beta_q)
# compute inferred standard deviation
factor = beta_q / (alpha_q * (1.0 + alpha_q + beta_q))
inferred_std = inferred_mean * math.sqrt(factor)

print("\nbased on the data and our prior belief, the fairness " +
      "of the coin is %.3f +- %.3f" % (inferred_mean, inferred_std))

输出：

1
based on the data and our prior belief, the fairness of the coin is 0.532 +- 0.090

这个估计将用来和明确的后验均值对比，在此例中是 $16/30 = 0.53$。注意到最后这枚硬币的公平性估计在先验估计的0.5和原始的经验频率 $6/10=0.6$ 之间。

参考

[1] Automated Variational Inference in Probabilistic Programming, David Wingate, Theo Weber

[2] Black Box Variational Inference, Rajesh Ranganath, Sean Gerrish, David M. Blei

[3] Auto-Encoding Variational Bayes, Diederik P Kingma, Max Welling