Title: Entropy-based Training Methods for Scalable Neural Implicit Samplers
Venue: NeurIPS 2023
Accepted: yes
Keywords: implicit sampler

learning to sample

generative models
Abstract: Efficiently sampling from un-normalized target distributions is a fundamental problem in scientific computing and machine learning. Traditional approaches such as Markov Chain Monte Carlo (MCMC) guarantee asymptotically unbiased samples from such distributions but suffer from computational inefficiency, particularly when dealing with high-dimensional targets, as they require numerous iterations to generate a batch of samples. In this paper, we introduce an efficient and scalable neural implicit sampler that overcomes these limitations. The implicit sampler can generate large batches of samples with low computational costs by leveraging a neural transformation that directly maps easily sampled latent vectors to target samples without the need for iterative procedures. To train the neural implicit samplers, we introduce two novel methods: the KL training method and the Fisher training method. The former method minimizes the Kullback-Leibler divergence, while the latter minimizes the Fisher divergence between the sampler and the target distributions. By employing the two training methods, we effectively optimize the neural implicit samplers to learn and generate from the desired target distribution. To demonstrate the effectiveness, efficiency, and scalability of our proposed samplers, we evaluate them on three sampling benchmarks with different scales. These benchmarks include sampling from 2D targets, Bayesian inference, and sampling from high-dimensional energy-based models (EBMs). Notably, in the experiment involving high-dimensional EBMs, our sampler produces samples that are comparable to those generated by MCMC-based methods while being more than 100 times more efficient, showcasing the efficiency of our neural sampler. Besides the theoretical contributions and strong empirical performances, the proposed neural samplers and corresponding training methods will shed light on further research on developing efficient samplers for various applications beyond the ones explored in this study.

Entropy-based Training Methods for Scalable Neural Implicit Sampler

Abstract

Efficiently sampling from un-normalized target distributions is a fundamental problem in scientific computing and machine learning. Traditional approaches such as Markov Chain Monte Carlo (MCMC) guarantee asymptotically unbiased samples from such distributions but suffer from computational inefficiency, particularly when dealing with high-dimensional targets, as they require numerous iterations to generate a batch of samples. In this paper, we introduce an efficient and scalable neural implicit sampler that overcomes these limitations. The implicit sampler can generate large batches of samples with low computational costs by leveraging a neural transformation that directly maps easily sampled latent vectors to target samples without the need for iterative procedures. To train the neural implicit samplers, we introduce two novel methods: the KL training method and the Fisher training method. The former method minimizes the Kullback-Leibler divergence, while the latter minimizes the Fisher divergence between the sampler and the target distributions. By employing the two training methods, we effectively optimize the neural implicit samplers to learn and generate from the desired target distribution. To demonstrate the effectiveness, efficiency, and scalability of our proposed samplers, we evaluate them on three sampling benchmarks with different scales. These benchmarks include sampling from 2D targets, Bayesian inference, and sampling from high-dimensional energy-based models (EBMs). Notably, in the experiment involving high-dimensional EBMs, our sampler produces samples that are comparable to those generated by MCMC-based methods while being more than 100 times more efficient, showcasing the efficiency of our neural sampler. Besides the theoretical contributions and strong empirical performances, the proposed neural samplers and corresponding training methods will shed light on further research on developing efficient samplers for various applications beyond the ones explored in this study.

1 Introduction

Efficiently sampling from un-normalized distributions is a fundamental challenge that arises in various research domains, including Bayesian statistics [1], biology and physics simulations [2, 3], as well as generative modeling and machine learning [4, 5]. The task at hand involves generating batches of samples from a target distribution defined by a differentiable un-normalized potential function, denoted as log q ( x ) . This problem seeks effective techniques for obtaining samples that accurately represent the underlying target distribution while minimizing computational costs.

Two main classes of methods exist for addressing this challenge. The first class comprises Markov Chain Monte Carlo (MCMC) algorithms [6-9]. MCMC methods involve the design of Markov Chains with stationary distributions that match the target distribution. By simulating these Markov

Chains, it becomes possible to transition a batch of initial samples towards samples approximately distributed following the target distribution. While MCMC methods offer asymptotically unbiased samples, they often suffer from computational inefficiency, particularly when dealing with highdimensional target distributions. This inefficiency stems from the need to perform a significant number of iterations to generate a single batch of samples.

The second class of methods, known as learning to sample (L2S) models [10-18], focuses on training neural components, such as generative models, to generate samples following the target distribution. L2S models aim to improve the quality and efficiency of sample generation compared to traditional training-free algorithms like MCMCs. Among the L2S models, the neural implicit sampler stands out due to its remarkable efficiency. The neural implicit sampler is an L2S model that incorporates a neural transformation x = g θ ( z ) , where z is an easy-to-sample latent vector from a known distribution p z ( z ) . This transformation allows the sampler to generate samples directly from the target distribution without the need for iterations or gradient operations, which are commonly required in MCMCs. Such a sampler offers significant advantages in terms of efficiency and scalability, particularly when dealing with high-dimensional distributions such as those found in image spaces. However, training a neural implicit sampler poses technical challenges due to the intractability of sampler distributions that are caused by complex neural architectures of g θ . Achieving effective training requires careful consideration of various factors and the development of appropriate training strategies.

In this paper, we introduce two innovative training approaches, namely the KL training methods and the Fisher training methods, for training implicit samplers. We theoretically demonstrate that our proposed KL and Fisher training methods are equivalent to minimizing the Kullback-Leibler divergence and the Fisher divergence [19], respectively, between the sampler and the target distribution. To evaluate the effectiveness, efficiency, and scalability of our proposed training approach and sampler, we conduct experiments on three different sampling benchmarks that span various scales. These benchmarks include sampling from 2D targets with low dimensions, Bayesian inference with moderate dimensions, and sampling from high-dimensional energy-based models (EBMs) trained on the MNIST dataset. The results of experiments consistently demonstrate the superior performance of our proposed method and the quality of the generated samples. Notably, in the EBM experiment, our sampler trained with the KL training method achieves sample quality comparable to that of the EBM model while being more than 100 times more efficient than traditional MCMC methods. These findings highlight the effectiveness, efficiency, and scalability of the proposed approach across different sampling scenarios.

2 Backgrounds

2.1 Generative Models for Sampling

Generative models have emerged as powerful tools in various domains, demonstrating their ability to produce diverse, high-quality samples. They have been successfully applied in tasks such as text-toimage generation [20-28], audio generation[29], video and 3D creation[30-33], and even molecule design [34, 35]. Recently, there has been a growing interest in leveraging generative models for sampling from target distributions.

Specifically, consider the problem that we have access to an un-normalized target distribution q ( x ) , or its logarithm log q ( x ) . Our objective is to train generative models that can effectively generate samples following the target distribution. By learning the underlying structure of the target distribution, these generative models can generate samples that capture the characteristics of the target distribution, facilitating tasks such as posterior inference in Bayesian statistics.

Three primary classes of generative models have been extensively studied for sampling tasks. The first class comprises normalizing flows (NFs) [36], while the second class consists of diffusion models (DMs) [37]. NFs employ invertible neural transformations to map Gaussian latent vectors z to obtain samples x . The strict invertibility of NF transformations enables the availability of likelihood values for generated samples, which are differentiable with respect to the model's parameters. Training NFs often involves minimizing the KL divergence between the NF and the target distribution [12]. On the other hand, DMs employ neural score networks to model the marginal score functions of a data-initialized diffusion process. The score functions are learned using techniques related to score-matching [38]. DMs have been successfully employed to enhance the sampler per-

formance of the annealed importance sampling algorithm [39], a widely recognized MCMC method for various sampling benchmarks. Despite the successes of NFs and DMs, both models have their limitations. The invertibility of NFs restricts their expressiveness, which can hinder their ability to effectively model high-dimensional targets. Moreover, DMs still require a considerable number of iterations for sample generation, resulting in computational inefficiency.

The third class is the implicit generative model. An implicit generative model (IGM) uses a flexible neural transform g θ ( . ) to push forward easy-to-sample latent vectors z ∼ p z to obtain samples x = g θ ( z ) . The main difference between IGMs and NFs is that IGMs' neural transformation is not required to be strictly invertible, which unlocks both the flexibility and the modeling power of deep neural networks. For instance, Hu et al. [11] proposed to train implicit samplers by minimizing the Stein discrepancy between the sampler and the target distribution. The Stein Discrepancy (SD) [40] between p and q is defined as

The notation D represents the dimension of data space and f d ( . ) is the d -th component of the vectorvalued function f . The calculation of Stein's discrepancy relies on solving a maximization problem w.r.t. the test function f . When the function class F is carefully chosen, the optimal f may have an explicit solution or easier formulation. For instance, Hu et al. [11] found that if F is taken to be F = { f : E x ∼ p ∥ f ( x ) ∥ 2 2 ≤ δ } , the SD is equivalent to a regularized representation

They used two neural networks: g θ to parametrize an implicit sampler and f η to parametrize the test function. Let p θ ( x ) denote the implicit sampler distribution induced by x = g θ ( z ) with z ∼ p z ( z ) . They solved a bi-level mini-max problem on parameter pair ( θ, η ) to obtain a sampler that tries to minimize the SD with

Hu et al. [11] opened the door to training implicit samplers by minimizing divergences which are implemented with a two-networks bi-level optimization problem. Since the general Stein's discrepancy does not have explicit formulas, in our paper we study training approaches of an implicit sampler by minimizing the KL divergence and the Fisher divergences. For the implementation of minimization of Fisher divergence, some backgrounds of score function estimation are needed as we put in Section 2.2.

2.2 Score Function Estimation

Since the implicit sampler does not have an explicit log-likelihood function or score function, training it with score-based divergence requires inevitably estimating the score function (or equivalent component). Score matching [41] and its variants provided practical approaches to estimating score functions. Assume one only has available samples x ∼ p and wants to use a parametric approximated distribution q ϕ ( x ) to approximate p . Such an approximation can be made by minimizing the Fisher Divergence between p and q ϕ with the definition

Under certain conditions, the equality

holds (usually referred to as Stein's Identity[42]). Here ∆ x log q ϕ ( x ) = ∑ i ∂ 2 ∂ x 2 i log q ϕ ( x ) denotes the Laplacian operator applied on log q ϕ ( x ) . Combining this equality and noting that the first term

of FD E x ∼ p ∥∇ x log p ( x ) ∥ 2 2 does not rely on the parameter ϕ , minimizing D FD ( p, q ϕ ) is equivalent to minimizing the following objective

This objective can be estimated only through samples from p , thus is tractable when q ϕ is welldefined. Moreover, one only needs to define a score network s ϕ ( x ): R D → R D instead of a density model to represent the parametric score function. This technique was named after Score Matching. Other variants of score matching were also studied [43-48]. Score Matching related techniques have been widely used in the domain of generative modeling, so we use them to estimate the score function of the sampler's distribution.

3 Training Approaches of Neural Implicit Samplers

The neural implicit sampler has the advantage of sampling efficiency and scalability. But the training approach for such samplers still requires to explore. In this section, we start by introducing the proposed KL and Fisher training methods for implicit samplers.

3.1 Training Neural Implicit Sampler by Minimizing the KL Divergence

Let g θ ( · ): R D Z → R D X be an implicit sampler (i.e., a neural transform), p z the latent distribution, p θ the sampler induced distribution x = g θ ( z ) , and q ( x ) the un-normalized target. Our goal in this section is to train the implicit sampler by minimizing the KL divergence between the sampler distribution and the un-normalized target. For a neural implicit sampler, we can efficiently sample from it but the density p θ has no explicit expression. Our goal is to minimize the KL divergence between p θ and target q in order to train the sampler g θ . The KL divergence between p θ and q is defined as

In order to update the sampler g θ , we need to calculate the gradient of θ for the divergence equation 2. So we take the gradient of θ with respect to the KL divergence equation 2 and obtain

The second term vanishes under loose conditions (we put detailed discussions in the Appendix A.1)

So the gradient of the KL term only remains the first term and the gradient can be calculated with

where s q := ∇ x log q ( x ) is the score function of target distribution so we have the explicit expression. The s p := ∇ x log p θ ( x ) is the score function of the implicit sampler but we do not have the explicit expression. However, thanks to advanced techniques of neural score function estimation as we introduced in 2.2, we can use another score neural network to estimate the score function s p ( x ) ≈ s ϕ ( x ) with samples consistently drawn from p θ . With the estimated s p ( x ) , the formula equation 3 provides a practical method for updating the implicit sampler that is shown to minimize the KL divergence.

More precisely, we can alternate between a score-estimation phase and a KL minimization phase in order to minimize the KL divergence between the sampler and the target distribution. The former phase uses score estimation techniques to train a score network s ϕ to approximate the sampler's

score function with samples consistently generated from the implicit sampler. In the latter phase, we can minimize the KL divergence with gradient-based optimization algorithms according to gradient formula equation 3. By alternating the optimization of a score network s ϕ and implicit sampler g θ , one can eventually train an implicit sampler that can approximately generate samples that are distributed according to the target distribution. We name such method the KL training method .

3.2 Training Neural Implicit Sampler by Minimizing Fisher Divergence

In the previous section, we proposed a KL training method for implicit samplers. In this section, we proposed an alternative training method motivated by minimizing the Fisher divergence instead of the KL. We use the same notation as the section 3.1. Recall the definition of the Fisher divergence,

The score functions s θ and s q remain the same meaning as in the equation 3. For an implicit sampler, we do not know the explicit expression of s θ . In order to minimize the Fisher divergence, we take the θ gradient of equation 4 and obtain

We use s θ † to represent a function that does not differentiate with respect to parameter θ , so the first term gradient equation 5 is equivalent to taking the gradient of an equivalent objective

As for the second term of equation 5, we notice that under regularity conditions

So the second term of equation 5 equals to

We put detailed derivation and required conditions in Appendix A.2. With the expression of the Grad ( F, 2) ( θ ) , minimizing the second part of D ( F ) ( θ ) is equivalent to minimizing an equivalent objective

Combining equivalent loss function equation 6 and equation 8, we obtain a final objective function for training implicit sampler that minimizes the Fisher divergence

And we formally define the gradient operator

Similar to the KL divergence case as discussed in Section 3.1, our overall training approach of such a Fisher divergence minimization consists of two phases: the score-estimation phase, and the Fisher divergence minimization phase. The score-estimation phase is the same as the one used for KL training, while the Fisher divergence minimization updates the implicit sampler by minimizing objective equation 9. Since the introduced method aims to minimize the Fisher divergence, we name our training the Fisher training for an implicit sampler.

Architecture choice of score network. Notice that for Fisher training, the objective equation 9 needs the calculation of the gradient of the scoring network ∇ x s ϕ ( x ) , so it requires the scoring network to be point-wisely differentiable. Commonly used activation functions such as ReLU or LeakyReLU do not satisfy the differentiability property. In our further experiments, we use differentiable activation functions for Fisher training.

We formally give a unified Algorithm for the training of implicit samplers with KL, Fisher, and Combine training in Algorithm 1. We use the standard score-matching objective as an example of score estimation phases, other score estimation methods are also suitable. The notation s ( d ) ϕ and x d represent the d -th component of the score network s ϕ ( . ) and the data x , and D means the data dimension. The notation sg( . ) represents stopping the parameter dependence of the input.

Algorithm 1: Training Algorithm for Implicit Samplers

3.3 Connection to Fisher-Stein's Sampler

In this section, we provide an analysis of the connection between our proposed Fisher training method and Fisher-Stein's sampler proposed in Hu et al. [11]. Surprisingly, we discover that the training method used in Fisher-Stein's sampler is implicitly equivalent to our proposed Fisher training. To establish this connection, we adopt the notations introduced in the introduction of FisherStein's sampler in Section 2.1. We summarize our findings in a theorem, and for a comprehensive understanding of the proof and analysis, we provide detailed explanations in Appendix A.5.

Proposition 1. Estimating the sampler's score function s ϕ ( . ) with score matching is equivalent to the maximization of test function f of Fisher-Stein's Discrepancy objective (15). More specially, the optimal score estimation S ∗ and Fisher-Stein's optimal test function f ∗ satisfy

With Proposition 2, we can prove a Theorem that states that the Fisher-Stein training coincides with our proposed Fisher training.

Theorem 1. Assume f ∗ is the optimal test function that maximizes Fisher-Stein's objective 15. Then the minimization part of objective 15 for sampler G θ with the gradient-based algorithm is equivalent to the Fisher training, i.e. the minimization objective of 15 shares the same gradient with Grad ( F ) ( θ ) .

Theorem 1 states that Fisher-Stein's sampler and its training method are equivalent to our proposed Fisher training. The two methods differ only in the parametrization of the score network (and the test function f ) and the motivation that they are proposed. It is one of our contributions to point out

this equivalence. Our analysis reveals the underlying connection between the two training methods, shedding light on the similarities and shared principles between them. By establishing this connection, we provide a deeper understanding of the theoretical foundations of our proposed Fisher training approach and its relationship to existing techniques in the field.

4 Experiments

In previous sections, we have established the KL and Fisher training methods for implicit samplers. Both methods only require access to an un-normalized potential function of the target distribution. In this section, we aim to apply our introduced training methods on three sampling benchmarks whose scales vary across both low (2-dimensional targets) to high (784-dimensions image targets). We compare both the implicit sampler and our training methods to other competitor methods and samplers. All experiments demonstrate the significant advantage of the implicit sampler and the proposed training methods on both the sampling efficiency and the sample quality.

4.1 2D Synthetic Target Sampling

Sample quality on 2D targets. In this section, we refer to the open-source implementation of Sharrock and Nemeth [49] 4 and train samplers on eight 2D target distributions. We compare our neural samplers with 3 MCMC baselines: Stein variational gradient descent (SVGD) [50], Langevin dynamics (LD) [51], and Hamiltonian Monte Carlo (HMC) [52]; 1 explicit baseline: coupling normalizing flow [53]; and 2 implicit samplers: KSD neural sampler (KSD-NS) [11] and SteinGan [54]. All implicit samplers have the same neural architectures, i.e. four layer MLP with 400 hidden units at each layer and ELU activation functions for sampler and score network (if necessary).

The kernelized Stein's discrepancy (KSD) [55] is a popular metric for evaluating the sample quality of Monte Carlo samplers [55, 40]. We evaluate the KSD with IMQ kernel (implemented by an open-source package, the sgmcmcjax 5 ) on all target distributions as the metric reported in Table 1.

Settings. For all MCMC samplers, we set the number of iterations to 500, which we find is enough for convergence. For SVGD and LD, we set the sampling step size to 0.01. For the HMC sampler, we optimize and find the step size to be 0.1, and LeapFrog updates to 10 work the best. For Coupling Flow, we follow Dinh et al. [53], and use 3 invertible blocks, with each block containing a 4-layer MLP with 200 hidden units and Gaussian Error Linear Units (GELU) activations. The total parameters of flow are significantly larger than neural samplers. We find that adding more coupling blocks does not lead to better performance. For all targets, we train each neural sampler with the Adam optimizer with the same learning rate of 2e-5 and default bete. We use the same batch size of 5000 for 10k iterations when training all neural samplers. We evaluate the KSD for every 200 iterations with 500 samples with 20 repeats for each time. We pick the lowest mean KSD among 10k training iterations as our final results. Since our proposed Fisher neural sampler and KL neural sampler require learning the score network, we find that using 5-step updates of the score network for each update of the neural sampler works well.

Performance. Table 1 shows the numerical comparison between the mentioned samplers. The SVGD performs significantly the best among all samplers. However, the SVGD has a more heavy computational cost when the number of particles grows because its update requires matrix computation for a large matrix. The LD and HMC perform almost the same. The KL-NS performs the best across almost all targets, slightly better than LD and HMC on each target. The SteinGAN performs second and is closely comparable to KL-NS. In theory, both the KL-NS and the SteinGAN aim to minimize the KL divergence between the sampler and the target distribution in different ways, so we are not surprised by their similar performances. The Coupling Flow performs overall third, but it fails to correctly capture the Rosenbrock target. We believe more powerful flows, such as stochastic variants or flows with more complex blocks will lead to better performance, but these enhancements will inevitably bring in more computational complexity. The Fisher-NS performs the fourth, with 1 failure case on the Funnel target. We find that the KSD-NS is hard to tune in practice, which has two failure cases. Besides, the KSD-NS has a relatively high computational cost because it requires differentiation through the empirical KSD, which needs large matrix computation when the batch size is large. Overall, the one-shot KL-NS has shown strong performance, outperforming LD and HMC with multiple iterations. We also visualize the sample results on five distributions with hard-to-sample characteristics such as multi-modality and periodicity, as shown in Figure 2.

Comparison of Computational Costs. To demonstrate the significant advantage of our proposed implicit samplers over other samplers, we evaluate the efficiency of each sampler under the same environment to measure the wall-clock inference time. In below Table 2, we summarize the computational costs (wall-clock time) of SVGD, HMC, and LD together with our KL-NS neural sampler to compare the computational costs of each method. The eight targets are analytic targets while the EBM is a neural target. Each sampler can generate samples with comparable qualities (i.e. their KSD values are comparable). Table 2 records the wall-clock inference time (seconds) for each sampler when generating 1k samples.

Table 2: Inference Time Comparison of MCMC Samplers and Neural Samplers (seconds). The 2D experiment is conducted on an 8-CPU cluster with PyTorch of 1.8.1, while the EBM experiment is on 1 Nvidia Titan RTX GPU with PyTorch 1.8.1. Unless especially emphasizing, we set the stepsize=0.01, iterations=500, num particles=1000, num repeats=100.

The LD and HMC with 500 iterations have comparable (slightly worse) performance than the Neural Sampler, however, their wall-clock time for obtaining 1k samples is significantly larger than Neural samplers. Let's take the simple mixture-of-2-gaussian (MOG2) target (with analytic potential and scores) as an example, the neural sampler is 1.7125/0.0014=1223 times faster than HMC and 0.2168/0.0014=154 times faster than LD. If one wants to keep the same inference time for LD and HMCtobe comparable with the neural sampler, the LD will only have 500/154=3.2 iterations, while the HMC will only have 500/1223=0.4 iterations. It is impossible to obtain promising samples with only 3.2 or 0.4 MCMC iterations.

In conclusion, for both analytic and neural targets, the neural samplers show a significant efficiency advantage over MCMC samplers with better (comparable) performances. This shows that neural samplers have potentially wider use on many sampling tasks for which high efficiency and low computational costs are demanded.

4.2 Bayesian Regression

Test accuracy compared with Stein sampler. In the previous section, we validate the sampler on low-dimensional 2D target distributions. The Bayesian logistic regression provides a source of real-world target distributions with medium dimensions (10-100). In this section, we compare our sampler with the most related sampler, Stein's sampler that is proposed in Hu et al. [11]. We compare the test accuracy of our proposed sampler against Stein's samplers and MCMC samplers.

8

2

33

4

5

3

a

6

子

?

8

3

S

6

D

50

5

8

50

50

a

5

50

8

3

3

B

50

S

9

1

8

2

2

5

100

100

100

100

100

1P

8

3

3

150

B

2

2

150

150

9

2

150

5

3

150

2

4

$

7

3

4

S

3

5

8

9

200

200

200

200

200

3

8

3

3

B

?

5

5

3

3

3

一

5

8

250

22

1323

250

250

?

8

250

8

52e

250

6

39

S

9

0

50

100

150

200

250

0

50

100

150

200

250

0

50

T

100

150

200

250

50

T

100

150

200

250

0

50

T

100

150

200

250

Sampler1NFE

Ebm50NFE

Ebm100NFE

Ebm150NFE

Ebm200NFE

>>>

（)

Experiment settings. We follow the same setting as Hu et al. [11] on the Bayesian logistic regression tasks. The Covertype data set [11] has 54 features and 581,012 observations. It has been widely used as a benchmark for Bayesian inference. More precisely, following the Hu et al. [11], we assume prior of the weights is p ( w | α ) = N ( w ; 0 , α -1 ) and p ( α ) = Gamma( α ; 1 , 0 . 01) . The data set is randomly split into the training set (80%) and the testing set (20%). Our goal is to train an implicit sampler that can sample from posterior distribution efficiently. The sampler trained with Fisher training is the same as the Fisher-Neural Sampler, which achieves the best test accuracy, as we have proved in previous sections. The sampler trained with the KL method gives the secondary test accuracy in Table 3. This experiment shows that our proposed training method is capable of handling real-world Bayesian inference with a medium dimension. After training, the implicit sampler can result in the best accuracy while achieving hundreds of times faster than MCMC methods. For detailed settings, please refer to Appendix B.

4.3 Sampling from Energy-based Models

Amain advantage of an implicit sampler is its inference efficiency. This advantage benefits applications for which the default sampling method is inefficient, such as sampling from high-dimensional energy-based models. In this section, we use our proposed KL training method for an application of learning an efficient implicit sampler to sample from a pre-trained energy-based model.

Experiment settings. An energy-based model (EBM) [56] uses a neural network to parametrize a negative energy function, i.e. the logarithm of some un-normalized distribution. After training, obtaining samples from an EBM usually requires running an annealed Markov Chain Monte Carlo [4] which is computationally inefficient. In this section, we pre-train a deep EBM [57] on the MNIST dataset and apply our proposed KL training method 6 for learning a neural implicit sampler to learn to sample from the pre-trained EBM efficiently. To take the inductive biases of image data, we parametrize our sampler via the stacking of several convolutional neural blocks. The neural sampler takes a random Gaussian input with a dimension of 128 and outputs a 32x32 tensor which has the same size as the resolution on which the EBM has been trained. We visualize the samples drawn from EBM with 50+ evaluations of the forward and backward pass of EBMs with annealed Langevin dynamics and samples from our sampler which only requires a single forward pass of the sampler's neural network. We train an MNIST classifier and compute the Frechet Inception Distance (FID) [58] and the KID [59] as an evaluation metric.

Performance. As is shown in Table 4, our sampler generates indistinguishable samples from the EBM but is about 100+ more efficient than annealed Langevin dynamics. To demonstrate the sampling efficiency of our sampler, we compare the efficiency and computational costs of EBM and our sampler, by listing some major metrics on computational performance in Table 4. The EBM has a total amount of 9 . 06 M parameters and a base FLOPS of 0 . 58 G , while the sampler has 12 . 59 M parameters and a base FLOPS of 1 . 11 G . Generating samples from EBM requires running MCMC chains, which consist of up to 200 neural functional evaluations (NFEs) of the EBM. We compare the efficiency and performance of EBM with different NFEs and find that with less than 100 NFEs the EBM can not generate realistic samples, while our sampler only needs one single NFE for good samples. This high-dimensional EBM sampling experiment demonstrates the scalability of KL training methods and their potential for accelerating large-scale energy-based model sampling.

5 Limitations and Future Works

We have presented two novel approaches for training an implicit sampler to sampler from unnormalized density. We show in theory that our training methods are equivalent to minimizing the KL or the Fisher divergences. We systematically compare the proposed training methods against each other and against other samplers (NFs and MCMC). Besides, we conduct an experiment that trains a convolutional neural sampler to learn to sample from a pre-trained EBM. This experiment shows that our implicit sampler and corresponding training have great potential to improve the efficiency of modern machine learning models such as energy-based models or diffusion models.

However, there are still limitations to our proposed methods. First, the score estimation is not computationally cheap, so developing a more efficient training algorithm by eliminating the score estimation phase is an important research direction. Besides, now the sampler is only limited to sampling problems. How to extend the methodology for other applications such as generative modeling is another interesting direction.

Acknowledgments and Disclosure of Funding

This work has been supported by the National Natural Science Foundation of China (No. 12271011) and the Beijing Natural Science Foundation (Z190001).

References

P. J. Green, 'Reversible jump markov chain monte carlo computation and bayesian model determination,' Biometrika , vol. 82, no. 4, pp. 711-732, 1995.

C. Sch¨ utte, A. Fischer, W. Huisinga, and P. Deuflhard, 'A direct approach to conformational dynamics based on hybrid monte carlo,' Journal of Computational Physics , vol. 151, no. 1, pp. 146-168, 1999.

P. Olsson, 'Two phase transitions in the fully frustrated xy model,' Physical review letters , vol. 75, no. 14, p. 2758, 1995.

J. Xie, Y. Lu, S.-C. Zhu, and Y. Wu, 'A theory of generative convnet,' in International Conference on Machine Learning . PMLR, 2016, pp. 2635-2644.

C. Andrieu, N. De Freitas, A. Doucet, and M. I. Jordan, 'An introduction to mcmc for machine learning,' Machine learning , vol. 50, no. 1, pp. 5-43, 2003.

W. Hastings, 'Monte carlo sampling methods using markov chains and their applications,' Biometrika , vol. 57, no. 1, pp. 97-109, 1970.

G. O. Roberts and J. S. Rosenthal, 'Optimal scaling of discrete approximations to langevin diffusions,' Journal of the Royal Statistical Society: Series B (Statistical Methodology) , vol. 60, no. 1, pp. 255-268, 1998.

T. Xifara, C. Sherlock, S. Livingstone, S. Byrne, and M. Girolami, 'Langevin diffusions and the metropolis-adjusted langevin algorithm,' Statistics & Probability Letters , vol. 91, pp. 1419, 2014.

R. M. Neal, 'Mcmc using hamiltonian dynamics,' in Handbook of Markov Chain Monte Carlo . Chapman and Hall/CRC, 2011, pp. 139-188.

D. Levy, M. D. Hoffman, and J. Sohl-Dickstein, 'Generalizing hamiltonian monte carlo with neural networks,' in International Conference on Learning Representations , 2018.

T. Hu, Z. Chen, H. Sun, J. Bai, M. Ye, and G. Cheng, 'Stein neural sampler,' arXiv preprint arXiv:1810.03545 , 2018.

H. Wu, J. K¨ ohler, and F. No´ e, 'Stochastic normalizing flows,' Advances in Neural Information Processing Systems , vol. 33, pp. 5933-5944, 2020.

L. L. di Langosco, V. Fortuin, and H. Strathmann, 'Neural variational gradient descent,' arXiv preprint arXiv:2107.10731 , 2021.

M. Arbel, A. Matthews, and A. Doucet, 'Annealed flow transport monte carlo,' in International Conference on Machine Learning . PMLR, 2021, pp. 318-330.

Q. Zhang and Y. Chen, 'Path integral sampler: A stochastic control approach for sampling,' in International Conference on Learning Representations , 2021.

A. Matthews, M. Arbel, D. J. Rezende, and A. Doucet, 'Continual repeated annealed flow transport monte carlo,' in International Conference on Machine Learning . PMLR, 2022, pp. 15 196-15 219.

F. Vargas, W. S. Grathwohl, and A. Doucet, 'Denoising diffusion samplers,' in The Eleventh International Conference on Learning Representations , 2022.

S. Lahlou, T. Deleu, P. Lemos, D. Zhang, A. Volokhova, A. Hern´ andez-Garcıa, L. N. Ezzine, Y. Bengio, and N. Malkin, 'A theory of continuous generative flow networks,' in International Conference on Machine Learning . PMLR, 2023, pp. 18 269-18 300.

O. Johnson, 'Information theory and the central limit theorem,' 2004.

A. Brock, J. Donahue, and K. Simonyan, 'Large scale gan training for high fidelity natural image synthesis,' in International Conference on Learning Representations , 2018.

T. Karras, S. Laine, and T. Aila, 'A style-based generator architecture for generative adversarial networks,' in Proceedings of the IEEE/CVF conference on computer vision and pattern recognition , 2019, pp. 4401-4410.

T. Karras, S. Laine, M. Aittala, J. Hellsten, J. Lehtinen, and T. Aila, 'Analyzing and improving the image quality of stylegan,' in Proceedings of the IEEE/CVF conference on computer vision and pattern recognition , 2020, pp. 8110-8119.

T. Karras, M. Aittala, S. Laine, E. H¨ ark¨ onen, J. Hellsten, J. Lehtinen, and T. Aila, 'Alias-free generative adversarial networks,' Advances in Neural Information Processing Systems , vol. 34, pp. 852-863, 2021.

A. Q. Nichol and P. Dhariwal, 'Improved denoising diffusion probabilistic models,' in International Conference on Machine Learning . PMLR, 2021, pp. 8162-8171.

P. Dhariwal and A. Nichol, 'Diffusion models beat gans on image synthesis,' Advances in neural information processing systems , vol. 34, pp. 8780-8794, 2021.

A. Ramesh, P. Dhariwal, A. Nichol, C. Chu, and M. Chen, 'Hierarchical text-conditional image generation with clip latents.'

C. Saharia, W. Chan, S. Saxena, L. Li, J. Whang, E. L. Denton, K. Ghasemipour, R. Gontijo Lopes, B. Karagol Ayan, T. Salimans et al. , 'Photorealistic text-to-image diffusion models with deep language understanding,' Advances in Neural Information Processing Systems , vol. 35, pp. 36 479-36 494, 2022.

R. Rombach, A. Blattmann, D. Lorenz, P. Esser, and B. Ommer, 'High-resolution image synthesis with latent diffusion models,' in Proceedings of the IEEE/CVF conference on computer vision and pattern recognition , 2022, pp. 10 684-10 695.

R. Huang, J. Huang, D. Yang, Y. Ren, L. Liu, M. Li, Z. Ye, J. Liu, X. Yin, and Z. Zhao, 'Makean-audio: Text-to-audio generation with prompt-enhanced diffusion models,' arXiv preprint arXiv:2301.12661 , 2023.

A. Clark, J. Donahue, and K. Simonyan, 'Adversarial video generation on complex datasets,' arXiv preprint arXiv:1907.06571 , 2019.

J. Ho, W. Chan, C. Saharia, J. Whang, R. Gao, A. Gritsenko, D. P. Kingma, B. Poole, M. Norouzi, D. J. Fleet et al. , 'Imagen video: High definition video generation with diffusion models,' arXiv preprint arXiv:2210.02303 , 2022.

E. Molad, E. Horwitz, D. Valevski, A. R. Acha, Y. Matias, Y. Pritch, Y. Leviathan, and Y. Hoshen, 'Dreamix: Video diffusion models are general video editors,' arXiv preprint arXiv:2302.01329 , 2023.

B. Poole, A. Jain, J. T. Barron, and B. Mildenhall, 'Dreamfusion: Text-to-3d using 2d diffusion,' in The Eleventh International Conference on Learning Representations , 2022.

A. Nichol, P. Dhariwal, A. Ramesh, P. Shyam, P. Mishkin, B. McGrew, I. Sutskever, and M. Chen, 'Glide: Towards photorealistic image generation and editing with text-guided diffusion models,' in International Conference on Machine Learning , 2021.

J. Ho, T. Salimans, A. Gritsenko, W. Chan, M. Norouzi, and D. J. Fleet, 'Video diffusion models,' arXiv preprint arXiv:2204.03458 , 2022.

D. Rezende and S. Mohamed, 'Variational inference with normalizing flows,' in International conference on machine learning . PMLR, 2015, pp. 1530-1538.

Y. Song, J. Sohl-Dickstein, D. P. Kingma, A. Kumar, S. Ermon, and B. Poole, 'Score-based generative modeling through stochastic differential equations,' in International Conference on Learning Representations , 2021.

A. Doucet, W. S. Grathwohl, A. G. D. G. Matthews, and H. Strathmann, 'Score-based diffusion meets annealed importance sampling,' in Advances in Neural Information Processing Systems , 2022.

R. M. Neal, 'Annealed importance sampling,' Statistics and computing , vol. 11, no. 2, pp. 125-139, 2001.

J. Gorham and L. Mackey, 'Measuring sample quality with stein's method,' Advances in Neural Information Processing Systems , vol. 28, 2015.

A. Hyv¨ arinen and P. Dayan, 'Estimation of non-normalized statistical models by score matching.' Journal of Machine Learning Research , vol. 6, no. 4, 2005.

C. M. Stein, 'Estimation of the mean of a multivariate normal distribution,' The annals of Statistics , pp. 1135-1151, 1981.

Y. Song, S. Garg, J. Shi, and S. Ermon, 'Sliced score matching: A scalable approach to density and score estimation,' in Proceedings of the Thirty-Fifth Conference on Uncertainty in Artificial Intelligence, UAI 2019, Tel Aviv, Israel, July 22-25, 2019 , 2019, p. 204.

P. Vincent, 'A connection between score matching and denoising autoencoders,' Neural computation , vol. 23, no. 7, pp. 1661-1674, 2011.

T. Pang, K. Xu, C. Li, Y. Song, S. Ermon, and J. Zhu, 'Efficient learning of generative models via finite-difference score matching,' Advances in Neural Information Processing Systems , vol. 33, pp. 19 175-19 188, 2020.

C. Meng, L. Yu, Y. Song, J. Song, and S. Ermon, 'Autoregressive score matching,' Advances in Neural Information Processing Systems , vol. 33, pp. 6673-6683, 2020.

C. Lu, K. Zheng, F. Bao, J. Chen, C. Li, and J. Zhu, 'Maximum likelihood training for scorebased diffusion odes by high order denoising score matching,' in International Conference on Machine Learning . PMLR, 2022, pp. 14 429-14 460.

F. Bao, C. Li, K. Xu, H. Su, J. Zhu, and B. Zhang, 'Bi-level score matching for learning energy-based latent variable models,' Advances in Neural Information Processing Systems , vol. 33, pp. 18 110-18 122, 2020.

L. Sharrock and C. Nemeth, 'Coin sampling: Gradient-based bayesian inference without learning rates,' arXiv preprint arXiv:2301.11294 , 2023.

Q. Liu and D. Wang, 'Stein variational gradient descent: A general purpose bayesian inference algorithm,' Advances in neural information processing systems , vol. 29, 2016.

M. Welling and Y. W. Teh, 'Bayesian learning via stochastic gradient langevin dynamics,' in Proceedings of the 28th international conference on machine learning (ICML-11) , 2011, pp. 681-688.

R. M. Neal et al. , 'Mcmc using hamiltonian dynamics,' Handbook of markov chain monte carlo , vol. 2, no. 11, p. 2, 2011.

L. Dinh, J. N. Sohl-Dickstein, and S. Bengio, 'Density estimation using real nvp,' ArXiv , vol. abs/1605.08803, 2016.

D. Wang and Q. Liu, 'Learning to draw samples: With application to amortized mle for generative adversarial learning,' arXiv preprint arXiv:1611.01722 , 2016.

Q. Liu, J. Lee, and M. I. Jordan, 'A kernelized stein discrepancy for goodness-of-fit tests,' in International Conference on Machine Learning , 2016.

Y. LeCun, S. Chopra, R. Hadsell, A. Ranzato, and F. J. Huang, 'A tutorial on energy-based learning,' 2006.

Z. Li, Y. Chen, and F. T. Sommer, 'Learning energy-based models in high-dimensional spaces with multi-scale denoising score matching,' arXiv preprint arXiv:1910.07762 , 2019.

M. Heusel, H. Ramsauer, T. Unterthiner, B. Nessler, and S. Hochreiter, 'GANs trained by a two time-scale update rule converge to a local Nash equilibrium,' in Advances in Neural Information Processing Systems , 2017, pp. 6626-6637.

M. Bi´ nkowski, D. J. Sutherland, M. Arbel, and A. Gretton, 'Demystifying mmd gans,' in International Conference on Learning Representations , 2018.

L. K. Wenliang and H. Kanagawa, 'Blindness of score-based methods to isolated components and mixing proportions,' arXiv preprint arXiv:2008.10087 , 2020.

W. Gong and Y. Li, 'Interpreting diffusion score matching using normalizing flow,' in ICML Workshop on Invertible Neural Networks, Normalizing Flows, and Explicit Likelihood Models , 2021.

J. A. Blackard and D. J. Dean, 'Comparative accuracies of artificial neural networks and discriminant analysis in predicting forest cover types from cartographic variables,' Computers and Electronics in Agriculture , vol. 24, pp. 131-151, 1999.

M. K. Titsias and M. L´ azaro-Gredilla, 'Doubly stochastic variational bayes for non-conjugate inference,' in International Conference on Machine Learning , 2014.

D. Hendrycks and K. Gimpel, 'Gaussian error linear units (gelus),' arXiv: Learning , 2016.

K. He, X. Zhang, S. Ren, and J. Sun, 'Deep residual learning for image recognition,' 2016 IEEE Conference on Computer Vision and Pattern Recognition (CVPR) , pp. 770-778, 2015.

S. Zagoruyko and N. Komodakis, 'Wide residual networks,' arXiv preprint arXiv:1605.07146 , 2016.

M. F. Hutchinson, 'A stochastic estimator of the trace of the influence matrix for laplacian smoothing splines,' Communications in Statistics - Simulation and Computation , vol. 18, pp. 1059-1076, 1989.

A Theory

A.1 Detailed Derivations of Vanishing Term of KL Training Method

Recall the vanishing term

The equality holds if p θ ( x ) satisfies the conditions (1). p θ ( x ) is Lebesgue integrable for x with each θ ; (2). For almost all x ∈ R D , the partial derivative ∂p θ ( x ) /∂θ exists for all θ ∈ Θ . (3) there exists an integrable function g ( . ) : R D → R , such that p θ ( x ) ≤ g ( x ) for all x in its domain. Then the derivative w.r.t θ can be exchanged with the integral over x , i.e.

A.2 Detailed Derivations of Fisher Training Methods

Recall the definition of the Fisher divergence,

The score functions s θ and s q remain the same meaning as in the equation 3. For an implicit sampler, we do not know the explicit expression of s θ . In order to minimize the Fisher divergence, we take the θ gradient of equation 4 and obtain

We use s θ † to represent a function that does not differentiate with respect to parameter θ , so the first term gradient equation 5 is equivalent to taking the gradient of an equivalent objective

As for the second term of equation 5, we notice that under regularity conditions

So the second term of equation 5 equals to

The notation x ∈ Ω means the boundary of distribution density of p θ ( x ) and the value p θ vanishes on the boundary. Usually, the integral over the boundary Ω of the support of x vanishes if the density p θ decays fast enough and the logarithm of density has finite parameter derivatives ∥ ∂ ∂θ log p θ ( x ) ∥ 2 for all θ ∈ Θ , then the first term vanishes. In practice, we find that the first term of the above expression does not bother the performance of Fisher training, however, for complex circumstances, the first term can not be guaranteed to be exactly 0. Besides, the above derivations are for the Fisher training, so our proposed KL training method is not bothered by them.

A.3 Additional Discussions on Proposed Algorithms

The training algorithm 1 consists of two alternative phases for training implicit sampler: the score estimation phase and the generator update phase. The former phase uses score-matching-related techniques to update the score network to match the score function of the implicit distribution. The latter phase uses the score function of the target distribution and the learned score network to update the generator's parameter in order to minimize the KL divergence between the implicit distribution and the target distribution.

Concerns on blindness and ill-landscape. Our proposed methods incorporate a score estimation phase, where a score neural network is used to estimate the score function of the sampler. However, two concerns arise regarding this score estimation step: the blindness issue and the presence of an ill-landscape in score estimation.

Blindness. The blindness [60] is a pervasive practical issue of score-based methods that states that two multi-modal distributions could potentially have very small Fisher divergence even if they are highly dissimilar. For our proposed training method that minimizes the Fisher divergence, there exist concerns that blindness would cause the training to fail. However, our intensive study shows that, with the introduced annealing techniques that train samplers in a progressive way, the blindness issue does not occur.

Ill-landscape. As is pointed out in Gong and Li [61], the parameter landscape of the Fisher divergence to heavy-tail distribution can be ill-posed, having a unique minimizer but hard to be caught by gradient-based optimization algorithms. For our proposed Fisher training, we find that the annealing technique and suitable parametrization of the score network overcome the ill parameter landscape of training. For a straight demonstration, we train an implicit sampler to capture the 1D student-t distribution, a heavy-tail distribution that has been studied in Gong and Li [61] and find that the sampler is capable of learning such a heavy-tail distribution correctly.

A.4 Details on Ill-Landscape

Ill-landscape. As is pointed out in Gong and Li [61], the parameter landscape of the Fisher divergence to heavy-tail distribution can be ill-posed, having a unique minimizer but hard to be caught by gradient-based optimization algorithms. The student-t distribution, a heavy-tail distribution, is a representative distribution that has an ill landscape as is shown in Gong and Li [61]. Since our Fisher training aims to minimize the Fisher divergence, there is a concern about the ill landscape of the Fisher training. However, for both our proposed KL and Fisher training, we find that suitable parametrization of the score network overcomes the ill parameter landscape of training. For a straight demonstration, we train an implicit sampler to capture the 1D student-t distribution, a heavy-tail distribution that has been studied in Gong and Li [61] and finds that the sampler is capable of learning such a heavy-tail distribution correctly. More specifically, We plot a comparison of the studentt distribution and our learned sampler with Fisher training in Figure 3.

A.5 Proofs of Equivalence to Fisher-Stein Sampler

We provide proof of Proposition 1 here.

Proposition 2. Estimating the sampler's score function s ϕ ( . ) with score matching is equivalent to the maximization of test function f of Fisher-Stein's Discrepancy objective. More specially, the optimal score estimation S ∗ and Fisher-Stein's optimal test function f ∗ satisfy

Proof. With fixed p and known target q , the optimal test function f ∗ has representation

Where functional L ( f ) has integral representation

Here l ( x , f , ∇ f ) = ∫ p ( x ) ⟨∇ x log q ( x ) , f ( x ) ⟩ + p ( x ) ⟨∇ x , f ( x ) ⟩ -λp ( x )[ f T ( x ) f ( x )] . By EulerLagrange equation, the optimal function f satisfies

By calculation, we have

So the optimal f ∗ satisfies the Euler-Lagrange equation as

Divide the both side with p ( x ) and note that ∇ x p ( x ) /p ( x ) = ∇ x log p ( x ) , the equation turns to

Next, consider optimal S ∗ . The S ∗ is obtained by minimizing the Score Matching objective, which is equivalent to minimizing the Fisher divergence between p and S induced family, thus the optimal S ∗ ( x ) = ∇ x log p ( x ) . Substitute ∇ x log p ( x ) with S ∗ into f ∗ formula, we have

Next, we prove Theorem 2. Recall the definition of Fisher-Stein objective

Theorem. Assume f ∗ is the optimal test function that maximizes Fisher-Stein's objective 15. Then the minimization part of objective 15 for sampler G θ with the gradient-based algorithm is equivalent to the Fisher training, i.e. the minimization objective of 15 shares the same gradient with Grad ( F ) ( θ ) .

Proof. Notice that by Proposition 2, the maximization of equation 15 to get an optimal f ∗ is equivalent to learning a score function S ∗ with the score-matching related techniques, i.e.

Here s θ ( x ) denotes the unknown score function of the implicit sampler's distribution. Here we use Stein's identity which states

Put the equation 19 and f ∗ in objective 15 with S ∗ according to the relation 17, we obtain

Taking the θ gradient of the above objective, we have

In the above equation, some gradient terms vanish because S ∗ ( x ) = s θ ( x ) for all x . Since the p θ is implemented with the implicit sampler, i.e. x ∼ p θ means x = G θ ( z ) , z ∼ p z ( z ) . So above gradient term can be further simplified as

Here s q ( x ) := ∇ x log q ( x ) . So the θ gradient is the same as the Fisher training (equation (5)) in the paper.

B Experiments

B.1 Experiment Details on 2D Synthetic Sampling

Model architectures. For toy 2-dimensional data experiments, we use a 4-layer MLP neural network with 200 hidden units in each layer as the sampler. The activation of the sampler is chosen as LeakyReLU non-linearity with a 0.2 coefficient. The score network is a 4-layer MLP with 200 hidden units in each layer. The activation of the score network is GELU non-linearity.

Evaluation metrics. We calculate the Kernelized Stein Discrepancy with multi-scale RBM kernel as is implemented in codebase 7 . We set the bandwidth of the kernel to be 0.25.

Comparison of KL and Fisher training methods. In this work, we propose the KL training method and the Fisher training method for training neural implicit samplers. Both methods require two phases which update a neural score network to match the score function of implicit distribution and update the implicit sampler's parameter in order to minimize either the KL divergence or the Fisher divergence. To this end, we give some comparisons of the proposed two training methods. Empirically, we find that the KL training methods work better than Fisher training under the KSD for 2D targets. We also observe that the Fisher training tends to learn an implicit sampler that collapses on certain modes of the target distribution. Besides, since the Fisher training requires calculating the data derivative of the score function, i.e. ∇ x [ s q ( x ) -s θ † ( x ) ] , thus requires the score function to be differentiable. The non-differentiable activate functions such as ReLU and LeakyReLU are not suitable for implementing Fisher training. However, the KL training does not suffer from such issues. Besides, the computational costs of Fisher training are also more expensive than KL training for each iteration.

B.2 Experiment Details on Bayesian Regression

Experiment settings. Weuse the same setting as [11], where the prior of the weights is p ( w | α ) = N ( w ; 0 , α -1 ) and p ( α ) = Gamma( α ; 1 , 0 . 01) . The Covertype dataset [62] (581,012 samples with 54 features) is randomly split into the training set (80%) and testing set (20%). Our methods are compared with Stein GAN, SVGD, SGLD, doubly stochastic variational inference (DSVI) [63], KSD-NS, and Fisher-NS [11] on this data set. For SGLD, DSVI, and SVGD, the model is trained with 3 epochs of the training set (about 15k iterations), while for neural network-based methods (Fisher-NS, KSD-NS, and Stein GAN), the training is run until convergence. Since we have shown the equivalence of Fisher-NS [11], we only train the implicit sampler with the KL training method. We set the training learning rate for both the score network and the sampler to be 0 . 0002 . The target distribution (posterior distribution) is computed with 500 random data samples. We set the number of score estimation phases as 2. We evaluate the test accuracy of the logistic regression model with 100 samples from the trained sampler.

AverageTrainingAccuracy

0.75

0.70

(%)

accuracy(

0.65

0.60

0.55

0.50

0

20

40

60

80

100

numberofiterations

We use the same settings for Bayesian inference as in Hu et al. [11]. The neural samplers are implemented via a 4-layer MLP with 1024 hidden units in each layer and GELU activation functions. The output dimension of the sampler is 55 and the input dimension is set to be 55x10 = 550, following the same setting as in [1]. The score network has the same neural architecture as the sampler, but the input dimension is set to 55 which matches the output dimension. We use Adam optimizers with a learning rate of 0.0002 and default beta values for both sampler and score networks. We use a batch size of 100 for training the sampler and use 2 updates for score network for each update of the sampler network. We use the standard score matching for learning the score network. We train the sampler for 10k iterations for each repeat and use 30 independent repeats to compute the mean and std of the test accuracy. The learning rate of SGLD is chosen to be 0 . 1 / ( t +1) 0 . 55 as suggested in Welling and Teh [51], and the average of the last 100 points is used for evaluation. For DSVI, the learning rate is 1 e -07 , and 100 iterations are used for each stage. For SVGD, we use RBF kernel with bandwidth h calculated by the median trick as in Liu and Wang [50], and 100 particles are used for evaluation with step size being 0.05.

Model architectures. We use 4-layer MLP neural networks as the sampler and score network, respectively. The activation of the sampler is chosen as Gaussian Error Linear Units (GELU) [64] non-linearity for both the score network and the sampler. We set the hidden dimension to 1024. The input dimension of the sampler is set to 128.

B.3 Experiment Details on Sampling from EBMs

Datasets and model architectures. We pre-train a deep (multi-scale) EBM [57] denoted as E d ( . ) with 12 residual layers [65] as is used in DeepEBM 8 . We set the starting noise level and ending noise level to be σ min = 0 . 3 and σ max = 3 . 0 . The energy for different noises is parametrized with E σ ( x ) = f ( x ) /σ . The training learning rate is set to 0.001. We pre-train the EBM for 200k iterations. Samples from deep EBM are obtained by simulating annealed Langevin dynamics algorithm that is the same as the one used in Li et al. [57]. Our implicit sampler is also a neural network. The network consists of four inverse convolutional layers which hidden dimensions to be 1024, 512, 256, and 3. For each inverse convolutional layer, we use a 2D BatchNormalization and LeakyReLU non-linearity with a leak hyper-parameter to be 0.2. The prior distribution of the implicit generator is the standard Multi-variate Gaussian distribution. The score network follows a

UNet architecture that is adapted from the codebase 9 . In order to match the design of the multi-scale EBM, we parametrize the score function in the same way as S σ ( x ) := S ( x ) /σ .

Training details. Wepre-train deep EBM on the MNIST dataset in the house. Then the pre-trained EBM is viewed as a multi-scale un-normalized target distribution. We then use the KL training method to update the score network and the generator in order to let the generator match the target distribution. We initialize the generator and the score network randomly. For each training iteration, we randomly select a noise σ ∈ [ σ min , σ max ] . We then generate a sample from the generator and add a Gaussian noise with variance σ 2 . Then we update the score network with generated samples according to denoising score matching. Then we generate a batch of samples and add Gaussian noise with the same variance from the generator and use the gradient estimation in Algorithm 1 to update the generator's parameter. We use the Adam optimizer with a learning rate of 0.0001 for both the score network and the generator. For the score network, we set the hyper-parameter of the optimizer to be β 0 = 0 . 9 and β 1 = 0 . 99 . For the generator, we set β 0 = 0 and β 1 = 0 . 99 .

Evaluation metrics. To qualitatively evaluate the quality of generated samples, we adapt the definition of the Frechet Inception Distance (FID) [58]. We pre-train an MNIST image classifier with architecture the WideResNet [66] model with a depth of 16 and widen factor of 8. We then calculate the Wasserstein distance in the feature space of the pre-trained classifier with the same method as the FID.

Additional discussion. In this experiment, we use the KL training method to train a generator to sample from pre-trained multi-scale EBM. We also try the Fisher training method. But the Fisher training fails to train implicit samplers to sample from EBMs. One of the reasons is that the Fisher training requires calculating the trace of the Jacobian matrix of the score network. For high-dimensional data such as image data, such trace term is very expensive. We also try the stochastic trace estimation method such as the Skilling-Hutchison trace estimation technique [67]. But for the Fisher training, we find that stochastic trace estimation does not make the Fisher training work for high-dimensional samplers.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis article studies the problem of learning implicit samplers, for which the score function is not available, and therefore the KL (or the Fisher) divergence cannot be explicitly computed. They introduce an alternating procedure, where, for a fixed sampler parameter $\\theta$ they first learn the surrogate score $s_{\\phi}$ of the implicit sampler, and then use this score to minimize the divergence with respect to the target. An equivalent between Fisher divergence training and the min-max approach to minimizing Stein's discrepancy is proven. The attractivity of the method is illustrated on standard benchmarks.\n\nSTRENGTHS:\nThis article is an original application of score-matching methods to implicit samplers. The procedure is well-motivated, grounded, and, to the best of my knowledge novel. The method is also fairly significant in its somewhat general applicability. The sampling application is relevant to a large part of the community. The article is further pretty well written and only requires minimal changes from a presentation/clarity standpoint.\n\nWEAKNESSES:\nThe performance of the method is evaluated only in terms of sampling accuracy, never in terms of the total cost *including training*. As a consequence, I don't believe that HMC and SGLD are fair competitors given that they provide samples immediately, with no slow training procedure. \n\nMore related are the works considering warping the target distribution, such as https://arxiv.org/abs/1903.03704, https://arxiv.org/abs/2107.08001, https://invertibleworkshop.github.io/INNF_2020/accepted_papers/pdfs/24.pdf, https://arxiv.org/abs/2210.10644.\n\nUnderstanding the competitiveness of the proposed approach with respect to these more related methods both in terms of ease of training and in terms of samples quality post training would make the article a lot stronger.\n\nQUESTIONS:\n### What you can do to improve my score\nThe empirical evaluation is, in my opinion, the weakest part of this, otherwise, well-written and motivated article. Effectively comparing with the relevant sampling literature would improve the empirical soundness and make for a stronger acceptance case.\n\n### Some questions/remarks\n- I am very confused by the paragraph on performance l.287 onwards. Can you please reformulate it?\n- It would be good to explain how you mimic the FID exactly.\n\n#### Minor comments\n- l43-44: please add a specific reference when introducing the neural implicit sampler.\n- [24] is by Hyvärinen only. Dayan was the editor of the paper. This seems to come from Google Scholar. Please do verify the rest of the citations for similar issues with Google Scholar referencing.\n- In some places, you use \"sg\" vs. the dagger superscript to denote the stop-gradient operator, it would be better to stick to one only.\n- I don't really see the point of A.1, this is well-known (for instance from the log-derivative trick) and should be replaced by a citation in this effect. Also, what is the point of putting in the conditions for Lebesgue's theorem to apply if you don't verify them for your models?\n- Please use \\eqref rather than \\ref for equation referencing.\n\n\n### Typos\nA didn't notice many\n- Algorithm 1: dependecne\n- Capitalization in citations is off (people's names are often written without upper case: stein -> Stein, markov -> Markov, etc)\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nEfficient sampling from unnormalized target distribution is of crucial interest in Bayesian inference. Classic methods such as Markov Chains Monte Carlo sampler provide unbiased samples but could be computationally expensive. This paper proposes a novel approach called neural implicit sampler to employ a neural transformation and leverage generative models to sample from target distribution. The authors have developed two innovative inference methods: the KL training method, which minimizes the Kullback-Leibler divergence and the Fisher Training method, which minimizes the Fisher divergence. The methods are evaluated on three benchmark cases including 2D target sampling, Bayesian inference and high-dimensional energy-based models. Using metrics Frechet Inception Distance, the proposed method is over 100 times more efficient than EBM's sampler.\n\nSTRENGTHS:\nThis paper has derived two novel training approaches for neural implicit sampler via KL divergence training and Fisher divergence training. The authors have also provided a theoretical analysis of the connection between their proposed Fisher training method and Fisher-Stein's sampler proposed in Hu. The authors also have evaluated their methods via three empirical benchmark studies and it is impressive that the proposed sampler can produce samplers with similar quality while being 100 times more efficient, which demonstrates the efficiency of the proposed method.\n\nWEAKNESSES:\nI do not find clear weakness of the paper. I just have some minor comments:\n\n1. In Table 1, the authors use HMC as a competing method. Since HMC's performance is very sensitive to the two tuning parameters epsilon and L, have the authors picked these two parameters to optimize the performance of HMC?\n\n2. Can the authors make their code repository publicly open if the paper is accepted?\n\n3. In [1] , a metric KID similar to FID is proposed, which shows to converge quickly to its presumed true value than FID. Can the authors evaluate their method using KID as well to see if the conclusions are consistent with FID? \n\n[1] Bińkowski, Mikołaj, Danica J. Sutherland, Michael Arbel, and Arthur Gretton. \"Demystifying mmd gans\", 2018 ICLR.\n\nQUESTIONS:\n1. In Table 1, the authors use HMC as a competing method. Since HMC's performance is very sensitive to the two tuning parameters epsilon and L, have the authors picked these two parameters to optimize the performance of HMC?\n\n2. Can the authors make their code repository publicly open if the paper is accepted?\n\n3. In [1] , a metric KID similar to FID is proposed, which shows to converge quickly to its presumed true value than FID. Can the authors evaluate their method using KID as well to see if the conclusions are consistent with FID? \n\n[1] Bińkowski, Mikołaj, Danica J. Sutherland, Michael Arbel, and Arthur Gretton. \"Demystifying mmd gans\", 2018 ICLR.\n\nLIMITATIONS:\nCurrently, the sampling approach is only limited for sampling problems. It is interesting to see if the method can be applied for generative modeling or image translation problems.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

comments: SUMMARY:
This submission proposes two training algorithms for implicit samplers, which are based on KL divergence and Fisher divergence, respectively. Tractable objective estimator are derived and practical training algorithms are demonstrated. Numerical experiments are conduct in several different settings.

STRENGTHS:
The writing is clear for most of the parts. The idea of implicit sampler has been studied before but has not been studied under the perspective of this article, especially with Fisher divergence. The derivation is sound mathematically. The experiments are done in a proper way.

WEAKNESSES:
The reviewer has the following issues.

For the second part of gradient estimator of Fisher divergence, can we directly use
${sg}\[s_\theta(g_\theta(z) - s_q(g_\theta(z))\]^{T}s_\theta(x)$? Here "sg" means stop_gradient.

There are a bunch of missing baselines and benchmarks (at least citations). To name a few:
- PIS http://arxiv.org/abs/2111.15141
- AFT http://arxiv.org/abs/2102.07501
- CRAFT http://arxiv.org/abs/2201.13117
- DDS http://arxiv.org/abs/2302.13834
- GFlowNet https://arxiv.org/abs/2301.12594

QUESTIONS:
See above.

LIMITATIONS:
See above.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 2
DATE: 2023-07-17

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper is the area of approximate inference. The goal is to use approximating distributions where the density is unknown but samples can be drawn. The advantage of this approximating class is that it is potentially more expressive and easier to work with ones where the density is known/tractable.  The paper considers approximating distributions based on two distortion measures - namely the Fisher Divergence and the reverse KL divergence.  In both cases the gradient updates are written in terms of the score of the approximating distribution. This unknown score is then approximated using score matching using samples from the distribution. An equivalence is derived with the Fisher Stein Sampler of Hu et al. for one variant of the algorithm under optimality assumptions. Experiments are carried out including a comparison to EBM inference on MNIST.\n\nSTRENGTHS:\nI think the basic idea is interesting. Certainly it is true that implicit distributions offer some advantages over more tractable approximating classes if one can deal with the resulting challenges and although there is work in this area (see below) it is less explored than for example normalizing flows. Also I think there is room for improvement on the existing approaches and with more work this approach might provide such an improvement. \n\nWhilst it is not a big jump in methodology conceptually I’m not aware of this exact approach in the prior literature and the smaller details can make a big difference in terms of improved performance.\n\nI think equivalence described with the Fisher Stein sampler and the Fisher divergence variant of the proposed method is interesting. I also think it is not a dead end in terms of the potential of this proposed algorithm variant. Since the equivalence relies on idealized optimality of the witness function the two methods may well differ in practice and this one might be better. So I would encourage the authors to explore this further. \n\nThere are quite a few baselines compared to in the Bayesian logistic regression section although the details are somewhat sparse.\n\nWEAKNESSES:\n**Unfortunately, the literature review is missing quite a few relevant references.** The papers are significant enough to the proposed method that this is a significant limitation of the submission as it stands. Here are is a non-exhaustive list of examples:\n\n*Unbiased Implicit Variational Inference. Titsias and Ruiz. AISTATS 2019.*\n\n*Semi-implicit variational inference. Yin and Zhou. ICML 2018.*\n\n*Semi-Implicit Variational Inference via Score Matching. Yu and Zhang. ICLR 2023.*\n\nThis group of papers is on the topic of variational inference with implicit distributions. The methods are somewhat distinct from the proposed submission but they are close that they should be discussed in the text and possibly compared to.\n\n*Operator Variational Inference. Ranganath, Tran and Blei. NeurIPS 2016.* This highly cited work from NeurIPS is significantly earlier than some of the references in the submitted text. The paper uses Stein divergences and highlights the benefit of using implicit distributions. \n\n**Existence of approximating density** One theoretical concern is as to whether the implicit approximating distributions actually have a density at all. For example the experiment in Section 4.3 the dimension of the inputs random variables is 128 but the output space is of dimension 32 x 32. So the samples will fall on a lower dimensional manifold of the output space and no have density with respect to the full Lebesgue measure. Note that this is a different issue from the density existing but being unknown/intractable. I think with thought and effort this issue could probably be tamed but it should be considered and discussed in a paper with this topic. \n\n**The submitted experiments are not entirely reproducible.** Take for example the toy experiment in Section 4.1. No details are given for the baselines in either the main text or the supplement. Hamiltonian Monte Carlo is a gold standard asymptotically exact method which should work well for unimodal distributions in low dimensions. It is therefore not credible for the banana experiment that the proposed algorithm could be better unless the computational budget is constrained or the algorithm is mistuned. No details are given of the tuning or computational budget.  Similarly coupling flows are in principle capable of modelling simple low dimensional distributions. There are many possible reasons why they did not in this case but since there are few details given it is not possible to know why. For Section 4.2 the details of the baselines are lacking. The tuning of these algorithms is challenging and can substantially affect their performance. Submitting the code would have helped a lot here.\n\n**The experiments do not cover all relevant questions.** In Section 4.3 whilst I can see that the sampling is faster than annealed Langevin dynamics there are several methods that would plausibly be able to learn a one shot sampler and these are not compared to.\n\nWhilst I acknowledge some discussion of the computational bottlenecks of the method, I would have liked more. For example the training time in the EBM method must be large and is not included in the comparison to the annealed Langevin dynamics. \n\nSince learning point estimates of parameters along side inference is a big advantage of variational approaches I would have liked to see an example of this. For example, learning the parameters of the EBM rather than just using a pretrained one. \n\n**Smaller points:**\n\nEquation following Line 146: The result you are proving here is that the expectation of the score is zero. This is a very well known result. For instance it often comes up when discussing the Fisher Information matrix. So you could spend less of your precious space in the main text proving it. \n\nLine 305: \"Besides, now the sampling is only limited for sampling problems\" The meaning was unclear for this sentence. This has not affected my review score.\n\nQUESTIONS:\nI have no questions.\n\nLIMITATIONS:\nI see no negative societal impact.\n\nThere is some discussion of limitations but I already mentioned some things I would have liked to see more discussion of.  Please see the weaknesses section above.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-25"}