Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Exact closed-form Gaussian moments of residual layers

Simon Kuang, Xinfan Lin

TL;DR

This work introduces exact closed-form Gaussian moment propagation through deep neural networks by deriving layer-wise moment maps for activations including probit, GeLU, ReLU, Heaviside, and sine, and extending to residual connections. The authors define an analytic framework, Y_ana, that computes the exact mean and covariance after each layer and chain them across layers, while benchmarking against ground-truth and Gaussian proxies. They provide a theoretical Wasserstein-based guarantee and demonstrate substantial empirical improvements on random networks, regression and classification under input uncertainty, Bayesian network inference, and exploratory stochastic activations. The results indicate that layer-wise exact moment matching markedly improves uncertainty propagation, calibration, and predictive reliability in practical settings, with a priori error control and extensions to stochastic neurons.

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.

Exact closed-form Gaussian moments of residual layers

TL;DR

This work introduces exact closed-form Gaussian moment propagation through deep neural networks by deriving layer-wise moment maps for activations including probit, GeLU, ReLU, Heaviside, and sine, and extending to residual connections. The authors define an analytic framework, Y_ana, that computes the exact mean and covariance after each layer and chain them across layers, while benchmarking against ground-truth and Gaussian proxies. They provide a theoretical Wasserstein-based guarantee and demonstrate substantial empirical improvements on random networks, regression and classification under input uncertainty, Bayesian network inference, and exploratory stochastic activations. The results indicate that layer-wise exact moment matching markedly improves uncertainty propagation, calibration, and predictive reliability in practical settings, with a priori error control and extensions to stochastic neurons.

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.
Paper Structure (44 sections, 21 theorems, 121 equations, 120 figures, 223 tables)

This paper contains 44 sections, 21 theorems, 121 equations, 120 figures, 223 tables.

Table of Contents

  1. Introduction
  2. Methodology
  3. Our analytic method $Y_\mathrm{ana}$
  4. Ground truth(s) $Y_0$ and $Y_1$
  5. Theoretical Guarantees
  6. Adversarial examples
  7. Examples, applications, and extensions
  8. Random networks
  9. Input uncertainty: regression
  10. Input uncertainty: binary classification
  11. Weight uncertainty: variational Bayes networks
  12. Stochastic activations
  13. Novelty and significance
  14. Baselines
  15. The mean-field approximation $Y_\mathrm{mfa}$
  16. ...and 29 more sections

Key Result

Lemma 2.4

For some activation function $\sigma$, let $g$ be the function defined by $g_\sigma(x; A, b, C, d) = \sigma(A x+ b) + C x + d$. Let $X \sim \mathcal{N}(\mu, \Sigma)$. Then and where for all valid indices $(i, j)$,

Figures (120)

  • Figure 5.1: Probability distributions for Network(architecture=small, weights=trained, activation=probit residual), variance=large.
  • Figure 5.2: Comparison of goodness of approximation (lower KL divergence is better) for all random neural networks, grouped by approximation method, in the small input variance scenario.
  • Figure 5.3: Comparison of goodness of approximation (lower KL divergence is better) for all random neural networks, grouped by approximation method, in the large input variance scenario.
  • Figure 5.4: KL divergence (lower is better) between pseudo-true (ground truth moments) predictive distribution (by Monte Carlo) and approximations for the concrete compressive strength dataset. W24@$k$ means the $k$th partial sum of the GeLU covariance series of wright_analytic_2024.
  • Figure 5.5: Output distribution of a stochastic neural network, pseudo-true Normal distribution ("pseudo"), and layer-by-layer moment-matched Normal distribution ("analytic").
  • ...and 115 more figures

Theorems & Definitions (55)

  • Definition 2.1
  • Definition 2.2
  • Definition 2.3
  • Lemma 2.4
  • Definition 3.1
  • Example 1: for $Y_\mathrm{lin}$
  • Example 2: for $Y_\mathrm{u95}$ and $Y_\mathrm{u02}$
  • Example 3: for $Y_\mathrm{mfa}$
  • Example 4: for $Y_\mathrm{ana}$
  • Definition 2.1
  • ...and 45 more