谷歌大腦最新提出了一種簡單的方法，用于將概率編程嵌入到深度學習生態系統。這種簡單分布式、加速的概率編程，可以將概率程序擴展到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；以及多GPUNo-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分辨率的巨型概率模型的極限。