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

选自 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 或寻找条件独立性。

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

本文分享自微信公众号 - 机器之心(almosthuman2014)

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

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

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

发表于

我来说两句

0 条评论
登录 后参与评论

相关文章

来自专栏机器之心

专栏 | 阿里IJCAI 2017 Workshop论文:使用深度强化学习方法求解一类新型三维装箱问题

机器之心专栏 阿里菜鸟物流人工智能部 据机器之心了解,阿里巴巴有 11 篇论文入选如今正在墨尔本进行的 IJCAI 2017 大会,其中 6 篇来自阿里巴巴-浙...

1.2K60
来自专栏新智元

DeepMind丨深度学习最新生成记忆模型,远超RNN的GTMM

【新智元导读】DeepMind 的最新研究成果,对广泛使用于语音识别、图像识别、语义理解等领域的深度学习人工网络RNN性能带来显著提升(substantiall...

42760
来自专栏PPV课数据科学社区

【学习】数据可视化基础学习摘要

数据可视化参考流程 科学可视化的早期可视化流水线,描述了从数据空间到可视空间的映射,包含串行处理数据的各个阶段: 原始数据->数据分析->预处理数据->过滤->...

37760
来自专栏iOSDevLog

十个主题,最全的优秀 TensorFlow 相关资源列表

480110
来自专栏算法channel

铁粉巨献:某独角兽公司数据挖掘工程师岗位 2000字面试总结

感谢铁粉:地球球长,奉献出自己的面试经验,面经总结的非常认真、详细、用心。这是球长同学第二次无私地奉献他的面经了,上一篇面经 ( @all: 新浪 机器学习算...

12300
来自专栏量子位

干货分享 | 详解特征工程与推荐系统及其实践(附PPT)

云脑科技机器学习训练营11月13日正式开始啦!量子位作为合作媒体独家为大家分享课程干货内容。 本期内容简介 主题:亿级用户电商平台推荐系统挑战 主讲人:张本宇(...

77540
来自专栏星回的实验室

特征选择的评估方法

数据挖掘中,特征选择的过程就是计算特征与样本观测结果的相关性。卡方检验和互信息是用得较多的计算方法。

22210
来自专栏PPV课数据科学社区

用 TensorFlow 做个聊天机器人

本文结构: 聊天机器人的架构简图 用 TensorFlow 实现 Chatbot 的模型 如何准备 chatbot 的训练数据 Chatbot 源码解读 1. ...

35930
来自专栏ATYUN订阅号

机器学习中的问题解决方案:解析解vs数值解

这篇文章会让你了解为什么没人能告诉你要使用什么算法,或如何为特定的数据集配置算法。另外,找到好的数据、算法、配置实际上是应用机器学习的难点,也是你需要集中解决的...

15750
来自专栏机器学习AI算法工程

用TensorFlow实现文本分析模型,做个聊天机器人

不错的学习聊天机器人的资源,不知道小伙伴们有没有去学习呢。 自己动手做聊天机器人教程 http://www.shareditor.com/blogshow...

51250

扫码关注云+社区

领取腾讯云代金券