Simple, Distributed, and Accelerated Probabilistic Programming
We describe a simple, low-level approach for embedding probabilistic programming in a deep learning ecosystem. In particular, we distill probabilistic programming down to a single abstraction-the random variable. Our lightweight implementation in TensorFlow enables numerous applications: a model-parallel variational auto-encoder (VAE) with 2nd-generation tensor processing units (TPUv2s); a data-parallel autoregressive model (Image Transformer) with TPUv2s; and multi-GPU No-U-Turn Sampler (NUTS). For both a state-of-the-art VAE on 64x64 ImageNet and Image Transformer on 256x256 CelebA-HQ, our approach achieves an optimal linear speedup from 1 to 256 TPUv2 chips. With NUTS, we see a 100x speedup on GPUs over Stan and 37x over PyMC3.
def model(): p = ed.Beta(1., 1., name="p") x = ed.Bernoulli(probs=p, sample_shape=50, name="x") return x Figure 2: Variational program [35], available in eager mode. Python control flow is applicable to generative processes: given a coin flip, the program generates from one of two neural nets. Their outputs can have differing shape (and structure).
system: this makes them difficult to compose with arbitrary computation, and with the broader machine learning ecosystem, such as production platforms [5].
In this paper, we describe a simple approach for embedding probabilistic programming in a deep learning ecosystem; our implementation is in TensorFlow and Python, named Edward2. This lightweight approach offers a low-level modality for flexible modeling-one which deep learners benefit from flexible prototyping with probabilistic primitives, and one which probabilistic modelers benefit from tighter integration with familiar numerical ecosystems.
Contributions. We distill the core of probabilistic programming down to a single abstraction-the random variable. Unlike existing languages, there is no abstraction for learning: algorithms may for example be functions taking a model as input (another function) and returning tensors.
This low-level design has two important implications. First, it enables research flexibility: a researcher has freedom to manipulate model computation for training and testing. Second, it enables bigger models using accelerators such as tensor processing units (TPUs) [22]: TPUs require specialized ops in order to distribute computation and memory across a physical network topology.
We illustrate three applications: a model-parallel variational auto-encoder (VAE) [24] with TPUs; a data-parallel autoregressive model (Image Transformer [31]) with TPUs; and multi-GPU No-U-Turn Sampler (NUTS) [21]. For both a state-of-the-art VAE on 64x64 ImageNet and Image Transformer on 256x256 CelebA-HQ, our approach achieves an optimal linear speedup from 1 to 256 TPUv2 chips. With NUTS, we see a 100x speedup on GPUs over Stan [8] and 37x over PyMC3 [39].
The paper describes Edward2, a probabilistic programming language on top of TensorFlow and supports the TPU, GPU and CPU accelerators. It is the successor of the Edward-1 package.  Edward-2's main USP is that it has a single abstraction (random variable), and offers no separate abstractions to describe learning itself. The learning procedures are left as functions that are user-defined, making Edward-2 low-level (and powerful -- claim).  The paper describes Edward2's tracing mechanism, which allows one to do complex inference procedures, doing interventions on outcomes, etc.  The paper does experiments on No-U-Turn sampler, comparing to other packages in the ecosystem. Additionally it does experiments on VAEs, also comparing to other packages in the ecosystem.  The results showcase that Edward's usage of GPUs (in NUTS) vs Stan and PyMC3 give it considerable speedups (attributing to the use of GPUs). Additionally for VAEs, they show that their use of TPUs give it a 5x speedup over using GPUs, attributing to the TPU being faster.  ------------------------------  The paper is very incomprehensible. I am familiar with TensorFlow and a little bit of probabilistic programs, but I could barely comprehend the paper. I read it 2 and a half times, and I was often short of context. I think the authors seem to have struggled with the page limits and compressed content + removed context. Sections such as 2.2 and 3, 3.1, 3.2 are clearly shortened beyond comprehension for a reader that is not familiar with Edward-1 or whatever context that the authors assume that the reader has. While I have a superficial "tl;dr" of such sections, the point that was being made in the section didn't actually come across. I have very little understanding of what makes Edward2 better. All the code samples are not discussed in sufficient detail either, often "assuming" that the reader "gets" the superiority of what is being presented.  While the project seems to be an impressive effort: "PPL with CPU, GPU, TPU support and considerable integration with popular framework such as TensorFlow", the paper should be expanded and probably sent to a venue or journal where one does not have to fight the page limits and can express themselves better. I dont have confidence that as a poster, the authors can do a good job of explaining the framework well. This is my primary reason for rejection of the paper, and I will let the Area Chair make a decision on whether this is a valid reasoning line or not.  Lastly, something that seems to be a dishonest "wording" is their sentences around speedups against other packages. Starting from the abstract, and continuing into many places in the paper (8, 44 Table-{1, 2, 3} captions, 235, 243, 9, Table-1 caption, 219, 244) they keep repeating "Edward2 on TPU vs Pyro on GPU" or "Edward2 on GPU vs STAN on CPU". This is not only disingenous, it's dangerous. If you see Table-2, it is clear that Edward2 on TPU vs Edward2 on GPU is a 5x speed gain. So the main speed gain is actually TPU vs GPU, yet they clearly try to mix up the language so that "Edward2" is associated with the speedup rather than "TPU". There's more problematic issues with the benchmarking of TPUs vs GPUs. They compare the TPU to a 1080Ti GPU. They are comparing the latest generation TPU hardware to an older generation GPU hardware. This is not acceptable, especially if they want to make their claim that TPUs are faster. If they compare a TPU against Volta GPUs, will their speedup reduce by a significant amount? Definitely, as Volta's fp32 is a 2x number of flops and has more memory bandwidth. This is my added reason for rejecting the paper in it's current form.   Nit: 55, Fig.4, 73 refer Edward2 as "Edward", making it confusing.