[论文解读] Simple, Distributed, and Accelerated Probabilistic Programming

Dustin Tran, Matthew W. Hoffman+4 more|arXiv (Cornell University)|Jan 1, 2018

Machine Learning and Data Classification被引用 30

一句话总结

本文提出了一种轻量级、低层次的概率编程框架，嵌入于TensorFlow中，以单一抽象——随机变量为核心，实现了高效的模型并行与数据并行推理。该框架在256块TPUv2芯片上对VAE和Image Transformers实现了最优的线性加速，并在多GPU NUTS采样中实现了比Stan快100倍的性能提升。

ABSTRACT

We describe a simple, low-level approach for embedding probabilistic programming in a deep learning ecosystem. In particular, we distill probabilistic programming down to a single abstraction—the random variable. Our lightweight implementation in TensorFlow enables numerous applications: a model-parallel variational auto-encoder (VAE) with 2nd-generation tensor processing units (TPUv2s); a data-parallel autoregressive model (Image Transformer) with TPUv2s; and multi-GPU No-U-Turn Sampler (NUTS). For both a state-of-the-art VAE on 64x64 ImageNet and Image Transformer on 256x256 CelebA-HQ, our approach achieves an optimal linear speedup from 1 to 256 TPUv2 chips. With NUTS, we see a 100x speedup on GPUs over Stan and 37x over PyMC3.

研究动机与目标

• 通过将概率编程简化为单一抽象——随机变量，来实现简化。
• 在深度学习生态系统中实现高效可扩展的推理，且工程开销最小。
• 在数千块加速器上，对模型并行和数据并行的概率模型均实现线性加速。
• 通过多GPU无截断采样器（NUTS）加速马尔可夫链蒙特卡洛采样，超越现有工具的性能。

提出的方法

• 该框架仅使用随机变量作为核心抽象来建模概率程序，隐藏了复杂分布和推理机制的细节。
• 与TensorFlow原生集成，利用其自动微分和分布式执行能力实现高效计算。
• 在TPUv2硬件上，使用该抽象实现了模型并行的VAE和数据并行的自回归模型（如Image Transformer）。
• 通过参数分片和设备间梯度同步，实现了多GPU NUTS推理。
• 系统支持变分推断和MCMC采样，且仅需极少代码修改。
• 通过在1至256块TPUv2芯片上实现高效的数据并行和模型并行，实现了线性扩展。

实验结果

研究问题

• RQ1 是否可以在不牺牲表达能力或性能的前提下，将概率编程简化为单一抽象？
• RQ2 是否最小化、低层次的接口能够在模型并行和数据并行设置下，实现数千块加速器上的最优扩展？
• RQ3 此类系统是否能在MCMC采样中实现优于现有概率编程工具的性能？
• RQ4 该框架在分布式执行下，对SOTA模型（如VAE和Transformer）的支持程度如何？

主要发现

• 该框架在64×64 ImageNet上的SOTA VAE上，从1到256块TPUv2芯片实现了最优线性加速。
• 在256×256 CelebA-HQ上的Image Transformer也实现了最优线性加速，证明了其在数据并行设置下的可扩展性。
• 多GPU NUTS采样相比Stan实现了100倍加速，相比PyMC3实现了37倍加速。
• 该系统通过相同的低层次抽象，实现了对变分自编码器和自回归模型的高效可扩展推理。
• 该方法在多种模型架构和硬件配置（包括TPU和多GPU设置）下均保持了高性能。
• 随机变量抽象的简洁性使其具有广泛适用性，且不牺牲性能或可扩展性。