Exact closed-form Gaussian moments of residual layers
This provides a highly accurate method for uncertainty quantification in deep learning, particularly for residual networks, with broad applications in fields requiring reliable inference under epistemic uncertainty.
The paper tackled the problem of propagating Gaussian moments through deep residual networks by deriving exact closed-form moment matching for several activation functions, achieving up to a millionfold improvement in KL divergence error on random networks and hundredfold improvements over state-of-the-art methods on variational Bayes networks.
We study the problem of propagating the mean and covariance of a general multivariate Gaussian distribution through a deep (residual) neural network using layer-by-layer moment matching. We close a longstanding gap by deriving exact moment matching for the probit, GeLU, ReLU (as a limit of GeLU), Heaviside (as a limit of probit), and sine activation functions; for both feedforward and generalized residual layers. On random networks, we find orders-of-magnitude improvements in the KL divergence error metric, up to a millionfold, over popular alternatives. On real data, we find competitive statistical calibration for inference under epistemic uncertainty in the input. On a variational Bayes network, we show that our method attains hundredfold improvements in KL divergence from Monte Carlo ground truth over a state-of-the-art deterministic inference method. We also give an a priori error bound and a preliminary analysis of stochastic feedforward neurons, which have recently attracted general interest.