Scalable Energy-Based Models via Adversarial Training: Unifying Discrimination and Generation

Xuwang Yin, Claire Zhang, Julie Steele, Nir N Shavit, Tony T. Wang

Keywords: Joint Modeling, Energy-Based Models (EBMs), Adversarial Training, Robust Classification, Generative Modeling, PGD Attacks, explainability

TL;DR: Adversarial training replaces unstable SGLD in Joint Energy-Based Models, scaling EBM-based hybrids to ImageNet 256×256 with state-of-the-art robust classification and generation quality in a single model.

Abstract: Simultaneously achieving robust classification and high-fidelity generative modeling within a single framework presents a significant challenge. Hybrid approaches, such as Joint Energy-Based Models (JEM), interpret classifiers as EBMs but are often limited by the instability and poor sample quality inherent in training based on Stochastic Gradient Langevin Dynamics (SGLD). We address these limitations by proposing a novel training framework that integrates adversarial training (AT) principles for both discriminative robustness and stable generative learning. The proposed method introduces three key innovations: (1) the replacement of SGLD-based JEM learning with a stable, AT-based approach that optimizes the energy function through a Binary Cross-Entropy (BCE) loss that discriminates between real data and contrastive samples generated via Projected Gradient Descent (PGD); (2) adversarial training for the discriminative component that enhances classification robustness while implicitly providing the gradient regularization needed for stable EBM training; and (3) a two-stage training strategy that addresses normalization-related instabilities and enables leveraging pretrained robust classifiers, generalizing effectively across architectures. Experiments on CIFAR-10/100 and ImageNet demonstrate that our approach: (1) is the first EBM-based hybrid to scale to high-resolution datasets with high training stability, simultaneously achieving state-of-the-art discriminative and generative performance on ImageNet 256$\times$256; (2) uniquely combines generative quality with adversarial robustness, enabling faithful counterfactual explanations; and (3) functions as a competitive standalone generative model, matching autoregressive models and surpassing diffusion models while offering additional versatility.

Primary Area: unsupervised, self-supervised, semi-supervised, and supervised representation learning

Submission Number: 13551