JEM++: Improved Techniques for Training JEM

TL;DR

This paper introduces new training techniques and architectural improvements for JEM, enhancing its accuracy, stability, and efficiency, and enabling the use of batch normalization in energy-based models.

Contribution

The paper presents novel methods including proximal SGLD, a multi-step differential game framework, informative initialization, and batch normalization integration for JEM.

Findings

Improved training stability with proximal SGLD.

Significant acceleration of training via the extended YOPO framework.

Enabling batch normalization in JEM enhances model performance.

Abstract

Joint Energy-based Model (JEM) is a recently proposed hybrid model that retains strong discriminative power of modern CNN classifiers, while generating samples rivaling the quality of GAN-based approaches. In this paper, we propose a variety of new training procedures and architecture features to improve JEM's accuracy, training stability, and speed altogether. 1) We propose a proximal SGLD to generate samples in the proximity of samples from the previous step, which improves the stability. 2) We further treat the approximate maximum likelihood learning of EBM as a multi-step differential game, and extend the YOPO framework to cut out redundant calculations during backpropagation, which accelerates the training substantially. 3) Rather than initializing SGLD chain from random noise, we introduce a new informative initialization that samples from a distribution estimated from training data. 4) This informative initialization allows us to enable batch normalization in JEM, which further releases the power of modern CNN architectures for hybrid modeling. Code: https://github.com/sndnyang/JEMPP