Your Classifier Can Do More: Towards Bridging the Gaps in Classification, Robustness, and Generation

TL;DR

This work addresses the triad of accuracy, adversarial robustness, and generation in deep models. It introduces EB-JDAT, a unified energy-based framework that explicitly aligns the energy landscapes of clean, adversarial, and generated data through a joint distribution $p_{ heta}( extbf{x}, ilde{ extbf{x}}, y)$ and a min–max objective optimized with SGLD sampling. Empirical results on CIFAR-10, CIFAR-100, and an ImageNet subset show EB-JDAT achieves state-of-the-art robustness while preserving near-original clean accuracy and generation quality, effectively resolving the longstanding triple trade-off. The approach provides a principled, scalable path toward robust, generative discriminators suitable for real-world deployment.

Abstract

Joint Energy-based Models (JEMs) are well known for their ability to unify classification and generation within a single framework. Despite their promising generative and discriminative performance, their robustness remains far inferior to adversarial training (AT), which, conversely, achieves strong robustness but sacrifices clean accuracy and lacks generative ability. This inherent trilemma-balancing classification accuracy, robustness, and generative capability-raises a fundamental question: Can a single model achieve all three simultaneously? To answer this, we conduct a systematic energy landscape analysis of clean, adversarial, and generated samples across various JEM and AT variants. We observe that AT reduces the energy gap between clean and adversarial samples, while JEMs narrow the gap between clean and synthetic ones. This observation suggests a key insight: if the energy distributions of all three data types can be aligned, we might bridge their performance disparities. Building on this idea, we propose Energy-based Joint Distribution Adversarial Training (EB-JDAT), a unified generative-discriminative-robust framework that maximizes the joint probability of clean and adversarial distribution. EB-JDAT introduces a novel min-max energy optimization to explicitly aligning energies across clean, adversarial, and generated samples. Extensive experiments on CIFAR-10, CIFAR-100, and ImageNet subsets demonstrate that EB-JDAT achieves state-of-the-art robustness while maintaining near-original accuracy and generation quality of JEMs, effectively resolving the triple trade-off between accuracy, robustness, and generation.

Figures (5)

• Figure 1: Comparisons of SOTA AT-based methods on CIFAR-10 in terms of accuracy and robustness (AutoAttack). Our method achieves the best robustness while maintaining competitive standard accuracy.
• Figure 2: Distributions of the $E_\theta(\textbf{x},y)$ of adversarial samples for PGD-20 vs. generated samples vs. clean samples on CIFAR-10. indicates generated samples, indicates adversarial samples, indicates clean samples.
• Figure 3: Performance of generated samples on ImageNet subset (64x64) with EB-JDAT-JEM++.
• Figure 4: Performance comparison of generated samples from different methods, all methods are sampled by SGLD with informative initialization.
• Figure 5: Ablation study of adversarial sampling steps. ECO=100 indicates no collapse. Time denotes the per-epoch cost.