Year: 2021

Abstract:
We propose a new approach to propagating probability distributions through neural networks. To handle non-linearities, we use local linearization and show this to be an optimal approximation in terms of total variation for ReLUs. We demonstrate the advantages of our method over the moment matching approach popularized in prior works. In addition, we formulate new loss functions for training neural networks based on distributions. To demonstrate the utility of propagating distributions, we apply it to quantifying prediction uncertainties. In regression tasks we obtain calibrated confidence intervals, and in a classification setting we improve selective prediction on out-of-distribution data. We also show empirically that training with our uncertainty aware losses improve robustness to random and adversarial noise.

heading: 1 INTRODUCTION
text: Neural networks are routinely used in applications affecting our daily lives, including safety-critical domains. As a result, quantifying uncertainties in neural network decisions and improving their robustness against noise has become an important problem. A prominent example is autonomous driving [1], where a neural network is not only supposed to detect and classify various objects like other cars or pedestrians on the road, but also to know how certain it is about this decision and to allow, e.g., for human assistance for uncertain cases. As pointed out by many works, e.g., by Kendall and Gal [2] or by Kiureghian et al. [3], prediction uncertainties can arise from two different sources: systematic uncertainty in the data, which is referred to as epistemic uncertainty, or random uncertainties in the data, e.g., because of noisy sensors, which is referred to as aleatoric uncertainty [4]. Both types of uncertainty, epistemic and aleatoric, have received considerable attention. Works focusing on epistemic uncertainties include out-of-distribution (OOD) detection [5], e.g., via orthogonal certificates [6] or via Bayesian neural networks [7–10]. Other works focus on aleatoric uncertainties, e.g., via uncertainty propagation [1, 11–13], via ensembles [14], via simultaneous quantile regression [6], or also via Bayesian neural networks [15]. Note that many of the aforementioned methods can be applied to quantify both types of uncertainties. Similarly, in this work, we focus on both aleatoric and epistemic uncertainties.
We consider the problem of evaluating f(x+ ǫ), where f is a neural network, x is an input data point, and ǫ is a random noise variable. Therefore, f(x+ ǫ) is a random variable, where its mean is used for prediction, and its variance for quantifying the uncertainty. This perspective allows assessing the sensitivity of the neural network for uncertainty quantification. Aleatoric uncertainty due to input measurement errors can be estimated by modeling how much a prediction changes under a respective input uncertainty. Incorporating the sensitivity of a neural network during training can also allow epistemic uncertainty quantification: when the output variance is regularized (as in our proposed loss function), the variance of predictions for observed data is minimized. Thus, for the regions where data does not provide sufficient information, the variance can be larger. We visualize the uncertainty quantification on a toy problem in Fig. 1: on the left our method identifies high uncertainty in the regions where samples from two classes overlap corresponding to aleatoric uncertainty; on the right, we capture epistemic uncertainty due to many possible decision boundaries separating the two classes. See figure 6 in Hüllermeier et al. [16] for an analogous example with linear models.
Due to the high complexity of neural networks, specifically many non-linear functions, it is intractable to compute the distribution of f(x+ ǫ) analytically. A straightforward way to approximate the mean and (co-)variances of f(x+ǫ) is to use the Monte Carlo method. Unfortunately, this is not suitable for
the problems we consider: evaluating a neural network f many times is computationally prohibitive, and the quality of the approximation for the (co-)variances as well as the quality of the gradients deteriorates quickly when the data is high-dimensional, e.g., images. Instead, we consider analytical parametric approximations of f(x+ ǫ). The key challenge is to approximate the transformations of distributions by non-linearities such as ReLU. The widely used moment matching technique [17] computes mean and variance of the corresponding transformed distribution and uses a Gaussian distribution with the same mean and variance as the approximation. Moment matching can also be applied for other distributions of the exponential distribution family [18]. However, this procedure requires that the moments are defined and finite, which is not always the case, e.g., for Cauchy distributions. Also, it cannot be applied to conventionally pre-trained deep neural networks. In addition, in our empirical studies, we observed that training with moment matching is numerically unstable, requiring careful hyperparameter tuning to obtain meaningful results.

heading: Contributions
text: • We propose an approximation of a distribution transformed by a ReLU non-linearity that minimizes the total variation distance to the true distribution. Empirically, we show that this approximation is also effective in approximating transformations with other popular activation functions, and applicable to multivariate and multilayer cases.
• We propose a new loss function for learning with Gaussian and Cauchy input distributions.
• We show that our method can quantify aleatoric uncertainty by obtaining calibrated confidence intervals in regression and epistemic uncertainty by improving selective prediction in classification.
• We demonstrate that our method improves robustness against random and adversarial noise.
Our PyTorch-based [19] framework for propagating distributions will be publicly available.

heading: 2 RELATED WORK
text: Several approaches have been suggested to introduce some treatment of uncertainty into neural networks. They can be grouped into three broad categories, although there are also hybrid approaches.
The first category tries to model uncertainty via sampling. Here, we find, among others: Variational inference for neural networks [9], a tractable approximation to Bayesian inference for neural networks; Bayes by Backprop [8], which learns a probability distribution by Monte Carlo sampling to introduce uncertainty in the weights of the network; Monte Carlo dropout [20], which estimates models’ prediction uncertainties by applying dropout [21] at test time. An overview of current techniques for Bayesian deep learning in the context of computer vision models can be found in Gustafsson et al. [22]. For aleatoric uncertainties Bouchacourt et al. [23] minimize the dissimilarity coefficient between the true and an estimated distribution (modeled by a neural network). The drawback of these methods is the computationally expensive inference, i.e., to make a single prediction they require many forward passes of a neural network.
The second category of approaches concentrates on modeling uncertainties without sampling, i.e., analytically. These works approximately propagate normal distributions at individual training samples through the network to model the network response to perturbed inputs [11, 12, 18, 24]. Wang et al. [18] propose Natural-Parameter networks, which allow using exponential-family distributions to model weights and neurons. Similarly, Wang et al. [25] and Postels et al. [26] use uncertainty propagation to sampling-free approximate dropout. These methods use moment matching for propagating parametric probability distributions through neural networks, as we will discuss in greater detail subsequently. Wu et al. [27] propose an approximation to moment matching of Gaussian distributions with covariance for ReLU non-linearities to provide a deterministic variational inference approximation (DVIA) for Bayesian neural networks.
The third category of approaches estimates uncertainties by training a neural network to predict not only an output value but also the output value’s uncertainty. Lakshminarayanan et al. [14] use neural networks for predicting Gaussian distributions instead of individual values and extend this idea with deep ensembles. Tagasovska et al. [6] learn conditional quantiles for aleatoric uncertainty estimation and propose orthonormal certificates for quantifying epistemic uncertainty. We use several methods from this category as baselines in our experiments.
Our approach falls into the second category and is most similar to Gast et al. [11] and Shekhovtsov et al. [12]. Our approach differs in how distributions are passed through non-linearities such as ReLUs. Previous work relies on moment matching [17], also known as assumed density filtering. Moment matching is a technique where the first two moments of a distribution, such as the output distribution of ReLU (i.e., the mean and variance) are computed and used as parameters for a normal distribution to approximate the true distribution. As it is otherwise usually intractable, moment matching assumes diagonal covariances. In our approach, we propagate normal distributions by an approximation that minimizes the total variation (TV) distance for ReLU activations. We find that our approximation is faster to compute, a better approximation with respect to TV, and allows propagating full covariances. Further, as we do not rely on moment matching, we can also propagate Cauchy distributions for which the moments are not finitely defined. Another advantage of our method is that it can be applied to pre-trained models. On the other hand, moment matching changes the mean such that there are deviations from what a network without uncertainty propagation would compute, leading to poor predictions when applied to a pre-trained network. Lastly, instead of approximating the distribution of the output classes for the loss function [11, 12], we propose a new loss function suitable for distributional outputs.

heading: 3 PROPAGATING DISTRIBUTIONS THROUGH NEURAL NETWORKS
text: For propagating parametric distributions through neural networks, we consider affine transformations and non-linearities, such as ReLU, separately. For affine transformations, exact computation of the parametric output distribution is possible due to the reproductive property of Gaussian and Cauchy distributions. For non-linearities, we use local linearization and show that this approximation minimizes the TV distance to the true intractable distribution for ReLUs. Recall that in our setting, the input to a neural network is a random variable and the weights are learned parameters.

heading: 3.1 AFFINE TRANSFORMATIONS
text: Fully connected layers as well as convolutional layers are affine transformations of their inputs.
For fully connected layers, we use the notation y = xA⊤+ b where A ∈ Rm×n is the weight matrix and b ∈ R1×m is the bias vector with n,m ∈ N+. As convolutional layers can be expressed as fully connected layers, we discuss them in greater detail in Supplementary Material B and only discuss fully connected layers here.
Given a multivariate normal distribution X ∼ N (µ,Σ), µ ∈ R1×n,Σ ∈ Rn×n, X can be transformed by a fully connected layer via µ 7→ µA⊤ + b and Σ 7→ AΣA⊤. Given a multivariate normal distribution without covariances X ∼ N (µ,σ2),µ ∈ R1×n,σ ∈ R
1×n, X can be transformed by a fully connected layer as follows: µ 7→ µA⊤ + b and σ2 7→ σ2(A2)⊤ where · 2 denotes the element-wise square.
Given a multivariate Cauchy distribution X ∼ C(x0,γ), x0 ∈ R1×n,γ ∈ R1×n, X can be transformed by a fully connected layer by x0 7→ x0A⊤ + b and γ 7→ γ Abs(A)⊤. Notably, in all cases, the location µ/x0 does coincide with the values propagated in conventional network layers. This allows applying this method directly to conventionally trained neural networks.
Average pooling down-samples by averaging the values of pooling regions. As this is a linear combination, it can be expressed using matrix multiplications and thus needs no special treatment.

heading: 3.2 NON-LINEAR TRANSFORMATIONS
text: To handle non-linearities, we utilize local linearization. That is, we transform the mean / median and the variance / scale as follows: (µ, σ) 7→ (f(µ), f ′(µ) · σ) (1) for univariate distributions and as (µ,Σ) 7→ (f(µ), f ′(µ)Σf ′(µ)⊤) (2) for multivariate distributions. ReLU is the most common non-linearity for neural networks, and we now study it in more detail. Following local linearization (Eq. 1), our approximation for transforming distributions with ReLUs is
ReLU : (µ, σ) 7→ { (µ, σ) µ ≥ 0 (0, 0) otherwise
(3)
for distributions parameterized via µ and σ. In fact, for Gaussian and Cauchy distributions, this approximation is optimal wrt. TV. The TV between two probability distributions, i.e., between the true distribution Q and an approximation P , is defined as
TV (P ,Q) = sup A
∣ ∣ ∣ ∣ ∫
A
(p− q) dν ∣ ∣ ∣
∣
= 1
2
∫
|p− q| dν . (4)
To motivate our choice of the approximation quality metric, we note that TV is a proper distance metric on probability distributions. Specifically, a TV of 0 implies that two distributions are the same, i.e., all their moments (including means) are equal. Devroye et al. [28] show that TV upper bounds the maximum of differences between moments for Gaussian distributions.
In the following theorem, we formalize that our approximation of the ReLU non-linearity minimizes the TV. Note that when parameterizing a Gaussian or Cauchy distribution with (0, 0), we are referring to Dirac’s δ distribution. An illustration of the transformed distributions and their approximations can be found in Supplementary Material A.
Theorem 1. Local linearization provides the best Gaussian approximation of a Gaussian distribution transformed by a ReLU non-linearity with respect to the total variation:
argmin (µ̃,σ̃) TV (P ,Q) =
{
(µ, σ) µ ≥ 0 (0, 0) otherwise
(5)
where P = N (µ̃, σ̃2), Q = ReLU(N (µ, σ2))
Proof. We distinguish 3 cases, µ < 0, µ = 0, and µ > 0:
(µ < 0) In the first case, the true distribution has a probability mass of p0 > 0.5 at 0 because all values below 0 will be mapped to 0 and CDF(0) > 0.5. As the approximation is parameterized as (0, 0), it has a probability mass of 1 at 0. Therefore, the TV is |1 − p0| < 0.5. All parameterized distributions where σ > 0 have no probability mass at 0 and thus a TV of at least |0−p0| = p0 > 0.5. (µ = 0) In this case, ∫∞
0 |p− q| dν = 0 because the distributions are equal on this domain. The
true distribution has a probability mass of p0 = 0.5 at 0, therefore the TV is |0 − p0| = 0.5. In fact, all distributions with σ > 0 have a TV of at least 0.5 because |0− p0| = 0.5. The distribution parameterized by (0, 0) has the same TV at this point.
(µ > 0) In this case, p0 = CDF(0) < 0.5. Further, again ∫∞ 0 |p− q| dν = 0. Thus, the TV of our distribution is |0 − p0| = p0 < 0.5. All distributions with σ > 0 have a TV of at least p0 because |0− p0| = p0. The distribution parameterized by (0, 0) has a TV of |1− p0| > 0.5. Thus, the approximation by linearization (Eq. 3) is the optimal approximation wrt. total variation.
Corollary 1. Theorem 1 also applies to Cauchy distributions parameterized via x0, γ instead of µ, σ.
Proof. The proof of Theorem 1 also applies here.
To demonstrate the accuracy of this approximation in the multi-dimensional and multi-layer case, we simulate multi-layer neural networks with ReLU activations and evaluate the TV between parametric approximations and Monte Carlo approximation with
106 samples representing the oracle. We use a neural network with 4 hidden layers and 100 neurons per layer with ReLU non-linearities trained with softmax cross-entropy on the Iris data set. This data set has small input and output dimensionalities, allowing us to use Monte Carlo as a feasible and accurate oracle estimator of the truth. The results of this simulation are displayed in Tab. 1.
The simulation shows that propagating the covariances produces the best estimates of the output distribution. Further, the simulation shows that when discarding covariances after each layer, on average, the approximation from Eq. 3 performs better than moment matching. Note that moment matching with covariances is intractable. Therefore, we use the deterministic variational inference approximation method (DVIA) by Wu et al. [27] in this case as a baseline. As DVIA has a significantly larger computational complexity (its computational overhead is linear in the size of the largest layer; see Section 5.5 for additional details), we can compare our method only in these simulation experiments to DVIA. Overall, all methods are better for small input standard deviations σ. The reason for this is that larger variances cause a larger part of the distribution to be mapped to 0 (for positive means) or to positive values (for negative means), which are the sources of error in the approximations. Results for additional architectures and non-linearities are presented in Supplementary Material C.1. There, we evaluate our method also for Leaky-ReLU, GELU [29], and SiLU [30].

heading: 4 LEARNING WITH UNCERTAINTY PROPAGATION
text: With the ability to propagate parametric distributions through neural networks, we can shift our focus to learning with such uncertainty propagation. Recall that the output of a neural network is now a (multivariate) Gaussian or a Cauchy distribution. We want to incorporate the covariances corresponding to a neural network’s output into the loss function. This is simple to achieve for regression, i.e., by using log-likelihood as the loss, however, for classification it is more challenging.

heading: 4.1 REGRESSION
text: For regression (possibly with multiple outputs), we use the probability density of the predicted distribution at the corresponding value as our training objective. For normal distributions, our objective is maximizing the probability density by minimizing the negative log likelihood
log det(Σ) + (y − µ)⊤Σ−1(y − µ) . (6) Here, the k dimensional prediction is µ with covariance matrix Σ and the ground truth value is y. For Cauchy distributions, the respective probability density can be maximized analogously.

heading: 4.2 CLASSIFICATION
text: More interesting than regression is the classification case, for which we propose a new loss function. Previous work has proposed using moment matching of softmax [11], Dirichlet outputs [11], and an approximation to the (n− 1)-variate logistic distribution [12]. Instead of finding a surrogate that incorporates the variance of the prediction, we use exact probabilities for classification. This allows incorporating not only the variances but also the covariances of a neural network’s prediction.
To arrive at our loss, we compute the probability of correct classification, i.e., the probability that the score for a certain class is the maximum among all classes. As the exact probability of a score among n ≫ 2 scores being the maximum is intractable, we resort to computing the pairwise exact probability of correct classification, which also allows us to consider their covariance.
The exact probability of pairwise correct classification for two random variables X and Y is
P(X > Y ) = P(X − Y > 0) = ∫ ∞
0
PDFX−Y (x) dx = CDFY−X(0) . (7)
For a multivariate normal distribution (X,Y ) ∼ N ((µX , µY ), Σ) with covariance matrix Σ. P(X > Y ) = 12 ( 1 + erf ( µX−µY√ 2(σ2
XX +σ2 Y Y −2σXY )
))
(8)
For independent normal distributions X ∼ N (µX , σ2X) and Y ∼ N (µY , σ2Y ). P(X > Y ) = 12 ( 1 + erf ( µX−µY√ 2(σ2
X +σ2 Y )
))
(9)
For Cauchy distributions X ∼ C(xX , γX) and Y ∼ C(xY , γY ). P(X > Y ) = 1
π
( arctan (
xX−xY γX+γY
))
+ 12 (10)
With that, our classification training objective is maximizing ∑ e6=c 1 k−1 P(Zc > Ze), (11) where Zc takes on the distribution of the correct class score and Ze takes on the distributions of the corresponding (among k − 1) erroneous class scores. We refer to this loss function as the pairwise Gaussian and the pairwise Cauchy, respectively.

heading: 5 EXPERIMENTS
text: First, we investigate how well our method quantifies aleatoric uncertainty on UCI regression tasks. Second, we test how well our method can estimate epistemic uncertainty for detection of out-ofdistribution test data. Third, we investigate robustness against Gaussian noise and adversarial attacks.

heading: 5.1 ALEATORIC UNCERTAINTY QUANTIFICATION
text: To evaluate aleatoric uncertainty predicted though uncertainty propagation, we examine its capability to produce calibrated prediction intervals on eight UCI regression tasks. We follow the experimental setting of Tagasovska et al. [6] and compare our method to their Conditional Quantiles method as well as to their best baseline, Conditional Gaussians [14]. Conditional Quantiles enable learning all conditional quantiles of a given target variable. Conditional Gaussian fits a conditional normal distribution, i.e., a neural network with two outputs, the mean and the variance. In comparison, our approach does not let a neural network predict the variance, but instead propagates input uncertainties through a neural network and computes calibrated prediction intervals from the output covariances.
We use the same model architecture, hyper-parameters, and evaluation metrics as Tagasovska et al. [6] for all three methods. The evaluation metrics are Prediction Interval Coverage Probability (PICP), i.e., the fraction of test data points falling into the predicted intervals, and the Mean Prediction Interval Width (MPIW). In Tab. 2, following Tagasovska et al. [6], we report the test PICP and MPIW of those models where the validation PICP lies between 92.5% and 97.5%. The goal is to achieve a narrow interval (small MPIW) while the optimal test PICP is 95%. Our method achieves the narrowest well-calibrated prediction intervals. Specifically, for 5 out of 8 data sets, our method has the narrowest intervals while Conditional Gaussian as well as Conditional Quantile each achieve the narrowest well-calibrated prediction intervals on only 2 of the data sets. Notably, for the ‘naval’ task, our method has the best-calibrated coverage intervals while having by far the smallest MPIW.

heading: 5.2 EPISTEMIC UNCERTAINTY FOR OUT-OF-DISTRIBUTION SELECTIVE PREDICTION
text: Selective prediction [31] is a formulation where instead of only predicting a class, a neural network can also abstain from a prediction if it is not certain. We benchmark selective prediction on MNIST [32] and use EMNIST letters [33] as out-of-distribution data. EMNIST letters is a data set that contains letters from A to Z in the same format as MNIST. We train a neural network on the MNIST training data set and then combine the MNIST test data set (10 000 images) with 10 000 images from the EMNIST letter data set. This gives us a test data set of 20 000 images, 50% of which are out-of-distribution samples and for which the model should abstain from prediction.
In the first setting, we train a neural network with conventional softmax cross-entropy to simulate an existing off-the-shelf network. In the second setting, we train the neural network using uncertainty propagation with the pairwise Gaussian loss (ours) as well as using moment matching propagation and Dirichlet outputs [11]. We evaluate using five methods to compute certainty scores. First, we use the softmax entropy of the prediction, and apply Orthonormal Certificates from Tagasovska et al. [6]. Second, we propagate a distribution with σ = 0.1 through the network to obtain the covariances. Using these covariances, we use our pairwise Gaussian probabilities as well as the categorical probabilities of the Dirichlet outputs proposed by Gast et al. [11] to compute the entropies of the predictions. In Fig. 2, we provide risk-coverage plots [31] of the selective prediction based on these scores. Risk-coverage plots report the empirical risk (i.e., the error) for each degree of coverage α. That is, we select the α most certain predictions and report the error. This corresponds to a setting, where a predictor can abstain from prediction in 1− α of the cases. We use risk-coverage area-under-the-curve (AUC) to quantify the overall selective prediction performance. Smaller AUC implies that the network is accurate on in-distribution test data, while abstaining from making a wrong prediction on out-of-distribution examples. In this experiment, no correct prediction on OOD data is possible because the classifier can only predict numbers while the OOD data consists of letters.
We evaluate a variety of combinations of training losses, uncertainty propagation methods, and confidence scores for selective prediction. See Tab. 3 for the summary and the corresponding risk-coverage AUC results. We find that training with an uncertainty-aware training objective (pairwise Gaussian
and Cauchy losses and Dirichlet outputs) improves the out-of-distribution detection substantially. Further, we find that the pairwise Cauchy scores achieve the best out-of-distribution detection on a pre-trained model. Orthonormal certificates (OC) estimate null-space of data in the latent space to detect out-of-distribution examples. In our setting, digits and letters may share a similar null-space, making OC not as effective. The Dirichtlet distribution does not perform well on a pretrained model because it is not designed for this scenario but instead designed to work in accordance with moment matching. We also observe that, for a model trained with Dirichlet outputs and moment matching uncertainty propagation, selective prediction confidence scores of pairwise Gaussian and pairwise Cauchy offer an improvement in comparison to the Dirichlet output based scores. This suggest that pairwise Cauchy scores are beneficial across various training approaches. The overall best accuracy (18.3% AUC) is achieved by training with an uncertainty-aware objective and evaluating using pairwise Cauchy. We also report AUC of a perfect predictor, i.e., one that always abstains on out-of-distribution data and always predicts correctly in-distribution.

heading: 5.3 ROBUSTNESS AGAINST GAUSSIAN AND ADVERSARIAL NOISE
text: Robustness of neural networks has received significant attention in the past years [34–36]. There are many methods for robustifying neural networks against random noise in the inputs or adversarially crafted perturbations. Our method is not explicitly designed for these purposes, however, here we empirically demonstrate that it achieves meaningful robustness improvements. This is an additional benefit to its main purpose of quantifying uncertainties.
We evaluate the robustness of classification models trained with uncertainty propagation on the MNIST [32] and CIFAR-10 [37] data sets. We consider the following training scenarios: First, we propagate normal distributions with and without covariance and train using the pairwise Gaussian loss. Second, we train with Cauchy distribution propagation using the pairwise Cauchy loss. As for baselines, we compare to training with softmax cross-entropy as well as moment matching propagation with the Dirichlet output loss [11]. The network architectures, hyper-parameters, as well as the results on CIFAR-10 are presented in Supplementary Material C.2.
Random Gaussian Noise To evaluate the random noise robustness of models trained with uncertainty propagation, we add random Gaussian per-pixel noise during evaluation. Specifically, we use input noise with σ = 0.5 (see Supplementary Material C.2 for σ ∈ {0.25, 0.75}) and clamp the image to the original pixel value range (between 0 and 1). In Fig. 3 (left), we display results for input standard deviations / scales from 10−9 to 104. For every input standard deviation, we train separate models. For each method, the range of good hyperparameters is around 4− 6 orders of magnitude. We observe, when training with uncertainty propagation and the pairwise Gaussian / Cauchy losses, the models are substantially more robust than with conventional training on the MNIST data set.
Adversarial Noise To evaluate the adversarial robustness of models trained with uncertainty propagation, we use the projected gradient descent (PGD) attack [35] as it is a popular and powerful attack. In Fig. 3 (right), we demonstrate the adversarial robustness of CNNs on distributions. We can see that our approach is competitive with the accuracy of the softmax cross-entropy loss and can even outperform it. For measuring the adversarial robustness, we use L∞-bounded PGD, for
MNIST with ǫ = 16/255 (and with ǫ = 8/255 in the supplementary material), and for CIFAR with ǫ ∈ {3/255, 4/255} (results for CIFAR are presented in the supplementary material). Note that our uncertainty propagation, in combination with pairwise probability classification losses, outperform both baselines in each case. We find that the adversarially most robust training method is propagating normal distributions with covariances, where we achieve a robust accuracy of 85% with ǫ = 16/255 on the MNIST data set compared to a robust accuracy of 14% for models trained with softmax crossentropy. Note that the baseline Dirichlet loss with moment-matching [11] offers no robustness gains in this experiment. We observe that, when training with normal distributions without covariances, the models are only robust against adversarial noise for rather large input standard deviations. At the same time, this effect is inverse for random Gaussian noise, where robustness against random noise is achieved when input standard deviations are rather small.
We perform PGD on our models with our objective functions as described in Eq. 11 and PGD on the cross-entropy trained models with the softmax cross-entropy objective. In Supplementary Material C.2, we provide results where a softmax cross-entropy based objective is used by the adversary: our models are more robust against the cross-entropy based attack than against the true objective based attacks which are reported in the plot. Our methods are not as robust as methods specifically designed for that purpose, e.g., adversarial training [35]. However, our methods provide a substantial improvement over training with softmax cross-entropy.

heading: 5.4 PROPAGATING DISTRIBUTIONS THROUGH RESNETS
text: To validate that our method also performs well for large architectures, we apply it to learning CIFAR-10 with ResNet-18 and ResNet-34 [38] architectures. Here, we train each model using Adam [39] for 400 epochs and compare it to the same setting with the softmax cross-entropy loss
in Tab. 4. The results validate that training with the Pairwise Gaussian loss (w/ covariances) does not substantially damage the performance in comparsion to softmax cross-entropy.

heading: 5.5 RUNTIME ANALYSIS
text: In Tab. 5, we report runtimes for propagating distributions. For small models, propagating distributions without covariances is almost without overhead while propagating covariances is slower. Moment matching is slightly more expensive than the proposed propagation, which is because moment matching requires evaluating ad-
ditional (scalar) functions at the non-linearities. The complexity (as a factor of the cost of propagating a single point through the neural network) of propagating distributions with covariances is linear in the number of outputs, while for propagating distributions without covariances it is constant. As a reference, exact propagation of the distributions (and not a parametric approximation as in this work) has an exponential runtime factor in the number of neurons [40]. The runtime factor of DVIA [27] is linear in the number of neurons in the largest layer, which makes experiments exceeding Table 1 infeasible. For example, in the robustness experiments on MNIST dataset in Section 5.3, the dimension of the largest hidden layer is ≈ 12 500, i.e., computational time of a single forward pass with DVIA is over 12 500 regular forward passes.

heading: 6 CONCLUSION
text: We studied the problem of uncertainty propagation in neural networks. Uncertainty propagation allows quantifying both epistemic and aleatoric uncertainties, and promotes robustness to input noises, as demonstrated in our empirical studies. Our local linearization based methods, together with the pairwise Gaussian and Cauchy loss functions, performed well. Extending our theoretical results to better understand the empirical success could be an interesting direction for future work.

heading: REPRODUCIBILITY STATEMENT
text: We will make the source code and experiments of this work publicly available to foster future research in this direction. All data sets are publicly available. We specify all necessary hyperparameters for each experiment. Each experiment can be reproduced on a single GPU.

heading: Supplementary Material for Propagating Distributions through Neural Networks
text: A ILLUSTRATION OF PARAMETRIC APPROXIMATIONS OF RELU

heading: B TRANSFORMATION OF CONVOLUTIONS
text: For two-dimensional convolutions, we use the notation Y = X ∗ W +B where X ∈ Rc0×n0×n1 is the input (image), W ∈ Rc1×c0×k0×k1 is the weight tensor, and B,Y ∈ Rc1×m0×m1 are the bias and output tensors with c0, c1, k0, k1, n0, n1,m0,m1,m0,m1 ∈ N+. Further, let XU ∈ R
m0×m1×c0×k0×k1 be the unfolded sliding local blocks1 from tensor X such that Yc1m0m1 = ∑
c0k0k1 XUm0m1c0k0k1Wc1c0k0k1 +B = X ∗ W +B.
Normal Distribution with Covariances For the convolutional layer, µ ∈ Rc0×n0×n1 ,Σ ∈ R
c0×n0×n1×c0×n0×n1 . Note that ΣU ∈ Rm0×m1×c0×k0×k1×m0×m1×c0×k0×k1 is the unfolded sliding local blocks covariance of the covariance tensor Σ. µ 7→ µ ∗ W + B and Σc1m0m1c′1m′0m′1 7→ ∑
c0k0k1c ′ 0 k′ 0 k′ 1
Wc1c0k0k1Σ U
m0m1c0k0k1m ′ 0 m′ 1 c′ 0 k′ 0 k′ 1
Wc′ 1 c′ 0 k′ 0 k′ 1 .
Normal Distribution without Covariances For the convolutional layer, µ,σ ∈ Rc0×n0×n1 . µ 7→ µ ∗ W +B and σ2 7→ σ2 ∗ W2.
Cauchy Distribution For the convolutional layer, x0,γ ∈ Rc0×n0×n1 . x0 7→ x0 ∗ W + B and γ 7→ γ ∗Abs(W).
1Einstein summation notation. The unfolded sliding local blocks are equivalent to torch.unfold.

heading: C ADDITIONAL EXPERIMENTS

heading: C.1 SIMULATION OF PROPAGATING NORMAL DISTRIBUTIONS
text: Tab. 6–22 show additional simulations with 1, 2, 4, and 6 hidden layers as well as ReLU, Leaky-ReLU, GELU, SiLU, and logistic sigmoid activations. Here, the reported metric is intersection of probability mass, i.e., 1− TV . The baseline is computed via 106 samples. Tab. 23 considers the average ratio between predicted and true standard deviations.
As there we did not find a moment matching method for GELU and SiLU in the literature, we omit moment matching in these cases. DVIA [27] is only applicable to ReLU among the non-linearities we consider. For moment matching with logistic sigmoid, we use numerical integration as it does not have a closed form solution.

heading: C.2 GAUSSIAN AND ADVERSARIAL NOISE ROBUSTNESS
text: In Fig. 5, we present an extension to Fig. 3 which demonstrates random and adversarial noise robustness for alternative noise intensities.
In Fig. 9 and 10, we present an additional demonstration of the adversarial and random noise robustness of CNNs on distributions. Note that we do not include results for normal distributions with covariances (as in Figures 3 and 5). Here, the emphasis is set on using both MNIST and CIFAR.
In Fig. 6, we show the robustness where a cross-entropy based objective (green) is used by the adversary. A cross-entropy based objective could be used in practice in multiple scenarios: For example, publishing just the weights of the model without information about the objective function, or by using a standard library for adversarial attacks that does not require or support specification of the training objective. We emphasize that our models are more robust against the cross-entropy based attack (CE-PGD) than the pairwise distribution loss based attack (PA-PGD) which is reported in all other plots.
Note that the phenomenon of vanishing gradients can occur in our models (as it can also occur in any cross-entropy based model). Thus, we decided to consider all attempted attacks, where the gradient is zero or zero for some of the pixels in the input image, as successful attacks because the attack might be mitigated by a lack of gradients. In Fig. 7, we compare conventional robustness to robustness with the additional Gradient-Non-Zero (GNZ) requirement. We emphasize that we use the GNZ requirement also in all other plots.
For evaluating the robustness against random noise, we use input noise with σ ∈ {0.25, 0.5, 0.75} and clamp the image to the original pixel value range.

heading: C.3 PROPAGATING WITHOUT COVARIANCES
text: We analyze the accuracy of a fully connected neural network (FCNN) where a normal distribution is introduced in the kth out of 5 layers. Here, the normal distribution is modeled without covariances. The result is displayed in Fig. 8, where we can observe that there is an optimal scale such that the accuracy is as large as possible. Further, we can see the behavior if the distribution is not introduced in the first layer, but in a later layer instead. That is, the first layers are conventional layers and the distribution is introduced after a certain number of layers. Here, we can see that the more layers propagate distributions rather than single values, the larger the optimal standard deviation is. The reason for this is that the scale of the distribution decays with each layer such that for deep models a larger input standard deviation is necessary to get an adequate output standard deviation for training the model. This is because we modeled the distributions without covariances, which demonstrates that propagating without covariances underestimates the variances.
D IMPLEMENTATION DETAILS

heading: D.1 SIMULATION OF PROPAGATING NORMAL DISTRIBUTIONS
text: We trained a fully-connected neural network with 1, 2, or 4 hidden layers and ReLU / Leaky-ReLU activations for 5000 epochs on the Iris data set via the softmax cross-entropy loss. Here, each hidden layer has 100 hidden neurons. Results are averaged over 10 runs, and 10 inputs are propagated for each of these 10 runs. Across all methods of propagating distributions, each Monte Carlo baseline as well as model weights are shared.

heading: D.2 ALEATORIC UNCERTAINTY QUANTIFICATION
text: For the aleatoric uncertainty quantification experiment, we use the same hyperparameter and settings as Tagasovska et al. [6]. That is, we use a network with 1 ReLU activated hidden layer, with 64 hidden neurons and train it for 5000 epochs. We perform this for 20 seeds and for a learning rate η ∈ {10−2, 10−3, 10−4} and weight decay ∈ {0, 10−3, 10−2, 10−1, 1}. For the input standard deviation, we made a single initial run with input variance σ2 ∈ {10−8, 10−7, 10−6, 10−5, 10−4, 10−3, 10−2, 10−1, 100} and then (for each data set) used 11 variances at a resolution of 100.1 around the best initial variance.

heading: D.3 EPISTEMIC UNCERTAINTY FOR SELECTIVE PREDICTION ON OUT-OF-DISTRIBUTION DATA
text: Here we use the same network architecture as for the robustness experiments in the next subsection (A CNN composed of two ReLU-activated convolutional layers with a kernel size of 3, a stride of 2, and hidden sizes of 64 and 128; followed by a convolutional layer mapping it to an output size of 10) and also train it for 100 epochs, however, at a learning rate of 10−3. For training with our uncertainty propagation, we use an input standard deviation of σ = 0.1. For training with uncertainty propagation via moment matching and the Dirichlet loss as in Gast et al. [11], we use an input standard deviation of σ = 0.01, which is what they use in their experiments and also performs best in our experiment.

heading: D.4 GAUSSIAN AND ADVERSARIAL NOISE ROBUSTNESS
text: Here, our CNN is composed of two ReLU-activated convolutional layers with a kernel size of 3, a stride of 2, and hidden sizes of 64 and 128. After flattening, this is followed by a convolutional layer mapping it to an output size of 10. We have trained all models for 100 epochs with an Adam optimizer, a learning rate of 10−4, and a batch size of 128.

heading: D.5 PROPAGATING WITHOUT COVARIANCES
text: The architecture of the 5-layer ReLU-activated FCNN is 784-256-256-256-256-10.

heading: D.6 EVALUATION METICS

heading: D.6.1 COMPUTATION OF TOTAL VARIATION IN THE SIMULATION EXPERIMENTS
text: We compute the total variation using the following procedure: We sample 106 samples from the continuous propagated distribution. We bin the distributions into 10 × 10 × 10 bins (the output is 3-dimensional) and compute the total variation for the bins. Here, we have an average of 1000 samples per bin. For each setting, we compute this for 10 input points and on 10 models (trained with different seeds) each, so the results are averaged over 100 data point / model configurations. Between the different methods, we use the same seeds and oracles. To validate that this is reliable, we also tested propagating 107 samples, which yielded very similar results, but just took longer to compute the oracle. We also validated the evaluation methods by using a higher bin resolution, which also yielded very similar results.

heading: D.6.2 COMPUTATION OF MPIW
text: Given a predicted mean and standard deviation, the Prediction Interval Width (for 95%) can be computed directly. In a normal distribution, 68.27% of samples lie within 1 standard deviation of the mean. For covering 95% of a Gaussian distribution, we need to cover ±1.96 standard deviations. Therefore, the PIW is 3.92 standard deviations. We computed the reported MPIW on the test data set. This setup follows the setup by Tagasovska et al. [6].

heading: D.7 MAXPOOL
text: For MaxPool, we apply Equation 1 such that (µ, σ) 7→ (max(µ), argmax(µ) · σ) where argmax yields a one-hot vector. The equation maxpools all inputs to one output; applying it to subsets of neurons correspondingly maxpools the respective subsets.

Review summaries:
The paper proposes a method to propagate uncertainties through neural networks by performing a first-order approximation at nonlinear activation functions. By injecting uncertainty into the inputs, prediction uncertainties are obtained. These prediction uncertainties are exploited at training time by using new uncertainty-aware loss functions, and at test-time to obtained prediction uncertainties, out-of-distribution detection, and adversarial robustness.

This paper proposed an approach to propagate probability distribution through neural networks. In particular, the authors use local linearization to handle nonlinearity and show its optimality in terms of total variation for ReLU. Empirically they show that the proposed method can provide calibrated confidence intervals for regression problems and improve detection of OOD data in classification problems compared to some baselines. 

This paper studies the problem of propagating probability distributions through neural networks and applies the results to quantify prediction uncertainties. It proposes a local linearization method to approximate a distribution transformed by a ReLU network, as well as new loss function for learning with distribution-valued inputs. It also provides empirical results, showing that the method can quantify two kinds of uncertainty (aleatoric and epistemic) in classification and regression tasks, and training with the new loss function can improve robustness to random and adversarial perturbations.

The paper proposes a simple method that propagates the uncertainty in the data through a neural network with ReLU non-linearity. The proposed method simply transforms the mean and covariance of the input Gaussian/Cauchy distribution through linear layers, and applies a local linear approximation for the ReLU activation. The proposed transformation is shown to be optimal under the total variation criterion. The paper applies the method to train neural networks with regression or classification targets, and show that the proposed method can well estimate the uncertainty of the network's output. Moreover, they also show some improvements in the robustness of the trained networks.