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.