FEED-FORWARD PROPAGATION IN PROBABILISTIC NEURAL NETWORKS WITH CATEGORICAL AND MAX LAYERS
Probabilistic Neural Networks deal with various sources of stochasticity: input noise, dropout, stochastic neurons, parameter uncertainties modeled as random variables, etc. In this paper we revisit a feed-forward propagation approach that allows one to estimate for each neuron its mean and variance w.r.t. all mentioned sources of stochasticity. In contrast, standard NNs propagate only point estimates, discarding the uncertainty. Methods propagating also the variance have been proposed by several authors in different context. The view presented here attempts to clarify the assumptions and derivation behind such methods, relate them to classical NNs and broaden their scope of applicability. The main technical contributions are new approximations for the distributions of argmax and max-related transforms, which allow for fully analytic uncertainty propagation in networks with softmax and max-pooling layers as well as leaky ReLU activations. We evaluate the accuracy of the approximation and suggest a simple calibration. Applying the method to networks with dropout allows for faster training and gives improved test likelihoods without the need of sampling.
* Summary
The authors focus on the problem of uncertainty propagation DNN. The authors claim two main contributions: they revisit the assumptions of the feed forward method (proposed by several authors as an inference method for BNNs based on ADF/EP) and proposed a new approximation for argmax/max based functions that allows to propagated the first two moments analytically. 
* Comments:
The authors claim two main contributions: an analysis for the feed forward method (sections 2 and 3) previously proposed by several authors as an inference method for BNN based on ADF/EP, and a new method to propagate the uncertainty through argmax/max based operations (section 4).
Regarding the first contribution, I was expecting some new insights about the method that I did not find. I would suggest to focus on the second contribution and refactor this section as a background section. I would make it shorter, focusing on the representation of probabilities as latent variables trough a function, which is the important bit to understand the real contribution of the paper described in section 4. I would also remove some examples that do not seem critical to understand the rest of the paper and just increase its length.
The second contribution is quite novel. The authors propose a new approximation of argmax/max operations. The firstly proposed an approximation for argmax operations, e.g. latent variable view of the softmax, that avoids resorting to the normal cdf function that has numerical stability issues. Secondly, they suggest an approximation for max based operations, e.g. leaky relu, that again, does not depend on the gaussian cdf. 
In the experimental section, the authors test:
a)	The accuracy of the proposed method approximating the posterior of the neurons
b)	End-to-end training benefits
In a) they use MC to collect the ground truth statistics and compare the proposed method (AP2) with a classical NN (AP1). The analysis is nice but I miss a comparison with other state-of-the-art methods. In particular, the authors claim that the novelty of their method compared to other feed-forward methods is that they can propagate the uncertainty through argmax/max operations analytically. They do not compare with these other feed forwards methods to show the benefit of this.  This is shown in the end-to-end training experiments; however, I would like to see a direct comparison with the classical paper (Hern´andez-Lobato & Adams, 2015). Finally, one of the justifications of the approximations that they propose is to avoid the numerical issues of the standard cdf. Have the authors compared with this, e.g. eq 18a, 18b? Using a robust implementation of the normal cdf/pdf function and further truncating them to avoid negative variances?
typo: Shortly before eq. 12, Should not S_{n-1} be defined as the softmax operation?