谷歌大腦最新提出了一種簡單的方法，用于將概率編程嵌入到深度學習生態(tài)系統(tǒng)。這種簡單分布式、加速的概率編程，可以將概率程序擴展到512個TPUv2、1億+參數(shù)的模型。

將概率程序擴展到512個TPU、1億+參數(shù)的模型是什么概念？

谷歌大腦近日公開一篇論文“Simple, Distributed, and Accelerated Probabilistic Programming”，發(fā)表于NIPS 2018。論文描述了一種簡單、低級的方法，用于將概率編程嵌入到深度學習生態(tài)系統(tǒng)中。

該研究將概率編程提取為一種單一的抽象——隨機變量（random variable）。

在繼續(xù)介紹論文之前，讓我們先了解一下Edward，因為該研究是基于Edward2實現(xiàn)的。

Edward是哥倫比亞大學、谷歌大腦等在2017年提出的新深度概率編程語言，也是一個用于概率建模、推理和評估的Python 庫。Edward 融合了以下三個領域：貝葉斯統(tǒng)計學和機器學習、深度學習、概率編程。

研究人員在TensorFlow的輕量級實現(xiàn)證明該方法可支持多種應用：使用TPUv2的模型并行變分自動編碼器（VAE）；使用TPUv2的數(shù)據(jù)并行自回歸模型Image Transformer；以及多GPUNo-U-Turn Sampler(NUTS)。

對于64x64 ImageNet上最先進的VAE和256x256 CelebA-HQ上最先進的Image Transformer，該方法實現(xiàn)了從1 TPU到256 TPU的最佳線性加速。對于NUTS，相對Stan的GPU加速達到100倍，相對PyMC3的加速達到37倍。

從可微編程到概率編程

深度學習的許多進展可以被解釋為模糊了模型和計算之間的界限。有人甚至提出一種新的“可微編程”（differentiable programming）范式，在這種范式中，其目標不僅僅是訓練一個模型，而是執(zhí)行一般的程序合成。

在這種觀點里，注意力（attention）和門控（gating）描述了布爾邏輯；跳過連接（skip connections）和條件計算描述了控制流；外部存儲訪問函數(shù)內(nèi)部范圍之外的元素。學習算法也越來越動態(tài)：例如， learning to learn，神經(jīng)結(jié)構(gòu)搜索，以及層內(nèi)優(yōu)化等。

可微編程范式鼓勵人們考慮計算成本：不僅要考慮模型的統(tǒng)計特性，還必須考慮其計算、內(nèi)存和帶寬成本。這種理念使研究人員設計出深度學習系統(tǒng)，這些系統(tǒng)運行在最前沿的現(xiàn)代硬件上。

相比之下，概率編程社區(qū)傾向于在模型和計算之間劃清界限：首先，將概率模型指定為程序；其次，執(zhí)行“推理查詢”來自動訓練給定數(shù)據(jù)的模型。

這種設計使得很難真正大規(guī)模地實現(xiàn)概率模型，因為訓練具有數(shù)十億參數(shù)的模型需要跨加速器地拆分模型計算和調(diào)度通信。

在這篇論文中，我們描述了一種在深度學習生態(tài)系統(tǒng)中嵌入概率編程的簡單方法; 我們的實現(xiàn)基于TensorFlow和Python，名為Edward2。這種輕量級方法為靈活的建模提供了一種 low-level 的模式——深度學習者可以從使用概率原語進行靈活的原型設計中獲益，并且概率建模者可以從與數(shù)字生態(tài)系統(tǒng)更緊密的集成中獲益。

研究貢獻：

我們將概率編程的核心提煉為單個抽象——隨機變量（random variable）。

這種low-level的設計有兩個重要含義：

首先，它使研究具有靈活性：研究人員可以自由地操作模型計算，以進行訓練和測試。

其次，它可以使用加速器（例如TPU）來實現(xiàn)更大的模型：TPU需要專門的ops，以便在物理網(wǎng)絡拓撲中分配計算和內(nèi)存。

我們舉例說明了三種應用：使用TPUv2的模型并行變分自動編碼器（VAE）；使用TPUv2的數(shù)據(jù)并行自回歸模型（Image Transformer）；以及多GPU No-U-Turn Sampler (NUTS)。

對于64x64 ImageNet上最先進的VAE和256x256 CelebA-HQ上最先進的Image Transformer，我們的方法實現(xiàn)了從1 TPUv2到256 TPUv2的最佳線性加速。對于NUTS，GPU的加速比Stan快100倍，比PyMC3快37倍。

只需要隨機變量

Random Variables Are All You Need！

在這一節(jié)，我們概述了Edward2的概率程序。它們只需要一個抽象：一個隨機變量。然后，我們將描述如何使用跟蹤（tracing）來執(zhí)行靈活的、low-level 的操作。

概率程序、變分程序

Edward2將任何可計算的概率分布具體化為一個Python函數(shù)(程序)。通常，該函數(shù)執(zhí)行生成過程，并返回示例。程序的輸入—以及任何限定范圍的Python變量—表示分布條件的值。

要在程序中指定隨機選擇，我們使用了Edward的RandomVariables。RandomVariables提供了log_prob和sample等方法，包裝TensorFlow Distributions。此外，Edward隨機變量增加了TensorFlow操作的計算圖：每個隨機變量x與圖中的張量x?～p(x)相關聯(lián)。

圖1描述了一個示例：一個Beta-Bernoulli模型

圖1：Beta-Bernoulli program

重要的是，所有的分布——不管下游用什么——都是作為概率程序編寫的。

圖2描述了一個隱式變分程序，即允許采樣但可能不具有易于處理的密度的變分分布。

圖2：Variational program

一般而言，變分程序、proposal programs和對抗式訓練中的discriminators都是可計算的概率分布。如果我們有一個操縱這些概率程序的機制，那么就不需要引入任何額外的抽象來支持強大的推理范例。

下面我們將使用一個model-parallel VAE來展示這種靈活性。

示例：模型并行的變分自動編碼器

圖4實現(xiàn)了一個模型并行的變分自動編碼器(model-parallel VAE)，它由decoder、prior和encoder組成。decoder生成16位音頻；它采用一種自回歸的flow，用于訓練有效地在序列長度上并行化。encoder將每個樣本壓縮成粗分辨率，由一個壓縮函數(shù)參數(shù)化。

圖4：Model-parallel VAE with TPUs, generating 16-bit audio from 8-bit latents

TPU集群在環(huán)形網(wǎng)絡中布置核心，例如，512個核心可以布置為16x16x2的環(huán)面互連。為了利用集群， prior和decoder都應用分布式自回歸流(如圖3所示)。

圖3：Distributed autoregressive flows

概率程序很簡潔。它們利用了最近的進展，如autoregressive flows和multi-scale latent variables，并且實現(xiàn)了以前從未嘗試過的架構(gòu)，其中使用16x16 TPUv2芯片(512核心)，模型可以在4.1TB內(nèi)存中分割，并使用最多1016個FLOPS。VAE的所有元素——分布式、架構(gòu)和計算位置——都是可擴展的。

跟蹤

我們將概率程序定義為任意Python函數(shù)。為了實現(xiàn)靈活的訓練，我們應用了跟蹤（tracing），這是概率編程的經(jīng)典技術，以及自動微分（automatic differentiation）。

圖5顯示了核心實現(xiàn)：10行的代碼。

圖5：tracing的最小實現(xiàn)

圖6：程序執(zhí)行。

其他示例，包括數(shù)據(jù)并行Image Transformer，No-U-Turn Sampler，概率程序?qū)R，通過梯度下降的變分推理學習等，請閱讀原始論文。

結(jié)論

我們描述了一種簡單、低級別的方法，用于在深度學習生態(tài)系統(tǒng)中嵌入概率編程。對于64x64 ImageNet上的最先進的VAE和256x256 CelebA-HQ上的Image Transformer，我們實現(xiàn)了從1到256 TPUv2芯片的最佳線性加速。對于NUTS，相比其他系統(tǒng)速度提升100倍。

目前，我們正在推進這種設計，作為生成模型和貝葉斯神經(jīng)網(wǎng)絡基礎研究的一個階段。此外，我們的實驗依賴于數(shù)據(jù)并行性以得到大幅的加速加速。最近的一些研究改進了神經(jīng)網(wǎng)絡的分布式編程，用于模型并行性以及對大規(guī)模輸入(如超高分辨率圖像)的并行性。結(jié)合這項工作，我們希望突破超過1萬億參數(shù)和超過4K分辨率的巨型概率模型的極限。