用70行Python编写一个概率编程语言

1、简介

在这篇文章中，我将介绍概率编程语言（Probabilistic Programming Languages，简称PPL）的工作原理，并逐步演示如何用Python构建一个简单的概率编程语言。

本文主要面向的读者是统计学家、AI研究员和好奇的程序员，相信大家都熟悉 PPL 和贝叶斯统计，并掌握了基本的 Python知识。

我们将要构建的API如下：

mu =LatentVariable("mu",Normal, [0.0, 5.0])
y_bar =ObservedVariable("y_bar", Normal, [mu, 1.0],observed=3.0)

evaluate_log_density(y_bar,{"mu": 4.0})

前两行定义了统计模型：

最后一行求在条件下，在 μ = 4 处该模型定义的（未正规化的）概率分布。

希望本文能让读者理解 PPL 的工作原理，并了解如何用 Python 实现一门嵌入式领域专用语言（Embedded Domain-Specific Languages，简称EDSL）。

2、相关研究

据我所知，目前尚没有使用Python的PPL实现。

《TheDesign and Implementation of Probabilistic Programming Languages》一书的重点是编程语言理论，需要读者熟悉 continuation-passing style、协程，而且采用了JavaScript 作为实现语言。
《Anatomyof a Probabilistic Programming Framework》一文（https://www.georgeho.org/prob-prog-frameworks/）是很不错的高层概要，但并没有涉及具体实现细节和代码示例。
Junpeng Lao的演讲（https://www.youtube.com/watch?v=WHoS1ETYFrw&feature=youtu.be）和 PyMC3的开发者指南（https://docs.pymc.io/en/v3/developer_guide.html）描述了 PyMC 的具体实现细节，但想根据这些内容实现一个 PPL 并不容易。

3、实现

高层表示

我们使用下述模型作为基本的参考：

这两个表达式定义了一个联合概率分布，对应的概率密度函数（Probability Density Function，简称PDF）为：

可以用两种形式的图来表示这个表达式（和模型）：一种是图模型，一种是有向因子图。

左：用概率图模型（PGM）表示的模型。右：用有向因子图表示的模型。

虽然在各类文档中PGM方式更常见，但我认为LFG对于实现PPL的人来说更有用。这个图揭示了几个重要方面：

我们需要一种方式来表达两种变量：
一种是观测到的值（

，灰色背景）
一种是无法观测到的隐变量（μ，白色背景）
我们需要处理常数，以及每个变量的分布。
最后，我们需要一种方法将观测值、隐变量和常数连接起来。

分布

在文本中，分布指的是一个类，它带有一个函数，可以求在某个点上的对数概率密度。log_density 函数接受一个 float 类型的参数，表示分布支持的一个点；一个 List[float]，表示分布的参数；返回一个 float，是该分布在该点上的对数PDF。要实现新的分布，只需从 Distribution 虚类集成即可。暂时不支持向量或矩阵分布。

我们使用 SciPy 实现一个 Normal 分布，param[0] 是均值，param[1] 是标准差。

from scipy.stats import norm

classDistribution:
    @staticmethod
    def log_density(point, params):
       raise NotImplementedError("Mustbe implemented by a subclass")

classNormal(Distribution):
    @staticmethod
    def log_density(point, params):
        returnfloat(norm.logpdf(point, params[0], params[1]))

变量和DAG

现在我们来看看变量。这些变量包含三个方面：它们有关联的分布，可以为隐变量或观测变量，它们之间互相连接（即变量可以有子变量）。

dist_class字段就是变量关联的分布 Distribution。在需要求完整的对数密度时，我们会使用该字段访问变量的分布的log_density方法。

区分隐变量和观测变量的方式是使用 LatentVariable 和 ObservedVariable 类。观测变量有一个observed 字段，其中包含了观测的值。由于隐变量在建模阶段没有值，我们必须在运行时给其复制，才能求出全部的对数分布。要在运行时给隐变量复制，我们需要用唯一的字符串 name 来区分它们。

最后，我们可以将一个变量的分布的参数设置为变量或常量。在本例中，的平均值是 μ，一个正太随机变量，而它的标准差是常数 1。为了表示它，我们使用 dist_args 属性。mypy 中的 dist_args 的签名为 dist_args: Union[float, LatentVariable, ObservedVariable]。这就是说，一个隐变量或观测变量可以有“参数”，参数本身也可以是常量的隐变量或观测变量，这就创建了一个有向无环图（DAG）。

classLatentVariable:
    def __init__(self, name, dist_class, dist_args):
       self.name = name
       self.dist_class = dist_class
       self.dist_args = dist_args


classObservedVariable:
    def __init__(self, name, dist_class, dist_args, observed):
       self.name = name
       self.dist_class = dist_class
       self.dist_args = dist_args
        self.observed = observed

可视化该DAG，就会注意到它和隐因子图之间的一个重要区别：箭头方向是相反的。其原因是我们的建模API中的变量指定方式。而且，似乎将观测变量作为根节点，能更好地表示计算联合对数密度的过程。

左：将模型表示为有向因子图。右：内存中的DAG表示。

为了进一步说明，我们来看看模型中的 dist_args 是什么样子：

mu =LatentVariable("mu",Normal, [0.0, 5.0])
y_bar =ObservedVariable("y_bar", Normal, [mu, 1.0],observed=5.0)

print(mu)
#=> <__main__.LatentVariable object at 0x7f14f96719a0>
print(mu.dist_args)
#=> [0.0, 5.0]
print(y_bar)
#=> <__main__.ObservedVariable object at 0x7f14f9671940>
print(y_bar.dist_args)
# => [<__main__.LatentVariable object at 0x7f14f96719a0>,1.0]

求出对数密度

我们距离目标不远了，现在还缺一种通过DAG计算联合对数密度的方法。我们需要遍历DAG，将每个变量的对数密度加在一起。对数密度加法相当于密度的乘法，但加法在数值上更稳定。

遍历DAG需要使用一个递归算法，叫做深度优先搜索。collect_variables函数会访问所有变量一次，将所有非float变量收集到一个列表中。然后算法从根节点开始，递归地遍历收集到的每个变量的所有dist_args。

defevaluate_log_density(variable, latent_values):
   visited = set()
   variables = []

    def collect_variables(variable):
       if isinstance(variable, float):
           return

       visited.add(variable)
       variables.append(variable)

       for arg in variable.dist_args:
           if arg not in visited:
                collect_variables(arg)

    collect_variables(variable)

对于每个变量，我们需要获取它的每个参数的数值，然后用它求出分布的对数密度。由于float参数已经是数值了，而LatentVariables根据求值所在的位置不同而有不同的值。为了指定隐变量的值，我们需要传递一个包含了从变量名到数值的映射的字典，名为latent_values。注意ObservedVariables不能是参数，只能是根节点。

注意：

dist_args可以是float常数，或LatentVariables。
dist_params都是float，或者是常数，或者是通过latent_values在运行时（即计算对数密度时）赋给隐变量的值。

最后，我们有了从参数中提取出的分布参数，就可以更新整体的对数密度了。LatentVariable需要在latent_values指定的点上计算对数密度，而ObservedValues会在observed指定的点上计算对数密度。

  log_density = 0.0
    for variable in variables:
       dist_params = []
       for dist_arg in variable.dist_args:
           if isinstance(dist_arg, float):
                dist_params.append(dist_arg)
           if isinstance(dist_arg, LatentVariable):
               dist_params.append(latent_values[dist_arg.name])

       if isinstance(variable, LatentVariable):
           log_density +=variable.dist_class.log_density(
                latent_values[variable.name],dist_params
           )
       if isinstance(variable, ObservedVariable):
           log_density +=variable.dist_class.log_density(
                variable.observed, dist_params
           )

    returnlog_density

我们来检查一下整体的对数概率是否与期待的结果一致：

mu =LatentVariable("mu",Normal, [0.0, 5.0])
y_bar =ObservedVariable("y_bar", Normal, [mu, 1.0],observed=5.0)

latent_values = {"mu": 4.0}
print(evaluate_log_density(y_bar, latent_values))
#=> -4.267314978843446
print(norm.logpdf(latent_values["mu"], 0.0, 5.0)
    + norm.logpdf(5.0, latent_values["mu"], 1.0))
# => -4.267314978843446

4、结论和改进

分布、变量DAG和对数密度计算是概率编程语言的组成部分。变量可以是隐变量或观测变量，在计算对数密度时，每个变量必须单独处理。我们用Python实现了这些概念，从而实现了简单但强大的PPL。

下一步是增加对张量和随机变量变换的支持，从而支持更有用的模型，如线性回归、层次/混合效果模型等。另一个有用的特性是构建一个API用于先验预测采样。最后，我们可以不在Python中进行计算，而是使用像Theano/Aesara、Jax或TensorFlow之类的计算图框架，可以极大地改善性能。

计算图还可以通过反式自动微分的方式，利用对数密度计算梯度，这可以用于汉密尔顿·蒙特卡洛等更高级的采样器上。

5、篇外：后验网格近似

下面，我们讨论一下对数密度的用途。其用途之一就是找到后验分布的模，即找到参数的最有可能的值。

本例中，观测样本的平均值是1.5，在正态零均值先验的条件下，它会向 0 移动一点。这就是说，最大后验估计（Maximum A Posteriori estimate）大约是1.4。

import numpy as npimport pandas as pdimport altair as altfrom smolppl import Normal, LatentVariable, ObservedVariable,                    evaluate_log_density#Define model#Weakly informative mean priormu =LatentVariable("mu",Normal, [0.0, 5.0])#Observation model. I make some observations y_1, y_2, ..., y_n and compute the#sample mean y_bar. It is given that the sample mean has standard deviation 1.y_bar =ObservedVariable("y_bar", Normal, [mu, 1.0],observed=1.5)#Grid approximation for the posterior#Since the prior has mean 0, and the observations have some uncertainty, I#expect the mode to be a bit smaller than 1.5. Something like 1.4grid =np.linspace(-4, 4, 20)evaluations = [evaluate_log_density(y_bar, {"mu": mu}) for mu in grid]#Plottingdata =pd.DataFrame({"grid":grid, "evaluations": evaluations})chart =alt.Chart(data).mark_line(point=True).encode(    x=alt.X('grid', axis=alt.Axis(title="mu")),    y=alt.Y('evaluations', axis=alt.Axis(title="logdensity"))).interactive().configure_axis(   labelFontSize=16,   titleFontSize=16)chart