Propagating Distributions through Neural Networks

1 INTRODUCTION . 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 can not 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 . Contributions . • 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 . 2 RELATED WORK . 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 . 3 PROPAGATING DISTRIBUTIONS THROUGH NEURAL NETWORKS . 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 . 3.1 AFFINE TRANSFORMATIONS . 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 .

Summary: 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.