# Probabilistic computation and uncertainty quantification with emerging
covariance

###### Abstract

Building robust, interpretable, and secure AI system requires quantifying and
representing uncertainty under a probabilistic perspective to mimic human
cognitive abilities. However, probabilistic computation presents significant
challenges for most conventional artificial neural network, as they are
essentially implemented in a deterministic manner. In this paper, we develop
an efficient probabilistic computation framework by truncating the
probabilistic representation of neural activation up to its mean and
covariance and construct a moment neural network that encapsulates the
nonlinear coupling between the mean and covariance of the underlying
stochastic network. We reveal that when only the mean but not the covariance
is supervised during gradient-based learning, the unsupervised covariance
spontaneously emerges from its nonlinear coupling with the mean and faithfully
captures the uncertainty associated with model predictions. Our findings
highlight the inherent simplicity of probabilistic computation by seamlessly
incorporating uncertainty into model prediction, paving the way for
integrating it into large-scale AI systems.

Probabilistic computation and uncertainty quantification with emerging
covariance
Hengyuan Ma1, Yang Qi1,2,3, Li Zhang4, Jie Zhang1,2, Wenlian Lu1,2,5,6,7,8,
Jianfeng Feng1,2,$\ast$

1 Institute of Science and Technology for Brain-inspired Intelligence, Fudan
University, Shanghai 200433, China
2 Key Laboratory of Computational Neuroscience and Brain-Inspired Intelligence
(Fudan University), Ministry of Education, China
3 MOE Frontiers Center for Brain Science, Fudan University, Shanghai 200433,
China
4 School of Data Science, Fudan University, Shanghai 200433, China
5 School of Mathematical Sciences, Fudan University, No. 220 Handan Road,
Shanghai, 200433, Shanghai, China
6 Shanghai Center for Mathematical Sciences, No. 220 Handan Road, Shanghai,
200433, Shanghai, China
7 Shanghai Key Laboratory for Contemporary Applied Mathematics, No. 220 Handan
Road, Shanghai, 200433, Shanghai, China
8 Key Laboratory of Mathematics for Nonlinear Science, No. 220 Handan Road,
Shanghai, 200433, Shanghai, China
$\ast$<EMAIL_ADDRESS>
keywords: Stochastic neural network, Probabilistic computation, Uncertainty
quantification, Machine learning, Artificial intelligence

##  1 Introduction

Uncertainty is an inherent and omnipresent aspect of learning, arising from
various sources such as the environment, observed data, learning process, and
the model itself. While humans are capable of incorporating such uncertainty
into their decision-making [1], it is not easy for most of the current
artificial neural networks (ANNs) to evaluate the level of uncertainty
associated with their predictions [2, 3]. Emulating the brain’s approach to
information processing is one of the most viable ways towards more general
artificial intelligence [4]. The challenge arises from one of the fundamental
distinctions between information processing in the human brain and today’s
hardware used for implementing ANNs: the brain employs probabilistic
computation to process information [5]. In contrast to deterministic
computation based on von Neumann architecture [6], probabilistic computation
goes beyond making predictions and enables the quantification of prediction
uncertainty. This capability is essential for more robust, trustworthy, and
interpretable artificial intelligence [7].

Probabilistic computation performs a series of transformations of probability
distributions to generate a desired distribution whose mean represents the
prediction and whose covariance represents the prediction uncertainty
associated with the relevant input. However, learning general probabilistic
computations is numerically intractable due to the complexity of high-
dimensional joint probability density functions [8]. One of the most popular
solutions to this challenge is Bayesian variational inference (BVI) [9], which
employs a parameterized variational distribution and minimizes a variational
loss. However, evaluating the gradient of the variational loss for this
process is computationally costly. This problem is alleviated by Monte Carlo
(MC) sampling methods [10, 11, 12]. Nevertheless, MC methods require multiple
model calls or training ensembles of models in parallel, leading to increased
time and space complexity compared to deterministic models.

When an ANN is used as a model in BVI, stochasticity is introduced into the
network, and the mean and covariance of the network states are supervised
(either explicitly or implicitly) so that the model can learn to make both
predictions and to represent uncertainty [13, 14, 15, 16]. In biological
neuronal systems, the mean and covariance of neuronal spiking activities
exhibit nonlinear coupling [17, 18], and their propagation across two
populations of neurons shares the same synaptic weights [19]. This implies
that when a model optimizes the output mean for a task loss, it concurrently
adjusts the output covariance through this coupling. We wonder whether this
coupling can autonomously regulate the covariance, enabling it to capture
prediction uncertainty without additional supervision. This inspires us to
construct a probabilistic computation framework where only the mean is
supervised, without additional supervision on the covariance. This approach
allows the covariance to naturally emerge through nonlinear coupling. We
propose that this emerging covariance serves as a faithful representation of
uncertainty.

Through theoretical analysis and numerical experiments we reveal two key
properties of the proposed framework of probabilistic computation, one of
which is that a general probabilistic computation can be approximated by
another stochastic process underlied by a neural network with nonlinearly
coupled covariance, and the other is that prediction uncertainty can be
captured by the emerging covariance. We develop an efficient probabilistic
computation framework consisting of a model called moment neural network (MNN)
and a modified gradient descent optimizer called supervised mean unsupervised
covariance (SMUC). The MNN represents a probability distribution up to its
second cumulants, and conducts nonlinearly coupled mean and covariance via
moment activations (MAs). We train the MNN to learn a range of probabilistic
computation tasks with SMUC, which only back-propagates the gradient through
the mean but not the covariance. As learning progresses, covariance
automatically emerges due to the nonlinear coupling with the mean. By
computing the prediction entropy from emerging covariance, we demonstrate
exceptional capability of the MNN for uncertainty quantification including
misclassified and out-of-distribution sample detection, and adversarial attack
awareness. Moreover, our approach can achieve comparable or even superior
performance for uncertainty quantification compared to existing methods. Our
approach challenges the conventional belief that uncertainty quantification
requires supervision of the covariance, either explicitly or implicitly.
Instead, we show that when the model learns making predictions, the emerging
covariance faithfully captures the prediction uncertainty, suggesting that
learning of probabilistic computation can be fundamentally simpler than
previously believed.

Figure 1:  Schematic of probabilistic computation and learning with
unsupervised covariance. a, Schematic diagram of a stochastic neural network
(left) (Eq. (5), Appendix) and the corresponding moment neural network (MNN,
right) (Eq. (13)). b, Comparison of conventional neural activation (top) and
moment activations (MAs) (bottom). We illustrate the ReLU and ReLU MAs as
examples. c, We train MNNs using the supervised mean unsupervised covariance
(SMUC) scheme, where back-propagation only involves the mean while the
covariances is treated as constants. d, Schematics of a mixed MNN consisting
of shallow layers using conventional activations and deep layers using MAs.
The rectangles represent the layers, and their heights indicate the respective
layer dimensions.

##  2 Probabilistic computation framework for supervised learning

Given a dataset consisting of pairs of input $\mathbf{x}\in\mathbb{R}^{m}$ and
output (label) $\mathbf{y}\in\mathbb{R}^{n}$, we aim to learn a mapping from
an input $\mathbf{x}$ to a probability distribution $q(\mathbf{y}|\mathbf{x})$
of the output, which incorporates both prediction and prediction uncertainty.
In general, exact calculation of such a mapping requires the knowledge of
underlying generative model which is not always available. Moreover, computing
such a mapping over a high-dimensional space is intractable. Here we propose
that such a mapping can be approximated by a chain of simpler stages

$\displaystyle\mathbf{x}\mapsto q(\mathbf{v}^{(1)}|\mathbf{x})\mapsto
q(\mathbf{v}^{(2)}|\mathbf{v}^{(1)})\mapsto\cdots\mapsto
q(\mathbf{v}^{(L)}|\mathbf{v}^{(L-1)})\mapsto q(\mathbf{y}|\mathbf{v}^{(L)}),$
(1)

where $\mathbf{v}^{(l)}$ is the $l$-th hidden state. Equation (1) can be
rewritten in terms of its cumulant sequences [20]

$\displaystyle\mathbf{x}\mapsto\\{m_{z}^{(1)}\\}\mapsto\\{m_{z}^{(2)}\\}\mapsto\cdots\mapsto\\{m_{z}^{(L)}\\}\mapsto\\{m_{z,y}\\},$
(2)

where $m_{z}^{(l)}$ represents the $z$-th order cumulants of the the
distribution $q(\mathbf{v}^{(l)}|\mathbf{x})=\int
q(\mathbf{v}^{(l)}|\mathbf{v}^{(l-1)})p(\mathbf{v}^{(l-1)}|\mathbf{x})d\mathbf{v}^{(l-1)}$,
with $z=1,2,\ldots,+\infty$. When higher-order cumulants beyond $z=Z$ vanish,
the cumulant sequences can be truncated by considering only the first $Z$
cumulants for each distribution. When $Z=1$, the truncated process is written
as

$\displaystyle\mathbf{x}\mapsto\bm{\mu}^{(1)}\mapsto\bm{\mu}^{(2)}\mapsto\bm{\mu}^{(3)}\mapsto\cdots\mapsto\bm{\mu}^{(L)}\mapsto\bm{\mu}_{y},$
(3)

where $\bm{\mu}^{(l)}$ is the mean of the distribution
$q(\mathbf{v}^{(l)}|\mathbf{x})$, and $\bm{\mu}_{y}$ is the mean of model
prediction. We call Eq. (3) as deterministic computation, since the output
$\bm{\mu}_{y}$ is deterministic and the prediction uncertainty is always zero,
and the truncated process at $Z\geq 2$ as probabilistic computation, where
second- or higher- order cumulants account for the uncertainty of the process.
When $Z=2$, the truncated process can be interpreted as a Gaussian
approximation, where the transition of the mean $\bm{\mu}^{(l)}$ and
covariance $C^{(l)}$ across each layer can be expressed as

$\displaystyle\mathbf{x}\mapsto\\{\bm{\mu}^{(1)},C^{(1)}\\}\mapsto\\{\bm{\mu}^{(2)},C^{(2)}\\}\mapsto\cdots\mapsto\\{\bm{\mu}^{(L)},C^{(L)}\\}\mapsto\\{\bm{\mu}_{y},C_{y}\\},$
(4)

with $C_{y}$ representing the prediction covariance. When $3\leq Z<\infty$,
the truncated process cannot be regarded as a stochastic process, rendering
the entropy undefined. This is due to that the cumulant sequences of a
probability distribution either has all but the first two cumulants vanish or
has an infinite number of non-vanishing cumulants [20]. Therefore, we focus on
the $Z=2$ case as this is the minimal model that is capable of representing
uncertainty. The remaining issue is to develop a learnable model for
effectively conducting the truncated process and to devise a suitable training
approach for this model. For deterministic computation (Eq (3)), this can be
achieved by conventional ANNs via compositions of linear and nonlinear
transformations with trainable weights and biases. Due to its universal
approximation property [21], an ANN can learn to approximate arbitrary
continuous functional mapping. In the following, we show that a generalized
neural network incorporating second-order cumulants can be used to learn
probabilistic computation.

##  3 Training stochastic networks with emerging covariance

###  3.1 Moment activations (MAs) and moment neural networks (MNNs)

We propose an extension of conventional ANNs by incorporating stochasticity in
order to approximate mappings from vector spaces to probability spaces beyond
simple functions between vector spaces.

Consider a multi-layer stochastic neural network as follows

$\displaystyle\begin{aligned}
\frac{d\mathbf{x}^{(1)}}{dt}&=-\mathbf{x}^{(1)}+\mathbf{x}+\sqrt{2}\sigma_{1}\bm{\xi}^{(1)},\quad\mathbf{v}^{(1)}=\mathbf{x}^{(1)}\\\
\frac{d\mathbf{x}^{(l)}}{dt}&=-\mathbf{x}^{(l)}+W^{(l-1)}\mathbf{v}^{(l-1)}+\mathbf{b}^{(l-1)}+\sqrt{2}\sigma_{l}\bm{\xi}^{(l)},\quad\mathbf{v}^{(l)}=\mathbf{h}(\mathbf{x}^{(l)}),\quad
l=2,\ldots,L\\\
\frac{d\mathbf{y}}{dt}&=-\mathbf{y}+W^{(L)}\mathbf{v}^{(L)},\end{aligned}$ (5)

where $\mathbf{x},\mathbf{y}$ are the input and output of the system,
$\mathbf{x}^{(l)}$ is the neural state of layer $l$ corresponding to one stage
in Eq. (1), $\mathbf{v}^{(l)}$ is the signal transmitted to neurons at the
next layer, $W^{(l)},\mathbf{b}^{(l)}$ are learnable weights and biases,
$\bm{\xi}^{(l)}$ is the unit Gaussian white noise, $\sigma_{l}^{2}$ is
constant diffusion coefficients, and $\mathbf{h}(\cdot)$ is an element-wise
nonlinearity which transfers each coordinate $x_{i}$ through a real-valued
function $h_{i}(x_{i})$. The goal is to train this stochastic network to
approximate $\mathbf{x}\mapsto q(\mathbf{y}|\mathbf{x})$. Directly training
requires MC methods for estimating the gradient of the network which is
computationally costly. Instead, we approximate the system by constructing its
corresponding truncated process up to the first two cumulants (Eq. (4). To do
so, we assume that the system has reached its stationary distribution, then
estimate the stationary mean and covariance $\\{\bm{\mu}^{(l)},C^{(l)}\\}$ of
the signal vector $\mathbf{v}^{(l)}$ for $l=1,\ldots,L$ as well as the output
mean and covariance $\\{\bm{\mu}_{y},C_{y}\\}$. For the first layer, the mean
and covariance can be directly calculated from Eq. (5)

$\displaystyle\bm{\mu}^{(1)}=\mathbf{x},\quad C^{(1)}=\sigma_{1}^{2}I.$ (6)

We also have the relationship between the cumulants of the last two layers:

$\displaystyle\bm{\mu}_{y}=W^{(L)}\bm{\mu}^{(L)},\quad
C_{y}=W^{(L)}C^{(L)}(W^{(L)})^{\mathrm{T}}.$ (7)

For other layers, it is unfeasible to calculate the cumulants of the signal
$\mathbf{v}^{(l)}$ due to the nonlinearity $\mathbf{h}(\cdot)$. Nevertheless,
we can estimate the truncated transition
$\\{\bm{\mu}^{(l-1)},C^{(l-1)}\\}\mapsto\\{\bm{\mu}^{(l)},C^{(l)}\\}$ between
layers using a mean field approximation (see Sec.1-2, Supplementary
Information for details)

$\displaystyle\bm{\mu}^{(l)}=\mathbf{m}_{v}(\bar{\bm{\mu}}^{(l)},\bar{C}^{(l)}),\quad
C^{(l)}=C_{v}(\bar{\bm{\mu}}^{(l)},\bar{C}^{(l)}),\quad l=1,\ldots,L,$ (8)

where $\bar{\bm{\mu}}^{(l)}=W^{(l-1)}\bm{\mu}^{(l-1)}+\mathbf{b}^{(l-1)}$ and
$\bar{C}^{(l)}=W^{(l-1)}C^{(l-1)}(W^{(l-1)})^{\mathrm{T}}+\sigma_{l}^{2}I$ are
the linear summations of mean and covariance according to the weights and
biases. The two nonlinear mappings
$\mathbf{m}_{v}(\cdot,\cdot)\in\mathbb{R}^{n}\times\mathbb{R}^{n\times
n}\rightarrow\mathbb{R}^{n}$ and
$C_{v}(\cdot,\cdot)\in\mathbb{R}^{n}\times\mathbb{R}^{n\times
n}\rightarrow\mathbb{R}^{n\times n}$ are moment activations (MAs)

$\displaystyle\mathbf{m}_{v}(\bm{\bar{\mu}},\bar{C})_{i}=m_{v}(\bar{\mu}_{i},\bar{C}_{i}),\quad
1\leq i\leq n$ (9) $\displaystyle
C_{v}(\bm{\bar{\mu}},\bar{C})_{ij}=\left\\{\begin{array}[]{@{}ll@{}}C_{v}(\bar{\mu}_{i},\bar{C}_{i}),&\quad
1\leq i=j\leq n\\\
\chi(\bar{\mu}_{i},\bar{C}_{i})\chi(\bar{\mu}_{j},\bar{C}_{j})\tfrac{\bar{C}_{ij}}{\sqrt{\bar{C}_{i}\bar{C}_{j}}},&\quad
1\leq i\neq j\leq n,\end{array}\right.$ (12)

where the dimension of the $\bm{\bar{\mu}}^{(l)}$ is $n$, and three nonlinear
functions $m_{v}(\cdot,\cdot),C_{v}(\cdot,\cdot)$ and $\chi(\cdot,\cdot)$ can
be derived for concrete activation $h(\cdot)$. MAs encapsulate the nonlinear
transformation from the mean and covariance of the input to the output at each
layer, as illustrated in Fig. 1b. Combining Eq. (6)-(8), the stochastic neural
network (Eq. (5)) now is reduced to the following nonlinear deterministic
network

$\displaystyle\begin{aligned} \begin{aligned} &\bm{\mu}^{(1)}=\mathbf{x},\\\
&\begin{array}[]{@{}ll@{}}\bar{\bm{\mu}}^{(l)}=W^{(l-1)}\bm{\mu}^{(l-1)}+\mathbf{b}^{(l-1)},\\\
\bm{\mu}^{(l)}=\mathbf{m}_{v}(\bar{\bm{\mu}}^{(l)},{\color[rgb]{0.13671875,0.48046875,0.828125}\bar{C}^{(l)}}),\end{array}\\\
&\bm{\mu}_{y}=W^{(L)}\bm{\mu}^{(L)},\end{aligned}\begin{aligned}
&{\color[rgb]{0.13671875,0.48046875,0.828125}C^{(1)}=\sigma_{1}^{2}I},\\\
&\left.\begin{array}[]{@{}ll@{}}{\color[rgb]{0.13671875,0.48046875,0.828125}\bar{C}^{(l)}=W^{(l-1)}C^{(l-1)}(W^{(l-1)})^{\mathrm{T}}+\sigma_{l}^{2}I},\\\
{\color[rgb]{0.13671875,0.48046875,0.828125}C^{(l)}=C_{v}(\bar{\bm{\mu}}^{(l)},\bar{C}^{(l)})},\end{array}\right\\}l=2,\ldots,L\\\
&{\color[rgb]{0.13671875,0.48046875,0.828125}C_{y}=W^{(L)}C^{(L)}(W^{(L)})^{\mathrm{T}}}.\end{aligned}\end{aligned}$
(13)

We call the network Eq. (13) as moment neural networks (MNN). The MNN
effectively calculates the mean and covariance of the signals $v^{(l)}$ in
each layer of the stochastic neural network (Eq. (5)). In this way, the MNN
can estimate the output mean and covariance of the stochastic neural network,
without conducting its stochastic dynamics, as illustrated in Fig. 1a. When
the covariances (blue part in Eq. (13)) are fixed at zero, the MNN reduces to
the conventional ANNs.

###  3.2 Supervised mean unsupervised covariance (SMUC)

Suppose a specific configuration ($W^{(l)},\mathbf{b}^{(l)},\sigma_{l}$)
exists that enables the stochastic network Eq. (5) well fits the mapping
$\mathbf{x}\rightarrow q(\mathbf{y}|\mathbf{x})$ of the dataset. We refer to
this network as the ground-truth model, and denote the corresponding MNN as
$M_{GT}$. We aim to train an MNN denoted as $\tilde{M}$ with the same
architecture as Eq. (5) to fit $M_{GT}$, enabling both prediction and
uncertainty quantification. Since $\tilde{M}$ contains means and covariances
of each layer, the conventional perspective for training $\tilde{M}$
necessitates the supervision of both of them. This involves incorporating both
mean and covariance in the loss function and backpropagating the error
gradient through the mean and covariance of each layer during training. We
call this training approach as supervised mean supervised covariance (SMSC).
One drawback of SMSC is the computational cost for backpropagating through
covariances compared to ANNs involving only the mean, as it requires the
calculation and storage the gradient of a quadratic number of elements for
each layer.

Nevertheless, in many neural systems, both biological and artificial alike,
the mean and covariance of neural activity are not independent but nonlinearly
coupled [17, 18]. We wonder whether this nonlinear coupling can be exploited
to autonomously regulate the covariance. Based on this intuition, we propose a
modified gradient descent algorithm called supervised mean unsupervised
covariance (SMUC) for training the MNN, which only backpropagates the error
gradient through the mean while treating the gradients with respect to
covariances (blue part in Eq. (13)) as zero. Concretely, given input-output
pairs $(\mathbf{x},\mathbf{y})$ from the dataset $\mathcal{D}$, the task loss
function $\mathcal{L}$, and learning rate $\gamma_{t}$, the SMUC updates the
parameters $\theta=\\{W^{(l)},\mathbf{b}^{(l)}\\}$ of the MNN as follows

$\displaystyle\theta_{t+1}=\theta_{t}-\gamma_{t}\sum_{(\mathbf{x},\mathbf{y})\in\mathcal{D}}\frac{\breve{\partial}\mathcal{L}(\mathbf{y},\bm{\mu}_{y}(\mathbf{x},\theta_{t}))}{\breve{\partial}\theta_{t}},$
(14)

where $\breve{\partial}$ is a modified gradient operator. When calculating the
gradient using $\breve{\partial}$, we backpropagate the error gradient through
the chain rule only on the mean while treating the covariances as constants:

$\displaystyle\begin{aligned}
&\frac{\breve{\partial}\mathcal{L}}{\breve{\partial}\theta}=\frac{\partial\mathcal{L}}{\partial\bm{\mu}^{(L)}}\frac{\breve{\partial}\bm{\mu}^{(L)}}{\breve{\partial}\theta},\\\
&\frac{\breve{\partial}\bm{\mu}^{(l)}}{\breve{\partial}\theta}=\frac{\partial\bm{\mu}^{(l)}}{\partial\bm{\mu}^{(l-1)}}\frac{\breve{\partial}\bm{\mu}^{(l-1)}}{\breve{\partial}\theta}+\frac{\partial
C^{(l)}}{\partial
C^{(l-1)}}\frac{\breve{\partial}C^{(l-1)}}{\breve{\partial}\theta},\quad\forall
l=L,\ldots,2,\quad\frac{\breve{\partial}\bm{\mu}^{(1)}}{\breve{\partial}\theta}=\frac{{\partial}\bm{\mu}^{(1)}}{{\partial}\theta},\\\
&\frac{\breve{\partial}C^{(l)}}{\breve{\partial}\theta}=0,\quad\forall
l=L,\ldots,1,\end{aligned}$ (15)

A schematic diagram for this learning rule is illustrated in Fig. 1c. The
direct consequence of this approach is that only the mean is supervised for
minimizing the loss, while the covariance spontaneously emerges through the
nonlinear coupling with the mean during forward computation.

###  3.3 Theoretical analysis of the SMUC

We analyze the theoretical properties of the SMUC to demonstrate that the
nonlinearly coupling between the mean and covariance helps to capture the
uncertainty of the model. We find that SMUC is equivalent to stochastic
Riemannian gradient descent (SRGD) [22], a variant of the stochastic gradient
descent (SGD) algorithm. At each iteration of the SRGD, an SGD update is
conducted on the parameter, which is then projected onto a Riemannian
manifold. In our SMUC, it turns out that at each iteration, the parameter
$\theta$ is projected onto the Riemannian manifold where the variance at each
layer are identical to that of the ground-truth MNN $M_{GT}$. (See
Supplementary Information for more details). Based on this, we further prove
the convergence of the SMUC under mild conditions as shown in Thm. 1.

######  Theorem 1.

If the learning rate schedule $\\{\gamma_{t}\\}$ satisfies a modified usual
condition (Eq. (S38), Supplementary Information), the parameter $\theta$
converges to a point at which the gradient diminishes almost surely.

Next, we investigate what the MNN can learn from the dataset after
convergence. We prove that the MNN learns the mean and variance at each layer
of $M_{GT}$ after convergence, as shown in Thm. 2.

######  Theorem 2.

After convergence, $\tilde{M}$ learns the mean and variance of $M_{GT}$ at
each layer.

For the covariance, the situation is more complicated. Although there can be
error on the covariance at each layer of $\tilde{M}$, we prove that for
sufficiently deep networks, the output covariance of $\tilde{M}$ converges to
that of $M_{GT}$, even when the input covariance is randomly chosen, as shown
in Thm. 3.

######  Theorem 3.

After convergence, denote $\lambda_{\min}(C^{(l)})$ as the smallest eigenvalue
of $C^{(l)}$, and define the covariance deviation rate at the $l$-th layer

$\displaystyle
r^{(l)}=\frac{n^{(l)}(\max_{i}\chi_{i}^{(l)})^{2}}{\sqrt{\lambda_{\min}(C^{(l)})^{2}+2\lambda_{\min}(C^{(l)})\sigma_{l}^{2}}}$
(16)

with

$\displaystyle\chi_{i}^{(l)}=\chi(W_{i\bullet}^{(l)}\mathbf{v}^{(l)},W_{i\bullet}^{(l)}C^{(l)}(W_{i\bullet}^{(l)})^{\mathrm{T}}+\sigma_{l}^{2}).$
(17)

Then, as the total number of layers approaches infinity, the output covariance
of the $\tilde{M}$ converges to the output covariance of $M_{GT}$, given that
the series $-\sum_{l}\log r^{(l)}$ diverges to positive infinity.

As a result, the MNN’s ability of uncertainty quantification is maintained
through the output covariance $C_{y}$ without tracking of covariances in the
shallow layers. This allows us to construct a hybrid model called mixed MNN
(Fig. 1d), where we set the covariance in the shallow layers of the network to
zero, reducing them to a conventional ANNs, whereas only in the last few
layers the covariance are computed.

Since the dimensions of the deep layers are lower, the mixed MNN significantly
reduces the computational cost on calculating and storing the covariance
matrices. See Supplementary Information for the proof of Thm. 1-3.
Furthermore, the complexity for training MNNs can be further reduced by
sharing the covariances within a mini-batch (Appendix).

###  3.4 Instantialization of the MNNs

For cases where the activations of the stochastic network are Heaviside
function

$\displaystyle\Theta(x)=\begin{cases}1,\quad x\geq 0\\\ 0,\quad
x<0\end{cases}$ (18)

and ReLU function

$\displaystyle\mathrm{ReLU}(x)=\begin{cases}x,\quad x\geq 0\\\ 0,\quad
x<0\end{cases},$ (19)

we derive expressions of $m_{v}(\cdot,\cdot)$, $C_{v}(\cdot,\cdot)$, and
$\chi(\cdot,\cdot)$ in the corresponding MAs (Eq. (9)-(12)) as follows

Heaviside function.

$\displaystyle
m_{v}(\bar{\mu},\bar{C})=\frac{1}{2}+\frac{1}{\sqrt{2\pi}}\int_{0}^{\frac{\bar{\mu}}{\sqrt{\bar{C}}}}e^{-\frac{x^{2}}{2}}dx,$
(20) $\displaystyle
C_{v}(\bar{\mu},\bar{C})=\frac{1}{2}+\frac{1}{\sqrt{2\pi}}\int_{0}^{\frac{\bar{\mu}}{\sqrt{\bar{C}}}}e^{-\frac{x^{2}}{2}}dx-(\frac{1}{2}+\frac{1}{\sqrt{2\pi}}\int_{0}^{\frac{\bar{\mu}}{\sqrt{\bar{C}}}}e^{-\frac{x^{2}}{2}}dx)^{2},$
(21)
$\displaystyle\chi(\bar{\mu},\bar{C})=\frac{1}{\sqrt{2\pi}}\exp(-\frac{\bar{\mu}^{2}}{2\bar{C}}).$
(22)

ReLU function.

$\displaystyle
m_{v}(\bar{\mu},\bar{C})=\frac{\sqrt{\bar{C}}}{\sqrt{2\pi}}\exp(-\frac{\bar{\mu}^{2}}{2\bar{C}})+\bar{\mu}(\frac{1}{2}+\frac{1}{\sqrt{2\pi}}\int_{0}^{\frac{\bar{\mu}}{\sqrt{\bar{C}}}}e^{-\frac{x^{2}}{2}}dx),$
(23) $\displaystyle
C_{v}(\bar{\mu},\bar{C})=(\bar{C}+\bar{\mu}^{2})(\frac{1}{2}+\frac{1}{\sqrt{2\pi}}\int_{0}^{\frac{\bar{\mu}}{\sqrt{\bar{C}}}}e^{-\frac{x^{2}}{2}}dx)+\frac{\bar{\mu}\sqrt{\bar{C}}}{\sqrt{2\pi}}\exp(-\frac{\bar{\mu}^{2}}{2\bar{C}})-m_{v}(\bar{\mu},\bar{C})^{2},$
(24)
$\displaystyle\sqrt{\bar{C}}(\frac{1}{2}+\frac{1}{\sqrt{2\pi}}\int_{0}^{\frac{\bar{\mu}}{\sqrt{\bar{C}}}}e^{-\frac{x^{2}}{2}}dx).$
(25)

We also consider MAs derived from a type of biological spiking neuron model,
called the leaky integrate-and-fire (LIF) neuron model [19], and corresponding
MNN dubbed LIF MNN, see Eq. (27)-(29) in Appendix for details.

Figure 2: Emerging covariance faithfully captures prediction uncertainty. A
ReLU MNN and a mixed ReLU MNN are trained on MNIST and CIFAR-10 datasets for
image classification respectively. a, Upper panel: the test accuracy during
training. Shades indicate half standard deviation across trails; lower panel:
the entropy for correctly classified (orange) and misclassified (green) inputs
on the test set as learning progresses. Shades indicate half standard
deviation across input samples. b, The entropy of the network state in the
last two layers of the network for correctly classified and misclassified
inputs. c, Separability of entropy between correctly classified and
misclassified inputs increases with the layer index of the MNN. d,
Separability and test accuracy during training are positively correlated.
Dashed lines represent linear regression, with coefficients of determination
found to be 0.6803 for MNIST and 0.6068 for CIFAR-10.

###  3.5 Numerical verification

We numerically verify that the MNN (Eq. (13)) well approximates the output
mean and covariance of the corresponding stochastic network (Eq. (5)) in Fig.
8, Appendix. We train a fully connected ReLU MNN to perform image
classification on the MNIST hand-written digit dataset [23], then use the
trained parameters to construct the corresponding stochastic neural network
(Eq. (5)). When provided with an input (Fig. 8a), the stochastic network
generates corresponding dynamic patterns. (Fig. 8b). We simulate the
stochastic neural network using the Euler-Maruyama method [20], calculate
sample estimates of its output mean, variance and covariance, and compared
them with the output of the MNN (Fig. 8c,d). It is observed that MNN
approximates the stochastic network well in terms of the output mean and
covariance. We conduct the same experiment for Heaviside MNN and obtain
consistent results, as shown in Fig. 9, Appendix.

##  4 Emerging covariance faithfully captures uncertainty

Having provided the essential elements of our theoretical framework (the MNN
and the SMUC learning algorithm), we now turn to training MNN with SMUC to
quantify uncertainty on various tasks. For classification task, we demonstrate
the MNN’s ability for quantifying uncertainty by comparing its performances on
correctly classified, misclassified, out-of-distribution and adversarially
attacked samples. For regression task, we evaluate the performance of MNN
through the log-likelihood measurement.

###  4.1 In-distribution prediction uncertainty

We train a fully connected, multi-layered ReLU MNN on MNIST [24] and a mixed
ReLU MNN on CIFAR-10 [23] datasets for classification. The prediction is
determined by the output mean $\mathbf{m}_{y}$, which produces accuracy
comparable to ANNs’ as shown in Fig. 2a. However, unlike ANNs, the MNN also
expresses prediction uncertainty through the output covariance $C_{y}$. To
quantify the prediction uncertainty, we calculate the entropy of the
prediction distribution $\mathcal{N}(\mathbf{m}_{y},C_{y})$, which depends
solely on the output covariance $C_{y}$ (Appendix). As the training
progresses, the model exhibits higher entropy on misclassified inputs compared
to correctly classified ones (Fig. 2a), indicating that the model can
effectively express uncertainty in its output covariance. We assess the
statistical disparity in the entropy across layers between correctly
classified and misclassified inputs by calculating their separability
(Appendix), which quantifies the degree of separation between network
responses to correctly classified and misclassified inputs. As the layer index
increases, the separability also increases (Fig.2c), indicating an improved
quality of uncertainty representation with deeper layers. While it is widely
acknowledged that increasing the depth of a neural network improves its
fitting ability [25], our finding provides an additional perspective: deeper
networks enable more elaborate representation of uncertainty. We also find a
considerably positive correlation between the separability and accuracy on the
test set throughout the training procedure (Fig. 2d), suggesting that the
performance of prediction and uncertainty quantification are learned
concurrently, even without explicit supervision on covariance. Consistent
results are observed in the Heaviside MNN (Fig. 12, Appendix).

We further investigate if the emerging covariance of the LIF MNN can also
exhibit high uncertainty on the misclassified samples compared to the
correctly classified samples. As shown in Fig. 3, we train a LIF MNN using
SMUC, and observe the consistent results. This implies that the LIF network
can also utilize the nonlinear coupling of the mean and covariance for
efficient uncertainty quantification. Our finding sheds light on the mechanism
through which biological neural systems capture their prediction uncertainty
based on the stochastic neural fluctuation.

Figure 3: Emerging covariance of leaky integrate-and-fire (LIF) MNN faithfully
captures the prediction uncertainty. LIF MNN trained on MNIST for image
classification. a, The training of MNNs on MNIST, with the accuracy (red line)
and entropy (orange and green for correctly classified inputs and
misclassified inputs respectively) on the test set. b, The entropy of
correctly classified and misclassified inputs in the last two layers of MNIST.
The misclassified inputs result in relatively higher entropy in both layers.
c, The entropy separability of correctly classified and misclassified inputs
on MNIST increases as the layer index increases. d, The relationship between
the separability and test accuracy during training on MNIST was analyzed using
linear regression. The with coefficient of determination of 0.7760.

###  4.2 Out-of-distribution prediction uncertainty

Figure 4:  Out-of-distribution detection and adversarial attacks awareness in
the MNN . a, Distribution of entropy, maximum softmax probability (MSP), and
softmax entropy (SE) for correctly classified, misclassified, and out-of-
distribution inputs. b, The separability of three indicators on the correctly
classified, incorrectly classified and out-of-distribution (OoD) samples. The
negative value of separability implies that the MSP and SE incorrectly
indicate higher uncertainty for misclassified samples compared to out-of-
distribution samples. c, The test accuracy of different models under various
strengths $\epsilon$ of FGSM adversarial attacks. d, Box plot of the entropy
of the model prediction for samples in the test set, for the MNN with
$\sigma_{1},\sigma_{2}>0$ under various strengths $\epsilon$ of FGSM attacks.
Orange line: median; box: upper and lower quartile.

Another important aspect for uncertainty representation is detecting out-of-
distribution samples, which refers to data of a type different from that a
model is trained on. To investigate whether MNNs can reasonably express
uncertainty of out-of-distribution samples, we analyze the above MNN trained
on MNIST using out-of-distribution samples from the notMNIST dataset [26].
Specifically, we compare the uncertainty expressed by output covariance (as
measured by entropy) and that expressed by output mean (as measured by maximum
softmax probability and softmax entropy, two common out-of-distribution
indicators [27]). As shown in Fig. 4a, the distribution of covariance-based
indicator (entropy) on the correctly classified, misclassified and out-of-
distribution samples are better separated than that of the mean-based
indicators. Additionally, the separabilities of the covariance-based indicator
are significantly higher than that of the mean-based indicators (Fig. 4b).
These results show that the entropy is more effective at separating the
correctly classified, misclassified and out-of-distribution samples compared
to the other two indicators, implying that the output covariance is a more
appropriate uncertainty indicator than the those derived from the output mean.
Furthermore, it sheds light on why mean-based indicators can offer information
about uncertainty: since the mean and covariance are non-linearly coupled,
they share pertinent information regarding uncertainty.

We investigate the potential of MNNs to defend against gradient-based
adversarial attacks [28, 29], which can significantly degrade the performance
of a classifier by modifying input images imperceptible by human. We train an
Heaviside MNN on the MNIST dataset and then apply fast gradient sign method
(FGSM) [29] to generate adversarial samples. We find that the MNN is
susceptible to adversarial attacks (Fig. 4c). Nevertheless it indicates high
uncertainty on the attacked samples by exhibiting an increase in the output
entropy (Fig. 4d).

Here, we propose a simple modification to the network to combat gradient-based
adversarial attacks. During inference, we set $\sigma_{1}$ and $\sigma_{2}$ in
Eq.(13) to zero, making the Heaviside MA in the first layer (Eq. (20)) reduce
to the Heaviside function with zero derivatives. This results in significant
boost in robustness in exchange of only a marginal drop in test accuracy (from
$0.9797$ to $0.9745$). As shown in Fig. 4c, the test accuracy remains
unchanged regardless of the attack strength applied, implying that the FGSM
attack ineffective on this MNN. This is because FGSM and similar attack
methods utilize the error gradient, which becomes zero due to the Heaviside
function at the first layer, making it difficult to generate effective
adversarial samples.

###  4.3 Comparison with other uncertainty quantification approaches

We further compare our MNN with other representative methods for quantifying
uncertainty. We train MNNs on UCI regression dataset [30] with mean squared
error (MSE) as the training loss. We report the log-likelihood which evaluates
both the prediction (mean) and uncertainty quantification (covariance) of the
model. As shown in Tab. 1, our MNNs demonstrate comparable performance
compared to other methods, and in some cases, even outperform them on several
datasets. Importantly, although our learning algorithm (SMUC) and loss (MSE)
do not incorporate the covariance, it still effectively maximizes log-
likelihood. These results further verify that the emerging covariance can
faithfully capture the prediction uncertainty without supervision.

Table 1: Comparison of MNN and other uncertainty quantification methods. Log-likelihood of different models trained on the UCI regression datasets over random train/test splits. The data shows the mean (standard deviation) over 20 runs. MNN (R) represents ReLU MNN and MNN (H) represents Heaviside MNN. BP, probabilistic backpropagation [31]; MNVI, multiplicative noise variational inference [32]; DVI, deterministic variational inference [33]; Monte Carlo (MC) dropout [13]; deep ensembles [14]. Bold font indicates the highest log-likelihood value for each dataset.  | boston | concrete | kin8 | energy | power | protein | wine | yacht
---|---|---|---|---|---|---|---|---
PBP | -2.57 (0.09) | -3.16 (0.02) | 0.90 (0.01) | -2.04 (0.02) | -2.84 (0.01) | -2.97 (0.00) | -0.97 (0.01) | -1.63 (0.02)
MNVI | -2.43 (0.02) | -3.05 (0.01) | 1.15 (0.01) | -1.33 (0.05) | -2.66 (0.01) | -2.99 (0.01) | -0.96 (0.01) | -0.37 (0.02)
DVI | -2.41 (0.02) | -3.06 (0.02) | 1.13 (0.00) | -1.01 (0.06) | -2.80 (0.00) | -2.85 (0.01) | -0.90 (0.02) | -0.47 (0.03)
Ensemble | -2.41 (0.25) | -3.06 (0.18) | 1.20 (0.02) | -1.38 (0.22) | -2.79 (0.00) | -2.83 (0.02) | -0.94 (0.12) | -1.18 (0.21)
MC Dropout | -2.46 (0.25) | -3.04 (0.09) | 0.95 (0.03) | -1.99 (0.09) | -2.89 (0.01) | -2.80 (0.05) | -0.93 (0.05) | -1.55 (0.12)
MNN (R) | -2.44 (0.00) | -3.59 (0.01) | 0.92 (0.00) | -2.30 (0.01) | -2.92 (0.00) | -0.97 (0.00) | -0.75 (0.00) | -0.29 (0.06)
MNN (H) | -2.65 (0.03) | -3.28 (0.01) | 0.71 (0.00) | -1.31 (0.01) | -2.94 (0.00) | -0.92 (0.00) | -0.90 (0.01) | -0.40 (0.00)

##  5 Further analysis

###  5.1 Mechanism of uncertainty representation by emerging covariance

Figure 5: Mechanism of uncertainty representation by emerging covariance. a,
Illustration of a binary classification dataset, where a sample in
$\mathbb{R}^{2}$ belongs to class A (or class B) if the sign of its two
coordinates $x_{1},x_{2}$ are the same (or different). Solid lines indicate
the class boundary of the dataset. b, Scatter plot of training samples
randomly drawn from a Gaussian distribution. The misclassifed samples are
typically located close to the decision boundary. The size of the bubbles
indicates the magnitude of output entropy. c, Negative correlation is found
between the distance to the decision boundary and the output entropy, with a
coefficient of determination of 0.7196.

To better understand the mechanism underlying uncertainty representation by
emerging covariance, we train an MNN for a binary classification task on
$\mathbb{R}^{2}$. Since the input has a spatial embedding in $\mathbb{R}^{2}$,
this task allows us to directly analyze the relationship between the
uncertainty conveyed by the MNN about each input sample and their spatial
properties. As shown in Fig. 5a, the trained MNN learns the correct class for
most input samples (with an accuracy of $0.9948$). Moreover, misclassified
samples are typically located close to the decision boundary of the model
(Fig. 5b) and have larger entropy. Figure. 5c shows that the average entropy
decreases with samples’ distance to the decision boundary with a correlation
coefficient of $-0.85$, hence the closer a sample is to the decision boundary,
the higher the output entropy. This explains why the adversarial attacks
increase the entropy associated with each sample (Fig. 4d): they push the
samples toward the decision boundary of the model [34], thus increasing the
entropy of model prediction.

This observation leads to the question why samples closer to the decision
boundary have higher entropy. This can be understood conceptually by
considering the stochastic network (Eq. (5)) underlying the MNN. The
stochastic state at each layer serves as input to the remaining part of the
network, which acts as a classifier. When the input state is near the decision
boundary of this classifier, even small fluctuations can cause substantial
output variations. This leads to significant input-output sensitivity which
causes the network to amplify the stochastic fluctuation of its input,
resulting in high entropy. To provide a clearer theoretical understanding, we
derive an explicit expression for the entropy and the input’s distance to the
decision boundary for a binary logistic regression problem (Appendix).

Above insight allows us to formulate the following explanation for the high
entropy of model prediction on the out-of-distribution samples. Consider the
scenario where an MNN is trained for an arbitrary classification problem.
Initially, the decision boundaries of the model are randomly distributed
throughout the entire input space. As learning progresses, the model adapts
the decision boundaries to fit the actual class boundary of the dataset of in-
distribution samples. Since the entropy is lower when they are far away from
the decision boundaries (as we have shown in Fig. 5b), only the small set of
samples near the class boundary have high entropy. In the out-of-distribution
region, however, the decision boundaries lack a structured distribution and
are randomly scattered throughout the region, hence increasing the chance that
an out-of-distribution sample is close to one of the decision boundaries. As a
result, the entropy of out-of-distribution samples tends to be higher compared
to that of in-distribution samples.

###  5.2 Effect of the covariance

We further investigate the effect of the correlation between neurons on model
learning and uncertainty quantification. We repeat the experiments presented
in Fig. 2, while manually setting all off-diagonal elements of the covariance
matrices to zero during the training and inference of MNNs. As shown in the
left panel of Fig. 6a, for MNIST classification (pure MNN), the MNN without
correlation has similar performance in terms of test accuracy, However, for
CIFAR-10 classification (mixed MNN), MNN without correlation learns much
slower compared to the one with correlations, as shown in the right panel of
Fig. 6a. In terms of uncertainty quantification, we compare the separability
of the entropy on the correctly classified and misclassified samples across
layers for MNIST the left panel of (Fig. 6b) and CIFAR-10 the right panel of
(Fig. 6b) classification. It is observed that for CIFAR-10, there is only a
slight decrease in separability when omitting the correlation between neurons.
On the other hand, for MNIST, there is a relatively larger decrease in
separability. This difference can be attributed to the fact that in the mixed
MNN, the inputs are initially decorrelated by the convolution layers. As a
result, the correlations in the layers that utilize MAs become less crucial
for uncertainty quantification. However, as demonstrated in Figure 6b, these
correlations still play certain role in facilitating faster learning.

Figure 6: Comparison of MNNs with and without correlation. a, Test accuracy
across training procedure for MNNs with and without correlation on MNIST
(left) and CIFAR-10 (right) classification. b, Entropy separability across
layers procedure for MNNs with and without correlation on MNIST (left) and
CIFAR-10 (right) classification.

##  6 Discussion

###  6.1 Sufficiency of the second-order cumulant approximation

In this study, we have approximated a stochastic network up to its second-
order cumulant while omitting the higher-order cumulants. The validity of this
approximation is supported by both empirical evidence and theoretical
analysis. In ref. [35], the authors analyzed neural data of vertebrate retina
and discovered that the pairwise correlations of biological networks can
accurately predict the collective effects of the system without considering
higher-order interactions. In ref. [36], the authors discovered that in a
weakly correlated state, the contributions of higher-order cumulants on
network dynamics are significantly smaller, implying that the first two
cumulants provide a reasonably accurate representation of the network
dynamics.

###  6.2 Relation to different learning schemes

Several methods for uncertainty quantification represent the stochasticity of
neural networks in a deterministic fashion [31, 37, 33, 32]. Our approach
differs from their methods in that the neural network is trained with the SMUC
learning rule, in which only the mean but not the covariance is supervised,
while theirs fall under the BVI formulation using the SMSC learning rule, in
which both mean and covariance are supervised during training, as illustrated
in Appendix. The MNN trained with SMUC can demonstrate superior performance on
uncertainty quantification (Tab. 1). Although the present study has focused on
probabilistic learning under the supervised setting, future works can future
explore this idea for unsupervised learning, where labels are unavailable. For
example, Barlow Twins method [38] trains an encoder using a covariance-based
loss function to reduce redundancy. This suggests it is possible to train MNNs
using a unsupervised mean supervised covariance (UMSC) approach for
unsupervised learning tasks.

###  6.3 Inspiration for learning in the brain

Our results leads to an intriguing prediction about the learning rules
employed by the brain. Although biological neural systems is inherently
probabilistic due to the presence of various sources of noise and variability
[39], we expect that learning of probabilistic computation can be realized by
adjusting synaptic connectivity through gradient descent via the mean firing
rate of neurons. Our prediction is founded on two key observations: first, we
have demonstrated that the MNN derived from the leaky integrate-and-fire
neuron model, which is commonly used for modeling spiking neuronal activity in
the brain, can capture prediction uncertainty through its emerging covariance
(Fig. 3); second, it is commonly hypothesized that the brain adjusts synaptic
connectivity in a manner similar to gradient descent [40]. To further test our
predictions, we propose investigating whether neural dynamics observed in the
brain can be replicated by MNNs.

##  7 Conclusion

Our research introduces two perspectives on probabilistic computation. First,
we demonstrate that complex probabilistic computation can be conducted by
combinations of simple nonlinearities that captures truncated transition
probabilities up to their second-order cumulants. Second, we reveal that
output covariance spontaneously emerging from the nonlinear coupling between
the mean and covariance of network states faithfully captures the prediction
uncertainty. Our probabilistic computation framework acts as a natural
extension to conventional ANNs, seamlessly incorporating uncertainty
quantification into learning and inference. Our methodology demonstrates the
feasibility of probabilistic computation within the current deterministic
setting, paving the way for its integration into large-scale AI systems.

We derive a set of moment activations (MAs, Eq. (9)-(12)) that propagate
nonlinearly coupled cumulants and enable gradient-based learning of
probabilistic computation. We then utilize the MAs to construct a new class of
neural network called the moment neural network (MNN, Eq. (13)) which
approximates a parameterized stochastic neural network (Eq. (5)) up to its
second cumulants. Unlike conventional ANNs that primarily focus on nonlinear
transformations of the mean, MNNs also incorporate nonlinear coupling between
the mean and covariance. This capability enables direct training through
gradient descent, eliminating the need for MC sampling methods. Moreover, MNNs
offer a unique advantage in analyzing the dynamics of stochastic networks by
explicitly expressing the nonlinear coupling between the mean and covariance,
which is implicit in MC methods.

Next, we develop a learning algorithm called SMUC for training MNNs (Eq.
(13)-(14)), where the loss is backpropagated only through the mean while
treating the covariance as constants during gradient descent. Through
extensive experiments (Fig. 2-4, Tab. 1), we demonstrate that the covariance
spontaneously emerging from its nonlinear coupling with the mean effectively
captures the prediction uncertainty. This challenges the conventional wisdom
that it is necessary to supervise both the prediction and the prediction
uncertainty for uncertainty quantification.

## Code Availability

The code for training the moment neural network (MNN) through supervised mean
unsupervised covariance (SMUC) is available without restrictions on GitHub
(https://github.com/AwakerMhy/probabilistic-computing-mnn).

## Acknowledgments

Supported by STI2030-Major Projects (No. 2021ZD0200204); supported by Shanghai
Municipal Science and Technology Major Project (No. 2018SHZDZX01), ZJ Lab, and
Shanghai Center for Brain Science and Brain-Inspired Technology; supported by
the 111 Project (No. B18015).

## Competing interests

The authors declare no competing interests.

## Materials & Correspondence

Correspondence and requests for materials should be addressed to J.F.

## Appendix

### Moment neural network (MNN)

In general, it is hard to solve $\mathbf{v}^{(l)}$ in the stochastic neural
network Eq. (5) analytically. Nevertheless, under steady-state distribution,
we can approximate $\mathbf{v}^{(l)}$ by a Gaussian distribution
$\mathcal{N}(\bar{\bm{\mu}}^{(l)},C^{(l)})$ using mean field approximation
(Supplementary Information). In this way, the signals received by the neurons
at $l$-th layer obey a Gaussian distribution
$\mathcal{N}(\bar{\bm{\mu}}^{(l)},\bar{C}^{(l)})$, with
$\bar{\bm{\mu}}^{(l)}=W^{(l-1)}\bm{\mu}^{(l-1)}+\mathbf{b}^{(l-1)}$ and
$\bar{C}^{(l)}=W^{(l-1)}C^{(l-1)}(W^{(l-1)})^{\mathrm{T}}+\sigma_{l}^{2}I$.
Based on the linear response theory [41] (Supplementary Information), we
obtain a set of nonlinear mappings dubbed moment activations (MAs) describing
how the neurons transfer the first two cumulants of received signals to that
of the output signal. See Supplementary Information for the derivation of the
MAs (Eq. (S12), (S13), and (S16)).

### LIF MNN

In [42, 19], the authors analyzed the stochastic dynamics of a spiking
neuronal network (SNN) by a set of nonlinear mapping describing how the
neurons transfer the input spike signals from other neurons. Concretely, the
authors studied a network of $N$ neurons with the dynamics of each neuron
applying the leaky integrate-and-fire (LIF) neuron model

$\displaystyle\dfrac{dV_{i}}{dt}=-LV_{i}+\sum_{j=1}^{N}w_{ij}S_{j},\quad
i=1,\ldots,N,$ (26)

where $V_{i}$ is the membrane potential of the $i$-th neuron, $L$ is the leaky
conductance, $w_{ij}$ is the synaptic weight from neuron $j$ to neuron $i$,
and $S_{j}$ is the spike train from pre-synaptic neurons. When $V_{i}$ reaches
the firing threshold $V_{\text{th}}$ the neuron will release a spike which is
transmitted to other connected neurons, and then $V_{i}$ is reset to the
resting potential $V_{\text{res}}$ and enters a refractory period
$T_{\text{ref}}$. The corresponding moment activations, which we refer to as
LIF moment activations (LIF MAs), are as follows [19]

$\displaystyle m_{v}(\bar{\mu},\bar{C})=\big{(}T_{\rm
ref}+\frac{2}{L}\int_{I_{\rm lb}}^{I_{\rm ub}}g(x)dx\big{)}^{-1},$ (27)
$\displaystyle
C_{v}(\bar{\mu},\bar{C})=\frac{8}{L^{2}}m_{v}(\bar{\mu},\bar{C})^{3}\textstyle\int_{I_{\rm
lb}}^{I_{\rm ub}}h(x)dx,$ (28)
$\displaystyle\chi(\bar{\mu},\bar{C})=\frac{\partial}{\partial\bar{\mu}}m_{v}(\bar{\mu},\bar{C})$
(29)

where $T_{\rm ref}$ is the refractory period with integration bounds $I_{\rm
ub}(\bar{\mu},\bar{C})=\tfrac{V_{\rm th}L-\bar{\mu}}{\sqrt{L\bar{C}}}$ and
$I_{\rm lb}(\bar{\mu},\bar{C})=\tfrac{V_{\rm
res}L-\bar{\mu}}{\sqrt{L\bar{C}}}$. The constant $L$, $V_{\rm res}$, and
$V_{\rm th}$ are identical to those in Eq. (26). The pair of Dawson-like
functions $g(x)$ and $h(x)$ appearing in Eq. (27) and Eq. (28) are

$\displaystyle g(x)=e^{x^{2}}\int_{-\infty}^{x}e^{-u^{2}}du,\quad
h(x)=e^{x^{2}}\int_{-\infty}^{x}e^{-u^{2}}[g(u)]^{2}du.$ (30)

It is important to highlight that in the LIF MAs (Eq. (27)-(29)), the mean is
the average number of the spike in a unit time interval, the variance is the
variance of the number of the spike in a unit time interval, and the
covariance between two neuron $i$ and $j$ is the multiplication of their
variances and the coherence of their spike train [41]. Using the LIF MAs, we
can construct the corresponding MNN dubbed LIF MNN with learnable weight
matrix. After training, the MNN can be reverted into the corresponding SNN
[43], similar as the ReLU and Heaviside MNNs. In comparison to the ReLU and
Heaviside MAs, the computation of LIF MAs involves higher numerical complexity
due to the Dawson-like functions, which can present challenges in implementing
them for large-scale models. Nonetheless, the LIF MNN exhibits well-defined
biological plausibility, making it a valuable tool for analyzing the neural
circuit mechanisms underlying brain coding and learning processes [44]. For
implementation, we set $V_{\rm th}=20$ mV, $V_{\rm res}=0$ mV, $T_{\rm ref}=5$
ms, and $\tau=1/L=20$ ms, following [19]. and conduct the LIF MAs by the
efficient numerical algorithm developed in [45].

### Metrics for quantifying uncertainty

Entropy.  The output of the MNN can be considered as a multi-dimensional
Gaussian distribution $\mathcal{N}(\bm{\mu}_{y},C_{y})$, whose entropy is
calculated as

$\displaystyle
H\big{(}\mathcal{N}(\bm{\mu}_{y},C_{y})\big{)}=\frac{n}{2}(1+\log
2\pi)+\frac{1}{2}\log\mathrm{det}[C_{y}],$ (31)

where $n$ is the dimension of the output. If $C_{y}$ is singular, we have
$\mathrm{det}[C_{y}]=0$, which leads to entropy of infinite negative value. In
such scenarios, we compute the entropy in the linear subspace of
$\mathbb{R}^{n}$ where the covariance matrix $C_{y}$ is full-rank through SVD
and use the dimension of this subspace as the effective dimension for
calculation. It is important to note that the entropy indicators do not
incorporate the output mean $\bm{\mu}_{y}$, which is supervised during
training. Instead, this indicator solely rely on the output covariance
$C_{y}$, which emerges from the nonlinear coupling with the mean.
Additionally, we can calculate the entropy at each layer of the MNN. Given the
mean and covariance of the $l$-th layer as $\bm{\mu}^{(l)},C^{(l)}$, the
entropy at the $l$-th layer is

$\displaystyle
H\big{(}\mathcal{N}(\bm{\mu}_{v}^{(l)},C^{(l)})\big{)}=\frac{m^{(l)}}{2}(1+\log
2\pi)+\frac{1}{2}\log\mathrm{det}[C^{(l)}].$ (32)

Maximum softmax probability (MSP) Given the output mean $\bm{\mu}_{y}$, we can
define a probability distribution using softmax function

$\displaystyle
p_{i}=\frac{\exp\\{\mu_{y,i}\\}}{\sum_{j=1}^{n}\exp\\{\mu_{y,j}\\}},\quad
i=1,\ldots,n.$ (33)

The maximum softmax probability (MSP) has been widely employed as an indicator
for detecting out-of-distribution samples [27]. The MSP is determined by
calculating the maximum probability, denoted as $\max_{i}p_{i}$, from the
softmax output. Lower values of MSP indicate higher uncertainty in the model’s
predictions.

Softmax entropy (SE). Another indicator utilized is the entropy of the
distribution ${p_{i}}$, as proposed by [46]. The entropy, referred to as
softmax entropy (SE) in this context, is computed as $-\sum_{i=1}^{n}p_{i}\log
p_{i}$. It is important to distinguish this entropy from the entropy
calculated by MNNs. The SE provides a measure of the uncertainty in the
softmax output distribution. Higher SE values indicate greater uncertainty in
the model’s predictions.

### Separability

Suppose we aim to quantify the degree of separation for a variable between two
populations, labeled as 1 and 2. In such cases, we can calculate the
separability as

$\displaystyle\frac{\mu_{1}-\mu_{2}}{\sqrt{\sigma_{1}^{2}+\sigma_{2}^{2}}},$
(34)

where $\mu_{1}$ and $\mu_{2}$ are the means of the variable in population 1
and 2 respectively, and $\sigma_{1}^{2}$ and $\sigma_{2}^{2}$ are variances of
the variable in population 1 and 2 respectively. A higher absolute value of
the separability indicates a greater statistical distinction between the two
populations. When calculating the separability of entropy between two sample
populations in the main paper, we designate population 1 as the population we
expect to have a higher entropy (misclassified samples, for instance) compared
to the other population (correctly classified samples, for instance).

### Mean square error (MSE) and log likelihood (LL)

Supposed that there are $N_{\mathrm{test}}$ samples in the test set, and for
the $i$-th sample in the test set, the model given the output mean
$\bm{\mu}_{y,i}\in\mathbb{R}^{d}$, covariance $C_{y,i}\in\mathbb{R}^{d\times
d}$ and the corresponding ground-truth output is
$\mathbf{y}_{i}\in\mathbb{R}^{d}$. The mean square error (MSE) measurement is
calculated as

$\displaystyle\mathrm{MSE}=\frac{1}{N_{\mathrm{test}}}\sum_{i=1}^{N_{\mathrm{test}}}\left\|\mathbf{y}_{i}-\bm{\mu}_{y,i}\right\|^{2},$
(35)

which evaluates the accuracy of the model prediction while does not include
the uncertainty quantification ability of the model. The log likelihood (LL)
measurement is calculated as

$\displaystyle\mathrm{LL}=-\frac{1}{N_{\mathrm{test}}}\sum_{i=1}^{N_{\mathrm{test}}}\frac{1}{2}\big{(}\log\mathrm{det}[2\pi
C_{y,i}]+(\mathbf{y}_{i}-\bm{\mu}_{y,i})^{\mathrm{T}}C_{y,i}^{-1}(\mathbf{y}_{i}-\bm{\mu}_{y,i})\big{)},$
(36)

which evaluates both the prediction (mean) and uncertainty quantification
(covariance) of the model.

### Training implementations

By default, we utilize the cross-entropy loss function as the objective
function and employ the Adam optimizer [47] for training, without any
additional specifications. Additionally, we set $\sigma_{l}$ to a value of
$0.2$ for all layers.

Classification on MNIST [23].  We train ReLU and Heaviside MNNs respectively,
whose architecture consisting of six fully connected layers with dimensions of
$784,392,196,96,48,24,10$, respectively. We set the training batch size as
$128$, the learning rate as $0.0005$, weight decay factor as $0.001$ and the
training epochs as $50$. For LIF MNN, we construct a fully connected LIF MNN
on MNIST dataset for classification, where the dimension at each layer are
784, 392, 96 and 10, and we add a batch-normalization layer before each
nonlinearity, following [43]. We apply Adam [47] as the optimizer. The
training batch-size is 128, and epoch is 15, the learning rate is 0.0005, and
the factor of weight decay is 0.0001.

Classification on CIFAR-10 [23].  We train a mixed MNN with ReLU and Heaviside
MAs respectively. The network consists of 10 layers of convolutional layer and
three fully connected MA layers. The construction of these convolutional
layers follows the design of VGG13 [48]. The dimensions of the MA layers are
set as $64,32,10$, respectively. We set the training batch size as $128$, the
learning rate as $0.0005$, weight decay factor as $0.001$ and the training
epochs as $120$, using random crop and random horizontal flip as data
augmentation.

Out-of-distribution detection.  We train a fully connected ReLU MNN on the
MNIST dataset, and the dimension of each layer is $784,392,196,96,48,24,10$.
The training batch-size is 128, and epoch is 50, the learning rate is 0.0005,
and the factor of weight decay is 0.0001. After training, we use the notMNIST
[26] dataset as the out-of-distribution samples. NotMNIST dataset consists of
$28\times 28$ grayscale images of alphabets from $A$ to $J$.

Adversarial attack defense and awareness.  We train an fully connected
Heaviside MNN on MNIST, where the dimension of each layer is
$784,392,196,96,48,24,10$. The training batch-size is 128, and epoch is 50,
the learning rate is 0.0005, and the factor of weight decay is 0.0001. We
apply the fast gradient sign method (FGSM) [29] for generating adversarial
samples on the test set. FGSM adds an optimal max-norm constrained
perturbation to the input $\mathbf{x}$, and this perturbation is calculated as
$\epsilon\mathrm{sgn}(\bigtriangledown_{\mathbf{x}}\mathcal{L}(\theta,\mathbf{x},\mathbf{y}))$,
where $\mathbf{x},\mathbf{y}$ is an input-output pair from dataset, $\theta$
is the parameters of the trained model, and $\epsilon>0$ controls the strength
of the perturbation. The FGSM finds the direction along which the model output
changes the most and adds small perturbation along that direction.

Regression problem We train both ReLU MNN and Heaviside MNN through SMUC for
reducing the mean square error (MSE) loss on the predicted outputs and ground-
truth outputs. We normalize the dimensions of all the datasets so that their
mean are $0$ and standard deviation are $1$. The MNN is fully connected and
has one hidden layer of width $50$, which is the same as that in [32]. We
utilize the Adam optimizer with a weight decay factor of zero for
optimization. Across all datasets, the training batch size is fixed at 128,
the learning rate is set to 0.001, and the number of epochs is set to 500,
except for the power and protein datasets, where we set the epoch to 20. We
assign a value of 0 to $\sigma_{2}$ and choose the best results from among the
options for $\sigma_{1}$ which include 0.02, 0.05, and 0.1.

Binary classification problem.  We assign the categories of samples in the
training dataset as follows. Given a sample $\mathbf{x}\in\mathbb{R}^{2}$, it
belongs to class A if the sign of two coordinates
$\mathbf{x}_{1},\mathbf{x}_{2}$ are the same, otherwise it belongs to class B.
We train a fully connected ReLU MNN. The dimensions of each layer in the
network are $2,64,32,16,8,2$, respectively, and $\sigma_{l}$ are all set as
$0.5$. We generate the training set of 10,000 samples by randomly sampling
from a standard Gaussian distribution on $\mathbb{R}^{2}$. The training batch
is 128, the epoch is 20, the learning rate is 0.0005, and the factor of the
weight decay is 0.00001. Since the model fits well with the dataset (Fig. 5b),
we consider that the decision boundary learned by the model aligns with the
classification boundary of the dataset ($\mathbf{x}_{1}$ and $\mathbf{x}_{2}$
axes). The distance from a point $\mathbf{x}$ to the decision boundary
($\mathbf{x}_{1}$ and $\mathbf{x}_{2}$ axes) is
$d(\mathbf{x})=\min(\left|\mathbf{x}_{1}\right|,\left|\mathbf{x}_{2}\right|)$.

## Numerical verification of covariance activations and moment neural networks

Figure 7: Numerical verification for covariance activations. The covariance of
the signals produced by two neurons calculated using Heaviside or ReLU
covariance activations is consistent with that of Monte Carlo simulations. We
present the results with two different correlation levels: a small correlation
value of $c=0.01$ (top) and a relatively larger correlation value of $c=0.2$
(bottom).  Figure 8: The ReLU MNN (Eq. (13)) well approximates the
corresponding stochastic network (Eq. (5)) up to its first two cumlants. a,
The input data from MNIST. b, We train an MNN on MNIST and use the parameters
of the trained MNN to construct the corresponding stochastic network, and show
its dynamics. c, The mean and variance produced by MNN and estimated using
Monte Carlo methods from the stochastic network. d, The covariance matrix
produced by MNN and estimated using Monte Carlo methods from the stochastic
network.  Figure 9: The Heaviside MNN (Eq. (13)) approximates the
corresponding stochastic network (Eq. (5)) up to its first two cumlants. a,
The input data from MNIST. b, We train an MNN on MNIST and use the parameters
of the trained MNN to construct the corresponding stochastic network, and show
its dynamics. c, The mean and variance output by MNN and estimated using Monte
Carlo methods from the stochastic network. d, The covariance matrix produced
by MNN and estimated using Monte Carlo methods from the stochastic network.

We first numerically verify the accuracy of the covariance activations (Eq.
(S16)) for ReLU and Heaviside functions. We present the signals produced by
two neurons $C_{v,ij}$ calculated using the covariance activation, and compare
it with the covariance obtained through Monte Carlo simulation as shown in
Fig. 7. The consistency between the results obtained from MNNs and the
stochastic networks validate that MNNs well approximates the underlied
stochastic network up to its first two cumulants.

We then numerically verify that the ReLU and the Heaviside MNN can approximate
the output mean and covariance of the corresponding stochastic neural network.
We train a fully connected Heaviside MNN on MNIST for classification, and then
use the trained parameters to construct the corresponding stochastic neural
network (Eq. (5)). When provided with an input (Fig. 9a), the stochastic
network generates corresponding dynamic patterns. (Fig. 9b). We estimate the
output mean, variance and covariance of the stochastic network through Euler
Maruyama method, and compared it with the that calculated by the Heaviside MNN
(Fig. 9c,d). As shown in Fig. 9, it is observed that the Heaviside MNN
approximates the corresponding stochastic network well in terms of the output
mean and covariance.

## Theoretical comparison with Bayesian variational inference framework

We conducted a comparison between SMUC and Bayesian variational inference
(BVI) framework. To enable a direct comparison, we rewrite the MNN using the
steady-state of the stochastic network (Eq. (5)) under the mean field
approximation

$\displaystyle\begin{aligned}
&\mathbf{v}^{(1)}=\mathbf{x}+\sigma^{(1)}\mathbf{z}^{(1)}\\\
&\mathbf{v}^{(l)}=\mathbf{}(W^{(l-1)}\mathbf{v}^{(l-1)}+\mathbf{b}^{(l-1)}+\sigma_{l}\mathbf{z}^{(l)}),\quad
l=2,\ldots,L\\\ &\mathbf{y}=W^{(L)}\mathbf{v}^{(L)},\end{aligned}$ (37)

where $\mathbf{z}_{l}$ are independent standard Gaussian vectors. After
combining the bias $\mathbf{b}^{(l)}$ and the Gaussian vectors
$\sigma_{l}\mathbf{z}^{(l)}$ at each layer, we can regard the biases as
stochastic parameters, obeying
$\mathcal{N}(\mathbf{b}^{(l)},\sigma_{l}^{2}I)$. Hence, we can consider the
network is parameterized by stochastic parameters $\theta$, and the
distribution of $\theta$ is parameterized by the trainable parameters $\psi$
which collects all the weights $W^{(l)}$ and the bias $\mathbf{b}^{(l)}$.
Denote the distribution of $\theta$ as $q_{\psi}(\theta)$. Then the
variational loss is

$\displaystyle\mathcal{L}(\theta)=-\int q_{\psi}(\theta)\log
p(\mathbf{y}|\mathbf{x},\theta)d\theta+\mathrm{KL}(q_{\psi}(\theta)||p(\theta)),$
(38)

where $p(\theta)$ is the prior distribution of $\theta$. The right of Eq. (38)
are log-likelihood term and KL-divergence term respectively. We approximate
the prediction probability $\int
p(\mathbf{y}|\mathbf{x},\theta)q_{\psi}(\theta)d\theta$ up to its second
cumulants by a Gaussian distribution
$\mathcal{N}(\mu(\mathbf{x},\psi),C(\mathbf{x},\psi))$, where the mean and
covariance are functions of $\mathbf{x}$ and $\bm{\psi}$. Suppose that the
loss of the task is mean square error, then the log-likelihood term of
variational inference is

$\displaystyle\mathbb{E}_{\bm{\eta}\sim\mathcal{N}(\mu(\mathbf{x},\psi),C(\mathbf{x},\psi))}[\left\|\bm{\eta}-\mathbf{y}\right\|^{2}]=\left\|\mu(\mathbf{x},\psi)-\mathbf{y}\right\|^{2}+\mathrm{Tr}(C(\mathbf{x},\psi)).$
(39)

In our training procedure SMUC, however, only the mean is supervised for
minimizing $\left\|\mu(\mathbf{x},\psi)-\mathbf{y}\right\|^{2}$, and the
covariance spontaneously emerges via the nonlinear coupling with the mean.
Thus, we call BVI as supervised mean supervised covariance (SMSC), as it
supervises both mean and covariance.

## Sharing the covariances within a mini-batch during training

We empirically find that the complexity for training MNNs can be further
reduced without significantly sacrificing performance. Addressing
computational complexity is a common challenge in probabilistic computation
algorithms, particularly when dealing with high-dimensional inputs or large
datasets [2]. For MNN and other probabilistic computation algorithms that
explicitly calculate covariance matrices such as [33], the main overhead of
computation is on the calculating and storing the covariance matrices, which
arises $O(d^{3})$ time complexity and $O(d^{2})$ space complexity. In the case
of MNNs and other probabilistic computation algorithms that involve explicit
calculations of covariance matrices, such as the approach described in [33],
the primary computational overhead lies in the calculation and storage of
these covariance matrices. This process incurs a time complexity of $O(d^{3})$
and a space complexity of $O(d^{2})$, where $d$ represents the data
dimensionality. When the covariance matrices are involved into
backpropagation, the computational complexity becomes even higher. Luckily, we
do not backpropagation on the covariance in MNN. Hence the main complexity
comes from the feed-forward propagation. Supposed that the batch size during
training is $N_{\mathrm{batch}}$, then to compute a covariance matrix of
vectors of dimension $d$ for the whole batch, the space and time complexity
are $O(N_{\mathrm{batch}}d^{2})$ and $O(N_{\mathrm{batch}}d^{3})$
respectively. To make MNNs more scalable, instead of calculating the
covariance matrix for each input in the mini-batch during training, we
calculate a single covariance matrix for each mini-batch. Supposed that there
is a mini-batch of input $\\{x_{1},x_{2},\ldots,x_{N_{\mathrm{batch}}}\\}$, at
layer $l$, the corresponding means are
$\\{\mu_{1}^{(l)},\mu_{2}^{(l)},\ldots,\mu_{N_{\mathrm{batch}}}^{(l)}\\}$. We
calculate the mini-batch wise average mean

$\displaystyle\mu^{(l)}=\frac{1}{N_{\mathrm{batch}}}\sum_{k=1}^{N_{\mathrm{batch}}}\mu_{k}^{(l)},\quad
l=1,2,\ldots,L$ (40)

and use it to calculate the covariance matrix, which is shared for the inputs
in this mini-batch

$\displaystyle
C^{(l)}=C_{v}(W^{(l-1)}\mu^{(l-1)},W^{(l-1)}C^{(l-1)}(W^{(l-1)})^{\mathrm{T}}),\quad
l=2,\ldots,L,$ (41)

where $C^{(1)}=C^{\mathrm{in}}$, which is set as the same for all the input.
In this way, the space and time complexity are reduced to $O(d^{2})$ and
$O(d^{3})$ respectively.

We report the performance of MNN when sharing the covariance in a mini-batch
(Appendix). As shown in Fig. 10 and Fig. 11, the emerging covariance in the
MNNs when sharing the covariance in each mini-batch also faithfully captures
the prediction uncertainty, out-of-distribution detection, and adversarial
attack awareness, similar as the original MNNs (Fig. 2). We also report the
corresponding performance comparison on the regression and classification
tasks in Tab. 2 and Tab. 3 respectively, and we observed that the sharing
covariance within a mini-batch only causes marginal degradation compared to
the original MNNs (Tab. 4 and Tab. 5).

Figure 10: Emerging covariance of MNNs faithfully captures the prediction uncertainty when sharing the covariance within a mini-batch during training. (Mixed) MNN trained on MNIST (CIFAR-10) for image classification using the batch-wise covariance trick. a, The training of MNNs on (left) and CIFAR-10 (right), with the accuracy (red line) and entropy (orange and green for correctly classified inputs and misclassified inputs respectively) on the test set. b, The entropy of correctly classified and misclassified inputs in the last two layers of MNIST (left) and CIFAR-10 (right). c, The entropy separability of correctly classified and misclassified inputs on MNIST (left) and CIFAR-10 (right) increases as the layer index increases. d, The relationship between the separability and test accuracy during training on MNIST (left) and CIFAR-10 (right) was analyzed using linear regression. The coefficients of determination of 0.4334 for MNIST and 0.8994 for CIFAR-10.  Figure 11: The emerging covariance of MNN using the batch-wise covariance tricks for out-of-distribution detection and adversarial attack awareness. a, The distribution of entropy on correctly classified, misclassified and out-of-distribution input. b, If we do not turn off the variance, although the model is vulnerable to adversarial attack, it can be aware of the attack, reflecting on the increasing of the entropy on the attacked test set. Table 2: Regression results of MNN when sharing the covariance within a mini-batch during training. We report the average log-likelihood (LL) on the UCI regression datasets over random train/test splits. MNN (H) represents Heaviside MNN, and MNN (R) represents ReLU MNN. The data shows the mean (standard deviation) over 20 runs. | boston | concrete | kin8 | energy | power | protein | wine | yacht
---|---|---|---|---|---|---|---|---
MNN (R) | -2.46 (0.01) | -3.58 (0.00) | 1.09 (0.00) | -2.33 (0.01) | -2.92 (0.00) | -0.98 (0.00) | -0.72 (0.01) | -0.34 (0.07)
MNN (H) | -2.67 (0.04) | -3.28 (0.02) | 0.73 (0.00) | -1.45 (0.00) | -2.95 (0.00) | -0.92 (0.00) | -0.89 (0.01) | -0.45 (0.03)
Table 3: Classification results. The test accuracy of MNNs trained using batch-wise covariance trick on image classification. MNN (R) means ReLU MNN, and MNN (H) represents Heaviside MNN.  | MNN (R) | MNN (H)
---|---|---
MNIST | 98.99 (0.00) | 99.00 (0.00)
FashionMNIST | 89.66 (0.00) | 90.13 (0.00)
CIFAR-10 | 70.75 (0.01) | 69.86 (0.01)
CIFAR-100 | 35.84 (0.02) | 40.21 (0.02)

## Uncertainty quantification on Heaviside MNNs

We report the uncertainty quantification results of Heaviside MNNs for
classification in Fig. 12, which is consistent with that of ReLU MNNs in Fig.
2.

Figure 12: Emerging covariance of Heaviside MNNs faithfully captures the
prediction uncertainty. (Mixed) Heaviside MNN trained on MNIST (CIFAR-10) for
image classification. a, The training of MNNs on MNIST (left) and CIFAR-10
(right), with the accuracy (red line) and entropy (orange and green for
correctly classified inputs and misclassified inputs respectively) on the test
set. b, The entropy of correctly classified and misclassified inputs in the
last two layers of MNIST (left) and CIFAR-10 (right), and the misclassified
inputs result in relatively higher entropy in both layers. c, The entropy
separability of correctly classified and misclassified inputs on MNIST (left)
and CIFAR-10 (right) increases as the layer index increases. d, The
relationship between the separability and test accuracy during training on
MNIST (left) and CIFAR-10 (right) was analyzed using linear regression. The
coefficients of determination of 0.6803 for MNIST and 0.6068 for CIFAR-10.

## More performance comparison with other probabilistic computation methods

Table 4: Performance on the regression problems of MNN without correlations. Log-likelihood of different models trained on the UCI regression datasets over random train/test splits. The data shows the mean (standard deviation) over 20 runs. MNN (R) means ReLU MNN, MNN (H) represents Heaviside MNN. dMMM represents MNN without considering the correlation.  | boston | concrete | kin8 | energy | power | protein | wine | yacht
---|---|---|---|---|---|---|---|---
MNN (R) | -2.44 (0.00) | -3.59 (0.01) | 0.92 (0.00) | -2.30 (0.01) | -2.92 (0.00) | -0.97 (0.00) | -0.75 (0.00) | -0.29 (0.06)
dMNN (R) | -2.78 (0.54) | -3.27 (0.00) | 0.67 (0.00) | -2.36 (0.03) | -3.00 (0.00) | -2.63 (0.25) | -0.77 (0.00) | -0.06 (0.30)
MNN (H) | -2.65 (0.03) | -3.28 (0.01) | 0.71 (0.00) | -1.31 (0.01) | -2.94 (0.00) | -0.92 (0.00) | -0.90 (0.01) | -0.40 (0.00)
dMNN (H) | -2.63 (0.04) | -3.18 (0.11) | 0.72 (0.00) | -1.46 (0.01) | -2.90 (0.00) | -0.88 (0.00) | -0.92 (0.01) | -0.25 (0.04)
Table 5: Comparison of performance on the classification problems. Comparison of test accuracy on image classification. The data shows the mean (standard deviation) over 20 runs. We compare MNNs (using ReLU or Heaviside activations, with or without correlations) with three representative uncertainty quantification methods. VarOut, Variational dropout [49]; DVI, deterministic variational inference [33]; VBP, Variance Back-Propagation [37].  | MNN (R) | dMNN (R) | MNN (H) | dMNN (H) | VarOut | DVI | VBP
---|---|---|---|---|---|---|---
MNIST | 99.04 (0.00) | 99.03 (0.02) | 99.01 (0.00) | 99.05 (0.00) | 98.88 (0.05) | 99.03 (0.06) | 99.15 (0.06)
FashionMNIST | 89.85 (0.03) | 90.05 (0.00) | 90.08 (0.00) | 89.93 (0.00) | 89.19 (0.24) | 89.67 (0.07) | 89.90 (0.16)
CIFAR-10 | 71.04 (0.01) | 70.57 (0.48) | 70.36 (0.01) | 69.63 (0.02) | 66.60 (1.06) | 64.85 (1.13) | 68.67 (1.00)
CIFAR-100 | 35.56 (0.01) | 31.01 (0.03) | 40.30 (0.00) | 40.04 (0.01) | 37.15 (1.43) | 33.79 (1.14) | 39.03 (1.61)

We show that our MNNs can achieve comparable or even superior results than
representative probabilistic computation methods in terms of model
performance. For regression tasks, we train MNNs on UCI regression datasets
[30]. Note that, the MSE loss does not involve output covariance. We consider
MNN using different activations (ReLU and Heaviside function). To investigate
the effect of correlations, we construct dMMN, in which we ignore all the
correlations between neurons at each layer. We use log-likelihood (LL) on the
test set as the metrics. As shown in Table. 4, we compare the log-likelihood
of various methods with our MMN, dMNN and show that both MNN and dMNN
outperform other methods in most datasets. Additionally, MNN is superior to
dMNN for both ReLU and Heaviside functions in most cases, implying that
correlation is meaningful. For classification tasks, we train MNNs through
SMUC on MNIST, FashionMNIST [50], CIFAR-10 and CIFAR-100 [23] respectively. We
use modified LeNet5 as the architecture of MNNs, following the settings in
[37]. The training loss is the cross entropy, calculated by the prediction
mean and the ground-truth label. As shown in Table. 5, MNNs can achieve
comparable performance than other probabilistic computation methods.

## Theoretical interpretation of the relationship between sensitivity and
entropy

Figure 13: Entropy of the stochastic network captures the sensitivity. a, The
derivation of the sigmoid function reaches its maximum at $0$. b, We consider
a two-dimensional case for illustration. We set $\mathbf{w}=(1,0.5),b=0,C=I$.
We calculate the prediction entropy at each $\mathbf{x}\in\mathbb{R}^{2}$. The
closer a point to the decision boundary (black line), the higher the
prediction entropy (heatmap).

For illustrative purposes, we analyze a simple binary logistic regression
problem to demonstrate how the sensitivity related to the prediction entropy
in the stochastic network. Consider that a binary classification problem, in
which $\mathbf{x}\in\mathbb{R}^{d}$ belongs to the class A if and only if
$\mathbf{w}^{T}\mathbf{x}+b<0$, otherwise it belongs the class B, where
$\mathbf{w}\in\mathbb{R}^{d}$ and $b\in\mathbb{R}$ are ground-truth
parameters. We construct a stochastic network for logistic regression
$y=\Phi(\mathbf{w}^{T}(\mathbf{x}+\mathbf{z})+b),\quad\mathbf{z}\sim\mathcal{N}(\mathbf{0},C)$,
where $\Phi(s)=\frac{1}{1+\exp(-s)}$ is the sigmoid function,
$C\in\mathbb{R}^{d\times d}$ is the covariance matrix of the noise vector
$\mathbf{z}$. Although there is no analytical form of the prediction
distribution, we can estimate it using a linear approximation

$\displaystyle
y=\Phi(\mathbf{w}^{T}(\mathbf{x}+\mathbf{z})+b)\approx\Phi(\mathbf{w}^{T}\mathbf{x}+b)+{\Phi}^{\prime}(\mathbf{w}^{T}\mathbf{x}+b)\mathbf{w}^{T}\mathbf{z},$
(42)

where ${\Phi}^{\prime}(s)$ is the derivative of ${\Phi}(s)$. Hence, we have

$\displaystyle
p(y|\mathbf{x})\approx\mathcal{N}(\Phi(\mathbf{w}^{T}\mathbf{x}+b),({\Phi}^{\prime}(\mathbf{w}^{T}\mathbf{x}+b))^{2}\mathbf{w}^{T}C\mathbf{w}),$
(43)

and its entropy is calculated as

$\displaystyle H(p(y|\mathbf{x}))\approx\frac{1}{2}(1+\log
2\pi+\log[({\Phi}^{\prime}(\mathbf{w}^{T}\mathbf{x}+b))^{2}\mathbf{w}^{T}C\mathbf{w}]),$
(44)

which is dependent on $\mathbf{x}$. Noticed that
$d(\mathbf{x})=\frac{\mathbf{w}^{T}\mathbf{x}+b}{\left\|\mathbf{w}\right\|}$
is the (signed) distance from $\mathbf{x}$ to the decision boundary
$\mathbf{w}^{T}\mathbf{x}+b=0$. The approximated entropy can be rewritten as

$\displaystyle H(p(y|\mathbf{x}))\approx\frac{1}{2}(1+\log
2\pi+\log[{\Phi}^{\prime}(\left\|\mathbf{w}\right\|d(\mathbf{x}))^{2}\mathbf{w}^{T}C\mathbf{w}]).$
(45)

Additionally, notice that
${\Phi}^{\prime}(s)=\frac{1}{1+\exp(-s)}-\frac{1}{(1+\exp(-s))^{2}}$ reaches
its maximum at $s=0$, and monotonically decreases as $s$ goes far away from
$0$ (Fig. 13a). Therefore, when $\mathbf{x}$ closer to the decision boundary
(Fig. 13b), the entropy is higher. As a result, the entropy of the stochastic
model reflects the sensitivity.

## References

  * [1] Simon P Kelly and Redmond G O’Connell.  The neural processes underlying perceptual decision making in humans: recent progress and future directions.  Journal of Physiology-Paris, 109(1-3):27–37, 2015.
  * [2] Moloud Abdar, Farhad Pourpanah, Sadiq Hussain, Dana Rezazadegan, Li Liu, Mohammad Ghavamzadeh, Paul Fieguth, Xiaochun Cao, Abbas Khosravi, U Rajendra Acharya, et al.  A review of uncertainty quantification in deep learning: Techniques, applications and challenges.  Information Fusion, 76:243–297, 2021.
  * [3] Jakob Gawlikowski, Cedrique Rovile Njieutcheu Tassi, Mohsin Ali, Jongseok Lee, Matthias Humt, Jianxiang Feng, Anna Kruspe, Rudolph Triebel, Peter Jung, Ribana Roscher, et al.  A survey of uncertainty in deep neural networks.  arXiv e-prints, 2021.
  * [4] Ben Goertzel.  Human-level artificial general intelligence and the possibility of a technological singularity: A reaction to ray kurzweil’s the singularity is near, and mcdermott’s critique of kurzweil.  Artificial Intelligence, 171(18):1161–1173, 2007.
  * [5] Mark Sprevak.  Two kinds of information processing in cognition.  Review of Philosophy and Psychology, 11(3):591–611, 2020.
  * [6] John Von Neumann and Ray Kurzweil.  The computer and the brain.  Yale university press, 2012.
  * [7] Dominik Seuß.  Bridging the gap between explainable ai and uncertainty quantification to enhance trustability.  arXiv preprint, 2021.
  * [8] Zoubin Ghahramani.  Probabilistic machine learning and artificial intelligence.  Nature, 521(7553):452–459, 2015.
  * [9] Michael I Jordan, Zoubin Ghahramani, Tommi S Jaakkola, and Lawrence K Saul.  An introduction to variational methods for graphical models.  Machine learning, 37:183–233, 1999.
  * [10] Arnaud Doucet, Nando De Freitas, Neil James Gordon, et al.  Sequential Monte Carlo methods in practice.  Springer, 2001.
  * [11] John Paisley, David Blei, and Michael Jordan.  Variational bayesian inference with stochastic search.  arXiv preprint, 2012.
  * [12] Diederik P Kingma and Max Welling.  Auto-encoding variational bayes.  arXiv preprint, 2013.
  * [13] Yarin Gal and Zoubin Ghahramani.  Dropout as a bayesian approximation: Representing model uncertainty in deep learning.  In international conference on machine learning, pages 1050–1059. PMLR, 2016.
  * [14] Balaji Lakshminarayanan, Alexander Pritzel, and Charles Blundell.  Simple and scalable predictive uncertainty estimation using deep ensembles.  Advances in neural information processing systems, 30, 2017.
  * [15] Mattias Teye, Hossein Azizpour, and Kevin Smith.  Bayesian uncertainty estimation for batch normalized deep networks.  In International Conference on Machine Learning, pages 4907–4916. PMLR, 2018.
  * [16] Aryan Mobiny, Pengyu Yuan, Supratik K Moulik, Naveen Garg, Carol C Wu, and Hien Van Nguyen.  Dropconnect is effective in modeling uncertainty of bayesian deep networks.  Scientific reports, 11(1):1–14, 2021.
  * [17] Adrián Ponce-Alvarez, Alexander Thiele, Thomas D Albright, Gene R Stoner, and Gustavo Deco.  Stimulus-dependent variability and noise correlations in cortical mt neurons.  Proceedings of the National Academy of Sciences, 110(32):13162–13167, 2013.
  * [18] Stefano Panzeri, Monica Moroni, Houman Safaai, and Christopher D Harvey.  The structures and functions of correlations in neural population codes.  Nature Reviews Neuroscience, 23:551–567, 2022.
  * [19] Wenlian Lu, Enrico Rossoni, and Jianfeng Feng.  On a gaussian neuronal field model.  NeuroImage, 52(3):913–933, 2010.
  * [20] Crispin W Gardiner et al.  Handbook of stochastic methods, volume 3.  springer Berlin, 1985.
  * [21] Allan Pinkus.  Approximation theory of the mlp model in neural networks.  Acta numerica, 8:143–195, 1999.
  * [22] Silvere Bonnabel.  Stochastic gradient descent on riemannian manifolds.  IEEE Transactions on Automatic Control, 58(9):2217–2229, 2013.
  * [23] Alex Krizhevsky, Geoffrey Hinton, et al.  Learning multiple layers of features from tiny images.  2009\.
  * [24] Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner.  Gradient-based learning applied to document recognition.  Proceedings of the IEEE, 86(11):2278–2324, 1998.
  * [25] Matus Telgarsky.  Benefits of depth in neural networks.  In Conference on learning theory, pages 1517–1539. PMLR, 2016.
  * [26] Yaroslav Bulatov.  Notmnist dataset.  Google (Books/OCR), Tech. Rep.[Online]. Available: http://yaroslavvb. blogspot. it/2011/09/notmnist-dataset. html, 2, 2011.
  * [27] Dan Hendrycks and Kevin Gimpel.  A baseline for detecting misclassified and out-of-distribution examples in neural networks.  In International Conference on Learning Representations, 2016.
  * [28] Christian Szegedy, Wojciech Zaremba, Ilya Sutskever, Joan Bruna, Dumitru Erhan, Ian Goodfellow, and Rob Fergus.  Intriguing properties of neural networks.  arXiv preprint, 2013.
  * [29] Ian J Goodfellow, Jonathon Shlens, and Christian Szegedy.  Explaining and harnessing adversarial examples.  stat, 1050:20, 2015.
  * [30] Arthur Asuncion and David Newman.  Uci machine learning repository, 2007.
  * [31] José Miguel Hernández-Lobato and Ryan Adams.  Probabilistic backpropagation for scalable learning of bayesian neural networks.  In International conference on machine learning, pages 1861–1869. PMLR, 2015.
  * [32] Jannik Schmitt and Stefan Roth.  Sampling-free variational inference for neural networks with multiplicative activation noise.  In Pattern Recognition: 43rd DAGM German Conference, DAGM GCPR 2021, Bonn, Germany, September 28–October 1, 2021, Proceedings, pages 33–47. Springer, 2022.
  * [33] Anqi Wu, Sebastian Nowozin, Edward Meeds, Richard E Turner, Jose Miguel Hernandez-Lobato, and Alexander L Gaunt.  Deterministic variational inference for robust bayesian neural networks.  arXiv preprint, 2018.
  * [34] Xiaoyu Cao and Neil Zhenqiang Gong.  Mitigating evasion attacks to deep neural networks via region-based classification.  In Proceedings of the 33rd Annual Computer Security Applications Conference, pages 278–287, 2017.
  * [35] Elad Schneidman, Michael J Berry, Ronen Segev, and William Bialek.  Weak pairwise correlations imply strongly correlated network states in a neural population.  Nature, 440(7087):1007–1012, 2006.
  * [36] David Dahmen, Hannah Bos, and Moritz Helias.  Correlated fluctuations in strongly coupled binary networks beyond equilibrium.  Physical Review X, 6(3):031024, 2016.
  * [37] Manuel Haußmann, Fred A Hamprecht, and Melih Kandemir.  Sampling-free variational inference of bayesian neural networks by variance backpropagation.  In Uncertainty in Artificial Intelligence, pages 563–573. PMLR, 2020.
  * [38] Jure Zbontar, Li Jing, Ishan Misra, Yann LeCun, and Stéphane Deny.  Barlow twins: Self-supervised learning via redundancy reduction.  In International Conference on Machine Learning, pages 12310–12320. PMLR, 2021.
  * [39] Bruno B Averbeck, Peter E Latham, and Alexandre Pouget.  Neural correlations, population coding and computation.  Nature Reviews Neuroscience, 7(5):358–366, 2006.
  * [40] Ian Goodfellow, Yoshua Bengio, and Aaron Courville.  Deep learning.  MIT press, 2016.
  * [41] Jaime De La Rocha, Brent Doiron, Eric Shea-Brown, Krešimir Josić, and Alex Reyes.  Correlation between neural spike trains increases with firing rate.  Nature, 448(7155):802–806, 2007.
  * [42] Jianfeng Feng, Yingchun Deng, and Enrico Rossoni.  Dynamics of moment neuronal networks.  Physical Review E, 73(4):041906, 2006.
  * [43] Qi Yang, Zhu Zhichao, Wei Yiming, Cao Lu, Wang Zhigang, Lu Wenlian, and Feng Jianfeng.  Toward spike-based stochastic neural computing.  arXiv preprint, 2023.
  * [44] Hengyuan Ma, Yang Qi, Pulin Gong, Wenlian Lu, and Jianfeng Feng.  Dynamics of bump attractors in neural circuits with emergent spatial correlations.  arXiv preprint, 2022.
  * [45] Yang Qi.  An efficient numerical algorithm for the moment neural activation.  arXiv preprint, 2022.
  * [46] Jacob Steinhardt and Percy S Liang.  Unsupervised risk estimation using only conditional independence structure.  Advances in Neural Information Processing Systems, 29, 2016.
  * [47] Diederik P Kingma and Jimmy Ba.  Adam: A method for stochastic optimization.  arXiv preprint, 2014.
  * [48] Karen Simonyan and Andrew Zisserman.  Very deep convolutional networks for large-scale image recognition.  arXiv preprint, 2014.
  * [49] Durk P Kingma, Tim Salimans, and Max Welling.  Variational dropout and the local reparameterization trick.  Advances in neural information processing systems, 28, 2015.
  * [50] Han Xiao, Kashif Rasul, and Roland Vollgraf.  Fashion-mnist: a novel image dataset for benchmarking machine learning algorithms.  arXiv preprint, 2017.
  * [51] Nicolas Boumal.  An introduction to optimization on smooth manifolds.  Cambridge University Press, 2023.
  * [52] Max Welling and Yee W Teh.  Bayesian learning via stochastic gradient langevin dynamics.  In Proceedings of the 28th international conference on machine learning (ICML-11), pages 681–688, 2011.
  * [53] Halsey Lawrence Royden and Patrick Fitzpatrick.  Real analysis, volume 32.  Macmillan New York, 1988.

Supplementary Information
Probabilistic computation and uncertainty quantification with emerging
covariance
Hengyuan Ma1, Yang Qi1,2,3, Li Zhang4, Jie Zhang1,2, Wenlian Lu1,2,5,6,7,8,
Jianfeng Feng1,2,$\ast$

1 Institute of Science and Technology for Brain-inspired Intelligence, Fudan
University, Shanghai 200433, China
2 Key Laboratory of Computational Neuroscience and Brain-Inspired Intelligence
(Fudan University), Ministry of Education, China
3 MOE Frontiers Center for Brain Science, Fudan University, Shanghai 200433,
China
4 School of Data Science, Fudan University, Shanghai 200433, China
5 School of Mathematical Sciences, Fudan University, No. 220 Handan Road,
Shanghai, 200433, Shanghai, China
6 Shanghai Center for Mathematical Sciences, No. 220 Handan Road, Shanghai,
200433, Shanghai, China
7 Shanghai Key Laboratory for Contemporary Applied Mathematics, No. 220 Handan
Road, Shanghai, 200433, Shanghai, China
8 Key Laboratory of Mathematics for Nonlinear Science, No. 220 Handan Road,
Shanghai, 200433, Shanghai, China
$\ast$<EMAIL_ADDRESS>
##  1 Derivation of the moment activations (MAs)

In this section, we derive moment activations (MAs) from a general stochastic
network (Eq. (S1)). First, we employ a mean field approximation to simplify
the network dynamics and derive expressions that describe the relationship
between the mean and variance of the output signals based on the input signals
of the network. Subsequently, we utilize the linear response approach to
estimate the relationship between the covariance of the output signals and the
input signals of the network. Then, we return to the specific stochastic
network used in the main paper (Eq. (5)) and provide specific expressions for
the corresponding MAs.

###  1.1 Mean field approximation for a nonlinear stochastic network

Consider a nonlinear stochastic neural network

$\displaystyle\frac{d\mathbf{x}}{dt}=\mathbf{f}(\mathbf{x})+W\mathbf{v}+\mathbf{v}_{s},\quad\mathbf{v}=\mathbf{h}(\mathbf{x})$
(S1)

where $\mathbf{x}\in\mathbb{R}^{N}$ is the state of the neurons,
$\mathbf{f}(\cdot)=(f_{i}(\cdot))_{i=1}^{N}$ is an element-wise function which
governs the dynamics of $\mathbf{x}$, $\mathbf{v}$ is the signals produced by
the neurons, which are transmitted to other neurons through the weight
$W=(w_{ij})_{i,j=1}^{N}$, $\mathbf{h}(\cdot)=(h_{i}(\cdot))_{i=1}^{N}$ is an
element-wise signal function that maps the state of the neurons $\mathbf{x}$
to its output signal $\mathbf{v}$, $\mathbf{v}_{s}$ is the stochasticity of
the neurons, which is a Brownian motion with drift
$\bm{\mu}_{s}\in\mathbb{R}^{N}$ and diffusion coefficient
$\Lambda\in\mathbb{R}^{N\times N}$.

We assume that the system reaches a stationary distribution, i.e., the
distribution of $x_{i}$ and $v_{i}$ are independent of time $t$. Although this
stationary distribution can not be solved analytically in general, we can
approximate it via a mean field approximation. Concretely, we approximate the
weighted received signals $W\mathbf{v}+\mathbf{v}_{s}$ by a Brownian motion
$\bm{\xi}$. We further assume that the signals $\mathbf{v}$ are independent of
the neural stochasticity $\mathbf{v}_{s}$. Denote the mean and covariance of
$\mathbf{v}$ as

$\displaystyle\bm{\mu}=\mathbb{E}[\mathbf{v}],\quad
C=\mathbb{E}[(\mathbf{v}-\bm{\mu})(\mathbf{v}-\bm{\mu})^{\mathrm{T}}].$ (S2)

Then the drift $\bar{\bm{\mu}}$, the diffusion coefficient $\bar{C}$ of
$\bm{\xi}$ as

$\displaystyle\bar{\bm{\mu}}=W\bm{\mu}+\bm{\mu}_{s},\quad\bar{C}=WCW^{\mathrm{T}}+\Lambda.$
(S3)

Under this mean field approximation, Eq. (S1) can be rewritten as

$\displaystyle\frac{d\mathbf{x}}{dt}=\mathbf{f}(\mathbf{x})+\bm{\xi}.$ (S4)

Denote the stationary distribution of each $x_{i}$ as
$p(x_{i}|\bar{\mu}_{i},\bar{\sigma}^{2}_{i})$ with
$\bar{\sigma}^{2}_{i}=\bar{C}_{ii}$ and that of each pair $(x_{i},x_{j})$ as
$p(x_{i},x_{j}|\bar{\mu}_{i},\bar{\mu}_{j},\bar{\sigma}^{2}_{i},\bar{\sigma}^{2}_{j},\rho_{ij})$,
where $\rho_{ij}=\bar{C}_{ij}/(\bar{\sigma}_{i}\bar{\sigma}_{j})$ is the
correlation coefficient of the input signals of two neurons. The distribution
$p(x_{i}|\bar{\mu}_{i},\bar{\sigma}^{2}_{i})$ is the solution to the time-
independent Fokker-Planck equation [20]

$\displaystyle\frac{\partial}{\partial
x_{i}}\big{(}(f_{i}(x_{i})+\bar{\mu}_{i})p(x_{i})\big{)}=\bar{\sigma}^{2}_{i}\frac{\partial^{2}}{\partial
x_{i}^{2}}p(x_{i}).$ (S5)

and the joint distribution
$p(x_{i},x_{j}|\bar{\mu}_{i},\bar{\mu}_{j},\bar{\sigma}^{2}_{i},\bar{\sigma}^{2}_{j},\bar{C}_{ij})$
is the solution to the time-independent Fokker-Planck

$\displaystyle(\frac{\partial}{\partial
x_{i}}(f_{i}(x_{i})+\bar{\mu}_{i})+\frac{\partial}{\partial
x_{j}}(f_{i}(x_{j})+\bar{\mu}_{j}))p(x_{i},x_{i})=(\bar{\sigma}^{2}_{i}\frac{\partial^{2}}{\partial
x_{i}^{2}}+\bar{\sigma}^{2}_{j}\frac{\partial^{2}}{\partial
x_{i}^{2}}+\bar{C}_{ij}\frac{\partial^{2}}{\partial x_{i}\partial
x_{j}})p(x_{i},x_{j}).$ (S6)

After solving Eq. (S5)-(S6), we can calculate the first two cumulants of the
output signals $\mathbf{v}$

$\displaystyle\mu_{i}=\mathbb{E}[v_{i}]=\int
h_{i}(x_{i})p(x_{i}|\bar{\mu}_{i},\bar{C}_{i})dx_{i},$ (S7) $\displaystyle
C_{ij}=\mathrm{cov}[v_{i},v_{j}]=\int(h_{i}(x_{i})-\mathbb{E}[v_{i}])(h_{j}(x_{j})-\mathbb{E}[v_{j}])p(x_{i},x_{j}|\bar{\mu}_{i},\bar{\mu}_{j},\bar{\sigma}^{2}_{i},\bar{\sigma}^{2}_{j},\bar{C}_{ij})dx_{i}dx_{j}.$
(S8)

We call Eq. (S7)-(S8) as the mean activation and covariance activation
respectively, and they constitute the moment activations (MAs).

The computation complexity of the double integration in the covariance
activation (Eq.(S8)) is high. To address this, we simplify the covariance
activation using the linear response theory [41]. The covariance between the
output signals of two neurons $\mathrm{cov}[v_{i},v_{j}]$ can be considered as
a function of $\rho_{ij}$, denoted as $C_{v,ij}(\rho_{ij})$, Assuming that
$\rho_{ij}$ is small, we can approximate $C_{v,ij}(\rho_{ij})$ using a first-
order Taylor expansion

$\displaystyle
C_{v,ij}(\rho_{ij})\sim\frac{d}{d\rho_{ij}}C_{v,ij}(0)\rho_{ij},$ (S9)

where the derivative $\frac{d}{d\rho_{ij}}C_{v,ij}(0)$ can be considered as a
sort of susceptibility describing how two neurons transfer weak correlated
input signals [41].

###  1.2 Moment activations for Ornstein–Uhlenbeck processes

We now develop the MAs for the Ornstein–Uhlenbeck system (Eq. (5)) in the main
paper, where $\mathbf{f}(\mathbf{x})=-\mathbf{x}$. In this case, the
stationary distribution $p(x_{i}|\bar{\mu}_{i},\bar{\sigma}^{2}_{i})$ is a
Gaussian distribution with density

$\displaystyle\frac{1}{\sqrt{2\pi}\bar{\sigma}_{i}}\exp\\{-\frac{(x_{i}-\bar{\mu}_{i})^{2}}{2\bar{\sigma}_{i}^{2}}\\}$
(S10)

and
$p(x_{i},x_{j}|\bar{\mu}_{i},\bar{\mu}_{i},\bar{\sigma}^{2}_{i},\bar{\sigma}^{2}_{j},\bar{C}_{ij})$
is a two-dimensional Gaussian distribution with density

$\displaystyle\frac{1}{2\pi\sqrt{1-\rho_{ij}^{2}}\bar{\sigma}_{i}\bar{\sigma}_{j}}\exp(-\frac{1}{2(1-\rho_{ij}^{2})\bar{\sigma}_{i}^{2}\bar{\sigma}_{j}^{2}}(x_{i}-\bar{\mu}_{i},x_{j}-\bar{\mu}_{j})\begin{pmatrix}\bar{\sigma}_{j}^{2}&-\rho_{ij}\bar{\sigma}_{i}\bar{\sigma}_{j}\\\
-\rho_{ij}\bar{\sigma}_{i}\bar{\sigma}_{j}&\bar{\sigma}_{i}^{2}\end{pmatrix}(x_{i}-\bar{\mu}_{i},x_{j}-\bar{\mu}_{j})^{\mathrm{T}}).$
(S11)

For convenience, we will denote $\bar{\sigma}^{2}_{i}$ as $\bar{C}_{i}$ Using
Eq. (S7)-(S8), we have the mean and variance of the output signals $v_{i}$

$\displaystyle m_{v,i}(\bar{\mu}_{i},\bar{C}_{i}):=\mathbb{E}[v_{i}]=\int
h_{i}(x_{i})p(x_{i}|\bar{\mu}_{i},\bar{C}_{i})dx_{i},$ (S12) $\displaystyle
C_{v,i}(\bar{\mu}_{i},\bar{C}_{i}):=\mathbb{E}[(v_{i}-\mathbb{E}[v_{i}])^{2}]=\int
h_{i}^{2}(x_{i})p(x_{i}|\bar{\mu}_{i},\bar{C}_{i})dx_{i}-m_{v}(\bar{\mu}_{i},\bar{C}_{i})^{2},$
(S13)

Using Eq. (S9), the approximated covariance is

$\displaystyle
C_{v,ij}(\bar{\mu}_{i},\bar{\mu}_{j},\bar{C}_{i},\bar{C}_{j},\bar{C}_{ij})\approx\frac{1}{2\pi}\int
h_{i}(\sigma_{i}x_{i}+\mu_{i})\exp(-\frac{1}{2}x_{i}^{2})x_{i}dx_{i}\int
h_{j}(\sigma_{j}x_{j}+\mu_{j})\exp(-\frac{1}{2}x_{j}^{2})x_{j}dx_{j}\rho_{ij},$
(S14)

which can be rewritten as

$\displaystyle
C_{v,ij}(\bar{\mu}_{i},\bar{C}_{i},\bar{\mu}_{j},\bar{C}_{j},c)=\chi_{i}(\bar{\mu}_{i},\bar{C}_{i})\chi_{j}(\bar{\mu}_{j},\bar{C}_{j})\rho_{ij},$
(S15)

where

$\displaystyle\chi_{i}(\bar{\mu}_{i},\bar{C}_{i})=\frac{1}{\sqrt{2\pi}}\int
h_{i}(\sqrt{\bar{C}_{i}}x+\bar{\mu}_{i})e^{-\frac{x^{2}}{2}}xdx.$ (S16)

##  2 Theoretical properties of supervised mean unsupervised covariance (SMUC)

We train an MNN $\tilde{M}$ through SMUC for learning the output mean and
covariance of the ground-truth MNN $M_{GT}$. We prove the following three
theoretical properties about the SMUC under mild conditions:

  * •

SMUC is equivalent to the stochastic Riemannian gradient descent, and the
parameter of the $\tilde{M}$ converges to the point where the loss gradient is
vanished(Thm. 1),

  * •

$\tilde{M}$ can learn ground-truth mean and variance at each layer even if the
ground-truth hyper-parameters $\sigma_{l}$ are unknown (Thm. 2),

  * •

As the number of the layers goes to infinity, the output covariance of
$\tilde{M}$ converges to that of $M_{GT}$ (Thm. 3).

We begin by analyzing the case where the network has only one nonlinear layer
with Heaviside activations in Sec. 2.1-2.2, and then generalize the results to
the multiple layer cases in Sec. 2.3. Next, we analyze the covariance of MNN
in Sec. 2.4. Finally, we generalize the above results to case with general
nonlinear activation functions in Sec. 2.5.

Consider the network with one nonlinear layer as follows

$\displaystyle\begin{aligned}
\frac{d\mathbf{x}^{(1)}}{dt}&=-\mathbf{x}^{(1)}+\mathbf{x}+\sqrt{2}C_{\mathrm{in}}^{\tfrac{1}{2}}\bm{\xi}^{(1)},\quad\mathbf{v}^{(1)}=\mathbf{x}^{(1)},\\\
\frac{d\mathbf{x}^{(2)}}{dt}&=-\mathbf{x}^{(2)}+W\mathbf{v}^{(1)}+\mathbf{b}+\sqrt{2}\sigma\bm{\xi}^{(2)},\quad\mathbf{v}^{(2)}=\bm{\Theta}(\mathbf{x}^{(2)})\\\
\frac{d\mathbf{y}}{dt}&=-\mathbf{y}+\mathbf{v}^{(2)}\end{aligned}$ (S17)

where $\mathbf{x}=(x_{i})_{i=1}^{m}\in\mathbb{R}^{m}$ is the input,
$\mathbf{y}=(y_{k})_{k=1}^{n}\in\mathbb{R}^{n}$ is the output,
$\bm{\Theta}(\cdot)$ is the element-wise Heaviside function,
$W=(\mathbf{w}_{k}^{\mathrm{T}})_{k=1}^{n}$ with
$\mathbf{w}_{k}^{\mathrm{T}}\in\mathbb{R}^{m}$ are the ground-truth weights,
$\mathbf{b}=(b_{k})_{k=1}^{n}\in\mathbb{R}^{n}$ are the ground-truth biases,
$\bm{\xi}^{(l)}$ are unit Gaussian white noises, $\sigma^{2}$ is the constant
ground-truth diffusion coefficient, and $C_{\mathrm{in}}\in\mathbb{R}^{m\times
m}$ is the unknown diffusion matrix (positive definite). We denote the
corresponding MNN of this stochastic neural network as $M_{GT}$. Note that
here the covariance matrix of the noise added to the input $\mathbf{x}$ can be
arbitrary positive definite matrix, no need to be supposed as a scalar matrix
as that in Eq. (5) in the main paper. In this way, the result derived here can
be directly generalized to the cases of multiple nonlinear layer later. Under
the mean field approximation, we simplify the above stochastic via mean field
approximation and obtain

$\displaystyle
y_{k}=\Theta(\mathbf{w}_{k}^{\mathrm{T}}[\mathbf{x}+C_{\mathrm{in}}^{\tfrac{1}{2}}\mathbf{z}]+b_{k}+\sigma
u_{k}),\quad\mathbf{z}\sim\mathcal{N}(\mathbf{0},I_{m}),\quad
u_{k}\sim\mathcal{N}(0,1),\quad k=1,\ldots,n.$ (S18)

For convenience, we introduce an augmented input
$\mathbf{x}^{\prime}=(\mathbf{x},1)\in\mathbb{R}^{m+1}$, augmented weights
$\mathbf{w}^{\prime}_{k}=(\mathbf{w}_{k},b_{k})\in\mathbb{R}^{m+1}$ which
include the bias terms as their component, augmented covariance matrix

$\displaystyle\begin{pmatrix}C_{\mathrm{in}}&\mathbf{0}\\\
\mathbf{0}^{\mathrm{T}}&0\end{pmatrix},$ (S19)

and augmented noise term $\mathbf{z}^{\prime}=(\mathbf{z},z)$ with
$z\sim\mathcal{N}(0,1)$ and write

$\displaystyle
y_{k}=\Theta(\mathbf{w}_{k}^{\mathrm{T}}[\mathbf{x}+C_{\mathrm{in}}^{\tfrac{1}{2}}\mathbf{z}]+\sigma
u_{k}),\quad\mathbf{z}\sim\mathcal{N}(\mathbf{0},I_{m}),\quad
u_{k}\sim\mathcal{N}(0,1),\quad k=1,\ldots,n,$ (S20)

where we use the original notations and do not increase the dimension $m$ by
$1$ for convenient. Using the definition of the moment mapping and Eq. (S20),
we have the ground-truth mean and covariance (variance) are

$\displaystyle\begin{aligned}
&\mu_{y,k}(\mathbf{x})=m_{v}(\bar{\mu}_{k},\bar{C}_{k}),\quad
C_{y,k}(\mathbf{x})=C_{v}(\bar{\mu}_{k},\bar{C}_{k}),\quad k=1,\ldots,n,\\\
&C_{y,kj}(\mathbf{x})=\chi(\bar{\mu}_{k},\bar{C}_{k})\chi(\bar{\mu}_{j},\bar{C}_{j})c_{kj},\quad
k,j=1,\ldots,n,\quad k\neq j,\end{aligned}$ (S21)

with

$\displaystyle\bar{\mu}_{k}=\mathbf{w}_{k}^{\mathrm{T}}\mathbf{x},\quad\bar{C}_{kj}=\mathbf{w}_{k}^{\mathrm{T}}C_{\mathrm{in}}\mathbf{w}_{j}+\delta_{kj}\sigma^{2},\quad\bar{C}_{k}=\bar{C}_{kk},\quad
c_{kj}=\frac{\bar{C}_{kj}}{\sqrt{\bar{C}_{k}\bar{C}_{j}}},$ (S22)

where $\delta_{kj}=1$ if $k=j$ else $\delta_{kj}=0$. To learn the mapping
$\mathbf{x}\mapsto q(\mathbf{y}|\mathbf{x})$ (Eq. (S18)) up to its second
cumulants, we construct a single layer MNN denoted as $\tilde{M}$ in which the
feed-forward propagation for an input $\mathbf{x}$ is calculated as

$\displaystyle\begin{aligned}
&\tilde{m}_{y,k}(\mathbf{x})=m_{v}(\tilde{\mu}_{k},\tilde{C}_{k}),\quad\tilde{C}_{y,k}(\mathbf{x})=C_{v}(\tilde{\mu}_{k},\tilde{C}_{k}),\quad
k=1,\ldots,n,\\\
&\tilde{C}_{y,kj}(\mathbf{x})=\chi(\tilde{\mu}_{k},\tilde{C}_{k})\chi(\tilde{\mu}_{j},\tilde{C}_{j})\tilde{c}_{kj},\quad
k,j=1,\ldots,n,\quad k\neq j,\end{aligned}$ (S23)

with

$\displaystyle\tilde{\mu}_{k}=\tilde{\mathbf{w}}_{k}^{\mathrm{T}}\mathbf{x},\quad\tilde{C}_{kj}=\tilde{\mathbf{w}}_{k}^{\mathrm{T}}\tilde{C}_{\mathrm{in}}\tilde{\mathbf{w}}_{j}+\delta_{kj}\tilde{\sigma}^{2},\quad\tilde{C}_{k}=\tilde{C}_{kk},\quad\tilde{c}_{kj}=\frac{\tilde{C}_{kj}}{\sqrt{\tilde{C}_{k}\tilde{C}_{j}}}.$
(S24)

where $\tilde{y}$ is the predicted $y$ using its mean and
$(\tilde{\mathbf{w}}_{1},\ldots,\tilde{\mathbf{w}}_{n})$ are the weights to be
learned, $\tilde{C}_{\mathrm{in}}$ and $\tilde{\sigma}$ are the manually set,
which may be different from $C_{\mathrm{in}}$ and $\sigma$. Note that in the
case of Eq. (13) of the main paper, we set $\tilde{C}_{\mathrm{in}}$ as
$\sigma_{1}^{2}I$ and $\sigma$ here corresponds to $\sigma_{2}$. The choice of
$\tilde{C}_{\mathrm{in}}$ is arbitrary, as long as it is positive definite.
Denote the objective function as loss
$\mathcal{L}(\mathbf{y},\bm{\tilde{\mu}}_{y})$, where
$\bm{\tilde{\mu}}_{y}=(\tilde{m}_{y,k})_{k=1}^{m}$. The training procedure of
SMUC is

$\displaystyle\tilde{\mathbf{w}}_{k,t+1}=\tilde{\mathbf{w}}_{k,t}-\gamma_{t}\frac{\partial\mathcal{L}}{\partial\tilde{\mathbf{w}}_{k,t}},\quad,k=1,\dots,n,\quad
t=1,\ldots,T,$ (S25)

where $\gamma_{t}$ is the learning rate, and $T$ is the step number of
training. Once again, we would like to emphasize that all second cumulants are
treated as constants during gradient calculation and do not participate in
backpropagation.

Using the MAs, we can approximate $\mathbf{y}$ by a Gaussian distribution
using its first two cumulants as
$\mathbf{y}\sim\bm{\mu}_{y}+C_{y}^{\tfrac{1}{2}}\mathbf{z},\quad\mathbf{z}\sim\mathcal{N}(\mathbf{0},I)$,
where $\bm{\mu}_{y}=(\mu_{y,k})_{k=1}^{n}$, and
$C_{y}=(C_{y,kj})_{k,j=1}^{n}$. Using this approximation, the objective can be
rewritten as

$\displaystyle\mathcal{L}(\mathbf{y},\bm{\tilde{\mu}}_{y})\sim\mathcal{L}(\bm{\mu}_{y},\bm{\tilde{\mu}}_{y})+J(\bm{\mu}_{y},\bm{\tilde{\mu}}_{y})C_{y}^{\tfrac{1}{2}}\mathbf{z},$
(S26)

where $J(\bm{\mu}_{y},\bm{\tilde{\mu}}_{y})$ is the Jacobian matrix of the
function $\mathcal{L}(\mathbf{y},\bm{\tilde{\mu}}_{y})$ at
$\mathbf{y}=\bm{\mu}_{y}$. Then the gradient at the $t$ step can be
approximately rewritten as the summation of the deterministic part
$\frac{\partial}{\partial\tilde{\mathbf{w}}_{k,t}}\mathcal{L}(\bm{\mu}_{y},\bm{\tilde{\mu}}_{y})$
with a stochastic term
$\bm{\xi}_{k,t}=\frac{\partial\mathcal{L}}{\partial\tilde{\mathbf{w}}_{k,t}}J(\bm{\mu}_{y},\bm{\tilde{\mu}}_{y})C_{y}^{\tfrac{1}{2}}\mathbf{z}$.

###  2.1 SMUC is equivalent to stochastic Riemannian gradient descent

We now prove the Thm. 1 in the single-layer case.

###### Proof.

We first prove the equivalence between SMUC and the stochastic Riemannian
gradient descent. The main thread of our proof is as follows. We leverage the
scalar property of the Heaviside MAs (Eq. (S27)) to reformulate the update
equation for the SMUC (Eq. (S25)). This reformulation enables us to align the
variances of each layer in the MNN with the variances in the ground-truth
network (Eq. (LABEL:eq:truth_network)) during the gradient descent process on
the mean. Furthermore, the alignment operation can be viewed as projecting the
parameters onto a manifold where the variance matches the ground-truth value.
By employing this alignment technique, we establish a correspondence between
the SMUC and stochastic Riemannian gradient descent.

Recall that according to Eq. (20), the mean activation of Heaviside function
satisfies the scalar property

$\displaystyle m_{v}(\mu,C)=m_{v}(a\mu,a^{2}C),\quad\forall
a>0,\mu\in\mathbb{R},C>0.$ (S27)

This suggests that the mean mapping remains unchanged when simultaneously
scaling the mean by a factor of $a$ and the covariance by a factor of $a^{2}$.
Denoting a scaling factor
$\alpha_{k,t}=\sqrt{\frac{\bar{C}_{k}}{\tilde{C}_{k,t}}}$, where
$\tilde{C}_{k,t}$ is defined in Eq. (S24), we can rewrite the loss
$\mathcal{L}$ as

$\displaystyle\mathcal{L}(\bm{\mu}_{y},\bm{\tilde{\mu}}_{y})=\mathcal{L}\big{(}(m_{v}(\bar{\mu}_{k},\bar{C}_{k}))_{k=1}^{n},(m_{v}(\tilde{\mathbf{w}}_{k,t}^{\mathrm{T}}\mathbf{x},\tilde{C}_{k}))_{k=1}^{n}\big{)}=\mathcal{L}\big{(}(m_{v}(\bar{\mu}_{k},\bar{C}_{k}))_{k=1}^{n},(m_{v}(\hat{\mathbf{w}}_{k,t}^{\mathrm{T}}\mathbf{x},\bar{C}_{k}))_{k=1}^{n}\big{)},$
(S28)

where we define $\hat{\mathbf{w}}_{k,t}=\alpha_{k,t}\tilde{\mathbf{w}}_{k,t}$.
Hence, the gradient can be rewritten as

$\displaystyle\frac{\partial}{\partial\tilde{\mathbf{w}}_{k,t}}\mathcal{L}(\bm{\mu}_{y},\bm{\tilde{\mu}}_{y})=\alpha_{k,t}\frac{\partial}{\partial\hat{\mathbf{w}}_{k,t}}\mathcal{L}\big{(}(m_{v}(\bar{\mu}_{k},\bar{C}_{k}))_{k=1}^{n},(m_{v}(\hat{\mathbf{w}}_{k,t}^{\mathrm{T}}\mathbf{x},\bar{C}_{k}))_{k=1}^{n}\big{)}.$
(S29)

As a result, we can write the update equation Eq. (S25) as

$\displaystyle\alpha_{k,t}\tilde{\mathbf{w}}_{k,t+1}$
$\displaystyle=\hat{\mathbf{w}}_{k,t}-\gamma_{t}[\alpha^{2}_{k,t}\frac{\partial}{\partial\hat{\mathbf{w}}_{k,t}}\mathcal{L}\big{(}(m_{v}(\bar{\mu}_{k},\bar{C}_{k}))_{k=1}^{n},(m_{v}(\hat{\mathbf{w}}_{k,t}^{\mathrm{T}}\mathbf{x},\bar{C}_{k}))_{k=1}^{n}\big{)}+\alpha_{k,t}\bm{\xi}_{k,t}]$
(S30)

If we track the weights $\hat{\mathbf{w}}_{k,t}$ instead of
$\tilde{\mathbf{w}}_{k,t}$ at each step $t$, the update of weight in Eq. (S25)
is equivalent to two steps

$\displaystyle\breve{\mathbf{w}}_{k,t+1}=\hat{\mathbf{w}}_{k,t}-\gamma_{t}(\alpha^{2}_{k,t}\frac{\partial}{\partial\hat{\mathbf{w}}_{k,t}}\mathcal{L}_{\mathrm{mean}}+\alpha_{k,t}\bm{\xi}_{k,t})$
(S31)
$\displaystyle\hat{\mathbf{w}}_{k,t+1}=\frac{\alpha_{k,t+1}}{\alpha_{k,t}}\breve{\mathbf{w}}_{k,t+1},$
(S32)

where $\mathcal{L}_{\mathrm{mean}}$ is the loss on the mean taking the second
cumulants as that of $M_{GT}$

$\displaystyle\mathcal{L}_{\mathrm{mean}}=\mathcal{L}\big{(}(m_{v}(\bar{\mu}_{k},\bar{C}_{k}))_{k=1}^{n},(m_{v}(\hat{\mathbf{w}}_{k,t}^{\mathrm{T}}\mathbf{x},\bar{C}_{k}))_{k=1}^{n}\big{)}.$
(S33)

Next, we show that Eq. (S31) and Eq. (S32) is equivalent to the stochastic
Riemannian gradient descent [51]. Define the manifold
$\mathcal{M}=\mathcal{M}_{1}\otimes\mathcal{M}_{2}\otimes\cdots\otimes\mathcal{M}_{n}\in\mathbb{R}^{m\times
n}$, where $\otimes$ is the Cartesian product, and

$\displaystyle\mathcal{M}_{k}=\\{\hat{\mathbf{w}}\in\mathbb{R}^{m},\hat{\mathbf{w}}^{\mathrm{T}}\tilde{C}_{\mathrm{in}}\hat{\mathbf{w}}+\alpha^{2}_{k,t}\tilde{\sigma}^{2}=\bar{C}_{k}\\},\quad
k=1,\ldots,n.$ (S34)

Due to the definition of $\alpha^{2}_{k,t}$, we always have
$\alpha^{2}_{k,t}\tilde{\sigma}^{2}<\bar{C}_{k}$, hence $\mathcal{M}_{k}$ is
not empty. We call $\mathcal{M}$ as iso-variance sphere, as the weights
$(\hat{\mathbf{w}}_{k})$ on it ensure that the mean activation
$m_{v}(\cdot,\cdot)$ receives the ground-truth variance $\bar{C}_{k}$, as
described in Eq. (S22). At each iteration, we first update the parameters to
minimize $\mathcal{L}_{\mathrm{mean}}$ through stochastic gradient Langevin
dynamics [52] in the Euclidean space, and then we project the parameters back
to the manifold $\mathcal{M}$ (Eq. (S32)). Thus, Eq. (S31) and Eq. (S32) is
equivalent to a stochastic Riemannian gradient descent on the iso-variance
sphere $\mathcal{M}$

$\displaystyle\hat{\omega}_{t+1}=R_{\hat{\omega}_{t}}[-\gamma_{t}(\alpha_{t}^{2}\odot\frac{\partial}{\partial\hat{\omega}_{t}}\mathcal{L}_{\mathrm{mean}}+\alpha_{t}\odot\bm{\xi}_{t})],$
(S35)

where $\hat{\omega}_{t}\in\mathbb{R}^{m\times n}$ is the concatenation of
$(\hat{\mathbf{w}}_{k,t},k=1,\ldots,n)$, $\bm{\xi}_{t}$ is the concatenation
of $(\bm{\xi}_{k,t},k=1,\ldots,n)$, $R_{\hat{\omega}}$ is a retraction on
$\mathcal{M}$ [51], $\alpha_{t}$ is the concatenation of
$(\alpha_{k,t},k=1,\ldots,n)$, and $\odot$ is the block-wise multiplication
defined as

$\displaystyle\alpha_{t}^{2}\odot\frac{\partial}{\partial\hat{\omega}_{t}}\mathcal{L}_{\mathrm{mean}}=(\alpha^{2}_{k,t}\frac{\partial}{\partial\hat{\mathbf{w}}_{k,t}}\mathcal{L}_{\mathrm{mean}},k=1,\ldots,n),\quad\alpha_{t}\odot\bm{\xi}_{t}=(\alpha_{k,t}\bm{\xi}_{k,t},k=1,\ldots,n).$
(S36)

The retraction $R_{\hat{\omega}}$ is a first-order approximation of the
corresponding exponential map $\exp_{\hat{\omega}}$ on $\mathcal{M}$ [51], and
the update Eq. (S35) is a first-order approximation of the following process

$\displaystyle\hat{\omega}_{t+1}=\exp_{\hat{\omega}_{t}}[-\gamma_{t}(\alpha^{2}_{k}\odot\frac{\partial}{\partial\hat{\omega}_{t}}\mathcal{L}_{\mathrm{mean}}+\alpha_{t}\odot\bm{\xi}_{t})].$
(S37)

Next, we prove the convergence of SMUC. We follow the methodology outlined in
[22] to investigate the convergence of Eq. (S37). Denote
$H(\hat{\omega}_{t},\bm{\xi}_{t})=\frac{\partial}{\partial\hat{\omega}_{t}}\mathcal{L}_{\mathrm{mean}}+\bm{\xi}_{t}\oslash\alpha_{t}$,
where $\oslash$ is the block-wise division defined as
$\bm{\xi}_{t}\oslash\alpha_{t}=(\frac{\bm{\xi}_{k,t}}{\alpha_{k,t}},k=1,\ldots,n)$.
Supposed that we select the proper learning rate $\\{\gamma_{t}\\}$ such that
it satisfies the following modified usual condition:

$\displaystyle\sum_{t}\gamma_{t}^{2}\alpha_{k,t}^{4}<+\infty,\quad\sum_{t}\gamma_{t}\alpha_{k,t}^{2}=+\infty,\quad\forall
k=1,\ldots,n.$ (S38)

Since $\mathcal{M}$ is bounded, there exists a compact set
$K\subset\mathbb{R}^{m}$, such that $\hat{\omega}_{t}\in K,\quad\forall t>0$.
Since $\mathcal{M}$ is the product of spheres $\mathcal{M}_{k}$ embedded in
the Euclidean space, $\mathcal{M}$ is a Hadamard manifold, and there is a
point $\nu\in\mathcal{M}$ and $S>0$, such that

$\displaystyle\inf_{\omega:D(\omega,\nu)>S}\langle\exp_{\omega}^{-1}(\nu),\frac{\partial}{\partial\omega}\mathcal{L}_{\mathrm{mean}}\rangle<0,$
(S39)

where $D(\omega,\nu)$ is the squared geodesic distance of $\omega$ and $\nu$
on $\mathcal{M}$. Additionally, there exists a lower bound on the sectional
curvature denoted by $\kappa<0$. Since $\mathcal{M}$ is bounded and
$\bm{\xi}_{k,t}$ are Gaussian, there exists a continuous function
$f:\mathcal{M}\rightarrow\mathbb{R}$ that satisfies

$\displaystyle\begin{aligned}
f(\omega)^{2}\geq\max\\{1,\mathbb{E}_{\xi}[\left\|H(\omega,\xi)\right\|^{2}(1+\sqrt{|\kappa|}(\sqrt{D(\omega,\nu)}+\left\|H(\omega,\xi)\right\|))],\mathbb{E}_{\xi}[(2\left\|H(\omega,\xi)\right\|\sqrt{D(\omega,\nu)}+\left\|H(\omega,\xi)\right\|^{2})^{2}]\\}.\end{aligned}$

Then, according to Thm.3 of [22], the iterations

$\displaystyle\hat{\omega}_{t+1}=\exp_{\hat{\omega}_{t}}[-\frac{\gamma_{t}}{f(\hat{\omega}_{t})}(\alpha^{2}_{t}\odot\frac{\partial}{\partial\hat{\omega}_{t}}\mathcal{L}_{\mathrm{mean}}+\alpha_{t}\odot\bm{\xi}_{t})].$
(S40)

converges to the point where the gradient
$\frac{\partial}{\partial\hat{\omega}_{t}}\mathcal{L}_{\mathrm{mean}}$
vanishes almost surely. ∎

We notice that when
$\tilde{\mathbf{w}}_{k,t}^{\mathrm{T}}\tilde{C}_{\mathrm{in}}\tilde{\mathbf{w}}_{k,t}>\bar{C}_{k}$,
the step-size $\alpha_{t,k}\gamma_{t}$ is reduced; when
$\tilde{\mathbf{w}}_{k,t}^{\mathrm{T}}\tilde{C}_{\mathrm{in}}\tilde{\mathbf{w}}_{k,t}<\bar{C}_{k}$,
the step-size $\alpha_{t,k}\gamma_{t}$ is amplified. Hence $\alpha_{k,t}$ can
be considered as a regularizer: when $\tilde{M}$ underestimates the variance
($\tilde{\mathbf{w}}_{k,t}^{\mathrm{T}}\tilde{C}_{\mathrm{in}}\tilde{\mathbf{w}}_{k,t}>\bar{C}_{k}$),
$\alpha_{k,t}$ slows down the gradient descent; when the $\tilde{M}$
overestimate the variance
($\tilde{\mathbf{w}}_{k,t}^{\mathrm{T}}\tilde{C}_{\mathrm{in}}\tilde{\mathbf{w}}_{k,t}<\bar{C}_{k}$),
$\alpha_{k,t}$ accelerates the gradient descent.

###  2.2 MNNs learn ground-truth mean and variance

We now prove the Thm. 2 in the single-layer case.

###### Proof.

After convergence, the trained weight $\tilde{\mathbf{w}}_{k}$ of $\tilde{M}$
satisfies

$\displaystyle\tilde{m}_{y,k}(\mathbf{x})=m_{v}(\tilde{\mathbf{w}}_{k}^{\mathrm{T}}\mathbf{x},\tilde{\mathbf{w}}_{k}^{T}\tilde{C}_{\mathrm{in}}\tilde{\mathbf{w}}_{k}+\tilde{\sigma}^{2})=m_{v}(\alpha_{k}\mathbf{w}_{k}^{\mathrm{T}}\mathbf{x},\mathbf{w}_{k}^{\mathrm{T}}C_{\mathrm{in}}\mathbf{w}_{k}+\sigma^{2})=m_{v}(\bar{\mu}_{k},\bar{C}_{k})=\mu_{y,k},$
(S41)

where $\alpha_{k}=\sqrt{\frac{\bar{C}_{k}}{\tilde{C}_{k}}}$. Therefore, the
$\tilde{M}$ can fit the mean of the ground-truth MNN $M_{GT}$. For the
variance activation of the Heaviside function (Eq. (21)), we also have the
scalar property

$\displaystyle C_{v}(\mu,C)=C_{v}(a\mu,a^{2}C),\quad\forall
a>0,\mu\in\mathbb{R},C>0.$ (S42)

Therefore, similar as Eq. (S41), we have

$\displaystyle\tilde{C}_{y,k}(\mathbf{x})=C_{v}(\tilde{\mathbf{w}}_{k}^{\mathrm{T}}\mathbf{x},\tilde{\mathbf{w}}_{k}^{T}\tilde{C}_{\mathrm{in}}\tilde{\mathbf{w}}_{k}+\tilde{\sigma}^{2})=C_{v}(\alpha_{k}\mathbf{w}_{k}^{\mathrm{T}}\mathbf{x},\mathbf{w}_{k}^{\mathrm{T}}C_{\mathrm{in}}\mathbf{w}_{k}+\sigma^{2})=C_{v}(\bar{\mu}_{k},\bar{C}_{k})=C_{y,k},$
(S43)

and the $\tilde{M}$ can also fit the variance of $M_{GT}$. ∎

###  2.3 Multiple-layer cases

We extend the previous results about that $\tilde{M}$ can learn the mean and
variance of $M_{GT}$ from the case of a single non-linear layer to the multi-
layer case. Consider the network described in Eq. (5), with
$\mathbf{h}(\cdot)$ is specific as $\bm{\Theta}(\cdot)$, and $\sigma_{1}$ is
replaced by $C_{\mathrm{in}}^{\tfrac{1}{2}}$. Under the mean field
approximation, we absorb all the bias terms and noise terms as we did in Sec.
2 to obtain the corresponding MNN $M_{GT}$

$\displaystyle\begin{aligned}
\mathbf{v}^{(1)}&=\mathbf{x}+C_{\mathrm{in}}^{\frac{1}{2}}\mathbf{z}_{\mathrm{in}}\\\
\mathbf{v}^{(l)}&=\bm{\Theta}(W^{(l-1)}\mathbf{v}^{(l-1)}+\sigma_{l}\mathbf{z}^{(l)}),\quad
l=2,\ldots,L\\\ \mathbf{y}&=W^{(L)}\mathbf{v}^{(L)}.\end{aligned}$ (S44)

Denote $C^{(l)}$ as the covariance of $\mathbf{v}^{(l)}$ and
$(\tilde{W}^{(l)})_{l=0}^{L-1}$ as the weight matrices of the network.
According to the chain rule, we have
$\frac{\partial\mathcal{L}}{\partial\tilde{W}_{ij}}=\frac{\partial\mathcal{L}}{\partial
v^{(l+1)}_{i}}\frac{\partial v^{(l+1)}_{i}}{\partial\tilde{W}_{ij}}$ for each
layer $l$. Thus, we can replace $\mathbf{x}$ in Eq. (S20) with
$\mathbf{v}^{(l)}$, replace $C_{\mathrm{in}}$ in Eq. (S44) with $C^{(l)}$, and
replace $\tilde{C}_{\mathrm{in}}$ in Eq. (S20) with $\tilde{C}^{(l)}$. After
applying these modifications, the rest of the analysis for each layer $l$
remains the same as the single-layer case. Therefore, after trained an MNN
$\tilde{M}$ with the same architecture, $\tilde{M}$ learns the weight
matrices, means and variances for each layer of $M_{GT}$.

###  2.4 MNNs approximate the covariance under mild conditions

We analyze the covariance of the MNN $\tilde{M}$ after trained by SMUC.
According to Eq. (22), similar to the mean activation of Heaviside, we also
have the scalar property

$\displaystyle\chi(\mu,C)=\chi(a\mu,a^{2}C),\quad\forall
a>0,\mu\in\mathbb{R},C>0$ (S45)

and hence

$\displaystyle\chi(\tilde{\mathbf{w}}_{k}^{\mathrm{T}}\mathbf{x},\tilde{C}_{k})=\chi(\bar{\mu}_{k},\bar{C}_{k})$
(S46)

The covariance of $\tilde{M}$ is

$\displaystyle\tilde{C}_{y,kj}(\mathbf{x})=\chi(\bar{\mu}_{k},\bar{C}_{k})\chi(\bar{\mu}_{j},\bar{C}_{j})\frac{\mathbf{w}_{k}^{T}\tilde{C}_{\mathrm{in}}\mathbf{w}_{j}}{\sqrt{\bar{C}_{k}\bar{C}_{j}}},\quad
k,j=1,\ldots,n,\quad k\neq j.$ (S47)

It is observed that unless
$\mathbf{w}_{k}^{T}\tilde{C}_{\mathrm{in}}\mathbf{w}_{j}=\mathbf{w}_{k}^{T}C_{\mathrm{in}}\mathbf{w}_{j}$,
we cannot ensure that the output covariance of $\tilde{M}$ is as the same as
the ground-truth covariance $C_{y,kj}$ of $M_{GT}$. When the dimension of the
output is one ($n=1$), it is not an issue, since we only need to fit the
variance of the output. When the dimension of the output is multi-dimensional
($n>1$), however, it seems that we need to know the ground-truth input
covariance $C_{\mathrm{in}}$ for the learning the ground-truth output
covariance. Nevertheless, we show that under some mild conditions, when the
number of the layers goes to infinity, the difference between the output
covariance of $\tilde{M}$ and $M_{GT}$ goes to zero, i.e., we prove the Thm.
3.

###### Proof.

According to Eq. (S46)-(S47), we have

$\displaystyle\left|C^{(l+1)}_{ij}-\tilde{C}^{(l+1)}_{ij}\right|$
$\displaystyle=\left|\frac{\chi_{i}^{(l)}\chi_{j}^{(l)}}{\sqrt{\big{(}W_{i\bullet}^{(l)}C^{(l)}(W_{i\bullet}^{(l)})^{\mathrm{T}}+\sigma_{l}^{2}\big{)}\big{(}W_{j\bullet}^{(l)}C^{(l)}(W_{j\bullet}^{(l)})^{\mathrm{T}}+\sigma_{l}^{2}\big{)}}}\sum_{k,q}W^{(l)}_{ik}W^{(l)}_{jq}(C^{(l)}_{kq}-\tilde{C}^{(l)}_{kq})\right|$
(S48)
$\displaystyle\leq\left|\chi_{i}^{(l)}\chi_{j}^{(l)}\right|\frac{n^{(l)}\left\|W_{i\bullet}^{(l)}\right\|_{2}\left\|W_{j\bullet}^{(l)}\right\|_{2}}{\sqrt{\big{(}\lambda_{\min}(C^{(l)})\left\|W_{i\bullet}^{(l)}\right\|_{2}^{2}+\sigma_{l}^{2}\big{)}\big{(}\lambda_{\min}(C^{(l)})\left\|W_{j\bullet}^{(l)}\right\|_{2}^{2}+\sigma_{l}^{2}\big{)}}}\max_{k,q}\left|C^{(l)}_{kq}-\tilde{C}^{(l)}_{kq}\right|$
(S49)
$\displaystyle\leq\left|\chi_{i}^{(l)}\chi_{j}^{(l)}\right|\frac{n^{(l)}}{\sqrt{\lambda_{\min}(C^{(l)})^{2}+2\lambda_{\min}(C^{(l)})\sigma_{l}^{2}}}\max_{k,q}\left|C^{(l)}_{kq}-\tilde{C}^{(l)}_{kq}\right|,\quad
i\neq j$ (S50)

where $\left\|\cdot\right\|_{2}$ is the $l_{2}$ norm of vectors. In the above,
we have applied the Cauchy-Schwartz inequality

$\displaystyle\sum_{k,q}W^{(l)}_{ik}W^{(l)}_{jq}\leq
n^{(l)}\left\|W_{i\bullet}^{(l)}\right\|_{2}\left\|W_{j\bullet}^{(l)}\right\|_{2},$
(S51)

and the arithmetic mean inequality

$\displaystyle\frac{\left\|W_{i\bullet}^{(l)}\right\|_{2}^{2}}{\left\|W_{j\bullet}^{(l)}\right\|_{2}^{2}}+\frac{\left\|W_{j\bullet}^{(l)}\right\|_{2}^{2}}{\left\|W_{i\bullet}^{(l)}\right\|_{2}^{2}}\geq
2\sqrt{\frac{\left\|W_{i\bullet}^{(l)}\right\|_{2}^{2}}{\left\|W_{j\bullet}^{(l)}\right\|_{2}^{2}}\frac{\left\|W_{j\bullet}^{(l)}\right\|_{2}^{2}}{\left\|W_{i\bullet}^{(l)}\right\|_{2}^{2}}}=2$
(S52)

Denote the infinity norm of a vector or matrix as
$\left\|\cdot\right\|_{\infty}$, then we have

$\displaystyle\left\|C^{(l+1)}-\tilde{C}^{(l+1)}\right\|_{\infty}\leq
r^{(l)}\left\|C^{(l)}-\tilde{C}^{(l)}\right\|_{\infty},$ (S53)

where $\left\|C^{(l)}-\tilde{C}^{(l)}\right\|_{\infty}$ is the the deviation
of the trained network’s covariance $\tilde{C}^{(l+1)}$ from the ground-truth
one $C^{(l)}$ at the $l$-th layer, and $r^{(l)}$ is the covariance deviation
rate at the $l$-layer defined in Eq. (16) in the main paper. It suggests that
$r^{(l)}$ describes how this deviation reduces through the $l$-layer. If
$-\sum_{l}\log r^{(l)}$ diverges to positive infinity, as assumed by the
condition, we have:

$\displaystyle\lim_{l\rightarrow+\infty}\left\|C^{(l)}_{ij}-\tilde{C}^{(l)}_{ij}\right\|_{\infty}\leq\lim_{l\rightarrow+\infty}\exp(\sum_{q=1}^{l}\log
r^{(q)})\left\|C^{(1)}_{ij}-\tilde{C}^{(1)}_{ij}\right\|_{\infty}=0.$ (S54)

∎

According to Eq.(22) and Eq.(S53), to reduce the deviation in the output
covariance, it is beneficial to have large
$\frac{\mu^{(l)}}{\sqrt{\bar{C}^{(l)}}}$ and $\lambda_{\min}(C^{(l)})$ and
have small $n^{(l)}$. Note that $\frac{\mu^{(l)}}{\sqrt{\bar{C}^{(l)}}}$ is
the signal-to-noise ratio (SNR) at the $l$-th layer, which reflects the
quality of the inputs from the previous layer, it suggests that a high SNR
helps to decrease the deviation of covariance. Additionally, the amount
$\lambda_{\min}(C^{(l)})$ is related to the fluctuation of the output
covariance at the layer $l$. Accordingly, high output fluctuation may also
benefits to the decreasing deviation of covariance when keeping
$\frac{\mu^{(l)}}{\sqrt{\bar{C}^{(l)}}}$ as the same. Due to the exponential
term in the function $\chi$, the covariance rate $r^{(l)}$ exhibits an
exponential decrease as the signal-to-noise ratio (SNR) increases. On the
other hand, the covariance rate increases linearly with the dimension
$n^{(l)}$. Consequently, even if the dimension $n^{(l)}$ is high, it is a
moderately high SNR is sufficient to control the deviation rate.

###  2.5 General activation functions

The previous results require the scaling property of the MAs of (Eq. (S27)),
which may not be satisfied when the stochastic network applies general
nonlinearities such as ReLU function. Nevertheless, we can generalize the
above results for general activation function via the simple function
approximation theorem [53], which argues that we can approximate any
measurable nonlinearities $g$ by a series of shifted Heaviside functions i.e.,
$g_{n}(x)\rightarrow g(x)$ point-wisely as $n$ goes to infinity, where
$g_{n}(x)=\sum_{r=1}^{R_{n}}a_{r_{n}}h(x-b_{r_{n}})$ satisfying

$\displaystyle\left|g_{n}(x)\right|\leq\left|g(x)\right|,\quad\forall
x\in\mathbb{R},n\in\mathbb{N}$ (S55)

and $a_{r_{n}}$ and $b_{r_{n}}$ are constants. This theorem implies that
linear combination of the Heaviside functions with different biases $b_{r}$
can approximate any measurable activation function $g$. To demonstrate the
validity of this approximation in training MNNs, specifically regarding the
model’s output and the calculation of the derivative of the mean activation
with respect to the input mean during training, we need to verify

$\displaystyle\lim_{n\rightarrow+\infty}\int
g_{n}(x)\mathcal{N}(x;\mu,\sigma^{2})dx=\int
g(x)\mathcal{N}(x;\mu,\sigma^{2})dx,\quad\forall\mu,\sigma,\quad\text{(zero-
order approximation),}$ (S56)
$\displaystyle\lim_{n\rightarrow+\infty}\frac{d}{d\mu}\int
g_{n}(x)\mathcal{N}(x;\mu,\sigma^{2})dx=\frac{d}{d\mu}\int
g(x)\mathcal{N}(x;\mu,\sigma^{2})dx,\quad\forall\mu,\sigma,\quad\text{(first-
order approximation)},$ (S57)

where we use $\mathcal{N}(x;\mu,\sigma^{2})$ to represent the probabilistic
density function of a Gaussian distribution with mean $\mu$ and variance
$\sigma^{2}$. We prove the above claims for the cases where the activation
function $g$ satisfies that

$\displaystyle\sup_{x\in\mathbb{R}}\frac{|g(x)|}{|x|^{m}+1}<+\infty$ (S58)

for some $m>0$. This is a mild condition satisfied by most of common used
activation functions. The proof is as follows. For given $\mu$, $\sigma$, and
$\epsilon>0$, we show that there exists $N(\epsilon)$, such that $\forall
n>N(\epsilon)$, we have

$\displaystyle\left|\lim_{n\rightarrow+\infty}\int
g_{n}(x)\mathcal{N}(x;\mu,\sigma^{2})dx-\int
g(x)\mathcal{N}(x;\mu,\sigma^{2})dx\right|<\epsilon.$ (S59)

Using Eq. (S55), and the condition Eq. (S58), there exists a constant $A$
(independent of $n$) such that

$\displaystyle\frac{|g_{n}(x)-g(x)|}{|x|^{m}+1}\leq\frac{|g(x)|+|g_{n}(x)|}{|x|^{m}+1}\leq\frac{2|g(x)|}{|x|^{m}+1}\leq
A.$ (S60)

Since $(|x|^{m}+1)\mathcal{N}(x;\mu,\sigma^{2})$ is integrable, there exists
$B>0$ such that

$\displaystyle
A\int_{x\notin[-B,B]}(|x|^{m}+1)\mathcal{N}(x;\mu,\sigma^{2})dx<\frac{\epsilon}{2}.$
(S61)

Since $[-B,B]$ is bounded, we can require $g_{n}$ to uniformly converges to
$g$ on $[-B,B]$ [53]. Hence, there exists a $N(\epsilon)$ such that

$\displaystyle\sup_{x\in[-B,B]}|g_{n}(x)-g(x)|<\frac{\epsilon}{2}.$ (S62)

Therefore, $\forall n>N(\epsilon)$, we have

$\displaystyle\left|\int g_{n}(x)\mathcal{N}(x;\mu,\sigma^{2})dx-\int
g(x)\mathcal{N}(x;\mu,\sigma^{2})dx\right|$ (S63)
$\displaystyle\leq\int|g_{n}(x)-g(x)|\mathcal{N}(x;\mu,\sigma^{2})dx$ (S64)
$\displaystyle<\frac{\epsilon}{2}\int_{x\in[-B,B]}\mathcal{N}(x;\mu,\sigma^{2})dx+A\int_{x\notin[-B,B]}(|x|^{m}+1)\mathcal{N}(x;\mu,\sigma^{2})dx<\frac{\epsilon}{2}+\frac{\epsilon}{2}=\epsilon$
(S65)

The proof of the first-order approximation (Eq. (S57)) is as the same, except
that we need to replace $\mathcal{N}(x;\mu,\sigma^{2})$ by
$\frac{d}{d\mu}\mathcal{N}(x;\mu,\sigma^{2})$ in above procedure. Hence, we
can prove the approximation up to any finite order. As a result, we can
approximate an MNN that employs any measurable functions as activation
functions with an MNN that utilizes the Heaviside function as the activation
function. This approximation holds true not only for the model’s output but
also for the training process of the MNN. In doing so, the previous
theoretical results of SMUC can be directly generalized to the cases where MNN
uses MAs derived from other nonlinearity such as ReLU MAs by approximating
this MNN using corresponding Heaviside MNNs.

In contrast to MNNs, it is important to note that it is invalid to approximate
a deterministic conventional ANNs using the Heaviside function as the
activation function in terms of training. This is because the derivative of a
linear combination of a finite number of Heaviside functions is almost zero