arXiv:1709.05870v1 [stat.ML] 18 Sep 2017

ZhuSuan: A Library for Bayesian Deep Learning∗
Jiaxin Shi1
Jianfei Chen1
Jun Zhu1
Shengyang Sun2
Yucen Luo1
Yihong Gu1
Yuhao Zhou1

shijx15@mails.tsinghua.edu.cn
chenjian14@mails.tsinghua.edu.cn
dcszj@tsinghua.edu.cn
ssy@cs.toronto.edu
luoyc15@mails.tsinghua.edu.cn
gyh15@mails.tsinghua.edu.cn
zhouyh16@mails.tsinghua.edu.cn
1
Department of Computer Science & Technology, TNList Lab, CBICR Center
1
State Key Lab of Intelligent Technology & Systems
2
Department of Electronic Engineering
Tsinghua University
Beijing, 100084, China

Abstract
In this paper we introduce ZhuSuan, a python probabilistic programming library for Bayesian
deep learning, which conjoins the complimentary advantages of Bayesian methods and deep
learning. ZhuSuan is built upon Tensorflow. Unlike existing deep learning libraries, which
are mainly designed for deterministic neural networks and supervised tasks, ZhuSuan is
featured for its deep root into Bayesian inference, thus supporting various kinds of probabilistic models, including both the traditional hierarchical Bayesian models and recent deep
generative models. We use running examples to illustrate the probabilistic programming
on ZhuSuan, including Bayesian logistic regression, variational auto-encoders, deep sigmoid
belief networks and Bayesian recurrent neural networks.
Keywords: Bayesian Inference, Deep Learning, Probabilistic Programming, Deep Generative Models

1. Introduction
These years have seen great advances of deep learning (LeCun et al., 2015) and its success in many applications, such as speech recognition (Hinton et al., 2012), computer vision (Krizhevsky et al., 2012), language processing (Sutskever et al., 2014), and computer
games (Silver et al., 2016). One good lesson we learned from the practice is that a deeply
architected model can be well fit by leveraging advanced optimization algorithms (e.g.,
stochastic gradient decent), a large training set (e.g., ImageNet), and powerful computing
devices (e.g., GPUs). Although the great expressiveness of deep neural networks (DNNs)
has made them a first choice for many complex fitting problems, especially those tasks
involving a mapping from an input space to an output space (e.g., classification and regression), the results given by them are usually point estimates. Such a deterministic approach
does not account for uncertainty, which is essential in every part of a learning machine (e.g.,
data, estimation, inference, prediction and decision-making) due to the randomness of the
∗. J. Zhu is the corresponding author. S. Sun is now with Department of Computer Science, University of
Toronto.

1

Shi et al.

physical world, incomplete information, and measurement noise. It has been demonstrated
that deterministic neural networks can be vulnerable to adversarial attacks (Szegedy et al.,
2013; Nguyen et al., 2015) (partly because of the lack of uncertainty modeling), which may
hinder their applications in the scenarios where representing uncertainty is of crucial importance, such as automated driving (Bojarski et al., 2016) and healthcare (Miotto et al.,
2017).1 On the other hand, the probabilistic view of machine learning offers mathematically
grounded tools for dealing with uncertainty, i.e., Bayesian methods (Ghahramani, 2015).
Bayesian methods also provide a theoretically sound approach to incorporating structural
bias (Lake et al., 2015) and domain knowledge as prior or posterior constraints (Mei et al.,
2014) to achieve efficient learning with a small number of training samples. Thus it is beneficial to combine the complimentary advantages of deep learning and Bayesian methods,
which in fact has been drawing tremendous attention in recent years (Gal, 2016).
Moreover, as most of the success in deep learning comes from supervised tasks that require a large number of labeled data, research in this area is paying more and more attention
to unsupervised learning, a long-standing goal of artificial intelligence (AI). Probabilistic
models and Bayesian methods are commonly treated as principled approaches to modeling unlabeled data for pure unsupervised learning or its hybrid with supervised learning
(aka., semi-supervised learning). Recently, the popularity of deep generative models again
demonstrates the promise of combining deep neural networks with probabilistic modeling,
which has shown superior results in image generation (Kingma and Welling, 2013; Goodfellow et al., 2014; Radford et al., 2015), semi-supervised classification (Kingma et al., 2014;
Salimans et al., 2016) and one-shot learning (Rezende et al., 2016).
We call such an arising direction that conjoins the advantages of Bayesian methods and
deep learning as Bayesian Deep Learning (BDL). The scope of BDL covers the traditional
Bayesian methods, the deep learning methods where probabilistic inference plays a key role,
and their intersection. One unique feature of BDL is that the deterministic transformation
between random variables can be automatically learned from data under an expressive
parametric formulation typically using deep neural networks (Johnson et al., 2016), while
in traditional Bayesian models, the transformation tends to have a simple analytical form
(e.g., the exponential function or inner product). One key challenge for Bayesian deep
learning is on posterior inference, which is typically intractable for such models and needs
sophisticated approximation techniques. Although much progress has been made on both
variational inference and Monte Carlo methods (Zhu et al., 2017), it is still too difficult
for practitioners to pick-up. Moreover, although variational inference and Monte Carlo
methods have their general recipes, it is still quite involved for an experienced researcher to
derive the equations and implement every detail for a particular model. Such a procedure
is error prone and takes a long time to debug.
In this paper, we present ZhuSuan, a probabilistic programming library for Bayesian
Deep Learning.2 We build ZhuSuan upon Tensorflow (Abadi et al., 2016) to leverage its
computation graphs for flexible modeling. With ZhuSuan, users can enjoy powerful fitting
1. Some recent efforts have been made towards improving the robustness of deep neural networks against
adversarial samples by doing adversarial training (Szegedy et al., 2013), reverse cross-entropy training (Pang et al., 2017) or Bayesian averaging (Li and Gal, 2017).
2. ZhuSuan is named after the Chinese name of abacus, which is the oldest calculating machine and has
been recognized as the fifth greatest innovation in China.

2

ZhuSuan: A Library for Bayesian Deep Learning

and multi-GPU training of deep learning, while at the same time they can use probabilistic models to model the complex world, exploit unlabeled data and deal with uncertainty
by applying principled Bayesian inference. We provide running examples to illustrate how
intuitive it is to program with ZhuSuan, including Bayesian logistic regression, variational
auto-encoders, deep sigmoid belief networks and Bayesian recurrent networks. More examples can be found in our code repository https://github.com/thu-ml/zhusuan. The
key difference of ZhuSuan from other probabilistic programming libraries lies in its flexibility, benefited from the deep learning paradigm on which the library is built, and the
key treatment of model reuse. More detailed comparisons with two closely related works,
Edward (Tran et al., 2016) and PyMC3 (Salvatier et al., 2016), are provided in this paper.

2. Bayesian Deep Learning
We first overview the basics of Bayesian methods and their intersection with deep learning,
and define the notations.
2.1 Modeling
As a conjunction of graph theory and probability theory, probabilistic graphical models
(PGMs) (Koller and Friedman, 2009) provide a powerful language for probabilistic machine learning. A PGM defines a joint distribution of a set of random variables using a
graph, which has an intuitive and compact semantic meaning to read out the conditional
independence structures. There are two major types of PGMs, namely, directed ones (aka.
Bayesian networks) and undirected ones (aka. Markov random fields, MRFs). Though
Bayesian networks and MRFs have different expressiveness, sometimes one can transform
a Bayesian network to an MRF or vice versa (Koller and Friedman, 2009). In ZhuSuan, we
concentrate on Bayesian networks, which provide an intuitive data generation process (Gelman et al., 2014). This choice is also reflected in the literature of Bayesian deep learning.
For example, while modern deep generative models make their first remarkable step with
undirected graphs (Hinton et al., 2006), the use of directed models (Kingma and Welling,
2013) has dominated the area due to cheap generation and fast inference.

(a)

(b)

(c)

Figure 1: Bayesian networks
Figure 1a shows a Bayesian network with a directed acyclic graph (DAG) structure.
This represents a joint probability distribution over the nodes that factorizes according to
3

Shi et al.

the graph:
p(Z, H, U, S, A, N ) = p(Z)p(H)p(U |Z)p(S|Z, H)p(A|S)p(N |U, A).

(1)

The joint probability is the product of two types of factors: the prior distribution of variables
at the roots (e.g., p(Z)) and the conditional probability distribution of other variables given
their parents (e.g., p(S|Z, H)). The conditional distributions are also known as conditional
probability tables (CPTs) for discrete variables and conditional probability density (CPD)
functions for continuous variables.
We consider the general formulation, where there are two kinds of nodes in a Bayesian
network: deterministic nodes (Figure 1b, right, double line notation) and stochastic nodes
(Figure 1b, left). A deterministic node calculates its value by a deterministic function of its
parents, while a stochastic node is represented by a standard probabilistic distribution. Although deterministic nodes are often not represented explicitly to make the graph structure
concise, it is beneficial to do so especially when we consider Bayesian deep learning models.
This is because the transformation in a traditional model is typically in a simple analytical
form, while a Bayesian deep learning model learns it from a very flexible family (e.g., deep
neural networks). Hence, explicitly representing deterministic nodes helps on both model
definition and inference, as shall be clear soon.
If all the variables are observed in the dataset, it is easy to do learning and inference for
a Bayesian network. However, in reality it is more common to have partially observed data
due to physical randomness, missing information and measurement noise, thereby some
variables are hidden. Such models are known as latent variable models (LVMs), which
provide a suite of useful tools to unveil the underlying factors (e.g., topics or latent feature
representations). Figure 1c illustrates a generic latent variable model, where the gray nodes
represent the observed variables (x1:n ), the rest (h1:n , θ) are latent variables, and n is the
number of observed data samples. This class of models include both global and local latent
variables. By local we mean the latent variable only has impact on its paired observation
(e.g., h is paired with x), while global means it influences all observations (e.g., θ).
2.2 Inference
For a latent variable model, posterior inference is the key step to unveil the latent factors for
a particular input. To make the notations simple, for this subsection we consider a highly
abstract model p(z, x) = p(z)p(x|z), where z represents all latent variables, including the
global and local ones, i.e., z = (h, θ), and x denotes observations. The vanilla Bayes’ rule
provides a principle to derive the posterior distribution:
p(z|x) =

p(z, x)
p(z)p(x|z)
=
.
p(x)
p(x)

(2)

As analyzed in (Zhu et al., 2014), the vanilla Bayes’ rule can be generalized to regularized
Bayesian inference (RegBayes) by introducing posterior regularization under an information
theoretical formulation.
In general, the posterior inference using Bayes’ rule or the RegBayes formulation is
intractable, except a few simple examples. Therefore, we have to resort to approximate
Bayesian inference methods. Many years of developments in this area has led to many fast,
4

ZhuSuan: A Library for Bayesian Deep Learning

widely applicable approximate inference algorithms, mainly divided into two categories,
variational inference and Monte Carlo methods (Zhu et al., 2017). Both will be covered
below.
2.2.1 Variational Inference
Variational inference (VI) is an optimization-based method for posterior approximation,
in which a parametric distribution family qφ (z) is chosen to approximate the true posterior p(z|x) by minimizing the KL divergence between them (KL(qφ (z)kp(z|x)).3 The KLdivergence minimization is equal to maximizing a lower bound of the marginal log likelihood
log p(x):
L(x; φ) = log p(x) − KL(qφ (z)kp(z|x))
(3)
= Eqφ (z) [log p(x|z)] − KL (qφ (z)kp(z)) .
In the VI literature, qφ (z) is called the variational posterior, and L(x; φ) is called the
evidence lower bound (ELBO). Generally there are two steps for doing VI. First is to choose
the parametric family that will be used as the variational posterior. Then the second step
is to solve the optimization problem with respect to the variational parameters (φ).
In recent years, benefited from the joint effort of the Bayesian and the deep learning
community, variational inference is undergoing many significant changes, both in the algorithm and the choices of variational families. From the algorithm side, stochastic approximation with data sub-sampling has enabled VI to scale to large datasets (Hoffman et al.,
2013). Meanwhile, direct optimization of the variational lower bounds by gradient descent
is replacing traditional analytic updates, which makes VI applicable to a broader range of
models with non-conjugate dependencies (Graves, 2011; Titsias and Lázaro-Gredilla, 2014).
The key challenge that draws most attention in this direction is on the gradient estimates.
Many gradient estimators have been developed with different variance reduction techniques
(Paisley et al., 2012; Mnih and Gregor, 2014; Mnih and Rezende, 2016). Specially, for continuous latent variable models, an algorithm named Stochastic Gradient Variational Bayes
(SGVB) has been very successful (Kingma and Welling, 2013), due to a clever trick to pass
the gradient through a stochastic node, which is now well-known as the reparameterization
trick. Besides, better lower bounds than the ELBO have also been developed, which can be
optimized by gradient descent as well and have been applied to both discrete and continuous
latent variable models (Burda et al., 2015; Mnih and Rezende, 2016).
On the other side, considerable efforts have also been put into the design of variational
posteriors. Using neural networks to parameterize the variational posterior has been a
common practice adopted by aforementioned works (Mnih and Gregor, 2014; Kingma and
Welling, 2013), serving as one of the key technologies to efficiently train deep generative
models. The network is usually fed with observations and is expected to amortize the
inference cost for each data point by learning to do inference. This type of scheme is
often referred as amortized inference (Rezende and Mohamed, 2015). In summary, VI in
the Bayesian deep learning context is more stochastic, differentiable, and amortized than
before.
3. We do not explicitly condition z on x in q because it is not necessary to explicitly model this, though
many modern literatures do. In fact, even without explicit modeling, the optimization process will
connect them together.

5

Shi et al.

2.2.2 Monte Carlo Methods
Unlike VI that reformulates the inference as an optimization problem, Monte Carlo methods
(Robert and Casella, 2005) are more direct ways to simulate samples or estimate properties
of a particular distribution. Two main techniques extensively used for Bayesian inference
are Importance Sampling and Markov Chain Monte Carlo.
Importance Sampling The basic idea of importance sampling is as follows. To estimate µ = Ep f (x) where p is a probability density defined over Rd , a probability distribution
q called the proposal is introduced. It is required that q(x) > 0 whenever p(x)f (x) > 0.
Then


Z
Z
f (x)p(x)
f (x)p(x)
f (x)p(x) dx =
µ=
.
(4)
q(x) dx = Eq
q(x)
q(x)
Rd
Rd
The Monte Carlo estimate of µ given by importance sampling is
µ̂ =

N
1 X f (xi )p(xi )
,
N
q(xi )

xi ∼ q(x).

(5)

i=1

In the Bayesian inference context, it is often the case that we only have access to an
unnormalized version of p, which we denote as p̃. p = p̃/Z, where Z is the normalizing
constant which is intractable. Self-normalized Importance Sampling is particularly useful
in this situation, which gives the estimate
µ̃ =

N
X
i=1

w̃i f (xi ),

where

wi
w̃i = PN

j=1 wj

, wi =

p̃(xi )
.
q(xi )

(6)

It can be proved that µ̃ is an asymptotically unbiased estimator of µ as N → ∞. For more
detailed introduction, we refer the readers to (Owen, 2013).
As introduced above, importance sampling in the strict sense is not a sampling method
because it does not directly draw samples from the target distribution. Instead, it provides
a method for estimating properties of a certain distribution (normalized or not). Thus the
use of importance sampling can be everywhere that needs an estimate of some integral over
a probability measure. Though the computation in eqs. (5) and (6) is rather direct, some
of their application scenarios may not be obvious for users.
Here we focus on application scenarios in Bayesian deep learning, which concentrate on
model learning and evaluation. For model learning, it is shown in (Bornschein and Bengio, 2014) that self-normalized importance sampling can be used to estimate gradients of
marginal log likelihoods with respect to model parameters, where this technique was proposed to train deep generative models. Even if not used for learning, importance sampling
is still a simple and efficient method for evaluating marginal log likelihoods of latent variable models (Wallach et al., 2009). One of the drawbacks of using importance sampling
is that the estimation has large variance if the proposal is not good enough (Owen, 2013).
Therefore people have been exploring the use of adaptive proposals (Cappé et al., 2008;
Cornuet et al., 2012). This idea is recently reformed into a new technique called neural
adaptive proposals, i.e., a neural network is used to parameterize the proposal distribution,
which is adapted towards the optimal proposal by gradient descent (Bornschein and Bengio, 2014). This technique proves to be successful in applications like dynamic systems (Gu
6

ZhuSuan: A Library for Bayesian Deep Learning

et al., 2015), inference acceleration (Du et al., 2017; Paige and Wood, 2016) and learning
attention models (Ba et al., 2015).
Markov Chain Monte Carlo Markov Chain Monte Carlo (MCMC) is a classic method
to generate samples from the posterior distribution in Bayesian statistics. Unlike variational
inference, MCMC is known to be asymptotically unbiased and allows the user to trade off
computation for accuracy without limit. The basic idea of MCMC is to design a Markov
chain whose stationary distribution is the target distribution, then samples can be generated
by simulating the chain until convergence. Specifically, for a Markov chain specified by
the transition kernel T (z0 |z), the most common sufficient condition for p(z|x) being its
stationary distribution is the detailed balance condition:
p(z|x)T (z0 |z, x) = p(z0 |x)T (z|z0 , x),

for all z, z0 .

An arbitrarily chosen initial state z0 and the transition kernel define a joint distribution
p(z0 , . . . , zt , . . . |x) and marginal distributions p(zt |x). It can be shown that under mild
conditions, p(zt |x) converges to p(z|x) as t → ∞ (Robert and Casella, 2005). In practice,
it is common to throw away samples from the initial stage before the chain converges. This
stage is often referred as the burn-in (or warm-up) stage. Once we get samples from the
MCMC methods, there are many ways to use them. For parameter estimation of LVMs, the
samples can be used in the Monte-Carlo EM (MCEM) algorithm (Wei and Tanner, 1990).
Throughout the literature there are many ways to design the transition kernel. For
example, the simplest form of the Metropolis Hastings algorithm (Metropolis et al., 1953;
Hastings, 1970) combines a Gaussian random walk proposal with an accept-reject test for
correction, which scales poorly with increasing dimension and complexity of the target distribution. Gibbs sampling (Geman and Geman, 1984) utilizes the structure of the target
distribution by taking its element-wise conditional distribution as the transition proposal.
However, it requires the conditionals to be analytically computable, which limits its application scope, especially in Bayesian deep learning.
A powerful MCMC method that efficiently explores high-dimensional continuous distributions is Hamiltonian Monte Carlo (HMC) (Neal et al., 2011). To sample from p(z|x),
HMC introduces an auxiliary momentum variable p with the same dimensionality as z,
and effectively samples from the joint distribution p(z, p|x) = p(z|x) exp(− 12 p> M −1 p),
where M is called the mass matrix. Samples are generated by simulating the Hamiltonian
dynamics, which governs the evolution of the (z, p) system along continuous time t as:
∂z
= ∇p H,
∂t

∂p
= −∇z H.
∂t

(7)

Here H(z, p) = − log p(z, p|x) is the Hamiltonian of the system. The Hamiltonian dynamics defines a continuous-time transition kernel, and the stationary distribution of the
corresponding Markov chain is the desired p(z, p|x), due to the volume-preservation and
the Hamiltonian-conservation properties of the dynamics. In practice one has to simulate the dynamics in discrete time steps. The leapfrog integrator (Leimkuhler and Reich,
2004) is widely used, since it keeps the two vital properties of the Hamiltonian dynamics: it is volume-preserving and approximately Hamiltonian-conserving. Combined with
a Metropolis-Hastings algorithm to correct the approximating error of the Hamiltonian,
the discrete-time simulation comes to our satisfactory: the Markov chain converges to our
7

Shi et al.

desired distribution. The key advantage of HMC is that it exploits information about the
geometry of the target distribution (Betancourt, 2017) while only needs the distribution
density and the gradient to proceed.
People have been exploring the use of HMC in Bayesian deep learning. The early work
by Radford Neal (Neal, 1995) applied HMC to the inference of Bayesian neural networks,
one of the representative models that apply Bayesian methods to capture uncertainty in
deep learning. Recent works on deep generative models have applied HMC to improve the
variational posterior (Salimans et al., 2015) as well as used HMC-based MCEM algorithm
to directly learn the model parameters (Hoffman, 2017; Li et al., 2017; Titsias, 2017).

3. Design Overview
ZhuSuan has been designed towards the basic needs of Bayesian deep learning, which, as
stated in the last section, mainly include two parts: modeling and inference. As for the
modeling part, we follow the principle that the code reads like the model definition. This
requires model primitives that support:
• Direct building of probabilistic models as computation graphs.
• Easy reuse by changing the states of stochastic nodes (observed or not).
• Arbitrary deterministic modeling with the full power of Tensorflow.
On the side of inference, to leverage recent advances in Bayesian deep learning while staying
applicable to a broad class of models, ZhuSuan makes efforts towards:
• Unifying traditional and advanced differentiable inference methods in a deep learning
paradigm.
• Supporting inference for both continuous latent variables and discrete ones.
Also from users’ perspective, the design of ZhuSuan has been governed by two principles 4 :
• Modularity: Make flexible abstractions that allow all parts to be used independently.
• Transparency: Do not hide the whole inference procedure behind abstractions, to
enable heavy customization by users.
This makes ZhuSuan particularly different from existing probabilistic programming libraries, which will be discussed in later comparison (Section 5).

4. Features
In this section we introduce the basic features of ZhuSuan, including model primitives and
inference algorithms.
4. Both principles are learned from Lasagne (Dieleman et al., 2015).

8

ZhuSuan: A Library for Bayesian Deep Learning

4.1 Model Primitives
For the reasons we have explained in Section 2.1, ZhuSuan’s model primitives are designed
for Bayesian networks. Since the basic structure is a DAG, we will introduce the node
primitives first, and then will see how ZhuSuan stores and fetches information in the graph.
Deterministic nodes In ZhuSuan, users are enabled to use any Tensorflow operation
for deterministic transformations. This includes various arithmetic operators (e.g., tf.
add,tf.tensordot,tf.matrix_inverse), neural network layers (e.g., tf.layers.fully_
connected,tf.layers.conv2d) and even control flows (e.g., tf.while_loop,tf.cond). In
Tensorflow computation graph, the output of an operation is named Tensor. ZhuSuan does
not add any higher-level abstraction on Tensors but directly treat them as deterministic
nodes in the Bayesian networks. We will see that the other primitives work well directly
with Tensors.
Stochastic nodes For stochastic nodes, ZhuSuan provides an abstraction called Stocha
sticTensor, which is named following its deterministic counterpart. Many commonly used
probabilistic distributions are implemented and wrapped in StochasticTensors (e.g., Normal, Bernoulli, Categorical), together with some recently developed variants for Bayesian
deep learning, such as Gumbel-softmax or Concrete (Jang et al., 2016; Maddison et al.,
2016). StochasticTensors inherit most behaviors of Tensors. The former can be directly
fed into Tensorflow operations and can automatically be cast into Tensors when computing
with them. When cast into Tensors, the values they choose can be set to samples or observations, or can be automatically determined according to their states given in the context,
which will be introduced next.
The graph context Above we see a Bayesian network can be built transparently with
mixes of Tensorflow operations and StochasticTensors. Because large and sophisticated
models are generally common today, it is still painful for one to deal with all these nodes
individually. To help users manage the graph in a convenient way, ZhuSuan provides a
graph context called the BayesianNet, which keeps track of all named StochasticTensors
constructed in it.
To start a BayesianNet context, use the with statement in python:
import zhusuan as zs
with zs.BayesianNet() as model:
# Build the graph.

The BayesianNet context supports making queries on the inner stochastic nodes. The
query options include the current-state outputs and local probabilities (the current-state
value of CPDs) at certain nodes.
We use several representative examples to illustrate the process of building models in
ZhuSuan, ranging from the simple Bayesian logistic regression to modern Bayesian deep
models. In general, the programming of a probabilistic model on ZhuSuan is as intuitive as
reading the corresponding graphical model, which is provided side-by-side.
Example 1 (Bayesian Logistic Regression, BLR) The generative process of a Bayesian
logistic regression can be written as:
w ∼ N (0, α2 I),
yi ∼ Bernoulli(σ(wT xi )),
9

i = 1, . . . , n,

(8)

Shi et al.

where w, xi ∈ RD , yi ∈ {0, 1}, and σ(·) is the sigmoid function. The graphical model and
the corresponding code is shown in Figure 2. Note that data is modeled in a batch.
import tensorflow as tf
with zs.BayesianNet() as model:
# Define inputs and parameters
x = tf.placeholder([None, D], tf.float32)
log_alpha = tf.Variable(tf.zeros([D]))
# w ~ N(0, alpha^2*I)
w = zs.Normal('w', mean=0., logstd=log_alpha,
group_ndims=1)
# y_logit = w^Tx
y_logit = tf.reduce_sum(
tf.expand_dims(w, 0)*x, axis=1)
# y ~ Bernoulli(sigmoid(y_mean))
y = zs.Bernoulli('y', y_logit)
(a)

(b)

Figure 2: BLR: (a) graphical model; (b) programming on ZhuSuan.
In this example the deterministic transformation part is the linear model (wT x), which
is implemented by the Tensorflow operations, tf.expand_dims,tf.reduce_sum and tf.
multiply (*) to enable batch processing of the inputs (xi , i = 1, . . . , n). The two random
variables y and w are created as zs.Bernoulli and zs.Normal respectively, which are the
StochasticTensor for Bernoulli and normal distributions. The group_ndims argument
passed to specify the StochasticTensor for w means that the last dimension of w is treated
as a group, whose probabilities are computed together.
Example 2 (Variational Auto-Encoders, VAE) Variational autoencoders (VAE) is one
of the most popular products of Bayesian deep learning. The generative process of a VAE
for modeling binarized MNIST data is
z ∼ N (0, I),
xlogits = fN N (z),

(9)

x ∼ Bernoulli(σ(xlogits )),
where z ∈ RD , x ∈ R784 . This generative process is a stereotype for deep generative models,
which starts with a latent representation (z) sampled from a simple distribution (such as the
standard normal). Then the samples are forwarded through a deep neural network (fN N ) to
capture the complex generative process of high dimensional observations such as images (x).
Finally, some noise is added to the output to get a tractable likelihood for the model. For
binarized MNIST, the observation noise is chosen to be Bernoulli, with its parameters output
by the neural network. Because ZhuSuan is built upon Tensorflow, which has full support
for deep neural networks, the implementation of VAE is extremely easy and intuitive, as
shown in Figure 3.

10

ZhuSuan: A Library for Bayesian Deep Learning

from tensorflow.contrib import layers
with zs.BayesianNet() as model:
z_mean = tf.zeros([n, D])
z = zs.Normal('z', z_mean, std=1., group_ndims=1)
h = layers.fully_connected(z, 500)
x_logits = layers.fully_connected(h, 784,
activation_fn=None)
x = zs.Bernoulli('x', x_logits, group_ndims=1)
(a)

(b)

Figure 3: VAE: (a) graphical model; (b) programming on ZhuSuan.
Example 3 (Deep Sigmoid Belief Networks, DSBN) Sigmoid Belief Networks (SBN)
is a directed discrete latent variable model that has close connections with feed-forward neural networks and Boltzmann Machines (Neal, 1992). In recent years, the return of neural
networks has brought a new life to this old model. In fact, the well-known Deep Belief Networks (DBN), the earliest work on deep learning, is an infinite-layer tied-weight SBN with
the bottom layers untied (Hinton et al., 2006). The generative process of a DSBN with L
layers is
z(L) ∼ Bernoulli(σ(p(L) )),
z(l−1) ∼ Bernoulli(σ(w(l)T z(l) )),

l = L, . . . , 1

(10)

x = z(0) ,
where σ is the sigmoid function; p(L) is the top layer parameters; w are hidden weights; z
are latent variables, and x are observations. From the definition we can see that DSBN is a
model with multi-layer stochastic nodes. The implementation of a two-layer DSBN (L = 2)
is shown in Figure 4.
with zs.BayesianNet() as model:
z2_logits = tf.zeros([n, n_z])
z2 = zs.Bernoulli('z2', z2_logits, dtype=tf.float32,
group_ndims=1)
z1_logits = layers.fully_connected(z2, n_z,
activation_fn=None)
z1 = zs.Bernoulli('z1', z1_logits, dtype=tf.float32,
group_ndims=1)
x_logits = layers.fully_connected(z1, n_x,
activation_fn=None)
x = zs.Bernoulli('x', x_logits, group_ndims=1)
(a)

(b)

Figure 4: DSBN: (a) graphical model; (b) programming on ZhuSuan.
11

Shi et al.

To provide a more interesting example that utilizes the powerful control flow operations
from Tensorflow, we describe a Bayesian Recurrent Neural Network (Bayesian RNN) below,
which can also be intuitively programmed on ZhuSuan.
Example 4 (Bayesian RNN) We have mentioned previously that deterministic neural
networks lack the ability to account for the uncertainty of its own predictions. A solution to
this given by Bayesian methods is a model named Bayesian Neural Network (Bayesian NN),
which treats the network weights as random variables and infers a posterior distribution over
them given data. The generative process of a Bayesian NN for classification tasks is
W ∼ N (0, I),
π = fN N (x; W),

(11)

y ∼ Cat(softmax(π)),
where fN N is a neural network with W as the weights, π is the predicted unnormalized log
probabilities for all classes, Cat represents the categorical distribution, and y is the class
label. When fN N is chosen to be a recurrent network, the above process describes a Bayesian
RNN. The graphical model for a Bayesian RNN is shown in Figure 5. Below we consider
a model for a two-class sequence classification task. For the RNN part, it uses a Long
Short-Term Memory (LSTM) network.

Figure 5: Bayesian RNN

class BayesianLSTMCell(object):
def __init__(self, num_units, forget_bias=1.0):
self._forget_bias = forget_bias
w_mean = tf.zeros([2 * num_units + 1, 4 * num_units])
self._w = zs.Normal('w', w_mean, std=1., group_ndims=2)
def __call__(self, state, inputs):
c, h = state
batch_size = tf.shape(inputs)[0]
linear_in = tf.concat([inputs, h, tf.ones([batch_size, 1])], axis=1)
linear_out = tf.matmul(linear_in, self._w)
12

ZhuSuan: A Library for Bayesian Deep Learning

# i = input_gate, j = new_input, f = forget_gate, o = output_gate
i, j, f, o = tf.split(value=linear_out, num_or_size_splits=4, axis=1)
new_c = (c * tf.sigmoid(f + self._forget_bias) +
tf.sigmoid(i) * tf.tanh(j))
new_h = tf.tanh(new_c) * tf.sigmoid(o)
return new_c, new_h
def bayesian_rnn(cell, inputs, seq_len):
initializer = (tf.zeros([batch_size, 128]), tf.zeros([batch_size, 128]))
c_list, h_list = tf.scan(cell, inputs, initializer=initializer)
relevant_outputs = tf.gather_nd(
h_list, tf.stack([seq_len - 1, tf.range(batch_size)], axis=1))
logits = tf.squeeze(tf.layers.dense(relevant_outputs, 1), -1)
return logits
with zs.BayesianNet() as model:
cell = BayesianLSTMCell(128, forget_bias=0.)
logits = bayesian_rnn(cell, self.x, self.seq_len)
_ = zs.Bernoulli('y', logits, dtype=tf.float32)

The code splits into three parts. First we build a BayesianLSTMCell that applies transformations to the inputs and the hidden states at each time step. Its only difference from
the tf.nn.rnn_cell.BasicLSTMCell in Tensorflow is that the weights are generated from a
normal StochasticTensor. Then the tf.scan operation is used to apply the cell function to
input sequences. This part utilizes the power of control flows to deal with the variable-length
sequences (tf.scan is internally implemented by the tf.while_loop operator). Finally, a
Bernoulli StochasticTensor generates the class label given the outputs from the RNN.
Model reuse Unlike supervised neural networks, a key feature of probabilistic graphical
models is polymorphism, i.e., a stochastic node can have two states: observed or latent.
This is a major difficulty when designing programming primitives. For example, consider
the above VAE case. If z is a Tensor sampled from the Normal distribution, when the model
is again used in a case where z is observed, the common practice in Tensorflow programming
is to write another piece of code that builds a graph with the observation of z as the input.
When the set of stochastic nodes is large, the process is very painful.
The reusability problem is a shared problem for probabilistic programming libraries
that are based on computation graph toolkits, e.g., Theano (Al-Rfou et al., 2016) and Tensorflow (Abadi et al., 2016). Other libraries such as PyMC3 (Salvatier et al., 2016) and
Edward (Tran et al., 2016) address this problem by directly manipulating the created computation graph. Specifically, PyMC3 (based on Theano) relies on the officially supported
theano.clone() function to copy and recreate subgraphs that are influenced when the state
of a node changes. While Edward (based on Tensorflow) has to implement their own copy()
function by looking into nonpublic low-level APIs of the Tensorflow computation graph. As
we will see later, the solutions by manipulating the created graphs induce problems and
limitations for these libraries.
13

Shi et al.

ZhuSuan has been carefully designed towards reusability, but does not rely on altering
the created graphs. Specifically, an observation can be passed to StochasticTensor directly, which enables model reuse by repeated calls of a function with different arguments
passed. Below is an illustration example:
def build_model(x_obs=None):
x = zs.Normal('x', observed=a_obs)
return f(x)
# samples from the Normal distribution will be used for computation in f
out = build_model()
# x_obs will be used for computation in f
out = build_model(x_obs=x_obs)

For models that have many stochastic nodes (e.g., x1,...,xN), in principle the polymorphism can be achieved by
def build_model(x1_obs=None, ..., xN_obs=None):
...

However, in practice this is often found to be redundant and makes it hard for automatic
inference algorithms to fit each model. This is where the BayesianNet context makes a
big difference. BayesianNet accepts an argument named observed, which is a dictionary
mapping from names of StochasticTensors to their observations. Then the context will
be responsible for determining the states of StochasticTensors. This makes it easier
for users to handle state changes of large sets of stochastic nodes, and more importantly,
enables a unified form of model building functions that makes automatic inference possible.
An example is shown below.
def build_model(observed=None):
with zs.BayesianNet(observed=observed) as model:
z = zs.Normal('z')
x = zs.Normal('x', f(z))
...
return model
#
m
#
m

No observation
= build_model()
Observe z and x
= build_model({'z': z_obs, 'x': x_obs})

4.2 Inference Algorithms
In Section 2.2 we have covered the major inference methods used in Bayesian deep learning.
Although most of them are gradient-based, stochastic, and black-box (Ranganath et al.,
2014), which is suitable for building a general inference framework, there is currently no
software that provides a complete support of them. To bridging the gap, ZhuSuan leverages
recent advances of differentiable inference in Bayesian deep learning, and provides a wide
and flexible support for both traditional and modern inference algorithms. These algorithms
are provided in an API that fits well into deep learning paradigms, which makes inference
14

ZhuSuan: A Library for Bayesian Deep Learning

Objective

Gradient
estimator

Supported latent-variable types

Implementations
in zs.variational

• continuous and reparameterizable
ELBO

SGVB

• Concrete relaxation of discrete

elbo().
reinforce

REINFORCE all types

Importance
weighted
objective

KL(pkq)

SGVB
(IWAE)

elbo().sgvb

• continuous and reparameterizable
• Concrete relaxation of discrete

iw_objective().
sgvb

VIMCO

all types

iw_objective().
vimco

RWS

all types

klpq().rws

Table 1: Variational inference in ZhuSuan. Relevant references are SGVB (Kingma and
Welling, 2013), REINFORCE (Williams, 1992; Mnih and Gregor, 2014), IWAE (Burda
et al., 2015), VIMCO (Mnih and Rezende, 2016), and RWS (Bornschein and Bengio, 2014).

as easy as doing gradient descent in deterministic neural networks. Below, we use examples
to explain both variational inference and Monte Carlo methods in ZhuSuan.
4.2.1 Variational Inference
As reviewed in Section 2.2.1, a variational inference (VI) algorithm consists of two parts:
the construction of variational posteriors and the optimization of variational objectives.
Typical VI implementations in probabilistic programming languages have mostly restricted themselves to simple variational posteriors. For example, the ADVI algorithm
(Kucukelbir et al., 2017) that serves as the main VI method in Stan (Carpenter et al.,
2017) uses only Gaussian variational posteriors. In contrast, ZhuSuan supports building
very flexible variational posteriors by leveraging Bayesian networks. This opens up the
door for rich variational families with user-specified dependency structures.
As for the optimization side, as mentioned in Section 2.2.1, many gradient-based variational methods have emerged in recent progress of Bayesian deep learning (Kingma and
Welling, 2013; Mnih and Gregor, 2014; Burda et al., 2015; Mnih and Rezende, 2016). These
methods differ in the variational objectives and the gradient estimators they use. To make
them more automatic and easier to handle, ZhuSuan has wrapped them all into single functions, which computes a surrogate cost for users to directly take derivatives on. This means
that optimizing these surrogate costs is equally optimizing the corresponding variational
objectives using their well-developed gradient estimators. The currently supported varia15

Shi et al.

tional methods are summarized in Table 1. They are grouped by the objective they aim to
optimize. Currently there are three kinds of objectives supported: the ELBO (equivalent
to KL(qkp)), the importance weighted objective (Burda et al., 2015) and KL(pkq), where
p, q denotes the true and the variational posterior.
It is easy to program with ZhuSuan to implement a variational inference algorithm,
following these steps:
1) Build the variational posterior using ZhuSuan’s modeling primitives.
2) Get samples and their log probabilities from the variational posterior.
3) Provide the log joint function of the model and call variational objectives.
4) Choose the gradient estimator to use and get the surrogate cost to minimize.
5) Call Tensorflow optimizers to run gradient descent on the surrogate cost.
Below we use a simple example to guide the readers through all these steps.
Example 5 (BLR, continued) We consider applying variational inference to the BLR
model in Example 1, following the above steps:
1) Build the variational posterior. The true posterior of w is intractable and should have
correlations across dimensions. Here we follow the common practice to make
Q the mean-field
assumption that the variational posterior q(w) is factorized, i.e., q(w) = D
d=1 q(wd ), where
D is the dimension of the weights. The code for using factorized normal distribution as the
variational posterior is:
with zs.BayesianNet() as variational:
w_mean = tf.Variable(tf.zeros([D]))
w_logstd = tf.Variable(tf.zeros([D]))
w = zs.Normal('w', w_mean, logstd=w_logstd, group_ndims=1)

2) Get samples and their log probabilities from the variational posterior. This is by
querying the stochastic node w about its outputs and local log probabilities.
qw_samples, log_qw = variational.query(
'w', outputs=True, local_log_prob=True)

3) Provide the log joint function of the model and call variational objectives. The log
joint value can be computed easily by taking sum of log probabilities at individual nodes of
the model:
def log_joint(observed):
model = blr(observed, x)
log_py_xw, log_pw = model.local_log_prob(['y', 'w'])
return log_py_xw + log_pw

Then we call the ELBO objective, passing the log joint function, observed data and outputs
of the variational posterior.
lower_bound = zs.variational.elbo(log_joint,
observed={'y': y},
latent={'w': [qw_samples, log_qw]})
16

ZhuSuan: A Library for Bayesian Deep Learning

4) Choose the gradient estimator to use and get the surrogate cost to minimize. Because
w is continuous and can be reparameterized as w = σw + µw , where µw and σw are the
mean and the standard deviation of the normal distribution, we choose the SGVB gradient estimator (see Table 1 for the scope of application of each estimator). The computed
surrogate costs are for a batch of data, which is then averaged.
cost = tf.reduce_mean(elbo.sgvb())

5) Call Tensorflow optimizers to run gradient descent on the surrogate cost. As previously explained, this will be optimizing the ELBO objective with the SGVB gradient estimator. We can also fetch the ELBO value to verify it.
optimizer = tf.train.AdamOptimizer(learning_rate=0.001)
infer_op = optimizer.minimize(cost)
with tf.Session() as sess:
for i in range(iters):
_, elbo_value = sess.run([infer_op, elbo])

The above is a very simple example using factorized variational posteriors, while as
mentioned previously, ZhuSuan supports building very flexible variational posteriors. In
the following example, we will see how to do amortized inference (see Section 2.2.1) by
leveraging neural networks in the variational posterior.
Example 6 (VAE, continued) Consider the VAE model in Example 2. The key difference compared to BLR (Example 1) is that VAE’s latent variables (z) are local instead of
global. As the number of local latent variables scales linearly with the number of observations
(x), it is hard to afford a separate variational posterior for each of them. This is where the
amortized style of inference becomes useful. Specifically, we use a neural network with x
as input to generate parameters of the variational posterior for each corresponding z. This
network is often referred as decoder or recognition model in VAE. The code is shown in
Figure 6.

with zs.BayesianNet() as variational:
h = layers.fully_connected(x, n_h)
z_mean = layers.fully_connected(h, n_z,
activation_fn=None)
z_logstd = layers.fully_connected(h, n_z,
activation_fn=None)
z = zs.Normal('z', z_mean, logstd=z_logstd,
group_ndims=1)
(a)

(b)

Figure 6: VAE’s variational posterior: (a) graphical illustration; (b) programming on
ZhuSuan.
Having built the variational posterior, we can follow the above steps to finish the inference
procedure. Note that because the variational distribution for z is again normal, we can use
17

Shi et al.

the SGVB gradient estimator for optimizing the ELBO objective. We omit these steps
because they are very similar to those illustrated in the previous example.
In deep, hierarchical models, there are often conditional dependencies between latent
variables that need consideration during inference. We use the DSBN model from Example 3
to illustrate how to introduce structured dependencies when building variational posteriors
for hierarchical models. This example also demonstrates applying VI to discrete latent
variables in ZhuSuan.
Example 7 (DSBN, continued) A commonly used variational posterior in recent works
for the DSBN model is
q(z

(1:L)

)=

L
Y

q(z(l) |z(l−1) )q(z(1) |x).

(12)

l=2

which has the same conditional dependency structure with the original model given x is
observed. The graphical model as well as the code for this structured posterior (when L = 2)
are shown in Figure 7. It can be seen from the code that z(2) directly depends on z(1) by

with zs.BayesianNet() as variational:
z1_logits = layers.fully_connected(x, n_z,
activation_fn=None)
z1 = zs.Bernoulli('z1', z1_logits, dtype=tf.float32)
z2_logits = layers.fully_connected(z1, n_z,
activation_fn=None)
z2 = zs.Bernoulli('z2', z2_logits, dtype=tf.float32)
(b)

(a)

Figure 7: DSBN’s variational posterior: (a) graphical illustration; (b) programming on
ZhuSuan.
a fully connected layer in neural networks. Note that here z(2) and z(1) are discrete latent
variables. In ZhuSuan there are generally two ways for doing variational inference with
them.
One way is to use the Concrete relaxation, or Gumbel-softmax trick (Maddison et al.,
2016; Jang et al., 2016), namely, we use Concrete random variables to replace the Bernoulli
ones during training. Due to the Concrete distribution is continuous and reparameterizable,
we can keep on using the SGVB estimator. In test time, we switch back to the Bernoulli
random variables using the same input parameters. Both the model and the posterior in this
way are modified a little:
# model with Concrete relaxation
with zs.BayesianNet() as model:
z2_logits = tf.zeros([n, n_z])
18

ZhuSuan: A Library for Bayesian Deep Learning

if is_training:
z2 = zs.BinConcrete('z2', z2_logits, group_ndims=1)
else:
z2 = zs.Bernoulli('z2', z2_logits, group_ndims=1, dtype=tf.float32)
z1_logits = layers.fully_connected(z2, n_z, activation_fn=None)
if is_training:
z1 = zs.BinConcrete('z1', z1_logits, group_ndims=1)
else:
z1 = zs.Bernoulli('z1', z1_logits, group_ndims=1, dtype=tf.float32)
x_logits = layers.fully_connected(z1, n_x, activation_fn=None)
x = zs.Bernoulli('x', x_logits, group_ndims=1)
# posterior with Concrete relaxation
with zs.BayesianNet() as variational:
z1_logits = layers.fully_connected(x, n_z, activation_fn=None)
if is_training:
z1 = zs.BinConcrete('z1', z1_logits, group_ndims=1)
else:
z1 = zs.Bernoulli('z1', z1_logits, group_ndims=1, dtype=tf.float32)
z2_logits = layers.fully_connected(z1, n_z, activation_fn=None)
if is_training:
z2 = zs.BinConcrete('z2', z2_logits, group_ndims=1)
else:
z2 = zs.Bernoulli('z2', z2_logits, group_ndims=1, dtype=tf.float32)

The other approach is to directly use a gradient estimator that applies to discrete latent
variables. This includes REINFORCE and VIMCO. Here we present the code for using
VIMCO. Because VIMCO is optimizing the importance weighted bound, which is a multisample bound, the model and the variational posterior need to be changed to multi-sample
versions as shown below.
# model (multi-sample version)
with zs.BayesianNet() as model:
z2_logits = tf.zeros([n, n_z])
z2 = zs.Bernoulli('z2', z2_logits, group_ndims=1, dtype=tf.float32,
n_samples=n_samples)
z1_logits = layers.fully_connected(z2, n_z, activation_fn=None)
z1 = zs.Bernoulli('z1', z1_logits, group_ndims=1, dtype=tf.float32)
x_logits = layers.fully_connected(z1, n_x, activation_fn=None)
x = zs.Bernoulli('x', x_logits, group_ndims=1)
# posterior (multi-sample version)
with zs.BayesianNet() as variational:
z1_logits = layers.fully_connected(x, n_z, activation_fn=None)
z1 = zs.Bernoulli('z1', z1_logits, group_ndims=1, dtype=tf.float32,
n_samples=n_samples)
z2_logits = layers.fully_connected(z1, n_z, activation_fn=None)
z2 = zs.Bernoulli('z2', z2_logits, dtype=tf.float32)

Note the n_samples argument added in the code. This time we call the VIMCO gradient
estimator to get the surrogate cost:
19

Shi et al.

lower_bound = zs.variational.iw_objective(
log_joint, observed={'x': x}, latent=qz_outputs, axis=0)
cost = tf.reduce_mean(lower_bound.vimco())

All samples along the axis 0 will be used to obtain variance reduced gradient estimates.
4.2.2 Monte Carlo Methods
In addition to variational inference, efforts have also been made towards unifying Monte
Carlo methods in the Bayesian deep learning context. We have covered the basics in Section 2.2.2. Below we walk through ZhuSuan’s support for both importance sampling and a
black-box MCMC method: Hamiltonian Monte Carlo.
Importance Sampling As reviewed in Section 2.2.2, the major application scenarios
for importance sampling in Bayesian deep learning include model learning and evaluation.
For model learning, we have introduced that self-normalized importance sampling can be
used to estimate the gradients of marginal log likelihoods with respect to model parameters
(Bornschein and Bengio, 2014). It turns out that the gradients estimated in this way are
equivalent to exact gradients of the importance weighted variational objective (mentioned
in Section 4.2.1) with respect to model parameters. So the model learning procedure with
importance sampling can be implemented by directly optimizing the importance weighted
objective (zs.variational.iw_objective(), see Table 1) with respect to the model parameters:
lower_bound = zs.variational.iw_objective(
log_joint, observed, latent, axis=axis)
optimizer = tf.train.AdamOptimizer(learning_rate)
learning_op = optimizer.optimize(lower_bound, var_list=model_params)

The only difference is that the meaning of the variational posterior has changed to the
proposal distribution.
Besides learning model parameters, we have mentioned that importance sampling is
extensively used for model evaluation. Towards this need, ZhuSuan’s evaluation module
provides an is_loglikelihood() function for estimating the marginal log likelihoods of
any given observations using simple importance sampling:
ll = zs.is_loglikelihood(log_joint, observed, latent, axis=axis)

Until now we haven’t covered how to build a proposal distribution for importance sampling. In fact, this procedure is exactly the same as that of building a variational posterior.
With ZhuSuan’s modeling primitives, neural adaptive proposals mentioned in Section 2.2.2
can be easy to implement by leveraging neural networks in the proposal distribution. Adapting the proposal for the above two scenarios is also straightforward, where the true posterior
distribution is often an ideal choice. So the adaptation turns out a variational inference
problem, which can be solved by choosing an appropriate method in Table 1. Specially,
when the KL(pkq) objective is chosen and the gradient estimation is by importance sampling, this recovers the method used in Reweighted Wake-Sleep (RWS) (Bornschein and
Bengio, 2014), which is why this estimator is named rws in ZhuSuan. We illustrate the
process of training DSBN by importance sampling in the following example.
20

ZhuSuan: A Library for Bayesian Deep Learning

Example 8 (DSBN, continued) In this example we reproduce the algorithm proposed
in the Reweighted Wake-Sleep paper (Bornschein and Bengio, 2014), which learns DSBN
by importance sampling and a neural adaptive proposal. The code snippet below follows
from the multi-sample version of the model and the posterior (now used as the proposal) in
Example 7, and we omit the code to draw samples and compute their log probabilities from
the proposal distribution.
# learning model parameters
lower_bound = tf.reduce_mean(
zs.variational.importance_weighted_objective(
log_joint, observed={'x': x_obs}, latent=latent, axis=0))
model_grads = optimizer.compute_gradients(-lower_bound, model_params)
# adapting the proposal
klpq_obj = zs.variational.klpq(
log_joint, observed={'x': x_obs}, latent=latent, axis=0)
klpq_cost = tf.reduce_mean(klpq_obj.rws())
klpq_grads = optimizer.compute_gradients(klpq_cost, proposal_params)
infer_op = optimizer.apply_gradients(model_grads + klpq_grads)

We can see that the code clearly has two parts. The first part is for learning model parameters, by optimizing the importance weighted objective with respect to model parameters. The
second part is for adapting the proposal, by minimizing the inclusive KL divergence between
the true posterior (the ideal proposal) and the current proposal. As the training proceeds, the
adaptation part will help maintain a good proposal, reducing the variance of the marginal
log likelihood estimate by importance sampling.
Hamiltonian Monte Carlo In Section 2.2.2 we have briefly analyzed existing MCMC
methods and have identified HMC as a powerful tool to address the posterior inference problem in high-dimensional spaces and non-conjugate models, which is a perfect fit for Bayesian
deep learning. However, in practice this algorithm involves lots of technical details and can
be hard to implement in a fast and efficient way. Besides, despite all tuning parameters in
HMC have clear physical meanings, it is still hard for users to tune them by hand because
the optimal choice always depends on unknown statistics of the underlying distribution.
For example, the mass matrix describes the variance of the underlying distribution, which
is hard to know before we can draw samples from it.
In recent years there has been a rise of practical algorithms and high-performance softwares that target at these problems. The No-U-Turn Sampler, or NUTS (Hoffman and
Gelman, 2014) proposes to automatically determine the number of leapfrog steps. It also
comes along with a dual averaging scheme for automatically tuning the step size. The HMC
implementation in Stan (Carpenter et al., 2017) also includes a procedure that estimates
the mass matrix from the samples drawn in the warm-up stage.
ZhuSuan has learned from the existing libraries and provides a fast, automatic, and
deep-learning style implementation of HMC. Specifically, ZhuSuan’s HMC has the following
features:
• Support running multiple chains in parallel on CPUs or GPUs.
21

Shi et al.

• Provide options for automatically tuning parameters, including the step size and the
mass matrix. The NUTS algorithm for determining leapfrog steps is not included
because it’s a recursive algorithm and each separate chain can have different leapfrog
steps, thus hard to have a parallel implementation in static computation graphs.
• The algorithm is provided in an API very similar to Tensorflow optimizers, which is
illustrated in Figure 8. We hope it is easy to start with for people who are familiar
with deep learning libraries.
z = tf.Variable(0.)

z = tf.Variable(0.)

hmc = zs.HMC(step_size=1e-3,
n_leapfrogs=10)

optimizer = tf.train.AdamOptimizer(
learning_rate=1e-3)

sample_op, hmc_info = hmc.sample(
log_joint, observed={'x': x},
latent={'z': z})

optimize_op = optimizer.minimize(
cost(z))

with tf.Session() as sess:
for i in range(iters):
_ = sess.run(sample_op)

with tf.Session() as sess:
for i in range(iters):
_ = sess.run(optimize_op)

(a) Using HMC in ZhuSuan

(b) Using Tensorflow optimizers

Figure 8: ZhuSuan’s HMC vs. Tensorflow optimizers

5. Comparison
We compare ZhuSuan with the representatives of python probabilistic programming libraries, namely PyMC3 (Salvatier et al., 2016) and Edward (Tran et al., 2016). A detailed
feature comparison is shown in Table 2. As seen from the table, all the three libraries
build upon modern computation graph libraries and can transparently run on GPUs. The
other features can be mainly divided into three categories: modeling capabilities, inference
support as well as architecture and design.
Modeling capabilities All the three libraries use the primitives from the computation
graph toolkits they base on and are designed for directed graphs (Bayesian networks).
PyMC3 solves the model reuse problem during inference by theano.copy() to copy the
related subgraphs. Edward also does this but does not rely on the official Tensorflow
API. ZhuSuan avoids altering created graphs and build reuse on purely function reuse and
context management. As a result, PyMC3 and ZhuSuan can correctly deal with control
flow primitives like theano.scan() and tf.while_loop(), while Edward faces challenges
when applying variational inference to this kind of models because it will require graph
copying when replacing latent variables with samples from the variational posterior, which
is troublesome for control flow operations given there is no official support.
Inference support As for the range of inference methods, the three libraries have
different emphases. PyMC3 started from MCMC methods, which is demonstrated by its
22

ZhuSuan: A Library for Bayesian Deep Learning

Features

PyMC3

Edward

ZhuSuan

Based on

Theano

Tensorflow

Tensorflow

GPU support

3

3

3

Deterministic
modeling

Any Theano operation

Control flows (tf.
while_loop,tf.cond)
are not properly handled

Any Tensorflow operation

Customizable
variational
posterior

Only for reparameterizable settings

3

3

HMC: Parallel chains

3

7

3

Modularity

Modeling and inference
are tightly coupled

Modeling and inference
are tightly coupled

All parts can be
used independently

Transparency

7

7

3

Table 2: Comparison between ZhuSuan and other python probabilistic programming libraries.

name. It has put much efforts into sampling algorithms and have made them applicable
to a broad class of models in Bayesian statistics. On the other hand, PyMC3’s support
for variational inference is limited and specific to several independent algorithms (ADVI,
SVGD, etc.). Edward has a general inference abstraction and implements a large number of
algorithms as subclasses. However, as the form of abstraction induces constraints, many of
these implemented algorithms have limited behaviors and make too strong assumptions on
the constructed models (e.g., GANInference). ZhuSuan emphasizes more on the scenarios
from Bayesian deep learning thus puts more efforts into modern differentiable algorithms
that can be unified into the deep learning paradigm. Both Edward and ZhuSuan support
customizable variational posteriors for all the VI methods while PyMC3 has only made this
specific to the reparameterizable settings.
Architecture and design The design philosophy of ZhuSuan has been introduced
in Section 3, which emphasizes two principles—modularity and transparency. These two
principles are very different from the traditional designs of probabilistic programming, for
which PyMC3 can be a stereotype. The PyMC3 programs strictly follow the form with
model definition, inference function call, and predictive checks. Little customization can be
made except in the model definition and provided inference options. Edward also roughly
follows this framework but has made it more flexible by general programmable variational
posteriors. However, the inference procedure is also hidden due to its reliance on lots of
internal manipulations of the created computation graphs, which can be hard for plain
users to understand. Both the libraries build their modeling and inference functionalities
in a tightly coupled way. This implies that if the model cannot be described using their
23

Shi et al.

modeling primitives, then there is little possibility using their inference features. ZhuSuan
draws benefits from deep learning paradigms and makes inference just as doing gradient
descent on a cost term. And all parts in the library can be used independently. This proves
to be much more transparent and compositional, especially in models with a large set of
stochastic activities, various observation behaviors and deep hierarchies.

6. Conclusions
We have described ZhuSuan, a python probabilistic programming library for Bayesian deep
learning, built upon Tensorflow. ZhuSuan bridges the gap between Bayesian methods and
deep learning by providing deep learning style primitives and algorithms for building probabilistic models and applying Bayesian inference. Many examples are provided to illustrate the intuitiveness of programming with ZhuSuan. We have open-sourced ZhuSuan on
GitHub5 , aiming to accelerate research and applications of Bayesian deep learning.

Acknowledgments
We thank Haowen Xu for helpful discussions on the API design, Chang Liu for advices on
paper writing and Shizhen Xu, Dong Yan for initial attempts to try ZhuSuan on multi-card
and distributed systems. We would like to acknowledge support for this project from the National 973 Basic Research Program of China (No. 2013CB329403), National NSF of China
(Nos. 61620106010, 61322308, 61332007), the National Youth Top-notch Talents Support
Program, the NVIDIA NVAIL program, and Tsinghua Tiangong Institite for Intelligent
Technology.

References
Martı́n Abadi, Ashish Agarwal, Paul Barham, Eugene Brevdo, Zhifeng Chen, Craig Citro,
Greg S Corrado, Andy Davis, Jeffrey Dean, Matthieu Devin, et al. Tensorflow: Large-scale
machine learning on heterogeneous distributed systems. arXiv preprint arXiv:1603.04467,
2016.
Rami Al-Rfou, Guillaume Alain, Amjad Almahairi, Christof Angermueller, Dzmitry Bahdanau, Nicolas Ballas, et al. Theano: A Python framework for fast computation of
mathematical expressions. arXiv e-prints, abs/1605.02688, May 2016. URL https:
//arxiv.org/abs/1605.02688.
Jimmy Ba, Ruslan R Salakhutdinov, Roger B Grosse, and Brendan J Frey. Learning wakesleep recurrent attention models. In Advances in Neural Information Processing Systems,
pages 2593–2601, 2015.
Michael Betancourt. A conceptual introduction to hamiltonian monte carlo. arXiv preprint
arXiv:1701.02434, 2017.
5. GitHub address: https://github.com/thu-ml/zhusuan

24

ZhuSuan: A Library for Bayesian Deep Learning

Mariusz Bojarski, Davide Del Testa, Daniel Dworakowski, Bernhard Firner, Beat Flepp,
Prasoon Goyal, Lawrence D Jackel, Mathew Monfort, Urs Muller, Jiakai Zhang, et al.
End to end learning for self-driving cars. arXiv preprint arXiv:1604.07316, 2016.
Jörg Bornschein and Yoshua Bengio.
arXiv:1406.2751, 2014.

Reweighted wake-sleep.

arXiv preprint

Yuri Burda, Roger Grosse, and Ruslan Salakhutdinov. Importance weighted autoencoders.
arXiv preprint arXiv:1509.00519, 2015.
Olivier Cappé, Randal Douc, Arnaud Guillin, Jean-Michel Marin, and Christian P Robert.
Adaptive importance sampling in general mixture classes. Statistics and Computing, 18
(4):447–459, 2008.
Bob Carpenter, Andrew Gelman, Matthew Hoffman, Daniel Lee, Ben Goodrich, Michael
Betancourt, Marcus Brubaker, Jiqiang Guo, Peter Li, and Allen Riddell. Stan: A probabilistic programming language. Journal of Statistical Software, Articles, 76(1):1–32, 2017.
ISSN 1548-7660. doi: 10.18637/jss.v076.i01. URL https://www.jstatsoft.org/v076/
i01.
Jean Cornuet, JEAN-MICHEL MARIN, Antonietta Mira, and Christian P Robert. Adaptive multiple importance sampling. Scandinavian Journal of Statistics, 39(4):798–812,
2012.
Sander Dieleman, Jan Schlter, Colin Raffel, Eben Olson, Sren Kaae Snderby, Daniel Nouri,
et al. Lasagne: First release., August 2015. URL http://dx.doi.org/10.5281/zenodo.
27878.
C. Du, J. Zhu, and B. Zhang. Learning deep generative models with doubly stochastic
gradient mcmc. IEEE Transactions on Neural Networks and Learning Systems, PP(99):
1–13, 2017. ISSN 2162-237X. doi: 10.1109/TNNLS.2017.2688499.
Yarin Gal. Uncertainty in Deep Learning. PhD thesis, University of Cambridge, 2016.
Andrew Gelman, John B Carlin, Hal S Stern, and Donald B Rubin. Bayesian data analysis,
volume 2. Chapman & Hall/CRC Boca Raton, FL, USA, 2014.
Stuart Geman and Donald Geman. Stochastic relaxation, gibbs distributions, and the
bayesian restoration of images. IEEE Transactions on Pattern Analysis and Machine
Intelligence, (6):721–741, 1984.
Zoubin Ghahramani. Probabilistic machine learning and artificial intelligence. Nature, 521
(7553):452–459, 2015.
Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil
Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial nets. In Advances in
Neural Information Processing Systems, pages 2672–2680, 2014.
Alex Graves. Practical variational inference for neural networks. In Advances in Neural
Information Processing Systems, pages 2348–2356, 2011.
25

Shi et al.

Shixiang Gu, Zoubin Ghahramani, and Richard E Turner. Neural adaptive sequential monte
carlo. In Advances in Neural Information Processing Systems, pages 2629–2637, 2015.
W Keith Hastings. Monte carlo sampling methods using markov chains and their applications. Biometrika, 57(1):97–109, 1970.
Geoffrey Hinton, Li Deng, Dong Yu, George E Dahl, Abdel-rahman Mohamed, Navdeep
Jaitly, Andrew Senior, Vincent Vanhoucke, Patrick Nguyen, Tara N Sainath, et al. Deep
neural networks for acoustic modeling in speech recognition: The shared views of four
research groups. IEEE Signal Processing Magazine, 29(6):82–97, 2012.
Geoffrey E Hinton, Simon Osindero, and Yee-Whye Teh. A fast learning algorithm for deep
belief nets. Neural computation, 18(7):1527–1554, 2006.
Matthew D. Hoffman. Learning deep latent Gaussian models with Markov chain Monte
Carlo. In Proceedings of the 34th International Conference on Machine Learning, pages
1510–1519, 2017.
Matthew D Hoffman and Andrew Gelman. The no-u-turn sampler: adaptively setting
path lengths in hamiltonian monte carlo. Journal of Machine Learning Research, 15(1):
1593–1623, 2014.
Matthew D Hoffman, David M Blei, Chong Wang, and John William Paisley. Stochastic
variational inference. Journal of Machine Learning Research, 14(1):1303–1347, 2013.
Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterization with gumbelsoftmax. arXiv preprint arXiv:1611.01144, 2016.
Matthew Johnson, David K Duvenaud, Alex Wiltschko, Ryan P Adams, and Sandeep R
Datta. Composing graphical models with neural networks for structured representations
and fast inference. In Advances in Neural Information Processing Systems, pages 2946–
2954, 2016.
Diederik P Kingma and Max Welling. Auto-encoding variational bayes. arXiv preprint
arXiv:1312.6114, 2013.
Diederik P Kingma, Shakir Mohamed, Danilo Jimenez Rezende, and Max Welling. Semisupervised learning with deep generative models. In Advances in Neural Information
Processing Systems, pages 3581–3589, 2014.
Daphne Koller and Nir Friedman. Probabilistic graphical models: principles and techniques.
2009.
Alex Krizhevsky, Ilya Sutskever, and Geoffrey E Hinton. Imagenet classification with deep
convolutional neural networks. In Advances in Neural Information Processing Systems,
pages 1097–1105, 2012.
Alp Kucukelbir, Dustin Tran, Rajesh Ranganath, Andrew Gelman, and David M. Blei.
Automatic differentiation variational inference. Journal of Machine Learning Research,
18(14):1–45, 2017. URL http://jmlr.org/papers/v18/16-107.html.
26

ZhuSuan: A Library for Bayesian Deep Learning

Brenden M Lake, Ruslan Salakhutdinov, and Joshua B Tenenbaum. Human-level concept
learning through probabilistic program induction. Science, 350(6266):1332–1338, 2015.
Yann LeCun, Yoshua Bengio, and Geoffrey Hinton. Deep learning. Nature, 521(7553):
436–444, 2015.
Benedict Leimkuhler and Sebastian Reich. Simulating hamiltonian dynamics, volume 14.
Cambridge university press, 2004.
Yingzhen Li and Yarin Gal. Dropout inference in bayesian neural networks with alphadivergences. arXiv preprint arXiv:1703.02914, 2017.
Yingzhen Li, Richard E Turner, and Qiang Liu. Approximate inference with amortised
mcmc. arXiv preprint arXiv:1702.08343, 2017.
Chris J Maddison, Andriy Mnih, and Yee Whye Teh. The concrete distribution: A continuous relaxation of discrete random variables. arXiv preprint arXiv:1611.00712, 2016.
Shike Mei, Jun Zhu, and Jerry Zhu. Robust regbayes: Selectively incorporating first-order
logic domain knowledge into bayesian models. In Proceedings of the 31st International
Conference on Machine Learning (ICML-14), pages 253–261, 2014.
Nicholas Metropolis, Arianna W Rosenbluth, Marshall N Rosenbluth, Augusta H Teller,
and Edward Teller. Equation of state calculations by fast computing machines. The
journal of chemical physics, 21(6):1087–1092, 1953.
Riccardo Miotto, Fei Wang, Shuang Wang, Xiaoqian Jiang, and Joel T Dudley. Deep
learning for healthcare: review, opportunities and challenges. Briefings in Bioinformatics,
page bbx044, 2017.
Andriy Mnih and Karol Gregor. Neural variational inference and learning in belief networks.
In Proceedings of the 31st International Conference on Machine Learning, 2014.
Andriy Mnih and Danilo J Rezende. Variational inference for monte carlo objectives. arXiv
preprint arXiv:1602.06725, 2016.
Radford M Neal. Connectionist learning of belief networks. Artificial intelligence, 56(1):
71–113, 1992.
Radford M Neal. Bayesian learning for neural networks. PhD thesis, University of Toronto,
1995.
Radford M Neal et al. Mcmc using hamiltonian dynamics. Handbook of Markov Chain
Monte Carlo, 2(11), 2011.
Anh Nguyen, Jason Yosinski, and Jeff Clune. Deep neural networks are easily fooled: High
confidence predictions for unrecognizable images. In Proceedings of the IEEE Conference
on Computer Vision and Pattern Recognition, pages 427–436, 2015.
Art B. Owen. Monte Carlo theory, methods and examples. 2013.
27

Shi et al.

Brooks Paige and Frank Wood. Inference networks for sequential monte carlo in graphical
models. In International Conference on Machine Learning, pages 3040–3049, 2016.
John Paisley, David Blei, and Michael Jordan. Variational bayesian inference with stochastic
search. arXiv preprint arXiv:1206.6430, 2012.
Tianyu Pang, Chao Du, and Jun Zhu. Robust deep learning via reverse cross-entropy
training and thresholding test. arXiv preprint arXiv:1706.00633, 2017.
Alec Radford, Luke Metz, and Soumith Chintala. Unsupervised representation learning with
deep convolutional generative adversarial networks. arXiv preprint arXiv:1511.06434,
2015.
Rajesh Ranganath, Sean Gerrish, and David Blei. Black box variational inference. In
Artificial Intelligence and Statistics, pages 814–822, 2014.
Danilo Jimenez Rezende and Shakir Mohamed. Variational inference with normalizing flows.
arXiv preprint arXiv:1505.05770, 2015.
Danilo Jimenez Rezende, Shakir Mohamed, Ivo Danihelka, Karol Gregor, and Daan Wierstra. One-shot generalization in deep generative models. arXiv preprint arXiv:1603.05106,
2016.
Christian P. Robert and George Casella. Monte Carlo Statistical Methods (Springer Texts in
Statistics). Springer-Verlag New York, Inc., Secaucus, NJ, USA, 2005. ISBN 0387212396.
Tim Salimans, Diederik Kingma, and Max Welling. Markov chain monte carlo and variational inference: Bridging the gap. In Proceedings of the 32nd International Conference
on Machine Learning (ICML-15), pages 1218–1226, 2015.
Tim Salimans, Ian Goodfellow, Wojciech Zaremba, Vicki Cheung, Alec Radford, and
Xi Chen. Improved techniques for training gans. In Advances in Neural Information
Processing Systems, pages 2226–2234, 2016.
John Salvatier, Thomas V Wiecki, and Christopher Fonnesbeck. Probabilistic programming
in python using pymc3. PeerJ Computer Science, 2:e55, 2016.
David Silver, Aja Huang, Chris J Maddison, Arthur Guez, Laurent Sifre, George Van
Den Driessche, Julian Schrittwieser, Ioannis Antonoglou, Veda Panneershelvam, Marc
Lanctot, et al. Mastering the game of go with deep neural networks and tree search.
Nature, 529(7587):484–489, 2016.
Ilya Sutskever, Oriol Vinyals, and Quoc V Le. Sequence to sequence learning with neural
networks. In Advances in Neural Information Processing Systems, pages 3104–3112, 2014.
Christian Szegedy, Wojciech Zaremba, Ilya Sutskever, Joan Bruna, Dumitru Erhan, Ian
Goodfellow, and Rob Fergus. Intriguing properties of neural networks. arXiv preprint
arXiv:1312.6199, 2013.
28

ZhuSuan: A Library for Bayesian Deep Learning

Michalis Titsias and Miguel Lázaro-Gredilla. Doubly stochastic variational bayes for nonconjugate inference. In Proceedings of the 31st International Conference on Machine
Learning (ICML-14), pages 1971–1979, 2014.
Michalis K Titsias. Learning model reparametrizations: Implicit variational inference by
fitting mcmc distributions. arXiv preprint arXiv:1708.01529, 2017.
Dustin Tran, Alp Kucukelbir, Adji B. Dieng, Maja Rudolph, Dawen Liang, and David M.
Blei. Edward: A library for probabilistic modeling, inference, and criticism. arXiv
preprint arXiv:1610.09787, 2016.
Hanna M Wallach, Iain Murray, Ruslan Salakhutdinov, and David Mimno. Evaluation
methods for topic models. In Proceedings of the 26th International Conference on Machine
Learning, pages 1105–1112. ACM, 2009.
Greg CG Wei and Martin A Tanner. A monte carlo implementation of the em algorithm
and the poor man’s data augmentation algorithms. Journal of the American statistical
Association, 85(411):699–704, 1990.
Ronald J Williams. Simple statistical gradient-following algorithms for connectionist reinforcement learning. Machine Learning, 8(3-4):229–256, 1992.
Jun Zhu, Ning Chen, and Eric P Xing. Bayesian inference with posterior regularization
and applications to infinite latent svms. Journal of Machine Learning Research, 15:1799,
2014.
Jun Zhu, Jianfei Chen, Wenbo Hu, and Bo Zhang. Big learning with bayesian methods.
National Science Review, page nwx044, 2017.

29

