A Stochastic Polynomial Expansion for Uncertainty Propagation through Networks
Authors: Songhan Zhang
ShiNung Ching
Venue: Accepted by TMLR ()

Abstract: Network-based machine learning constructs are becoming more prevalent in sensing and decision-making systems. As these systems are implemented in safety-critical environments such as pedestrian detection and power management, it is crucial to evaluate confidence in their decisions. At the heart of this problem is a need to understand and characterize how errors at the input of networks become progressively expanded or contracted as signals move through layers, especially in light of the non-trivial nonlinearities manifest throughout modern machine learning architectures. When sampling methods become expensive due to network size or complexity, approximation is needed and popular methods include Jacobian (first order Taylor) linearization and stochastic linearization. However, despite computational tractability, the accuracy of these methods can break down in situations with moderate to high input uncertainty. 
Here, we present a generalized method of propagating variational multivariate Gaussian distributions through neural networks. We propose a modified Taylor expansion function for nonlinear transformation of Gaussian distributions, with an additional approximation in which the polynomial terms act on independent Gaussian random variables (which are identically distributed). With these approximated higher order terms (HOTs), we obtain significantly more accurate estimation of layer-wise distributions. Despite the introduction of the HOTs, this method can propagate a full covariance matrix with a complexity of $\boldsymbol{O}(n^2)$ (and $\boldsymbol{O}(n)$ if only propagating marginal variance), comparable to Jacobian linearization. Thus, our method finds a balance between efficiency and accuracy. We derived the closed form solutions for this approximate Stochastic Taylor expansion for seven commonly used nonlinearities and verified the effectiveness of our method in deep residual neural networks, Bayesian neural networks, and variational autoencoders. This general method can be integrated into use-cases such as Kalman filtering, adversarial training, and variational learning.

Published in Transactions on Machine Learning Research (06/2025)
A Stochastic Polynomial Expansion for Uncertainty Propaga- tion through Networks
Songhan Zhang Deportment of Electrical and Systems Engineering Washington University St. Louis
songhan.2@wustl.cmm edu
ShiNung Ching
shinungg Croustl edu
Department of Electricol and Systems Engineering Washington University in St. Louis
Reviewed on OpenReview: https: // openreveww. net/ pet//foum?? id= LyDRBhUshy
Abstract
Network-based machine learning constructs are becoming more prevalent in sensing and Docision-making systems. As these systems are implemented in safety-critical environments such as pedestrian detection and power management, it is crucial to evaluate confidence in their decisions. At the heart of this problem is a need to understand and characterize how errors the input of networks become progressively expanded or contracted as signals move through layers, especially in light of the non-trivial nonlinearities manifest throughout modern machine learning architectures. When sampling methods become expensive due to network size or complexity, approximation is needed and popular methods include Jacobian (first order Taylor) lincarization and stochastic lincarization. However, despite computational tractability, the accuracy of these methods can break down in sit uations with moderate to high imput uncertainty. Here, we present a generalized method of propagating variational multivariate Gaussian distributions through neural networks. We propose a modified Taylor expansion function for nonlinear transformation of Gaussian distributions, with an additional approximation in which the polynomial terms act independent Gaussian random variables (which are identically distributed). With these approximated higher order terms (HOTs), we obtain significantly more accurate estimation of layer-wise distributions. Despite the introduction of the HOTs, this method can propagate a full covariance matrix with complexity of On?)) (and (nn) only propagating marginal variance), comparable to Jacobian linearization. Thus, our method finds a balance between efficiency and accuracy. We derived the closed form solutions for this approximate Stochastic Taylor expansion for seven commonly used nonlinearities and verified the effectiveness of our method in deep residual neural networks, Bayesian neural networks, and variational autoencoders. This general method can be integrated into use-cases such as Kalman filtering, adversarial training, and variational learning.
1 Introduction
A fundamental problem in uncertainty estimation and verification is to characterize how given input distribution becomes transformed by the operant function (succinctly, Y= When takes the form of modern machine lcarning (ML) architecture, this problem quickly becomes analytically intractable, necessitating her sampling methods or approximation. The predominant approximation ion technique remains Jacobian linearization (JL), i.e., deterministic first order Taylor expansion around the mean of input distribution. To give few examples in ML contexts, Gandhi al. (2018): Dera et al. (2021): Petersen et al. (2024) used Jacobian lincarization for uncertainty propagation through networks, and Beiu et al. (1994); A Abdelaziz et al. (2015) used a piece- wise linear approximation of sigmoid functions prior to JT.. Perhaps

---

Published in Transactions on Machine Learning Research (06/2025)
unsurprisingly, these methods work best in low-uncertainty regimes, because taking derivate at input mean ignores the uncertainty of the derivative itself.
In another line of research, Bayesian neural networks try to do uncertainty awareness training by fitting Gaussians to their weights, so they will naturally output distributions rather than a deterministic results, avoiding overconfidence. In exact Bayesian inference, one needs to solve the posterior distribution of the p(D)0)p(0) The marginal distribution is an integral of due to the nonlincarity, this integral has to be approximated.
network paramcters f given data D, i.. D) = p(D(0')p(0))d''
the likelihood p(D over all possible combinations of network parameters. Since the likelihood is intractable
Early approaches attempled to precompute posterior mean and variance lookup tables to avoid Monte Carlo sampling at runtime (Hinton and van Camp 1993), but this necessitated discretizing continuous variables, introducing imprecision. Monte Carlo-based methods remained popular for some time Graves (2011) proposed Monte Carlo Variational Inference (MCVI) to approximate the evidence lower bound, though gradient estimation via sampling remained a limitation.
Between fully sampling and fully deterministic methods lies a class of numerical integration techniques that seek greater efficiency than naive Monte Carlo. Causs-Hermite quadrature is classical example: instead of drawing random samples, it deterministically chooses quadrature points and weights based on Hermite polynomials to efficiently approximate integrals of Gaussian-weieee underpins algorithms such as the cubalure Kalman filter (Arasaratnam and Haykin 2009 Särkkä 2013),
functions. This approach
widely used for nonlinear state estimation in signal processing and control. In Section A.13 we analyze the relationship between our method and Causs-Hermite quadrature, and demonstrate experimentally that our approach which is fully deterministic and does not rely on quadrature sampling achieves faster convergence.
At the fully deterministic end of the spectrum are cooment-matchig approaches, which propagate distributions by analytically computing the first and second moments of the output. The exact moments of Gaussian
distributions transformed by ReLU were first derived in Frey and Hinton 1999] and later reused in several works, such as Gast and Roth (2018). Spiegelhalter and Lauritzen (1990) and MacKay (1992b) provided approximations for sigmoid and tanh functions using different reformulations, while Wang et al. (2016) presented another similar approach. (Shekhovtsov and Flach (2019) developed an analytical method for propagating uncertainty through argmax and softmax, assuming logistic and Gumbel priors.
To integrate these ideas into BNN training, Jylänki et al. (2014) introduced expectation propagation. A more scalable alternative, probabilistic backpropagation (PBP), was later proposed by Hernández-Looaal and Adams 2015), which propagates moments forward through the network to estimate the marginal likelihood, then backpropagates gradients with respect to approximated posterior paramcters. However, PBP only propagates the diagonal of covariance. To address this limitation, Wu et al. (2019) introduced a formulation that enables full covariance propagation through ReLU and Heaviside activations.
Contributions To summarize, current uncortainty propagation of variational Gaussian distributions through nonlinear layers relies on one of the following (1) Monte Carlo sampling nonlinearities or (3) direct derivation (or approximation) of the first two moments of output distributions under the assumption of zero correlation. The first is plausible thanks to powerful tensor calculation with GPUs, but is prone to undersampling, especially in high dimensional settings. The second allows propagating full covariance matrices, but introduces errors from ignoring higher order moments' contribution to the covariance calculation. The last one introduces errors from ignoring the correlations" contributions to the variance calculation.
(2) local linearizat ion of the
To address these shortcomings, we here postulate a generalized framework of using a stochastic polynomial expansion as a surrogate nonlincarity, and derive the closed-form solutions of the mcan, covariance, and cross- covariance of propagated multi-variate distributions, for seven nonlinearities that are ubiquitous in modern ML network constructs. This is achieved with a computational complexity of Onn)), comparable to that of first-order Taylor expansion. Our methodology is inspired by stochastic lincarization (SL.), which uses expected value of the first derivative as gain and mean of output as bias, or Y "= VV, (X))())) (X E[X]) + E[)())
Stochastic linearization minimizes mean square of the residual (Booton 1953 Kazakov 1954 and has been
2

---

Published in Transactions on Machine Learning Research (06/2025)
used in the context of feedback control systems (see, e.g., Zhang et al. 2025 Ching et al.
2010 Elishakoff
and Crandall 2017
2 Theory
To ensure consistency and clarity, we explicitly define the notation and convention used throughout this paper. All vectors are vertical. All instances of the o operation denote element-wise (Hadamard) operations, including both element-wise multiplication and exponentiation. For example, given a vector X, we define X = X X, whereas X =²=XX² xx"" which represents the outer product. This convention extends to differntial operators well. Specifically, the dlement-wise second derivative of function f is given by:
18 Ou
f Ox
Vo2
diag (V² f)
Let X = (X,, T be Gaussian random vector following Wt, Σ), and - E E[X] Let f()) be smooth, univariate function. We define the vector function f(X) f(x)), f(X,))) Here, we have in mind that f(-) describes the activation function at the output of a leedforward layer. Now, let define a set of iidd. surrogate random vectors (s). which are mutually independent as random vectors, but the components within each vector may be correlated.
(a) Pw)-ina wiidd (O
1,,,
We now propose the central construct in this paper, the maeudo-Taylo polynomial expansion (PTPE) of fX) as:
9(X) E[f()) +
(()) s!
(5 - 3
(1)
We postulate and later show that this expansion provides a tractable and accurato approximation of f(X) for the purposes of propagating uncertainty through feedforward nelwork architectures.
In the above, we use the form of Taylor polynomial expansion to describe the behavior (in expectation) of the function f(X) subject to the stochastic input X. The choice of the i.i.d. surrogate polynomials, (1) 3 (3) ...., }, is made to simplify the ensuing derivations. Note that if taking only the first two terms, this expansion is equivalent to stochastic lincarization, because E [ equivalent to E[X]. It is straight.frwart to observe that g(X) has the same first moment as s(X) because all terms after the first one are designed to have zero mean. In the following, we will provide empirical evidence that the second moment is well-captured for many common activation functions.
First, we derive the solution for covariance and cross-covariance using the proposed stochastic polynomial expansion.
Lemma 1. Define
E [Vx((X)
E
(
Ao = [X(()]
A
A
1!
2!
Then, the covariance mairia of g( ( is
Σ
A, o (.coo(XX)) A
(2)
#=1
for an S-th order expansion. For S = 3,
Σg(x) = 5 AT
A₂ (25 A 1 +
Ag [6503 diag E) E diag T] AA
3

---

Published in Transactions on Machine Learning Research (06/2025)
Proof. The expected values be solved using central moments of Gaussian distributions and Isserlis' theorem. All power operations are element-wise. Note that since A, are dimensional vector, and all the power and product operations are element-wise, the complexity of calculating covariance is O(n²). For detailed derivation, see Appendix A.1.2
It is useful to note that an addition (rosidual or recurrent) layer sums the activation of two (or more) layers, c.g. X and g(Y). In thise case, the covariance X g(Y) the sum of their covariances and cross-covariances, i.e.
Ex+ Σx 2 + Exg(Y) EX
It thus helpful to postulate an additional lemma for the purpose of calculating covariance after addition.
Lemma 2. Let Y = (Y,, Y,..... Yn)T be onother Gaussian random vector that is cross-corvelated to X with Σvx. Then, the coss-covariance matriz between Y and Z == =g(X) is
Eyz
Σ
A cov coovYYXX;
(=1, odd
(3)
Σzy Σ
A, o cov(Xut Y)
s=1,8 is odd
for an S-th order expansion. For S=
Σyz A Eyx + +3A Eyx o diag Σx) Σzy A₁ Σxy + 3A, Exy diag Ex)
Proof. The expected value be calculated using Isserlis' theorem. Note that this term is nonzero only if and are odd. For details of derivation, see Appendix A.1.3
With these results, to find the covariance of the output of nonlinear layer, assuming the input follows a multi-variate normal distribution, one needs to derive the coefficients of the PTPE, i.., Ag, A1, Ag, etc., for the nonlinearity of interest. Note that these coefficients only depend on mean ju and variance 6² = ddiag Σ), not correlations, rendering the computational complexity (n)). We briefly discuss some of the techniques we adopted solve for these polynomial coofficients and list the final results in Table Table and Table For
detailed derivation for all nonlinearities, Appendix A.2 2.2
A.7
Tanh, Sigmoid, and Softplus. Because the integral / (f)n((()((())) is not tractable analytically, so we make a further approximation by substituting tanh with the error function which is very similar but more tractable. Specifically, we propose
1 p
tanh(x)
erf[
erf
p
where .. .pp is a set of scaling factors obtained by numerical optimization (see Eq in Appendix), and the relationship between approximation accuracy and the number of scaling factors is discussed in A.12 The error function defined as
/
n exp( t2)
crf(x)= = V
Then, the integral f (x)))()))((( can be approximated as 1 EE= /erf((eer))) which is tractable analytically. The higher order derivatives the error function are simply derivatives of Gaussian functions (p(x), which are related to Hermite polynomials H, (x) through
-1=
1
r
dx
=
V20²
H.
22
0
(4)
4

---

Published in Transactions on Machine Learning Research (06/2025)
where
(z)
H (z) 2x
H2(0) = 4x²
We show the Speudo-Tayloo coefficients are convolutions of Gaussian derivatives and Gaussian pdf, which are analytically tractable (sce Appendix A.2 We used a similar treatment for sigmoid function (sec Appendix Using error functions as an approximation was first. suggested by Spiegelhalter and Lauritzen (9900,, but we use a linear combination, which is easily parallelizable and enhanced approximation accuracy. The derivation for softplus can reuse the results of sigmoid, because the derivative of softplus is just sigmoid with scaling factor 3 A.4
A.3
ReLU and Lea. It is obvious that we cannot apply our method directly on ReLU, because it is not continuously disferentiable. Hence, we modified the results for softplus at the limit of 3 + oc, considering the relationship (A.5)
lim 800
log (1 + elt) (0,
Similarly, leaky ReLU and any piece- wise linear activation function can be described a combination of ReLU functions with different scaling, shifting, and/or mirroring.
GELU and SILU. The derivatives of GELU can be expressed using derivatives of Gaussian edf ((x)
e GELU(x)
g-11
0"
s Oz (x)
Ox ((x)
dr
The expected value of the GELU derivative involves integrating the product of Hermite polynomials and Gaussian functions, which analytically tractable (A.6). Using normal odf approximate sigmoid function, the derivations for a SiLU function becomes similar that of the GELU A.7).
Before presenting the final expression, Let 6² diag(S) Σ) denote the input variance and 63 represent constant dependent on the nonlincarity (see Table column 3). Defining oद σ² + 63. then have
introduce few element-wine operators to simplify the formulation.
1 Φj 4
L(ū ;ōz)
Likelihood of observing x = X NV
z=
Flu :0,)
+
Cumulative likelihood of observing x S p
p
I(p ;0,) =
x
dr = 1 ;0,) 03 ;0))
Expected value of excess (ī X) (only when j > x, and X NO.. 03))
We also define an element-wise derivative operator De such that
D'
C(c ((x;)))
H
C(p
200
26
Notice that
D I(p ;0;) F(j ;6,) D2 I(p :0,) = DF :,,) C(p ;0;)
Then, it becomes evident that all of these aseudo-Taylor coefficients A. can be written as derivatives of I, F, and L (Table. []. This a result of the Gaussian assumption and the reformulation tanh and sigmoid.
fAll operations the right hand side of the equation are element-wise.
5

---

Published in Transactions on Machine Learning Research (06/2025)
Table 1: General solutions for the pseudo- Taylor coefficients. For notational simplicity, all the product, division, and exponentiation operations are clement-wise, For expanded solutions of the first four coefficients, see Table in the Appendix.
Nonlinearity
General Solution
83 5² + +63
Definition T
Tanh
A. D [2F(A ;,, - -1]
1
63 215
tanh(2)
Σ erl((734)
8.
j=1
Sigmoid
A. 5.
D'F(p a
67 = 27
sigmoid(z) Pd
Σ
27g x)
p
Softplus
A, 8. j=l
1
D'I(j ;6,)
6) 277
softplus(2;
B)
Σ
(v (V223
da
p j=l
1
ReLU
4. D'I(ā; :0)
8²
8
1
(0)))
A,
D (0a I(p :ō]
8
GELU
A,
D [T :6) Liū ; 6)]
6²
8
1
SILU
A,
D' [[ b, 63 p( :b,]]
63 = 27g
same Sigmoid
8.
3
Results
3.1
PTPE significantly improves estimation accuracy when exposed to higher input variance
As an initial empirical test and demonstration of concept, we applied PTPE to single, univariate nonlinearity subject to a parameterized normally distributed input. We varied the input mean and variance and examined how the output mean and variance compared to those predicted by PTPE. For this comparison, the true output statistics were obtained through 107 Monte Carlo sampling across all input parameters. As expected, PTPE far outstrips Jacobian linearization, and this effect is prominent especially when input variance is high. With up-to third order PTPE, the estimated variance by our method is already very close to the ground truth (Fig. col 4).
3.2 PTPE accurately quantifies uncertainty in canonical network architectures
To benchmark PTPE for uncertainty estimation in neural networks, trained residual neural networks (He et al. (2016)) with three depths (13, 33, and 65 layers) and 3 three typical nonlinearities (Tanh, ReLU, GELU) on CIFAR10 (Krizhevsky (2009)). We corrupted each input image with additive Gaussian noise to simulate noise in low light conditions (first type of corruption in Hendrycks and Dietterich (2019)), then compared the PTPE-predicted and reference (via 10? Monte Carlo sampling) logits distributions. Four levels of corruption, with noise variance values of [1, 10, 100, 1000], were applied to RBG values (I0, 255)) of the input image. If z-scored, the corresponding noise variance scales are [1c-5, 1e-4. 1c-3, 1e-2). The visualization of the corrupted images are shown in (Fig. 2 The layerwise application of PTPE is outlined in Algorithm with accompanying pscudo code.
We measure the estimation accuracy of moments in three ways: the Euclidean distance from the reference mean to the predicted mean, || Pret 2, the Frobenius norm of the covariance residuals ||Eest - Ever and the 2- Wasserstein distance (or Clantorovich-Rubinston metric) between the reference and estimated dist ribut ions, assuming both distributions were Gaussian. This 2- Wasserstein distance is defined as
W2 VIl
Post
Preel12 + trace Enst + Eref 2 Ex//2> Eres 2
1/2)

---

Published in Transactions on Machine Learning Research (06/2025)
Msan
Variance Gain, 1st Order
Variance Gain, 2nd Order
Variance Gain, 3rd Order
10
Tanh
10
-0.5
10
-1
0.8
ReLu
0.2
1.5
GELU
0.5
0
center of input dist.
center input dist.
center input dist.
center input dist.
True 2o (mput 0.1)
True (o
True (o2
10)
True
(o input 100)
inpul
irrpt
PTPE ("p 0.1) PTPE (o impu."
PTPE (o2
Irput 10)
PTPE (o
100)
mumi. Jacobian Linearization
Figure Solid lines: mean (column 1) and variance gain (output variance divided by input variance) (column 2-5) obtained by sampling 1e7 datapoints from Gaussian distributions with centers ranging from to approximated mean and variance gain predicted with 1st, 2nd, and 3rd order pseudo Taylor polynomial expansion (column 2 4). Colors correspond to different input variances (blue: 0.1, yellow: 1, red: 10, green: 100). Dotted lines: approximated mean and variance gain using Jacobian linearization (first
3. Dashed lines:
order deterministic Taylor expansion around input mean, e.g. Petersen el al. 2024
For other nonlinearities,
see Fig.
Estimated
0²
G² 10
Uncertainty Propagation Methods
l.d. Gaussian Noise
a²
Ground Truth
Large-scale
Monte Carlo
Figure Schematic of experiment setup. We inject i.i.d. Gaussian noise to input image to simulate sensor noises, then compare the estimated output distribution to the ground truth obtained by large-scale simulation (sampling 107 noisy images).

---

Published in Transactions on Machine Learning Research (06/2025)
ResNet13 tanh
10
10
10a
Wasserstein2 distance
10
101
Frobenius norm
100
L2 norm
109
100
10
JL
Wang al. (2016)
10-1
10
PTPE 1st
10
PTPE 3rd
10
10
100
10
10
10a
10
10
108
10
10
10
108
ResNet33 - tanh
10a
102
102
Wasserstein2 distance
10
10
Frobenius norm
109
L2 norm
109
109
10
10-1
10a1
10
20
10
100
10
10a
10a
10
10
10a
103
10
10
10
10
ResNet85 tanh
102
10
10a
Wasserstein2 distance
10
101
Frobenius norm
109
L2 norm
10
100
1
1011
10a1
10-1
10 2
10a 2
10
10
10a
10
10!
10
10
10
10'
10a
10a
Input Variance
Input Variance
Input Variance
Figure 3: Estimation accuracies of expectation propagation methods evaluated on residual networks of three different depth, and four different input variance. The smaller the better. We chose Petersen et al. (2024) example of JL and Wang et al. (2016) ReLU and GELU, see Fig. 9 and 10 in the Appendix. Differences between methods are systematic rather than random; we assessed variability by repeating experiments 10 times and computing error bars, but the resulting variances were too small to be visible at this scale. For transparency and reproducibility, all data and code used to generate these results are publicly available in our GitHub repository (Soc.
alternative Tanh approximation. For non-lineaiites
We summarize the results in Fig. 3 and 10 Overall, the experimental results align with expectations: (1) Jacobian linearization degrades dramatically in moderate to high variance regime. (2) Direct derivation is not suitable for this task due to the assumption of independence, since the overlapping convolution kernels and residual layers introduce substantial correlation. (3) Introducing up-to the third order PTPE typically outperformes stochastic and Jacobian linearization by a large margin. We also compared using 4 scaling factors versus using 8, and the results showed no significant difference (Fig. 12 justifying our choice of four scaling factors rather than a larger number in this case.
While higher-order approximations generally yield better performance, they can occasionally underperform in practice (e.g., for the 1.2 norm in ResNet13-fanl with input variance 1), due to two factors:
8

---

Published in Transactions on Machine Learning Research (06/2025)
Numerical Instability at Low Variance: When imput variance is small (for example, due to batch normalization), higher-order terms become increasingly sensitive to numerical instability. This instability can introduce errors that outweigh the potential benefits of higher-order corrections.
Imprecise Reference Distributions: Our evaluation relics on reference distributions estimated using 10 samples. While increasing the sample size from 10 to 107 led to lower Wasserstein-2, I,2, and Frobenius distance metrics, it was not feasible to further increase the sample size to 10 due to prohibitive runtime and memory requirements. We expect that such an increase would further reduce these metrics, but this could not be tested directly. Notably, estimating the statistics (e.g., mean and covariance) of a high-dimensional distribution (e.g., 10D or more) ideally requires an exponentially larger number of samples (on the order of 10a0 for 10 dimensions), which is computationally infeasible. Thus, some of the observed underperformance of higher-order terms may be attributable to imperfections in the estimated reference.
While the socond issue is inherent to the problem setup and cannot be fully climinated, the first issue is, in principle, addressable. In practice, we recommend adaptively choosing the order of approximation based on input variance; for example, using second-order PTPE when variance is low to avoid instability, and higher orders when the input variance is larger.
3.3 PTPE addresses the limitations of DVI by incorporating non con-piecewise-iness activations.
One major contribution of this work is to address the lack of accurate deterministic moment estimation for general nonlinearities in the field of variational inference. The method of deterministic variational inference (DVT) proposed by Wu et al. (2019) shares the same goal, which is to find deterministic method to approximate moments in neural networks, thus eliminating gradient variance. However, closed form solutions of posterior mean and covariance were solved only for piecewise linear activations such as ReLU and Heaviside. We know Var(f(X)) = [[f(X)² [[]) and the first term, [[ = pxx)xxxxxx)))))))) becomes arduous to solve for more complex nonlinearities f(-). Our approach circumvents this issue by taking derivatives inside the integrals, which provides tractability for general nonlinearities.
Table 2: Averaged test performance in RMSE. The smaller the better. The best value is highlighted in bold.
Yacht Hydro-
Dataset
Concrcte Strength
Energy Elticiency
Kin8nm
Naval Propulsion
Power
Protcin Structure
Quality (Red)
dynamics
Data points
768
8192
45730
1599 11
308
Dimensions
8
16
9
6
MCVI
7.128
0.129
2,668
1: 0.081
0.099 1: 0.001
0.005 0.001
4.327 1. 0.035
4,842
:: 0.031
0.648
0.008
6,887 1. 0.675
PBP
5.867
0.000
1,804
0.048
0.098
0.001
0.006
0.000
4.124
0.035
4,732
0.013
0.635
0.008
1.015 + 0.054
Dropout
5.23 + 0.12
1.66 = 0.04
0.10 + 0.00
0.01 + 0.00
4.02 :: 0.04
4.36
= 0.01
0.62
0.01
1.11 : 0.09
Ensemble
6.00
0.58
0.29
0.09
4.11 0.17
0.64 0.04
PTPE R+LU
5.196
0.206
0.615
0.024
0.072 + 0.001
0.003 + 0.000
3.925
0.025
4.445 = 0.042
0.633 0.010
0.640 + 0.067
PTPE GELU
5.068
0.153
0.570
1: 0.021
0.071
1: 0.000
0.004 0.000
3.915
1: 0.024
4.415 :: 0.043
0.634 0.010
0.623 1: 0.049
PTPE Tunh
5.574 0.148
0.580
0.022
0.076 0.001
0.006 0.000
4.073
4.364
0.628 0.010
1.678 + 0.193
Table 3: Averaged test performance in average log- lielliicl The larger the better. The best value is highlighted in bold, and the second best is underlined.
Wine Quality (Red)
Yacht. Hydro- dynamics
Dataset
Coucrete Strength
Energy Efficiency
KinSum
Naval Propulsion
Power Plant
Protein Structure
MCVI
-3.391
0.017
-2.391 + 0.020
0.897 + 0.010
3.734
0.116
-2.890 + 0.010
-2.992 + 0.006
-0.980
0.013
-3,499 = 0.163
PBP
-3.161
0.019
-2.012
: 0.019
0.896
: 0.006
3.731
0.006
-2.37
= 0.009
-2993 = 0.003
-0.98
0.014
-1.631 = 0.016
Dropout
-3.04 0.02
-1.99
0.02
0.95
0.01
3.80
0.01
-2.80
0.01
-2.9
0.01
0.01
-1.55 0.03
Ensemble
-3.06 0.18
-1.38
4 0.22
1.20 1: 0.02
5.63
0.05
.2.99 1 0.04
.283 1: 0.02
.0.94
0.12
-1.18 1: 0.21
DVI
-3.06
0.01
-1.01
1 0.06
1.13 1 0.00
6.29
0.04
2.80 = 0.00
-2.5 1: 0.01
-0.90
0.01
-0.47 = 0.03
PTPE ReLU
-3.010
0.037
-1.045
1 0.044
1.251
1 0.009
5.751
0.086
.2.79
0.007
-2821
0.024
-0.66
0.029
.0.910
1. 0.044
PTPE GELU
-3.022 = 0.066
-0.789
4 0.039
1,278 4 0.007
5.858
131
-2.750
= 0.006
-2.001
+ 0.019
-0.982
0.029
-0.236 + 0.052
PTPE Tanh
-3.559
0.039
-0.827
+ 0.043
1.234
+ 0.008
6.050
0.0 0.028
-2.825
0.000
22802 0.013
-0.93-
0.016
-0.009 0.067
LL Tanh
-3.07 0.07
-0.05 : 0.05
1.29 4 0.01
6.29
0.19
22.79 0.01
0.00 -2.79
0.01 -0.98
-0.92 0.03
9

---

Published in Transactions on Machine Learning Research (06/2025)
We provide additional context and quantification by replacing the forward-passing functions of Gaussian moments in DVT with PTPE and conducting regression experiments on eight UCT datasets. Following the methodology suggested by Hernáádezz-oobe containing 50 hidden units (100 for the larger Protein Structure datasct) and report the best test performance. We randomly set aside 10% of the data as lest samples, and the error bars reflect the results from 20 random splits.
and Adams (2015), we search over MLPs with up to four layers
We evaluate RMSE and log-likelihood on the held-out data and summarize the results in Tables and For comparison, include reported statistics (where available) from Monte Carlo Variational Inference (MCVI) Graves 2011), Probabilistic Backpropagation (PBP) Hernández--oobale and Adams 2015), Dropout (Gal
and Ghahramani 2016), Ensemble [Lakshminarayanan et al. 2017), DVI (Wu et al. 2019), and Linearized
Laplace (MacKay 1992a Foong et al. 2019). The results indicate that PTPE-DV achieves competent
accuracy, demonstrating the effectiveness of PTPE.
We extend the evaluation to out-of-distribution (OOD) detection in MNIST. Here, we test how models trained with PTPE respond to rotated and OOD images, using FashionMNIST (iiao et al. 2017 as OOD data. The format of our analysis is similar to that of Ovadia et al. (0199).
The results (Fig. show two main findings. First, the DVI+PTPE model achieves the highest accuracy on shifted images and is the least overconfident among all models tested, though this comes at the cost of a small drop in performance on undistorted images. One possible explanation is that the DVI+PTPE model may have learned a more robust representation of handwritten digits, which helps it generalize better rotated images; however, we leave more detailed investigation of this effect to future work. Second, all models perform similarly in terms of OOD detection capability.
(a) Brier Score, Rotated MNIST
(b) Log Likelihood, Rotated MNIST
(c) Confidence vs Count, Rotated 90"
0.5
0.150
-2
Accuracy on (ny(p) T
0.125
Log Likelihood
-4
Brier Score
0.100
-6
0.075
0.3
-8
0.050
Vanilla
Dropout Ensemble DVI +PTPE
-10
0.2
0.025
-12
0.000
0.1
30
60
90 120 150 180
0
150 Intensity (degrees)
30
60
90
120
180
0.00
0.25
0,50
0.75
1.00
Intensity of Shift (degrees)
Confidence Threshold (t)
(d) Confidence vs Count, Rotated 90"
(e) Entropy Distributien, FashionMNIST (DOD)
(f) Confidence vs Count, FashionMNIST (OOD)
10000
Vanilla
10000
3500
Dropout
# examplees es confidence 2
# examples es w/ confidence
Number of examples
3000
Ensemble le PVI+
8000
2500
6000
6000
2000
4000
1500
4000
1000
2000
2000
500
0.00
0.25 0.50 0.75
1.00
0.00
0.25
0,50
0.75 1.00
Contidence Threshold (x)
Entropy (Nats)
Confidence Threshold (t)
Figure 4: (a) and (b) show accuracy measured by Brier score and log likelihood as the test images are increasingly rotated. We explore the predictive distributions in 90 degree rotation of each method by looking at the confidence of the predictions in (c) and (d). We also explore the entropy and confidence of each method on entirely OOD data in (e) and (f).
10

---

Published in Transactions on Machine Learning Research (06/2025)
3.4 PTPE enabled Variational Autoencoder demonstrates improved performance with O(n) complexity.
Since PTPE is fundamentally designed accurately propagate Gaussian moments through nonlinearities, it can be incorporated into various applications, one of which is the decoder of a Variational Autoencoder (VAE) [Kingma and Welling 2014). In the decoding stage, VAE propagates the Gaussian means and variances encoded by the encoder through layers of nonlinearity in the decoder. The original model, which refer to as the 'vanilla VAE," accomplishes this by propagating Monte Carlo samples. Instead, we replace the decoder with PTPE while keeping the trainable parameters unchanged and train the model to reconstruct MNIST handwritten digits (LeCun et al. 1998). As shown in Fig. [ the PTPE-enabled VAE achieves a higher Evidence Lower Bound (ELBO) and improved reconstruction accuracy compared to the vanilla VAE.
A key bottleneck in applying many well-established Bayesian methods to VAEs is scalability. Many Bayesian approaches require sampling, which suffers from curse of dimensionality. PTPE offers an alternative solution: the РТРPEVVÉ shown in Fig. propagates only the diagonal of the covariance, resulting in computational complexity of (n). Moreover, the training procedure remains identical to that of the vanilla VAE.
2D latent space
5D latent
train (vanila)
-80
test (vanila)
-100
train (PTPE)
-100
ELBO
-120
120
-140
-440
-160
-160
-180
-180
50
100
150
200
50
100
150
200
epoch
epoch
10D latent space
20D latent space
30
-80
-80
Reconstruction Loss
Vanilla PTPE
-100
-100
ELBO
-120
-120
20
-140
15
-160
-160
10
1100
100
150
epoch
epoch
Latent Space
Figure 5: (Left four) Training and testing objective, measured as ELBO, for both vanilla and РТРE implemented VAE. (Right) Reconstruction loss (MSE) of the two types of models.
4 Discussion
One immediate potential limitation of PTPE is its reliance on the assumption that inputs Gaussian. It has been well-cstablished that at the limit of infinite width, deep neural network with Gaussian input is equivalently a Gaussian process Neal (1994); Williams (1996); de G. Matthews et al. (2018) Lee et al. 2018); Gao et al. (2023 ), and a similar phenomenon also reported in Bayesian neural networks with Gaussian weights (Goulet et al. (2021); Nguyen and Goulet (2022)) ). Based on this observation, we assume a 'wide enough' neural network will have approximate Gaussianity in each layer, so that the error of using variational Gaussian distributions to approximate layer-wise distributions becomes negligible. We verify this assumption through simulation (sce c.g., Fig.
In this paper, we focus on the propagation of Gaussian distributions. This choice is due to their prevalence in machine learning and their convenient property of being Lévy alpha-stable, meaning linear combination of Gaussian random variables remains Gaussian. This makes Gaussian distributions pertinent to our objectives. Consequently, our method could potentially be extended to other types within the Lévy alpha-stable family.
11

---

Published in Transactions on Machine Learning Research (06/2025)
For instance, Petersen et al. (0024) demonstrated the propagation of Cauchy distributions through neural networks. A more comprehensive survey is provided in Wang et al. 2016). where the authors ex- amined the propagation of exponential family distri- butions (including Beta, Rayleigh, Gamma, Poisson, and Gaussian), though this requires more intricate derivations.
A strength of PTPE is its generality. As mentioned in the introduction, several immediate motivating use-cases are in the training robust networks in- cluding probabilistic network models. Furthermore, our proposed method may also find application in safety-critical engineering systems that require timates uncertainty. Recently, researchers com- bined an LSTM and Kalman filtering to monitor the states of plasma inside a nuclear fusion device Pavone et al. 2023 The Kalman filter, by construction, requires statistics on the output of the LSTM or- der to generate control signals. Such statistics were generated by using probabalistic architcture within the LSTM, i.., where parameters specified by learned distribution. PTPE provides patential alternative path for such problems (we discuss in A.13), but enabling uncertainty propagation through deterministic learned architectures.
Figure 6: Empirical distributions of the first three logits of resnet, obtained by Monte Carlo sampling 10 images with additive Gaussian noises. For distributions of all 10 logits, Fig. in Appendix.
5 Conclusion
In this article, developed stochastic polynomial expansion approach, PTPE, to perform uncertainty propagation in neural networks. Our method offers significant advantages in accurately propagating the full covariance matrix of an input distribution compared to state-ofthhe-art methods, without substantially sacrificing computational efficiency. We derived analytical solutions for the first two moments of the output distributions for seven commonly used nonlinearities, demonstrating remarkable accuracy in predicting univariate mean and variance, particularly under high uncertainty. Additionally, we assessed multivariate accuracy in deep residual neural networks trained on image categorization tasks. By incorporating up to third-order polynomial expansion, method generally outperformed others, except in scenarios with minimal uncertainty in which the performance of competing methods is comparable. Overall, our proposed method provides a tractable framework for solving uncertainty propagation problems. It can potentially be effectively applied in various domains, including adversarial training, Bayesian inference, generative models, and safety-critical applications, offering a versatile tool for enhancing the reliability and robustness of neural networks.
6
Code Availability
All code for reproducing the experiments and figures is publicly available at httts:///github... com com.ooggan/n////// Stochastic_PPlyynomaa. Expaassiiiiiiiiiin
7
Acknowledgement
This rescarch was, part, funded by the U.S. Government (award no. HR00112290113). The views and conclusions contained in this document are those the authors and should not be interpreted as representing the official policies, either expressed or implied, of the U.S. Government.
12

---

Published in Transactions on Machine Learning Research (06/2025)
References
A. Н. Abdelaziz, S. Watanabe, J. R. Hershey, E. Vincent, and D. Kolossa. Uncertainty propagation through deep neural networks. In Interspeech 2015, Dresden, Germany, September 2015.
Arasaratnam and S. Haykin. Cubature kalman filters. IEEE Transactions on Automatic Control, 54(6): 1254-1269, 2009. doi: 10. TL..//AA.2009.0000110
V. Beiu, J. Peperstracte, J. Vandewalle, and R. Lauwereins. Vlsi complexity reduction by piece-wise approximation of the sigmoid function. In The European Symposium on Artificial Neural Networks, pages 181 186, 1994.
R. C. Booton. The analysis of nonlinear control systems with random inputs. In Proceedings, symposium nomlinear circuit analysis, volume 2, pages 369 391, 1953. see also IEEE Transactions on Circuit Theory CT-1 (1954), pp 9-18.
P. Bromiley. Products and convolutions of gaussian probability density functions. Technical report, Tina-Vision Vision Memo, 2003.
S. Ching, Y. Eun, C. Gokcek, P. Kabamba, and S. Meerkov. Quasilinear Control: Performance Analysis and Design of Feedback Systems with Nonlinear Sensors and Actuators. Cambridge University Press, 01 2010. ISBN 9781107000568. doi: 10.1017/C100978011199666
W. J. Cody. Rational chebyshev approximations for the error function. Mathematics of Computation, 23 (107):631 1-7, 1969. ISSN 00255718, 10886842.
G. de G. Matthews, J. Hron, M. Rowland, E. Turner, and Z. Ghahramani. Gaussian process behaviour in wide deep neural networks. In International Conference on Learning Representations, 2018.
D. Dera, N. Bouaynaya, G. Rasool, R. Shterenberg, and H. Pathallah-Shaykh. Premium-enn: Propagating uncertainty towards robust convolutional neural networks. IEEE Transactions on Signal Processing, PP: 1-1, 07 2021. doi: 20.1099/SSP.0021.30000
I. Elishakoff and S. H. Crandall. Sixty years of stochastic lincarization technique. Meccanica, 52(():299-305, Jan 2017. ISSN 1572-9648. doi: 10.1007/s11012-000 x09900000000000
A. K. Foong, Y. Li, M. Hernández-Lobato, and R. E. Turner. 'in-between' uncertainty bayesian neural networks. In ICML Workshop on Uncertainly and Robustness in Deep Learning, 2019. arXiv:1906.11 11537.
B. J. Frey and G. F. Hinton. Variational learning in nonlinear gaussian belief net Neural Compulation, 11(1):193-213, Jan 1999.
Fubini. Sugli integrali multipli. Atti della Reale Accademia dei Lincei, 16:608 16:608-614, 1907.
Y. Gal and Z. Ghahramani. Dropout as bayesian approximation: Representing model uncertainty in deep learning. In Proceedings of The 33rd International Conference on Machine Learning, volume 48 of Proceedings of Machine Learning Research, pages PMLR.
New York, New York, USA, 20 22 Jun 2016.
M. S. Gandhi, K. Lee, Y. Pan, and E. A. Theodorou. Propagating uncertainty through the tanh function with application to reservoir computing. arXiv preprint arXiv: 1806. 09431, 2018.
Gao, X. Huo, H. Liu, and Gao. Wide neural networks as gaussian processes: Lessons from deep equilibrium models. In Seirth-seventh Conference on Neural Information Processing Systems, 2023.
J. Gast and S. Roth. Lightweight probabilistic deep networks. In Proceedings of the IEEE Conference on Compuler Vision and Pallern Recognition (CVPR), June 2018.
J.-A. Goulet, L.-H. Nguyen, and S. Amiri. Tractable approximate gaussian inference for bayesian neural networks. Journal of Machine Learning Research, 2021.
13

---

Published in Transactions on Machine Learning Research (06/2025)
I. S. Gradshteyn and I. M. Ryzhik. Table of Integrals, Series, and Products. Academic Press, 8th edition edition, 2015.
A. Graves. Practical variational inference for neural networks. In Advances in Neural Information Processing Systems, volume 24, pages 2348 2356, 2011.
K. He, X. Zhang, S. Ren, and J. Sun. Deep residual learning for image recognition. In Proceedings of the IEEE Conference Computer Vision and Pattern Recognition (CVPR), June 2016.
D. Hendrycks and T. Dietterich. Benchmarking neural network robustness to common corruptions and perturbations. In International Conference on Learning Representations, 2019.
J. M. Hernández-Lobato and Adams. Probabilistie backpropagation for scalable learning of bayesian neural networks. In Proceedings of the 32nd International Conference on Machine Learning (ICML), pages 1861-1869, 2015.
G. E. Hinton and D. van Camp. Keeping neural networks simple by minimizing the description length of the weights. In Proceedings of the sirth annuaal conference on Compulational learning theory (COLT), pages 5-13. ACM, 1993.
L. Isserlis. On formula for the product-moment coefficient of any order of normal frequency distribution in any number of variables. Biomelrika, 12(1/2):134 139, 1918. ISSN 00063444, 14643510.
S. Janson. Gaussian Hilbert Spaces / Svande Janson. Cambridge Tracts in Mathematics no. 129. Cambridge University Press, Cambridge, 1997. ISBN 9780511526169.
P. Jylänki, A. Nummenmaa, and A. Vehtari. Expectation propagation for neural networks with sparsity- promoting priors. Journal of Machine Learning Research, 15:1849-1901, 2014.
I. E. Kazakov. Approximate method of statistical investigations of nonlinear systems. In Proceedings, Voenmo- Vozdushnaya Akudemia imeni N.I. Zlokkuokogo, volume 394, pages 1 52, 1954.
D. P. Kingma and M. Welling, Auto-encoding variational bayes. In International Conference on Learning Representations (ICLR), 2014.
A. Krizhevsky. Learning multiple layers of features from tiny images. Technical report, University of Toronto, 2009.
B. Lakshminarayanan, A. Pritzel, and C. Blundell. Simple and scalable predictive uncertainty estimation using deep ensembles. In I. Guyon, U. V. Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan and R. Garnett, editors, Advances in Neurol Information Processing Systems, volume 30. Curran Associates, Inc., 2017.
Y. LeCun, L. Bottou, Y. Bengio, and P. Haffner. Gradient-based learning applied to document recognition. Proceedings of the IEEE, 86(11):22782322 1998.
J. Lee, Y. Bahri, R. Novak, S. S. Schocnholz, J. Pennington, and J. Sohl-Dickstein. Deep neural networks as gaussian processes. In International Conference on Learning Representations, 2018.
D. C. MacKay. A practical bayesian framework for backpropagation networks. Neural computation, 4(3): 448 472, 1992a.
D. J. C. MacKay. The evidence framework applied to classification networks. Newral Computation, 4(5): 720-736, 1992b. doi: http://ccco.199....7..7
J. V. Michalowicz, J. M. Nichols, F. Bucholtz, and C. C. Olson. A general isserlis theorem for naxsdagaussiann random variables. Statistics & Probahilily Leteers, 81:1233-1244 2011.
R. M. Neal. Priors for infinite networks. Technical Report CRG-TR-94-1, University of Toronto, 1994.
14

---

Published in Transactions on Machine Learning Research (06/2025)
E. W. Ng and M. Geller. A table of integrals of the error functions. Journal of Research of the National Bureau of Slandards- B. Mathematical Sciences, (3B(1):11 January-March 1969. doi: 10.6028/jres.07330 /rt.r/esss.3.000... Paper 73BI-281.
L.-H. Nguyen and J.-A. Goulet. Analytically tractable hidden-states inference in bayesian neural networks. Journal of Machine Learning Research, 2022.
Y. Ovadia, E. Fertig, J. Ren, Z. Nado, D. Sculley, Nowozin, J. V. Dillon, B. Lakshminarayanan, and
J. Snock. Can you trust your model's uncertainty? evaluating predictive uncertainty under dataset shift. In Advances in Neural Information Processing Systems, volume 32, 2019.
A. Pavone, A. Merlo, S. Kwak, and J. Svensson. Machine learning and bayesian inference in nuclear fusion research: an overview. Plasma Physics and Controlled Fusion, 65(5), 2023. doi:
6587/acc60f.
F. Petersen, A. A. Mishra, H. Kuehne, C. Borgelt, O. Deussen, and M. Yurochkin. Uncertainty quantification via stable distribution propagation. In The Twelfth International Conference on Learning Representations, 2024.
A. Shekhovtsov and B. Flach. Feod-forward propagation in probabilistic neural networks with categorical and max layers. In International Conference on Learning Representations (ICLR), 2019.
J. Spiegelhalter and Lauritzen. Sequential updating conditional probabilities on directed graphical structures. Networks, 20(5):559-660 605, 1990. doi: ww.1002/net.3230222
S. Särkkä. Bayesian Fillering and Smoothing. Cambridge University Press, Cambridge, UK, 2013. ISBN 978-1-107-03065-7.
H. Wang, X. Shi, and D. Y. Yeung. Natural- parameter networks: A class of probabilistic neural networks. In Neural Information Processing Systems, pages 118-126, 2016.
C. Williams. Computing with infinite networks. In Advances in Neural Information Processing Systems, volume 9, 1996.
A. Wu, S. Nowozin, Meeds, R. E. Turner, J. M. Heráández-Lobato, and A. L. Gaunt. Deterministic variational inference for robust bayesian neural networks. International Conference on Learning Representations, 2019.
H. Xiao, K. Rasul, and R. Vollgraf. Fashion-mnist: novel image dataset for benchmarking machine learning algorithms, 2017.
S. Zhang, M. Singh, D. Menolascino, and S. Ching. Estimating uncertainty from feed-forward network based sensing using quasi-linear approximation. Neural Networks, 188:107376, 2025. ISSN 0893-6080. doi: 10.1016/j 2t2p
15

---

Published in Transactions on Machine Learning Research (06/2025)
A Appendix
A.1 Derivation of mean, covariance, and cross-covariance propagated through a univariate nonlinear function
Revisit the definitions of notations.
U multivariate Gaussian input ~ ~,
W weight matrix (constant)
bias vector (constant)
X W b ~
i W uu+b
Σ W EW
= X ~ (0, (0,
6² diag {
{7 In positive scaling factor set obtained through numerical optimizations
All oprations are element-wise, including the Hadamard product and Hadamard exponentiation. However, for notationaly simplicity, all product, division, and power operations are element-wise starting from A.2
In the context machine learning, all non-linearities are applied element-wise they are univarite. Thus, the off-diagonal entries of their Hessian matrices (second order partial derivatives) are zero, and it has similar effect on higher order partial derivatives. This makes PTPE on multivariate input casier to write down.
Given smooth nonlinear function f(-) of univariate random variables, define vector operation f(X) (f(X)), f(X)), "(.)).))) We define approximation g(-), which stochastically expands f(-) under the Taylor scheme, such that f()) and g(-) have approxinately the same first and second moment. This expansion uses i.i.d. surrogate polynomials, {I, 2 (3) ... and such choice reduces computational complexity of covariance. Notably, the i..d. assumption is not a concession, but rather yields more accurate estimates than the non-i.idd altcrnative within our framework.
E V (X) g (X) = (X))
E 2
(E
(1))
=2 (2)
E (2)
1!
and denote
Ao E (X)]
E Vx. (X)
A
1!
E V02 (X)
4
0 f (X)
A3
3!
such that
C Σ
g
A,
08
E
sero mean
16

---

Published in Transactions on Machine Learning Research (06/2025)
A.1.1 Mean
Ao is simply the mean of the output. In the later sections we show this value can either be analytically solved or approximated using similarly behaving nonlinear functions.
A.1.2
Covariance
For clarity of reading, we omit the subscript of E, but will revisit the independence of polynomial basis. The covariance function of f(X) with S-th order expansion
(= =
Cov (o(X)) = E Σ A.
-
Σ A. o
s=1
(=% -E (=% - '
>
(A.A))or oE
(A.A!)
( 1E
which is an x n matrix, and the {i.j]-th entry is
Ar. E Foo -E [Eto,a
[
Since {1)) 333 becomes
are independent, the off-diagonal entries (s + t) are zero, then the {i.j]-th entry
8
Σ
Ans
E
(i))) - E E E
Bichd)) ))
Rewrite in matrix form
S
E EG S - [ E [ 1") o A!
Cov (o(X)) Σ
A,
8=1
(5)
Σ A, cov(X A!
where A, and (s) are both dimensional vertical vectors. Using contral moments of normal distributions,
0
if s is odd
E
k),a
(s - 1)!! if s is even
With application of Isscrlis' thcorem Isserlis
1918),
E
Fi
Σ
II
PEP
{e.d)Ep
where c, E (i,j), and Pod is correlation. The sum is over all the pairings of the set B ii..,....,,,,.... ا
i.e. all distinct (suppose each is different. from other i's or j's) ways of partitioning B into pairs {c. d),
17

---

Published in Transactions on Machine Learning Research (06/2025)
and the product is over the pairs contained in (Janson 1997 (2s 1)!! pairs in the partition, or (2s 1)!! terms in the sum. For example, the first four terms of Eqn. 5 are
Michalowicz et al. 2011 there exists
A E A
A (2 A
A3 [6 +9 diag( Σ) o E diag( ( [ o AA3
A4 [245 + 72 diag( D0 diag Σ) A
A, and diag(E) (É) are both n dimensional vertical vector. With this result, to find the covariance of the output of nonlinear layer, assuming the imput follows a multi-variate normal distribution, one just needs to derive the factors of Taylor polynomial, Ao: A1, A2, etc., for the nonlinearity.
A.1.3 Cross-covariance
Let Y = (Y1, ... Yn)TT be a Gaussian random vector that is cross-correlated to X, and Y- -Y If X undergoes non-lincar transformation via function f()),
Z f(X
The cross-covariance between Y and Z can be written
Σyz E Ω Σ A.
(=6
- B -B
-F()))
( (
which is an matrix, and the {i,j} j-tt entry is
Σ Als E
E [2, (io..
E [S4]] is zero by definition. The terms with product of odd mumber of Gaussian random variables are zero by Isserlis' theorem. It can be simplified as
Exz(i,j) =
Ar, E
0,
t=l, odd
S
Ezy(i,j) =
Σ
A
8=1,8 is odd
Rewrite in matrix form
s
Σyz =
Σ
A o
E
Ω
ot
Σ
A o Cov
(1)
t=1,t odd
1=1, is odd
(6)
S Σzy Σ
A, o E
05
Ω
Σ
A, o cov(X Y)
#= is odd
#=
and the expected value can be calculated using Isserlis' theorem mentioned above. Note that this term is nonzero only if and s are odd, so the first two terms are
Σyz AT Evx + +3AJ o Σyx o diag Σx)
Ezy A Σxy + 3Ag Exy diag (iag(Ex)
An addition (e.g. residual) layer outputs the summation of activation of (or more) layers, Y and Z. Thus the covariance of Y + Z is the sum of their covariance and cross-covariance.
Σ(Y Z) Ex + Eyz + Exy Σz
18

---

Published in Transactions on Machine Learning Research (06/2025)
A.2 Tanh layers f
We use linear combination of independent error functions with different scaling factors to approximate tanh function. In our experiments, we choose set of four scaling parameters, {0.5583, 0.8596, 0.8596, 1.2612), using fmincon in MATLAB. In practice, one can add more terms for even higher accuracy without losing efficiency (depending on the computing resources), because the extra terms can be easily paralleled. We define variance term considering the relation between error function and Gaussian edf, such that
Lanh(X)
R
erf (7₂
P j=1
(7)
1
63
2n
Thus, the factors of the pscudo Taylor polynomials are
p
X
Ag
erf
2263
X
A
E
V
erf
263
j=1
1 p
X
A
E
erf
2!
p j=1
v 26²
Since all the operations in Ao, A1, A2. are element-wise, we only show the derivation for univariate case for notational simplicity in the following sections
A.2.1
Find Ao
X
erf
erf
4
dx
266
262
This is a known integral
Ng and Geller
1969
erf
2202 + 263
We define
87 82 + 63
(8)
Thus,
A
Ao
erf
(9)
20
For notational simplicity, all the product, division, and power operations that appear and after this section wlesent-wise.
19

---

Published in Transactions on Machine Learning Research (06/2025)
The usage of error function instead of Gaussian edf may give Ag very distinctive form from those of the other factors. The reasons behind are purely out of considerations of numerical computing: calculating Gaussian edf is computationally demanding, while the approximation algorithm of the error function is available Cody 19999..
A.2.2 Find A1
Notice that
① Ox
X
0
2
erf
exp( exp(-12)d))
/269
dr
o
/
by Leibniz integral rule
x22
exp
π
263
26}
V
2
4p
where is the standard normal pdf. With the identity that the convolution of two Gaussians is still a Gaussian. (Bromiley (2003)
X
20
2
4
E
erf
6
dr
0x
2663
6
6
j
4
Ø2
+
62
82 +
6²
Therefore,
2 A
(10)
and each term of the summation is a Gaussian function written in its standardized form.
A.2.3 Find A2 and beyond
In previous section, we show that the first derivative of the error function a Gaussian, thus the expected value of which is the convolution of two Gaussians. Similarly, we can obtain A2, As, etc. by convolving the second, third, and higher order Gaussian derivatives with another Gaussian.
Gaussian derivatives can be represented by Hermite polynomial H2(z).
d
1
T
H.
4
dx
2
20
0
For examples,
20

---

Published in Transactions on Machine Learning Research (06/2025)
Ho(æ) 1
H (z) 2xx
(x) =4x²
Ha(z) = &2 12x
There are implemented functions for this from various scientific computing tools, such as hermiteH() from MATLAB and siipy.spocial.ermittee from SciPy.
Hence,
0
X
E
erf
Ox
1263
0r
erf
p(x)drx
dr
V
26
00
0
1
1
x
4
4
dr
Or
6
1
p
-2
H
dx
26
26
26
6
-1
1
=2
H
dx
262
8j
263
jj
62 p
where p
02
The convolution of Hermite polynomial and a Gaussian pdf known
0
integral Gradshteyn and Ryzhik 2015
8-1
///
/203 263
1
20
H
v
263
202
(1-2022 2A7
-1-2 1
1-1
-1 1
2
H
263
20
Therefore, we can write the formula of A, for
s-1
1
A,
2
H.-1
4
(11)
s!
20²
283
21

---

Published in Transactions on Machine Learning Research (06/2025)
To give few examples,
A2
1
Σ
P
2
1
4
21 p
0
6
6g
1
P
1
1
11
A3
4
31
P
0
(12)
1=1
1
f
+ 3'0d
1
j
A4 = 4! p
2
06
0
4
Note that we will reuse this relation in the following section
1
1
x
-1
1
Hg
4
4
dr
H
(13)
20²
26
6
v 20²
26
,
A.3 Sigmoid layers
We can apply the same framework on sigmoid layers, with modifications
sigmoid(z) = 1 exp (-x)
erf (1zz)
2p ==1
Using fmincon in MATLAB, find set of (0.2791, 0.4298, 0.4298, 0.6306) Then the first four factors of the Taylor polynomials are listed below. Ag is represented in complementary error function erfc to avoid subtractive cancellation that leads to inaccuracy in the tails. Note that except for Ao, all Ag of sigmoid layers are just 1/2 of those of tanh layers.
j
Ao
erfc
P j=1
26
A1
6 4
j=1
A
A
p 1=1
03
8
6,)
(14)
1
172
63
A3
3 J=1
1
+ 3nd
1
i1
A4 =
x
1,
4
A.4 Softplus layers
The derivation of psoudo-Taylo polynomials for a softplus layer is related to that for a sigmoid layer, since the derivative of the softplus function is the sigmoid function with scaling factor B, and the latter can be approximated with linear combination of Gaussian edf (or error functions like did in the previous
22

---

Published in Transactions on Machine Learning Research (06/2025)
section). We have
softplus(x)
log (1 + e
B
Then we use the approximation of sum of independent standard Gaussian edf +
0
1
1
softplus(x)
(15)
d
1 B-
P
where we re-define
1
63
(16)
32
Note that of changes definition and should not be confused with that in the tanh and sigmoid sections.
A.4.1 Find Ao
First we apply substitution of variables E, then
softplus(x) = softplus {) - log (1 el(tt))
0 0 a
Notice that Ox
since jú is constant, then
Ao (oftplus(A,)) E) p(E) dE
p(E)dg 1" a
a So.t dc o
by Fubini's theorem
Fuubini 1907
dc softplus((,)) p(E)dE
H
i
(S+
1 4
{
dc
dg
p
-
0j
using the known Gaussian integral identity
_ D b))((x) +
VI+a
p
dc
P
1=1
+
6jp
0,,
θj
Σ
erfc
+ +
11 t;
2
/26
j=1
Or, with simplification
1 P
Au = Aijt +
1
(17)
P
23

---

Published in Transactions on Machine Learning Research (06/2025)
A.4.2 Find A1
Since the first derivative of the softplus function is just a sigmoid function with scaling factor B, we can immediately write A1 using previous results
P A1
ji [63
erfc
(18)
j=1
A.4.3 Find A2 and beyond
In previous section, we find that V softplus(1) is approximately a Gaussian edf. Subscquently, V softplus(2) is approximately a Gaussian. Since Gaussian function is infinitely differentiable, all A, > 2) can be found using Gaussian derivatives, which be represented by Hermite polynomial (x(x) introduced above.
1 A,
0 softplus(x)
E s! 2r
yx-2
p(a) dx
8 p
dx-~2
6,
40
6,
x-2
1
1
H 2-
4
dx
88 p j=l
26
26²
6,
we solved this integral in tanh section
8-2
1
H₂-2
4
8s p
20²
20²
6
To summarize, As 2 2) can be expressed as
5-2
A.(s 2 2)
H₂-2
(19)
s!
20²
288
For examples,
1
A
4
21 p
8j
j=1
1
71
Ag
31 p
67
0
(20)
9
2
02
A
4
3,
A.5 ReLU, Leaky ReLU, and Piece-wise Linear layers
Since ReLU function is only first-order differentiable (x > 0), we cannot do PTPE directly. However, given its relation to softplus function,
lim 3 log (1 + elr) max(0,x] 8+000
24

---

Published in Transactions on Machine Learning Research (06/2025)
can reuse the results for softplus layers by applying the limit
lim 8
8j =
and
lim 03 = 8²
Therefore,
Ap AI + ōg
A1 2
erfc
V272
11
A
4
2 o
(21)
1
A3
4
31
2
A1
=
4!
0
and for we have the general form of
1
A, 2 2)
H,-2
(22)
8!
20J
20²
203
On the other hand, leaky ReLU be considered superposition of two ReLU functions consider leaky ReLU with negative slope of 0
LeakyReLU(z 0) ReLU(x) 0 ReLU(-x)
(23)
which can also be written
lim softplus(z) softplus(-z) 4000
Therefore,
P
j1
Ao lim
+ t
-
+ 8
θ,
0ī + (1 0)
+ 0.
To find the expected value the derivative LeakyReLU, first we find the derivative
a
dr LeakyReLU(r: 0) lim softplus(z)
softplus(-(x))
80000 dr
dr
lim
$ + e- B(w)
+ (B))
+ 0Φ
p -=1
6
1
lim
0 +
f
3
j=1
25

---

Published in Transactions on Machine Learning Research (06/2025)
Then write A for LeakyReLU
1
A1 lim
0
+
4
dx
3-000
P
4
I
= 0 +
lim
dr
0
j=1
0
0 + lim
Σ
+
$+000
8
0 + (1
Rewrite in complementary error function
1-8
A
A = + erfc
(24)
2
Note that we can also rewrite Ag using the result of A₁ to improve computational efficiency.
it
Ao Ai (1 0)āp
(25)
à
Note that starting from the second order, the derivative of LeakyReLU is just that of ReLU scaled by - 8. Therefore,
1 1
01 1
A2
4
1 0 1
1
A3
31 03
(26)
1
0
7
A4
5
and for s 2 2, we have the general form of
1-
8-2
A, 2 2)
H.
x
(27)
v 22
V252
Similarly, any piece-wise linear activation function can be described a combination of ReLU functions with different scaling, shifting, and/or mirroring. Thus, their pscudo Taylor cocfficients can be found using the same methodology.
A.6 GELU layers
GELU (Gaussian Error Linear Unit) is defined the product of input and a standard Gaussian edf
GELU(>) = rr (x))
and we can write the derivatives (with order s 2 1) of GELU as
8--1
8
GELU(x)
d---I (xx) (z) + 0
((x)
d
26

---

Published in Transactions on Machine Learning Research (06/2025)
A.6.1 Find Ag
Au = E (ELU
r 11 dr
( 1
1
5 ds
(++ g)
50 + +3)
1
5
dp Sp + I / E 4 ( + {)
{
{
de dg
dc + P
5
-
V1 + 52 10
V1+
v1 02 4
V1 0²
1 dp
-00
/ [T
V1
VI+
+
We re-define 2²
=
1+ ²²
(28)
and re-write the result with complementary error function
11 erfc
(-
0²
Ao
6
(29)
2
220
A.6.2
Find A
0
A =
GELU(x)
Or
$
40l
)dx + S rup(x) _
1
======= )d
--
using results of previous section
+
4
02
Therefore,
j
A erfc
+
(30)
A.6.3 Find A, and beyond
Higher order coefficients (A,(s 2 2)) all consist of two parts: (i) a term of expected value of a Gaussian derivative, (ii) term of expected value of the product of x and Gaussian derivative. We have already found general form of the first term in the tanh scction
dz pp())) = (
a--2
i
E
8
1262
H
V262
6
7
solve the second part, we need to use the Hermite polynomial recurrence relation:
1 H.(x)+ + H Hx)
x H (x)
27

---

Published in Transactions on Machine Learning Research (06/2025)
0
0--1
E
(x)
Or
1
11
2
H
9 x
6
dr
2
8-1
j
2
H
8 (x)
4
dx
V2
V2
V2
0
1
I
2
H
+
1)H, 2-2
40
ā 4
dr
22
V2
v2
T
x 1
H
4
22
V2
00
2-2
x j
(s
1)
/2
H.
-
2
8 a
dr
by equation 13
-1
-1 V28²
5-2
Φ 4
H.
(s
1)
H
H
V202
20
V202
Sum the two integral together, we get. the general form of A,(s 2 2)
1 s!
5-2
ji
Az(s 2 2) =
H
H
(31)
20
V2
2272
/2
For examples,
1
11
1
"
A2
+
82
84
8
1
P1
37
3
1
Ag
4
3!
22
01
(32)
7:2
6772
1
A4 =
4
4!
02
06
08
A.7 SiLU layers
SiLU (Sigmoid Lincar Unit), equivalent to Swish when / = 1, is defined as the product of input and sigmoid function
Si)U Sigmoid(z)
In the previous section, approximate Sigmoid function with error functions that we reuse derivations from the Tanh section. Here we approximate Sigmoid function with Gaussian edf's in order to reuse derivations from the GELU section. With as a numerically optimized scalar vector, let
2 1 jE {1. ...
28

---

Published in Transactions on Machine Learning Research (06/2025)
Then, approximate SiLU as and we can write the derivatives (with order 2 1) of SiLU
SiLU(r)
= -
+
j=1
8%h SiLU(x)=
8 1
I
O
+
Ou
P due-1
Oj
Ou
The rest of the derivation is very similar to that of GELU, so we only list the final results. With
03= 0² + 87
i
á
0²
A
Ao
erfc
2
203
0j
1=1
1
A
crfc
4
p
2
v 20²
63
p
62
A2 2!
1
1
(33)
1=1
1
P
3/6
T 6
1
A3
P p
02
6
6
4
6,
363
6122
p
A4 =
4
4! p
02
01
69
8
3
and the general form of As(s 2 2) is
2
j
j
A, = s! p
H
H,
(34)
26
V282
2022
V20²
à
29

---

Table 4: First four coefficients of the polynomials for seven commonly used nonlincarities. For notational simplicity, all the product, division, and power operations are elesent-wise. The method for determining 7/ outlined in Table !
Published in Transaction on Machine Learning Research (06 (
Ao
A1
A2
As
63 = + +
Tanh
(2F -) 1)
P
2C
1
26 n
P
2C
à
83 = 21
p j=1
2 j=1
03
3! p
j=1
j=1
1 Σ
1
P
1
1
Sigmoid
F
63 273
p j=1
p
31 p
6
j=1
9 Σ j=1
1/2
p
1 Σ4
9 Σ
1
Softplus
Fî+ + CG3
63 277
p
ド=
2p j=1
31 j=1
63
30
1 31
C 11
ReLU
F + L8²
F
c
-
6²=0
2
02
1- 0
0
LeakyReLU (0) 0 + (1- 0) (Fū+ + CO")
0+ (1-11
C
C G²
0==0
31
F+ C u
1
1
11
3
GELU
Fū + L02
C
1+
C
8² =1
02
2
02
04
3
02
0
06
SiLU
p Fū L0²
1 p F
L 1OO
1
p
c
1+
1I'03
1 p
Σ
C
j
63
1
P j=1
p
0
0
31 p
6
の )
273
j=1
j=1
1
A
where L 40
and
F= =
$
erfc
6,
θ,
2
/207

---

Published in Transactions on Machine Learning Research (06/2025)
A.8 Emperical distributions of ResNets show Gaussianity
Figure 7: Empirical distributions of all units before the final softmax layer of the resnet 13 (ReLU).
A.9 Approximation accuracy on other non-linearity
31

---

Published in Transactions on Machine Learning Research (06/2025)
Variance Gain
Variance Gain 2nd Order
Variance Gain 3rd Order
Variance Gain 4th Order
Mean
1st Order
1
100
Tanh
0
1
1
100
Sigmoid
0.5
10
0
Softplus
4
0.5
2
0
ReLU
4
0.5
2
LeakyReLU(0.5)
4
2
0.5
0
-2
1.5
GELU
2
0.5
1.5
6
4
2
0.5
0
-5 center of input dist.
center of input dist.
center of input dist.
center of input dist.
center of input dist.
Figure 8: The meaning of different colors and line styles is the same as Fig. The higher the input variance, the more significant is the benefit. of using higher order Stochastic Taylor expansion.
32

---

Published in Transactions on Machine Learning Research (06/2025)
A.10 Pseudo Code
Algorithm Propagating multi-variate Gaussian distribution through pretrained ResNet
μ Input mcan
Σ <- Inpul covariance
Pres + Storage for mean of the output of residual layer
Eres < Storage for covariance of the output residual layer
Deross + Storage for cross-covariance between the two input of the residual layer for layer neural network do if layer is linear then if layer is addition (residual) then
+ + Pres
+ + Mres + Eeross + XT
Ceross empty press Eres: Eeross
else
find effective weight W and bias b a t WT bb W T- EW
Meross t- WT if residual connection starts from here then
Pres 4 Eeross +
end if end if
else
μ. D, Meross PTPE[monlinearity μt, E, Meoos)
if residual connection starts from here then
Ures f Eves, Epross t Σ
end end if
end for
Houtput 1
Doutput f- E
33

---

Published in Transactions on Machine Learning Research (06/2025)
A.11
Additional results comparing PTPE and other expectation propagation methods.
ResNet13 relu
108
10
1010
Wasserstaind distancs
Frey (1999)
10"
PTPE 1st
Frobenius norm
105
PTPE Srd
L2 norm
100
102
109
10
10:2
10
10.5
10
10
100
10
10
10
100
101
10
10
ResNet33 relu
108
102
1010
Wassersteing distance
10
Frabenius norm
105
L2 norm
10
10
100
100
10
10
10.5
10
10
109
10'
102
10
100
102
ResNet65 relu
108
10
10
Wassersteind distance
104
Frabenius norm
105
12 norm
100
102
109
109
10
10
10:5
10
100
10
2
10
10
10
10a
10
Input Variance
Input Variance
Input Variance
Figure 9: Similar to Figure by Frey and Hinton (1999) performed poorly as it disregards correlation, critical factor in networks with overlapping convolutional kernels and residual connections.
except that the models employ ReLU the nonlinearity. The method proposed
34

---

Published in Transactions on Machine Learning Research (06/2025)
ResNet13 gelu
102
102
10a
Wassersteind distance
Frobenius norm
L2 norm
109
10
109
PTPE 1s1 PTPE 3rd
10
10
10
10
101
102
10a
10
101
102
10
100
10
10
10a
ResNet33 gelu
10a
10
102
Wasserstein2 distance
Frobenius norm
L2 norm
100
10
109
10
10
2
10
10P
10
10
103
10
10
102
103
10
10
102
10
ResNet65 gelu
10
10a
Wasserstein2 distance
Frobenius norm
L2 norm
100
109
10
10
1
10
10
1
103
100
10
10a
103
100
10
102
103
Input Variance
Input Variance
Input Variance
Figure 10: Similar to Figure 3 except that the models employ GELU as the nonlinearity.
35

---

Published in Transactions on Machine Learning Research (06/2025)
A.12 Additional results the selection of the number of scaling factors in the reformulation of the tanh function
x
scaling factors
Error by reformulation
scaling factore
6 scaling factors
scaling factors
10 scaling factors
-1
-10
88
-6
-4
-2
2
4
8
10
Input location
10
Optimization
1
10
2
3
7
8
9 10
Number of scaling factors
Figure 11: (Top) The error approximating the tanh function using linear combination error functions with different scaling factors analyzed. We selected four terms, as this configuration yields maximum approximation error the scale of 10- Importantly, increasing the number of terms enhances accuracy while only increasing computational complexity linearly, making this approach both flexible and computationally efficient. (Bottom) The optimization loss for determining the scaling factors was computed using fmincon in MATLAB. This loss defined as the root mean square error (RMSE) of the approximation.
36

---

Published in Transactions on Machine Learning Research (06/2025)
ResNet13 tanh
x10-33
0.02
0.02
PTPE 1st PTPE
Improv. in W2 distanco
Improv. in L2 norm
0.01
PTPE 3rd
0.01
Improv. Frob.
-0.01
-0.11
109
10
102
103
10
101
10
103
109
101
10
10
ResNet33 tanh
x10
0.02
0.02
Improv. W2 distance
Improv. Frob. norm
Improv. in L2 nor
0.01
0.01
2
08
-0.01
0.01
10
10
102
103
10
10
10
10
109
101
10a
10
ResNet65 tanh
x10-3
0.02
0.02
Improv. in W2 distance
Improv. Frob. norm
Improv. in L2 norm
0.01
0.01
2
-0.01
4
10
10 Input Variance
10
103
10
10?
103
10
102
10
Input Variance
Input Variance
Figure 12: Improvements in the three metrics Wasserstein-22 distance, Euclidean distance, and Frobenius norm when using eight scaling factors (7;) to approximate the tanh function, compared to using four, are reported. Positive values indicate improvement, while negative values denote deterioration. Notably, there is no clear dichotomy in predictive accuracy, suggesting that the approximation error of our method is not a dominant factor. Alternatively, Monte Carlo sampling with 10t data points insufficient to accurately estimate the true covariance of a 10-dimensional multivariate distribution. As a result, the reference moments may deviate from the true values. We argue This finding supports the justification for using four scaling factors rather than larger number.
37

---

Published in Transactions on Machine Learning Research (06/2025)
A.13 Comparison with Causs-Hermice Quadrature
Gauss-Hermice quadrature (GHQ) is a numerical integration method used to approximate integrals of the form:
[ f(z)e dx =0 (x;(
where are scalar weights, and n is the number of quadrature points (also referred to as sigma points). In this context, f(-) can represent a nonlinearity or its square, corresponding to the mean and covariance integrals. While infinitely many linear combinations of f(x) can approximate this integral, GHQ selects quadrature points and weights based on Hermite polynomials. These polynomials are orthogonal, making GHQ computationally elficient. This method also serves the foundation of the cubature Kalman filter (CKF).
We compare the predictive accuracy of PTPE and GHQ across the first four orders (Fig. 13 to 16 using
absolute error as the evaluation metric against a reference mean and variance obtained from 10' Monte Carlo
samples. Since Monte Carlo estimates with 10" samples typically fluctuate on the order of 10 -4 absolute errors of magnitude Σ 10-4 are considered negligible.
PTPE outperforms GHQ in the first three orders. At the fourth order, PTPE provides evidently more accurate estimates than GHQ when the input variance is 1 higher. Notably, 4th-order PTPE surpasses 4th-order GHQ on ReLU and LeakyReLU even at variance of 0.1.
To determine how many GHQ orders are needed to surpass 4th-order PTPE, we compare 10th- 20th-; 30th- and 40th-order GHQ to 4th-order PTPE (Figures 17 to 20 There is no clear threshold at which GHQ consistently outperforms 4th-order PTPE, as performance depends on the type of noulinearity. Among the seven types, PTPE is particularly effective for ReLU and LeakyReLU. However, we observe a general trend: higher-order GHQ is required to match PTPE, particularly at higher input variances.
In summary, although GHQ elficiently selects sampling points, it requires an increasing number of points for accurate estimation at higher input variances. More critically, the number of GHQ sampling points grows rapidly with dimensionality, making high-dimensional integration computationally expensive. In contrast, PTPE's complexity is at most O(n?), making it more scalable. However, GHQ is expected to outperform PTPE in narrow but deep neural networks. Unlike PTPE, GHQ does not require layer-wise propagation of moments. Furthermore, in narrow networks, the Gaussianity assumption of hidden layers often breaks down, making it difficult for PTPE to provide accurate estimates.
38

---

Published in Transactions on Machine Learning Research (06/2025)
Input Variance 0.1
Abs Err of Mean Order
Abs Err of Var Gain Order
Abs Err of Moan
Abs Err of Var Gain
Abs Err of Mean
Abs Err of Var Gain
Abs Err Mean 4th Order
Abs Err of Var Gain 4th Order
0.04
2nd Order
0.1
2nd
10
3rd Order
3rd Order
PTPE GHQ
0.01
Tanh
0.02
1
0.5
0.05
0.005
x10+
x10
Sigmold
0.05
10
10
x10
10
x10
Boftokus
0.1
0.01
11
0.1
0.01
0.05
ReUU
0.5
0.1
0.005
0.005
39
LeskyRReLUjg33
0.05
0.01
x100-
0.06
$=100
0.02
0.5
0.005
0.01
10
10
9
0.04
0.05
in
GELU
0.02
0.5
0.5
0.5
4
0.02
0.5
0.01
5
%
center Input dist.
center Input dist.
center Input dist.
center input dist.
center input dist.
center of input dist.
center of input dist.
center of input dist.
Figure 13: Comparing the absolute approximation error of PTPE and CHQ across the first four orders. The odd number of columns show the error of mean, and even number of columns show the error of variance, the smaller the better. In this panel, all input variance is 0.1.

---

Published in Transactions on Machine Learning Research (06 (26))
Input Variance 1
Abs Err of Mean Order
Abs Eir Ver Gain
Abe Eir of Mean
Abs Eir Var Gain
Abe Eir of Mean
Abs Err of Var Gain
Abs Err Mean
Aba Err of Gain 4th Order
Order
2nd Order
2nd
3rd Order
3rd Order
4th Order
04
0.1
0.
0.04
0.1
0.02
0.05
PTPE GHO
Tanh
fit
0.2
0.05
0.1
0.02
0.05
0.01
x10
+10
3
Sigmold
002
0,005
11
0.5
0.5
x100
x.10
1100
x10
Softpuus
0.05
0.5
0.5
0.6
0.4
0.1
0.2
0.1
0.05
0.05
0.5
0.1
0.05
40
02
0.05
0.04
0.05
0.02
0.1
0.5
0.05
0.02
All
0.01
x.10
10-3
0.-
0.02
0.05
0.02
0.5
0.1
0.01
210
9
9
0.2
0.1
0.02
0.1
0.5
0.05
0.01
0.5
0
0-
0-
0-
.8
center input dist.
conter Input dist.
conter Input dist.
center Input dist.
center Input dist.
center of Input dist.
center of input dist.
center of input dist.
Figure 14: Comparing the absolute approximation error of PTPE and CHQ across the first four orders. The odd number of columns show the error of mean, and even number of columns show the error of variance, the smaller the better. In this panel, all input variance is 1.

---

Published in Transactions on Machine Learning Research (06/ (
Input Variance 10
Als Err of Mean Order
Abs Eir Ver Gain Order
Aba Eir of Mean
Abs Eir Var Gain
Abs Eir of Mean
Abs Err of Var Gain 3rd Order
Abs Eir Mean
Aba Err of Var Gain
2nd Order
2nd
3rd Order
4th Order
4th Order
0.4
0.06
Tanh
PTPE GHO
0.2
0.02
0.2
0.01
110
Stgmma
0.05
0.02
0.00
0.2
0.05
0.04
0.1
0.05
0.5
0.01
0.1
0.2
02
0.5
022
0.
0.05
19
0.1
0.05
0.1
0.1
0.1
0.5
0.05
0.4
0.2
0.1
02
0.2
0.5
0.2
0.1
0.1
0.05
0.1
17
0.05
02
02
0.1
0.05
0.5
02
0.1
0.1
0.05
0.05
05
center input dist.
conter Input dist.
center Input dist.
center Input dist.
center Input dist.
center of input dist.
center of Input dist.
center of linuut dist.
Figure 15: Comparing the absolute approximation error of PTPE and CHQ across the first four orders. The odd number of columns show the error of mean, and even number of columns show the error of variance, the smaller the better. In this panel, all input variance is 10.

---

Published in Transaction on Machine Learning Research (06 (06/2025)
Input Variance 100
Abs Err of Mean
Abs Eir Var Gain Order
Abe Eir Mean
Abs Err Var Gain 2nd
Abs Err Mean
Aba Err Var Gain
Abs Eir Mean
Aba Err of Var Gain Order
2nd Order
3rd Order
10
3rd Order
4th Order
04
0.5
0.4
PTPE CHO
Taah
0.2
0.2
Sigmolu
0.1
0.1
0.1
0.8
0.
0.1
0.25
0.04
0.05
05
0.1
0.5
0.05
0.00
0.5
42
05
0.5
0.5
0.04
022
0.02
0.01
0.5
0.1
0.5
0.05
0.5
0-
0.05
0.5
0.5
0-
00
input
conter Input
center Input
Input dist.
dist.
Input
Input linput
Figure 16: Comparing the absolute approximation error of PТРE and CHQ across the first four orders. The odd number of columns show the error of mean, and even number of columns show the error of variance, the smaller the better. In this panel, all input variance is 100.

---

Published in Transactionn on Machine Learning Research (06 (
Input Variance 0.1
Abs Err of Mean 4th PTPE 10th GHQ
Abs Err of Var Gain 4th PTPE 10th GHQ
Abs Err of Mean 4th PTPE 20th GHQ
Abs Err of Var Gain 4th PTPE 20th GHQ
Abe Err of Mean 4th PTPE 30th GHQ 10
Abs Err of Var Gain 4th PTPE 30th GHQ 10
Abs Err Mean 4th PTPE 40th GHQ
Abs Err of Var Gain
4th PTPE 40th GHQ 100
x10
10
PTPE GHQ
Tanh
x110
10
10
10
x0
10
10
10
Sigmoid
Softplus
x10
2
x
210
ReLU
0.01
0.01
0.01
0.01
0.006
0.005
LedkyRsLUj033
x11
0.01
10
-10
-10
110
0.005
0.5
x10
111
10
10
210
OEL
x10
10
10
10
=10
SIU
0.5
0.5
0.5
0.5
center of input dist.
center of input dist.
center of input dist.
center of Input dist.
center of input dist.
center of input dist.
center of input dist.
center of input dist.
Comparing the absolute approximation error of 4th-order PTPE and 10th-, 20th- 30th- and 40th-order GHQ. The odd nunber of columns show the error of mean, and even number of columns show the error of variance, the smaller the better. In this panel, all input variance is 0.1.
Figure 17:

---

Published in Transactionn on Machine Learning Research (06/ (06/2025)
Input Variance 1
Abs Err of Mean 4th PTPE 10th GHQ
Abs Err of Var Gain 4th PTPE 10th GHQ
Abe Err of Mean
Abs Err of Var Gain 4th PTPE 20th GHQ
Abs Err of Mean 4th PTPE 30th GHQ 10
Abs Err of Var Gain 4th PTPE 30th GHQ
Abs Err Mean PTPE 40th GHQ 10
Abs Err of Var Gain
4th PTPE 20th GHQ 10
4th PTPE 40th GHQ
-10
PTPE
Tanh
GHQ 0.02
0.02
0.02
0.02
x10
x11
10
x310
100
10
-10
Sigmoid
Soffolue
0.5
0.5
0.5
0.5
x10
0.01
0.01
0.01
44
LeskyBeLU(0.5)
.110
x100
10
x100
0.01
0.005
4/1
x10
x100
x100
GELU
0.01
0.01
0.01
0.01
0.005
0.005
x100
10
10
10
10
-10
>100
SIUU
center Input dist.
center of input dist.
center of input dist.
center of input dist.
center of input dist.
center of input dist.
center of input dist.
center of input dist.
Comparing the absolute approximation error of 4th-order PTPE and 10th-, 20th-, 30th-, and 40th-order GHQ. The odd nunber of columns show the error of mean, and even number of columns show the error of variance, the smaller the better. In this panel, all input variance is 1.
Figure 18:

---

Published in Transaction on Machine Learning Research (06 2
Input Variance 10
Abs Err of Mean 4th PTPE 10th GHQ
Abs Err Var Gain 4th PTPE 10th GHQ
Abe Err of Mean
Abs Err Var Gain 4th PTPE 20th GHQ
Abs Err of Mean 4th PTPE 30th GHQ 10
Abs Err of Var Gain PTPE 50th GHQ
Abs Err Mean PTPE 40th GHQ 10
Abs Err of Var Gain
4th PTPE 20th GHQ
PTPE 40th GHQ
0.02
0.01
0.01
0.01
Tann
0.05
PIFC GHD
0.01
005
0,028
0.5
0.000
11101
x10
10
-10
10
-10
-100
-1002
Sigmald
0.5
0.5
0.5
0.5
0.01
0.01
0.01
10
0.01
Sofplus
2
0.000
0.5
0.005
0.005
0.005
0
ReLu
0.02
0.01
0.01
12 M
LcakyRoLU(e 5t
0.05
0.02
0.01
10
0.01
10
0.02
0.01
0.005
0.005
3 Net
x110
210
0.02
A
GELU
10
2.10
2210"
0.02
0.02
0.005
001
0.5
0.01
0.5
0.11
center input dist.
conter of input dist.
center of input dist.
center of Input dist.
center of Input dist.
center of Input dist.
center of Input dist.
center of input dist.
Figure 19: Comparing the absolute approximation error of 4th-order PTPE and 10th-, 20th- 30th- and 40th-order GHQ. The odd nunber of columns show the error of mean, and even number of columns show the error of variance, the smaller the better. In this panel, all input variance is 10.

---

Published in Transactions on Machine Learning Research (06/ (06/2025)
Input Variance 100
Abs Err of Mean 4th PTPE 10th GHQ
Abs Err Var Gain 4th PTPE 10th GHQ 10
Abe Err of Mean
Abs Err of Var Gain 4th PTPE 20th GHQ 10
Abs Err of Mean
Abs Err of Var Gain PTPE 50th GHQ 10
Abs Err of Mean PTPE 40th GHQ
Abs Err of Var Gain
4th PTPE 20th GHQ
4th PTPE 30th GHQ
PTPE 40th GHQ
10
0.1
Tanb
PTFT GHD
0.05
0.06
14
1
0.05
1
05
10-4
-0
10
10
0.01
Sigmeld
0.05
0.02
0.005
1
Sonthus
0.1
0.02
0.04
0.02
0.01
0.01
0.01
0.01
0.01
0.02
0.01
0.005
0.006
0.006
0.005
BeLL
0.1
0.01
0.05
0.01
0.05
1
1/
11
46
LcakyRoLu(j.s) 5)
0.02
10
0.05
5100
0.04
20
6as
0.01
0.02
0.1
0.01
0.05
0.01
0.05
x
0.04
0.1
0.02
0.02
0.02
0.62
0.00
0.01
0.01
0.01
0
center of input dist.
center of input dist.
center of input dist.
center of input dist.
center of input dist.
center of Input dist.
conter of input dist.
center of input dist.
Figure 20: Comparing the absolute approximation error of 4th-order PTPE and 10th-, 20th- 30th- and 40th-order GHQ. The odd nunber of columns show the error of mean, and even number of columns show the error of variance, the smaller the better. In this panel, all input variance is 100.


Reviews:
[{"review_id": "zLvDBEPaja", "reviewer": "Reviewer_LDk4", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Generally, the weaknesses from the previous submission are well addressed, and the paper is overall strong.\n\nStrengths:\n\n* (major) The empirical evaluation is quite thorough, and the authors are generally quite clear about the contributions\n* (major) The idea is clearly communicated, and there are helpful comparisons with similar methods \n\nWeaknesses:\n* (minor) It remains the case that the experiments look at multiple different orders of approximation which are sometimes quite similar - in Figures 3, 8, 9. It's not clear if the differences here are random (error bars?) or not - in particular, sometimes higher order approximations perform worse. Is that predictable? The final figures in the appendix help the reader to get a bit more insight, but it would be better if the text drew these together and gave some recommendations for the reader.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This work proposes a method for propagating Gaussian inputs through a non-linearity $f$ by local polynomial approximation. The authors focus on the setting where each co-ordinate of $f$ depends only on the corresponding co-ordinate of $X$, as in the non-linearity in a neural network, so that the derivatives of $f$ are all diagonal. They then use a truncated Taylor series approximation, at each step using the expectation of the derivative, and due to the diagonality, the cost is only $𝒪︀(n^2)$ for length $n$ inputs.\n\nThey provide approximations to permit calculation of the expected derivatives of various widely used non-linearities, and explicit expressions for the expansion.\n\nExperimentally, they show that the third order expansion performs empirically much better than existing first order expansions (local linearisations), in the sense that the propagated distribution is much closer to the true distribution, at least when the input variance is fairly large.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Some minor points:\n\n* Citations in sentences (as opposed to in parenthesis) sometimes have years and sometime don't. I think they should always have years.\n* Near end of section 2: 'v.s.' -> 'versus', 'vs', or 'against'\n* Figure 3, etc: It's fairly easy to guess what the legend is supposed to say, but it's both too small and overlapping with the lines - it would be good to tidy this up. You could put the legend at the bottom (in 1-2 rows; probably easiest), or elsewhere outside of the subplots.", "broader_impact_concerns": "none", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "CPBM0IG01L", "reviewer": "Reviewer_VBVu", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "## Strengths\n* The paper tackles an interesting and challenging problem.\n* The paper provides a novel solution to this problem and approximated moments for several commonly used activation functions. These moments could be useful outside of this line of work.\n\n## Weaknesses\n* The experimental validation is skewed towards toy settings that don't necessarily reflect the potential real-world usage of the method. For example, the comparisons to existing uncertainty qualification methods are primarily based on UCI regression datasets rather than large-scale image classification tasks, and the Gaussian noise applied to the CIFAR datasets is only one of several corruptions suggested by  Hendrycks and Dietterich (2019).  \n* Deep ensembles, which are a key uncertainty quantification baseline, have not been included.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper proposed a novel method for propagating uncertainty from the input of a neural network to its output, based on a \"stochastic polynomial expansion\" of the nonlinearities in the network. The authors provide closed-form expansions for the first four polynomial coefficients for seven commonly used activation functions and provide a general methodology for computing additional coefficients for general non-linearities. The authors provide experimental evidence for the accuracy of their approximations based on the UCI regression datasets, CIFAR-10 image classification, and MNIST digit generation with a VAE.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "I reviewed a previous version of the paper and I am pleased to see that many of my concerns have been addressed to some extent.  However, as mentioned above I still feel that the experiments in this paper do not reflect realistic settings. Thus, I have several recommendations to strengthen the paper. In my opinion, this is currently a borderline submission, so while none of the following suggestions are critical on their own, making even a few of these changes would be enough for me to recommend acceptance. \n\n1.[strengthen] Add a wider range of (more realistic) uncertainty quantification benchmarks:\n    * E.g., rotated-MNIST, corrupted-CIFAR, and OOD detection tasks (see \"Can you trust your model's uncertainty? Evaluating predictive uncertainty under dataset shift\" by Ovadia et al.).\n    * In particular, provide classification accuracy for the CIFAR-100 experiment, and add some \"more realistic\" corruptions from Hendrycks and Dietterich (2019).\n2. [strengthen] Add deep ensembles as a baseline in tables 2 and 3.\n3. [strengthen] Add an ablation for comparing the number of approximating terms (*and* MC estimation with different numbers of samples) based on a realistic uncertainty quantification benchmark (rather than the toy Gaussian setting shown in Fig 11.)", "broader_impact_concerns": "None", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "gLyhyOV6hC", "reviewer": "Reviewer_WbCR", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "__Strenghts.__\n\n- This work provides comprehensive derivations for a range of activation functions, including GELU and Swish, which are commonly used in modern architectures.\n\n- The aim of ensuring tractability through polynomial approximations is well-motivated.\n\n__Weaknesses.__\n\n- There seems to be insufficient discussion on the errors introduced by the proposed approximations. For instance, activations like tanh, sigmoid, and SiLU are first approximated using erf, and then the resulting approximations are further approximated using polynomial expansions. Isn’t there any concern with this two-step approximation process?\n\n- As there is no theoretical analysis regarding how accurate the proposed polynomial expansions, experimental validation is necessary. However, the current experimental results appear to be insufficient; more details are in the Requested Changes.\n\n- Currently, the contributions of this work appear to be relatively modest. However, the final sentence of the conclusion outlines a clear direction: \"It can potentially be effectively applied in various domains, including adversarial training, Bayesian inference, generative models, and safety-critical applications, offering a versatile tool for enhancing the reliability and robustness of neural networks.\" These ideas should not be left merely as future work; for the paper to be more substantial, it should either address them comprehensively or delve deeply into one of them.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This work addresses the uncertainty in neural network outputs by examining the propagation of Gaussian inputs through non-linear activation functions. The authors derive Taylor expansions for seven widely-used activation functions (Tanh, Sigmoid, Softplus, ReLU, LeakyReLU, GELU, and SiLU) and present experiments that demonstrate the effectiveness of higher order terms (up to third-order) in these expansions.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "__Major.__\n\n- In CIFAR-10 experiments, is there any specific reason for choosing ResNet depths of 13, 33, and 65? The standard CIFAR-10 settings, originally proposed by He et al. (2016) and widely followed, are 20, 32, 44, and 56. Deviating without a clear reason may make comparisons with existing and future work more difficult, so I recommend following the common convention if possible.\n\n- Section 3.2 demonstrates how well the proposed PTPE method predicts the mean and variance of the output. It would be helpful to also clarify how the method performs with respect to common uncertainty estimation metrics in the machine learning literature, such as negative log-likelihood, Brier score, and expected calibration error, when evaluating ResNet on CIFAR-10. The current results do not provide any indication of whether the trained neural network is practically useful, that is, how well it performs in terms of classification accuracy and how well-calibrated its prediction is.\n\n- In the UCI experiments, was a fair comparison made with the baseline methods? Tables 2 and 3 show results for PTPE using three activations: ReLU, GELU, and Tanh. The choice of activation function should also be considered for the baselines. Therefore, for a fair comparison, comparative results between PTPE and the baseline methods should be provided for each activation function. However, it appears that the reported baseline values were taken directly from the original papers without considering such experimental factors.\n\n__Minor.__\n\n- On page 4, there is an unknown reference; “the number of scaling factors is discussed in ??.”\n\n- On page 5, there is a formatting issue; “the reformulation oftanhand sigmoid.”\n\n- It would be helpful to include the meaning of each column in Figure 3, 8, and 9 either in the subplots or in the captions; while this information is provided in the main text of Section 3.2, adding it would improve readability.\n\n- I could not find the code at the provided anonymous GitHub link in the paper; only the DVI+PTPE, VAE+PTPE folders and LICENSE file are visible.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Decision: Accept (poster)