Expanding the Role of Diffusion Models for Robust Classifier Training

Pin-Han Huang; Shang-Tse Chen; Hsuan-Tien Lin

cs.LGFeb 23, 2026

Expanding the Role of Diffusion Models for Robust Classifier Training

Authors: Pin-Han Huang, Shang-Tse Chen, Hsuan-Tien Lin

Organizations: National Taiwan University

Abstract

Incorporating diffusion-generated synthetic data into adversarial training (AT) has been shown to substantially improve the training of robust image classifiers. In this work, we extend the role of diffusion models beyond merely generating synthetic data, examining whether their internal representations, which encode meaningful features of the data, can provide additional benefits for robust classifier training. Through systematic experiments, we show that diffusion models offer representations that are both diverse and partially robust, and that explicitly incorporating diffusion representations as an auxiliary learning signal during AT consistently improves robustness across settings. Furthermore, our representation analysis indicates that incorporating diffusion models into AT encourages more disentangled features, while diffusion representations and diffusion-generated synthetic data play complementary roles in shaping representations. Experiments on CIFAR-10, CIFAR-100, and ImageNet validate these findings, demonstrating the effectiveness of jointly leveraging diffusion representations and synthetic data within AT.

Figures & tables

Figure 1: We plot robust accuracy and representation similarity scores Huh et al. (2024) for CIFAR-10 $\ell_{\infty}$ -robust models from RobustBench Croce et al. (2021) . Similarity scores are measured with respect to representations extracted from the diffusion model. Implementation details and discussion are in Appendix A .

Figure 2: Overview of Diffusion Representation Alignment (DRA). We leverage an auxiliary projection head to align classifiers with the extracted diffusion representations.

Panel of Figure 3 — Figure 3: (a) The frequency saliency analysis of the linear-probed diffusion representation, adversarial trained robust model, and standard trained non-robust model. Low frequencies are being centered. (b) The CIFAR-10 robust accuracy across perturbation budgets for the linear probed diffusion representation ( DR ), adversarial trained robust model ( AT ), and standard trained non-robust model ( ST ).

Figure 4: Alignment and uniformity metrics on CIFAR-10 for the standard-trained model ( ST ), the adversarially trained model ( AT ), and the diffusion representations ( DR ).

Dataset	Model	Method	Clean Acc.	AutoAttack Acc.
CIFAR-10 ( $\ell_{\infty}=8/255$ )	WRN-28-10	DM-AT	92.44	67.31
	WRN-28-10	DM-AT + DRA	93.14	67.83
	ViT-B/2	DM-AT	94.35	71.31
	ViT-B/2	DM-AT + DRA	95.22	71.77
CIFAR-100 ( $\ell_{\infty}=8/255$ )	WRN-28-10	DM-AT	68.34	35.72
	WRN-28-10	DM-AT + DRA	69.85	36.27

Table 1: Clean and Robust Accuracy incorporating Diffusion Representation Alignment on CIFAR-10, CIFAR-100, and ImageNet.

Method	Architecture	Synthetic	Batch	Epoch	Clean	AA
AT	WRN-34-10	-	128	200	84.33	55.25
AT+ADR Wu et al. (2024)	WRN-34-10	-	128	200	86.11	55.26
AT+IKL Cui et al. (2024)	WRN-34-10	-	128	200	84.80	57.09
AT+DRA (Ours)	WRN-34-10	-	128	200	88.54	57.29
DM-AT Wang et al. (2023b)	WRN-28-10	1M	512	400	91.12	63.35
DM-AT Wang et al. (2023b)	WRN-28-10	1M	1024	800	91.43	63.96

Table 2: Comparison with state-of-the-art adversarial robustness methods across different settings on CIFAR-10.

Figure 5: We plot the alignment and uniformity metrics, along with clean and robust accuracy on CIFAR-10 ( $\ell_{\infty}=8/255$ ) shown in parentheses.

Figure 6: We train Top-K SAEs on ImageNet with sparsity $K\in\{8,16,32\}$ , using ViT-B checkpoints trained with AT, DM-AT, and DM-AT+DRA. Normalized SAE reconstruction losses are reported for comparison Gao et al. (2025) .

Figure 7: Clean and robust accuracy evaluated with a CIFAR-10 robust WRN-28-10 model when projecting features from natural or adversarial images onto the top- $K$ principal components before the classification head.

Method	CLS-95 Dim	CLS-99 Dim	Robust Dim
AT	9	14	9
AT+DRA	15	42	22
DM-AT	10	11	11
DM-AT+DRA	12	15	23

Table 3: We evaluate Classification Dimension , with CLS-95 and CLS-99 refering to the number of components required to recover 95% and 99% of the original classification performance, and Robust-aware dimension, the number of principal components where robust accuracy peaks, across different training methods on CIFAR-10 with WRN-28-10.

Figure 8: Clean and Robust accuracy on CIFAR-100 for DM-AT+DRA with different DRA regularization strength $\lambda$ .

Appendix figures & tables2 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 9: Similarity scores are measured with respect to representations extracted from the diffusion model ( red ) and from the RobustBench standard-trained model ( gray ).

Panel of Figure 10 — Figure 10: The input gradient frequency difference maps between DM-AT+DRA and DM-AT. (a) DRA aligned to noisy-input discriminative pretrained representations. (b) DRA aligned to the original diffusion representations. Positive values indicate increased sensitivity. Low-frequency components are in the upper-left.

Explore similar work

Representation-Conditioned Diffusion Models for Guided Training Data Generation
May 26, 2026Nithesh Chandher Karthikeyan, Jonas Unger, Gabriel EilertsenSynthetic Data Augmentation Latent Diffusion Models
Linkoping University
PAPT++: Risk-Aware Adversarial Tuning and Generation for Single Domain Generalization
Sep 4, 2026Zhipeng Xu, De Cheng, Xinyang Jiang +5Synthetic Data Augmentation Diffusion Model Guidance
Xidian University, Xi’an, China · Microsoft Research Asia, Shanghai, China
TEMPO-Diffusion: Temporally Exposed Malicious Poisoning of Diffusion Models
Jun 24, 2026William Aiken, Paula Branco, Guy-Vincent Jourdan +1Backdoor Attacks Diffusion-Based Image Generation
School of Electrical Engineering and Computer Science, University of Ottawa, Ottawa, ON K1N 6N5, Canada

May 26, 2026J/K move · Enter open · S save

Representation-Conditioned Diffusion Models for Guided Training Data Generation

Nithesh Chandher Karthikeyan, Jonas Unger, Gabriel Eilertsen

Linkoping University

Synthetic Data Augmentation Latent Diffusion Models Synthetic Data Generation Conditional Diffusion Models Image Classification

Open paper AskPDF arXiv

Data availability remains a critical bottleneck in many deep learning applications. Large-scale datasets are often expensive to collect, curate and annotate, which can limit the scalability and applicability of supervised learning methods. In this work, we evaluate the classification performance of models trained on synthetic image datasets produced by generative deep learning. In particular, we use latent diffusion models conditioned on learned representations from DINOv2, DINOv3, and CLIP. Our results demonstrates that this representation-conditioned formulation significantly outperforms class-conditioned generation by a large margin (+10.76 p.p. top-1 accuracy on ImageNet100), by improving sample quality and mode coverage. Furthermore, by scaling the size of the synthetic dataset, we are able to outperform a classifier trained on the real data (+2.0 p.p top-1 accuracy). We also demonstrate how generated images can be used for augmentation purposes, outperforming classical augmentation methods, and how the conditioning space can be used for sample filtering to further improve training value. Collectively, these findings highlight that representation-conditioned diffusion models provide a promising approach for augmenting, complementing, or potentially replacing real-world datasets in large-scale visual learning tasks.

Abstract

Explore similar work

Representation-Conditioned Diffusion Models for Guided Training Data Generation

PAPT++: Risk-Aware Adversarial Tuning and Generation for Single Domain Generalization

TEMPO-Diffusion: Temporally Exposed Malicious Poisoning of Diffusion Models