Expanding the Role of Diffusion Models for Robust Classifier Training

会议: NeurIPS 2026
arXiv: 2602.19931
OpenReview: https://openreview.net/forum?id=oRtG25ar63
领域: 图像生成
关键词: 扩散模型、对抗训练、对抗鲁棒性、表示对齐、特征解耦

一句话总结

打破扩散模型在对抗训练中仅用于合成数据采样的常规范式，提出扩散表示对齐（DRA）技术，将冻结扩散模型的中间层去噪表示作为辅助表征先验蒸馏给判别分类器，在零推理开销下显著提升模型的干净准确率、AutoAttack 鲁棒性以及特征空间的可解耦性。

研究背景与动机

深度图像分类器极易受到对抗样本的欺骗——极具隐蔽性、人类感知不可觉察的微小输入扰动足以彻底颠覆模型的预测结果。在众多经验防御机制中，基于投影梯度下降（PGD）扰动输入图像的对抗训练（Adversarial Training, AT）依然是公认最有效建立鲁棒分类器的主流方案。然而，对抗训练普遍遭受严重的“鲁棒过拟合”（robust overfitting）困扰，即在训练后期随着训练损失持续收敛，测试集上的对抗鲁棒性反而急剧恶化。近期的突破性进展（如 DM-AT 训练范式）证明，引入由预训练扩散模型大规模生成的逼真合成数据能够有效缓解数据饥渴与过拟合，成为当前 RobustBench 榜单上顶尖模型的标准配方。

然而，现存的扩散辅助对抗训练几乎完全将扩散模型视作纯粹的“离线黑盒样本发生器”。这种单向的数据流视角忽略了扩散模型内部极其丰富的表征结构：扩散模型在多步逐步去噪任务中所学习到的中间层激活，本质上编码了高度鲁棒的高层语义特征和全局上下文结构。已有研究表明，纯粹依靠像素重构的目标（如 MAE 等自编码器）由于过度依赖高频纹理信息，直接迁移至对抗训练时往往严重恶化下游鲁棒性；而扩散模型在带噪去噪过程中天然过滤了非本质的脆弱高频噪声，形成了更平滑、更具判别性的语义流形。学术界此前对能否在合成数据生成之外、直接挖掘并复用扩散模型自身的表征先验来赋能对抗鲁棒性，仍缺乏系统性探索。

针对该问题，本文提出扩展扩散模型在鲁棒分类器训练中的功能定位：核心 idea：将预训练冻结扩散模型的中间去噪表征作为对抗训练过程中的辅助正则化学习信号，通过轻量投影头显式实施特征对齐（DRA），与扩散合成数据协同作用，引导分类器在零额外推理代价下学习更具抗扰能力、更易解耦的低频鲁棒特征。

方法详解

整体框架

扩散表示对齐（Diffusion Representation Alignment, DRA）旨在不改变分类器本身推理结构与推断延迟的前提下，在对抗训练迭代中注入强鲁棒的生成式特征先验。方法整体流程为：给定带标签的训练输入图像，分类器首先通过多步 PGD 生成对抗扰动样本，并由分类器主干编码器提取其高维特征向量；同时，将原始图像加噪至去噪信噪比最优的时间步，输入至预训练且权重完全冻结的扩散模型（如 EDM 或 LightningDiT）中，抽取其瓶颈层或中间 Transformer 层的去噪激活表示。随后，通过一个可训练的轻量多层感知机（MLP）投影头将分类器对抗表征映射至扩散表征空间，并以余弦相似度形式施加正则化惩罚。最终目标函数将经典对抗鲁棒损失（如 TRADES）与 DRA 辅助对齐损失按固定权重线性加权，端到端联合更新分类器及投影头。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    In["输入图像与标签 (x, y)"] --> Gen["PGD 对抗攻击生成<br/>计算扰动样本 x_hat"]
    In --> DNoise["扩散前向加噪<br/>加噪至最优时间步 t 得 x_t"]
    Gen --> ClsEnc["分类器编码器<br/>抽取特征 h_CLS(x_hat)"]
    DNoise --> DiffFrozen["冻结扩散模型<br/>抽取中间层表征 h_DR(x_t, t, y)"]
    ClsEnc --> Proj["投影头表征映射<br/>MLP 将分类特征投射至扩散空间"]
    ClsEnc --> ClsHead["分类预测头<br/>计算任务交叉熵 / TRADES 损失"]
    DiffFrozen --> CosAlign["扩散特征对齐损失<br/>最大化与扩散表征的余弦相似度"]
    Proj --> CosAlign
    ClsHead --> TotalLoss["联合优化目标<br/>L_AT-DRA = L_AT + λ * L_DRA"]
    CosAlign --> TotalLoss

关键设计

1. 扩散特征先验提取：优选时间步与大激活异常校准

传统重构式自监督表征因过度保留图像中与分类无关的高频像素信号，迁移至对抗场景时容易引发频域敏感性放大。相比之下，扩散模型在适度噪声水平下的中间层激活展现出优异的均匀度（uniformity）与固有轻度鲁棒性。本文在提取扩散表征时，严格选取经验上线性探针判别质量最高的最优时间步（对于 EDM 选取噪声尺度 \(\sigma_t=0.1\)；对于 DiT 架构选取相对时间步 \(t=0.8\)）。针对最新 DiT 模型（如 LightningDiT）在特征提取中常出现的极少数通道异常巨大激活（large activations）现象，设计了通道范数异常检测过滤机制，对特定通道（如通道 1053 与 259）置零后执行自适应层归一化（AdaLN），从而为判别模型提供干净、平滑的生成式特征锚点。

2. 投影头特征对齐与联合正则化目标

为避免将分类器架构强制绑定于沉重的扩散生成器骨干，DRA 在分类器编码器 \(f_\theta\) 之上外挂了一个三层 MLP 构成的可学习投影头 \(g_{\text{proj}}\)，用于消除不同网络维度及语义空间的不匹配。在训练过程中，针对由 PGD 内部最大化生成的对抗样本 \(\hat{\mathbf{x}}\)，计算其在投影空间中与真实条件扩散特征 \(\mathbf{h}^{\text{DR}}_{\mathbf{x}_t, t, y}\) 的对齐损失：

\[\mathcal{L}_{\text{DRA}} = -\text{sim}\big(g_{\text{proj}}(f_\theta(\hat{\mathbf{x}})),\, \mathbf{h}^{\text{DR}}_{\mathbf{x}_t, t, y}\big)\]

其中相似度度量 \(\text{sim}(\cdot,\cdot)\) 采用余弦相似度。该对齐损失与分类任务的对抗目标函数（如 10 步 TRADES 损失 \(\mathcal{L}_{\text{AT}}\)）复合：

\[\mathcal{L}_{\text{AT-DRA}} = \mathcal{L}_{\text{AT}} + \lambda \mathcal{L}_{\text{DRA}}\]

正则化系数设定为 \(\lambda=1.2\)。这种机制驱使分类器在极度恶劣的对抗扰动下，其表征轨迹仍牢牢被拉向扩散模型所刻画的鲁棒语义流形，从特征空间源头阻断脆弱扰动的传播。

3. 频域偏置修正与特征解耦协同

频率显著性分析与稀疏自编码器（TopK-SAE）探针揭示了 DRA 的深层机制：扩散表征对齐有效压制了模型对中高频扰动的敏感度，赋予分类器类似低通滤波的频域不变性。在此基础上，本文深入剖析了合成数据与表征对齐的不同角色：通过主成分分析（PCA）测量分类维数（Classification Dimension）与鲁棒维数（Robust Dimension），发现纯扩散合成数据通过提供丰富支撑样本使特征空间偏向低秩泛化，而 DRA 则促使模型在非低秩的前提下更高效地激活各特征维度来编码鲁棒语义，二者产生互补的协同效应，使模型在极高稀疏度约束下也能大幅降低 SAE 重构误差，实现特征的高可解耦性。

损失函数 / 训练策略

对于 CIFAR-10 和 CIFAR-100，对抗扰动约束在 \(\ell_\infty\) 球内，扰动半径 \(\varepsilon=8/255\)，步长 \(\alpha=2/255\)，采用 10 步 PGD 构造对抗样本，主干网络采用 TRADES 准则（权重系数 \(\beta=5\)）。优化器针对 WideResNet（WRN-28-10, WRN-34-10）采用带动量的 SGD（学习率 0.2，余弦退火，权重衰减 \(5\times 10^{-4}\)，EMA 动量 \(\tau=0.995\)），针对 ViT-B 采用 Lion 优化器（学习率 \(10^{-4}\)，权重衰减 0.5）。ImageNet 实验在输入分辨率 \(224\times 224\) 下微调由 DINOv3 预训练初始化的 ViT-B/16 与 ConvNeXt-B，扰动预算 \(\varepsilon=4/255\)，采用 3 步 TRADES（\(\beta=10\)）与 Lion 优化器微调 100 个 epoch。

实验关键数据

主实验

在 CIFAR-10、CIFAR-100 以及 ImageNet 基准上，评估了在 SOTA 对抗训练配方（DM-AT）中引入 DRA 后的干净样本准确率（Clean Acc）与基于 AutoAttack（AA）的标准对抗鲁棒准确率。

数据集	模型架构	训练方法	干净准确率 (%)	AutoAttack 鲁棒率 (%)	提升幅度 (AA)
CIFAR-10 (\(\ell_\infty=8/255\))	WRN-28-10 (20M 合成)	DM-AT (Baseline)	92.44	67.31	-
CIFAR-10 (\(\ell_\infty=8/255\))	WRN-28-10 (20M 合成)	DM-AT + DRA (本文)	93.14	67.83	+0.52%
CIFAR-10 (\(\ell_\infty=8/255\))	ViT-B/2 (50M 合成)	DM-AT (Baseline)	94.35	71.31	-
CIFAR-10 (\(\ell_\infty=8/255\))	ViT-B/2 (50M 合成)	DM-AT + DRA (本文)	95.22	71.77	+0.46%
CIFAR-100 (\(\ell_\infty=8/255\))	WRN-28-10 (1M 合成)	DM-AT (Baseline)	68.34	35.72	-
CIFAR-100 (\(\ell_\infty=8/255\))	WRN-28-10 (1M 合成)	DM-AT + DRA (本文)	69.85	36.27	+0.55%
CIFAR-100 (\(\ell_\infty=8/255\))	ViT-B/2 (1M 合成)	DM-AT (Baseline)	68.53	36.52	-
CIFAR-100 (\(\ell_\infty=8/255\))	ViT-B/2 (1M 合成)	DM-AT + DRA (本文)	69.95	37.43	+0.91%
ImageNet (\(\ell_\infty=4/255\))	ConvNeXt-B (4M 合成)	DM-AT (Baseline)	74.49	54.44	-
ImageNet (\(\ell_\infty=4/255\))	ConvNeXt-B (4M 合成)	DM-AT + DRA (本文)	76.03	56.07	+1.63%
ImageNet (\(\ell_\infty=4/255\))	ViT-B/16 (4M 合成)	DM-AT (Baseline)	74.62	54.64	-
ImageNet (\(\ell_\infty=4/255\))	ViT-B/16 (4M 合成)	DM-AT + DRA (本文)	76.87	55.16	+0.52%

消融实验

下表给出了 CIFAR-10 上不同对抗训练组件与表征目标的对比消融，展示了合成数据与对齐机制的独立贡献及替代设计的性能差异。

训练配置	基础架构	合成数据量	干净准确率 (%)	AutoAttack 鲁棒率 (%)	实验说明
标准对抗训练 (AT Baseline)	WRN-34-10	无 (0)	84.33	55.25	无合成数据、无表征对齐的标准 AT
AT + DRA (仅表征对齐)	WRN-34-10	无 (0)	88.54	57.29	无合成数据下 DRA 带来 +2.04% AA 提升
直接端到端微调扩散编码器	UNet Enc	无 (0)	87.77	55.76	直接微调扩散主干，速度慢 1.35x 且表现逊色
DM-AT (1M 合成数据)	WRN-28-10	1M	91.12	63.35	仅使用 1M 扩散数据标准训练
DM-AT + 判别式带噪预训练对齐	WRN-28-10	1M	90.41	62.62	将对齐目标换为普通带噪判别特征，性能退化
DM-AT + DRA (本文方案)	WRN-28-10	1M	92.36	64.12	结合 1M 数据与扩散表征对齐，提升 0.77% AA

关键发现

• 对齐生成式表征优于判别式带噪表征：单纯用带噪图像进行传统判别式交叉熵预训练所抽取的特征作为对齐目标，不仅无法提升鲁棒性，反而使 AA 鲁棒率从 64.12% 暴跌至 62.62%，这证明扩散模型逐步去噪的生成机制本身是诱导鲁棒特征先验的决定性根源。
• 解耦内在表征维度：主成分分析（PCA）表征维数评估显示，普通 AT 的鲁棒峰值仅集中在 9 个主成分上，而 DM-AT + DRA 的有效鲁棒维数扩展至 23 维；TopK-SAE 评估表明其在各稀疏度级别下的重构误差均显著更低，证明模型减少了多特征混叠（superposition）对脆弱维度的依赖。
• 零推理代价与跨架构普适性：DRA 在 WideResNet、ConvNeXt 与 ViT 等多种完全不同的骨干网络上均实现一致且显著的性能收益，由于投影头在推理时可完全丢弃，实际部署不增加任何浮点运算与显存占用。

亮点与洞察

• 重新定义扩散模型的赋能角色：跳出了业界将扩散模型单纯看作“离线高成本制图机”的局限思维，从表征学习角度首次确立了去噪生成特征作为对抗防御强力先验的有效性。
• 机制互补的双轨驱动：深刻揭示了合成数据侧重于样本覆盖以促使低秩泛化，而特征对齐侧重于激活高维鲁棒语义流形，两者并非相互替代，而是正交互补。
• 频域敏感性的自适应抑制：巧妙利用扩散模型在中间时间步已滤除无关高频毛刺、保留主体语义轮廓的天然物理特性，使判别模型在训练中自动获得了抗高频对抗扰动的“免疫力”。

局限与展望

• 作者承认的局限：DRA 在训练阶段需要对冻结扩散模型执行前向推理以抽取特征，在不采用离线特征缓存的前提下，会增加训练阶段的显存占用和前向计算开销。
• 自身观察与评估：论文在时间步选择上主要依赖离线固定的经验经验步长（如 EDM 的 \(\sigma_t=0.1\)），未能探索动态时间步自适应调度或多时间步多尺度联合特征金字塔对齐策略。
• 未来改进方向：可进一步探索自适应特征通道重加权技术，或将 DRA 拓展至大尺度视觉语言模型（VLM）的多模态跨域对抗安全防御中。

相关工作与启发

• vs DM-AT (Wang et al., 2023b)：DM-AT 仅将扩散模型产出的合成图像作为额外数据喂给分类器，属于输入空间的简单扩充；本文在此基础上通过 DRA 在隐表征空间建立直接几何约束，在多项基准上打破了单纯数据扩展的收益边际。
• vs 扩散对抗净化 (Diffusion Purification, Nie et al., 2022)：对抗净化方法在测试期必须对输入图像实施耗时极长的迭代加噪-去噪采样，且容易受到基于 EOT 的自适应白盒攻击瓦解；本文方法在测试推理期完全不运行扩散模型，既保持了纯判别模型的毫秒级推理速度，又杜绝了随机混淆梯度的虚假安全风险。

评分

• 新颖性: ⭐⭐⭐⭐ [从数据级扩展深入至特征级对齐，为扩散模型辅助防御开辟了全新视角]
• 实验充分度: ⭐⭐⭐⭐⭐ [覆盖 CIFAR-10/100 与 ImageNet，涵括 CNN 与 ViT 架构，并进行了 SAE 与频域深度剖析]
• 写作质量: ⭐⭐⭐⭐⭐ [逻辑严密，理论铺垫到实证层层递进，洞察深入透彻]
• 价值: ⭐⭐⭐⭐⭐ [即插即用且零推理开销，为工业界与学术界构建前沿对抗鲁棒分类器提供了极具价值的新范式]

相关论文

• [NeurIPS 2026] Diffusion Models without Classifier-free Guidance
• [NeurIPS 2026] Robust Diffusion Models via Divergence-Induced Weighted Denoising
• [NeurIPS 2025] Why Diffusion Models Don't Memorize: The Role of Implicit Dynamical Regularization in Training
• [NeurIPS 2026] ParetoSlider: Diffusion Models Post-Training for Continuous Reward Control
• [NeurIPS 2026] DiffRatio: Training One-Step Diffusion Models Without Teacher Supervision