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