Adversarial Distillation of Bayesian Neural Network Posteriors

TL;DR

This paper introduces Adversarial Posterior Distillation, a method using GANs to efficiently compress Bayesian neural network posteriors learned via SGLD, enabling effective uncertainty estimation and application in various tasks.

Contribution

It presents a novel framework that distills SGLD samples into a GAN, preserving performance and enabling posterior and uncertainty computations for large BNNs.

Findings

No performance loss in anomaly detection, active learning, adversarial defense

Efficient posterior sampling with GANs reduces storage costs

First application of MCMC-based BNNs to these downstream tasks

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.