Dirichlet Wrapper to Quantify Classification Uncertainty in Black-Box Systems

Abstract: Nowadays, machine learning models are becoming a utility in many sectors. AI companies deliver pre-trained encapsulated models as application programming interfaces (APIs) that developers can combine with third party components, their models, and proprietary data, to create complex data products. This complexity and the lack of control and knowledge of the internals of these external components might cause unavoidable effects, such as lack of transparency, difficulty in auditability, and the emergence of uncontrolled potential risks. These issues are especially critical when practitioners use these components as black-boxes in new datasets. In order to provide actionable insights in this type of scenarios, in this work we propose the use of a wrapping deep learning model to enrich the output of a classification black-box with a measure of uncertainty. Given a black-box classifier, we propose a probabilistic neural network that works in parallel to the black-box and uses a Dirichlet layer as the fusion layer with the black-box.  This Dirichlet layer yields a distribution on top of the multinomial output parameters of the classifier and enables the estimation of aleatoric uncertainty for any data sample.  
      Based on the resulting uncertainty measure, we advocate for a rejection system that selects the more confident predictions, discarding those more uncertain, leading to an improvement in the trustability of the resulting system. We showcase the proposed technique and methodology in two practical scenarios, one for NLP and another for computer vision, where a simulated API based is applied to different domains. Results demonstrate the effectiveness of the uncertainty computed by the wrapper and its high correlation to wrong predictions and misclassifications.

Reviews and author responses:

In this paper, the authors propose to use a dirichlet prior over the multinomial distribution outputted by blackbox DL models, to quantify uncertainty in predictions. The main contribution is to learn the parameters of the prior and use it as a wrapper over the black box, to adjudicate whether to retain or reject a particular prediction made by the model. The dirichlet parameters are learnt in conjunction with the model parameters as a fine tuning step in transfer learning tasks. Experiments on NLP and CV domains and multiple datasets demonstrate the efficacy of the method. 

The paper is well written, and easy to understand. The main motivation of the paper seems to be to learn what samples to drop, but the authors do not address what can be done about the dropped samples (i.e what happens if we end up having to drop 80% of the samples?) . Adding the dirichlet prior to quantify uncertainty has also been studied in the context of VAEs before, (and LDA back in the day) so conceptually there's limited novelty. Nonetheless, the method seems to provide impressive results on multiple datasets, and I think this is interesting enough to warrant an accept. 

MINOR:
1. Figure 1: the values of \beta is the same in all subfigures
2. Sec 2.2 line 1: SImilarly --> Similar



This paper presents a method for learning a “wrapper” model which endows a multiclass predictor with an estimate of model uncertainty. The base model is treated as a black box which emits a categorical distribution, while the wrapper model estimates the parameters of a Dirichlet distribution. The wrapper is trained against silver labels from the base model, and the sampled predictive entropy is used to threshold predictions so as to withhold a decision on uncertain examples.

I believe this paper should not be accepted, as the main contribution is not particularly novel and some experimental details (Section 3) are not well-motivated. 

In particular, the idea of learning a Dirichlet prior is very similar to that proposed in (Malinin and Gales, NeurIPS 2018). The main contribution in this work seems to be the application to an existing black-box model, but this seems like a straightforward application of knowledge distillation, another well-established technique (e.g. Hinton et al. 2015, though oddly, this paper doesn’t mention this).

The details of the modeling set-up (section 2.2 and 3) are also not entirely clear. It seems from the preceding discussion that the main goal of the wrapper model is to estimate \beta, but it is not clear how the sampling procedure (2.2) allows for this. It seems that the value of the “sampled predictive entropy” (no equation number, but see end of Section 3) is just that of the mean predicted distribution, and it’s not clear why this should be different (assuming the wrapper model converges) than the predictive entropy of the base model.

This paper could be improved by a more thorough comparison to the literature, and by a clearer motivation for the training procedure used. Additionally, it would help to present a summary table of results, and to frame the metrics in terms of standard classification metrics such as precision and recall where applicable.


This manuscript proposes to train a wrapper to assess the confidence of a black box classifier decision on new samples. The resulting uncertainty prediction is used to reject decision with low confidence, such that the accuracy of the prediction on retained samples remains high. The idea, although a bit incremental, is potentially useful to practitioners, as argued in the introduction, and some empirical result tend to suggest that the method can be useful. However, I am not convinced the approach and its implementation are appropriate. 
Main comments:
Writing: the model seems strongly based on references such as Kendall & Gal (2017), however lack of introduction of notations and explanations prevent the paper to be self contained. For example: 
(1)	Here is a list of quantities for which I could not find a definition: y_m, \omega, W. 
(2)	\hat{y} appears with several indexing styles (up to three indices), sometime bold and sometimes not: the meaning of the indexing (which I could not find) is difficult to infer from the text
Objective:
I was unable to understand the rational behind the objective to be minimized, introduced in page 4. This is introduced as a cross-entropy loss, however, taking the expression of the Dirichlet distribution, it is not obvious to me how to reach this very simple expression. Is it possible that the authors simply mimic the expression of Kendall and Gal (2017, eq. (5)), that was designed for the Gaussian case, for which the cross-entropy expression is correct?
As a consequence, I do not see how this objective is supposed to learn properly the correct beta parameter.

Results:
Figures 4-8 shows convincing evidence that the procedure improves with respect to a baseline that consist (as far as I understood) in ranking decision based on the entropy of the output probability vector of the classifier. However, given that I am unsure about what the proposed optimization does, it remains unclear to me whether these results reflect a true achievement. For example, one can argue that the chosen baseline is unlikely to be a good estimate of the entropy of the decision due to the fluctuations of the output probability for the unlikely classes. Those low probability values are the classical source of variance and bias in entropy estimates, and a classifier is not designed to get these low probabilities right (as they are low anyways). As a consequence, an already better baseline might be achieved for a given reference class by cropping the probability vector to keep only classes that non-negligible probability over the training set (here they can be many alternative approaches to test). 
A trivial explanation for the proposed approached to work better than the currently chosen baseline is that the noise introduced by the sampling from the Dirichlet distribution leads to larger probabilities for the cases where the probability given by the classifier is small, which would reduce the variance of the entropy estimator based on the formula of page 5 (top). Overall, extensively checking many simpler baselines (which do not require training!) is a first step to see if the achieved result is not easy to get.   
Minor comments:
(1)	Please check for typos. 
(2)	Please avoid remove the multiple parenthesis for successive citations (use a single “\citep”).
(3)	In Fig. 1, subtitles are inconsistent.


Motivated by real-world challenges in applying pre-trained models, the authors propose a model for selective prediction (prediction with an option for abstention) that wraps an existing black-box classification model. The resulting model output is a Dirichlet distribution with mean equal to the categorical distribution produced by the black-box and concentration parameter specified by a separate auxiliary model. This additional model is trained to minimize negative log-likelihood of observations under categorical distributions sampled from the aforementioned Dirichlet along with an L1 regularization term on the concentration parameter. To infer the model’s level of uncertainty, the authors propose computing the entropy of the average of sampled categorical distributions.

The authors evaluated this model on several pairs of sentiment-analysis NLP tasks and one pair of image datasets where a base model is trained on the source dataset, and the auxiliary model is trained on the second, target dataset. Using metrics proposed in (Condessa et al. 2015), the results show positive results at nearly all thresholds compared to a simple entropy baseline.

The paper addresses an important practical challenge in machine learning, but a confusing problem-framing and lack of robust baselines make me skeptical that it is suitable for publication at ICLR 2020.

A primary concern with this work is its framing of the problem as one of measuring aleatoric (irreducible) uncertainty, but the motivation in transfer learning and interdependence in production ML systems requires models that can characterize epistemic (reducible) uncertainty. A black box model that yields distributions over classes expresses aleatoric uncertainty via that distribution, and uncertainty due to a shifted data distribution is epistemic as additional data from the new domain would reduce it.

More problematic is the study’s lack of robust baselines. The authors only present the predictive entropy baseline, but numerous methods exist for out-of-distribution detection and selective classification. Though the specific case of selective classification from a blackbox base model is perhaps more niche, other methods from related problems can either be adapted accordingly or used as upper / lower bounds on what we can expect for this problem. 

Some simple baselines I would expect to see include:
* Training a new classification model entirely on the new domain.
* Training an auxiliary classifier to predict if the base model will be correct (similar to SelectiveNet but without a shared network body). Ideally this model should also have access to the base-model’s prediction as input.
* “Confidence score” (i.e. the probability assigned to the base-model’s predicted class) -- this is a common baseline for OOD detection.

Additional questions / concerns:
* Do I understand correctly that when drawing samples for the entropy calculation, $E[\hat{y}]$ will equal the black-box model’s prediction in the limit of sample size? If so, this looks like an inefficient and complicated proxy for measuring the concentration of categorical entropies produced by the Dirichlet.
* Since the paper is focused on scenarios where one is taking advantage of a pre-trained model, one might wonder if the wrapping scheme is indeed less expensive than a new model trained from scratch on the new domain (e.g. with comparable capacity to the proposed wrapping model). Authors should include experiments / baselines to assess this.
* In section 2, the paper asserts that assuming access to logits breaks the blackbox assumption, but these are computable from softmax values (up to constant factors).
* Is the beta regularization term is theoretically required to prevent unbounded growth or is it simply an empirically practical necessity?
* In the image transfer experiments, STL-10 has relatively very few labeled examples (hundreds per class), so why was this only used as the source domain? Realistic transfer learning generally entails an expensive model trained on a source domain with plentiful data transferred to a target domain with scarce data.
* Authors should mention that STL-10 images represent a distribution shift from CIFAR-10 as the two were not generated identically. The paper currently only highlights differences in dataset size.


Meta-review: This paper proposes adding a Dirichlet distribution as a wrapper on top of a black box classifier in order to better capture uncertainty in the predictions.  This paper received four reviews in total with scores (1,1,1,6).  The reviewer who gave the weak accept found the paper well written, easy to follow and intuitive.  The other reviewers, however, were primarily concerned about the empirical evaluation of the method.  They found the baselines too weak and weren't convinced that the method would work well in practice.  The reviewers also cited a lack of comparison to existing literature for their scores.  One reviewer noted that while the method addresses aleatoric uncertainty, it doesn't provide any mechanism for epistemic uncertainty, which would be necessary for the applications motivating the work.  

The authors did not provide a response and thus there was no discussion. 

Decision: reject