[论文解读] Adversarial Distillation of Bayesian Neural Network Posteriors

Kuan-Chieh Wang, Paul Vicol+4 more|arXiv (Cornell University)|Jun 27, 2018

Adversarial Robustness in Machine Learning参考文献 36被引用 25

一句话总结

该论文提出对抗后验蒸馏（Adversarial Posterior Distillation, APD），一种利用生成对抗网络（Generative Adversarial Network, GAN）从贝叶斯神经网络（Bayesian Neural Networks, BNNs）的随机梯度朗之万动力学（Stochastic Gradient Langevin Dynamics, SGLD）中蒸馏后验样本的方法。GAN生成器学习生成高质量的后验样本，实现在测试阶段高效推理的同时存储开销极低，同时在异常检测、主动学习和对抗防御等对不确定性敏感的任务中保持优异性能。

ABSTRACT

Bayesian neural networks (BNNs) allow us to reason about uncertainty in a principled way. Stochastic Gradient Langevin Dynamics (SGLD) enables efficient BNN learning by drawing samples from the BNN posterior using mini-batches. However, SGLD and its extensions require storage of many copies of the model parameters, a potentially prohibitive cost, especially for large neural networks. We propose a framework, Adversarial Posterior Distillation, to distill the SGLD samples using a Generative Adversarial Network (GAN). At test-time, samples are generated by the GAN. We show that this distillation framework incurs no loss in performance on recent BNN applications including anomaly detection, active learning, and defense against adversarial attacks. By construction, our framework not only distills the Bayesian predictive distribution, but the posterior itself. This allows one to compute quantities such as the approximate model variance, which is useful in downstream tasks. To our knowledge, these are the first results applying MCMC-based BNNs to the aforementioned downstream applications.

研究动机与目标

• 解决在贝叶斯神经网络（BNNs）中维护多个SGLD样本带来的高存储成本问题，该问题限制了大规模模型的可扩展性。
• 在测试阶段实现高效、参数化的后验近似，同时不牺牲不确定性估计的质量。
• 证明基于MCMC的BNN（特别是SGLD）在不确定性敏感应用中可优于MC正则化等简单方法。
• 表明基于GAN的蒸馏能够保留完整的后验结构（包括模型方差），这对下游任务至关重要。
• 建立一个实用框架，使基于MCMC的BNN能够在不确定性量化至关重要的现实应用中实际部署。

提出的方法

• 使用随机梯度朗之万动力学（SGLD）在模型参数上生成一组后验样本，以表示真实后验分布。
• 训练一个生成对抗网络（GAN），其中生成器学习生成与SGLD后验样本分布匹配的样本。
• 判别器负责区分真实SGLD样本与生成样本，而生成器则通过优化以欺骗判别器。
• 采用WGAN-GP结合梯度惩罚以稳定训练并提升样本质量，确保更优的后验近似。
• 在测试阶段，从训练好的GAN生成器生成样本，而非存储SGLD样本，从而大幅降低内存使用。
• 利用蒸馏后的GAN样本计算不确定性度量（如熵和BALD），支持异常检测和主动学习等下游任务。

实验结果

研究问题

• RQ1 GAN能否有效蒸馏通过SGLD获得的贝叶斯神经网络后验分布，同时保留其不确定性特征？
• RQ2 通过GAN蒸馏得到的后验分布是否在异常检测和对抗防御等不确定性敏感任务上达到与SGLD样本相当的性能？
• RQ3 与更简单的近似方法（如高斯混合模型，MoG）相比，基于GAN的蒸馏在准确性和存储效率方面表现如何？
• RQ4 基于GAN的后验蒸馏是否能在显著降低存储成本的同时维持高质量的不确定性估计，相比直接存储SGLD样本？
• RQ5 何种训练形式（如原始GAN、WGAN、WGAN-GP）能为BNN的后验蒸馏提供最稳定和高效的性能？

主要发现

• 在notMNIST OOD异常检测任务中，APD保留了SGLD样本99.8%的性能，优于60分量高斯混合模型（MoG，99.3%），且参数量显著更少（1.67M vs. 9.54M）。
• APD性能随生成样本数量增加而提升，仅用20个生成样本即可达到50个SGLD样本的性能，实现2.5倍的存储节省。
• WGAN-GP结合梯度惩罚的训练收敛更快，且训练过程振荡更少，相比原始GAN或使用权重裁剪的WGAN，展现出更稳定的后验蒸馏效果。
• 单分量高斯混合模型（MoG）在异常检测中表现不佳，表明SGLD后验分布具有多模态特性，无法被简单的因子化近似捕捉。
• APD在测试阶段实现了对完整后验分布的访问，支持模型方差及其他不确定性度量的计算，这对主动学习和对抗鲁棒性至关重要。
• 该框架表明，尽管以往因存储成本而被回避，基于MCMC的BNN可通过对抗蒸馏实现实际部署，且在关键应用中优于MC正则化等简单方法。