谷歌大腦提出了一種簡(jiǎn)單的方法，用于將概率編程嵌入到深度學(xué)習(xí)生態(tài)系統(tǒng)

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

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

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

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

在繼續(xù)介紹論文之前，讓我們先了解一下Edward，因?yàn)樵撗芯渴腔贓dward2實(shí)現(xiàn)的。

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

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

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

從可微編程到概率編程

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

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

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

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

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

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

研究貢獻(xiàn)：

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

這種low-level的設(shè)計(jì)有兩個(gè)重要含義：

首先，它使研究具有靈活性：研究人員可以自由地操作模型計(jì)算，以進(jìn)行訓(xùn)練和測(cè)試。

其次，它可以使用加速器（例如TPU）來(lái)實(shí)現(xiàn)更大的模型：TPU需要專(zhuān)門(mén)的ops，以便在物理網(wǎng)絡(luò)拓?fù)渲蟹峙溆?jì)算和內(nèi)存。

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

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

只需要隨機(jī)變量

Random Variables Are All You Need！

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

概率程序、變分程序

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

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

圖1描述了一個(gè)示例：一個(gè)Beta-Bernoulli模型

圖1：Beta-Bernoulli program

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

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

圖2：Variational program

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

下面我們將使用一個(gè)model-parallel VAE來(lái)展示這種靈活性。

示例：模型并行的變分自動(dòng)編碼器

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

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

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

圖3：Distributed autoregressive flows

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

跟蹤

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

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

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

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

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

結(jié)論

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

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