NIPS 2018:谷歌大脑提出简单、分布式概率编程,可用TPU大规模训练

新智元报道

来源:arXiv

作者:Google Brain 编译:肖琴

【新智元导读】谷歌大脑最新提出了一种简单的方法,用于将概率编程嵌入到深度学习生态系统。这种简单分布式、加速的概率编程,可以将概率程序扩展到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分辨率的巨型概率模型的极限。

论文:

https://arxiv.org/pdf/1811.02091.pdf

本文分享自微信公众号 - 新智元(AI_era)

原文出处及转载信息见文内详细说明,如有侵权,请联系 yunjia_community@tencent.com 删除。

原始发表时间:2018-11-13

本文参与腾讯云自媒体分享计划,欢迎正在阅读的你也加入,一起分享。

发表于

我来说两句

0 条评论
登录 后参与评论

相关文章

来自专栏机器之心

Diss所有深度生成模型,DeepMind说它们真的不知道到底不知道什么

深度学习在应用层面获得了巨大成功,这些实际应用一般都希望利用判别模型构建条件分布 p(y|x),其中 y 是标签、x 是特征。但这些判别模型无法处理从其他分布中...

11510
来自专栏大数据文摘

迷人又诡异的辛普森悖论:同一个数据集是如何证明两个完全相反的观点的?

在辛普森悖论中,餐馆可以同时比竞争对手更好或更差,锻炼可以降低和增加疾病的风险,同样的数据集能够用于证明两个完全相反的论点。

15030
来自专栏GreenLeaves

TFS2018环境搭建一硬件要求

TFS可以安装在Windows Server和Windows PC操作系统中,但是TFS2018和2018只支持64位操作系统中,早期的版本没有操作系统的位数限...

40730
来自专栏smy

一张图解释负载均衡

首先当大量用户访问时候,先请求到nignx服务器,因为nignx对于高并发支持较好,所以由nignx服务器将访问需求分配给不同的apache服务器,apache...

20430
来自专栏Python专栏

200行代码,一行行教你自制微信机器人

1) 用一个windows客户端工具运营公众号,真的很局限。虽然工具的功能很强大,能自动添加好友,自动拉好友入群,关键字回复等等,但是有一个绕不开的点,它是一款...

61120
来自专栏数据结构笔记

python基础类型(一):字符串和列表

注意到最后三个的单双引号是嵌套使用的,但是最后一个的使用方法是错误的,因为当我们混合使用两种引号时必须有一种用来划分字符串的边界,即在两边的引号不能出现在字符串...

13720
来自专栏编程坑太多

『高级篇』docker之Mesos集群架构图(23)

11640
来自专栏chenssy

多线程:为什么在while循环中加入System.out.println,线程可以停止

这个我们都知道,由于 stopReqested 的更新值在主内存中,而线程栈中的值不是最新的,所以会一直循环,线程并不能停止。加上 Volatile 关键字后,...

21240
来自专栏苦逼的码农

一些常用的算法技巧总结

数组的下标是一个隐含的很有用的数组,特别是在统计一些数字,或者判断一些整型数是否出现过的时候。例如,给你一串字母,让你判断这些字母出现的次数时,我们就可以把这些...

21030
来自专栏我是攻城师

理解BitMap算法的原理

位图:一种常用的数据结构,代表了有限域中的稠集(dense set),每一个元素至少出现一次,没有其他的数据和元素相关联。在索引,数据压缩,海量数据处理等方面有...

24330

扫码关注云+社区

领取腾讯云代金券

年度创作总结 领取年终奖励