Joint Discriminative-Generative Modeling via Dual Adversarial Training

TL;DR

This paper tackles the longstanding challenge of building models that are simultaneously robust discriminators and high-fidelity generative storytellers. It introduces Dual Adversarial Training (DAT), a fusion of adversarial training for classification with a stable BCE-based energy-learning objective for generation within a Joint Energy-Based Model, replacing unstable SGLD-based learning with PGD-driven contrastive sampling. A two-stage training protocol and decoupled data augmentation enable scaling to high-resolution data (ImageNet 256×256) and leveraging pretrained robust classifiers, achieving state-of-the-art generative and discriminative performance while maintaining robustness and offering useful counterfactual explanations. Empirically, DAT matches or surpasses specialized generative models in quality, provides competitive robustness, and delivers superior throughput relative to diffusion models, marking a significant advance in unified discriminative-generative modeling.

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 Stochastic Gradient Langevin Dynamics (SGLD)-based training. 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 by discriminating between real data and Projected Gradient Descent (PGD)-generated contrastive samples using the BCE loss; (2) synergistic adversarial training for the discriminative component that enhances classification robustness while eliminating the need for explicit gradient penalties; and (3) a two-stage training strategy that addresses normalization-related instabilities and enables leveraging pretrained robust classifiers, generalizing effectively across diverse 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 critical applications like robust counterfactual explanations; and (3) functions as a competitive standalone generative model, matching the generative quality of autoregressive methods (VAR-d16) and surpassing diffusion models while offering unique versatility.

Figures (15)

• Figure 1: Counterfactual FIDs and classifier confidences under different perturbations.
• Figure 2: Samples produced by different augmentations on CIFAR-10.
• Figure 3: CIFAR-10 training curves under different data augmentations during stage 2 joint training.
• Figure 4: Calibration diagrams on CIFAR-10 (without temperature scaling).
• Figure 5: Calibration diagrams on CIFAR-100 (without temperature scaling).

Theorems & Definitions (2)

• Proposition A.1: Optimal class logits
• proof