NIPS 2018：谷歌大腦提出簡單、分散式概率編程，可用TPU大規模訓練

新智元報道

來源：arXiv

作者：Google Brain 編譯：肖琴

【新智元導讀】谷歌大腦最新提出了一種簡單的方法，用於將概率編程嵌入到深度學習生態系統。這種簡單分散式、加速的概率編程，可以將概率程序擴展到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解析度的巨型概率模型的極限。

論文：

https://arxiv.org/pdf/1811.02091.pdf

【加入社群】

新智元 AI 技術 + 產業社群招募中，歡迎對 AI 技術 + 產業落地感興趣的同學，加小助手微信號：aiera2015_3入群；通過審核後我們將邀請進群，加入社群後務必修改群備註（姓名 - 公司 - 職位；專業群審核較嚴，敬請諒解）。

喜歡這篇文章嗎？立刻分享出去讓更多人知道吧！