We propose a method for quantifying uncertainty in neural network regression models when the targets are real values on a d-dimensional simplex, such as probabilities. We show that each target can be modeled as a sample from a Dirichlet distribution, where the parameters of the Dirichlet are provided by the output of a neural network, and that the combined model can be trained using the gradient of the conditional likelihood. This approach provides interpretable predictions in the form of multidimensional distributions, rather than point estimates, from which one can obtain confidence intervals or quantify risk in decision making. Furthermore, we show that the same approach can be used to model targets in the form of empirical counts as samples from the Dirichlet-multinomial compound distribution. In experiments, we verify that our approach provides these benefits without harming the performance of the point estimate predictions on two diverse applications: (1) distilling deep convolutional networks trained on CIFAR-100, and (2) predicting the location of particle collisions in the XENON1T Dark Matter detector.
This paper considers parameterizing Dirichlet, Dirichlet-multinomial, and Beta distributions with the outputs of a neural network. They present the distributions and gradients, discuss appropriate activation functions for the output layer, and evaluate this approach on synthetic and real datasets with mixed results. Overall, I found the writing very clear, the main idea sound, and paper generally well executed, but I have serious concerns about the significance of the contributions that lead me to recommend rejection. It would be very useful to me if the authors would provide a concise list of what they consider the main contributions to be and why they are significant. As I see it, the paper does three main things:
1. In section 2, the authors consider parameterizing Dirichlet, Dirichlet-multinomial, and Beta distributions with the outputs of a neural network (Section 2). As the authors note, parameterizing an exponential family distribution with the outputs of a neural network is not a novel contribution (e.g. Rudolph et al. (2016) and David Belanger's PhD thesis (2017)) and though I have never personally seen the Dirichlet, Dirichlet-multinomial, and Beta distributions used, the conceptual leap required is small. Most of section 2 is dedicated to writing down, simplifying, and deriving gradient equations for these three distributions. The simplifications and gradient derivations are well known and appear in many places (e.g. http://jonathan-huang.org/research/dirichlet/dirichlet.pdf, https://arxiv.org/pdf/1405.0099.pdf) and should not be considered contributions in the age of automatic differentiation (see Justin Domke's blog post on autodiff).
2. In section 3, the authors consider the unique challenges of using the proposed networks. They propose targeted activation functions that will improve the stability of learning. I found this to be the most interesting portion of the paper and the most significant contribution. Unfortunately, it is short on details and empirical results are referenced that do not appear in the paper (i.e. the second to last paragraph on page 5). If I were to rewrite this paper, I would focus on answering the question "What are the unique challenges of parameterizing Dirichlet, Dirichlet-multinomial, and Beta distributions with the outputs of a neural network and how can we address them?", replacing section 2 with an expanded section 3.
3. In section 4, the authors evaluate the proposed networks on a collection of synthetic and real tasks. In the end, the results are mixed, with the Dirichlet network performing best on the XENON1T task and the standard softmax network performing best on the CIFAR-100 task. In general, I don't mind mixed results and I appreciate that the authors included both sets of experiments; however, it is important that there is a convincing argument for why one would prefer the proposed solution even when accuracy is the same (e.g. it is faster, it is interpretable, etc.). The authors briefly argue that the proposed methods are superior because they provide uncertainty estimates for the output distributions. This may be true, but they only perform evaluations on tasks where the primary goal is accuracy. If the main benefit of the proposed networks is proper uncertainty quantification, then the evaluations (even if they are qualitative) should reflect that.
In summary, I do not think the models proposed in section 2 are sufficiently novel to justify publication alone which means that the authors need to either: (1) evaluate novel methods that are critical for use of these models or (2) present a convincing evaluation that strongly motivates the proposed model's use or that provides some novel insight into the model's behavior. I think that the authors are on their way to achieving (1), but do not achieve (2). I would suggest finding an application that requires uncertainty estimates for the distribution and centering the paper around that application.
Minor comments:
- Figure 2 (right) should include a y-axis label (e.g. "parameter value").
- In Figure 3 (right), it is not obvious what the "Sigmoid" line corresponds to. 
- It is not clear what the authors are trying to show in section 4.1. The EL activation function is smooth and monotone and the likelihood is convex, so there should be no question that the distribution will concentrate around y.
- Section 4.4 was interesting, but would have been more convincing if paired with an evaluation on real data.