Exact closed-form Gaussian moments of residual layers

2026-01-29

AI Summary

This work addresses the challenge of accurately computing the mean and covariance of Gaussian distributions during forward propagation in deep residual networks. To this end, the authors propose a layer-wise moment-matching approach and, for the first time, derive exact closed-form expressions for moment propagation through feedforward and generalized residual layers with Gaussian inputs under a variety of activation functions—including Probit, GeLU, ReLU, Heaviside, and Sine. The method integrates analytical moment matching, asymptotic analysis, variational Bayesian inference, and stochastic neuron modeling, substantially enhancing theoretical completeness. Experiments demonstrate up to a 10⁶-fold reduction in KL divergence error in random networks, superior calibration of epistemic uncertainty on real-world data, and up to two orders of magnitude higher accuracy in variational Bayesian neural networks compared to existing approaches.

Technology Category

Machine Learning: Calibration & Uncertainty QuantificationReasoning under Uncertainty: Relational Probabilistic ModelsComputer Vision: Generative Adversarial Networks (GANs) for Vision

Application Category

Graph Algorithms and Modeling for the Web: Graph neural networks and deep learning approaches for Web-related graphsSemantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMsSocial Networks and Social Media: Influence propagation, information diffusion, and the prediction on networks

Abstract

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.

Problem

• Gaussian moments
• residual networks
• moment matching
• epistemic uncertainty
• neural network inference

Innovation

• exact moment matching
• Gaussian propagation
• residual networks
• closed-form solutions
• epistemic uncertainty