谷歌大脑提出了一种简单的方法，用于将概率编程嵌入到深度学习生态系统

谷歌大脑最新提出了一种简单的方法，用于将概率编程嵌入到深度学习生态系统。这种简单分布式、加速的概率编程，可以将概率程序扩展到512个TPUv2、1亿+参数的模型。

将概率程序扩展到512个TPU、1亿+参数的模型是什么概念？

谷歌大脑近日公开一篇论文“Simple, Distributed, and Accelerated Probabilistic Programming”，发表于NIPS 2018。论文描述了一种简单、低级的方法，用于将概率编程嵌入到深度学习生态系统中。

该研究将概率编程提取为一种单一的抽象——随机变量（random variable）。

在继续介绍论文之前，让我们先了解一下Edward，因为该研究是基于Edward2实现的。

Edward是哥伦比亚大学、谷歌大脑等在2017年提出的新深度概率编程语言，也是一个用于概率建模、推理和评估的Python 库。Edward 融合了以下三个领域：贝叶斯统计学和机器学习、深度学习、概率编程。

研究人员在TensorFlow的轻量级实现证明该方法可支持多种应用：使用TPUv2的模型并行变分自动编码器（VAE）；使用TPUv2的数据并行自回归模型Image Transformer；以及多GPU No-U-Turn Sampler (NUTS)。

对于64x64 ImageNet上最先进的VAE和256x256 CelebA-HQ上最先进的Image Transformer，该方法实现了从1 TPU到256 TPU的最佳线性加速。对于NUTS，相对Stan的GPU加速达到100倍，相对PyMC3的加速达到37倍。

从可微编程到概率编程

深度学习的许多进展可以被解释为模糊了模型和计算之间的界限。有人甚至提出一种新的“可微编程”（differentiable programming）范式，在这种范式中，其目标不仅仅是训练一个模型，而是执行一般的程序合成。

在这种观点里，注意力（attention）和门控（gating）描述了布尔逻辑；跳过连接（skip connections）和条件计算描述了控制流；外部存储访问函数内部范围之外的元素。学习算法也越来越动态：例如， learning to learn，神经结构搜索，以及层内优化等。

可微编程范式鼓励人们考虑计算成本：不仅要考虑模型的统计特性，还必须考虑其计算、内存和带宽成本。这种理念使研究人员设计出深度学习系统，这些系统运行在最前沿的现代硬件上。

相比之下，概率编程社区倾向于在模型和计算之间划清界限：首先，将概率模型指定为程序；其次，执行“推理查询”来自动训练给定数据的模型。

这种设计使得很难真正大规模地实现概率模型，因为训练具有数十亿参数的模型需要跨加速器地拆分模型计算和调度通信。

在这篇论文中，我们描述了一种在深度学习生态系统中嵌入概率编程的简单方法; 我们的实现基于TensorFlow和Python，名为Edward2。这种轻量级方法为灵活的建模提供了一种 low-level 的模式——深度学习者可以从使用概率原语进行灵活的原型设计中获益，并且概率建模者可以从与数字生态系统更紧密的集成中获益。

研究贡献：

我们将概率编程的核心提炼为单个抽象——随机变量（random variable）。

这种low-level的设计有两个重要含义：

首先，它使研究具有灵活性：研究人员可以自由地操作模型计算，以进行训练和测试。

其次，它可以使用加速器（例如TPU）来实现更大的模型：TPU需要专门的ops，以便在物理网络拓扑中分配计算和内存。

我们举例说明了三种应用：使用TPUv2的模型并行变分自动编码器（VAE）；使用TPUv2的数据并行自回归模型（Image Transformer）；以及多GPU No-U-Turn Sampler (NUTS)。

对于64x64 ImageNet上最先进的VAE和256x256 CelebA-HQ上最先进的Image Transformer，我们的方法实现了从1 TPUv2到256 TPUv2的最佳线性加速。对于NUTS，GPU的加速比Stan快100倍，比PyMC3快37倍。

只需要随机变量

Random Variables Are All You Need！

在这一节，我们概述了Edward2的概率程序。它们只需要一个抽象：一个随机变量。然后，我们将描述如何使用跟踪（tracing）来执行灵活的、low-level 的操作。

概率程序、变分程序

Edward2将任何可计算的概率分布具体化为一个Python函数(程序)。通常，该函数执行生成过程，并返回示例。程序的输入—以及任何限定范围的Python变量—表示分布条件的值。

要在程序中指定随机选择，我们使用了Edward的RandomVariables。RandomVariables提供了log_prob和sample等方法，包装TensorFlow Distributions。此外，Edward随机变量增加了TensorFlow操作的计算图：每个随机变量x与图中的张量x∗∼p(x)相关联。

图1描述了一个示例：一个Beta-Bernoulli模型

图1：Beta-Bernoulli program

重要的是，所有的分布——不管下游用什么——都是作为概率程序编写的。

图2描述了一个隐式变分程序，即允许采样但可能不具有易于处理的密度的变分分布。

图2：Variational program

一般而言，变分程序、proposal programs和对抗式训练中的discriminators都是可计算的概率分布。如果我们有一个操纵这些概率程序的机制，那么就不需要引入任何额外的抽象来支持强大的推理范例。

下面我们将使用一个model-parallel VAE来展示这种灵活性。

示例：模型并行的变分自动编码器

图4实现了一个模型并行的变分自动编码器(model-parallel VAE)，它由decoder、prior和encoder组成。decoder生成16位音频；它采用一种自回归的flow，用于训练有效地在序列长度上并行化。encoder将每个样本压缩成粗分辨率，由一个压缩函数参数化。

图4：Model-parallel VAE with TPUs, generating 16-bit audio from 8-bit latents

TPU集群在环形网络中布置核心，例如，512个核心可以布置为16x16x2的环面互连。为了利用集群， prior和decoder都应用分布式自回归流(如图3所示)。

图3：Distributed autoregressive flows

概率程序很简洁。它们利用了最近的进展，如autoregressive flows和multi-scale latent variables，并且实现了以前从未尝试过的架构，其中使用16x16 TPUv2芯片(512核心)，模型可以在4.1TB内存中分割，并使用最多1016个FLOPS。VAE的所有元素——分布式、架构和计算位置——都是可扩展的。

跟踪

我们将概率程序定义为任意Python函数。为了实现灵活的训练，我们应用了跟踪（tracing），这是概率编程的经典技术，以及自动微分（automatic differentiation）。

图5显示了核心实现：10行的代码。

图5：tracing的最小实现

图6：程序执行。

其他示例，包括数据并行Image Transformer，No-U-Turn Sampler，概率程序对齐，通过梯度下降的变分推理学习等，请阅读原始论文。

结论

我们描述了一种简单、低级别的方法，用于在深度学习生态系统中嵌入概率编程。对于64x64 ImageNet上的最先进的VAE和256x256 CelebA-HQ上的Image Transformer，我们实现了从1到256 TPUv2芯片的最佳线性加速。对于NUTS，相比其他系统速度提升100倍。

目前，我们正在推进这种设计，作为生成模型和贝叶斯神经网络基础研究的一个阶段。此外，我们的实验依赖于数据并行性以得到大幅的加速加速。最近的一些研究改进了神经网络的分布式编程，用于模型并行性以及对大规模输入(如超高分辨率图像)的并行性。结合这项工作，我们希望突破超过1万亿参数和超过4K分辨率的巨型概率模型的极限。