[Paper Review] Reverse KL-Divergence Training of Prior Networks: Improved Uncertainty and Adversarial Robustness

Andrey Malinin, Mark Gales|arXiv (Cornell University)|May 31, 2019

Adversarial Robustness in Machine LearningComputer Science37 references69 citations

TL;DR

The paper proposes training Prior Networks with reverse KL-divergence to better model uncertainty and improve out-of-distribution detection, while enabling a generalized adversarial training approach that increases robustness to adaptive whitebox attacks.

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.

Motivation & Objective

• Motivate uncertainty estimation in neural networks and the need for reliable confidence under misclassification, OOD inputs, and adversarial threats.
• Introduce Prior Networks as a Dirichlet-prior model over output distributions to emulate ensembles efficiently.
• Show that reverse KL-divergence is the appropriate training criterion for Prior Networks, enabling scalability to many classes and better OOD detection.
• Explore a generalized adversarial training framework using the reverse KL criterion to make adaptive attacks harder to craft against Prior Networks.

Proposed method

• Define Prior Networks that parameterize a Dirichlet distribution over output distributions.
• Contrast forward KL with reverse KL training criteria for Dirichlet targets, arguing RKL yields a single high-precision mode in high data uncertainty regions.
• Derive and compare loss formulations: forward KL (original PN) vs. reverse KL (proposed PN-RKL) with explicit mixtures in expectation (geometric vs. arithmetic mixtures).
• Train PN-RKL on image datasets without auxiliary losses and evaluate on in-domain accuracy and out-of-distribution detection.
• Extend the framework to adversarial attack detection by formulating a generalized adversarial training loss that uses RKL to shape uncertainty for adversarial inputs.
• Evaluate robustness to adaptive whitebox attacks using targeted PGD-MIM and compare against DNN, adversarially trained DNN, and MC-dropout baselines.

Experimental results

Research questions

• RQ1 Does reverse KL-divergence provide the correct training signal for Prior Networks across datasets with varying numbers of classes?
• RQ2 Can PN-RKL maintain competitive classification performance while improving out-of-distribution detection compared to PN-KL?
• RQ3 Does PN-RKL enable effective detection of adversarial attacks and increase robustness to adaptive whitebox attacks?
• RQ4 How does the proposed adversarial training formulation influence the space of successful attacks against Prior Networks?
• RQ5 What are the comparative limitations of PN-RKL with different OOD training data selections on complex datasets?

Key findings

• PN-RKL yields uncertainty measures that align with dataset structure: high data uncertainty in overlap regions and high knowledge uncertainty for out-of-distribution inputs.
• On synthetic high-uncertainty data, PN-RKL produces a more accurate decomposition of total vs. data vs. knowledge uncertainty than PN-KL.
• PN-RKL achieves classification error rates comparable to standard DNNs and to ensembles, while PN-KL degrades with more complex datasets.
• PN-RKL improves OOD detection AUROC over PN-KL and competes with or surpasses ensembles on CIFAR-10/ CIFAR-100 when using appropriate OOD data.
• Adversarial training under PN-RKL (beta_in and beta_adv settings) makes adaptive whitebox attacks significantly more computationally demanding to succeed than against standard DNNs, DNN-ADV, or MC-dropout defenses.
• The approach yields reduced transferability of adaptive attacks, with black-box attacks often failing against PN-RKL.