您的位置:易优娱乐 > 编程语言 >
编程语言

NIPS 2018 Edward22一种可以用TPU大规模训练的概率编

2018-11-16作者:admin来源:未知次阅读

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

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

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

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

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

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

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

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

  摘要:我们描述了一种用于在深度学习生态系统中嵌入概率编程的简单、底层方法。具体来说,我们将概率编程归结为一个抽象:随机变量。我们在 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 函数(程序),通常该函数执行生成过程并返回示例。

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

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

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

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

NIPS 2018 Edward22一种可以用TPU大规模训练的概率编 相关的内容:

关于 NIPS 2018 Edward22一种可以用TPU大规模训练的概率编 的评论