PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Andrey Malinin, Mark Gales

arXiv:1905.13472v2stat.MLcs.LG

TL;DR

Prior Networks need a scalable training criterion for uncertainty estimation on complex, many-class datasets. The paper replaces forward KL with reverse KL for Dirichlet priors and uses the resulting control to generalize adversarial training. The proposed models improve out-of-distribution detection and make adaptive whitebox attacks more computationally challenging on CIFAR-10 and CIFAR-100.

  • Problem

    Prior Networks are difficult to train on complex datasets with many classes using the originally proposed forward KL criterion, which also requires auxiliary losses.

  • Method

    The paper trains Prior Networks with reverse KL-divergence between Dirichlet distributions and uses it to train models that remain correct but more uncertain on adversarial inputs.

  • Results

    The proposed criterion improves out-of-distribution detection, while adaptive attacks are more computationally challenging against the resulting Prior Networks than against comparison defenses.

  • Takeaways & Limitations

    Reverse KL training enables Prior Networks to handle arbitrary class counts and supports adversarial-attack detection through uncertainty over output distributions.

  • Takeaways & Limitations

    Forward KL remains sensible when parameters or samples from the correct target distribution over output distributions are available for every input.

Abstract

from arXiv · show

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.

1 Introduction

The introduction motivates Prior Networks as efficient uncertainty-estimation models and proposes reverse KL training plus generalized adversarial training to improve scalability, detection, and attack resistance.

  • Prior Networks efficiently emulate ensembles by parameterizing distributions over output distributions, supporting uncertainty-based misclassification, out-of-distribution, and adversarial-attack detection.
  • Forward KL training requires auxiliary losses, struggles with many-class datasets, and motivates replacing it with reverse KL-divergence training.
  • Reverse KL training is presented as yielding desired Prior Network behaviors without auxiliary losses while preserving classification performance on complex image-classification tasks.
  • The paper proposes generalized adversarial training so Prior Networks predict the correct class on adversarial inputs while assigning them greater uncertainty than natural inputs.
  • Adaptive whitebox attacks that both alter predictions and evade detection are reported as more computationally challenging against the proposed Prior Networks than against several defended baselines.

2 Prior Networks

Prior Networks parameterize Dirichlet priors over categorical output distributions, enabling closed-form ensemble-like uncertainty measures. Their predictive and uncertainty behaviors distinguish confidence, data uncertainty, and knowledge uncertainty.

  • A Prior Network explicitly parameterizes a prior distribution over output distributions, emulating an ensemble without sampling.
  • The model typically uses a Dirichlet distribution because it is conjugate to the categorical distribution and has tractable analytic properties.
  • Dirichlet concentration parameters α determine precision α0, with larger α0 producing sharper and more confident distributions.
  • The predictive distribution is the expected categorical distribution under the conditional Dirichlet prior, with class probabilities given by normalized concentration parameters.
  • The desired behaviors distinguish confident predictions, data uncertainty from severe class overlap, and knowledge uncertainty for out-of-distribution inputs.
  • Mutual information between y and π provides a measure of prediction uncertainty, while predictive entropy represents total uncertainty that can be decomposed into data and knowledge uncertainty.

3 Forward and Reverse KL-Divergence Losses

The section explains why forward KL-divergence induces unsuitable target mixtures for Prior Networks and motivates reverse KL-divergence as a better criterion. Reverse KL produces a unimodal Dirichlet target whose concentration parameters reflect the desired data uncertainty.

  • Forward KL-divergence: Training Prior Networks from class labels requires inducing a target distribution over output distributions from empirical samples.The original setup uses class-dependent target concentration parameters β(c), with out-of-domain targets assigned βout = 0 for a flat Dirichlet distribution.
  • Forward KL-divergence: Forward KL-divergence induces an arithmetic mixture with modes at simplex corners, which can force low precision in regions of substantial data uncertainty.This spreads the model across multiple modes and can compromise predictive performance, whereas the desired target is a single high-precision mode at the simplex center.
  • Reverse KL-divergence: The proposed reverse KL-divergence avoids the arithmetic mixture by inducing a geometric mixture of target Dirichlet distributions.Its concentration parameters are arithmetic mixtures of the target concentration parameters for each class.
  • Reverse KL-divergence: The resulting target is always a standard unimodal Dirichlet with a mode reflecting the correct level of data uncertainty.The paper further notes that forward KL remains sensible when the correct target distribution or samples from it are directly available for every input.

4 Experiments on Synthetic Data

Synthetic-data experiments compare Prior Networks trained with forward and reverse KL-divergence using uncertainty measures derived from the model distributions. Reverse KL yields a more appropriate decomposition of data and knowledge uncertainty than forward KL.

  • Experimental setup: The experiments use a synthetic three-class dataset with high uncertainty and out-of-distribution training data sampled as a thin shell around the training data.The study evaluates both KL criteria using manually constructed target concentration parameters.
  • Uncertainty comparison: Forward-KL training overestimates total uncertainty in-domain, making it similarly high near decision boundaries, in class-overlap regions, and out-of-domain.Its mutual information is also higher in-domain along decision boundaries than out-of-domain.
  • Uncertainty comparison: Reverse-KL training decomposes uncertainty correctly: data uncertainty peaks in class-overlap regions, while mutual information is low in-domain and high out-of-domain.These results support the preceding analysis of the two training criteria.

5 Image Classification Experiments

Image-classification experiments show that reverse-KL Prior Networks retain classification performance as dataset complexity and class count increase, while improving out-of-distribution detection relative to forward-KL training. On CIFAR-10, reverse KL outperforms forward KL and ensembles across all reported OOD test datasets.

  • Classification performance: Forward-KL Prior Networks exhibit increasingly worse classification error as datasets become more complex and contain more classes.Reverse-KL Prior Networks have error rates similar to standard DNNs and ensembles, without auxiliary losses.
  • Out-of-distribution detection: Reverse-KL Prior Networks achieve better CIFAR-10 OOD detection than forward-KL Prior Networks and ensembles on SVHN, LSUN, and TinyImageNet.Performance is assessed using AUROC from mutual information across five random initializations.
  • Out-of-distribution detection: On CIFAR-100, Prior Networks outperform ensembles against LSUN and SVHN but have difficulty distinguishing CIFAR-10 from CIFAR-100 test sets.The paper attributes this limitation to the classification model and OOD training data rather than the training criterion.

6 Adversarial Attack Detection

The paper generalizes adversarial training for Prior Networks by training them to classify adversarial inputs correctly while assigning them higher uncertainty. Against adaptive whitebox attacks, this approach makes successful attacks more computationally difficult.

  • Adversarial training: Adversarially perturbed inputs serve as out-of-distribution training data, with Prior Networks trained for correct predictions and high uncertainty.Natural inputs receive sharp Dirichlet targets, whereas adversarial inputs receive wider targets.
  • Adversarial training: The proposed training generalizes adversarial training by shaping the entire distribution over output distributions, rather than only decision boundaries.The target concentration parameters can differ between natural and adversarial data, such as βin = 1e2 and βadv = 1.
  • Attack evaluation: Adaptive attacks target both prediction changes and unchanged uncertainty, using targeted iterative PGD-MIM for evaluation and targeted FGSM during training.The attack target is the second most likely class to encourage a less unnatural output perturbation.
  • Attack objectives: DIR attacks impose more explicit constraints than LKL attacks and were found harder to optimize, producing less aggressive attacks.Constraints on the full Dirichlet distribution include both relative and absolute logit values.
  • Results: At 1000 iterations on CIFAR-10, Prior Networks had JSR 0.25, while the other models reached JSR 0.5; CIFAR-100 showed a more modest but significant robustness increase.DNN and DNN-ADV attacks changed predictions and evaded detection, while Prior Networks required substantially more gradient-descent iterations to do so.

7 Conclusion

The paper concludes that reverse KL-divergence enables Prior Networks to scale to complex classification problems and improves OOD detection. It also enables adversarially trained Prior Networks that make adaptive whitebox attacks more computationally expensive, while OOD-training-data choice remains important.

  • Conclusion: Reverse KL-divergence yields the desired Prior Network behaviours and supports training on complex datasets with arbitrary numbers of classes.The conclusion reports both theoretical and empirical support for this result.
  • Conclusion: Reverse KL-divergence improves OOD detection on CIFAR-10 and CIFAR-100 relative to forward KL-divergence training.The paper cautions that inappropriate OOD training data can limit performance on complex datasets.
  • Conclusion: Adversarially trained Prior Networks detect whitebox adaptive attacks and make successful attack construction more computationally expensive than standard adversarial training or Monte-Carlo dropout.The conclusion frames this method as both a generalization of and a drop-in replacement for standard adversarial training.

Appendix A Further Analysis of reverse KL-divergence Loss

The reverse KL loss decomposes into a reverse cross-entropy and negative differential entropy, yielding a classification-like objective with explicit control over Dirichlet precision. Under the chosen target parameters, it drives the predicted precision toward β + K.

  • Loss decomposition: The reverse KL-divergence loss decomposes into reverse cross-entropy and negative differential entropy.This decomposition is used to analyze the loss's effect on the predicted Dirichlet distribution.
  • Reverse cross-entropy: The reverse cross-entropy becomes a scaled upper bound on discrete cross-entropy when target concentration parameters follow equation 7.The bound is related to a parallel Dirichlet Prior Network formulation.
  • Reverse cross-entropy: The upper-bound loss matches negative log-likelihood except for digamma functions replacing natural logarithms.An asymptotic digamma approximation exposes the loss's additional precision-control term.
  • Precision control: The additional term drives each class concentration parameter ˆαc as large as possible.This follows from the asymptotic analysis of the digamma-function expression.
  • Precision control: The complete loss is negative log-likelihood weighted by β plus a term driving Dirichlet precision ˆα0 toward β + K.Here K denotes the number of classes.

Appendix B Synthetic Experiments

The synthetic dataset contains three Gaussian class clusters with nearby out-of-distribution samples, enabling comparison of uncertainty behavior under forward and reverse KL training.

  • The artificial dataset samples three normally distributed clusters with tied isotropic covariances and equidistant means, each representing a separate class.
  • Class-conditional probabilities are obtained from the cluster distributions and class priors using Bayes’ rule.
  • The dataset uses green out-of-distribution training points sampled near the in-domain region, with experiments conducted at σ = 4.
  • Both losses produce low differential entropy in-domain and high differential entropy out-of-distribution.
  • Reverse KL-divergence captures more dataset structure, suggesting its differential entropy measures total uncertainty rather than only knowledge uncertainty.

Appendix C Experimental Setup

The appendix describes datasets, model implementations, adversarial attack formulations, and training procedures for DNN and Prior Network comparisons.

  • The experiments use VGG-16 models with 2048-unit fully connected layers, training DNNs with negative log-likelihood and Prior Networks with forward or reverse KL losses.
  • Adversarial examples are generated by constrained optimization balancing target-loss minimization against a perceptual perturbation distance.
  • FGSM, FGM, BIM, MIM, PGD, C&W, and EAD provide attack variants across L1, L2, and L∞ perturbation settings.
  • C.2 Adversarial Training of DNNs and Prior Networks: DNN adversarial training jointly maximizes likelihood on natural data and dynamically generated targeted L∞ FGSM examples with randomly sampled perturbation strengths.
  • C.2 Adversarial Training of DNNs and Prior Networks: Prior Network adversarial training uses targeted FGSM examples and reverse KL-divergence to targeted Dirichlet distributions, jointly training on natural and perturbed minibatches.
  • C.2 Adversarial Training of DNNs and Prior Networks: Setting βadv = 1 yields a wide Dirichlet whose mode and mean favor the target class while maintaining high predictive uncertainty.
  • FGSM is used during training because it is computationally cheap, whereas iterative attacks would make training much slower.

Appendix D Jointly Assessing Adversarial Attack Robustness and Detection

The appendix defines adversarial success jointly with detection evasion, extending standard attack and uncertainty-detection metrics into a single operating-point analysis.

  • Adversarial attacks are evaluated in deployment-like settings where success requires both changing predictions and evading detection.
  • Standard attack success is summarized as S and plotted against maximum perturbation ϵ measured with an L1, L2, or L∞ distance.
  • Threshold-based detection labels samples using an uncertainty measure H(x), such as entropy or mutual information, relative to threshold T.
  • Detection performance is evaluated across thresholds with true-positive and false-positive rates and summarized using ROC analysis.
  • The Joint Success Rate counts attacks that are successfully generated and evade detection, with the Equal Error-Rate point providing one operating point.

Appendix E Additional Adversarial Attack Detection Experiments

Additional experiments compare adaptive attacks targeting categorical versus Dirichlet outputs and compare L2 with L∞ momentum-iterative attacks on Prior Networks.

  • KL PMF attacks are more successful than KL DIR attacks at switching predictions and evading detection on CIFAR-10 Prior Networks.
  • The authors suggest that limited common support between sharp Dirichlet distributions at different simplex corners makes KL DIR optimization more difficult.
  • L2 PGD Momentum Iterative attacks minimizing KL PMF are marginally more successful than their L∞ versions.
  • Adapting C&W L2 attacks to these loss functions remains necessary for evaluating a more aggressive set of L2 attacks.
Loading 1905.13472v2…