ZhuSuan: Bayesian Deep Learning Library

• ZhuSuan is a Python probabilistic programming library that integrates Bayesian inference with deep learning, utilizing TensorFlow's graph execution and GPU support.
• It supports complex models such as hierarchical Bayesian models, latent-variable models, and Bayesian neural architectures through its modular BayesianNet interface.
• The library enables scalable variational inference, importance sampling, and Hamiltonian Monte Carlo for efficient uncertainty quantification in deep learning applications.

ZhuSuan is a Python probabilistic programming library for Bayesian deep learning, built on top of TensorFlow and designed to make it practical to define probabilistic models with the expressiveness of deep learning while carrying out Bayesian inference with scalable modern algorithms. Its stated role is to join the complementary advantages of Bayesian methods and deep neural networks: neural networks supply highly expressive deterministic transformations, whereas Bayesian methods provide principled uncertainty modeling over latent variables, parameters, predictions, and decisions. In that framing, ZhuSuan targets model classes that conventional deep learning libraries do not naturally support, including hierarchical Bayesian models, latent-variable models, deep generative models, and Bayesian neural architectures such as Bayesian recurrent neural networks (Shi et al., 2017).

1. Conceptual setting within Bayesian deep learning

The motivating premise is that standard deep learning toolkits are optimized for deterministic neural networks and mostly supervised learning, so they are not well aligned with models in which some variables are random, some are observed conditionally, and posterior inference is intractable. ZhuSuan is positioned at the intersection of probabilistic programming and Bayesian deep learning: it is a probabilistic programming library, but one explicitly organized around the needs of BDL rather than classical Bayesian statistics alone (Shi et al., 2017).

The modeling framework centers on Bayesian networks, that is, directed acyclic graphical models whose joint distribution factorizes into priors at root nodes and conditional distributions at child nodes. The paper illustrates this with

$$p(Z, H, U, S, A, N)=p(Z)p(H)p(U\mid Z)p(S\mid Z,H)p(A\mid S)p(N\mid U,A).$$

This factorized representation is paired with a latent-variable perspective in which observations $x$ are only partially explained by latent variables $z$:

$$p(z,x)=p(z)p(x\mid z), \qquad p(z\mid x)=\frac{p(z,x)}{p(x)}=\frac{p(z)p(x\mid z)}{p(x)}.$$

The paper explicitly notes that $z$ may include both global latent variables and local latent variables.

Because ZhuSuan is built on TensorFlow, it inherits computation graphs, automatic differentiation, GPU support, multi-GPU training, neural network layers, and control-flow operators such as tf.while_loop and tf.scan. The intended consequence is that deterministic computation remains ordinary TensorFlow, while probabilistic structure and inference are layered on top. This suggests that ZhuSuan’s central design commitment is not a special-purpose modeling DSL, but rather the direct coexistence of TensorFlow operations with probabilistic nodes inside a single programming model.

2. Probabilistic programming interface

ZhuSuan’s core modeling abstraction is the BayesianNet context. Inside a block of the form with zs.BayesianNet() as model:, stochastic nodes are introduced through constructors such as zs.Normal, zs.Bernoulli, and zs.Categorical, whereas deterministic transformations are represented by ordinary TensorFlow tensors and operations. Stochastic nodes are encapsulated by the StochasticTensor abstraction, while deterministic nodes are simply TensorFlow objects (Shi et al., 2017).

A central design point is the distinction between observed and latent states. Rather than rewriting or copying computation graphs when a random variable changes role, ZhuSuan allows an observed dictionary to be passed into BayesianNet, mapping node names to observed values. The same model-building function can therefore be reused with different subsets of variables clamped. The paper treats this as a major advantage for model reuse and compositionality, especially relative to graph-copying approaches.

The BayesianNet maintains named stochastic nodes and supports querying outputs and local log probabilities. This arrangement makes model definition explicit without tying it to a single inference procedure. The paper presents that separation as architectural rather than incidental: modeling and inference are intended to be modular, transparent, and individually customizable.

3. Variational objectives and gradient estimators

Posterior inference is treated as generally intractable in Bayesian deep learning, particularly when flexible neural networks appear in the generative model or recognition model. ZhuSuan therefore supports variational inference with multiple objectives and stochastic gradient estimators. The central variational target is to approximate $p(z\mid x)$ with a tractable family $q_\phi(z)$, typically by maximizing an evidence lower bound:

$$\mathcal{L}(x;\phi)=\log p(x)-\mathrm{KL}\bigl(q_\phi(z)\,\|\,p(z\mid x)\bigr) =\mathbb{E}_{q_\phi(z)}[\log p(x\mid z)]-\mathrm{KL}\bigl(q_\phi(z)\,\|\,p(z)\bigr),$$

or, equivalently,

$$\mathcal{L}(q)=\mathbb{E}_{q(\mathbf{z})}\bigl[\log p(\mathbf{x},\mathbf{z})-\log q(\mathbf{z})\bigr].$$

The lower-bound relation

$\log p(x)\ge \mathcal{L}(x;\phi)$

is the organizing principle behind the variational interface (Shi et al., 2017).

Objective	Estimator	Supported latent-variable types
ELBO	SGVB	Continuous reparameterizable latents; discrete variables under Concrete/Gumbel-softmax relaxation
ELBO	REINFORCE	All latent-variable types
Importance weighted objective	SGVB (IWAE)	Continuous reparameterizable variables; Concrete relaxations
Importance weighted objective	VIMCO	All latent-variable types
$x$0	RWS	All latent-variable types

These methods are exposed through zs.variational.elbo(), zs.variational.iw_objective(), and zs.variational.klpq(), followed by estimator-specific methods such as .sgvb(), .reinforce(), .vimco(), or .rws(). ZhuSuan computes a surrogate cost whose gradient corresponds to the desired objective under the selected estimator, after which ordinary TensorFlow optimizers are used. In the paper’s formulation, inference is therefore made to look like deep-learning optimization.

The examples also show that the variational family need not be restricted to mean-field form. In the deep sigmoid belief network example, ZhuSuan uses the same BayesianNet abstraction to express a structured posterior

$x$1

which mirrors the hierarchy of the generative model. For discrete latent variables, the paper highlights two routes: continuous relaxation via Concrete/Gumbel-softmax to enable SGVB, or unbiased/discrete-variable estimators such as REINFORCE and VIMCO.

4. Importance sampling, evaluation, and Hamiltonian Monte Carlo

ZhuSuan also supports Monte Carlo methods, both for learning and for evaluation. The paper reviews importance sampling through the identity

$x$2

with estimator

$x$3

and the self-normalized form

$x$4

Within ZhuSuan, importance sampling is used both for marginal likelihood evaluation and for learning procedures such as Reweighted Wake-Sleep. The helper is_loglikelihood() is provided for marginal log-likelihood estimation, while klpq().rws() supports minimizing the inclusive divergence $x$5 to adapt a proposal distribution (Shi et al., 2017).

For MCMC, the library focuses on Hamiltonian Monte Carlo. The paper motivates HMC as suitable for high-dimensional, continuous, non-conjugate posteriors common in Bayesian deep learning. HMC augments the target with a momentum variable:

$x$6

with Hamiltonian $x$7 and dynamics

$x$8

The paper also records the detailed-balance relation

$x$9

as the general stationary condition for a Markov kernel.

ZhuSuan’s HMC interface is intentionally presented in a deep-learning style. An HMC object is created, .sample(log_joint, observed=..., latent=...) is called, and the returned sampling operation is executed repeatedly in a TensorFlow session. The implementation supports parallel chains on CPU or GPU and automatic tuning of step size and mass matrix. A stated limitation is that NUTS is not included, because its recursive, per-chain adaptive control flow is difficult to parallelize in static computation graphs.

5. Representative models and case studies

The library’s scope is illustrated by four running examples that span several BDL regimes (Shi et al., 2017).

Bayesian logistic regression is used as the simplest case:

$z$0

The paper uses this model to demonstrate mean-field variational inference, with

$z$1

constructed as a BayesianNet, followed by log_joint, zs.variational.elbo(...), sgvb(), and optimization with Adam.

The variational autoencoder serves as the canonical deep latent-variable model, combining a deep neural decoder with amortized inference. Although the paper does not emphasize the formula in the code example, the posterior implemented is of the form

$z$2

Because the latent variable is continuous and Gaussian, SGVB is used.

The deep sigmoid belief network demonstrates multiple layers of discrete latent variables:

$z$3

Its role in the paper is to exhibit structured variational posteriors, discrete-latent inference, and the use of importance-weighted multi-sample objectives with VIMCO.

The Bayesian recurrent neural network example is a Bayesian LSTM-style sequence classifier with random weights:

$z$4

This example is significant because it shows that stochastic parameters can be embedded inside sophisticated sequence architectures using TensorFlow control flow. The paper presents that capability as evidence that ZhuSuan is not confined to shallow Bayesian models or feedforward latent-variable models.

6. Architecture, workflow, and position in the software landscape

The practical workflow described in the paper is stable across examples: define a model-building function with BayesianNet; specify observations through the observed dictionary; define a variational posterior or proposal, again as a BayesianNet, if required; query samples and local log probabilities; write a log_joint function; construct an objective such as ELBO or an importance-weighted bound; choose a gradient estimator; and then use standard TensorFlow optimization or sampling loops. Core constructs are zs.BayesianNet, zs.Normal, zs.Bernoulli, zs.Categorical, the zs.variational objectives, is_loglikelihood(), and zs.HMC (Shi et al., 2017).

The paper characterizes the library by four architectural principles. First, modeling and inference are modular. Second, the interface favors transparency rather than full automation, so users can inspect and customize model structure, posterior samples, and probability terms. Third, deterministic computation reuses deep-learning primitives from TensorFlow rather than reimplementing them. Fourth, the observed mechanism enables compositionality and model reuse without graph surgery.

Relative to prior systems, the paper compares ZhuSuan mainly with PyMC3 and Edward. Its claim is not maximal method coverage, but particular suitability for Bayesian deep learning. Against PyMC3, it is described as offering more support for customizable variational posteriors beyond reparameterizable settings and as being more aligned with modern differentiable inference. Against Edward, it is argued to avoid inference-time graph copying and therefore to be more robust for TensorFlow control-flow constructs such as tf.while_loop and tf.scan.

The paper also states several limitations. ZhuSuan focuses on Bayesian networks rather than undirected graphical models. Its HMC implementation omits NUTS because of static-graph constraints. It does not claim full automation of inference, and its method set is centered on the Bayesian deep learning landscape of ELBO-based variational inference, IWAE/VIMCO/RWS-style methods, importance sampling, and HMC. A common misconception would therefore be to read ZhuSuan as a universal probabilistic programming environment or as an automated inference engine; the paper instead presents it as a flexible, transparent, BDL-oriented layer that preserves the programming style of deep learning while making uncertainty-aware models substantially easier to build and study.