We propose a simple algorithm to train stochastic neural networks to draw samples from given target distributions for probabilistic inference. Our method is based on iteratively adjusting the neural network parameters so that the output changes along a Stein variational gradient (Liu & Wang, 2016) that maximumly decreases the KL divergence with the target distribution. Our method works for any target distribution specified by their unnormalized density function, and can train any black-box architectures that are differentiable in terms of the parameters we want to adapt. As an application of our method, we propose an amortized MLE algorithm for training deep energy model, where a neural sampler is adaptively trained to approximate the likelihood function. Our method mimics an adversarial game between the deep energy model and the neural sampler, and obtains realistic-looking images competitive with the state-of-the-art results.
1 INTRODUCTION
Modern machine learning increasingly relies on highly complex probabilistic models to reason about uncertainty. A key computational challenge is to develop efficient inference techniques to approximate, or draw samples from complex distributions. Currently, most inference methods, including MCMC and variational inference, are hand-designed by researchers or domain experts. This makes it difficult to fully optimize the choice of different methods and their parameters, and exploit the structures in the problems of interest in an automatic way. The hand-designed algorithm can also be inefficient when it requires to make fast inference repeatedly on a large number of different distributions with similar structures. This happens, for example, when we need to reason about a number of observed datasets in settings like online learning, or need fast inference as inner loops for other algorithms such as maximum likelihood training. Therefore, it is highly desirable to develop more intelligent probabilistic inference systems that can adaptively improve its own performance to fully the optimize computational efficiency, and generalize to new tasks with similar structures.
Specifically, denote by p(x) a probability density of interest specified up to the normalization constant, which we want to draw sample from, or marginalize to estimate its normalization constant. We want to study the following problem:
Problem 1. Given a distribution with density p(x) and a function f(η; ξ) with parameter η and random input ξ, for which we only have assess to draws of the random input ξ (without knowing its true distribution q0), and the output values of f(η; ξ) and its derivative ∂ηf(η; ξ) given η and ξ. We want to find an optimal parameter η so that the density of the random output variable x = f(η; ξ) with ξ ∼ q0 closely matches the target density p(x).
Because we have no assumption on the structure of f(η; ξ) and the distribution of random input, we can not directly calculate the actual distribution of the output random variable x = f(η; ξ); this makes it difficult to solve Problem 1 using the traditional variational inference (VI) methods. Recall that traditional VI approximates p(x) using simple proposal distributions qη(x) indexed by parameter η, and finds the optimal η by minimizing KL divergence KL(qη || p) = Eqη [log(qη/p)], which requires to calculate the density qη(x) or its derivative that is not computable by our assump-
tion (even when the Monte Carlo gradient estimation and the reparametrization trick (Kingma & Welling, 2013) are applied).
In fact, it is this requirement of calculating qη(x) that has been the major constraint for the designing of state-of-the-art variational inference methods with rich approximation families; the recent successful algorithms (e.g., Rezende & Mohamed, 2015b; Tran et al., 2015; Ranganath et al., 2015, to name only a few) have to handcraft special variational families to ensure the computational tractability of qη(x) and simultaneously obtain high approximation accuracy, which require substantial mathematical insights and research effects. Methods that do not require to explicitly calculate qη(x) 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 will use the term wild variational inference to refer to new variants of variational methods that require no tractability qη(x), to distinguish with the black-box variational inference (Ranganath et al., 2014) which refers to methods that work for generic target distributions p(x) without significant model-by-model consideration (but still require to calculate the proposal density qη(x)).
A similar problem also appears in importance sampling (IS), where it requires to calculate the IS proposal density q(x) in order to calculate the importance weight w(x) = p(x)/q(x). However, there exist methods that use no explicit information of q(x), which, seemingly counter-intuitively, give better asymptotic variance or converge rates than the typical IS that uses the proposal information (e.g., Liu & Lee, 2016; Briol et al., 2015; Henmi et al., 2007; Delyon & Portier, 2014). Discussions on this phenomenon dates back to O’Hagan (1987), who argued that “Monte Carlo (that uses the proposal information) is fundamentally unsound” for violating the Likelihood Principle, and developed Bayesian Monte Carlo (O’Hagan, 1991) as an example that uses no information on q(x), yet gives better convergence rate than the typical Monte Carlo O(n−1/2) rate (Briol et al., 2015). Despite the substantial difference between IS and VI, these results intuitively suggest the possibility of developing efficient variational inference without calculating q(x) explicitly.
In this work, we propose a simple algorithm for Problem 1 by iteratively adjusting the network parameter η to make its output random variable changes along a Stein variational gradient direction (SVGD) (Liu & Wang, 2016) that optimally decreases its KL divergence with the target distribution. Critically, the SVGD gradient includes a repulsive term to ensure that the generated samples have the right amount of variability that matches p(x). In this way, we “amortize SVGD” using a neural network, which makes it possible for our method to adaptively improve its own efficiency by leveraging fast experience, especially in cases when it needs to perform fast inference repeatedly on a large number of similar tasks. As an application, we use our method to amortize the MLE training of deep energy models, where a neural sampler is adaptively trained to approximate the likelihood function. Our method, which we call SteinGAN, mimics an adversarial game between the energy model and the neural sampler, and obtains realistic-looking images competitive with the state-of-theart results produced by generative adversarial networks (GAN) (Goodfellow et al., 2014; Radford et al., 2015).
Related Work 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 data-driven proposals for (sequential) Monte Carlo methods (e.g., Paige & Wood, 2016), to name only a few. Most of these methods, however, require to explicitly calculate q(x) (or its gradient). One exception is a very recent paper (Ranganath et al., 2016) that avoids calculating q(x) using an idea related to Stein discrepancy (Gorham & Mackey, 2015; Liu et al., 2016; Oates et al., 2014; Chwialkowski et al., 2016). There is also a raising interest recently on a similar problem of “learning to optimize” (e.g., Andrychowicz et al., 2016; Daniel et al., 2016; Li & Malik, 2016), which is technically easier than the more general problem of “learning to sample”. In fact, we show that our algorithm reduces to “learning to optimize” when only one particle is used in SVGD.
Generative adversarial network (GAN) and its variants have recently gained remarkable success on generating realistic-looking images (Goodfellow et al., 2014; Salimans et al., 2016; Radford et al., 2015; Li et al., 2015; Dziugaite et al., 2015; Nowozin et al., 2016). All these methods are set up to train latent variable models (the generator) under the assistant of the discriminator. Our SteinGAN instead performs traditional MLE training for a deep energy model, with the help of a neural sampler that learns to draw samples from the energy model to approximate the likelihood
function; this admits an adversarial interpretation: we can view the neural sampler as a generator that attends to fool the deep energy model, which in turn serves as a discriminator that distinguishes the real samples and the simulated samples given by the neural sampler. This idea of training MLE with neural samplers was first discussed by Kim & Bengio (2016); one of the key differences is that the neural sampler in Kim & Bengio (2016) is trained with the help of a heuristic diversity regularizer based on batch normalization, while SVGD enforces the diversity in a more principled way. Another method by Zhao et al. (2016) also trains an energy score to distinguish real and simulated samples, but within a non-probabilistic framework (see Section 5 for more discussion). Other more traditional approaches for training energy-based models (e.g., Ngiam et al., 2011; Xie et al., 2016) are often based on variants of MCMC-MLE or contrastive divergence (Geyer, 1991; Hinton, 2002; Tieleman, 2008), and have difficulty generating realistic-looking images from scratch.
2 STEIN VARIATIONAL GRADIENT DESCENT (SVGD)
Stein variational gradient descent (SVGD) (Liu & Wang, 2016) is a general purpose Bayesian inference algorithm motivated by Stein’s method (Stein, 1972; Barbour & Chen, 2005) and kernelized Stein discrepancy (Liu et al., 2016; Chwialkowski et al., 2016; Oates et al., 2014). It uses an efficient deterministic gradient-based update to iteratively evolve a set of particles {xi}ni=1 to minimize the KL divergence with the target distribution. SVGD has a simple form that reduces to the typical gradient descent for maximizing log p when using only one particle (n = 1), and hence can be easily combined with the successful tricks for gradient optimization, including stochastic gradient, adaptive learning rates (such as adagrad), and momentum.
To give a quick overview of the main idea of SVGD, let p(x) be a positive density function on Rd which we want to approximate with a set of particles {xi}ni=1. SVGD initializes the particles by sampling from some simple distribution q0, and updates the particles iteratively by
xi ← xi + φ(xi), ∀i = 1, . . . , n, (1) where is a step size, and φ(x) is a “particle gradient direction” chosen to maximumly decrease the KL divergence between the distribution of particles and the target distribution, in the sense that
φ = arg max φ∈F { − d d KL(q[ φ] || p) ∣∣ =0 } , (2)
where q[ φ] denotes the density of the updated particle x′ = x + φ(x) when the density of the original particle x is q, and F is the set of perturbation directions that we optimize over. We choose F to be the unit ball of a vector-valued reproducing kernel Hilbert space (RKHS)Hd = H×· · ·×H with eachH associating with a positive definite kernel k(x, x′); note thatH is dense in the space of continuous functions with universal kernels such as the Gaussian RBF kernel.
Critically, the gradient of KL divergence in (2) equals a simple linear functional of φ, allowing us to obtain a closed form solution for the optimal φ. Liu & Wang (2016) showed that
− d d KL(q[ φ] || p) ∣∣ =0 = Ex∼q[Tpφ(x)], (3)
with Tpφ(x) = ∇x log p(x)>φ(x) +∇x · φ(x), (4) where Tp is considered as a linear operator acting on function φ and is called the Stein operator in connection with Stein’s identity which shows that the RHS of (3) equals zero if p = q:
Ep[Tpφ] = Ep[∇x log p>φ +∇x · φ] = 0. (5) This is a result of integration by parts assuming the value of p(x)φ(x) vanishes on the boundary of the integration domain.
Therefore, the optimization in (2) reduces to
D(q || p) def= max φ∈Hd {Ex∼q[Tpφ(x)] s.t. ||φ||Hd ≤ 1}, (6)
where D(q || p) is the kernelized Stein discrepancy defined in Liu et al. (2016), which equals zero if and only if p = q under mild regularity conditions. Importantly, the optimal solution of (6) yields a closed form
φ∗(x′) ∝ Ex∼q[∇x log p(x)k(x, x′) +∇xk(x, x′)].
Algorithm 1 Amortized SVGD for Problem 1 Set batch size m, step-size scheme { t} and kernel k(x, x′). Initialize η0. for iteration t do
Draw random {ξi}mi=1, calculate xi = f(ηt; ξi), and the Stein variational gradient ∆xi in (7). Update parameter η using (8), (9) or (10).
end for
By approximating the expectation under q with the empirical average of the current particles {xi}ni=1, SVGD admits a simple form of update:
xi ← xi + ∆xi, ∀i = 1, . . . , n, where ∆xi = Êx∈{xi}ni=1 [∇x log p(x)k(x, xi) +∇xk(x, xi)], (7)
and Êx∼{xi}ni=1 [f(x)] = ∑ i f(xi)/n. The two terms in ∆xi play two different roles: the term with the gradient ∇x log p(x) drives the particles toward the high probability regions of p(x), while the term with ∇xk(x, xi) serves as a repulsive force to encourage diversity; to see this, consider a stationary kernel k(x, x′) = k(x − x′), then the second term reduces to Êx∇xk(x, xi) = −Êx∇xik(x, xi), which can be treated as the negative gradient for minimizing the average similarity Êxk(x, xi) in terms of xi. Overall, this particle update produces diverse points for distributional approximation and uncertainty assessment, and also has an interesting “momentum” effect in which the particles move collaboratively to escape the local optima.
It is easy to see from (7) that ∆xi reduces to the typical gradient ∇x log p(xi) when there is only a single particle (n = 1) and ∇xk(x, xi) when x = xi, in which case SVGD reduces to the standard gradient ascent for maximizing log p(x) (i.e., maximum a posteriori (MAP)).
3 AMORTIZED SVGD: TOWARDS AN AUTOMATIC NEURAL SAMPLER
SVGD and other particle-based methods become inefficient when we need to repeatedly infer a large number different target distributions for multiple tasks, including online learning or inner loops of other algorithms, because they can not improve based on the experience from the past tasks, and may require a large memory to restore a large number of particles. We propose to “amortize SVGD” by training a neural network f(η; ξ) to mimic the SVGD dynamics, yielding a solution for Problem 1.
One straightforward way to achieve this is to run SVGD to convergence and train f(η; ξ) to fit the SVGD results. This, however, requires to run many epochs of fully converged SVGD and can be slow in practice. We instead propose an incremental approach in which η is iteratively adjusted so that the network outputs x = f(η; ξ) changes along the Stein variational gradient direction in (7) in order to decrease the KL divergence between the target and approximation distribution.
To be specific, denote by ηt the estimated parameter at the t-th iteration of our method; each iteration of our method draws a batch of random inputs {ξi}mi=1 and calculate their corresponding output xi = f(η; ξi) based on ηt; here m is a mini-batch size (e.g., m = 100). The Stein variational gradient ∆xi in (7) would then ensure that x′i = xi + ∆xi forms a better approximation of the target distribution p. Therefore, we should adjust η to make its output matches {x′i}, that is, we want to update η by
ηt+1 ← arg min η m∑ i=1 ||f(η; ξi)− x′i||22, where x′i = xi + ∆xi. (8)
See Algorithm 1 for the summary of this procedure. If we assume is very small, then (8) reduces to a least square optimization. To see this, note that f(η; ξi) ≈ f(ηt; ξi) + ∂ηf(ηt; ξi)(η − ηt) by Taylor expansion. Since xi = f(ηt; ξi), we have
||f(η; ξi)− x′i||22 ≈ ||∂ηf(ηt; ξi)(η − ηt)− ∆xi||22. As a result, (8) reduces to the following least square optimization:
ηt+1 ← ηt + ∆ηt, where ∆ηt = arg min δ m∑ i=1 ||∂ηf(ηt; ξi)δ −∆xi||22. (9)
Update (9) can still be computationally expensive because of the matrix inversion. We can derive a further approximation by performing only one step of gradient descent of (8) (or (9)), which gives
ηt+1 ← ηt + m∑ i=1 ∂ηf(η t; ξi)∆xi. (10)
Although update (10) is derived as an approximation of (8)-(9), it is computationally faster and we find it works very effectively in practice; this is because when is small, one step of gradient update can be sufficiently close to the optimum.
Update (10) also has a simple and intuitive form: (10) can be thought as a “chain rule” that backpropagates the Stein variational gradient to the network parameter η. This can be justified by considering the special case when we use only a single particle (n = 1) in which case ∆xi in (7) reduces to the typical gradient ∇x log p(xi) of log p(x), and update (10) reduces to the typical gradient ascent for maximizing Eξ[log p(f(η; ξ))], in which case f(η; ξ) is trained to maximize log p(x) (that is, learning to optimize), instead of learning to draw samples from p for which it is crucial to use Stein variational gradient ∆xi to diversify the network outputs.
Update (10) also has a close connection with the typical variational inference with the reparameterization trick (Kingma & Welling, 2013). Let qη(x) be the density function of x = f(η; ξ), ξ ∼ q0. Using the reparameterization trick, the gradient of KL(qη || p) w.r.t. η can be shown to be
∇ηKL(qη || p) = −Eξ∼q0 [∂ηf(η; ξ)(∇x log p(x)−∇x log qη(x))].
With {ξi} i.i.d. drawn from q0 and xi = f(η; ξi), ∀i, the standard stochastic gradient descent for minimizing the KL divergence is
ηt+1 ← ηt + ∑ i ∂ηf(η t; ξi)∆̃xi, where ∆̃xi = ∇x log p(xi)−∇x log qη(xi). (11)
This is similar with (10), but replaces the Stein gradient ∆xi defined in (7) with ∆̃xi. The advantage of using ∆xi is that it does not require to explicitly calculate qη , and hence admits a solution to Problem 1 in which qη is not computable for complex network f(η; ξ) and unknown input distribution q0. Further insights can be obtained by noting that
∆xi ≈ Ex∼q[∇x log p(x)k(x, xi) +∇xk(x, xi)] = Ex∼q[(∇x log p(x)−∇x log q(x))k(x, xi)] (12) = Ex∼q[(∆̃x)k(x, xi)],
where (12) is obtained by using Stein’s identity (5). Therefore, ∆xi can be treated as a kernel smoothed version of ∆̃xi.
4 AMORTIZED MLE FOR GENERATIVE ADVERSARIAL TRAINING
Our method allows us to design efficient approximate sampling methods adaptively and automatically, and enables a host of novel applications. In this paper, we apply it in an amortized MLE method for training deep generative models.
Maximum likelihood estimator (MLE) provides a fundamental approach for learning probabilistic models from data, but can be computationally prohibitive on distributions for which drawing samples or computing likelihood is intractable due to the normalization constant. Traditional methods such as MCMC-MLE use hand-designed methods (e.g., MCMC) to approximate the intractable likelihood function but do not work efficiently in practice. We propose to adaptively train a generative neural network to draw samples from the distribution during MLE training, which not only provides computational advantage, and also allows us to generate realistic-looking images competitive with, or better than the state-of-the-art generative adversarial networks (GAN) (Goodfellow et al., 2014; Radford et al., 2015) (see Figure 1-5).
Algorithm 2 Amortized MLE as Generative Adversarial Learning Goal: MLE training for energy model p(x|θ) = exp(−φ(x, θ)− Φ(θ)). Initialize η and θ. for iteration t do
Updating η: Draw ξi ∼ q0, xi = f(η; ξi); update η using (8), (9) or (10) with p(x) = p(x|θ). Repeat several times when needed. Updating θ: Draw a mini-batch of observed data {xi,obs}, and simulated data xi = f(η; ξi), update θ by (13).
end for
To be specific, denote by {xi,obs} a set of observed data. We consider the maximum likelihood training of energy-based models of form
p(x|θ) = exp(−φ(x, θ)− Φ(θ)), Φ(θ) = log ∫ exp(−φ(x, θ))dx,
where φ(x; θ) is an energy function for x indexed by parameter θ and Φ(θ) is the log-normalization constant. The log-likelihood function of θ is
L(θ) = 1
n n∑ i=1 log p(xi,obs|θ),
whose gradient is
∇θL(θ) = −Êobs[∂θφ(x; θ)] + Eθ[∂θφ(x; θ)],
where Êobs[·] and Eθ[·] denote the empirical average on the observed data {xi,obs} and the expectation under model p(x|θ), respectively. The key computational difficulty is to approximate the model expectation Eθ[·]. To address this problem, we use a generative neural network x = f(η; ξ) trained by Algorithm 1 to approximately sample from p(x|θ), yielding a gradient update for θ of form
θ ← θ + ∇̂θL(θ), ∇̂θL(θ) = −Êobs[∂θφ(x; θ)] + Êη[∂θφ(x; θ)], (13)
where Êη denotes the empirical average on {xi} where xi = f(η; ξi), {ξi} ∼ q0. As θ is updated by gradient ascent, η is successively updated via Algorithm 1 to follow p(x|θ). See Algorithm 2. We call our method SteinGAN, because it can be intuitively interpreted as an adversarial game between the generative network f(η; ξ) and the energy model p(x|θ) which serves as a discriminator: The MLE gradient update of p(x|θ) effectively decreases the energy of the training data and increases the energy of the simulated data from f(η; ξ), while the SVGD update of f(η; ξ) decreases the energy of the simulated data to fit better with p(x|θ). Compared with the traditional methods based on MCMC-MLE or contrastive divergence, we amortize the sampler as we train, which gives much faster speed and simultaneously provides a high quality generative neural network that can generate realistic-looking images; see Kim & Bengio (2016) for a similar idea and discussions.
5 EMPIRICAL RESULTS
We evaluated our SteinGAN on four datasets, MNIST, CIFAR-10, CelebA (Liu et al., 2015), and Large-scale Scene Understanding (LSUN) (Yu et al., 2015), on which we find our method tends to generate realistic-looking images competitive with, sometimes better than DCGAN (Radford et al., 2015) (see Figure 2 - Figure 3). Our code is available at https://github.com/DartML/ SteinGAN.
Model Setup In order to generate realistic-looking images, we define our energy model based on an autoencoder:
p(x|θ) ∝ exp(−||x−D(E(x; θ); θ)||), (14)
where x denotes the image. This choice is motivated by Energy-based GAN (Zhao et al., 2016) in which the autoencoder loss is used as a discriminator but without a probabilistic interpretation. We
assume f(η; ξ) to be a neural network whose input ξ is a 100-dimensional random vector drawn by Uniform([−1, 1]). The positive definite kernel in SVGD is defined by the RBF kernel on the hidden representation obtained by the autoencoder in (14), that is,
k(x, x′) = exp(− 1 h2 ||E(x; θ)− E(x′; θ)||2).
As it is discussed in Section 3, the kernel provides a repulsive force to produce an amount of variability required for generating samples from p(x). This is similar to the heuristic repelling regularizer in Zhao et al. (2016) and the batch normalization based regularizer in Kim & Bengio (2016), but is derived in a more principled way. We take the bandwidth to be h = 0.5 ×med, where med is the median of the pairwise distances between E(x) on the image simulated by f(η; ξ). This makes the kernel change adaptively based on both θ (through E(x; θ)) and η (through bandwidth h).
Some datasets include both images x and their associated discrete labels y. In these cases, we train a joint energy model on (x, y) to capture both the inner structure of the images and its predictive relation with the label, allowing us to simulate images with a control on which category it belongs to. Our joint energy model is defined to be
p(x, y|θ) ∝ exp { − ||x−D(E(x; θ); θ)|| −max[m, σ(y, E(x; θ))] } , (15)
where σ(·, ·) is the cross entropy loss function of a fully connected output layer. In this case, our neural sampler first draws a label y randomly according to the empirical counts in the dataset, and then passes y into a neural network together with a 100 × 1 random vector ξ to generate image x. This allows us to generate images for particular categories by controlling the value of input y.
Stabilization In practice, we find it is useful to modify (13) to be θ ← θ − Êobs[∇θφ(x, θ)] + (1− γ)Êη[∇θφ(x, θ)]. (16)
where γ is a discount factor (which we take to be γ = 0.7). This is equivalent to maximizing a regularized likelihood:
max θ {log p(x|θ) + γΦ(θ)}
where Φ(θ) is the log-partition function; note that exp(γΦ(θ)) is a conjugate prior of p(x|θ). We initialize the weights of both the generator and discriminator from Gaussian distribution N (0, 0.02), and train them using Adam (Kingma & Ba, 2014) with a learning rate of 0.001 for the generator and 0.0001 for the energy model (the discriminator). In order to keep the generator and discriminator approximately aligned during training, we speed up the MLE update (16) of the discriminator (by increasing its learning rate to 0.0005) when the energy of the real data batch is larger than the energy of the simulated images, while slow down it (by freezing the MLE update of θ in (16)) if the magnitude of the energy difference between the real images and the simulated images goes above a threshold of 0.5. We used the bag of architecture guidelines for stable training suggested in DCGAN (Radford et al., 2015).
Discussion The MNIST dataset has a training set of 60, 000 examples. Both DCGAN and our model produce high quality images, both visually indistinguishable from real images; see figure 1.
CIFAR-10 is very diverse, and with only 50,000 training examples. Figure 2 shows examples of simulated images by DCGAN and SteinGAN generated conditional on each category, which look equally well visually. We also provide quantitively evaluation using a recently proposed inception score (Salimans et al., 2016), as well as the classification accuracy when training ResNet using 50, 000 simulated images as train sets, evaluated on a separate held-out testing set never seen by the GAN models. Besides DCGAN and SteinGAN, we also evaluate another simple baseline obtained by subsampling 500 real images from the training set and duplicating them 100 times. We observe that these scores capture rather different perspectives of image generation: The inception score favors images that look realistic individually and have uniformly distributed labels; as a result, the inception score of the duplicated 500 images is almost as high as the real training set. We find that the inception score of SteinGAN is comparable, or slightly lower than that of DCGAN. On the other hand, the classification accuracy measures the amount information captured in the simulated image sets; we find that SteinGAN achieves the highest classification accuracy, suggesting that it captures more information in the training set.
Figure 3 and 4 visualize the results on CelebA (with more than 200k face images) and LSUN (with nearly 3M bedroom images), respectively. We cropped and resized both dataset images into 64×64.
6 CONCLUSION
We propose a new method to train neural samplers for given distributions, together with a new SteinGAN method for generative adversarial training. Future directions involve more applications and theoretical understandings for training neural samplers.

This paper considers the energy-based model interpretation of GAN, where the discriminator is an unnormalized model for the likelihood of a generative model p(x|theta) and the generator is a directed model that approximates this distribution. The generator is used to draw approximate negative phase samples that are used in stochastic maximum likelihood / contrastive divergence learning of the EBM / discriminator.
The main idea in the paper is to fit the generator by following the Stein variational gradient. In practice this gradient consists of the usual gradient provided by the discriminator with an added term that provides a repulsive force between the sampled data points to increase sample diversity.
The idea of using a kernel to push apart the sampled points is interesting, and will work in low dimensions, but it is hard to see how it can work in full scale images. For high dimensional samples x, the proposed kernel is unlikely to provide a useful distance measure between points. There are no convincing experiments in the paper that show otherwise. Specifically:
- There is no experiment that compares between standard GAN and GAN + repulsion, using the same architecture. (please address this in the rebuttal)
- If the Stein variational idea is taken literally, the right thing to do would be to fully optimize the generator at every step, and then taking a single optimization step on the discriminator. Instead, each is updated in turn, and the learning rates of both steps are adjusted to keep the two "in line".
- The kernel used to fit the generator is defined in the auto-encoder space of the discriminator, and thus depends on the discriminator parameters. The objective that is used to fit the generator thus changes at every step, and the procedure can no longer be interpreted as stochastic gradient descent with respect to any single well defined objective.
The authors obtain good results: The generated images clearly look better than those generated by DCGAN. However, their approach has a number of changes compared to DCGAN, so it is not clear where the improvement comes from. In addition, by now the DCGAN is no longer a very strong baseline, as various other techniques have been proposed.
Note: The use of phi for both the "particle gradient direction" and energy function is confusing


RECURRENT NORMALIZATION PROPAGATION
We propose an LSTM parametrization that preserves the means and variances of the hidden states and memory cells across time. While having training benefits similar to Recurrent Batch Normalization and Layer Normalization, it does not need to estimate statistics at each time step, therefore, requiring fewer computations overall. We also investigate the parametrization impact on the gradient flows and present a way of initializing the weights accordingly. We evaluate our proposal on language modelling and image generative modelling tasks. We empirically show that it performs similarly or better than other recurrent normalization approaches, while being faster to execute.
1 INTRODUCTION
Recurrent neural network have shown remarkably good performances for sequential modelling tasks including machine translation (Bahdanau et al., 2015), visual captioning (Xu et al., 2015; Yao et al., 2015) or question answering (Hermann et al., 2015). However, such models remain notoriously hard to train with gradient backpropagation. As the number of time steps in the input sequence increases, the contractive or expanding effects associated with the state-to-state transformation at each time step can shrink or grow exponentially, leading respectively to vanishing or exploding gradients (Hochreiter, 1991; Bengio et al., 1994; Pascanu et al., 2012). In particular, with gradient vanishing, states at a given time are not influenced by changes happening much earlier in the sequence, preventing the model from learning long-term dependencies.
While the long-term dependencies problem is unsolvable in absolute (Hochreiter, 1991; Bengio et al., 1994), different RNN parameterizations, such as LSTM or GRU (Hochreiter & Schmidhuber, 1997; Cho et al., 2014) can help mitigate it. Furthermore, the LSTM parametrization has been recently extended to include layer-wise normalization (Cooijmans et al., 2016; Ba et al., 2016), building upon Batch Normalization (BN) (Ioffe & Szegedy, 2015). By normalizing the hidden state distributions to a fix scale and shift through the different time steps, normalized LSTMs have been shown to ease training, resulting in a parametrization that converges faster than a standard LSTM.
However, normalized LSTM introduces extra-computations as it involves standardizing the hidden states, enforcing their means and variances at each time step. By contrast, we propose an LSTM reparametrization that allows by construction to cheaply preserve the normalization of the hidden states through time. Our approach can be seen as the recurrent counterpart to the recent normalization propagation applied in feed-forward network (Arpit et al., 2016). It results in faster training convergence similar to Layer Normalization (LN) and Recurrent Batch Normalization while requiring fewer operations per time step and generalizing naturally to variable length sequences.
In addition, we investigate the impact of our parametrization, and more generally of normalized LSTM, on the vanishing and exploding gradient problems. We observe that layer-wise normalization provides a direct way to orient LSTM behaviour toward either gradient explosion or vanishing, and therefore biases the LSTM either towards reliably storing bits of information throughout time or allowing it to be more sensitive to new input changes.
∗Associate Fellow, Canadian Institute For Advanced Research (CIFAR)
We empirically validate our proposal on character-level language modelling on the Penn Treebank corpus (Marcus et al., 1993) and on image generative modelling, applying our normalisation to the DRAW architecture (Gregor et al., 2015).
The paper is structured as follows: section 2 provides a brief overview of the Batch-Normalized LSTM, in section 3 we derive our Normalized LSTM, section 4 investigates the impact of such normalization on the gradient flow, section 5 presents some experimental results, and we conclude in section 5.
2 PRE-REQUISITES
2.1 BN-LSTM
Batch-Normalized Long Short-Term Memory (BN-LSTM) (Cooijmans et al., 2016) is a reparametrization of LSTM that takes advantage of Batch Normalization (BN) to address the Covariate Shift (Shimodaira, 2000) occurring between time steps. Changes in the LSTM output at one time-step are likely to cause correlated changes in the summed inputs of the sequence next timesteps. This Temporal Covariate Shift can slow down the training process as the parameters of the model must not only be updated to minimize the cost of the task at hand but also adapt to the changing distribution of the inputs. In other words, the latter time steps in an LSTM need to account for the shifting distribution of the previous hidden states.
BN-LSTM proposes to reduce this temporal covariate shift by fixing the mean and the variance at each time step, relying on the BN transform
BN(x; γ, β) = γ x− Ê[x]√ V̂ar[x] + + β (1)
where Ê[x], V̂ar[x] are the activation mean and variance estimated from the mini-batch samples. Given an input sequence X = (x1,x2, . . . ,xT ), the BN-LSTM defines a sequence of hidden states ht and memory cell states ct according to
ĩt f̃t õt g̃t  = BN(Wxxt; γx, βx) + BN(Whht−1; γh, βh) + b (2) ct = σ(̃it) tanh(g̃t) + σ(f̃t) ct−1 (3) ht = σ(õt) tanh(BN(ct; γc, βc)), (4)
where Wh ∈ Rdh×4dh ,Wx ∈ Rdx×4dh ,b ∈ R4dh and the initial states h0 ∈ Rdh , c0 ∈ Rdh are model parameters. σ is the logistic sigmoid function, and denotes the Hadamard product. Ba et al. (2016) latter extended this parametrization by estimating the normalizing statistics (Ê[x], V̂ar[x]) using the different feature channels rather than mini-batch samples in order to naturally generalize to variable length sequences.
2.2 NORMALIZATION PROPAGATION
While increasing the training convergence speed relatively to a standard LSTM (Cooijmans et al., 2016), BN-LSTM needs to perform more computations per sample as it requires to compute 3x the BN transform at each time step.
On the other hand, Normalization Propagation (Norm Prop) (Arpit et al., 2016) aims at preserve the normalization of the input throughout the network. Unlike BN, the normalization doesn’t rely on the statistics of the mini-batch. Instead, it is the structure of the network itself that maintains the normalization. We therefore propose an LSTM reparametrization that preserves the normalization through the different time steps in order to avoid those extra computation.
3 NORMALIZED LSTM
While Norm Prop properties are appealing for recurrent models, its application to LSTM is not straightforward due to the memory cell structure. In this section we show how to derive a LSTM reparametrization that preserves normalization of the state ht through time.
3.1 CONSTRUCTION OF THE NORMALIZED LSTM
Following (Arpit et al., 2016; Salimans & Kingma, 2016), we will attempt to ensure, through an analytical reparametrization, that several intermediate quantities in the computation remain approximately standardized. We first compensate for the distribution changes induced by the weight matrices in the gates and cell candidate gt computations
ĩt f̃t õt g̃t  = γx Wx||Wx,i||2xt + γh Wh||Wh,i||2ht−1 + b. (5) where ||W·,i||2 is the vector of L2-norm of each line of the matrix and γx and γh are the trainable rescaling factors that restore the representation power lost in the rescaling of the weight matrices. To preserve the constant error carousel mechanism of the LSTM, we use the usual cell update,
ct = σ(̃it) tanh(g̃t) + σ(f̃t) ct−1 (6) Let us now construct an approximate analytical estimate of Var(ct). The evolution of ct through time can bee seen as a geometric series, with σ(f̃t) as constant ratio. Since σ(·) is upper-bounded by (and in practice smaller than) 1, ct will converge in expectation to a fixed value. This is the reason why in BN-LSTM the mini-batch statistics converge to a fixed value after a few time steps (Cooijmans et al., 2016). Moreover, if we consider that ĩt, f̃t, g̃t and ct−1 are (as a rough approximation) independent1, we can use the variance product rule of two independent random variables X and Y
Var[XY ] = Var[X] Var[Y ] + Var[X]E[Y ]2 +Var[Y ]E[X]2 (7) to compute Var[ct]. Considering that E[tanh(g̃t)] ≈ 0 and assuming that the cell has converged i.e. Var[ct] = Var[ct−1], we have
Var[ct] = Var[tanh(g̃t)] Var[σ(̃it)] + E[σ(̃it)]2
1−Var[σ(f̃t)]− E[σ(f̃t)]2 (8)
We can therefore analytically or numerically compute the mean and variance of each of those elements, assuming that both input xt and hidden state ht−1 are independent drawn from N (0, 1)
E[it] = E[σ(γxzx + γhzh)] (9) Var[it] = Var[σ(γxzx + γhzh)] (10) E[gt] = E[tanh(γxzx + γhzh)] (11)
Var[gt] = Var[tanh(γxzx + γhzh)] (12) where zx, zh ∼ N (0, 1). The statistics of the gates ot and ft can be computed in a similar way. We can then compute the value to which Var[ct] converges. Using this variance estimate, we compensate ct in order to compute the next hidden state ht
ht = σ(õt) tanh ( γcct√ Var[ct] ) (13)
Since we assumed that Var[ht−1] = 1, to ensure that we need to correct for the variance induced by the product of the tanh with the output gate. Using again the variance product rule (equation 7) we obtain
Var[ht] = Var
[ tanh ( γcct√ Var[ct] )] (Var[σ(õt)] + E[σ(õt)]2) (14)
We can estimate this variance through similar computation than equation 12. Scaling ht with 1/ √ Var[ht] ensure that its variance is 1 and so the propagation is maintained throughout the recurrence. 1This assumption is strong, but we don’t have any easy way to model the covariance between those terms without estimating it from the data.
3.2 PROPOSED REPARAMETRIZATION
Using equations 5, 6 and 13, we propose the following reparametrization of the LSTM, simply called the Normalized LSTM
ĩt f̃t õt g̃t  = γx Wx||Wx,i||2xt + γh Wh||Wh,i||2ht−1 + b (15) ct = σ(̃it) tanh(g̃t) + σ(f̃t) ct−1 (16)
ht = 1√
Var[ht]
[ σ(õt) tanh ( γcct√ Var[ct] )] (17)
where Var[ct] and Var[ht] are computed using equations 8 and 14, respectively. Those two variances are estimated at the initialization of the network (eq. 10 to eq. 12), and are then kept fixed during the training as in Norp Prop. γx, γh and γc are parameters learned via gradient descent.
Note that the reparametrization of equation 15 is identical to Weight Normalization (Weight Norm) (Salimans & Kingma, 2016). The main difference comes from equation 17, where we compensate for the variance of ct, the tanh and σ(õt), which ensures a normalized propagation. Overall, this reparametrization is equivalent in spirit to the BN-LSTM, but it benefits from the same advantages that Norm Prop has over BN: There is no dependence on the mini-batch size and the computation is the same for training and inference. Also, the rescaling of the matrices Wx and Wh can be done before the recurrence, leading to computation time closer to a vanilla LSTM.
3.3 WEIGHTS INITIALIZATION
With such reparametrization of the weight matrices, one can think that the scale of the initialization of the weights doesn’t matter in the learning process anymore. It is actually true for the forward and backward computation of the layer
yi = aWi ||aWi||2 x = Wi ||Wi||2 x (18)
∂yi ∂x = aWi ||aWi||2 = Wi ||Wi||2
(19)
and since the variance of both forward and backward passes is fixed, using an initialization scheme such as Glorot (Glorot & Bengio, 2010) doesn’t make sense with Norm Prop. However, the update of the parameters is affected by their scale:
∂yi ∂Wij ∂L ∂yi = 1 ||Wi||2
[ xj − yi
Wij ||Wi||2 ] ∂L ∂yi
(20)
The scale of the parameters affect the learning rate of the layer: the bigger the weights, the smaller the update. This induces a regularization effect in Norm Prop that is also present in BN (Ioffe & Szegedy, 2015). However, this could possibly be an issue for such parametrization: different initializations lead to different learning rates, and it is true even with adaptive step rules, such as Adam (Kingma & Ba, 2014). Moreover, the parameters that are not normalized (such as γ and b) aren’t affected by this effect, and so they are not regularized. This is the reason why forcing the weight matrices to have a unit L2 norm of the lines, as proposed in Arpit et al. (2016), helps the training procedure.
To still benefit from the reduction of the learning rate, which is know to ease the optimization (Vogl et al., 1988), we propose to simply force the unit L2 norm of the lines of the matrices and combine it with a global learning rate decay schedule.
4 GRADIENT PROPAGATION IN NORMALIZED LSTM
In this section we study the gradient flow in the Normalized LSTM. Since this reparametrization is similar to the BN-LSTM, the analysis we do here can be transposed to the BN-LSTM case.
4.1 THE EXPLODING AND VANISHING GRADIENTS PROBLEM
Given an input sequence X = (x1,x2, . . . ,xT ), we consider a recurrent network, parametrized by θ, that defines a sequence of hidden states ht = fθ(ht−1,xt) and cost function L which evaluates the model performance on a given task. Such network is usually trained using backpropagation through time, where the backpropagation is applied on the time-unrolled model. The chain rule can be applied in order to compute the derivative of the loss L with respect to parameters θ.
∂L ∂θ
= ∑
1≤t≤T
∂Lt ∂θ
= ∑
1≤t≤T ∑ 1≤k≤t ∂Lt ∂hk ∂hk ∂ht ∂ht ∂θ . (21)
The factors ∂hk∂ht = ∏ k≤l≤t ∂hl ∂hl−1
transports the error “in time” from step t back to step k and are also the cause of vanishing or exploding gradient in RNN (Pascanu et al., 2012). Indeed, if the Jacobian ∂hl∂hl−1 has singular value different from 1, the factor ∂hk ∂ht
, which is a product of t − k Jacobian matrices will either explode or vanish.
4.2 GRADIENT OF THE NORMALIZED LSTM
To study the gradient propagation of the Normalized LSTM, we first need to derive it. Using equation 15-17, we can write the gradient of ht with respect to ht−1
at = 1√
Var[ht] tanh ( γcct√ Var[ct] ) (22)
∂ht ∂ht−1 = ∂ot ∂ht−1 at + ot ∂at ∂ht−1
[
∂it ∂ht−1 gt + it ∂gt ∂ht−1 + ∂ft ∂ht−1 ct−1
] (23)
As we can see in equation 23 with the normalization, the gradient depends not only on the derivative of the cell candidate, the gates and the output tanh, but also on on the variance of ht and ct.
If we assume that ht−1 and xt are independent, we can compute the variance of ct. Neglecting the weight matrices and the effect of the gates, we can write from equations 8 and 14
Var[ct] ≈ Var[gt] = Var[tanh(z)], z ∼ N (0, γ2x + γ2h) (24) Var[ht] = Var[tanh(z)], z ∼ N (0, γ2c (γ2x + γ2h)) (25)
In both cases, the variance depends explicitly on the value of the different γ: The bigger the γ, the higher the variance. Neglecting again the weight matrices, we can now write the equations of the cell candidates gt and the gates it,ot and ft with respect to ht−1
∂gt ∂ht−1 = ∂ tanh(g̃t) ∂g̃t ∂g̃t ∂ht−1
= ( 1− tanh(γxxt + γhht−1)2 ) γh (26)
∂it ∂ht−1 = ∂σ(̃it)
∂ ĩt
∂ ĩt ∂ht−1 = σ(γxxt + γhht−1)(1− σ(γxxt + γhht−1))γh (27)
The gradients of ot and ft can be computed similarly. The effect of the γ here is double: They appear both in the activation function, where they control the saturation regime, and γh also appears as a multiplicative term in the gradient. They should therefore be small enough to prevent the activation from saturating too much, but at the same time γh can’t be too small, because it can also make the gradients vanish. Putting it all together, we have
∂ht ∂ht−1 = ∂ot ∂õt γh at + ot ∂at ∂ãt γc√ Var[ct]
γh [ ∂it
∂ ĩt gt + it ∂gt ∂g̃t + ∂ft ∂ f̃t ct−1
] (28)
In this equations we can see that the different γ directly scale the gradient, and they also control the saturation of the activation functions. Bad initialization of γ could thus lead to saturation or explosion regimes. Figure 1 shows the norm of the gradient with respect to γx and γh in a simulated LSTM. As we can see, one important parameter is the ratio between γh and γx: They control most of the propagation of the gradients. If γx > γh, the network will focus more on the input and so the gradients will tend to vanish more. On the other hand, if γh > γx, the network will tend have less vanishing gradients, but will focus less on its inputs.
5 EXPERIMENTS
5.1 CHARACTER-LEVEL LANGUAGE MODELLING
The first task we explore is character-level language modelling on the Penn Treebank corpus (Marcus et al., 1993). The goal is to predict the next character of the sequence given the previous ones. We use the same splits as Mikolov et al. (2012) and the same training procedure as Cooijmans et al. (2016), i.e. we train on sequences of length 100, with random starting point. The model is a 1000 units LSTM followed by a Softmax classifier. We use orthogonal initialization for the weight matrices. Because Norm Prop requires normalized inputs, we multiply the one-hot inputs vector with an untrained but fixed orthogonal matrix. This tricks does not only help the optimization of Norm Prop, but also all other variants.
To compare the convergence properties of Norm Prop against LN and BN, we first ran experiments using Adam (Kingma & Ba, 2014) with learning rate 2e-3, exponential decay of 1e-3 and gradient clipping at 1.0. As explained in section 3.3, we rescale the matrices such that they have a unit norm on the lines. For Norm Prop, we use γx = γh = 2 and γc = 1, for LN all the γ = 1.0 and for BN all the γ = 0.1. The results are presented in Table 1 and in Figure 2.
To show the potential of Norm Prop against other state-of-the-art system, we followed Ha et al. (2016) and apply dropout on both the input and output layer (p = 0.1) and recurrent dropout inside the LSTM (p = 0.1). We also used the Batch Data Normalization scheme presented by Arpit et al. (2016), so we standardize each input example using the mini-batch statistics and use population statistics at inference time. Finally, we also reduce the learning rate decay to 1e-4, to compensate for the fact that a network with dropout needs more time to train. The results are presented in Table 2.
As we can see in Figure 2 and in Table 1, Norm Prop compares really well against the other reparametrization. Also Norm Prop is roughly 30 % computationally faster2 than BN and LN. LN shows better optimization performances, but also overfits more. We also see that both optimization and generalization are better than the ones from Weight Norm, which shows the importance of compensating for the variance of ct and ht. Moreover, although Norm Prop doesn’t combine well with
2The GPU used is a NVIDIA GTX 750.
dropout in feed-forward networks (Arpit et al., 2016), it works will with recurrent dropout, as we can see in Table 2. We believe it is because recurrent dropout is less affecting its output distribution than dropout in feed forward networks, because we copy the variable at the previous time step instead of setting it to 0. With such regularization, Norm Prop compares well with other state-of-the-art approaches.
5.2 DRAW
The second task we explore is a generative modelling task on binarized MNIST (Larochelle & Murray, 2011) using the Deep Recurrent Attentive Writer (DRAW) (Gregor et al., 2015) architecture. DRAW is a variational auto-encoder, where both encoder and decoder are LSTMs, and has two attention mechanisms to select where to read and where to write.
We use Jörg Bornschein’s implementation3, with the same hyper-parameters as Gregor et al. (2015), ie the read and write size are 2x2 and 5x5 respectively, the number of glimpses is 64, the LSTMs have 256 units and the dimension of z is 100. We use Adam with learning rate of 1e-2, exponential decay of 1e-3 and mini-batch size of 128. We use orthogonal initialization and force the norm of the lines of the matrices to be 1. For Norm Prop, we use γx = γh = γc = 0.5. The test variational bound for the first 100 epochs is presented in Figure 3.
As we can see in Figure 3, both Weight Norm and Norm Prop outperform the baseline network by a significant margin. Also, as expected, Norm Prop performs better than Weight Norm, showing one again the importance of the compensation of the variance of ct and ht. Table 3 shows the test variational bound after 200 epochs of training. Norm Prop also compares favorably against LN.
3https://github.com/jbornschein/draw
6 CONCLUSION
Based on the BN-LSTM, we have shown how to build a Normalized LSTM that is able to preserve the variance of its output at each time step, by compensating for the variance of the cell and the hidden state. Such LSTM can be seen as the Norm Prop version of the BN-LSTM, and thus benefits from the same advantages that Norm Prop has over BN, while being way faster to compute. Also, we propose a scheme to initialize the weight matrices that takes into account the reparametrization. Moreover, we have derived the gradients of this LSTM and pointed out the importance of the initialization of the rescaling parameters. We have validated the performances of the Normalized LSTM on two different tasks, showing similar performances than BN-LSTM and LN-LSTM, while being significantly faster in computation time. Also, unlike the feed-forward case, this architecture works well with recurrent dropout, leading to close to state-of-the-art performances on the character-level language modelling task.
Future work includes trying this architecture on more challenging tasks and also studying the impact of not keeping the variance estimates of the cell and the hidden states fixed during the learning process.
ACKNOWLEDGMENTS
Part of this work was funded by Samsung. We used Theano (Theano Development Team, 2016), Blocks and Fuel (van Merriënboer et al., 2015) for our experiments. We also want to thanks Caglar Gulcehre and Tim Cooijmans for the talks and Jörg Bornschein for his DRAW implementation.

The authors show how the hidden states of an LSTM can be normalised in order to preserve means and variances. The method’s gradient behaviour is analysed. Experimental results seem to indicate that the method compares well with similar approaches.
Points
1) The writing is sloppy in parts. See at the end of the review for a non-exhaustive list.
2) The experimental results show marginal improvements, of which the the statistical significance is impossible to asses. (Not completely the author’s fault for PTB, as they partially rely on results published by others.) Weight normalisation seems to be a viable alternative in the: the performance and runtime are similar. The implementation complexity of weight norm is, however, arguably much lower. More effort could have been put in by the authors to clear that up. In the current state, practitioners as well as researchers will have to put in more effort to judge whether the proposed method is really worth it for them to replicate.
3) Section 4 is nice, and I applaud the authors for doing such an analysis.
List of typos etc.
- maintain -> maintain
- requisits -> requisites
- a LSTM -> an LSTM
- "The gradients of ot and ft are equivalent to equation 25.” Gradients cannot be equivalent to an equation.
- “beacause"-> because
- One of the γx > γh at the end of page 5 is wrong.