Learning Kernel Stein Discrepancy for Training Energy-Based Models

Abstract

The primary challenge in unsupervised learning is training unnormalized density models and then generating similar samples. Few traditional unnormalized models know what the quality of the trained model is, as most models are evaluated by downstream tasks and often involve complex sampling processes. Kernel Stein Discrepancy (KSD), a goodness-of-fit test method, can measure the discrepancy between the generated samples and the theoretical distribution; therefore, it can be employed to measure the quality of trained models. We first demonstrate that, under certain constraints, KSD is equal to Maximum Mean Discrepancy (MMD), a two-sample test method. PT KSD GAN (Kernel Stein Discrepancy Generative Adversarial Network with a Pulling-Away Term) is produced to compel generated samples to approximate the theoretical distribution. The generator, functioning as an implicit generative model, employs KSD as loss to avoid tedious sampling processes. In contrast, the discriminator is trained to identify the data manifold, also known as an explicit energy-based model. To demonstrate the effectiveness of our approach, we undertook experiments on two-dimensional toy datasets. Our results highlight that our generator adeptly captures the accurate density distribution, while the discriminator proficiently recognizes the unnormalized approximate distribution shape. When applied to linear Independent Component Analysis datasets, the log likelihoods of PT KSD GAN improve by about 5‰ over existing methods when the data dimension is less than 30. Furthermore, our tests on image datasets reveal that the PT KSD GAN excels in navigating high-dimensional challenges, yielding authentically genuine samples.

1. Introduction

In unsupervised learning, training unnormalized density models and generating similar samples are hot topics, which are also the fundamental work in probability statistics [1]. In density estimation, we often achieve an unnormalized model, which contains a normalization constant. The computation of this normalization constant has been proven to be an NP-hard problem and intractable in high-dimensional datasets [2]. Only with the density function of given samples can we conduct some downstream tasks, such as sample generation and quickest change detection [3].

The traditional unnormalized density models approximate the theoretical probability distribution of given datasets by calculating the maximum likelihood estimate (MLE), which derives a variety of inference algorithms, like variational inference and MCMC [4]. But, the initializing probability distribution and Markov chain structures are hand-designed [5]; therefore, these algorithms lack universality. Simultaneously, high-dimensional curse and integral items are barriers in the sampling process [6], greatly restricting the application scenarios.

Various unnormalized models have been proposed in recent years. Energy-based models [7] are widely used in unnormalized density models because of their freedom and simplicity in designing energy function. But, in order to normalize the model, the denominator in the density function is usually an integral term. To calculate the tedious integral item and its gradient, diversiform deep energy models [8,9,10] have been proposed. Their models indirectly eliminate integral terms, and the process of generating samples from the Gibbs distribution is very complex. Few of these models know what the quality of the trained model is, as most models are evaluated via downstream tasks [11] or by the quality of generated samples from trained models [12].

The generative adversarial network (GAN) [13] indirectly evaluates generative models using generated samples. GAN removes the probability density function estimation process in the unnormalized model and directly simulates the generation of samples. By avoiding the complex density estimation process and the bias caused by the estimation, GAN makes sample generation no longer a difficult task. GAN has brought about revolutionary evolution in image generation [14] and other domains [15], but the minimax loss function leads to mode collapse and a vanishing gradient during GAN training [16]. Various solutions have been proposed to alleviate these challenges in GAN, such as substituting the regularization terms and loss functions: WGAN [17], Fisher GAN [18], CLC-GAN [19], etc. GAN does not consider the probability distribution of given samples, nor can it provide an explicit estimate. As a result, GAN cannot conduct approximate inference and related downstream tasks, such as Independent Component Analysis.

From a probabilistic statistical point of view, it is natural to introduce classical hypothesis testing methods into machine learning. Applying hypothesis testing to train unnormalized models is usually a novel method favored by researchers. GMMN [20] replaces the critic of GAN with a two-sample test based on MMD [21]. DEAN [22] optimizes models using goodness-of-fit test statistical methods. ${\chi }^2$ GAN [23] measures the discrepancy between two probability distributions by ${\chi }^2$ distance. However, because the density distribution of samples needs to be explicitly given, KSD is rarely used in training unnormalized models.

In this paper, we propose a new unnormalized model to address the problems mentioned above without sampling from the model, and we can measure the quality of the model directly during the training process. We introduce the PT KSD GAN, an unnormalized model that merges Kernel Stein Discrepancy (KSD) [24,25] as a goodness-of-fit test method with the GAN framework. Our discriminator leverages a deep energy network to assign energy values to each sample, effectively distinguishing between genuine and synthetic samples. Simultaneously, our generator aims to minimize the Kernel Stein Discrepancy between generated samples and the theoretical density distribution. Our work has several main contributions:

(1)

To summarize, this paper provides evidence that, under certain assumptions, Kernel Stein Discrepancy as a goodness-of-fit test method is equivalent to Maximum Mean Discrepancy as a two-sample test method.

(2)

An adversarial learning algorithm, PT KSD GAN, is proposed. We select a deep energy network as the discriminator so that the density distribution of given samples is explicitly obtained through the energy model. The generator of PT KSD GAN can produce samples that minimize Kernel Stein Discrepancy between generated samples and the obtained density distribution. We eliminate the tedious calculation of integral terms in the energy model and can easily generate the desired samples through the generator.

(3)

For the first time, we demonstrate experimentally that KSD is valid in high-dimensional datasets. In linear Independent Component Analysis datasets, where the data dimension is less than 30, the performance of PT KSD GAN outperforms other unnormalized training methods. Our tests on image datasets mark a significant advancement in addressing high-dimensional data challenges.

2. Relative Work

2.1. Score-Matching Generative Models

This paper is about density estimation, where the general problem description is as below:

Problem 1.

Given a finite set of samples ${\left\{x_i\right\}}_{i=1}^n\subset \chi \subset R^d$(also named true samples) from an unknown theoretical density distribution $q\left(x\right)$, which we approximate by $p\left(\theta ;z\right)$, where z is a random variable, which is sampled from a known simple distribution $q\left(z\right)$, we optimize parameters $\theta$so that the output values (also named as false samples) of $p\left(\theta ;z\right)$match the theoretical distribution $q\left(x\right)$.

The energy model is a commonly used explicit density estimation algorithm. In the energy model, energy function $e$ transforms every observation into a single scalar, named as “energy”. Then, the Gibbs distribution can be expressed:

$$q\left(x\right)=\frac{\exp \left(-e\left(x\right)\right)}{Z}$$

(1)

where $e:R^d\to R$, and $Z={\int }_x^{}\exp \left(-e\left(x\right)\right)$ is the integral item, which is also known as the partition function. When using Formula (1) to calculate the explicit density distribution of samples, the partition function is intractable or even unsolvable. In order to solve the intractable problem, Hyvärinen [26] proposed score matching to eliminate computation of the partition function. The score is the derivative of the log density of the Gibbs distribution in Formula (1):

$$s_q\left(x\right)={\nabla }_x\log q\left(x\right)={\nabla }_x\log \left(\frac{\exp \left(-e\left(x\right)\right)}{{\rm Z}}\right)={\nabla }_x\left(-e\left(x\right)\right)-{\nabla }_x\log \left({\rm Z}\right)$$

(2)

The second term ${\nabla }_x\log \left(Z\right)$ is a constant term, which can be ignored; we just need to calculate the first term ${\nabla }_x\log \left(-e\left(x\right)\right)$. Score matching minimizes the squared distance of two density distributions’ scores, which can enable the fake distribution to approach the true theoretical distribution. Score-matching generative models [27,28,29] directly train the score network to approach the score of given samples. When the training is completed, the unnormalized density distribution is achieved, which Langevin dynamics can sample from. However, this generation model has some limitations on the density distribution. The Langevin dynamics sampler is not suitable for mixed distributions, and the results are not accurate in low-density regions.

2.2. Goodness-of-Fit Test and Generative Models

In the goodness-of-fit test, we are given a finite set of samples ${\left\{y_j\right\}}_{j=1}^n$, which are samples from an unknown distribution model $p\left(\theta ;z\right)$ and an unnormalized theoretical model $q\left(x\right)$, which we already know. The goal of the goodness-of-fit test is to determine whether samples ${\left\{y_j\right\}}_{j=1}^n$ generated from distribution $p\left(\theta ;z\right)$ coincide with the theoretical density distribution $q\left(x\right)$. So, there are two hypotheses:

$$H_0:p=q{H}_1:p\ne q$$

(3)

We reject hypothesis $H_0$ or $H_1,$ based on the selected thresholds and statistics. It is rational to apply the goodness-of-fit test to generative models to measure the discrepancy between fake samples and the theoretical density distribution. References [22,30] both introduce the goodness-of-fit test into deep generative models, but they know nothing about what the quality of the trained model is and can only be measured via the quality of samples generated. Simultaneously, without an accurate density function, they cannot perform the related tasks of approximate inference. Also, they construct complicated autoencoder networks; we only use simple convolutional architectures mentioned in DCGAN [31], which make sample generation and approximate inference simple tasks.

3. Maximum Mean Discrepancy and Kernel Stein Discrepancy

In this section, we first introduce two methods of hypothesis testing: Maximum Mean Discrepancy (MMD) and Kernel Stein Discrepancy (KSD). We then demonstrate the relationship between them.

3.1. Maximum Mean Discrepancy: A Two-Sample Test Method

From the perspective of hypothesis testing, the most popular solution paradigm of density distribution is the two-sample test. During the two-sample test, equivalent samples are picked from two density distributions separately. Then, we verify whether $p\left(\theta ;z\right)$ equals $q\left(x\right)$ by calculating the values of the statistics. Due to the spatial mapping capability, Maximum Mean Discrepancy (MMD) [21] is a valid and widely used two-sample test method. MMD maps the samples into high-dimensional even infinite feature space RKHS (also named Reproducing Kernel Hilbert Space) to measure their distance. For all $f\in \mathrm{RKHS}$, $E_{x}f=\langle f,{\mu }_p\rangle$ is called the mean embedding of $p$, among ${\mu }_p\in \mathrm{RKHS}$. If the distance between the mean embeddings of two density distributions is less than a certain threshold value, we can be sure that the two sets of samples are generated from the same distribution, namely $q\left(x\right)=p\left(\theta ;z\right)$.

MMD employs kernel functions to transform spatial coordinate bases. We denoted that $k\left(x, x^{\prime }\right)$ is a positive definite kernel and also a function: $\chi \times \chi \to R$. ${\mathcal{H}}_k$ is the Reproducing Kernel Hilbert Space (RKHS) defined by $k\left(x, x^{\prime }\right)$. RKHS is an inner product space with reproducibility, separableness, and completeness. In density distribution, a finite set of samples ${\left\{x_i\right\}}_{i=1}^n\subset \chi \subset R^d$ are sampled from an unknown density distribution $q\left(x\right).$ Another set of samples are defined as ${\left\{y_j\right\}}_{j=1}^m\subset Y\subset R^d$ (also named fake samples), and $y_j$ is an independent identically distributed sample from the density distribution $p\left(\theta ;z\right)$. The square of MMD can be represented by the mean embeddings ${\mu }_p$ and ${\mu }_q$, and they can be written as follows:

$$MM{D}_k^2={\Vert {\mu }_p-{\mu }_q\Vert }_{{\mathscr{H}}_k}^2=E_q\left[k\left(x_i,x_{i^{\prime }}\right)\right]-2E_{q,p}\left[k\left(x_i,y_j\right)\right]+E_p\left[k\left(y_j,y_{j^{\prime }}\right)\right]$$

(4)

3.2. Kernel Stein Discrepancy: A Goodness-of-Fit Test Method

In Formula (2), $S_q\left(x\right)={\nabla }_x\log q\left(x\right)$ is also named a Stein score function. The condition is that $q\left(x\right)$ is a continuous and differentiable density function and that the derivative of $q\left(x\right)$ exists. In [25], a continuous differentiable function $f:\chi \to R$ is defined in the Stein class of $p$ when it satisfies the following:

$${\displaystyle \underset{x\in X}{\int }\left(f\left(x\right)p\left(x\right)\right)dx}=0$$

(5)

When $F\left(x\right)=\left[f_1\left(x\right), \cdots ,f_r\left(x\right)\right]$ is a vector function, $F\left(x\right)$ is in the Stein class of $p\left(x\right),$ if $\forall i\in r$ and $f_i$ meets the condition (5). Assume $\chi =R^d$. Through integrating, by parts, the term on the left-hand side of Equation (5), Formula (5) is transformed into the following:

$${\lim }_{\Vert x\Vert \to \infty }f\left(x\right)p\left(x\right)=0$$

(6)

The Stein operator of $p$ is named $A_{p}f\left(x\right)$; it is a linear operator of $f\left(x\right)$. Among them, $f\left(x\right)$ is in the Stein class of $p\left(x\right). A_{p}f\left(x\right)$ is defined as follows:

$$A_{p}f\left(x\right)=s_p\left(x\right)f\left(x\right)+{\nabla }_{x}f\left(x\right)$$

(7)

Inspired by the equation $E_p\left[A_{q}f\left(x\right)\right]=E_p\left[\left(s_q\left(x\right)-s_p\left(x\right)\right)f{\left(x\right)}^T\right],$ Kernelized Stein Discrepancy (KSD) [24,25] is proposed to measure the distance between the two distributions. KSD is denoted as follows:

$$S\left(q,p\right)=E_p\left[{\left(s_q\left(y_j\right)-s_p\left(y_j\right)\right)}^{T}k\left(y_j,y_j^{\prime }\right)\left(s_q\left(y_j^{\prime }\right)-s_p\left(y_j^{\prime }\right)\right)\right]$$

(8)

In the above equation, $s_q\left(y_i\right)$ and $s_p\left(y_j\right)$ are the score function of $q$ and $p$; they are unknown in most unnormalized models, and we only have their samples in hand. In [25], the author eliminated the explicit expression of $p$ by applying the constraints of the Stein class. Kernel function $k\left(y_j, y_j^{\prime }\right)$ is a special continuous differentiable function $f$. And, it is easy to prove that if $k\left(y_j,y_j^{\prime }\right)$ has continuous second-order partial derivatives, $k\left(y_j,y_j^{\prime }\right)$ is in the Stein class of $p$. In other words, $k\left(y_j,y_j^{\prime }\right)$ meets the condition in Equation (5). Then:

$$E_p\left[k\left(y_j,y_j^{\prime }\right)\left(s_q\left(y_j^{\prime }\right)-s_p\left(y_j^{\prime }\right)\right)\right]=E_p\left[k\left(y_j,y_j^{\prime }\right)s_q\left(y_j^{\prime }\right)+{\nabla }_{y_j^{\prime }}k\left(y_j,y_j^{\prime }\right)\right]$$

(9)

Let us plug Formula (9) into Formula (8) and simplify:

$$\begin{array}{c}S\left(q,p\right)=E_p\left[u_q\left(y,y^{\prime }\right)\right]\hfill \\ u_q\left(y,y^{\prime }\right)=s_q{\left(y_j\right)}^{T}k\left(y_j,y_j^{\prime }\right)s_q\left(y_j^{\prime }\right)+s_q{\left(y_j\right)}^T{\nabla }_{y_j^{\prime }}k\left(y_j,y_j^{\prime }\right)\\ +{\nabla }_{y_j}k\left(y_j,y_j^{\prime }\right)s_q\left(y_j^{\prime }\right)+trace\left({\nabla }_{y_j,y_j^{\prime }}k\left(y_j,y_j^{\prime }\right)\right)\end{array}$$

(10)

We can compute $S\left(q,p\right)$ without the density distribution $p$ and only need score function $s_q\left(y_i\right)$ of the density distribution $q$. And, when $k\left(y_j,y_j^{\prime }\right)$ is a strictly positive definite kernel, it is easy to demonstrate that $p=q$, when and only when $S\left(q,p\right)=0$. In [25], KSD is applied to bootstrap goodness-of-fit testing to measure differences between two probability distributions. That is, we can calculate the KSD between the density distribution $q$ and samples from $p$ to determine whether $q$ equals $p$.

3.3. The Connection between MMD and KSD

In this section, we will discuss and demonstrate the connection between MMD and KSD below.

Lemma 1.

If $k\left(y_j,y_j^{\prime }\right)$ is a strictly positive definite kernel, the kernel $k^{\prime }\left(y_j,y_j^{\prime }\right)$ is still a strictly positive definition.

$$k^{\prime }\left(y_j,y_j^{\prime }\right)={\left(s_q\left(y_j\right)-s_p\left(y_j\right)\right)}^{T}k\left(y_j,y_j^{\prime }\right)\left(s_q\left(y_j^{\prime }\right)-s_p\left(y_j^{\prime }\right)\right)$$

(11)

Proof.

It is easy to demonstrate this using the definition of strictly positive definite kernel and eigenvalue decomposition. □

Assume $k^{\prime }\left(y_j,y_j^{\prime }\right)$ is a manually assigned and specific positive definite kernel. Replace the kernel function $k$ in Formula (4) with $k^{\prime }\left(y_j,y_j^{\prime }\right)$; the square of MMD can be written as below:

$$\begin{array}{c}MM{D}_{k^{\prime }}^2=E_q\left[s_q{\left(s_q\left(x_i\right)-s_p\left(x_i\right)\right)}^{T}k\left(x_i,x_i^{\prime }\right)\left(s_q\left(x_i^{\prime }\right)-s_p\left(x_i^{\prime }\right)\right)\right]\\ -2E_{p,q}\left[{\left(s_q\left(x_i\right)-s_p\left(x_i\right)\right)}^{T}k\left(x_i,y_j\right)\left(s_q\left(y_i\right)-s_p\left(y_j\right)\right)\right]\\ +E_p\left[{\left(s_q\left(y_j\right)-s_p\left(y_j\right)\right)}^{T}k\left(y_j,y_j^{\prime }\right)\left(s_q\left(y_j^{\prime }\right)-s_p\left(y_j^{\prime }\right)\right)\right]\end{array}$$

(12)

Lemma 2.

If $k\left(y_j,y_j^{\prime }\right)$is both in the Stein class of $p$ and $q$, then:

$$MM{D}_{k^{\prime }}^2=E_p\left[{\left(s_q\left(y_j\right)-s_p\left(y_j\right)\right)}^{T}k\left(y_j,y_j^{\prime }\right)\left(s_q\left(y_j^{\prime }\right)-s_p\left(y_j^{\prime }\right)\right)\right]$$

(13)

Lemma 2 is proved in Appendix A. We can see the right-hand sides of Formula (8) and Formula (13) make no odds. So, the Maximum Mean Discrepancy is equivalent to the Kernelized Stein Discrepancy when the kernel satisfies the conditions in Lemma 2.

4. Learning the Kernel Stein Discrepancy of Energy-Based Models

In the above section, we introduced two methods of hypothesis testing. MMD has been widely used in various unnormalized models [32,33]. But, due to the need to estimate the theoretical density distribution of $q$, KSD’s use in unnormalized models is extremely limited, only applied in the mixed Gaussian model [34]. In this article, we propose an algorithm for solving Problem 1 by using Kernel Stein Discrepancy to train unnormalized density models.

The bottleneck of measuring KSD as the explicit density estimation of $q$ is unsuspected. In [35], two generative models were proposed for mutual calibration between the explicit and implicit generative model. The difference between explicit energy-based models and the implicit generative model is measured through the Stein discrepancy. But, its objective function includes three kinds of statistical distances, which increases the complexity of the calculation. And, in the experimental part, no experimental results on likelihood inference are given. Inspired by [35] and the right-hand side of the third equation in Formula (2), it is convenient to obtain the explicit Gibb distribution of samples in explicit energy-based models. Enlightened by this, we replaced $s_q\left(y_j\right)$ in Formula (10) with $-{\nabla }_{y_j}e\left(y_j\right)$. The Kernelized Stein Discrepancy of energy-based models is defined as follows:

$$\begin{array}{c}{\mathrm{KSD}}_e\left(q,p\right){=E}_p[u_q\left(y,y^{\prime }\right)]\hfill \\ u_q\left(y,y^{\prime }\right)={\left[{\nabla }_{y_j}e\left(y_j\right)\right]}^{T}k\left(y_j,y_j^{\prime }\right){\nabla }_{y_j}e\left(y_j\right)-{\left[{\nabla }_{y_j}e\left(y_j\right)\right]}^T{\nabla }_{y_j^{\prime }}k\left(y_j,y_j^{\prime }\right)\\ -{\nabla }_{y_j}k\left(y_j,y_j^{\prime }\right){\nabla }_{y_j^{\prime }}e\left(y_j^{\prime }\right)+trace\left({\nabla }_{y_j,y_j^{\prime }}k\left(y_j,y_j^{\prime }\right)\right)\end{array}$$

(14)

In the above formula, the explicit density estimation of $q$ is limited to the Gibb distribution of the energy-based model. And, $e$ is the energy value of input samples. Although designing the energy function $e$ is unconstrained, it is hard to find the right energy function for unknown density distribution. Because we know nothing about the distribution of given samples, we have to look for the right energy function $e$ within a larger function range. In the vast majority of cases, a deep neural network is selected as the energy function, attributed to its strong nonlinear fitting ability and large value range. So, our model also selects a deep energy network as the energy function $e$. $k\left(y_j,y_j^{\prime }\right)$ in Formula (14) is a strictly positive definite kernel, which maps input samples to the Reproducing Kernel Hilbert Space and returns their inner product. The choice of kernel function for high-dimensional data is always a difficult issue. Different kernel functions map input samples to different feature spaces, and sample features and their inner product values vary accordingly. Here, we select the RBF (Radial Basis Function) as the kernel function, also known as the Gaussian kernel function. The RBF kernel function is defined as follows:

$$k\left(x,x^{\prime }\right)=\exp \left(-\frac{1}{2{\sigma }^2}{\Vert x-x^{\prime }\Vert }^2\right)$$

(15)

The RBF kernel meets the conditions we described in Lemma 2 and has a continuous second-order partial derivative. Hyperparameter σ needs to be set manually. The RBF kernel maps samples to infinite dimensional space, has fewer parameters, and has simple model complexity. Therefore, it is selected as the kernel function to be used in various application scenarios. Then, we can obtain the kernel matrix of input samples by calculating the RBF kernel function. ${\nabla }_{y_j^{\prime }}k\left(y_j,y_j^{\prime }\right)$, ${\nabla }_{y_j}k\left(y_j,y_j^{\prime }\right),$ and $trace({\nabla }_{y_j,y_j^{\prime }}k\left(y_j,y_j^{\prime }\right))$ can be established by taking the derivative of the matrix. Through the reverse derivation of the deep neural network, the value of ${\nabla }_{y_j^{\prime }}e\left(y_j^{\prime }\right)$, ${\nabla }_{y_j}e\left(y_j\right)$ can be obtained.

So, in Formula (14), each term of $KS{D}_e\left(q,p\right)$ can be calculated and measures the degree of matching between the input samples and density distribution $q$ represented by the energy-based model. The Gibbs distribution of the energy-based model removes the integral term by derivation and greatly simplifies the operation. Then, we analyze the application of KSD in training unnormalized models. The input samples ${\left\{y_j\right\}}_{j=1}^m\subset Y\subset R^d$ are a finite set of fake samples from a density distribution $p$. Samples of the density distribution $p$ should be readily available. And, the density distribution $q$ is represented by the energy-based model, which $p$ aims to approximate. Since we know nothing about $q$ at the beginning, the initial $p$ must be far away from $q$. We calculate the $KS{D}_e$ of the energy-based model in Formula (14) and update the fake samples of $p$ and $p$ iteratively, which will make $p$ slowly approach $q$. Finally, when the value of $KS{D}_e$ is less than a given small value, we think $p\mathrm{and}$ $q$ are the same density distribution.

5. Training Unnormalized Models with KSD

In Formula (14), fake samples $y_j,y_j^{\prime }$ and the energy model $e$ need to be provided. In this paper, to avoid bias caused by samplers and to reduce the cost of sample generation, we adapt a specialized generative network of the generative adversarial network to generate fake samples $y_j,y_j^{\prime }$. Meanwhile, we choose the explicit energy-based model as an estimation of unnormalized density distribution $q$, which can also be employed as the discriminator in generative adversarial networks.

5.1. Generative Adversarial Network (GAN)

The generative adversarial network is a popular method for training generative models. The generative adversarial network (GAN) [13] involves a generator G and a discriminator (or critic) D. The goal of the generator is to transform latent codes $z\sim p\left(z\right)$ into generated samples, which the discriminator cannot distinguish from real data samples, and the goal of the discriminator is to distinguish real data samples and generated samples as well as possible. With i.i.d. data samples, the optimal generator parameter ${\theta }_g$ and discriminator parameter ${\theta }_d$ are the solution to the following minmax game:

$$\underset{{\theta }_g}{L=min}\underset{{\theta }_d}{max}V\left({\theta }_g,{\theta }_d\right)={\mathrm{E}}_{x\sim q_{}\left(x\right)}\left[\log D\left(x,{\theta }_d\right)\right]+{\mathrm{E}}_{z\sim q\left(z\right)}\left[\log \left(1-D\left(G\left(z,{\theta }_g\right),{\theta }_d\right)\right)\right]$$

(16)

The minimax game highlights the adversarial nature of the training process. The discriminator, parameterized by ${\theta }_d$, aims to maximize its ability to differentiate between real and generated data, while the generator, parameterized by ${\theta }_g$, endeavors to minimize the discriminator’s accuracy by producing increasingly realistic data. ${\mathrm{E}}_{x\sim q_{}\left(x\right)}\left[\log D\left(x,{\theta }_d\right)\right]$ term signifies the expected log probability over real data samples $x$ drawn from the true distribution $q\left(x\right)$ that the discriminator D correctly identifies as genuine. The goal for the discriminator is to minimize this value, ensuring a high probability of correctly classifying real samples. ${\mathrm{E}}_{z\sim q\left(z\right)}\left[\log \left(1-D\left(G\left(z,{\theta }_g\right),{\theta }_d\right)\right)\right]$ denotes the expected log probability, for noise variables z sampled from a distribution $q\left(z\right)$, that the discriminator D correctly labels the data generated by the generator G (using the noise variables z) as fake. The discriminator aims to maximize this value to efficiently identify and label the artificially generated data, while the generator seeks to minimize this, thereby improving its ability to generate data that the discriminator misclassifies as real. Together, these terms formulate the objective function for the GAN, capturing the adversarial relationship between the generator and discriminator during the training process.

5.2. Generative Adversarial Networks for Unnormalized Density Models

The discriminator of traditional GAN solves classification problems. The discriminator is a single or multiple classification network. We study the use of the discriminator in unnormalized density problems.

In [36], EBGAN was proposed to find the data manifold based on the energy model. This assigns low-energy values to data manifold and high-energy values to other data regions. The mission of the discriminator is to discern whether the data are from the true samples manifold. Nevertheless, the generator is trained to approximate the data manifold to confuse the discriminator, so that the discriminator loses discernment. The antagonistic process can be defined by the loss:

$$\begin{array}{c}{L}_D=D\left(x_i\right)+{\left[m-D\left(G\left(z\right)\right)\right]}^{+}\\ L_G=D\left(G\left(z\right)\right)\hfill \end{array}$$

(17)

$L_D$ is the loss for the discriminator; $D\left(x_i\right)$ is the discriminator’s output for a real data sample $x_i$. $D\left(G\left(z\right)\right)$) is the discriminator’s output for a generated sample $G\left(z\right)$, m is a positive margin, and ${\left[\right]}^{+}$ denotes the positive part function, defined as ${\left[y\right]}^{+}=\max \left(0,y\right)$. Following this adaptation, the model is equipped to tackle issues pertaining to unnormalized density problems. However, the autoencoder employed within the discriminator proved computationally intensive. The Stein score function $S_q\left(x\right)$ of the autoencoder is hard to obtain.

5.3. PT KSD GAN (Kernel Stein Discrepancy Generative Adversarial Network with a Pulling-Away Term)

In the above section, we analyzed the disadvantages of the autoencoder network architecture in EBGAN. Consequently, we optimized it to a standard neural network architecture. To ensure stability, we adjusted the hyperparameter $m$ in Formula (17) to 1. The output of the discriminator is regarded as energy value of samples, which is the numerator of the explicit Gibbs density estimation. As stated in [25], the Kernel Stein Discrepancy (KSD) can measure the distance of different density distributions. That is, minimizing the $KS{D}_e$ of the generator is equivalent to approaching fake samples, coinciding with the theoretical distribution of true samples. The modified formulation of $L_D$ and $L_G$ is presented as follows:

$$\begin{array}{c}{L}_D=D\left(x_i\right)+\left(1-D\left(G\left(z\right)\right)\right)\\ L_G={\mathit{KSD}}_e\left(q,p\right)\hfill \end{array}$$

(18)

In the above formula, the generator and the discriminator are iteratively updated until, finally, the discriminator cannot tell whether the input sample is real or fake. We call this process of adversarial generation in Formula (18), KSD GAN. The discriminator is conducted as an energy-based model during training to discover the explicit Gibbs distribution of real samples. At the beginning, fake samples generated by the generator are very far from the density distribution region of real samples, and the discriminator can easily identify fake samples from real ones. During alternating iterations, samples generated by the generator are constantly narrowing $KS{D}_e$ from the theoretical distribution of real samples, which is represented by the energy model distribution of the discriminator. As the number of iterations increases and the Kernel Stein Discrepancy shrinks, samples generated by the generator approximate the Gibbs distribution of true samples, fooling the discriminator into thinking they are real samples.

In fact, this iterative process of KSD GAN can also be regarded as the confrontation process between the implicit generation model and the explicit generation model. The discriminator is trained as an explicit generation model, and its density distribution is the Gibbs distribution based on the energy model. Of course, we can use the sampler to sample from this Gibbs distribution, which we want to obtain. As an implicit generation model, the generator bypasses the estimation of the density distribution and generates imitation samples directly by reducing $KS{D}_e$.

GAN is notoriously difficult to train, and mode collapse often occurs. This is due to the widely dispersed density distribution of high-dimensional datasets, and there is a huge span between regions. It is very possible for the generator to find a local optimal and then stop optimizing. Then, samples generated by the generator are only one or several types, and the sample diversity is insufficient. Inspired by [36,37], reducing the similarity and increasing orthogonalization of generated samples in energy-based models can avoid model collapse and speed up convergence. We add a repelling regularizer named the Pulling-away Term (PT) to the generator loss:

$$PT\left(S\right)=\frac{1}{M\left(M-1\right)}{\displaystyle \sum _i{{\displaystyle \sum _{j\ne i}\left(\frac{S_i^T{S}_j}{‖S_i‖‖S_j‖}\right)}}^2}$$

(19)

After adding the regular term to the generator, the loss $L_D$ and $L_G$ of KSD GAN are defined as follows:

$$\begin{array}{c}{L}_D=D\left(x_i\right)+\left(1-D\left(G\left(z\right)\right)\right)\\ L_G=\mathrm{KSD}\left(q,p\right)+\alpha *PT\hfill \end{array}$$

(20)

In Formula (20), KSD GAN with the regular term PT added is called PT KSD GAN. Our model structure of PT KSD GAN is shown in Figure 1.

The figure above shows that our model is actually an enhanced GAN framework incorporating the Kernelized Stein Discrepancy. The data processing steps are described below:

1

Two data streams, denoted as $x_i$ and $y_j$, are added to the system. The former is directed towards the discriminator, while the latter interfaces with the generator.

2

Acting as the GAN’s energy model, the discriminator, $D\left(x\right)$ processes the data $x$. Its primary responsibility is to differentiate between the real sample $x_i$ and fake sample $y_j$ generated by the generator.

3

Simultaneously, the generator $G\left(z\right)$, creates synthetic samples based on an input noise vector z. This noise vector serves as a seed, ensuring diversity in generated samples.

4

The fake samples are subjected to the PT Kernel Stein Discrepancy module. The integration of KSD, a non-parametric measure of discrepancy between distributions, aims to quantitatively evaluate the deviation of the generated samples from the target distribution. This evaluation imparts valuable feedback to the generator, steering it towards optimal data synthesis.

5

The outcomes from the PT Kernel Stein Discrepancy and Mean Absolute Error modules are collated to furnish a composite gradient that guides the generator’s training. By assimilating feedback not only from the discriminator but also from an external discrepancy measure (KSD), the generator’s training becomes more informed, focused, and robust against mode collapse.

6. Experiment

We introduce the PT KSD GAN for training unnormalized models (Figure 1). Using an energy-based discriminator, it assigns energy values to samples, while the generator, minimizing Kernel Stein Discrepancy, produces samples reflecting true density. Initial tests on 2D toy datasets (Figure 2) validate the generator’s alignment with true density and sample diversity.

We further validate using linear Independent Component Analysis (ICA), a standard for unnormalized model training. ICA decomposes signals into independent sources. Our focus was on retrieving original latent variables from mixtures, with the results shown in

Table 1. Linear ICA training results.

Method	Data Dimension
	10	20	30	40	50
Max. Likelihood	−10.98	−18.48	−21.49	−23.43	−25.53
PT KSD GAN	−10.90	−18.35	−21.21	−25.18	−25.49
LSD	−10.95	−18.37	−21.23	−25.14	−25.36
Score Matching	−11.13	−27.20	−21.48	NaN	NaN
NCE	−10.92	−22.52	−30.33	−55.53	−73.62
CNCE	−11.00	−18.77	−24.47	−37.64	−36.31

.

Finally, PT KSD GAN’s performance on image datasets like MNIST and CIFAR-10, using a DCGAN framework, is depicted in Figure 3, indicating realistic image generation with room for refinement.

6.1. Toy Datasets

We initiate our study by assessing PT KSD GAN’s performance on two-dimensional toy datasets. These datasets, derived from Sklearn’s library, provide a range of shapes in two-dimensional data. For the generator, we use inputs $\{z_1,z_2,\dots ,z_N$} sampled from a standard normal distribution. Both the generator and discriminator are structured as a two-layer multi-layer perceptron, producing two-dimensional outputs. Key hyperparameters include an RBF kernel function bandwidth $\sigma$ of 1.0; a regularized term coefficient α of 0.8; and Adam Optimizer settings with a learning rate of 0.002, ${\beta }_1$ of 0.5, and ${\beta }_2$ of 0.99. We standardize the batch size to 1000 across all toy datasets and set a training cap at 10,000 iterations. The outcomes of this training are illustrated in Figure 2.

Figure 2 delineates the density distributions: real samples are shown on the left, the energy-based model’s output (or the discriminator) in the center, and the generator-produced sample density on the right. Notably, the energy-based model provides unnormalized energy values, leading to density estimates that capture only the approximate shape of the true samples’ density distribution. While potential exists to use samplers like the Hamiltonian Monte Carlo, biases from initial probability distributions and other constraints deterred us from this approach.

Our approach employs specialized networks to produce fake samples, aiming to minimize the Kernel Stein Discrepancy against the true distribution. As evident in Figure 2, the generator’s outputs align closely with the true density distribution, ensuring sample diversity and authenticity. The PT KSD GAN training process can be viewed as a competitive interaction between two generative networks. While the discriminator captures the true samples’ explicit Gibbs distribution, the generator focuses on producing realistic samples without explicit density distribution searches.

6.2. Linear ICA

Linear Independent Component Analysis (ICA) [26] is a statistical method designed to separate a multivariate signal into independent non-Gaussian signals. It is particularly useful in scenarios where multiple signals overlap, and the objective is to discern the original signals from their mixtures. Building on this, we further explored the efficacy of PT KSD GAN within the context of the linear ICA model. This model serves as a benchmark in the evaluation of algorithms tailored for training unnormalized models.

The generative process of true samples is relatively simple. The latent variables $s_i,i=1,\dots ,n$ are independent and identically distributed, and they are sampled from the standard Laplace distribution, that is, $\mathbf{s}~Laplace\left(0,1\right)$. Then, take a linear transformation:

$$x=Ws$$

(21)

The model parameter $W$ is an $N\times N$ matrix, and the determinant is not zero. Under Formula (21), we can deduce $\mathit{x}$’s log density:

$$\log p(x;W)=\log p_s(W^{-1}x)-\log \left|W\right|$$

(22)

$p_s\left(·\right)$ is the probability density function of latent variables s. Thus, (22) can be used as non-normalized likelihood to evaluate generative models. There are two hyperparameters: the bandwidth parameter σ and the regularizer $\alpha$. The bandwidth parameter σ of the RBF kernel function is self-adapting according to the input sample number and sample feature dimension. Sample dimensions $N$ range from 2 to 50. The value of regularized term coefficients is 0.8. We set the learning rate of the Adam Optimizer as $0.001$ and the two hyperparameters of the Adam Optimizer as ${\beta }_1=0.3$, ${\beta }_2=0.9$. The number of iterations is 100,000, and the batch size is set to 1000. See Table 1 for the results of linear ICA with PT KSD GAN.

Table 1 shows the log likelihoods of various algorithms: score matching [26], maximum likelihood [38], Learned Stein Discrepancy (LSD) [39], and Noise-Contrastive Estimation (NCE) [40]. For low-dimensional samples, we find that the power of PT KSD GAN is optimal. As the dimensions increase, the performance of PT KSD GAN is slightly inferior to ML and LSD. The reason is that the manual selection of kernel functions makes it impossible to select the optimal kernel function for the current dataset. On the whole, our approach can achieve the same performance as maximum likelihood.

6.3. Preliminary Image Results

To validate our approach’s applicability to train more expressive unnormalized models and to generate samples, we conducted experiments assessing PT KSD GAN’s performance on high-dimensional image datasets, specifically MNIST and CIFAR-10, with a training capacity of 50K images. Our experimental setup adhered to the DCGAN [31] model architecture, employing simple convolutional neural networks for both the generator and the discriminator. The discriminator computes the energy value for each image vector, and we subsequently averaged these values to determine the energy for an image batch. Optimization was performed using the Adam Optimizer, set with a learning rate of 0.00005 and hyperparameters ${\beta }_1=0.5$ and ${\beta }_2=0.999$. We standardized the batch size to 64 across all datasets.

Given the intricate nature of the image datasets, we restricted our regularization term calculations to the initial twenty images, enhancing computational efficiency. This means that the value of $M$ in Formula (20) remains consistent at 20 for each training batch. This approach significantly elevates the diversity of the generated images. The results, as depicted in Figure 3, show that the images produced from MNIST and CIFAR-10 training are convincingly realistic, though there is potential for refining image precision in future iterations.

The images displayed in the figure represent the generator’s outputs within the adversarial network. While it is feasible to use a sampler to derive samples from the discriminator’s trained energy model, we prioritize efficiency in large-scale sample generation and, thus, opt for the generation network over the sampler. Admittedly, the quality of these generated samples is not the best in the field. However, they mark a pioneering demonstration of KSD’s ability to manage high-dimensional data evaluation and generation. Future endeavors could explore intricate network architectures and sophisticated samplers to handle even more complex image datasets.

7. Conclusions

In this paper, we demonstrated that Kernel Stein Discrepancy is an efficient way to train unnormalized density models and to generate samples from these models. In Section 3, we connected the two-sample test method and the goodness-of-fit test by demonstrating that Kernel Stein Discrepancy is equal to Maximum Mean Discrepancy under certain assumptions. We introduced a new adversarial learning paradigm to train the unnormalized density model, named PT KSD GAN. The discriminator was trained as an energy model to determine the explicit Gibbs density distribution of real samples. Low energies were assigned to regions within the true data manifold, while high energies were assigned to other areas. The generator was responsible for generating fake samples to fool the discriminator. The quality of the generated samples was measured using a goodness-of-fit test method, KSD.

Three experiments were conducted to prove the feasibility and robustness of our algorithm. In toy datasets, an unnormalized energy model of PT KSD GAN can capture the approximate and blurry shape of the density distribution of true samples. The generated samples lie strictly within the true density distribution. The generator strictly and precisely identifies the true data manifold. In the linear Independent Component Analysis task, the performance of PT KSD GAN outperformed other unnormalized training methods when the data dimension was less than 30. The manual selection of kernel functions and hyperparameters limits the application of KSD in high-dimensional datasets. In more complex image datasets, the results showed that KSD is no longer cursed by high-dimensional datasets; however, the sharpness of the generated pictures should be improved in future work.

Unsupervised models are a critical area of machine learning, where the estimation of the density distribution is important for unlabeled samples. With the density distribution, it is possible to generate similar samples and to determine whether the input sample is abnormal. In future work, we will study the application of adversarial generation networks and hypothesis testing methods in anomaly detection.