微信扫一扫
分享到朋友圈

NIPS 2018 | Edward2.2,一种可以用TPU大规模训练的概率编程

作者:机器之心 来源:机器之心 公众号
分享到:

11-12

选自 arXiv

作者:Dustin Tran等

机器之心编译


本文描述了一种用于在深度学习生态系统中嵌入概率编程的简单、底层方法。具体来说,作者将概率编程归结为一个抽象:随机变量。


深度学习的很多研究结果都模糊了模型和计算之间的界限,有的甚至表明是一种「可微分编程」的新范式,它们的目标不仅仅是训练模型,同时还希望实现一般的程序综合体。在这一观点下,注意力机制和门控机制可以描述布尔逻辑运算符,残差连接和条件计算可以描述控制流,外部记忆可以访问函数内部作用范围外的元素。此外,学习算法也将变得越来越动态,例如学习如何学习、神经架构搜索和层级内的最优化等。


可微分编程范式鼓励建模者明确考虑计算成本,我们不仅需要考虑模型的统计属性,例如模型到底能不能拟合真实数据分布,拟合得又有多好。同时我们还需要考虑计算量、内存占用和传输带宽的成本等,这里主要关注训练和预测的效率究竟怎样。


相比之下,概率编程社区主要倾向于在模型和计算之间划清界限:首先我们可以指定一个概率模型作为程序,其次我们可以执行「inference query」以在给定数据的情况下自动训练模型。这种设计使得概率编程很难在大规模真实环境下实现概率模型,其中大规模可能需要训练有数十亿参数量的模型,并要求通过多个加速器和调度通信切分模型计算量。


Edward 等最近的研究能更好地控制深度学习中的推断过程,然而它们都将推断过程作为一个封闭系统,因此很难构成任意的计算或更广泛的机器学习生态系统。


在这篇论文中,作者描述了一种在深度学习生态系统中嵌入概率编程的简单方法,且他们在 TensorFlow 和 Python 实现了新的库 Edward2.2。这种轻量方法为灵活建模提供了底层模块,深度学习研究者可以使用概率基元灵活地设计原型。本论文主要的贡献有三点:


  • 作者将概率编程的核心提炼为单个抽象:随机变量。

  • 这种底层设计有两个重要意义:首先它允许研究具有足够的灵活性;其次它允许使用张量计算单元(TPU)等加速器实现更大的模型。

  • 作者后面还举例了三种应用:通过 TPU 训练的模型并行化变分自编码器;由 TPU 训练的数据并行化自回归模型;和多 GPU No-U-Turn Sampler (NUTS)。本文并没有介绍这三种应用,感兴趣的读者可查看原论文。


论文:Simple, Distributed, and Accelerated Probabilistic Programming

 


论文地址:https://arxiv.org/abs/1811.02091


摘要:我们描述了一种用于在深度学习生态系统中嵌入概率编程的简单、底层方法。具体来说,我们将概率编程归结为一个抽象:随机变量。我们在 TensorFlow 中的轻量级实现支持许多应用:通过 TPU(TPUv2)训练的模型并行变分自编码器(VAE);通过 TPUv2 训练的数据并行自回归模型(Image Transformer);多 GPU 且无 U-Turn 的取样器(NUTS)。对于 64x 64 ImageNet 上的最新 VAE 和 256 x256 CelebA-HQ 上的 Image Transformer,我们的方法实现了从 1 到 256 块 TPUv2 芯片的线性加速。通过 NUTS,我们发现在 GPU 我们的方法比 Stan 快 100 倍,比 PyMC3 快 37 倍。



随机变量


我们在 Edward2 中概述了概率编程。它们只需要一个抽象:随机变量。然后,我们描述了如何利用追踪来执行灵活的底层操作。Edward 2 将任何可计算的概率分布具化为 Python 函数(程序),通常该函数执行生成过程并返回示例。


图 1:Beta-Bernoulli 程序。在即时运行模式中,model() 生成一个包含 50 个因素的二进制向量,model() 返回一个 op,该 op 在 TensorFlow 会话中进行评估。


图 2:变分程序 [35],可在即时运行模式下获得。Python 控制流可用于生成流程:给定硬币翻转,该程序可根据两个神经网络中的一个 生成流程。其输出具有不同的形状和结构。


图 3:分布式自回归流。(右)默认长度是 8,每个有 4 个独立的流。每个流通过关于自回归排序的层来转换输入。(左)流跨越 4x4 内核(矩形)的虚拟拓扑进行划分;每个内核计算两个流,并在本地连接;最后会汇集一个内核。虚拟拓扑与物理 TPU 拓扑对齐:对于 4x4 TPU,它很精确;对于 16x16 TPU,它要为了数据并行性而复制。


图 4:通过 TPU 训练模型并行的 VAE,从 8 位延迟生成 16 位 latents。先验和解码器根据分布式自回归流来分割计算。编码器可能会根据压缩器分割计算;为了节省空间我们选择忽略它。


图 6:一个程序执行。它是有向无环图,并追踪累计对数概率等各种 op 或寻找条件独立性。



本文为机器之心编译,转载请联系本公众号获得授权

✄------------------------------------------------

加入机器之心(全职记者 / 实习生):hr@jiqizhixin.com

投稿或寻求报道:content@jiqizhixin.com

广告 & 商务合作:bd@jiqizhixin.com

阅读8572
举报0
关注机器之心微信号:almosthuman2014

用微信扫描二维码即可关注
声明

1、头条易读遵循行业规范,任何转载的稿件都会明确标注作者和来源;
2、本文内容来自“机器之心”微信公众号,文章版权归机器之心公众号所有。

评论
更多

文章来自于公众号:

机器之心

微信号:almosthuman2014

邮箱qunxueyuan#163.com(将#换成@)
微信编辑器
免责声明
www.weixinyidu.com   免责声明
版权声明:本站收录微信公众号和微信文章内容全部来自于网络,仅供个人学习、研究或者欣赏使用。版权归原作者所有。禁止一切商业用途。其中内容并不代表本站赞同其观点和对其真实性负责,也不构成任何其他建议。如果您发现头条易读网站上有侵犯您的知识产权的内容,请与我们联系,我们会及时修改或删除。
本站声明:本站与腾讯微信、微信公众平台无任何关联,非腾讯微信官方网站。