Edward2.2，一種可以用TPU大規模訓練的概率編程

選自 arXiv

作者：Dustin Tran等

機器之心編譯

本文描述了一種用於在深度學習生態系統中嵌入概率編程的簡單、底層方法。具體來說，作者將概率編程歸結為一個抽象：隨機變數。

深度學習的很多研究結果都模糊了模型和計算之間的界限，有的甚至表明是一種「可微分編程」的新範式，它們的目標不僅僅是訓練模型，同時還希望實現一般的程序綜合體。在這一觀點下，注意力機制和門控機制可以描述布爾邏輯運算符，殘差連接和條件計算可以描述控制流，外部記憶可以訪問函數內部作用範圍外的元素。此外，學習演算法也將變得越來越動態，例如學習如何學習、神經架構搜索和層級內的最優化等。

可微分編程範式鼓勵建模者明確考慮計算成本，我們不僅需要考慮模型的統計屬性，例如模型到底能不能擬合真實數據分布，擬合得又有多好。同時我們還需要考慮計算量、內存佔用和傳輸帶寬的成本等，這裡主要關注訓練和預測的效率究竟怎樣。

相比之下，概率編程社區主要傾向於在模型和計算之間劃清界限：首先我們可以指定一個概率模型作為程序，其次我們可以執行「inference query」以在給定數據的情況下自動訓練模型。這種設計使得概率編程很難在大規模真實環境下實現概率模型，其中大規模可能需要訓練有數十億參數量的模型，並要求通過多個加速器和調度通信切分模型計算量。

Edward 等最近的研究能更好地控制深度學習中的推斷過程，然而它們都將推斷過程作為一個封閉系統，因此很難構成任意的計算或更廣泛的機器學習生態系統。

在這篇論文中，作者描述了一種在深度學習生態系統中嵌入概率編程的簡單方法，且他們在 TensorFlow 和 Python 實現了新的庫 Edward2.2。這種輕量方法為靈活建模提供了底層模塊，深度學習研究者可以使用概率基元靈活地設計原型。本論文主要的貢獻有三點：

作者將概率編程的核心提煉為單個抽象：隨機變數。

這種底層設計有兩個重要意義：首先它允許研究具有足夠的靈活性；其次它允許使用張量計算單元（TPU）等加速器實現更大的模型。

作者後面還舉例了三種應用：通過 TPU 訓練的模型並行化變分自編碼器；由 TPU 訓練的數據並行化自回歸模型；和多 GPU No-U-Turn Sampler (NUTS)。本文並沒有介紹這三種應用，感興趣的讀者可查看原論文。

論文：Simple, Distributed, and Accelerated Probabilistic Programming

論文地址：https://arxiv.org/abs/1811.02091

摘要：我們描述了一種用於在深度學習生態系統中嵌入概率編程的簡單、底層方法。具體來說，我們將概率編程歸結為一個抽象：隨機變數。我們在 TensorFlow 中的輕量級實現支持許多應用：通過 TPU（TPUv2）訓練的模型並行變分自編碼器（VAE）；通過 TPUv2 訓練的數據並行自回歸模型（Image Transformer）；多 GPU 且無 U-Turn 的取樣器（NUTS）。對於 64x 64 ImageNet 上的最新 VAE 和 256 x256 CelebA-HQ 上的 Image Transformer，我們的方法實現了從 1 到 256 塊 TPUv2 晶元的線性加速。通過 NUTS，我們發現在 GPU 我們的方法比 Stan 快 100 倍，比 PyMC3 快 37 倍。

隨機變數

我們在 Edward2 中概述了概率編程。它們只需要一個抽象：隨機變數。然後，我們描述了如何利用追蹤來執行靈活的底層操作。Edward 2 將任何可計算的概率分布具化為 Python 函數（程序），通常該函數執行生成過程並返回示例。

圖 1：Beta-Bernoulli 程序。在即時運行模式中，model() 生成一個包含 50 個因素的二進位向量，model() 返回一個 op，該 op 在 TensorFlow 會話中進行評估。

圖 2：變分程序 [35]，可在即時運行模式下獲得。Python 控制流可用於生成流程：給定硬幣翻轉，該程序可根據兩個神經網路中的一個 生成流程。其輸出具有不同的形狀和結構。

圖 3：分散式自回歸流。（右）默認長度是 8，每個有 4 個獨立的流。每個流通過關於自回歸排序的層來轉換輸入。（左）流跨越 4x4 內核（矩形）的虛擬拓撲進行劃分；每個內核計算兩個流，並在本地連接；最後會彙集一個內核。虛擬拓撲與物理 TPU 拓撲對齊：對於 4x4 TPU，它很精確；對於 16x16 TPU，它要為了數據並行性而複製。

圖 4：通過 TPU 訓練模型並行的 VAE，從 8 位延遲生成 16 位 latents。先驗和解碼器根據分散式自回歸流來分割計算。編碼器可能會根據壓縮器分割計算；為了節省空間我們選擇忽略它。

圖 6：一個程序執行。它是有向無環圖，並追蹤累計對數概率等各種 op 或尋找條件獨立性。

本文為機器之心編譯，轉載請聯繫本公眾號獲得授權。

------------------------------------------------

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