Reverse KL-Divergence Training of Prior Networks: Improved Uncertainty and Adversarial Robustness

TL;DR

This paper introduces a new training criterion for Prior Networks based on reverse KL-divergence, enhancing their ability to detect out-of-distribution data and adversarial attacks, especially on complex datasets.

Contribution

It demonstrates that reverse KL-divergence is the optimal training criterion for Prior Networks, enabling scalable training and improved detection of out-of-distribution and adversarial inputs.

Findings

Reverse KL-divergence improves training on complex datasets.

Prior Networks outperform other ensemble methods in out-of-distribution detection.

Adversarial attacks against these networks require more computational effort.

Abstract

Ensemble approaches for uncertainty estimation have recently been applied to the tasks of misclassification detection, out-of-distribution input detection and adversarial attack detection. Prior Networks have been proposed as an approach to efficiently \emph{emulate} an ensemble of models for classification by parameterising a Dirichlet prior distribution over output distributions. These models have been shown to outperform alternative ensemble approaches, such as Monte-Carlo Dropout, on the task of out-of-distribution input detection. However, scaling Prior Networks to complex datasets with many classes is difficult using the training criteria originally proposed. This paper makes two contributions. First, we show that the appropriate training criterion for Prior Networks is the \emph{reverse} KL-divergence between Dirichlet distributions. This addresses issues in the nature of the training data target distributions, enabling prior networks to be successfully trained on classification tasks with arbitrarily many classes, as well as improving out-of-distribution detection performance. Second, taking advantage of this new training criterion, this paper investigates using Prior Networks to detect adversarial attacks and proposes a generalized form of adversarial training. It is shown that the construction of successful \emph{adaptive} whitebox attacks, which affect the prediction and evade detection, against Prior Networks trained on CIFAR-10 and CIFAR-100 using the proposed approach requires a greater amount of computational effort than against networks defended using standard adversarial training or MC-dropout.