Paper:
TWO METHODS FOR WILD VARIATIONAL INFERENCE
Variational inference provides a powerful tool for approximate probabilistic inference on complex, structured models. Typical variational inference methods, however, require to use inference networks with computationally tractable probability density functions. This largely limits the design and implementation of variational inference methods. We consider wild variational inference methods that do not require tractable density functions on the inference networks, and hence can be applied in more challenging cases. As an example of application, we treat stochastic gradient Langevin dynamics (SGLD) as an inference network, and use our methods to automatically adjust the step sizes of SGLD, yielding significant improvement over the hand-designed step size schemes.
1 INTRODUCTION
Probabilistic modeling provides a principled approach for reasoning under uncertainty, and has been increasingly dominant in modern machine learning where highly complex, structured probabilistic models are often the essential components for solving complex problems with increasingly larger datasets. A key challenge, however, is to develop computationally efficient Bayesian inference methods to approximate, or draw samples from the posterior distributions. Variational inference (VI) provides a powerful tool for scaling Bayesian inference to complex models and big data. The basic idea of VI is to approximate the true distribution with a simpler distribution by minimizing the KL divergence, transforming the inference problem into an optimization problem, which is often then solved efficiently using stochastic optimization techniques (e.g., Hoffman et al., 2013; Kingma & Welling, 2013). However, the practical design and application of VI are still largely restricted by the requirement of using simple approximation families, as we explain in the sequel.
Let p(z) be a distribution of interest, such as the posterior distribution in Bayesian inference. VI approximates p(z) with a simpler distribution q∗(z) found in a set Q = {qη(z)} of distributions indexed by parameter η by minimizing the KL divergence objective:
min η
{ KL(qη || p) ≡ Ez∼qη [log(qη(z)/p(z))] } , (1)
where we can get exact result p = q∗ if Q is chosen to be broad enough to actually include p. In practice, however, Q should be chosen carefully to make the optimization in (1) computationally tractable; this casts two constraints on Q: 1. A minimum requirement is that we should be able to sample from qη efficiently, which allows us to make estimates and predictions based on qη in placement of the more intractable p. The samples from qη can also be used to approximate the expectation Eq[·] in (1) during optimization. This means that there should exist some computable function f(η; ξ), called the inference network, which takes a random seed ξ, whose distribution is denoted by q0, and outputs a random variable z = f(η; ξ) whose distribution is qη .
2. We should also be able to calculate the density qη(z) or it is derivative in order to optimize the KL divergence in (1). This, however, casts a much more restrictive condition, since it requires us to use only simple inference network f(η; ξ) and input distributions q0 to ensure a tractable form for the density qη of the output z = f(η; ξ).
In fact, it is this requirement of calculating qη(z) that has been the major constraint for the design of state-of-the-art variational inference methods. The traditional VI methods are often limited to
using simple mean field, or Gaussian-based distributions as qη and do not perform well for approximating complex target distributions. There is a line of recent work on variational inference with rich approximation families (e.g., Rezende & Mohamed, 2015b; Tran et al., 2015; Ranganath et al., 2015, to name only a few), all based on handcrafting special inference networks to ensure the computational tractability of qη(z) while simultaneously obtaining high approximation accuracy. These approaches require substantial mathematical insights and research effects, and can be difficult to understand or use for practitioners without a strong research background in VI. Methods that allow us to use arbitrary inference networks without substantial constraints can significantly simplify the design and applications of VI methods, allowing practical users to focus more on choosing proposals that work best with their specific tasks.
We use the term wild variational inference to refer to variants of variational methods working with general inference networks f(η, ξ) without tractability constraints on its output density qη(z); this should be distinguished with the black-box variational inference (Ranganath et al., 2014) which refers to methods that work for generic target distributions p(z) without significant model-by-model consideration (but still require to calculate the proposal density qη(z)). Essentially, wild variational inference makes it possible to “learn to draw samples”, constructing black-box neural samplers for given distributions. This enables more adaptive and automatic design of efficient Bayesian inference procedures, replacing the hand-designed inference algorithms with more efficient ones that can improve their efficiency adaptively over time based on past tasks they performed.
In this work, we discuss two methods for wild variational inference, both based on recent works that combine kernel techniques with Stein’s method (e.g., Liu & Wang, 2016; Liu et al., 2016). The first method, also discussed in Wang & Liu (2016), is based on iteratively adjusting parameter η to make the random output z = f(η; ξ) mimic a Stein variational gradient direction (SVGD) (Liu & Wang, 2016) that optimally decreases its KL divergence with the target distribution. The second method is based on minimizing a kernelized Stein discrepancy, which, unlike KL divergence, does not require to calculate density qη(z) for the optimization thanks to its special form.
Another critical problem is to design good network architectures well suited for Bayesian inference. Ideally, the network design should leverage the information of the target distribution p(z) in a convenient way. One useful perspective is that we can view the existing MC/MCMC methods as (hand-designed) stochastic neural networks which can be used to construct native inference networks for given target distributions. On the other hand, using existing MC/MCMC methods as inference networks also allow us to adaptively adjust the hyper-parameters of these algorithms; this enables amortized inference which leverages the experience on past tasks to accelerate the Bayesian computation, providing a powerful approach for designing efficient algorithms in settings when a large number of similar tasks are needed.
As an example, we leverage stochastic gradient Langevin dynamics (SGLD) (Welling & Teh, 2011) as the inference network, which can be treated as a special deep residential network (He et al., 2016), in which important gradient information ∇z log p(z) is fed into each layer to allow efficient approximation for the target distribution p(z). In our case, the network parameter η are the step sizes of SGLD, and our method provides a way to adaptively improve the step sizes, providing speed-up on future tasks with similar structures. We show that the adaptively estimated step sizes significantly outperform the hand-designed schemes such as Adagrad.
Related Works The idea of amortized inference (Gershman & Goodman, 2014) has been recently applied in various domains of probabilistic reasoning, including both amortized variational inference
(e.g., Kingma & Welling, 2013; Rezende & Mohamed, 2015a) and date-driven designs of Monte Carlo based methods (e.g., Paige & Wood, 2016), to name only a few. Most of these methods, however, require to explicitly calculate qη(z) (or its gradient).
One well exception is a very recent work (Ranganath et al., 2016) that also avoids calculating qη(z) and hence works for general inference networks; their method is based on a similar idea related to Stein discrepancy (Liu et al., 2016; Oates et al., 2017; Chwialkowski et al., 2016; Gorham & Mackey, 2015), for which we provide a more detailed discussion in Section 3.2.
The auxiliary variational inference methods (e.g., Agakov & Barber, 2004) provide an alternative way when the variational distribution qη(z) can be represented as a hidden variable model. In particular, Salimans et al. (2015) used the auxiliary variational approach to leverage MCMC as a variational approximation. These approaches, however, still require to write down the likelihood function on the augmented spaces, and need to introduce an additional inference network related to the auxiliary variables.
There is a large literature on traditional adaptive MCMC methods (e.g., Andrieu & Thoms, 2008; Roberts & Rosenthal, 2009) which can be used to adaptively adjust the proposal distribution of MCMC by exploiting the special theoretical properties of MCMC (e.g., by minimizing the autocorrelation). Our method is simpler, more generic, and works efficiently in practice thanks to the use of gradient-based back-propagation. Finally, connections between stochastic gradient descent and variational inference have been discussed and exploited in Mandt et al. (2016); Maclaurin et al. (2015).
Outline Section 2 introduces background on Stein discrepancy and Stein variational gradient descent. Section 3 discusses two methods for wild variational inference. Section 4 discuss using stochastic gradient Langevin dynamics (SGLD) as the inference network. Empirical results are shown in Section 5.
2 STEIN’S IDENTITY, STEIN DISCREPANCY, STEIN VARIATIONAL GRADIENT
Stein’s identity Stein’s identity plays a fundamental role in our framework. Let p(z) be a positive differentiable density on Rd, and φ(z) = [φ1(z), · · · , φd(z)]> is a differentiable vector-valued function. Define∇z · φ = ∑ i ∂ziφ. Stein’s identity is
Ez∼p[〈∇z log p(z), φ(z)〉+∇z · φ(z)] = ∫ X ∇z · (p(z)φ(z))dx = 0, (2)
which holds once p(z)φ(z) vanishes on the boundary of X by integration by parts or Stokes’ theorem; It is useful to rewrite Stein’s identity in a more compact way:
Ez∼p[Tpφ(z)] = 0, with Tpφ def = 〈∇z log p, φ〉+∇z · φ, (3)
where Tp is called a Stein operator, which acts on function φ and returns a zero-mean function Tpφ(z) under z ∼ p. A key computational advantage of Stein’s identity and Stein operator is that they depend on p only through the derivative of the log-density ∇z log p(z), which does not depend on the cumbersome normalization constant of p, that is, when p(z) = p̄(z)/Z, we have ∇z log p(z) = ∇z log p̄(z), independent of the normalization constant Z. This property makes Stein’s identity a powerful practical tool for handling unnormalized distributions widely appeared in machine learning and statistics.
Stein Discrepancy Although Stein’s identity ensures that Tpφ has zero expectation under p, its expectation is generally non-zero under a different distribution q. Instead, for p 6= q, there must exist a φ which distinguishes p and q in the sense that Ez∼q[Tpφ(z)] 6= 0. Stein discrepancy leverages this fact to measure the difference between p and q by considering the “maximum violation of Stein’s identity” for φ in certain function set F :
D(q || p) = max φ∈F
{ Ez∼q[Tpφ(z)] } , (4)
where F is the set of functions φ that we optimize over, and decides both the discriminative power and computational tractability of Stein discrepancy. Kernelized Stein discrepancy (KSD) is a special
Stein discrepancy that takes F to be the unit ball of vector-valued reproducing kernel Hilbert spaces (RKHS), that is,
F = {φ ∈ Hd : ||φ||Hd ≤ 1}, (5)
where H is a real-valued RKHS with kernel k(z, z′). This choice of F makes it possible to get a closed form solution for the optimization in (4) (Liu et al., 2016; Chwialkowski et al., 2016; Oates et al., 2017):
D(q || p) = max φ∈Hd
{ Ez∼q[Tpφ(z)], s.t. ||φ||Hd ≤ 1 } , (6)
= √ Ez,z′∼q[κp(z, z′)], (7)
where κp(z, z′) is a positive definite kernel obtained by applying Stein operator on k(z, z′) twice:
κp(z, z ′) = T z
′
p (T zp ⊗ k(z, z′)), = sp(z)sp(z
′)k(z, z′) + sp(z)∇z′k(z, z′) + sp(z′)∇zk(z, z′) +∇z · (∇z′k(z, z′)), (8)
where sp(z) = ∇z log p(z) and T zp and T zp denote the Stein operator when treating k(z, z′) as a function of z and z′, respectively; here we defined T zp ⊗k(z, z′) = ∇x log p(x)k(z, z′)+∇xk(z, z′) which returns a d× 1 vector-valued function. It can be shown that D(q || p) = 0 if and only if q = p when k(z, z′) is strictly positive definite in a proper sense (Liu et al., 2016; Chwialkowski et al., 2016). D(q || p) can treated as a variant of maximum mean discrepancy equipped with kernel κp(z, z
′) which depends on p (which makes D(q || p) asymmetric on q and p). The form of KSD in (6) allows us to estimate the discrepancy between a set of sample {zi} (e.g., drawn from q) and a distribution p specified by∇z log p(z),
D̂2u({zi} || p) = 1 n(n− 1) ∑ i6=j [κp(zi, zj)], D̂2v({zi} || p) = 1 n2 ∑ i,j [κp(zi, zj)], (9)
where D̂2u(q || p) provides an unbiased estimator (hence called a U -statistic) for D2(q || p), and D̂2v(q || p), called V -statistic, provides a biased estimator but is guaranteed to be always nonnegative: D̂2v({zi} || p) ≥ 0.
Stein Variational Gradient Descent (SVGD) Stein operator and Stein discrepancy have a close connection with KL divergence, which is exploited in Liu & Wang (2016) to provide a general purpose deterministic approximate sampling method. Assume that {zi}ni=1 is a sample (or a set of particles) drawn from q, and we want to update {zi}ni=1 to make it “move closer” to the target distribution p to improve the approximation quality. We consider updates of form
zi ← zi + φ∗(zi), ∀i = 1, . . . , n, (10)
where φ∗ is a perturbation direction, or velocity field, chosen to maximumly decrease the KL divergence between the distribution of updated particles and the target distribution, in the sense that
φ∗ = arg max φ∈F { − d d KL(q[ φ] || p) ∣∣ =0 } , (11)
where q[ φ] denotes the density of the updated particle z′ = z + φ(z) when the density of the original particle z is q, and F is the set of perturbation directions that we optimize over. A key observation (Liu & Wang, 2016) is that the optimization in (11) is in fact equivalent to the optimization for KSD in (4); we have
− d d KL(q[ φ] || p) ∣∣ =0 = Ez∼q[Tpφ(z)], (12)
that is, the Stein operator transforms the perturbation φ on the random variable (the particles) to the change of the KL divergence. Taking F to be unit ball of Hd as in (5), the optimal solution φ∗ of (11) equals that of (6), which is shown to be (e.g., Liu et al., 2016)
φ∗(z′) ∝ Ez∼q[T zp k(z, z′)] = Ez∼q[∇z log p(z)k(z, z′) +∇zk(z, z′)].
Algorithm 1 Amortized SVGD and KSD Minimization for Wild Variational Inference for iteration t do
1. Draw random {ξi}ni=1, calculate zi = f(η; ξi), and the Stein variational gradient ∆zi in (13). 2. Update parameter η using (14) or (15) for amortized SVGD, or (17) for KSD minimization.
end for
By approximating the expectation under q with the empirical mean of the current particles {zi}ni=1, SVGD admits a simple form of update that iteratively moves the particles towards the target distribution,
zi ← zi + ∆zi, ∀i = 1, . . . , n, ∆zi = Êz∈{zi}ni=1 [∇z log p(z)k(z, zi) +∇zk(z, zi)], (13)
where Êz∼{zi}ni=1 [f(z)] = ∑ i f(zi)/n. The two terms in ∆zi play two different roles: the term with the gradient ∇z log p(z) drives the particles towards the high probability regions of p(z), while the term with ∇zk(z, zi) serves as a repulsive force to encourage diversity; to see this, consider a stationary kernel k(z, z′) = k(z − z′), then the second term reduces to Êz∇zk(z, zi) = −Êz∇zik(z, zi), which can be treated as the negative gradient for minimizing the average similarity Êzk(z, zi) in terms of zi.
It is easy to see from (13) that ∆zi reduces to the typical gradient ∇z log p(zi) when there is only a single particle (n = 1) and ∇zk(z, zi) when z = zi, in which case SVGD reduces to the standard gradient ascent for maximizing log p(z) (i.e., maximum a posteriori (MAP)).
3 TWO METHODS FOR WILD VARIATIONAL INFERENCE
Since the direct parametric optimization of the KL divergence (1) requires calculating qη(z), there are two essential ways to avoid calculating qη(z): either using alternative (approximate) optimization approaches, or using different divergence objective functions. We discuss two possible approaches in this work: one based on “amortizing SVGD” (Wang & Liu, 2016) which trains the inference network f(η, ξ) so that its output mimic the SVGD dynamics in order to decrease the KL divergence; another based on minimizing the KSD objective (9) which does not require to evaluate q(z) thanks to its special form.
3.1 AMORTIZED SVGD
SVGD provides an optimal updating direction to iteratively move a set of particles {zi} towards the target distribution p(z). We can leverage it to train an inference network f(η; ξ) by iteratively adjusting η so that the output of f(η; ξ) changes along the Stein variational gradient direction in order to maximumly decrease its KL divergence with the target distribution. By doing this, we “amortize” SVGD into a neural network, which allows us to leverage the past experience to adaptively improve the computational efficiency and generalize to new tasks with similar structures. Amortized SVGD is also presented in Wang & Liu (2016); here we present some additional discussion.
To be specific, assume {ξi} are drawn from q0 and zi = f(η; ξi) the corresponding random output based on the current estimation of η. We want to adjust η so that zi changes along the Stein variational gradient direction ∆zi in (13) so as to maximumly decrease the KL divergence with target distribution. This can be done by updating η via
η ← arg min η n∑ i=1 ||f(η; ξi)− zi − ∆zi||22. (14)
Essentially, this projects the non-parametric perturbation direction ∆zi to the change of the finite dimensional network parameter η. If we take the step size to be small, then the updated η by (14) should be very close to the old value, and a single step of gradient descent of (14) can provide a
good approximation for (14). This gives a simpler update rule: η ← η + ∑ i ∂ηf(η; ξi)∆zi, (15)
which can be intuitively interpreted as a form of chain rule that back-propagates the SVGD gradient to the network parameter η. In fact, when we have only one particle, (15) reduces to the standard gradient ascent for maxη log p(f(η; ξ)), in which fη is trained to “learn to optimize” (e.g., Andrychowicz et al., 2016), instead of “learn to sample” p(z). Importantly, as we have more than one particles, the repulsive term ∇zk(z, zi) in ∆zi becomes active, and enforces an amount of diversity on the network output that is consistent with the variation in p(z). The full algorithm is summarized in Algorithm 1.
Amortized SVGD can be treated as minimizing the KL divergence using a rather special algorithm: it leverages the non-parametric SVGD which can be treated as approximately solving the infinite dimensional optimization minq KL(q || p) without explicitly assuming a parametric form on q, and iteratively projecting the non-parametric update back to the finite dimensional parameter space of η. It is an interesting direction to extend this idea to “amortize” other MC/MCMC-based inference algorithms. For example, given a MCMC with transition probability T (z′|z) whose stationary distribution is p(z), we may adjust η to make the network output move towards the updated values z′ drawn from the transition probability T (z′|z). The advantage of using SVGD is that it provides a deterministic gradient direction which we can back-propagate conveniently and is particle efficient in that it reduces to “learning to optimize” with a single particle. We have been using the simple L2 loss in (14) mainly for convenience; it is possible to use other two-sample discrepancy measures such as maximum mean discrepancy.
3.2 KSD VARIATIONAL INFERENCE
Amortized SVGD attends to minimize the KL divergence objective, but can not be interpreted as a typical finite dimensional optimization on parameter η. Here we provide an alternative method based on directly minimizing the kernelized Stein discrepancy (KSD) objective, for which, thanks to its special form, the typical gradient-based optimization can be performed without needing to estimate q(z) explicitly.
To be specific, take qη to be the density of the random output z = f(η; ξ) when ξ ∼ q0, and we want to find η to minimize D(qη || p). Assuming {ξi} is i.i.d. drawn from q0, we can approximate D2(qη || p) unbiasedly with a U-statistics:
D2(qη || p) ≈ 1 n(n− 1) ∑ i 6=j κp(f(η; ξi), f(η; ξj)), (16)
for which a standard gradient descent can be derived for optimizing η:
η ← η − 2 n(n− 1) ∑ i 6=j ∂ηf(η; ξi)∇ziκp(zi, zj), where zi = f(η; ξi). (17)
This enables a wild variational inference method based on directly minimizing η with standard (stochastic) gradient descent. See Algorithm 1. Note that (17) is similar to (15) in form, but replaces ∆zi with a ∆̃zi ∝ − ∑ j : i 6=j ∇ziκp(zi, zj). It is also possible to use the V -statistic in (9), but we find that the U -statistic performs much better in practice, possibly because of its unbiasedness property.
Minimizing KSD can be viewed as minimizing a constrastive divergence objective function. To see this, recall that q[ φ] denotes the density of z′ = z + φ(z) when z ∼ q. Combining (11) and (6), we can show that
D2(q || p) ≈ 1
(KL(q || p)−KL(q[ φ] || p)).
That is, KSD measures the amount of decrease of KL divergence when we update the particles along the optimal SVGD perturbation direction φ given by (11). If q = p, then the decrease of KL
divergence equals zero and D2(q || p) equals zero. In fact, as shown in Liu & Wang (2016) KSD can be explicitly represented as the magnitude of a functional gradient of KL divergence:
D(q || p) = ∣∣∣∣∣∣ d dφ KL(q[φ] || p) ∣∣ φ=0 ∣∣∣∣∣∣ Hd ,
where q[φ] is the density of z = z +φ(z) when z ∼ q, and ddφF (φ) denotes the functional gradient of functional F (φ) w.r.t. φ defined in RKHSHd, and ddφF (φ) is also an element inH
d. Therefore, KSD variational inference can be treated as explicitly minimizing the magnitude of the gradient of KL divergence, in contract with amortized SVGD which attends to minimize the KL divergence objective itself.
This idea is also similar to the contrastive divergence used for learning restricted Boltzmann machine (RBM) (Hinton, 2002) (which, however, optimizes p with fixed q). It is possible to extend this approach by replacing z′ = z + φ(z) with other transforms, such as these given by a transition probability of a Markov chain whose stationary distribution is p. In fact, according the so called generator method for constructing Stein operator (Barbour, 1988), any generator of a Markov process defines a Stein operator that can be used to define a corresponding Stein discrepancy.
This idea is related to a very recent work by Ranganath et al. (2016), which is based on directly minimizing the variational form of Stein discrepancy in (4); Ranganath et al. (2016) assumes F consists of a neural network φτ (z) parametrized by τ , and find η by solving the following min-max problem:
min η max τ Ez∼q[Tpφτ (z)].
In contrast, our method leverages the closed form solution by taking F to be an RKHS and hence obtains an explicit optimization problem, instead of a min-max problem that can be computationally more expensive, or have difficulty in achieving convergence.
Because κp(x, x′) (defined in (8)) depends on the derivative ∇x log p(x) of the target distribution, the gradient in (17) depends on the Hessian matrix ∇2x log p(x) and is hence less convenient to implement compared with amortized SVGD (the method by Ranganath et al. (2016) also has the same problem). However, this problem can be alleviated using automatic differentiation tools, which be used to directly take the derivative of the objective in (16) without manually deriving its derivatives.
4 LANGEVIN INFERENCE NETWORK
With wild variational inference, we can choose more complex inference network structures to obtain better approximation accuracy. Ideally, the best network structure should leverage the special properties of the target distribution p(z) in a convenient way. One way to achieve this by viewing existing MC/MCMC methods as inference networks with hand-designed (and hence potentially suboptimal) parameters, but good architectures that take the information of the target distribution p(z) into account. By applying wild variational inference on networks constructed based on existing MCMC methods, we effectively provide an hyper-parameter optimization for these existing methods. This allows us to fully optimize the potential of existing Bayesian inference methods, significantly improving the result with less computation cost, and decreasing the need for hyper-parameter tuning by human experts. This is particularly useful when we need to solve a large number of similar tasks, where the computation cost spent on optimizing the hyper-parameters can significantly improve the performance on the future tasks.
Stochastic Gradient Langevin Dynamics We first take the original stochastic gradient Langevin dynamics (SGLD) algorithm (Welling & Teh, 2011) as an example. SGLD starts with a random initialization z0, and perform iterative update of form
zt+1 ← zt + ηt ∇z log p̂(zt; Mt) + √ 2ηt ξt, ∀t = 1, · · ·T, (18)
where log p̂(zt; Mt) denotes an approximation of log p(zt) based on, e.g., a random mini-batch Mt of observed data at t-th iteration, and ξt is a standard Gaussian random vector of the same size as z, and ηt denotes a (vector) step-size at t-th iteration; here “ ” denotes element-wise product. When running SGLD for T iterations, we can treat zT as the output of a T -layer neural network
parametrized by the collection of step sizes η = {ηt}Tt=1, whose random inputs include the random initialization z0, the mini-batch Mt and Gaussian noise ξt at each iteration t. We can see that this defines a rather complex network structure with several different types of random inputs (z0, Mt and ξt). This makes it intractable to explicitly calculate the density of zT and traditional variational inference methods can not be applied directly. But wild variational inference can still allow us to adaptively improve the optimal step-size η in this case.
General Langevin Networks Based on the original formula of SGLD, we proposed a more general langevin network structure, and each layer of the network has a form
zt+1 ← Atzt + h(BtBt>∇z log p̂(zt; Mt) +Btξt +Dt), ∀t = 1, · · ·T, (19)
where At, Bt and Dt are network parameters at t-th iteration(whose size is d× d, and d is the size of zt), and h(·) denotes a smooth element-wise non-linearity function; here ξt is still a standard gaussian random vector with the same size as z. With this more complex network, we can use fewer layers to construct more powerful back-box samplers.
5 EMPIRICAL RESULTS
5.1 SGLD INFERENCE NETWORK
We first test our algorithm with SGLD inference network with (18) formula on both a toy Gaussian mixture model and a Bayesian logistic regression example. We find that we can adaptively learn step sizes that significantly outperform the existing hand-designed step size schemes, and hence save computational cost in the testing phase. In particular, we compare with the following step size schemes, for all of which we report the best results (testing accuracy in Figure 3(a); testing likelihood in Figure 3(b)) among a range of hyper-parameters:
1. Constant Step Size. We select a best constant step size in {1, 2, 23, . . . , 229} × 10−6. 2. Power Decay Step Size. We consider t = 10a × (b + t)−γ where γ = 0.55, a ∈ {−6,−5, . . . , 1, 2}, b ∈ {0, 1, . . . , 9}. 3. Adagrad, Rmsprop, Adadelta, all with the master step size selected in {1, 2, 23, . . . , 229}× 10−6, with the other parameters chosen by default values.
Gaussian Mixture We start with a simple 1D Gaussian mixture example shown in Figure 2 where the target distribution p(z) is shown by the red dashed curve. We use amortized SVGD and KSD to optimize the step size parameter of the Langevin inference network in (18) with T = 20 layers (i.e., SGLD with T = 20 iterations), with an initial z0 drawn from a q0 far away from the target distribution (see the green curve in Figure 2(a)); this makes it critical to choose a proper step size to achieve close approximation within T = 20 iterations. We find that amortized SVGD and KSD allow us to achieve good performance with 20 steps of SGLD updates (Figure 2(b)-(c)), while the result of the best constant step size and power decay step-size are much worse (Figure 2(d)-(e)).
Bayesian Logistic Regression We consider Bayesian logistic regression for binary classification using the same setting as Gershman et al. (2012), which assigns the regression weights w with a Gaussian prior p0(w|α) = N (w,α−1) and p0(α) = Gamma(α, 1, 0.01). The inference is applied on the posterior of z = [w, logα]. We test this model on the binary Covertype dataset1 with 581,012 data points and 54 features.
To demonstrate that our estimated learning rate can work well on new datasets never seen by the algorithm. We partition the dataset into mini-datasets of size 50, 000, and use 80% of them for training and 20% for testing. We adapt our amortized SVGD/KSD to train on the whole population of the training mini-datasets by randomly selecting a mini-dataset at each iteration of Algorithm 1, and evaluate the performance of the estimated step sizes on the remaining 20% testing mini-datasets.
Figure 3 reports the testing accuracy and likelihood on the 20% testing mini-datasets when we train the Langevin network with T = 10, 50, 100 layers, respectively. We find that our methods outperform all the hand-designed learning rates, and allow us to get performance closer to the fully converged SGLD and SVGD with a small number T of iterations.
Figure 4 shows the testing accuracy and testing likelihood of all the intermediate results when training Langevin network with T = 100 layers. It is interesting to observe that amortized SVGD and KSD learn rather different behavior: KSD tends to increase the performance quickly at the first few iterations but saturate quickly, while amortized SVGD tends to increase slowly in the beginning and boost the performance quickly in the last few iterations. Note that both algorithms are set up to optimize the performance of the last layers, while need to decide how to make progress on the intermediate layers to achieve the best final performance.
5.2 GENERAL LANGEVIN INFERENCE NETWORK
We further test our algorithm with general Langevin inference network. We firstly construct one single layer general Langevin network to approach the posterior of Bayesian logistic regression parameters and we can achieve 74.58% average accuracy and −0.5216 average testing log-likelihood in 100 repeat experiments. This result proves the proposed general Langevin Inference Network is quite competitive and worth to explore. Moreover, we use it as a black-box sampler to approach more complicate Gaussian Mixture distributions.
Gaussian Mixture We consider 10 components Gaussian Mixture Models with mean and covariance matrix of each component randomly drawed from a uniform distribution, and we test our methods on different dimensions models.
We construct 6 layers of general Langevin networks as a black-box sampler, and our proposed two methods to train the black-box sampler to approximate the target distribution. Figure 5 shows our
1https://www.csie.ntu.edu.tw/˜cjlin/libsvmtools/datasets/binary.html
results on 50 dimension Gaussian Mixture case and figure 6 shows results of different dimensions of Gaussian Mixture. From the figures we can know that our proposed sampling structure is quite competive comparing with NUT sampler(Hoffman & Gelman, 2014), and these two variational inference methods can both train a good black-box sampler.
6 CONCLUSION
We consider two methods for wild variational inference that allows us to train general inference networks with intractable density functions, and apply it to adaptively estimate step sizes of stochastic gradient Langevin dynamics. More studies are needed to develop better methods, more applications and theoretical understandings for wild variational inference, and we hope that the two methods we discussed in the paper can motivate more ideas and studies in the field.

Review:
The authors propose two variational methods based on the theme of posterior approximations which may not have a tractable density. The first is from another ICLR submission on "amortized SVGD" (Wang and Liu, 2016), where here the innovation is in using SGLD as the inference network. The second is from a NIPS paper (Ranganath et al., 2016) on minimizing the Stein divergence with a parametric approximating family, where here the innovation is in defining their test functions to be an RKHS, obtaining an analytic solution to the inner optimization problem.
The methodology is incremental. Everything up to Section 3.2 is essentially motivation, background, or related work. The notion of a "wild variational approximation" was already defined in Ranganath et al. (2016), termed a "variational program". It would be useful for the authors to comment on the difference, if any.
Section 3.2 is at first interesting because it analytically solves the maximum problem that is faced in Ranganath et al. (2016). However, this requires use of a kernel which will certainly not scale in high dimensions, so it is then equivalent in practice to having chosen a very simple test function family. To properly scale to high dimensions would require a deeper kernel and also learning its parameters; this is not any easier than parameterizing the test function family as a neural network to begin with, which Ranganath et al. (2016) do.
Section 4 introduces a Langevin inference network, which essentially chooses the variational approximation as an evolving sequence of Markov transition operators as in Salimans et al. (2015). I had trouble understanding this for a while because I could not understand what they mean by inference network. None of it is amortized in the usual inference network sense, which is that the parameters are given by the output of a neural network. Here, the authors simple define global parameters of the SGLD chain which are used across all the latent variables (which is strictly worse?). (What then makes it an "inference network"?) Is this not the variational approximation used in Salimans et al. (2015), but using a different objective to train it?
The experiments are limited, on a toy mixture of Gaussians posterior and Bayesian logistic regression. None of this addresses the problems one might suspect on high-dimensional and real data, such as the lack of scalability for the kernel, the comparison to Salimans et al. (2015) for the Langevin variational approximation, and any note of runtime or difficulty of training.
Minor comments
+ It's not clear if the authors understood previous work on expressive variational families or inference networks. For example, they argue Rezende & Mohamed, 2015b; Tran et al., 2015; Ranganath et al., 2015 require handcrafted inference networks. However, all of them assume use of any neural network for amortized inference. None of them even require an inference network. Perhaps the authors mean handcrafted posterior approximations, which to some extent is true; however, the three mentioned are all algorithmic in nature: in Rezende & Mohamed (2015), the main decision choice is the flow length; Tran et al. (2015), the size of the variational data; Ranganath et al. (2015), the flow length on the auxiliary variable space. Each works well on different problems, but this is also true of variational objectives which admit intractable q (as the latter two consider, as does Salimans et al. (2015)). The paper's motivation could be better explained, and perhaps the authors could be clearer on what they mean by inference network.
+ I also recommend the authors not term a variational inference method based on the class of approximating family. While black box variational inference in Ranganath et al. (2014) assumes a mean-field family, the term itself has been used in the literature to mean any variational method that imposes few constraints on the model class.