summaries: In this paper the authors propose a new variance-reduction technique to use when computing an expected loss gradient where the expectation is with respect to independent binary random variables, e.g. for training VAEs with a discrete latent space. The paper is interesting, highly relevant, simple to implement, suggests many possible extensions, and shows good results on the experiments performed. However the exposition leaves a lot to be desired.

summaries: For binary layers, how to calculate and backpropagate gradients is a big problem, particularly for the binary neural networks. To solve the problem, this paper proposes an unbiased and low variance augment-REINFORCE-merge (ARM) estimator. With the help of an appropriate reparameterization, the antithetic sampling in an augmented space can be used to drive a variance-reduction mechanism. The experimental results show that ARM estimator converges fast, has low computational complexity, and provides advanced prediction performance.

abstractText: To backpropagate the gradients through stochastic binary layers, we propose the augment-REINFORCE-merge (ARM) estimator that is unbiased, exhibits low variance, and has low computational complexity. Exploiting variable augmentation, REINFORCE, and reparameterization, the ARM estimator achieves adaptive variance reduction for Monte Carlo integration by merging two expectations via common random numbers. The variance-reduction mechanism of the ARM estimator can also be attributed to either antithetic sampling in an augmented space, or the use of an optimal anti-symmetric “self-control” baseline function together with the REINFORCE estimator in that augmented space. Experimental results show the ARM estimator provides state-of-the-art performance in auto-encoding variational inference and maximum likelihood estimation, for discrete latent variable models with one or multiple stochastic binary layers. Python code for reproducible research is publicly available.

authors:
Yoshua Bengio
Nicholas Léonard
Aaron Courville

title: Estimating or propagating gradients through stochastic neurons for conditional computation

venue: arXiv preprint arXiv:1308.3432,

title: Neural Networks for Pattern Recognition

venue: Oxford university press,

authors:
David M. Blei
Alp Kucukelbir
Jon D. McAuliffe

title: Variational inference: A review for statisticians

venue: Journal of the American Statistical Association,

authors:
Junyoung Chung
Sungjin Ahn
Yoshua Bengio

title: Hierarchical multiscale recurrent neural networks

venue: arXiv preprint arXiv:1609.01704,

title: Gradient estimation

venue: Handbooks in operations research and management science,

title: Likelihood ratio gradient estimation for stochastic systems

venue: Communications of the ACM,

authors:
Will Grathwohl
Dami Choi
Yuhuai Wu
Geoff Roeder
David Duvenaud

title: Backpropagation through the Void: Optimizing control variates for black-box gradient estimation

venue: In ICLR,

authors:
Karol Gregor
Ivo Danihelka
Andriy Mnih
Charles Blundell
Daan Wierstra

title: Deep autoregressive networks

venue: arXiv preprint arXiv:1310.8499,

authors:
Shixiang Gu
Sergey Levine
Ilya Sutskever
Andriy Mnih

title: MuProp: Unbiased backpropagation for stochastic neural networks

authors:
Geoffrey E Hinton
Peter Dayan
Brendan J Frey
Radford M Neal

title: The “wake-sleep” algorithm for unsupervised neural networks

authors:
Michael I Jordan
Zoubin Ghahramani
Tommi S Jaakkola
Lawrence K Saul

title: An introduction to variational methods for graphical models

venue: Machine learning,

authors:
Diederik P Kingma
Jimmy Ba

authors:
Diederik P Kingma
Max Welling

title: Auto-encoding variational Bayes

venue: arXiv preprint arXiv:1312.6114,

authors:
Alp Kucukelbir
Dustin Tran
Rajesh Ranganath
Andrew Gelman
David M Blei

title: Automatic differentiation variational inference

venue: Journal of Machine Learning Research,

authors:
Hugo Larochelle
Iain Murray

title: The neural autoregressive distribution estimator

venue: In Proceedings of the Fourteenth International Conference on Artificial Intelligence and Statistics,

authors:
Andrew L Maas
Awni Y Hannun
Andrew Y Ng

title: Rectifier nonlinearities improve neural network acoustic models

venue: In ICML,

authors:
Chris J Maddison
Andriy Mnih
Yee Whye Teh

title: The Concrete distribution: A continuous relaxation of discrete random variables

authors:
Andriy Mnih
Karol Gregor

title: Neural variational inference and learning in belief networks

venue: In ICML, pp

authors:
Andriy Mnih
Danilo J Rezende

title: Variational inference for Monte Carlo objectives

venue: arXiv preprint arXiv:1602.06725,

authors:
Christian Naesseth
Francisco Ruiz
Scott Linderman
David Blei

title: Reparameterization gradients through acceptance-rejection sampling algorithms

venue: In AISTATS,

title: Connectionist learning of belief networks

venue: Artificial Intelligence,

title: Monte Carlo Theory, Methods and Examples, chapter 8 Variance Reduction

authors:
John Paisley
David M Blei
Michael I Jordan

title: Variational Bayesian inference with stochastic search

authors:
Tapani Raiko
Mathias Berglund
Guillaume Alain
Laurent Dinh

title: Techniques for learning binary stochastic feedforward neural networks

venue: arXiv preprint arXiv:1406.2989,

authors:
Rajesh Ranganath
Sean Gerrish
David Blei

title: Black box variational inference

venue: In AISTATS, pp

authors:
Danilo Jimenez Rezende
Shakir Mohamed
Daan Wierstra

title: Stochastic backpropagation and approximate inference in deep generative models

title: Introduction to Probability Models

authors:
Francisco J.R. Ruiz
Michalis K. Titsias
David M. Blei

title: The generalized reparameterization gradient

venue: In NIPS,

authors:
Ruslan Salakhutdinov
Iain Murray

title: On the quantitative analysis of deep belief networks

authors:
Lawrence K Saul
Tommi Jaakkola
Michael I Jordan

title: Mean field theory for sigmoid belief networks

venue: Journal of Artificial Intelligence Research,

authors:
Yichuan Tang
Ruslan R Salakhutdinov

title: Learning stochastic feedforward neural networks

venue: In NIPS, pp

authors:
Martin A Tanner
Wing Hung Wong

title: The calculation of posterior distributions by data augmentation

venue: J. Amer. Statist. Assoc.,

authors:
Michalis K Titsias
Miguel Lázaro-Gredilla

title: Local expectation gradients for black box variational inference

authors:
George Tucker
Andriy Mnih
Chris J Maddison
John Lawson
Jascha Sohl-Dickstein

title: REBAR: Low-variance, unbiased gradient estimates for discrete latent variable models

authors:
Aaron van den Oord
Oriol Vinyals
Koray Kavukcuoglu

title: Neural discrete representation learning

authors:
David A Van Dyk
Xiao-Li Meng

title: The art of data augmentation

venue: Journal of Computational and Graphical Statistics,

title: Simple statistical gradient-following algorithms for connectionist reinforcement learning

venue: In Reinforcement Learning,

authors:
Mingzhang Yin
Mingyuan Zhou

title: Semi-implicit variational inference

authors:
Mingyuan Zhou
Lawrence Carin

title: Negative binomial process count and mixture modeling

venue: arXiv preprint arXiv:1209.3442v1,

heading: 1 INTRODUCTION

text:
Given a function f(z) of a random variable z = (z1, . . . , zV )T , which follows a distribution qφ(z) parameterized by φ, there has been significant recent interest in estimating φ to maximize (or minimize) the expectation of f(z) with respect to z ∼ qφ(z), expressed as
E(φ) = ∫ f(z)qφ(z)dz = Ez∼qφ(z)[f(z)]. (1)
In particular, this expectation objective appears in both maximizing the evidence lower bound (ELBO) for variational inference (Jordan et al., 1999; Blei et al., 2017) and approximately maximizing the log marginal likelihood of a hierarchal Bayesian model (Bishop, 1995), two fundamental problems in statistical inference. To maximize (1), if ∇zf(z) is tractable to compute and z ∼ qφ(z) can be generated via reparameterization as z = Tφ( ), ∼ p( ), where are random noises and Tφ(·) denotes a deterministic transform parameterized by φ, then one may apply the reparameterization trick (Kingma & Welling, 2013; Rezende et al., 2014) to compute the gradient as
∇φE(φ) = ∇φE ∼p( )[f(Tφ( ))] = E ∼p( )[∇φf(Tφ( ))]. (2)
This trick, however, is often inapplicable to discrete random variables, as widely used to construct discrete latent variable models such as sigmoid belief networks (Neal, 1992; Saul et al., 1996).
To maximize (1) for discrete z, using the score function∇φ log qφ(z) = ∇φqφ(z)/qφ(z), one may compute ∇φE(φ) via REINFORCE (Williams, 1992) as
∇φE(φ) = Ez∼qφ(z)[f(z)∇φ log qφ(z)] ≈ 1
K ∑K k=1 f(z(k))∇φ log qφ(z(k)),
where z(k) iid∼ qφ(z) are independent, and identically distributed (iid). This unbiased estimator is also known as (a.k.a.) the score-function (Fu, 2006) or likelihood-ratio estimator (Glynn, 1990). While it is unbiased and only requires drawing iid random samples from qφ(z) and computing∇φ log qφ(z(k)), ∗Corresponding author.
its high Monte-Carlo-integration variance often limits its use in practice. Note that if f(z) depends onφ, then we assume it is true that Ez∼qφ(z)[∇φf(z)] = 0. For example, in variational inference, we need to maximize the ELBO as Ez∼qφ(z)[f(z)], where f(z) = log[p(x | z)p(z)/qφ(z)]. In this case, although f(z) depends on φ, as Ez∼qφ(z)[∇φ log qφ(z)] = ∫ ∇φqφ(z)dz = ∇φ ∫ qφ(z)dz = 0, we have Ez∼qφ(z)[∇φf(z)] = 0. To address the high-variance issue, one may introduce an appropriate baseline (a.k.a. control variate) to reduce the variance of REINFORCE (Paisley et al., 2012; Ranganath et al., 2014; Mnih & Gregor, 2014; Gu et al., 2016; Mnih & Rezende, 2016; Ruiz et al., 2016; Kucukelbir et al., 2017; Naesseth et al., 2017). Alternatively, one may first relax the discrete random variables with continuous ones and then apply the reparameterization trick to estimate the gradients, which reduces the variance of Monte Carlo integration at the expense of introducing bias (Maddison et al., 2017; Jang et al., 2017). Combining both REINFORCE and the continuous relaxation of discrete random variables, REBAR of Tucker et al. (2017) and RELAX of Grathwohl et al. (2018) both aim to produce a low-variance and unbiased gradient estimator by introducing a continuous relaxation based baseline function, whose parameters, however, need to be estimated at each mini-batch by minimizing the sample variance of the estimator with stochastic gradient descent (SGD). Estimating the baseline parameters often clearly increases the computation. Moreover, the potential conflict, between minimizing the sample variance of the gradient estimate and maximizing the expectation objective, could slow down or even prevent convergence and increase the risk of overfitting. Another interesting variance-control idea applicable to discrete latent variables is using local expectation gradients, which estimates the gradients based on REINFORCE, by performing Monte Carlo integration using a single global sample together with exact integration of the local variable for each latent dimension (Titsias & Lázaro-Gredilla, 2015).
Distinct from the usual idea of introducing baseline functions and optimizing their parameters to reduce the estimation variance of REINFORCE, we propose the augment-REINFORCE-merge (ARM) estimator, a novel unbiased and low-variance gradient estimator for binary latent variables that is also simple to implement and has low computational complexity. We show by rewriting the expectation with respect to Bernoulli random variables as one with respect to augmented exponential random variables, and then expressing the gradient as an expectation via REINFORCE, one can derive the ARM estimator in the augmented space with the assistance of appropriate reparameterization. In particular, in the augmented space, one can derive the ARM estimator by using either the strategy of sharing common random numbers between two expectations, or the strategy of applying antithetic sampling. Both strategies, as detailedly discussed in Owen (2013), can be used to explain why the ARM estimator is unbiased and could lead to significant variance reduction. Moreover, we show that the ARM estimator can be considered as improving the REINFORCE estimator in an augmented space by introducing an optimal baseline function subject to an anti-symmetric constraint; this baseline function can be considered as a “self-control” one, as it exploits the function f itself and correlated random noises for variance reduction, and adds no extra parameters to learn. This “self-control” feature makes the ARM estimator distinct from both REBAR and RELAX, which rely on minimizing the sample variance of the gradient estimate to optimize the baseline function.
We perform experiments on a representative toy optimization problem and both auto-encoding variational inference and maximum likelihood estimation for discrete latent variable models, with one or multiple binary stochastic layers. Our extensive experiments show that the ARM estimator is unbiased, exhibits low variance, converges fast, has low computation, and provides state-of-the-art out-of-sample prediction performance for discrete latent variable models, suggesting the effectiveness of using the ARM estimator for gradient backpropagation through stochastic binary layers. Python code for reproducible research is available at https://github.com/mingzhang-yin/ARM-gradient.

heading: 2 ARM: AUGMENT-REINFORCE-MERGE ESTIMATOR

text:
In this section, we first present the key theorem of the paper, and then provide its derivation. With this theorem, we summarize ARM gradient ascent for multivariate binary latent variables in Algorithm 1, as shown in Appendix A. Let us denote σ(φ) = eφ/(1 + eφ) as the sigmoid function and 1[·] as an indicator function that equals to one if the argument is true and zero otherwise. Theorem 1 (ARM). For a vector of V binary random variables z = (z1, . . . , zV )T , the gradient of
E(φ) = Ez∼∏Vv=1 Bernoulli(zv ;σ(φv))[f(z)] (3)
with respect to φ = (φ1, . . . , φV )T , the logits of the Bernoulli probability parameters, can be expressed as
∇φE(φ) = Eu∼∏Vv=1 Uniform(uv ;0,1) [( f(1[u>σ(−φ)])− f(1[u<σ(φ)]) ) (u− 1/2) ] , (4)
where 1[u>σ(−φ)] := ( 1[u1>σ(−φ1)], . . . ,1[uV >σ(−φV )] )T .

heading: 2.1 UNIVARIATE ARM ESTIMATOR

text:
Below we will first present the ARM estimator for a univariate binary latent variable (i.e., V = 1), and then generalize it to a multivariate one (i.e., V > 1). In the univariate case, we need to evaluate the gradient of E(φ) = Ez∼Bernoulli(σ(φ))[f(z)] with respect to φ, which has an analytic expression as
∇φE(φ) = ∇φ[σ(φ)f(1) + σ(−φ)f(0)] = σ(φ)σ(−φ)[f(1)− f(0)]. (5)
Since ∫ σ(φ)
0 (1− 2u)du = σ(φ)σ(−φ) and ∫ 1 σ(φ) (1− 2u)du = −σ(φ)σ(−φ), we can rewrite (5) as
∇φE(φ) = ∫ σ(φ) 0 f(1)(1− 2u)du+ ∫ 1 σ(φ) f(0)(1− 2u)du = ∫ 1 0 f(1[u<σ(φ)])(1− 2u)du
= Eu∼Uniform(0,1)[f(1[u<σ(φ)])(1− 2u)]. (6)
We refer to (6) as the univariate augment-REINFORCE (AR) estimator, as we initially derived it by combining variable augmentation (Tanner & Wong, 1987; Van Dyk & Meng, 2001), REINFORCE, and Rao Blackwellization (Casella & Robert, 1996); we defer the details to Appendix B.
Applying antithetic sampling (Owen, 2013) to the AR estimator in (6), with ũ = 1− u, we have
∇φE(φ) = Eu∼Uniform(0,1)[f(1[u<σ(φ)])(1/2− u)] + Eũ∼Uniform(0,1)[f(1[ũ<σ(φ)])(1/2− ũ)] = Eu∼Uniform(0,1)[f(1[u<σ(φ)])(1/2− u) + f(1[ũ<σ(φ)])(1/2− ũ)] = Eu∼Uniform(0,1) [( f(1[u>σ(−φ)])− f(1[u<σ(φ)]) ) (u− 1/2) ] , (7)
which provides the proof for Theorem 1 for V = 1. We refer to (7) as the univariate ARM estimator, as we initially derived it by improving the AR estimator with an additional merge step, which shares random numbers between two expectations to reduce the variance of Monte Carlos integration; we defer the details to Appendix B.
Note that since Eu∼Uniform(0,1)[f(1[u<σ(φ)])(1/2−u)] = −Eu∼Uniform(0,1)[f(1[u>σ(−φ)])(1/2−u)], we have Eu∼Uniform(0,1) [( f(1[u<σ(φ)]) + f(1[u>σ(−φ)]) ) (1/2− u) ] = 0, subtracted which from the AR estimator in (6) leads to the ARM estimator in (7). For this reason, we can also consider the ARM estimator as the AR estimator subtracted by a zero-mean baseline function as
b(u) = ( f(1[u<σ(φ)]) + f(1[u>σ(−φ)]) ) (1/2− u).
This baseline function is distinct from previously proposed ones in being parameterized by the function f itself over two correlated binary latent variables and satisfying b(u) = −b(1− u). From this point of view, the ARM estimator can be considered as a “self-control” gradient estimator that exploits the function f itself to control the variance of Monte Carlo integration .

heading: 2.2 MULTIVARIATE GENERALIZATION VIA THE LAW OF TOTAL EXPECTATION

text:
Although the ARM estimator for univariate binary is of little use in practice, as the true gradient, shown in (5), has an analytic expression, it serves as the foundation for generalizing it to multivariate settings. Let us denote (·)\v as a vector whose vth element is removed. For the expectation in (3), applying the univariate ARM estimator in (37) together with the law of total expectation, we have
∇φvE(φ) = Ez\v∼∏j 6=v Bernoulli(zj ;σ(φj)){∇φvEzv∼Bernoulli(σ(φv)[f(z)]} = Ez\v∼∏j 6=v Bernoulli(zj ;σ(φj)){Euv∼Uniform(0,1)[(uv − 1/2) × ( f(z\v, zv = 1[uv>σ(−φv)])− f(z\v, zv = 1[uv<σ(φv)]) )]} . (8)
Since z\v ∼ ∏ j 6=v Bernoulli(zj ;σ(φj)) can be equivalently generated as z\v = 1[u\v<σ(φ\v)] or
as z\v = 1[u\v>σ(−φ\v)], where u\v ∼ ∏ j 6=v Uniform(uj ; 0, 1), exchanging the order of the two
expectations in (8) and applying reparameterization, we conclude the proof for (4) of Theorem 1 with ∇φvE(φ) = Euv∼Uniform(0,1) { (uv − 1/2) Ez\v∼∏j 6=v Bernoulli(zj ;σ(φj))[ f(z\v, zv = 1[uv>σ(−φv)])− f(z\v, zv = 1[uv<σ(φv)])
]} = Eu∼∏Vv=1 Uniform(uv ;0,1) [ (uv − 1/2)f(z\v = 1[u\v>σ(−φ\v)], zv = 1[uv>σ(−φv)])
] − Eu∼∏Vv=1 Uniform(uv;0,1) [ (uv − 1/2)f(z\v = 1[u\v<σ(φ\v)], zv = 1[uv<σ(φv)])
] = Eu∼∏Vv=1 Uniform(uv ;0,1) [ (uv − 1/2) ( f(1[u>σ(−φ)])− f(1[u<σ(φ)]) )] . (9)
Alternatively, instead of generalizing the univariate ARM gradient as in (8) and (9), we can first do a multivariate generalization of the univariate AR gradient in (6) as
∇φvE(φ) = Ez\v∼∏j 6=v Bernoulli(zj ;σ(φj)){∇φvEzv∼Bernoulli(σ(φv)[f(z)]} = Ez\v∼∏j 6=v Bernoulli(zj ;σ(φj)){Euv∼Uniform(0,1)[(1− 2uv)f(z\v, zv = 1[uv<σ(φv)])]} = Eu∼∏Vv=1 Uniform(uv ;0,1) [ (1− 2uv)f(1[u<σ(φ)]) ] . (10)
The same as the derivation of the univariate ARM estimator, here we can arrive at (4) from (10) by either adding an antithetic sampling step, or subtracting the AR estimator by a baseline function as
b(u) = ( f(1[u<σ(φ)]) + f(1[u>σ(−φ)]) ) (1/2− u), (11)
which has zero mean, satisfies b(u) = −b(1−u), and is distinct from previously proposed baselines in taking advantage of “self-control” for variance reduction and adding no extra parameters to learn.

heading: 2.3 EFFECTIVENESS OF ARM FOR VARIANCE REDUCTION

text:
For the univariate case, we show below that the ARM estimator has smaller worst-case variance than REINFORCE does. The proof is deferred to Appendix C. Proposition 2 (Univariate gradient variance). For the objective function Ez∼Bernoulli(σ(φ))[f(z)], with a single Monte-Carlo sample u ∼ Uniform(0, 1) or z ∼ Bernoulli(σ(φ)), the ARM gradient is expressed as gARM(u, φ) = ( f(1[u>σ(−φ)]) − f(1[u<σ(φ)]) ) (u − 1/2), and the REINFORCE gradient as gR(z, φ) = f(z)∇φ log Bernoulli(z;σ(φ)) = f(z)(z − σ(φ)). Assuming f ≥ 0 (or f ≤ 0), then supφvar[gARM(u,φ)]supφvar[gR(u,φ)] ≤ 16 25 (1− 2 f(0) f(0)+f(1) ) 2 ≤ 1625 . In the general setting, with u(1), . . . ,u(K) iid∼ ∏V v=1 Uniform(0, 1), we define the ARM estimate of ∇φvE(φ) with K Monte Carlo samples, denoted as gARMK ,v, and the AR estimate with 2K Monte Carlo samples, denoted as gAR2K ,v , using
gARMK ,v = 1
2K ∑K k=1(gv(u (k)) + gv(1− u(k))), gAR2K ,v = 12K ∑2K k=1 gv(u (k)), (12)
where gv(u(k)) = f(1[u(k)<σ(φ)])(1− 2u (k) v ). Similar to the analysis in Owen (2013), the amount of variance reduction brought by the ARM estimator can be reflected by the ratio as
var[gARMK ,v] var[gAR2K ,v] = var[gv(u)]− Cov(−gv(u), gv(1− u)) var[gv(u)] = 1−ρv, ρv = Corr(−gv(u), gv(1−u)).
Note−gv(u) = f(1[u<σ(φ)])(2uv−1), gv(1−u) = f(1[u>σ(−φ)])(2uv−1), and P (1[uv<σ(φv)] = 1[uv>σ(−φv)]) = σ(|φv|) − σ(−|φv|). Therefore a strong positive correlation (i.e., ρv → 1) and hence noticeable variance reduction are likely, especially if φv moves far away from zero during training. Concretely, we have the following proposition. Proposition 3 (Variance reduction). For the ARM estimate gARMK ,v and AR estimate gAR2K ,v shown in (12), the variance of gARMK ,v is guaranteed to be lower than that of gARK ,v; moreover, if f ≥ 0 (or f ≤ 0), then the variance of gARMK ,v is guaranteed to be lower than that of gAR2K ,v .
We show below that under the anti-symmetric constraint b(u) = −b(1 − u), which implies that Eu∼∏Vv=1 Uniform(uv ;0,1)[b(u)] is a vector of zeros, Equation (11) is the optimal baseline function to be subtracted from the AR estimator for variance reduction. The proof is deferred to Appendix C.
Proposition 4 (Optimal anti-symmetric baseline). For the gradient of Ez∼qφ(z)[f(z)], the optimal anti-symmetric baseline function to be subtracted from the AR estimator gAR(u) = f(1[u<σ(φ)])(1− 2u), which minimizes the variance of Monte Carlo integration, can be expressed as
arg min bv(u)∈B
var[gAR,v(u)− bv(u)] = 1
2 (gAR,v(u)− gAR,v(1− u)), (13)
where B = {b(u) : bv(u) = −bv(1 − u) for all v} is the set of all zero-mean anti-symmetric baseline functions. Note the optimal baseline function shown in (13) is exactly the same as (11), which is subtracted from the AR estimator in (10) to arrive at the ARM estimator in (4). Corollary 5 (Lower variance than constant baseline). The optimal anti-symmetric baseline function for the AR estimator, as shown in (13) and also in (11), leads to lower estimation variance than any constant based baseline function as bv(u) = cv(1/2−uv), where cv is a dimension-specific constant whose value can be optimized for variance reduction.

heading: 3 BACKPROPAGATION THROUGH DISCRETE STOCHASTIC LAYERS

text:
A latent variable model with multiple stochastic hidden layers can be constructed as
x ∼ pθ0(x | b1), b1 ∼ pθ1(b1 | b2), . . . , bt ∼ pθt(bt | bt+1), . . . , bT ∼ pθT (bT ), (14)
whose joint likelihood given the distribution parameters θ0:T = {θ0, . . . ,θT } is expressed as p(x, b1:T |θ0:T ) = pθ0(x | b1) [∏T−1
t=1 pθt(bt | bt+1)
] pθT (bT ). (15)
In comparison to deterministic feedforward neural networks, stochastic ones can represent complex distributions and show natural resistance to overfitting (Neal, 1992; Saul et al., 1996; Tang & Salakhutdinov, 2013; Raiko et al., 2014; Gu et al., 2016; Tang & Salakhutdinov, 2013). However, the training of the network, especially if there are stochastic discrete layers, is often much more challenging. Below we show for both auto-encoding variational inference and maximum likelihood estimation, how to apply the ARM estimator for gradient backpropagation in stochastic binary networks.

heading: 3.1 ARM VARIATIONAL AUTO-ENCODER

text:
For auto-encoding variational inference (Kingma & Welling, 2013; Rezende et al., 2014), we construct a variational distribution as
qw1:T (b1:T |x) = qw1(b1 |x) [∏T−1
t=1 qwt+1(bt+1 | bt)
] , (16)
with which the ELBO can be expressed as
E(w1:T ) = Eb1:T∼qw1:T (b1:T |x) [f(b1:T )] , where f(b1:T ) = log pθ0(x | b1) + log pθ1:T (b1:T )− log qw1:T (b1:T |x).
(17)
Proposition 6 (ARM backpropagation). For a stochastic binary network with T binary stochastic hidden layers, constructing a variational auto-encoder (VAE) defined with b0 = x and
qwt(bt | bt−1) = Bernoulli(bt;σ(Twt(bt−1))) (18)
for t = 1, . . . , T , the gradient of the ELBO with respect to wt can be expressed as
∇wtE(w1:T ) = Eq(b1:t−1) [Eut∼Uniform(0,1)[f∆(ut, Twt(bt−1), b1:t−1)(ut − 1/2)]∇wtTwt(bt−1)] , where f∆(ut, Twt(bt−1), b1:t−1) = Ebt+1:T∼q(bt+1:T | bt), bt=1[ut>σ(−Twt (bt−1))] [f(b1:T )]
− Ebt+1:T∼q(bt+1:T | bt), bt=1[ut<σ(Twt (bt−1))] [f(b1:T )]. (19)
The gradient presented in (19) can be estimated with a single Monte Carlo sample as
f̂∆(ut, Twt(bt−1), b1:t−1) =
{ 0, if b(1)t = b (2) t ,
f(b1:t−1, b (1) t:T )− f(b1:t−1, b (2) t:T ), otherwise,
(20)
where b(1)t = 1[ut>σ(−Twt (bt−1))], b (1) t+1:T ∼ q(bt+1:T | b (1) t ), b (2) t = 1[ut<σ(Twt (bt−1))], and b (2) t+1:T ∼ q(bt+1:T | b (2) t ). The proof of Proposition 6 is provided in Appendix C. Suppose the computation complexity of vanilla REINFORCE for a stochastic hidden layer isO(1), which involves a single evaluation of the function f and gradient backpropagation as ∇wtTwt(bt−1), then for a T -stochastic-hidden-layer network, the computation complexity of vanilla REINFORCE is O(T ). By contrast, if evaluating f is much less expensive in computation than gradient backpropagation, then the ARM estimator also has O(T ) complexity, whereas if evaluating f dominates gradient backpropagation in computation, then its worst-case complexity is O(2T ).

heading: 3.2 ARM MAXIMUM LIKELIHOOD ESTIMATION

text:
For maximum likelihood estimation, the log marginal likelihood can be expressed as
log pθ0:T (x) = logEb1:T∼pθ1:T (b1:T )[pθ0(x | b1)] ≥ E(θ1:T ) = Eb1:T∼pθ1:T (b1:T )[log pθ0(x | b1)]. (21)
Generalizing Proposition 6 leads to the following proposition. Proposition 7. For a stochastic binary network defined as
pθt(bt | bt+1) = Bernoulli(bt;σ(Tθt(bt+1))), (22)
the gradient of the lower bound in (21) with respect to θt can be expressed as
∇θtE(θ1:T ) = Ep(bt+1:T ) [Eut∼Uniform(0,1)[f∆(ut, Tθt(bt+1), bt+1:T )(ut − 1/2)]∇θtTθt(bt+1)] , where f∆(ut, Tθt(bt+1), bt+1:T ) = Eb1:t−1∼p(b1:t−1 | bt), bt=1[ut>σ(−Tθt (bt+1))])[log pθ0(x | b1)]
− Eb1:t−1∼p(b1:t−1 | bt), bt=1[ut<σ(Tθt (bt+1))])[log pθ0(x | b1)].

heading: 4 EXPERIMENTAL RESULTS

text:
To illustrate the working mechanism of the ARM estimator, related to Tucker et al. (2017) and Grathwohl et al. (2018), we consider learning φ to maximize
E(φ) = Ez∼Bernoulli(σ(φ))[(z − p0)2], where p0 ∈ {0.49, 0.499, 0.501, 0.51}.
The optimal solution is σ(φ) = 1[p0<0.5]. The closer p0 is to 0.5, the more challenging the optimization becomes. We compare both the AR and ARM estimators to the true gradient as
gφ = (1− 2p0)σ(φ)(1− σ(φ)) (23)
and three previously proposed unbiased estimators, including REINFORCE, REBAR (Tucker et al., 2017), and RELAX (Grathwohl et al., 2018). Since RELAX is closely related to REBAR in introducing stochastically estimated control variates to improve REINFORCE, and clearly outperforms REBAR in our experiments for this toy problem (as also shown in Grathwohl et al. (2018) for p0 = 0.49), we omit the results of REBAR for brevity. With a single random sample u ∼ Uniform(0, 1) for Monte Carlo integration, the REINFORCE and AR gradients can be expressed as
gφ,REINFORCE = (1[u<σ(φ)] − p0)2(1[u<σ(φ)] − σ(φ)), gφ,AR = (1[u<σ(φ)] − p0)2(1− 2u),
while the ARM gradient can be expressed as gφ,ARM = [ (1[u>σ(−φ)] − p0)2 − (1[u<σ(φ)] − p0)2 ] (u− 1/2). (24)
See Grathwohl et al. (2018) for the details on RELAX.
As shown in Figure 1, the REINFORCE gradients have large variances. Consequently, a REINFORCE based gradient ascent algorithm may diverge if the gradient ascent stepsize is not sufficiently small. For example, when p0 = 0.501, the optimal value for the Bernoulli probability σ(φ) is 0, but the algorithm with 0.1 as the stepsize infers it to be close to 1 at the end of 2000 iterations of a random trial. The AR estimator behaves similarly as REINFORCE does. By contrast, both RELAX and ARM exhibit clearly lower estimation variance. It is interesting to note that the trace plots of the estimated probability σ(φ) with the univariate ARM estimator almost exactly match these with the
-0.01
0.01
Gr ad
ien t
True
0.2
0.2 REINFORCE
0.2
0.2 AR
0.2
0.2 RELAX
0.01
0.01 ARM
0.49 0.499 0.501 0.51
0.00
0.50
1.00
Pro ba
bil ity
0.0
0.5
1.0
0.0
0.5
1.0
0.0
0.5
1.0
0.0
0.5 1.0 0.49 0.499 0.501 0.51
0 2000 -0.05
0.00
0.05
G_ va
r
0 2000 1.6
1.8
2.0
2.2 1e 2
0 2000 1.6
1.8
2.0
2.2 1e 2
0 2000 0.0
0.5
1.0 1e 3
0 2000
1.00
1.25
1.50
1e 5
0.49 True
Iteration
Figure 1: Comparison of a variety of gradient estimators in maximizing E(φ) = Ez∼Bernoulli(σ(φ))[(z−p0)2] via gradient ascent, where p0 ∈ {0.49, 0.499, 0.501, 0.51}; the optimal solution is σ(φ) = 1(p0 < 0.5). Shown in Rows 1 and 2 are the trace plots of the true/estimated gradients∇φE(φ) and estimated Bernoulli probability parameters σ(φ), with φ updated via gradient ascent. Shown in Row 3 are the gradient variances for p0 = 0.49, estimated using K = 5000 Monte Carlo samples at each iteration; the theoretical gradient variances are also shown if they can be analytically calculated (see Appendices C and D and for related analytic expressions).
true gradients, despite that the trace plots of the ARM gradients are distinct from these of the true gradients. More specifically, while the true gradients smoothly evolve over iterations, the univariate ARM gradients are characterized by zeros and random spikes; this distinct behavior is expected by examining (38) in Appendix C, which suggests that at any given iteration, the univariate ARM gradient based on a single Monte Carlo sample is either exactly zero, which happens with probability σ(|φ|) − σ(−|φ|), or taking |[f(1)− f(0)](1/2− u)| as its absolute value. These observations suggest that by adjusting the frequencies and amplitudes of spike gradients, the univariate ARM estimator very well approximates the behavior of the true gradient for learning with gradient ascent.
In Figure 4 of Appendix D, we plot the gradient estimated with multiple Monte Carlo samples against the true gradient at each iteration, further showing the ARM estimator has the lowest estimation variance given the same number of Monte Carlo samples. Moreover, in Figure 5 of Appendix D, for each estimator specific column, we plot against the value of φ the sample mean ḡ, sample standard deviation sg, and the gradient signal-to-noise ratio defined as SNRg = |ḡ|/sg; for each φ value, we use K = 1000 single-Monte-Carlo-sample gradient estimates to calculate ḡ, sg, and SNRg. Both figures further show that the ARM estimator outperforms not only REINFORCE, which has large variance, but also RELAX, which improves REINFORCE with an adaptively estimated baseline.
In Figure 5 of Appendix D, it is also interesting to notice that the gradient signal-to-noise ratio for the ARM estimator appears to be only a function of φ but not a function of p0; this can be verified to be true using (23) and (39) in Appendix C, as the ratio of the absolute value of the true gradient |gφ| to√
var[gφ,ARM], the standard deviation of the ARM estimate in (24), can be expressed as σ(φ)(1− σ(φ)) /√
1 16 (1− t)(t3 + 7 3 t 2 + 13 t+ 1 3 ), t = σ(|φ|)− σ(−|φ|). (25)
We find that the values of the ratio shown above are almost exactly matched by the values of SNRg = |ḡ|/sg under the ARM estimator, shown in the bottom right subplot of Figure 5. Therefore, for this example optimization problem, the ARM estimator exhibits a desirable property in providing high gradient signal-to-noise ratios regardless of the value of p0.

heading: 4.1 DISCRETE VARIATIONAL AUTO-ENCODERS

text:
To optimize a variational auto-encoder (VAE) for a discrete latent variable model, existing solutions often rely on biased but low-variance stochastic gradient estimators (Bengio et al., 2013; Jang et al., 2017), unbiased but high-variance ones (Mnih & Gregor, 2014), or unbiased REINFORCE combined with computationally expensive baselines, whose parameters are estimated by minimizing the sample variance of the estimator with SGD (Tucker et al., 2017; Grathwohl et al., 2018). By contrast, the ARM estimator exhibits low variance and is unbiased, efficient to compute, and simple to implement.
For discrete VAEs, we compare ARM with a variety of representative stochastic gradient estimators for discrete latent variables, including Wake-Sleep (Hinton et al., 1995), NVIL (Mnih & Gregor, 2014), LeGrad (Titsias & Lázaro-Gredilla, 2015), MuProp (Gu et al., 2016), Concrete (GumbelSoftmax) (Jang et al., 2017; Maddison et al., 2017), REBAR (Grathwohl et al., 2018), and RELAX (Tucker et al., 2017). Following the settings in Tucker et al. (2017) and Grathwohl et al. (2018), for the encoder defined in (15) and decoder defined in (16), we consider three different network architectures, as summarized in Table 1, including “Nonlinear” that has one stochastic but two Leaky-ReLU (Maas et al., 2013) deterministic hidden layers, “Linear” that has one stochastic hidden layer, and “Linear two layers” that has two stochastic hidden layers. We consider a widely used binarization (Salakhutdinov & Murray, 2008; Larochelle & Murray, 2011; Yin & Zhou, 2018), referred to as MNIST-static and available at http://www.dmi.usherb.ca/∼larocheh/mlpython/ modules/datasets/binarized mnist.html, making our numerical results directly comparable to those reported in the literature. In addition to MNIST-static, we also consider MNIST-threshold (van den Oord et al., 2017), which binarizes MNIST by thresholding each pixel value at 0.5, and the binarized OMNIGLOT dataset.
We train discrete VAEs with 200 conditionally iid Bernoulli random variables as the hidden units of each stochastic binary layer. We maximize a single-Monte-Carlo-sample ELBO using Adam (Kingma & Ba, 2014), with the learning rate selected from {5, 1, 0.5} × 10−4 by the validation set. We set the batch size as 50 for MNIST and 25 for OMNIGLOT. For each dataset, using its default
training/validation/testing partition, we train all methods on the training set, calculate the validation log-likelihood for every epoch, and report the test negative log-likelihood when the validation negative log-likelihood reaches its minimum within a predefined maximum number of iterations.
We summarize the test negative log-likelihoods in Table 2 for MNIST-static. We also summarize the test negative ELBOs in Table 4 of the Appendix, and provide related trace plots of the training and validation negative ELBOs on MNIST-static in Figure 2, and these on MNIST-threshold and OMNIGLOT in Figures 6 and 7 of the Appendix, respectively. For these trace plots, for a fair comparison of convergence speed between different algorithms, we use publicly available code and setting the learning rate of ARM the same as that selected by RELAX in Grathwohl et al. (2018). Note as shown in Figures 2(a,d) and 7(a,d), both REBAR and RELAX exhibit clear signs of overfitting on both MNIST-static and Omniglot using the “Nonlinear” architecture; as ARM runs much faster per iteration than both of them and do not exhibit overfitting given the same number of iterations, we allow ARM to run more stochastic gradient ascent steps under these two scenarios to check whether it will eventually overfit the training set.
These results show that ARM provides state-of-the-art performance in delivering not only fast convergence, but also low negative log-likelihoods and negative ELBOs on both the validation and test sets, with low computational cost, for all three different network architectures. In comparison to the vanilla REINFORCE on MNIST-static, as shown in Table 2 (a), ARM achieves significantly lower test negative log-likelihoods, which can be explained by having much lower variance in its gradient
estimation, while only costing 20% to 30% more computation time to finish the same number of iterations.
The trace plots in Figures 2, 6, and 7 show that ARM achieves its objective better or on a par with the competing methods in all three different network architectures. In particular, the performance of ARM on MNIST-threshold is significantly better, suggesting ARM is more robust, better resists overfitting, and has better generalization ability. On both MNIST-static and OMNIGLOT, with the “Nonlinear” network architecture, both REBAR and RELAX exhibit severe overfitting, which could be caused by their training procedure, which updates the parameters of the baseline function by minimizing the sample variance of the gradient estimator using SGD. For less overfitting linear and two-stochastic-layer networks, ARM overall performs better than both REBAR and RELAX and runs significantly faster (about 6-8 times faster) in terms of the computation time per iteration.
To understand why ARM has the best overall performance, we examine the trace plots of the logarithm of the estimated variance of gradient estimates in Figure 3. On the MNIST-static dataset with the “Nonlinear” network, the left subplot of Figure 3 shows that both REBAR and RELAX exhibit lower variance than ARM does for their single-Monte-Carlo-sample based gradient estimates; however, the corresponding trace plots of the validation negative ELBOs, shown in Figure 2(a), suggest they both severely overfit the training data as the learning progresses; our hypothesis for this phenomenon is that REBAR and RELAX may favor suboptimal solutions that are associated with lower gradient variance; in other words, they may have difficulty in converging to local optimal solutions that are associated with high gradient variance. For the “Linear” network architecture, the right subplot of Figure 3 shows that ARM exhibits lower variance for its gradient estimate than both REBAR and RELAX do, and Figure 2(b) shows that none of them exhibit clear signs of overfitting; this observation could be used to explain why ARM results in the best convergence for both the training and validation negative ELBOs, as shown in Figure 2(b).

heading: 4.2 MAXIMUM LIKELIHOOD ESTIMATION FOR A STOCHASTIC BINARY NETWORK

text:
Denoting xl,xu ∈ R394 as the lower and upper halves of an MNIST digit, respectively, we consider a standard benchmark task of estimating the conditional distribution pθ0:2(xl |xu) (Raiko et al., 2014; Bengio et al., 2013; Gu et al., 2016; Jang et al., 2017; Tucker et al., 2017), using a stochastic binary network with two stochastic binary hidden layers, expressed as
xl ∼ Bernoulli(σ(Tθ0(b1))), b1 ∼ Bernoulli(σ(Tθ1(b2))), b2 ∼ Bernoulli(σ(Tθ2(xu))). (26)
We set the network structure as 392-200-200-392, which means both b1 and b2 are 200 dimensional binary vectors and the transformation Tθ are linear so the results are directly comparable with those in Jang et al. (2017). We approximate log pθ0:2(xl |xu) with log 1K ∑K k=1 Bernoulli(xl;σ(Tθ0(b (k) 1 ))), where b(k)1 ∼ Bernoulli(σ(Tθ1(b (k) 2 ))), b (k) 2 ∼ Bernoulli(σ(Tθ2(xu))). We perform training with K = 1, which can also be considered as optimizing on a single-Monte-Carlo-sample estimate of the lower bound of the log marginal likelihood shown in (21). We use Adam (Kingma & Ba, 2014), with the learning rate set as 10−4, mini-batch size as 100, and number of epochs for training as 2000. Given the inferred point estimate of θ0:2 after training, we evaluate the accuracy of conditional density estimation by estimating the negative log-likelihood as − log pθ0:2(xl |xu), averaging over the test set using K = 1000. We show example results of predicting the activation probabilities of the pixels of xl given xu in Figure 8 of the Appendix.
As shown in Table 3, optimizing a stochastic binary network with the ARM estimator, which is unbiased and computationally efficient, achieves the lowest test negative log-likelihood, outperforming previously proposed biased stochastic gradient estimators on similarly structured stochastic networks, including DARN (Gregor et al., 2013), straight through (ST) (Bengio et al., 2013), slope-annealed ST (Chung et al., 2016), and ST Gumbel-softmax (Jang et al., 2017), and unbiased ones, including score-function (SF) and MuProp (Gu et al., 2016).

heading: 5 CONCLUSION

text: To train a discrete latent variable model with one or multiple stochastic binary layers, we propose the augment-REINFORCE-merge (ARM) estimator to provide unbiased and low-variance gradient estimates of the parameters of Bernoulli distributions. With a single Monte Carlo sample, the estimated gradient is the product of uniform random noises and the difference of a function of two vectors of correlated binary latent variables. Without relying on estimating a baseline function with extra learnable parameters for variance reduction, it maintains efficient computation and avoids increasing the risk of overfitting. Applying the ARM gradient leads to not only fast convergence, but also low test negative log-likelihoods (and low test negative evidence lower bounds for variational inference), on both auto-encoding variational inference and maximum likelihood estimation for stochastic binary feedforward neural networks. Some natural extensions of the proposed ARM estimator include generalizing it to multivariate categorical latent variables, combining it with a baseline or local-expectation based variance reduction method, and applying it to reinforcement learning whose action space is discrete.

heading: ACKNOWLEDGMENTS

text: M. Zhou acknowledges the support of Award IIS-1812699 from the U.S. National Science Foundation, and support from the McCombs Research Excellence Grant. The authors acknowledge the support of NVIDIA Corporation with the donation of the Titan Xp GPU used for this research, and the computational support of Texas Advanced Computing Center.

heading: A THE ARM GRADIENT ASCENT ALGORITHM

text:
We summarize the algorithm to compute ARM gradient for binary latent variables. Here we show the gradient with respect to the logits associated with the probability of Bernoulli random variables. If the logits are further generated by deterministic transform such as neural networks, the gradient with respect to the transform parameters can be directly computed by the chain rule. For stochastic transforms, the implementation of ARM gradient is discussed in detail in Section 3 and summarized in Algorithm 2.
Algorithm 1: ARM gradient for a V -dimensional binary latent vector input : Bernoulli distribution {qφv (zv)}v=1:V with probability {σ(φv)}v=1:V , target f(z); z = (z1, · · · , zV ), φ = (φ1, · · · , φV ) output :φ and ψ that maximize E(φ,ψ) = Ez∼∏Vv=1 qφv (zv)[f(z;ψ)] Initialize φ, ψ randomly; while not converged do
Sample a mini-batch of x from the data; Sample zv ∼ Bernoulli(σ(φv)) for v = 1, · · · , V ; sample uv ∼ Uniform(0, 1) for v = 1, · · · , V , u = (u1, · · · , uV ) ; gψ = ∇ψf(z;ψ) ; f∆(u,φ) = f ( 1[u>σ(−φ)] ) − f ( 1[u<σ(φ)] ) ; gφ = f∆(u,φ)(u− 0.5) φ = φ+ ρtgφ, ψ = ψ + ηtgψ, with stepsizes ρt, ηt
end
Algorithm 2: ARM gradient for a T -stochastic-hidden-layer binary network input : Inference network qw1:T (b1:T |x) = qw1(b1 |x) [∏T−1 t=1 qwt+1(bt+1 | bt) ] where qwt(bt | bt−1) = Bernoulli(bt;σ(Twt(bt−1))), target f(b1:T ;ψ) output : w1:T and ψ that maximize E(w1:T ,ψ) = Eb1:T∼qw1:T [f(b1:T ;ψ)]
Initializew1:T , ψ randomly; while not converged do
Sample a mini-batch of x from data; for t = 1:T do
If t ≥ 2, sample bt−1 ∼ q(bt−1|bt−2), if t = 2, b1:t−1 = b1, else b1:t−1 = [b1:t−2, bt−1] ; sample ut ∼ ∏ Uniform(0, 1) ; b1t = 1[ut>σ(−Twt (bt−1))] ; b2t = 1[ut<σ(Twt (bt−1))] ; if b1t = b2t then
f∆(b1:t−1, b 1 t:T , b 2 t:T ) = 0 ;
else b1t+1:T ∼ q(bt+1:T |b1t ) ; b2t+1:T ∼ q(bt+1:T |b2t ) ; f∆(b1:t−1, b 1 t:T , b 2 t:T ) = f(b1:t−1, b 1 t:T )− f(b1:t−1, b2t:T ) ;
gwt = f∆(b1:t−1, b 1 t:T , b 2 t:T )(ut − 12 ) T∇wtTwt(bt−1) ; end wt = wt + ρtgwt with step-size ρt
end ψ = ψ + ηt∇ψf(b1:T ;ψ) with step-size ηt
end

heading: B ORIGINAL DERIVATION OF THE AR AND ARM ESTIMATORS

text:
Let us denote t ∼ Exp(λ) as an exponential distribution, whose probability density function is defined as p(t |λ) = λe−λt, where λ > 0 and t > 0. The mean and variance are E[t] = λ−1 and var[t] = λ−2, respectively. The exponential random variable t ∼ Exp(λ) can be reparameterized as
t = /λ, ∼ Exp(1). It is well known, e.g., in Ross (2006), that if t1 ∼ Exp(λ1) and t2 ∼ Exp(λ2) are two independent exponential random variables, then the probability that t1 is smaller than t2 can be expressed as P (t1 < t2) = λ1/(λ1 + λ2); moreover, since t1 ∼ Exp(λ1) is equal in distribution to 1/λ1, 1 ∼ Exp(1) and t2 ∼ Exp(λ2) is equal in distribution to 2/λ2, 2 ∼ Exp(1), we have
P (t1 < t2) = P ( 1/λ1 < 2/λ2) = P ( 1 < 2λ1/λ2) = λ1/(λ1 + λ2). (27)
B.1 AUGMENTATION OF A BERNOULLI RANDOM VARIABLE AND REPARAMETERIZATION
From (27) it becomes clear that the Bernoulli random variable z ∼ Bernoulli(σ(φ)) can be reparameterized by comparing two augmented exponential random variables as
z = 1[ 1< 2eφ], 1 ∼ Exp(1), 2 ∼ Exp(1). (28)
Consequently, the expectation with respect to the Bernoulli random variable can be reparameterized as one with respect to two augmented exponential random variables as
E(φ) = Ez∼Bernoulli(σ(φ))[f(z)] = E 1, 2iid∼ Exp(1)[f(1[ 1e−φ< 2])]. (29)
B.2 REINFORCE ESTIMATOR IN THE AUGMENTED SPACE
Since the indicator function 1[ 1e−φ< 2] is not differentiable, the reparameterization trick in (2) is not directly applicable to computing the gradient of (29). Fortunately, as t1 = 1e−φ, 1 ∼ Exp(1) is equal in distribution to t1 ∼ Exp(eφ), the expectation in (29) can be further reparameterized as
E(φ) = E 1, 2 iid∼ Exp(1) [f(1[ 1e−φ< 2])] = Et1∼Exp(eφ), 2∼Exp(1)[f(1[t1< 2])], (30)
and hence, via REINFORCE and then another reparameterization, we can express the gradient as
∇φE(φ) = Et1∼Exp(eφ), 2∼Exp(1)[f(1[t1< 2])∇φ log Exp(t1; e φ)]
= Et1∼Exp(eφ), 2∼Exp(1)[f(1[t1< 2])(1− t1e φ)]
= E 1, 2 iid∼ Exp(1) [f(1[ 1e−φ< 2])(1− 1)]. (31)
Similarly, we have E(φ) = E 1, 2 iid∼ Exp(1) [f(1[ 1< 2eφ])] = E 1∼Exp(1), t2∼Exp(−eφ)[f(1[ 1<t2])], and hence can also express the gradient as
∇φE(φ) = E 1∼Exp(1), t2∼Exp(e−φ)[f(1[ 1<t2])∇φ log Exp(t2; e −φ)]
= −E 1∼Exp(1), t2∼Exp(e−φ)[f(1[ 1<t2])(1− t2e −φ)]
= −E 1, 2 iid∼ Exp(1) [f(1[ 1e−φ< 2])(1− 2)]. (32)
Note that letting 1, 2 iid∼ Exp(1) is the same in distribution as letting
1 = u, 2 = (1− u), where u ∼ Uniform(0, 1), ∼ Gamma(2, 1), (33)
which can be proved using Exp(1) d= Gamma(1, 1), (u, 1 − u)T d= Dirichlet (12), where u ∼ Uniform(0, 1)}, together with Lemma IV.3 of Zhou & Carin (2012); we use “ d=” to denote “equal in distribution.” Thus, (B.2) can be reparameterized as
∇φE(φ) = Eu∼Uniform(0,1), ∼Gamma(2,1) [ f(1[u<σ(φ)])(1− u) ] ,
Applying Rao-Blackwellization (Casella & Robert, 1996), we can further express the gradient as ∇φE(φ) = Eu∼Uniform(0,1) [ f(1[u<σ(φ)])(1− 2u) ] . (34)
Therefore, the gradient estimator shown above, the same as (6), is referred to as the AugmentREINFORCE (AR) estimator.
B.3 MERGE OF REINFORCE GRADIENTS
A key observation of the paper is that by swapping the indices of the two iid standard exponential random variables in (32) , the gradient∇φE(φ) can be equivalently expressed as
∇φE(φ) = −E 1, 2 iid∼ Exp(1) [f(1[ 2e−φ< 1])(1− 1)]. (35)
As the term inside the expectation in (31) and that in (35) could be highly positively correlated, we are motivated to merge (31) and (35) by sharing the same set of standard exponential random variables for Monte Carlo integration, which provides a new opportunity to well control the estimation variance (Owen, 2013). More specifically, simply taking the average of (31) and (35) leads to
∇φE(φ) = E 1, 2 iid∼ Exp(1)
[( f(1[ 1e−φ< 2])− f(1[ 2e−φ< 1]) ) (1/2− 1/2) ] . (36)
Note one may also take a weighted average of (31) and (35), and optimize the combination weight to potentially further reduce the variance of the estimator. We leave that for future study.
Note that (36) can be reparameterized as ∇φE(φ) = Eu∼Uniform(0,1), ∼Gamma(2,1) [( f(1[u>σ(−φ)])− f(1[u<σ(φ)]) ) ( u/2− 1/2) ] ,
Applying Rao-Blackwellization (Casella & Robert, 1996), we can further express the gradient as ∇φE(φ) = Eu∼Uniform(0,1) [( f(1[u>σ(−φ)])− f(1[u<σ(φ)]) ) (u− 1/2) ] . (37)
Therefore, the gradient estimator shown above, the same as (7), is referred to as the AugmentREINFORCE-merge (ARM) estimator.

heading: C PROOFS

text:
Proof of Proposition 2. Since the gradients gARM(u, φ), gAR(u, φ), and gR(z, φ) are all unbiased, their expectations are the same as the true gradient gtrue(φ) = σ(φ)(1−σ(φ))[f(1)−f(0)]. Denote f∆(u, φ) = f(1[u>σ(−φ)])− f(1[u<σ(φ)]). Since
f∆(u, φ) =  0, if σ(−|φ|) < u < σ(|φ|), f(1)− f(0), if u > σ(|φ|), f(0)− f(1), if u < σ(−|φ|),
(38)
The second moment of gARM(u, φ) can be expressed as
Eu∼Uniform(0,1)[g2ARM(u, φ)] = Eu∼Uniform(0,1)[f2∆(u, φ)(u− 1/2)2]
= ∫ 1 σ(|φ|) [f(1)− f(0)]2(u− 1/2)2du+ ∫ σ(−|φ|) 0 [f(0)− f(1)]2(u− 1/2)2du
= 1
12 [1− (σ(|φ|)− σ(−|φ|))3][f(1)− f(0)]2
Denoting t = σ(|φ|)− σ(−|φ|), we can re-express gtrue(φ) = 14 (1− t 2)[f(1)− f(0)]. Thus, the variance of gARM(u, φ) can be expressed as
var[gARM(u, φ)] = 1
4
[ 1
3 (1− t3)− 1 4 (1− t2)2
] [f(1)− f(0)]2
= 1 16 (1− t)(t3 + 7 3 t2 + 1 3 t+ 1 3 )[f(1)− f(0)]2 (39) ≤ 1 25 [f(1)− f(0)]2,
which reaches its maximum at 0.039788[f(1)− f(0)]2 when t = √
5−1 2 .
For the REINFORCE gradient, we have Ez∼Bernoulli(σ(φ))[g2R(z, φ)] =Ez∼Bernoulli(σ(φ)) [ f2(z)(z(1− σ(φ))− σ(φ)(1− z))2 ] =σ(φ)(1− σ(φ))[(1− σ(φ))f2(1) + σ(φ)f2(0)].
Therefore the variance can be expressed as
var[gR(u, φ)] =σ(φ)(1− σ(φ)) [ (1− σ(φ))f2(1) + σ(φ)f2(0)− σ(φ)(1− σ(φ))[f(1)− f(0)]2 ] =σ(φ)(1− σ(φ))[(1− σ(φ))f(1) + σ(φ)f(0)]2.
The largest variance satisfies
sup φ
var[gR(z, φ)] ≥ var[gR(z, 0)] = 1
16 [f(1) + f(0)]2,
and hence when f is always positive or negative, we have
supφ var[gARM(z, φ)] supφ var[gR(z, φ)] ≤ 16 25
( 1− 2 f(0)
f(0) + f(1) )2 ≤ 16 25 .
In summary, the ARM gradient has a variance that is bounded by 125 (f(1)−f(0)) 2, and its worst-case variance is smaller than that of REINFORCE.
Proof of Proposition 3. We only need to prove for K = 1 and the proof for K > 1 automatically follows. Since Eu[f(1[u<σ(φ)])2(uv − 1/2)2] = Eu[f(1[u>σ(−φ)])2(uv − 1/2)2], we have
var(gARM1,v)− var(gAR1,v) = −3Eu[f(1[u<σ(φ)])2(uv − 1/2)2] + Eu[f(1[u>σ(−φ)])2(uv − 1/2)2] − 2Eu[f(1[u>σ(−φ)])f(1[u<σ(φ)])(uv − 1/2)2]
= −Eu[f(1[u<σ(φ)])2(uv − 1/2)2]− Eu[f(1[u>σ(−φ)])2(uv − 1/2)2] − 2Eu[f(1[u>σ(−φ)])f(1[u<σ(φ)])(uv − 1/2)2]
= −Eu [( f(1[u>σ(−φ)]) + f(1[u<σ(φ)]) )2 (uv − 1/2)2 ] ≤ 0,
which shows that the estimation variance of gARMK ,v is guaranteed to be lower than that of the gARK ,v , unless f(1[u>σ(−φ)]) + f(1[u<σ(φ)]) = 0 almost surely. Furthermore, since
var(gARM1,v)− var(gAR2,v) =Eu[(f(1[u<σ(φ)])− f(1[u>σ(−φ)]))2(uv − 1/2)2] − Eu(1),u(2) [(f(1[u(1)<σ(φ)])(u(1)v − 1/2) + f(1[u(2)<σ(φ)])(u(2)v − 1/2))2]
=− 2Eu(1) [f(1[u(1)<σ(φ)])(u(1)v − 1/2)]Eu(2) [f(1[u(2)<σ(φ)])(u(2)v − 1/2))] − 2Eu[f(1[u<σ(φ)])f(1[u>σ(−φ)]))(uv − 1/2)2]
=− 2 ( Eu[f(1[u<σ(φ)])(uv − 1/2)] )2 − 2Eu[f(1[u<σ(φ)])f(1[u>σ(−φ)]))(uv − 1/2)2],
when f is always positive or negative, the variance of gARMK ,v is lower than that of gAR2K ,v .
Proof of Proposition 4. Denoting g(u) = gAR(u)− b(u), we have
var[gv(u)]− var[gAR,v(u)] = −2Eu[gAR,v(u)bv(u)] + Eu[b2v(u)].
To maximize the variance reduction, it is equivalant to consider the constrained optimization problem
min bv(u)
− 2Eu[gAR,v(u)bv(u)] + Eu[b2v(u)]
subject to: bv(u) = −bv(1− u),
which is the same as a Lagrangian problem as
min bv(u),λv(u)
L(bv(u), λv(u)) = −2Eu[gAR,v(u)bv(u)] + Eu[b2v(u)] + ∫ λv(u)(bv(u) + bv(1− u))du.
Setting δLδλv = 0 gives bv(u) + bv(1 − u) = 0. By writing ∫ λv(u)(bv(u) + bv(1 − u))du =∫
(λv(u) + λv(1− u))bv(u)du and setting δLδbv = 0, we have
[2gAR,v(u)− 2bv(u)]p(u) = λv(u) + λv(1− u). (40) Interchange u and 1− u gives
[2gAR,v(1− u)− 2bv(1− u)]p(1− u) = λv(1− u) + λv(u). (41) Solving (40) and (41) with bv(u) + bv(1 − u) = 0 and p(u) = p(1 − u), we have the optimal baseline function as b∗v(u) = 1 2 (gAR,v(u)− gAR,v(1− u)). The proof is completed by noticing that gAR(u)− b∗(u) is the same as the single sample gradient estimate under the ARM estimator.
Proof of Corollary 5. Since bv(u) = cv(1−2u) satisfies the anti-symmetric property, we can directly arrive at Corollary 5 using Proposition 4. Alternatively, since Eu[f(1[u<σ(φ)])2(uv − 1/2)2] = Eu[f(1[u>σ(−φ)])2(uv − 1/2)2] and Eu[f(1[u<σ(φ)])(uv − 1/2)2] = Eu[f(1[u>σ(−φ)])(uv − 1/2)2], for gC,v = (f(1[u<σ(φ)])− cv)(1− 2uv), we have
var(gC,v)− var(gARM,v) = Eu[(f(1[u<σ(φ)])− cv)2(1− 2uv)2]− Eu[(f(1[u<σ(φ)])− f(1[u>σ(−φ)]))2(uv − 1/2)2] = Eu [( 4c2v − 8cvf(1[u<σ(φ)]) + 2f2(1[u<σ(φ)]) + 2f(1[u<σ(φ)])f(1[u>σ(−φ)]) ) (uv − 1/2)2
] = Eu [( f(1[u>σ(−φ)]) + f(1[u<σ(φ)])− 2cv )2 (uv − 1/2)2
] ≥ 0.
Proof of Proposition 6. First, to compute the gradient with respect to w1, since
E(w1:T ) = Eq(b1)Eq(b2:T | b1)[f(b1:T )], (42) we have
∇w1E(w1:T ) = Eu1∼Uniform(0,1)[f∆(u1, Tw1(x))(u1 − 1/2)]∇w1Tw1(x), (43) where
f∆(u1, Tw1(x)) = Eb2:T∼q(b2:T | b1), b1=1[u1>σ(−Tw1 (x))])[f(b1:T )]
− Eb2:T∼q(b2:T | b1), b1=1[u1<σ(Tw1 (x))])[f(b1:T )]. (44)
Second, to compute the gradient with respect to wt, where 2 ≤ t ≤ T − 1, since E(w1:T ) = Eq(b1:t−1)Eq(bt | bt−1)Eq(bt+1:T | bt)[f(b1:T )], (45)
we have ∇wtE(w1:T )=Eq(b1:t−1) [ Eut∼Uniform(0,1)[f∆(ut, Twt(bt−1), b1:t−1)(ut − 1/2)]∇wtTwt(bt−1) ] ,
(46)
where
f∆(ut, Twt(bt−1), b1:t−1) = Ebt+1:T∼q(bt+1:T | bt), bt=1[ut>σ(−Twt (bt−1))])[f(b1:T )]
− Ebt+1:T∼q(bt+1:T | bt), bt=1[ut<σ(Twt (bt−1))])[f(b1:T )]. (47)
Finally, to compute the gradient with respect to wT , we have ∇wT E(w1:T ) = Eq(b1:T−1) [ EuT∼Uniform(0,1)[f∆(uT , TwT (bT−1), b1:T−1)(uT − 1/2)] × ∇wT TwT (bT−1)] , (48)
where
f∆(uT , TwT (bT−1), b1:T−1) = f(b1:T−1, bT = 1[uT>σ(−TwT (bT−1))]) − f(b1:T−1, bT = 1[uT<σ(TwT (bT−1))]). (49)

heading: D ADDITIONAL EXPERIMENTAL RESULTS

text:
For the univariate AR gradient, we have Eu∼Uniform(0,1)[g2AR(u, φ)] = ∫ σ(φ)
0
f2(1)(1− 2u)2du+ ∫ 1 σ(φ) f2(0)(1− 2u)2du
= 1
6 (f2(0) + f2(1)) +
1 6 (1− 2σ(φ))3(f2(0)− f2(1)).
Thus its variance can be expressed as
var[gAR(u, φ)] = 1
6 (f2(0)+f2(1))+
1 6 (1−2σ(φ))3(f2(0)−f2(1))−σ2(φ)(1−σ(φ))2[f(1)−f(0)]2,
which is used for related plots in Figures 1 and 5.
0.01
0.00
0.01
me an
REINFORCE
0.01
0.00
0.01 AR
0.01
0.00
0.01 RELAX
0.01
0.00
0.01 ARM
0.49 0.499 0.501 0.51 True
0.06
0.08
0.10
0.12
std
0.140
0.145
0.00
0.05
0.10
0.002 0.004 0.49 0.499 0.501 0.51 True
2 0 2 0.00
0.05
0.10
SN R
2 0 2 0.00
0.05
0.10
2 0 2 0.0
0.5
1.0
1.5
2.0
2 0 2 0.0
0.5
1.0
1.5 2.0 0.49 0.499 0.501 0.51 True
Phi
Figure 5: Comparison of a variety of gradient estimators for estimating the gradient of E(φ) = Ez∼Bernoulli(σ(φ))[(z − p0)2], where p0 ∈ {0.49, 0.499, 0.501, 0.51} and the values of φ range from −2.5 to 2.5. For each φ value, we compute for each estimator K = 1000 single-Monte-Carlo-sample gradient estimates, and use them to calculate their sample mean ḡ, sample standard deviation sg , and gradient signal-tonoise ratio SNRg = |ḡ|/sg . In each estimator specific column, we plot ḡ, sg , and SNRg in Rows 1, 2, and 3, respectively. The theoretical gradient standard deviations and gradient signal-to-noise ratios are also shown if they can be analytically calculated (see Eq. 25 and Appendices C and D and for related analytic expressions).

title: ARM: AUGMENT-REINFORCE-MERGE GRADIENT FOR STOCHASTIC BINARY NETWORKS


summaries: The paper presents an extension of the MXFusion language that allows the use of probabilistic modules. These modules are defined as a set of random variables and a specific probabilistic distribution. The modules also contain dedicated inference methods. Using these modules, one can use probabilistic distributions with inference methods tailored to the distribution, which are usually more efficient than generic inference systems.

summaries: In this paper authors present a new Probabilistic Programming Language (PPL) MXFusion. Similarly to the languages for the deep learning (TensorFlow, PyTorch, etc.), this language introduce probabilistic modules that are used as building blocks for complex probabilistic models. Introducing modularity to the probabilistic programming, raises the problem of inference for probabilistic models. Since, we cannot obtain the exact solution on practice we have to resort to approximate inference methods. The approximate inference methods can be either generic, thus, being suitable for many probabilistic models but resulting in poor approximation, or specific, thus, having good approximation quality, but only for specific probabilistic models. Authors propose to address this problem by encapsulating specific inference methods in corresponding probabilistic modules. Doing so, one can perform approximate inference for every module with the best suitable inference technique. Authors demonstrate interface of MXFusion for three well known probabilistic models: Bayesian linear regression, deep kernel learning, Bayesian Gaussian process latent variable model.

abstractText: Modularity is a key feature of deep learning libraries but has not been fully exploited for probabilistic programming. We propose to improve modularity of probabilistic programming language by offering not only plain probabilistic distributions but also sophisticated probabilistic model such as Bayesian nonparametric models as fundamental building blocks. We demonstrate this idea by presenting a modular probabilistic programming language MXFusion, which includes a new type of re-usable building blocks, called probabilistic modules. A probabilistic module consists of a set of random variables with associated probabilistic distributions and dedicated inference methods. Under the framework of variational inference, the pre-specified inference methods of individual probabilistic modules can be transparently used for inference of the whole probabilistic model. We demonstrate the power and convenience of probabilistic modules in MXFusion with various examples of Gaussian process models, which are evaluated with experiments on real data.

authors:
Bob Carpenter
Andrew Gelman
Matthew D Hoffman
Daniel Lee
Ben Goodrich
Michael Betancourt
Marcus Brubaker
Jiqiang Guo
Peter Li
Allen Riddell

title: Stan: A probabilistic programming language

authors:
Bob Carpenter
Andrew Gelman
Matthew Hoffman
Daniel Lee
Ben Goodrich
Michael Betancourt
Marcus Brubaker
Jiqiang Guo
Peter Li
Allen Riddell

venue: Journal of Statistical Software, Articles,

authors:
Marco F. Cusumano-Towner
Vikash K. Mansinghka

title: Encapsulating models and approximate inference programs in probabilistic modules

venue: In workshop on probabilistic programming semantics,,

authors:
Zhenwen Dai
Andreas Damianou
Javier González
Neil Lawrence

title: Variational auto-encoded deep Gaussian processes

venue: International Conference on Learning Representations (ICLR),

title: Uber ai labs open sources pyro, a deep probabilistic programming language

authors:
Noah Goodman
Vikash Mansinghka
Daniel M Roy
Keith Bonawitz
Joshua B Tenenbaum

title: Church: A language for generative models

venue: In Uncertainty in Artificial Intelligence,

authors:
Diederik P. Kingma
Jimmy Ba

title: Adam: A method for stochastic optimization

venue: In International Conference on Learning Representations,

authors:
Vikash K. Mansinghka
Daniel Selsam
Yura N. Perov

title: Venture: a higher-order probabilistic programming platform with programmable inference

venue: ArXiv,

title: The Bayes net toolbox for Matlab

venue: Computing Science and Statistics,

title: Ibal: A probabilistic rational programming language

venue: In Proceedings of International Joint Conference on Artificial Intelligence,

title: Figaro: An object-oriented probabilistic programming language

authors:
John Salvatier
Thomas Wiecki
Christopher Fonnesbeck

title: Probabilistic programming in python using PyMC

authors:
John Salvatier
Thomas V. Wiecki
Christopher Fonnesbeck

title: Probabilistic programming in Python using PyMC3

venue: PeerJ Computer Science,

authors:
David J Spiegelhalter
Andrew Thomas
Nicky G Best
Wally R Gilks

title: BUGS: Bayesian inference using Gibbs sampling

venue: MRC Biostatistics Unit,

title: Variational learning of inducing variables in sparse Gaussian processes

authors:
Michalis K. Titsias
Neil D. Lawrence

title: Bayesian gaussian process latent variable model

venue: In International Conference on Artificial Intelligence and Statistics,

authors:
Dustin Tran
Alp Kucukelbir
Adji B. Dieng
Maja Rudolph
Dawen Liang
David M. Blei

title: Edward: A library for probabilistic modeling, inference, and criticism

venue: arXiv preprint arXiv:1610.09787,

authors:
Dustin Tran
Matthew D. Hoffman
Rif A. Saurous
Eugene Brevdo
Kevin Murphy
David M. Blei

title: Deep probabilistic programming

authors:
Andrew Gordon Wilson
Zhiting Hu
Ruslan Salakhutdinov
Eric P. Xing

title: Deep kernel learning

authors:
Frank Wood
Jan Willem Meent
Vikash Mansinghka

title: A New Approach to Probabilistic Programming Inference

venue: In Proceedings of the Seventeenth International Conference on Artificial Intelligence and Statistics,

heading: 1 INTRODUCTION

text:
Deep neural networks (DNNs) have seen great successful in various areas such as computer vision, speech recognition, and artificial intelligence. A major driver of such success is the rise of deep learning libraries, starting from Theano and following by many popular libraries such as TensorFlow, PyTorch. The unique benefits that these libraries bring are: (1) completely transparent training and prediction, (2) modular DNN construction with interchangeable components. A DNN component can vary dramatically in terms of size and complexity, ranging from a single matrix operator, to a highly complex deep neural network such as convolutional neural networks (CNN) or long short-term memory (LSTM). All the components share the same interface, which enables a user to replace one component by another with almost no changes on other parts of the code. In particular, modularity in the programming languages used by deep learning libraries revolutionizes the development, communication and deployment of DNNs. The ability to package a state-of-the-art implementation into a standard module and distribute enables rapid development with much lower chance of having bugs and performance issues comparing with re-implementing from scratch, which makes research results more reproducible. Communication about ongoing research method and results becomes much easier, as running a different experiment or modifying a network only requires to change a few lines of code. The standardized interface of components enables better testing and maintainability. It eases the transition of a research implementation into an industrial software.
Probabilistic programming language (PPL) is a similar domain specific programming language that aims at describing probabilistic models and automating inference in those models. Despite the different of DNNs and probabilistic models in mathematical formulations1, a similar modularity structure exist in both probabilistic models and DNNs. However, the benefits of modularity has not been fully exploited in currently popular PPLs such as Stan (Carpenter et al., 2017), Edward (Tran et al., 2016), PyMC (Salvatier et al., 2016), which is mostly about re-using plain probabilistic distributions and (configurable) generic inference algorithms such as Markov Chain Monte Carlo (MCMC) or Stochastic Variational Inference (SVI).
In this paper, we propose to further exploit the modularity in probabilistic models by offering not only plain probabilistic distributions but also sophisticated probabilistic model such as Bayesian
1A DNN represents a deterministic function, while a probabilistic model describes a probabilistic distribution.
non-parametric models as fundamental building blocks in a PPL. This would bring the a lot of benefits of modularities that we enjoy with deep learning libraries and improve the usabilities of PPLs on real world problems. We demonstrate this idea by presenting a modular probabilistic programming language MXFusion, which includes a new type of re-usable building blocks for probabilistic models called probabilistic modules. A probabilistic module consists of a set of random variables with associated probabilistic distributions and dedicated inference methods designed specifically for efficiency and accuracy on that set of random variables. This allows to package a sophisticated probabilistic model with their proposed inference methods. The reason for packaging inference methods within a probabilistic module is due to the fact that the exact inference of most of probabilistic models are intractable. Given a specific probabilistic model, a generic approximate inference method that is applicable to a wide range of probabilistic models often performs worse than a dedicated inference method. Most of current PPLs only rely on generic inference methods such as MCMC or black-box variational inference. We propose a framework of probabilistic modules based on variational inference. With the proposed framework, one could seamlessly combine generic inference methods with dedicated inference methods, which bridges the performance gap.
The rest of the paper is organized as follows: Section 2 describes relevant ideas in the literature of PPL; Section 3 provides details of probabilistic modules; Section 4 and ?? presents probabilistic modules in MXFusion; Section 5 demonstrates experiments of presented models with real data.

heading: 2 RELATED WORKS

text:
Probabilistic programming languages typically focus on offering an expressive programming language that describes probabilistic models and automating inference of probabilistic models through a “compilation” step. Due to the intractability of inference in probabilistic models, existing PPLs either focus on the expressiveness and employs a generic inference engine which suffers from scalability issues (Pfeffer, 2001; 2009; Goodman et al., 2012), or focus on efficiency by restricting down to a specific class of models and inference algorithms (Spiegelhalter et al., 1995; Murphy, 2001; Salvatier et al., 2015; Carpenter et al., 2016). By exploiting automatic-differentiation and hardware acceleration, Edward (Tran et al., 2017) and Pyro (Goodman, 2017) take a middle approach by offering the capability of customized inference methods, which is often referred to as an inference model. Both propose a clear separation between model and inference and typically rely on stochastic variational inference methods such as amortized inference.
The idea of encapsulating models and approximate inference programs in probabilistic modules has also been proposed by Cusumano-Towner & Mansinghka (2017), in which they demonstrate with Markov Chain Monte Carlo (MCMC) and Sequential Monte Carlo (SMC) inference programs. Venture and Anglican define inference as a collection of program fragments corresponding to local inference problems, although they do not support customizable posterior approximation (Mansinghka et al., 2014; Wood et al., 2014).

heading: 3 MODULARITY IN PROBABILISTIC PROGRAMMING

text:
DNN has a very nice modularity property that is exploited by deep learning libraries, which is a result of function composition. The formulation of a function can be broken into the composition of several sub-functions, e.g., f(x) = g1(· · · gk(x)). This enables a deep learning library to implement building blocks at various granularities as individual sub-functions. Then, a user can easily construct a DNN by putting these building blocks via function composition.
A similar modularity property also exists in probabilistic models, but in sightly different form. The joint probability distribution of a probabilistic model can often be decomposed into the product of the conditional distribution of individual variables with all the latent variables being marginalized out, e.g., p(y|z) = ∫ x p(y|x)p(x|z). This allows us to implement these individual conditional distributions as building blocks for a probabilistic model. So far, most of PPLs exploits this modularity and provides standard probabilistic distribution as building blocks. In this paper, we present an approach to provide more sophisticated probabilistic models as building blocks.
Inference for a DNN is mostly based on gradient optimization, in which a gradient can be computed transparently through automatic-differentiation. For a probabilistic model, inference is very chal-
lenging, as the exact inference for most of probabilistic models is intractable. Research has been focused on developing approximate inference or sampling methods for resolving this intractability. The literature of approximate inference methods can be broadly categorized into two groups: (1) proposing a better approximate inference method for a specific probabilistic model (2) developing a generic inference algorithm that can be applied to a family of probabilistic models as big as possible, sometimes being referred to as black-box inference. Both types of inference methods have received great attention and been developed by a lot of researchers. Obviously, given a specific model, a specialized inference method will always give a better performance than a generic inference method, as it is able to take advantages of the specific mathematical properties of the given model. Even among generic inference methods, a method only applicable to a small family of models performs better than a more generic method. For example, stochastic variational inference with reparameterization trick that is only applicable to a location-scale family of distributions produces a lower variance estimate of gradient than the score function approach that is applicable to mostly of distribution including discrete distributions.
Currently existing PPLs only use generic inference methods. Although those generic inference methods may be configurable, e.g., providing a customized variational posterior, it often performs worse for the probabilistic models, of which a specialized inference method exists. In MXFusion, we bridge the performance gap between generic and specialized inference methods by further exploiting the modularity of probabilistic models and offering probabilistic modules as building blocks. A probabilistic module encapsulates a sophisticated probabilistic model, e.g., a Gaussian process, and a set of specialized inference algorithms. With the latent variables of a probabilistic module being invisible from outside, a probabilistic module defines a conditional distribution over its exposed variables, which has no difference from a plain probabilistic distribution. A probabilistic module offers the same interface as a plain probabilistic distribution, so that an external inference method can transparently treat a probabilistic module as a plain probabilistic distribution. Internally, the interface of a probabilistic module is realized by calling a pre-specified inference method.
In the following section, we demonstrate the idea of a probabilistic module with variational inference. This idea can also be applied to other inference algorithms.

heading: 3.1 PROBABILISTIC MODULE WITH VARIATIONAL INFERENCE

text:
The concept of hiding the details of inference of a probabilistic module by specialized inference methods is nice. The challenge is that not all the approximate inference method are compatible with each other. In this section, we present an approach to implement probabilistic module with the framework of variational inference.
The main idea of variational inference is to approximate an intractable posterior distribution of latent variables with a parametric approximation, referred to as a variational posterior distribution. VI is often framed as a lower bound of the logarithm of the marginal distribution, e.g,
log p(y|z) = log ∫ x p(y|x)p(x|z) ≥ ∫ x q(x|y, z) log p(y|x)p(x|z) q(x|y, z) = L(y, z), (1)
where p(y|x)p(x|z) forms a probabilistic model with x as a latent variable, q(x|y) is the variational posterior distribution, and the lower bound is denoted as L(y, z). By then taking a natural exponentiation of L(y, z), we get a lower bound of the marginal probability denoted as p̃(y|z) = eL(y,z). Assume we are interested in plugging p(y|z) into another probabilistic model p(l|y)p(y|z) where y is a latent variable. With variational inference, the lower bound of the overall model can be derived as
log p(l|z) ≥ ∫ y q(y) log p(l|y)p(y|z) q(y) ≥ ∫ y q(y) log p(l|y)p̃(y|z) q(y) , (2)
where the usual variational lower bound is further lower bounded by replacing p(y|z) with its lower bound p̃(y|z). By substituting (1) into (2), we then derive the variational lower bound for the whole model,
log p(l|z) ≥ ∫ y,x q(x|y, z)q(y|z) log p(l|y)p(y|x)p(x|z) q(x|y, z)q(y|z) = L(l, z). (3)
This example shows that variational inference has a recursive property that enables inference modularity. A technical challenge with VI is that the integral of the lower bound of a probabilistic module
with respect to external latent variables, such as (2), may not always be tractable. Stochastic variational inference (SVI) offers an approximated solution to this new intractability by applying Monte Carlo Integration. Monte Carlo Integration is applicable to generic probabilistic distributions and lower bounds as long as we are able to draw samples from variational posterior.
In this case, the lower bound is approximated as
L(l, z) ≈ 1 N ∑ i log p(l|yi)eL(yi,z) q(yi|z) , L(yi, z) ≈ 1 M ∑ j log p(yi|xj)p(xj |z) q(xj |yi, z) , (4)
where yi|z ∼ q(y|z), xj |yi, z ∼ q(x|yi, z) andN is the number of samples of y andM is the number of samples of x given y. Note that if there is a closed form solution of p̃(yi|z), the calculation of L(yi, z) can be replaced with the closed-form solution. MXFusion offers modularity via probabilistic modules, which combine the definition of probabilistic distributions and specialized inference methods in a concise interface. If using VI as the primary inference algorithm, probabilistic modules used in a model automatically compute variational lower bounds, such as p̃(y|z) in the above example.

heading: 4 MODEL CONSTRUCTION WITH PROBABILISTIC MODULES

text:
In PPLs, a probabilistic model is often represented as a graph of random variables, where connections between random variables denotes corresponding probabilistic distributions. In such a representation, a probabilistic distribution is a fundamental re-usable component for constructing various probabilistic models. A probabilistic distribution typically have two reusable operations required by an inference method: drawing a sample of the random variable and computing the logarithm of its probabilistic density/mass function. Then, a user specifies a probabilistic model by defining individual variables and their corresponding probabilistic distributions if they are random variables. Each specified probabilistic distribution includes the form of the distribution and all the variables that the distribution is conditioned on, which corresponds to all the directed connections towards the random variable.
MXFusion follows this design of model specification. Figure 1 implements a Bayesian linear regression (BLR) model using a deep neural network as a feature extractor. There are N data points, each of which contains an input vector xn ∈ RQ and an output variable yn ∈ R. Each input vector is first applied to a DNN, hn = fφ(xn), where φ denotes the parameters of the DNN. The output variable, y = (y1, . . . , yN ), follows a multi-variate normal distribution, y ∼ N (0,H>H+σ2I), where H is the stack of the outcome of DNN, H = (h1, . . . ,hN )>. As there are no latent variables in BLR, a MAP inference method is created to estimate the model parameters with maximum likelihood given the observed data, i.e., φ∗, σ∗ = argmaxφ,σ log p(y|fφ(X), σ). The maximum likelihood estimate is achieved through optimizing the log-likelihood via a gradient-based optimization algorithm. The log-likelihood of BLR is computed via the log pdf method of the probabilistic distribution of the variable y, which is a multi-variate normal distribution. The gradient of the log-likelihood is obtained through auto-differentiation using MXNet.
A noticeable difference from PyMC and Edward is that all the variables in a probabilistic model are defined as explicit members of a Model class, which offers a clean memory management by
avoiding any global states. Like Edward, the information about which variables are observed is not given until inference time. This allows reuse a model specification across multiple inference methods. When applying a machine learning model to a problem, one almost always uses a model with at least two inference methods that have different observed variables.: one for estimating model parameters/posterior distribution from training data and the other for making prediction with new data. Taking BLR as an example, during training, the input variable X and output variable y are both observed and the aim of inference is to estimate the model parameters that maximizes the loglikelihood, while at prediction stage, only the input variable X is observed and the inference is to estimate the output variable with the optimal model parameters p(y|X, σ∗, θ∗).

heading: 4.1 PROBABILISTIC MODULE AS MODEL COMPONENT

text:
A stronger modularity for PPL is to enable reusing probabilistic models as building blocks. Ideally, one should be able to construct a probabilistic model by putting together well-known sophisticated probabilistic models such as Gaussian process, just like building a deep neural work with convolutional layers or LSTM layers. MXFusion offers such modularity by encapsulating the specification of a probabilistic model and a set of specialized inference algorithms, which is called a probabilistic module. A probabilistic module implements the same interface that a probabilistic distribution has, i.e., drawing a sample of output random variables and computing the logarithm of the probabilistic density/mass function. By following the same interface, a probabilistic module can be transparently used by a probabilistic model just like an usual probabilistic distribution.
In Figure 2, we demonstrate the convenience of a probabilistic module by modifying a BLR into deep kernel learning (DKL) (Wilson et al., 2016). This only requires to change one line of code, which replaces a multi-variate normal distribution with a sparse Gaussian process regression (SGPR) module. The syntax of specifying a probabilistic module is similar to the one for a normal distribution. The input arguments include the input variable to the SGPR module and the choice of kernel function for GP, which is radius basis function (RBF) in this case, and the noise variance for a Gaussian likelihood distribution. The function returns a variable representing the output of the SGPR module, which is assigned as the variable m.y in this example. In (Wilson et al., 2016), a particular GP approximation, known as KISS-GP, is used for scalability, which is replaced by a variational sparse GP approximation (Titsias, 2009) in this example.
Variational sparse GP speeds up the computation of log-likelihood by replacing it with a variational lower bound. It introduces a set of auxiliary variables u ∈ RM , known as inducing data, and a set of corresponding inducing inputs Z ∈ RM×Q. The resulting variational lower bound is
log p(y|X,Z, θ) ≥ LSGP(y,X,Z, θ) = ∫ f ,u p(f |u,X,Z)q(u) log p(y|f)(( (((p(f |u,X,Z)p(u|Z) (((( (p(f |u,X,Z)q(u)
where q(f ,u) = p(f |u,X,Z)q(u) is the variational posterior. The above variational lower bound has a closed form solution with a computational complexity O(NM2), which is significantly lower than the cubic complexity of GP, O(N3). The log pdf method of the SGPR module compute the above variational lower bound. The MAP inference method of the DKL model finds the parameters of SGPR θ, Z and the parameters of DNN φ that maximize the lower bound of DKL, i.e., θ∗, φ∗,Z∗ = argmaxθ,φ,Z LSGP(y, fφ(X),Z, θ).
The variational inference method of SGPR transparently provides a specialized implementation that exploits all the knowledge about the SGPR module and delivers good accuracy and efficiency. To achieve the same task with a standard PPL, one needs to explicit the prior distribution of a sparse Gaussian process regression model and the corresponding variational posterior. Even by doing so, the same quality of inference is not achievable with a generic variational inference method. Stochastic variational inference with reparameterization tricks would be the best choice in this case. The inference result will be slow and has high variance. This is because: (1) The inference algorithm is not able to figure out the cancellation of p(f |u,X,Z) in both the nominator and denominator. A direct evaluation of p(f |u,X,Z) results into a cubic complexity O(N3), which fails to deliver any speed-up comparing with GP. (2) The variational distribution p(f |u,X,Z) is high dimensional RN and highly correlated. The variance of the Monte Carlo integration via sampling will be high. (3) The optimal value of q(u) in the variational lower bound has a closed form solution, which is exploited by variational sparse GP (Titsias, 2009). With SVI, one needs to explicitly optimize it with gradient optimization. The speed and quality of the two inference approaches reflects the gap between a generic inference algorithm and a specialized inference algorithm. The probabilistic module aims at bridging this gap.

heading: 4.2 NESTED VARIATIONAL INFERENCE

text:
A probabilistic module can be treated transparently as a probabilistic distribution. Then, it is straightforward to construct a probabilistic model consists of multiple probabilistic modules such as Deep GPs (Dai et al., 2016). In these models, some of the exposed variables of probabilistic modules is also latent variables. As shown in Section 3.1, as long as a variational inference method is used for the whole probabilistic model, the inference of individual probabilistic modules can be transparently handled.
Figure 3 implements Bayesian Gaussian process latent variable model (BGPLVM) (Titsias & Lawrence, 2010), which is an example of this kind of models. BGPLVM can be constructed by assign the input variable X a Gaussian distribution with zero-mean and unit-variance. As the input variable X is a latent variable, the marginal log-likelihood is not tractable anymore. A variational lower bound can be written as
log p(Y) ≥ ∫ X q(X) log p(Y|X)p(X) q(X) ≥ ∫ X q(X) log LSGP(y,X,Z, θ)p(X) q(X)
where p(Y|X) is a GP and q(X) = N (X|µ, diag(s)) is the variational posterior of X, which is assumed to be a Gaussian distribution with a diagonal covariance matrix. By applying a variational sparse GP (Titsias, 2009) approximation to p(Y|X), we further lower bound the original lower bound by replace p(Y|X) with the variational sparse GP lower bound mentioned above LSGP(y,X,Z, θ). For the expectation with respect to q(X), we can apply stochastic variational inference (SVI) by drawing samples from q(X). In this way, we result into a nested variational inference combining a generic inference method (SVI) with a specialized inference method (variational sparse GP). Following the same approach, it is also straight-forward to extend BGPLVM into a variationally auto-encoded GPLVM/deep GP (Dai et al., 2016) by parameterizing µ and s in q(X) as the outcome of a DNN.

heading: 5 EXPERIMENTS

text: In the previous examples, we demonstrate that various sophisticated probabilistic models can be easily constructed with a few lines of codes by using probabilistic modules. In this section, we evaluate the models that are constructed with probabilistic modules deliver good performance on real data.

heading: 5.1 GAUSSIAN PROCESS REGRESSION AND DEEP KERNEL LEARNING

text:
We evaluate the variational sparse GP (SGP) (Titsias, 2009) module implemented in MXFusion and a variant of deep kernel learning (DKL) (Wilson et al., 2016) shown in Figure 2, comparing with a standard sparse GP implementation from GPy (2012). All the GP models use RBF kernels. By taking the advantage of GPU acceleration that is provided by MXNet, it allows to scale to a significantly higher number of inducing points. For DKL, we use one hidden layer with 20 hidden units and tanh activation. It outputs a 5-dimensional representation and, then, feeds into our SGP module.
In this experiment, we compare all the methods on four standard regression benchmarks. Naval1 and Naval2 are the same regression dataset with two different regression targets, they have 11934 data points and 16 dimensional inputs. Kin8nm has 8192 data points and 8 dimensional inputs. We used 90% data points for training and 10% for testing. All the datasets are normalized column-wise. We use 50 inducing points for GPy implementation which is a standard choice in the literature, and use 50 and 3200 inducing points for our SGP module. The Adam optimizer (Kingma & Ba, 2015) was used for training in MXFusion, while GPy models are trained with L-BFGS. Results with 10-fold averages measured in terms of root mean squared error (RMSE) and test set log likelihood (TLL) are provided in Table 1. GPy and MXFusion with 50 inducing points give similar performance for both RMSE and TLL. With 3200 inducing points, it gives significantly better performance on all the datasets for both RMSE and TLL. For DKL, with a relative small number of inducing points (1000), it produces a comparable RMSE with the help of a neural network, but does not give as good TLL.

heading: 6 DISCUSSION

text:
Despite the expressiveness and flexibility provided by a PPL, a major limitation of PPL is that they rely on generic inference algorithms. Recent work such as Edward and Pyro support customizable inference algorithms, in particular, stochastic variational inference with a customized variational posterior. Once an inference algorithm is chosen, it remains the same across a probabilistic model. However, given a specific probabilistic model, e.g., a conjugate model, a specialized inference algorithm that exploits the mathematical properties of that particular model will always produce inference results that are as good or better than the generic inference in terms of both accuracy and efficiency. In practice, when doing probabilistic modeling there is a tradeoff between implementing a specialized inference algorithm for improved speed and performance at the cost of maintainability
and flexibility by introducing specialized code for each model. MXFusion aims at closing the gap between having specialized, highly performant algorithms and generic, easily maintained generic algorithms by introducing probabilistic modules. A probabilistic module consists of both a model and inference definition, defined together and wrapped up in a modular, plug-and-play package. This allows specialized inference algorithms for corresponding probabilistic modules to be smoothly integrated into the inference algorithm of a larger probabilistic model. By bringing modularity into probabilistic programming, MXFusion offers a flexible maintainable framework for doing probabilistic modeling while keeping the accuracy and efficiency of specialized inference algorithms.