Title: GRADIENT ESTIMATORS FOR IMPLICIT MODELS

Abstract: Implicit models, which allow for the generation of samples but not for point-wise evaluation of probabilities, are omnipresent in real-world problems tackled by machine learning and a hot topic of current research. Some examples include data simulators that are widely used in engineering and scientific research, generative adversarial networks (GANs) for image synthesis, and hot-off-the-press approximate inference techniques relying on implicit distributions. The majority of existing approaches to learning implicit models rely on approximating the intractable distribution or optimisation objective for gradient-based optimisation, which is liable to produce inaccurate updates and thus poor models. This paper alleviates the need for such approximations by proposing the Stein gradient estimator, which directly estimates the score function of the implicitly defined distribution. The efficacy of the proposed estimator is empirically demonstrated by examples that include gradient-free MCMC, meta-learning for approximate inference and entropy regularised GANs that provide improved sample diversity.

1. INTRODUCTION:

Modelling is fundamental to the success of technological innovations for artificial intelligence. A powerful model learns a useful representation of the observations for a specified prediction task, and generalises to unknown instances that follow similar generative mechanics. A well established area of machine learning research focuses on developing prescribed probabilistic models (Diggle & Gratton, 1984) , where learning is based on evaluating the probability of observations under the model. Implicit probabilistic models, on the other hand, are defined by a stochastic procedure that allows for direct generation of samples, but not for the evaluation of model probabilities. These are omnipresent in scientific and engineering research involving data analysis, for instance ecology, climate science and geography, where simulators are used to fit real-world observations to produce forecasting results. Within the machine learning community there is a recent interest in a specific type of implicit models, generative adversarial networks (GANs) (Goodfellow et al., 2014) , which has been shown to be one of the most successful approaches to image and text generation (Radford et al., 2016; Yu et al., 2017; Arjovsky et al., 2017; Berthelot et al., 2017) . Very recently, implicit distributions have also been considered as approximate posterior distributions for Bayesian inference, e.g. see Liu & Feng (2016) ; Wang & Liu (2016) ; Li & Liu (2016) ; Karaletsos (2016) ; Mescheder et al. (2017) ; Huszár (2017) ; Li et al. (2017) ; Tran et al. (2017) . These examples demonstrate the superior flexibility of implicit models, which provide highly expressive means of modelling complex data structures.

Whilst prescribed probabilistic models can be learned by standard (approximate) maximum likelihood or Bayesian inference, implicit probabilistic models require substantially more severe approximations due to the intractability of the model distribution. Many existing approaches first approximate the model distribution or optimisation objective function and then use those approximations to learn the associated parameters. However, for any finite number of data points there exists an infinite number of functions, with arbitrarily diverse gradients, that can approximate perfectly the objective function at the training datapoints, and optimising such approximations can lead to unstable training and poor results. Recent research on GANs, where the issue is highly prevalent, suggest that restricting the representational power of the discriminator is effective in stabilising training (e.g. see Arjovsky et al., 2017; Kodali et al., 2017) . However, such restrictions often intro- Figure 1 : A comparison between the two approximation schemes. Since in practice the optimiser only visits finite number of locations in the parameter space, it can lead to over-fitting if the neural network based functional approximator is not carefully regularised, and therefore the curvature information of the approximated loss can be very different from that of the original loss (shown in (a)). On the other hand, the gradient approximation scheme (b) can be more accurate since it only involves estimating the sensitivity of the loss function to the parameters in a local region.

duce undesirable biases, responsible for problems such as mode collapse in the context of GANs, and the underestimation of uncertainty in variational inference methods (Turner & Sahani, 2011) .

In this paper we explore approximating the derivative of the log density, known as the score function, as an alternative method for training implicit models. An accurate approximation of the score function then allows the application of many well-studied algorithms, such as maximum likelihood, maximum entropy estimation, variational inference and gradient-based MCMC, to implicit models. Concretely, our contributions include:

• the Stein gradient estimator, a novel generalisation of the score matching gradient estimator (Hyvärinen, 2005) , that includes both parametric and non-parametric forms; • a comparison of the proposed estimator with the score matching and the KDE plug-in estimators on performing gradient-free MCMC, meta-learning of approximate posterior samplers for Bayesian neural networks, and entropy based regularisation of GANs.

2. LEARNING IMPLICIT PROBABILISTIC MODELS:

Given a dataset D containing i.i.d. samples we would like to learn a probabilistic model p(x) for the underlying data distribution p D (x). In the case of implicit models, p(x) is defined by a generative process. For example, to generate images, one might define a generative model p(x) that consists of sampling randomly a latent variable z ∼ p 0 (z) and then defining x = f θ (z). Here f is a function parametrised by θ, usually a deep neural network or a simulator. We assume f to be differentiable w.r.t. θ. An extension to this scenario is presented by conditional implicit models, where the addition of a supervision signal y, such as an image label, allows us to define a conditional distribution p(x|y) implicitly by the transformation x = f θ (z, y). A related methodology, wild variational inference (Liu & Feng, 2016; Li & Liu, 2016) assumes a tractable joint density p(x, z), but uses implicit proposal distributions to approximate an intractable exact posterior p(z|x). Here the approximate posterior q(z|x) can likewise be represented by a deep neural network, but also by a truncated Markov chain, such as that given by Langevin dynamics with learnable step-size.

Whilst providing extreme flexibility and expressive power, the intractability of density evaluation also brings serious optimisation issues for implicit models. This is because many learning algorithms, e.g. maximum likelihood estimation (MLE), rely on minimising a distance/divergence/discrepancy measure D[p||p D ], which often requires evaluating the model density (c.f. Ranganath et al., 2016; Liu & Feng, 2016) . Thus good approximations to the optimisation procedure are the key to learning implicit models that can describe complex data structure. In the context of GANs, the Jensen-Shannon divergence is approximated by a variational lower-bound represented by a discriminator (Barber & Agakov, 2003; Goodfellow et al., 2014) . Related work for wild variational inference (Li & Liu, 2016; Mescheder et al., 2017; Huszár, 2017; Tran et al., 2017) uses a GAN-based technique to construct a density ratio estimator for q/p 0 (Sugiyama et al., 2009; 2012; Uehara et al., 2016; Mohamed & Lakshminarayanan, 2016) and then approximates the KL-divergence term in the variational lower-bound:

EQUATION

In addition, Li & Liu (2016) and Mescheder et al. (2017) exploit the additive structure of the KLdivergence and suggest discriminating between q and an auxiliary distribution that is close to q, making the density ratio estimation more accurate. Nevertheless all these algorithms involve a minimax optimisation, and the current practice of gradient-based optimisation is notoriously unstable.

The stabilisation of GAN training is itself a recent trend of related research (e.g. see Salimans et al., 2016; Arjovsky et al., 2017) . However, as the gradient-based optimisation only interacts with gradients, there is no need to use a discriminator if an accurate approximation to the intractable gradients could be obtained. As an example, consider a variational inference task with the approximate posterior defined as z ∼ q φ (z|x) ⇔ ∼ π( ), z = f φ ( , x). Notice that the variational lower-bound can be rewritten as

EQUATION

the gradient of the variational parameters φ can be computed by a sum of the path gradient of the first term (i.e.

E π ∇ f log p(x, f ( , x)) T ∇ φ f ( , x)

) and the gradient of the entropy term

∇ φ H[q(z|x)].

Expanding the latter, we have

EQUATION

in which the first term in the last line is zero (Roeder et al., 2017) . As we typically assume the tractability of ∇ φ f , an accurate approximation to ∇ z log q(z|x) would remove the requirement of discriminators, speed-up the learning and obtain potentially a better model. Many gradient approximation techniques exist (Stone, 1985; Fan & Gijbels, 1996; Zhou & Wolfe, 2000; De Brabanter et al., 2013) , and in particular, in the next section we will review kernel-based methods such as kernel density estimation (Singh, 1977) and score matching (Hyvärinen, 2005) in more detail, and motivate the main contribution of the paper.

3. GRADIENT APPROXIMATION WITH THE STEIN GRADIENT ESTIMATOR:

We propose the Stein gradient estimator as a novel generalisation of the score matching gradient estimator. Before presenting it we first set-up the notation. Column vectors and matrices are boldfaced. The random variable under consideration is x ∈ X with X = R d×1 if not specifically mentioned. To avoid misleading notation we use the distribution q(x) to derive the gradient approximations for general cases. As Monte Carlo methods are heavily used for implicit models, in the rest of the paper we mainly consider approximating the gradient g(x k ) := ∇ x k log q(x k ) for x k ∼ q(x), k = 1, ..., K.

We use x i j to denote the jth element of the ith sample x i . We also denote the matrix form of the collected gradients as

G := ∇ x 1 log q(x 1 ), • • • , ∇ x K log q(x K )

T ∈ R K×d , and its approximation

Ĝ := ĝ(x 1 ), • • • , ĝ(x K ) T with ĝ(x k ) = ∇ x k log q(x k

) for some q(x).

3.1. STEIN GRADIENT ESTIMATOR: INVERTING STEIN'S IDENTITY:

We start from introducing Stein's identity that was first developed for Gaussian random variables (Stein, 1972; 1981) then extended to general cases (Gorham & Mackey, 2015; Liu et al., 2016) . Let h : R d×1 → R d ×1 be a differentiable multivariate test function which maps x to a column vector h(x) = [h 1 (x), h 2 (x), ..., h d (x)] T . We further assume the boundary condition for h:

EQUATION

This condition holds for almost any test function if q has sufficiently fast-decaying tails (e.g. Gaussian tails). Now we introduce Stein's identity (Stein, 1981; Gorham & Mackey, 2015; Liu et al., 2016 )

EQUATION

in which the gradient matrix term

∇ x h(x) = (∇ x h 1 (x), • • • , ∇ x h d (x)) T ∈ R d ×d .

This identity can be proved using integration by parts: for the ith row of the matrix h(x)∇ x log q(x) T , we have

E q [h i (x)∇ x log q(x) T ] = h i (x)∇ x q(x) T dx = q(x)h i (x)| ∂X -q(x)∇ x h i (x) T dx = -E q [∇ x h i (x) T ]. (6)

Observing that the gradient term ∇ x log q(x) of interest appears in Stein's identity ( 5), we propose the Stein gradient estimator by inverting Stein's identity. As the expectation in ( 5) is intractable, we further approximate the above with Monte Carlo (MC):

EQUATION

with err ∈ R d ×d the random error due to MC approximation, which has mean 0 and vanishes as K → +∞. Now by temporarily denoting 7) can be rewritten as -1 K HG + err = ∇ x h. Thus we consider a ridge regression method (i.e. adding an 2 regulariser) to estimate G:

H = h(x 1 ), • • • , h(x K ) ∈ R d ×K , ∇ x h = 1 K K k=1 ∇ x k h(x k ) ∈ R d ×d , equation (

ĜStein:

EQUATION

with || • || F the Frobenius norm of a matrix and η ≥ 0. Simple calculation shows that ĜStein

EQUATION

where

K := H T H, K ij = K(x i , x j ) := h(x i ) T h(x j ), ∇, K := KH T ∇ x h, ∇, K ij = K k=1 ∇ x k j K(x i , x k ).

One can show that the RBF kernel satisfies Stein's identity (Liu et al., 2016) . In this case h(x) = K(x, •), d = +∞ and by the reproducing kernel property (Berlinet & Thomas-Agnan, 2011) 

, h(x) T h(x ) = K(x, •), K(x , •) H = K(x, x ).

3.2. STEIN GRADIENT ESTIMATOR MINIMISES THE KERNELISED STEIN DISCREPANCY:

In this section we derive the Stein gradient estimator again, but from a divergence/discrepancy minimisation perspective. Stein's method also provides a tool for checking if two distributions q(x) and q(x) are identical. If the test function set H is sufficiently rich, then one can define a Stein discrepancy measure by S(q, q) := sup

EQUATION

see Gorham & Mackey (2015) for an example derivation. When H is defined as a unit ball in an RKHS induced by a kernel K(x, •), Liu et al. (2016) and Chwialkowski et al. (2016) showed that the supremum in ( 10) can be analytically obtained as (with K xx shorthand for K(x, x )):

EQUATION

) which is also named the kernelised Stein discrepancy (KSD). Chwialkowski et al. (2016) showed that for C 0 -universal kernels satisfying the boundary condition, KSD is indeed a discrepancy measure: S 2 (q, q) = 0 ⇔ q = q. Gorham & Mackey (2017) further characterised the power of KSD on detecting non-convergence cases. Furthermore, if the kernel is twice differentiable, then using the same technique as to derive ( 16) one can compute KSD by

EQUATION

) In practice KSD is estimated with samples {x k } K k=1 ∼ q, and simple derivations show that the Vstatistic of KSD can be reformulated as

S 2 V (q, q) = 1 K 2 Tr( ĜT K Ĝ + 2 ĜT ∇, K ) + C. Thus the l 2 error in (8) is equivalent to the V-statistic of KSD if h(x) = K(x, •

), and we have the following:

EQUATION

One can also minimise the U-statistic of KSD to obtain gradient approximations, and a full derivation of which, including the optimal solution, can be found in the appendix. In experiments we use Vstatistic solutions and leave comparisons between these methods to future work.

3.3. COMPARISONS TO EXISTING KERNEL-BASED GRADIENT ESTIMATORS:

There exist other gradient estimators that do not require explicit evaluations of ∇ x log q(x), e.g. the denoising auto-encoder (DAE) (Vincent et al., 2008; Vincent, 2011; Alain & Bengio, 2014) which, with infinitesimal noise, also provides an estimate of ∇ x log q(x) at convergence. However, applying such gradient estimators result in a double-loop optimisation procedure since the gradient approximation is repeatedly required for fitting implicit distributions, which can be significantly slower than the proposed approach. Therefore we focus on "quick and dirty" approximations and only include comparisons to kernel-based gradient estimators in the following.

3.3.1. KDE GRADIENT ESTIMATOR: PLUG-IN ESTIMATOR WITH DENSITY ESTIMATION:

A naive approach for gradient approximation would first estimate the intractable density q(x) ≈ q(x) (up to a constant), then approximate the exact gradient by ∇ x log q(x) ≈ ∇ x log q(x). Specifically, Singh (1977) considered kernel density estimation (KDE)

q(x) = 1 K K k=1 K(x, x k ) × C.

, then differentiated through the KDE estimate to obtain the gradient estimator:

EQUATION

Interestingly for translation invariant kernels K(x, x ) = K(x -x ) the KDE gradient estimator (14) can be rewritten as ĜKDE = -diag (K1) -1 ∇, K . Inspecting and comparing it with the Stein gradient estimator (9), one might notice that the Stein method uses the full kernel matrix as the pre-conditioner, while the KDE method computes an averaged "kernel similarity" for the denominator. We conjecture that this difference is key to the superior performance of the Stein gradient estimator when compared to the KDE gradient estimator (see later experiments). The KDE method only collects the similarity information between x k and other samples x j to form an estimate of ∇ x k log q(x k ), whereas for the Stein gradient estimator, the kernel similarity between x i and x j for all i, j = k are also incorporated. Thus it is reasonable to conjecture that the Stein method can be more sample efficient, which also implies higher accuracy when the same number of samples are collected.

3.3.2. SCORE MATCHING GRADIENT ESTIMATOR: MINIMISING MSE:

The KDE gradient estimator performs indirect approximation of the gradient via density estimation, which can be inaccurate. An alternative approach directly approximates the gradient ∇ x log q(x) by minimising the expected 2 error w.r.t. the approximation ĝ

(x) = (ĝ 1 (x), • • • , ĝd (x))

T :

EQUATION

)

It has been shown in Hyvärinen (2005) that this objective can be reformulated as

EQUATION

The key insight here is again the usage of integration by parts: after expanding the 2 loss objective, the cross term can be rewritten as E q ĝ(x) T ∇ x log q(x) = -E q [ ∇, ĝ(x) ] , if assuming the boundary condition (4) for ĝ (see ( 6)). The optimum of ( 16) is referred as the score matching gradient estimator. The 2 objective ( 15) is also called Fisher divergence (Johnson, 2004) which is a special case of KSD (11) by selecting K(x, x ) = δ x=x . Thus the Stein gradient estimator can be viewed as a generalisation of the score matching estimator.

The comparison between the two estimators is more complicated. Certainly by the Cauchy-Schwarz inequality the Fisher divergence is stronger than KSD in terms of detecting convergence (Liu et al., 2016) . However it is difficult to perform direct gradient estimation by minimising the Fisher divergence, since (i) the Dirac kernel is non-differentiable so that it is impossible to rewrite the divergence in a similar form to (12), and (ii) the transformation to (16) involves computing ∇ x ĝ(x). So one needs to propose a parametric approximation to G and then optimise the associated parameters accordingly, and indeed Sasaki et al. (2014) and Strathmann et al. (2015) derived a parametric solution by first approximating the log density up to a constant as log q(x) := K k=1 a k K(x, x k ) + C, then minimising ( 16) to obtain the coefficients âscore k and constructing the gradient estimator as

Ĝscore:

EQUATION

)

Therefore the usage of parametric estimation can potentially remove the advantage of using a stronger divergence. Conversely, the proposed Stein gradient estimator ( 9) is non-parametric in that it directly optimises over functions evaluated at locations {x k } K k=1 . This brings in two key advantages over the score matching gradient estimator: (i) it removes the approximation error due to the use of restricted family of parametric approximations and thus can be potentially more accurate; (ii) it has a much simpler and ubiquitous form that applies to any kernel satisfying the boundary condition, whereas the score matching estimator requires tedious derivations for different kernels repeatedly (see appendix).

In terms of computation speed, since in most of the cases the computation of the score matching gradient estimator also involves kernel matrix inversions, both estimators are of the same order of complexity, which is O(K 3 + K 2 d) (kernel matrix computation plus inversion). Low-rank approximations such as the Nyström method (Smola & Schökopf, 2000; Williams & Seeger, 2001 ) can enable speed-up, but this is not investigated in the paper. Again we note here that kernel-based gradient estimators can still be faster than e.g. the DAE estimator since no double-loop optimisation is required. Certainly it is possible to apply early-stopping for the inner-loop DAE fitting. However the resulting gradient approximation might be very poor, which leads to unstable training and poorly fitted implicit distributions.

3.4. ADDING PREDICTIVE POWER:

Though providing potentially more accurate approximations, the non-parametric estimator (9) has no predictive power as described so far. Crucially, many tasks in machine learning require predicting gradient functions at samples drawn from distributions other than q, for example, in MLE q(x) corresponds to the model distribution which is learned using samples from the data distribution instead. To address this issue, we derive two predictive estimators, one generalised from the nonparametric estimator and the other minimises KSD using parametric approximations.

Predictions using the non-parametric estimator. Let us consider an unseen datum y. If y is sampled from q, then one can also apply the non-parametric estimator (9) for gradient approximation, given the observed data X = {x 1 , ..., x K } ∼ q. Concretely, if writing ĝ(y) ≈ ∇ y log q(y) ∈ R d×1 then the non-parametric Stein gradient estimator computed on

X ∪ {y} is ĝ(y) T Ĝ = -(K * + ηI) -1 ∇ y K(y, y) + K k=1 ∇ x k K(y, x k ) ∇, K + ∇ y K(•, y) , K * = K yy K yX K Xy K ,

with ∇ y K(•, y) denoting a K × d matrix with rows ∇ y K(x k , y), and ∇ y K(y, y) only differentiates through the second argument. Then we demonstrate in the appendix that, by simple matrix calculations and assuming a translation invariant kernel, we have (with column vector 1 ∈ R K×1 ):

EQUATION

In practice one would store the computed gradient ĜStein V , the kernel matrix inverse (K + ηI) -1 and η as the "parameters" of the predictive estimator. For a new observation y ∼ p in general, one can "pretend" y is a sample from q and apply the above estimator as well. The approximation quality depends on the similarity between q and p, and we conjecture here that this similarity measure, if can be described, is closely related to the KSD.

Fitting a parametric estimator using KSD. The non-parametric predictive estimator could be computationally demanding. Setting aside the cost of fitting the "parameters", in prediction the time complexity for the non-parametric estimator is O(K 2 + Kd). Also storing the "parameters" needs O(Kd) memory for ĜStein V . These costs make the non-parametric estimator undesirable for high-dimensional data, since in order to obtain accurate predictions it often requires K scaling with d as well. To address this, one can also minimise the KSD using parametric approximations, in a similar way as to derive the score matching estimator in Section 3.3.2. More precisely, we define a parametric approximation in a similar fashion as (17), and in the appendix we show that if the RBF kernel is used for both the KSD and the parametric approximation, then the linear coefficients a = (a 1 , ..., a K ) T can be calculated analytically:

EQUATION

with X the "gram matrix" that has elements X ij = (x i ) T x j . Then for an unseen observation y ∼ p the gradient approximation returns ∇ y log q(y) ≈ (â Stein V ) T ∇ y K(•, y). In this case one only maintains the linear coefficients âStein V and computes a linear combination in prediction, which takes O(K) memory and O(Kd) time and therefore is computationally cheaper than the non-parametric prediction model ( 27).

4. APPLICATIONS:

We present some case studies that apply the gradient estimators to implicit models. Detailed settings (architecture, learning rate, etc.) are presented in the appendix. Implementation is released at https://github.com/YingzhenLi/SteinGrad.

4.1. SYNTHETIC EXAMPLE: HAMILTONIAN FLOW WITH APPROXIMATE GRADIENTS:

We first consider a simple synthetic example to demonstrate the accuracy of the proposed gradient estimator. More precisely we consider the kernel induced Hamiltonian flow (not an exact sampler) (Strathmann et al., 2015) on a 2-dimensional banana-shaped object: x ∼ B(x; b = 0.03, v = 100) ⇔ x 1 ∼ N (x 1 ; 0, v), x 2 = + b(x 2 1 -v), ∼ N ( ; 0, 1). The approximate Hamiltonian flow is constructed using the same operator as in Hamiltonian Monte Carlo (HMC) (Neal et al., 2011) , except that the exact score function ∇ x log B(x) is replaced by the approximate gradients. We still use the exact target density to compute the rejection step as we mainly focus on testing the accuracy of the gradient estimators. We test both versions of the predictive Stein gradient estimator (see section 3.4) since we require the particles of parallel chains to be independent with each other. We fit the gradient estimators on K = 200 training datapoints from the target density. The bandwidth of the RBF kernel is computed by the median heuristic and scaled up by a scalar between [1, 5]. All three methods are simulated for T = 2, 000 iterations, share the same initial locations that are constructed by target distribution samples plus Gaussian noises of standard deviation 2.0, and the results are averaged over 200 parallel chains.

We visualise the samples and some MCMC statistics in Figure 2 . In general all the resulting Hamiltonian flows are HMC-like, which give us the confidence that the gradient estimators extrapolate reasonably well at unseen locations. However all of these methods have trouble exploring the extremes, because at those locations there are very few or even no training data-points. Indeed we found it necessary to use large (but not too large) bandwidths, in order to both allow exploration of those extremes, and ensure that the corresponding test function is not too smooth. In terms of quantitative metrics, the acceptance rates are reasonably high for all the gradient estimators, and the KSD estimates (across chains) as a measure of sample quality are also close to that computed on HMC samples. The returned estimates of E[x 1 ] are close to zero which is the ground true value. We found that the non-parametric Stein gradient estimator is more sensitive to hyper-parameters of the dynamics, e.g. the stepsize of each HMC step. We believe a careful selection of the kernel (e.g. those with long tails) and a better search for the hyper-parameters (for both the kernel and the dynamics) can further improve the sample quality and the chain mixing time, but this is not investigated here. 

4.2. META-LEARNING OF APPROXIMATE POSTERIOR SAMPLERS FOR BAYESIAN NNS:

One of the recent focuses on meta-learning has been on learning optimisers for training deep neural networks, e.g. see (Andrychowicz et al., 2016) . Could analogous goals be achieved for approximate inference? In this section we attempt to learn an approximate posterior sampler for Bayesian neural networks (Bayesian NNs, BNNs) that generalises to unseen datasets and architectures. A more detailed introduction of Bayesian neural networks is included in the appendix, and in a nutshell, we consider a binary classification task: p(y = 1|x, θ) = sigmoid(NN θ (x)), p 0 (θ) = N (θ; 0, I). After observing the training data D = {(x n , y n )} N n=1 , we first obtain the approximate posterior q φ (θ) ≈ p(θ|D) ∝ p 0 (θ) N n=1 p(y n |x n , θ), then approximate the predictive distribution for a new observation as p(y

* = 1|x * , D) ≈ 1 K K k=1 p(y * = 1|x * , θ k ), θ k ∼ q φ (θ).

In this task we define an implicit approximate posterior distribution q φ (θ) as the following stochastic normalising flow (Rezende & Mohamed, 2015) θ t+1 = f (θ t , ∇ t , t ): given the current location θ t and the mini-batch data {(x m , y m )} M m=1 , the update for the next step is

θ t+1 = θ t + ζ∆ φ (θ t , ∇ t ) + σ φ (θ t , ∇ t )

t , t ∼ N ( ; 0, I),

EQUATION

)

The coordinates of the noise standard deviation σ φ (θ t , ∇ t ) and the moving direction ∆ φ (θ t , ∇ t ) are parametrised by a coordinate-wise neural network. If properly trained, this neural network will learn the best combination of the current location and gradient information, and produce approximate posterior samples efficiently on different probabilistic modelling tasks. Here we propose using the variational inference objective (2) computed on the samples {θ k t } to learn the variational parameters φ. Since in this case the gradient of the log joint distribution can be computed analytically, we only approximate the gradient of the entropy term H[q] as in (3), with the exact score function replaced by the presented gradient estimators. We report the results using the non-parametric Stein gradient estimator as we found it works better than the parametric version. The RBF kernel is applied for gradient estimation, with the hyper-parameters determined by a grid search on the bandwidth σ 2 ∈ {0.25, 1.0, 4.0, 10.0, median trick} and η ∈ {0.1, 0.5, 1.0, 2.0}.

We briefly describe the test protocol. We take from the UCI repository (Lichman, 2013) six binary classification datasets (australian, breast, crabs, ionosphere, pima, sonar), train an approximate sampler on crabs with a small neural network that has one 20-unit hidden layer with ReLU activation, and generalise to the remaining datasets with a bigger network that has 50 hidden units and uses sigmoid activation. We use ionosphere as the validation set to tune ζ. The remaining 4 datasets are further split into 40% training subset for simulating samples from the approximate sampler, and 60% test subsets for evaluating the sampler's performance.

Figure 3 presents the (negative) test log-likelihood (LL), classification error, and an estimate of the KSD U-statistic S 2 U (p(θ|D), q(θ)) (with data sub-sampling) over 5 splits of each test dataset. Besides the gradient estimators we also compare with two baselines: an approximate posterior sampler trained by maximum a posteriori (MAP), and stochastic gradient Langevin dynamics (SGLD) Figure 3 : Generalisation performances for trained approximate posterior samplers. (Welling & Teh, 2011 ) evaluated on the test datasets directly. In summary, SGLD returns best results in KSD metric. The Stein approach performs equally well or a little better than SGLD in terms of test-LL and test error. The KDE method is slightly worse and is close to MAP, indicating that the KDE estimator does not provide a very informative gradient for the entropy term. Surprisingly the score matching estimator method produces considerably worse results (except for breast dataset), even after carefully tuning the bandwidth and the regularisation parameter η. Future work should investigate the usage of advanced recurrent neural networks such as an LSTM (Hochreiter & Schmidhuber, 1997) , which is expected to return better performance.

4.3. TOWARDS ADDRESSING MODE COLLAPSE IN GANS USING ENTROPY REGULARISATION:

GANs are notoriously difficult to train in practice. Besides the instability of gradient-based minimax optimisation which has been partially addressed by many recent proposals (Salimans et al., 2016; Arjovsky et al., 2017; Berthelot et al., 2017) , they also suffer from mode collapse. We propose adding an entropy regulariser to the GAN generator loss. Concretely, assume the generative model p θ (x) is implicitly defined by x = f θ (z), z ∼ p 0 (z), then the generator's loss is defined by

EQUATION

where J gen (θ) is the original loss function for the generator from any GAN algorithm and α is a hyper-parameter. In practice (the gradient of) ( 21) is estimated using Monte Carlo.

We empirically investigate the entropy regularisation idea on the very recently proposed boundary equilibrium GAN (BEGAN) (Berthelot et al., 2017) method using (continuous) MNIST, and we refer to the appendix for the detailed mathematical set-up. In this case the non-parametric V-statistic Stein gradient estimator is used. We use a convolutional generative network and a convolutional auto-encoder and select the hyper-parameters of BEGAN γ ∈ {0.3, 0.5, 0.7}, α ∈ [0, 1] and λ = 0.001. The Epanechnikov kernel K(x, x ) := 1

d d j=1 (1-(x j -x j ) 2

) is used as the pixel values lie in a unit interval (see appendix for the expression of the score matching estimator), and to ensure the boundary condition we clip the pixel values into range [10 -8 , 1 -10 -8 ]. The generated images are visualised in Figure 4 . BEGAN without the entropy regularisation fails to generate diverse samples even when trained with learning rate decay. The other three images clearly demonstrate the benefit of the entropy regularisation technique, with the Stein approach obtaining the highest diversity without compromising visual quality.

We further consider four metrics to assess the trained models quantitatively. First 500 samples are generated for each trained model, then we compute their nearest neighbours in the training set using l 1 distance, and obtain a probability vector p by averaging over these neighbour images' label vectors. In Figure 5 we depict the entropy of p (top left), averaged l 1 distances to the nearest neighbour (top right), and the difference between the largest and smallest elements in p (bottom right). The error bars are obtained by 5 independent runs. These results demonstrate that the Stein Figure 5 : Quantitative evaluation on entropy regularised BEGAN. The higher the better for the LHS panels and the other way around for the RHS ones. See text for details. approach performs significantly better than the other two, in that it learns a better generative model not only faster but also in a more stable way. Interestingly the KDE approach achieves the lowest average l 1 distance to nearest neighbours, possibly because it tends to memorise training examples. We next train a fully connected network π(y|x) on MNIST that achieves 98.16% text accuracy, and compute on the generated images an empirical estimate of the inception score (Salimans et al., 2016) E p(x) [KL[π(y|x)||π(y)]] with π(y) = E p(x) [π(y|x)] (bottom left panel). High inception score indicates that the generate images tend to be both realistic looking and diverse, and again the Stein approach out-performs the others on this metric by a large margin.

Concerning computation speed, all the three methods are of the same order: 10.20s/epoch for KDE, 10.85s/epoch for Score, and 10.30s/epoch for Stein. 1 This is because K < d (in the experiments K = 100 and d = 784) so that the complexity terms are dominated by kernel computations (O(K 2 d)) required by all the three methods. Also for a comparison, the original BEGAN method without entropy regularisation runs for 9.05s/epoch. Therefore the main computation cost is dominated by the optimisation of the discriminator/generator, and the proposed entropy regularisation can be applied to many GAN frameworks with little computational burden.

5. CONCLUSIONS AND FUTURE WORK:

We have presented the Stein gradient estimator as a novel generalisation to the score matching gradient estimator. With a focus on learning implicit models, we have empirically demonstrated the efficacy of the proposed estimator by showing how it opens the door to a range of novel learning tasks: approximating gradient-free MCMC, meta-learning for approximate inference, and unsupervised learning for image generation. Future work will expand the understanding of gradient estimators in both theoretical and practical aspects. Theoretical development will compare both the V-statistic and U-statistic Stein gradient estimators and formalise consistency proofs. Practical work will improve the sample efficiency of kernel estimators in high dimensions and develop fast yet accurate approximations to matrix inversion. It is also interesting to investigate applications of gradient approximation methods to training implicit generative models without the help of discriminators. Finally it remains an open question that how to generalise the Stein gradient estimator to non-kernel settings and discrete distributions.

A SCORE MATCHING ESTIMATOR: REMARKS AND DERIVATIONS:

In this section we provide more discussions and analytical solutions for the score matching estimator. More specifically, we will derive the linear coefficient a = (a 1 , ..., a K ) for the case of the Epanechnikov kernel.

A.1 SOME REMARKS ON SCORE MATCHING Remark. It has been shown in Särelä & Valpola (2005) ; Alain & Bengio (2014) that de-noising autoencoders (DAEs) (Vincent et al., 2008) , once trained, can be used to compute the score function approximately. Briefly speaking, a DAE learns to reconstruct a datum x from a corrupted input x = x+σ , ∼ N (0, I) by minimising the mean square error. Then the optimal DAE can be used to approximate the score function as ∇ x log p(x) ≈ 1 σ 2 (DAE * (x)-x). Sonderby et al. ( 2017) applied this idea to train an implicit model for image super-resolution, providing some promising results in some metrics. However applying similar ideas to variational inference can be computationally expensive, because the estimation of ∇ z log q(z|x) is a sub-routine for VI which is repeatedly required. Therefore in the paper we deploy kernel machines that allow analytical solutions to the score matching estimator in order to avoid double loop optimisation.

Remark. As a side note, score matching can also be used to learn the parameters of an unnormalised density. In this case the target distribution q would be the data distribution and q is often a Boltzmann distribution with intractable partition function. As a parameter estimation technique, score matching is also related to contrastive divergence (Hinton, 2002) , pseudo likelihood estimation (Hyvärinen, 2006) , and DAEs (Vincent, 2011; Alain & Bengio, 2014) . Generalisations of score matching methods are also presented in e.g. Lyu (2009) ; Marlin et al. (2010) .

A.2 THE RBF KERNEL CASE:

The derivations for the RBF kernel case is referred to (Strathmann et al., 2015) , and for completeness we include the final solutions here. Assume the parametric approximation is defined as log q(x) = K k=1 a k K(x, x k ) + C, where the RBF kernel uses bandwidth parameter σ. then the optimal solution of the coefficients âscore = (Σ + ηI)

-1 v, with v = d i=1 σ 2 K1 -K(x i x i ) + diag(x i )K1 -2diag(x i )Kx i , Σ = d i=1 [diag(x i )K -Kdiag(x i )] [Kdiag(x i ) -diag(x i )K] , x i = (x 1 i , x 2 i , ..., x K i ) T ∈ R K×1 .

A.3 THE EPANECHNIKOV KERNEL CASE:

The Epanechnikov kernel is defined as K

(x, x ) = 1 d d i=1 (1 -(x i -x i ) 2 )

, where the first and second order gradients w.r.t. x i is

EQUATION

Thus the score matching objective with log q(

x) = K k=1 a k K(x, x k ) + C is reduced to F(a) = 1 K K j=1 || K k=1 a k 2 d (x k -x j )|| 2 2 -2 K k=1 a k 2 d d = 4 K K j=1 1 d 2 K k=1 K k =1 a k a k (x k -x j ) T (x k -x j ) -a T 1 := 4(a T Σa -a T 1), ∆ φ (θ t , ∇ t )(i)

and the noise variance σ φ (θ t , ∇ t )(i). Softplus non-linearity is used for the hidden layer and to compute the noise variance we apply ReLU activation to ensure non-negativity. The step-size ζ is selected as 1e-5 which is tuned on the KDE approach. For SGLD step-size 1e-5 also returns overall good results.

The training process is the following. We simulate the approximate sampler for 10 transitions and sum over the variational lower-bounds computed on the samples of every step. Concretely, the maximisation objective is

L(φ) = T t=1

L VI (q t ),

where T = 100 and q t (θ) is implicitly defined by the marginal distribution of θ t that is dependent on φ. In practice the variational lower-bound L VI (q t ) is further approximated by Monte Carlo and data sub-sampling:

L VI (q t ) ≈ N M M m=1

log p(y m |x m , θ t ) + log p 0 (θ t ) -log q t (θ t ).

The MAP baseline considers an alternative objective function by removing the log q t (θ t ) term from the above MC-VI objective.

Truncated back-propagation is applied for every 10 steps in order to avoid vanishing/exploding gradients. The simulated samples at time T are stored to initialise the Markov chain for the next iteration, and for every 50 iterations we restart the simulation by randomly sampling the locations from the prior. Early stopping is applied using the validation dataset, and the learning rate is set to 0.001, the number of epochs is set to 500.

We perform hyper-parameter search for the kernel, i.e. a grid search on the bandwidth σ 2 ∈ {0.25, 1.0, 4.0, 10.0, median trick} and η ∈ {0.1, 0.5, 1.0, 2.0}. We found the median heuristic is sufficient for the KDE and Stein approaches. However, we failed to obtain desirable results using the score matching estimator with median heuristics, and for other settings the score matching approach underperforms when compared to KDE and Stein methods.

C.2 BEGAN EXPERIMENTS:

In this section we describe the experimental details of the BEGAN experiment, but first we introduce the mathematical idea and discuss how the entropy regulariser is applied.

Assume the generator is implicitly defined: x ∼ p θ (x) ↔ x = f θ (z), z ∼ p 0 (z). In BEGAN the discriminator is defined as an auto-encoder D ϕ (x) that reconstructs the input x. After selecting a ratio parameter γ > 0, a control rate β 0 initialised at 0, and a "learning rate" λ > 0 for the control rate, the loss functions for the generator x = f θ (z), z ∼ p 0 (z) and the discriminator are:

J (x) = ||D ϕ (x) -x||, || • || = || • || 2 2 or || • || 1 , J gen (θ; ϕ) = J (f θ (z)

), z ∼ p 0 (z) J dis (ϕ; θ) = J (x) -β t J gen (θ; ϕ), x ∼ D β t+1 = β t + λ(γJ (x) -J (f θ (z))).

(34)

The main idea behind BEGAN is that, as the reconstruction loss J (•) is approximately Gaussian distributed, with γ = 1 the discriminator loss J dis is (approximately) proportional to the Wasserstein distance between loss distributions induced by the data distribution p D (x) and the generator p θ (x). In practice it is beneficial to maintain the equilibrium γE p D [J (x)] = E p θ [J (x)] through the optimisation procedure described in (34) that is motivated by proportional control theory. This approach effectively stabilises training, however it suffers from catastrophic mode collapsing problem (see the left most panel in Figure 4 ). To address this issue, we simply subtract an entropy term from the generator's loss function, i.e.

Jgen (θ; ϕ) = J gen (θ; ϕ) -αH[p θ ],

where the rest of the optimisation objectives remains as in (34). This procedure would maintain the equilibrium γE p D [J (x)] = E p θ [J (x)] -αH [p] . We approximate the gradient ∇ θ H[p θ ] using the estimators presented in the main text. For the purpose of updating the control rate β t two

:

All the methods are timed on a machine with an NVIDIA GeForce GTX TITAN X GPU.

Figures and Tables:

Figure 2: Kernel induced Hamiltonian flow compared with HMC. Top: samples generated from the dynamics, training data (in cyan), an the trajectory of a particle for T = 1 to 200 starting at the star location (in yellow). Bottom: statistics computed during simulations. See main text for details.

Figure 4: Visualisation of generated images from trained BEGAN models.

Peer review (iclr2018):
 This paper deals with the estimation of the score function, i.e., the derivative of the log likelihood. Some methods were introduced and a new method using Stein identity was proposed. The setup of the trasnductive learning was introduced to add the prediction power to the proposed method. The method was used to several applications. This is an interesting approach to estimate the score function for location models in a non-parametric way. I have a couple of minor comments below. - Stein identity is the formula that holds for the class of ellipsoidal distribution including Gaussian distribution. I*m not sure the term "Stein identity" is appropriate to express the equation (8). - Some boundary condition should be assumed to assure that integration by parts works properly. Describing an explicit boundary condition to guarantee the proper estimation would be nice.