††thanks: Correspondence should be addressed to P.S. and Y.T.

]Tsien Excellence in Engineering Program, Xingjian College, Tsinghua
University, Beijing, 100084, China. ]Tsien Excellence in Engineering Program,
Xingjian College, Tsinghua University, Beijing, 100084, China. ]Laboratory of
Advanced Computing and Storage, Central Research Institute, 2012 Laboratories,
Huawei Technologies Co. Ltd., Beijing, 100084, China. ]Department of
Psychology & Tsinghua Brain and Intelligence Lab, Tsinghua University,
Beijing, 100084, China. ]Department of Psychology & Tsinghua Laboratory of
Brain and Intelligence, Tsinghua University, Beijing, 100084, China. Also
at]Laboratory of Advanced Computing and Storage, Central Research Institute,
2012 Laboratories, Huawei Technologies Co. Ltd., Beijing, 100084, China.

# Statistical Physics of Deep Neural Networks: Initialization toward Optimal
Channels

Kangyu Weng<EMAIL_ADDRESS>[  Aohua Cheng
<EMAIL_ADDRESS>[  Ziyang Zhang<EMAIL_ADDRESS>[  Pei Sun
<EMAIL_ADDRESS>[  Yang Tian<EMAIL_ADDRESS>[  [

###### Abstract

In deep learning, neural networks serve as noisy channels between input data
and its representation. This perspective naturally relates deep learning with
the pursuit of constructing channels with optimal performance in information
transmission and representation. While considerable efforts are concentrated
on realizing optimal channel properties during network optimization, we study
a frequently overlooked possibility that neural networks can be initialized
toward optimal channels. Our theory, consistent with experimental validation,
identifies primary mechanics underlying this unknown possibility and suggests
intrinsic connections between statistical physics and deep learning. Unlike
the conventional theories that characterize neural networks applying the
classic mean-filed approximation, we offer analytic proof that this
extensively applied simplification scheme is not valid in studying neural
networks as information channels. To fill this gap, we develop a corrected
mean-field framework applicable for characterizing the limiting behaviors of
information propagation in neural networks without strong assumptions on
inputs. Based on it, we propose an analytic theory to prove that mutual
information maximization is realized between inputs and propagated signals
when neural networks are initialized at dynamic isometry, a case where
information transmits via norm-preserving mappings. These theoretical
predictions are validated by experiments on real neural networks, suggesting
the robustness of our theory against finite-size effects. Finally, we analyze
our findings with information bottleneck theory to confirm the precise
relations among dynamic isometry, mutual information maximization, and optimal
channel properties in deep learning. Our work may lay a cornerstone for
promoting deep learning in terms of network initialization and suggest general
statistical physics mechanisms underlying diverse deep learning techniques.

††preprint: APS/123-QED

##  I Introduction

###  I.1 Neural networks are information channels

In deep learning, neural networks attempt to identify an optimal latent
representation (e.g., a low-dimensional feature space) of the data such that
subsequent learning tasks can be solved more efficiently [1]. Below, we not
only review latest advances in deep learning but also suggest a general
perspective to unfiy them.

Let us consider a sample set $\mathbf{X}$ and an associated learning target
set $\mathbf{Y}=\gamma\left(\mathbf{X}\right)$ (e.g., labels), where $\gamma$
is a mapping defined by the learning task. A neural network parameterized by
$\phi$ is expected to optimize a representation $\phi\left(\mathbf{X}\right)$
with
$\operatorname{dim}\left(\phi\left(\mathbf{X}\right)\right)<\operatorname{dim}\left(\mathbf{X}\right)$
(here $\operatorname{dim}\left(\cdot\right)$ measures the dimensionality) such
that an ideal mapping
$\gamma_{\phi}:\phi\left(\mathbf{X}\right)\rightarrow\mathbf{Y}$ can be
readily learned to solve the task. This objective requires an appropriate
evaluation of the optimality of neural network representation
$\phi\left(\mathbf{X}\right)$.

Figure 1: Conceptual illustrations of research objectives. (a) The learning
target $\mathbf{Y}$ (e.g., a distribution of different classes of objects) is
represented by a sample set $\mathbf{X}$, where different classes of samples
(denoted by nodes with different colors) are promiscuously distributed in the
original sample space. A neural network is expected to learn an appropriate
neural representation $\phi\left(\cdot\right)$ such that all classes are
distributed following clear patterns in $\phi\left(\mathbf{X}\right)$ to
optimally capture the information of $\mathbf{Y}$. During this process, a
joint channel consisting of two sub-channels is defined between $\mathbf{Y}$
and $\phi\left(\mathbf{X}\right)$. The fist sub-channel, bridging between
$\mathbf{Y}$ and $\mathbf{X}$, determines how two Shannon entropies,
$\mathcal{H}\left(\mathbf{Y}\right)$ and $\mathcal{H}\left(\mathbf{X}\right)$,
share a common part of information measured by
$\mathcal{I}\left(\mathbf{X};\mathbf{Y}\right)$ given the information loss
measured by conditional entropies,
$\mathcal{H}\left(\mathbf{Y}\mid\mathbf{X}\right)$ and
$\mathcal{H}\left(\mathbf{X}\mid\mathbf{Y}\right)$. The second sub-channel,
bridging between $\mathbf{Y}$ and $\phi\left(\mathbf{X}\right)$, is
optimizable for neural networks. These sub-channels jointly define the upper-
bound of $\mathcal{I}\left(\phi\left(\mathbf{X}\right);\mathbf{Y}\right)$
shown in Eq. (2). (b) Because
$\mathcal{I}\left(\phi\left(\mathbf{X}\right);\mathbf{Y}\right)$ is
unmodifiable for neural networks, deep learning studies primarily explore
maximizing $\mathcal{I}\left(\phi\left(\mathbf{X}\right);\mathbf{X}\right)$.
This objective is partly related with, not directly equivalent to, maximizing
$\mathcal{I}\left(\phi\left(\mathbf{X}\right);\mathbf{Y}\right)$. While
previous works mainly focus on maximizing
$\mathcal{I}\left(\phi\left(\mathbf{X}\right);\mathbf{X}\right)$ during
optimization (e.g., training), we shall suggest a possibility to realize this
objective during neural network initialization by proposing a statistical
physics theory of neural networks.

Although the optimality of $\phi\left(\mathbf{X}\right)$ can be evaluated by
diverse metrics according to concrete task demands (e.g., see instances in
reinforcement [2], graph [3], and causal [4] representation learning
frameworks), a mainstream idea is to consider the neural network as a noisy
channel between $\mathbf{X}$ and its representation
$\phi\left(\mathbf{X}\right)$. This perspective naturally leads to the
consideration of two cases:

  * (1)

In unsupervised learning, the information of learning target $\mathbf{Y}$ is
not known to the neural network [5]. To make the learning task resolvable,
mapping $\gamma$ is assumed as a bijective function from sample $\mathbf{X}$
to target $\mathbf{Y}$ such that the neural network can learn the distribution
of $\mathbf{Y}$ by representing the distribution of $\mathbf{X}$ (e.g., the
target distribution is exactly an optimal representation of sample
distribution in clustering tasks [6, 7] and unsupervised representation
learning [8, 9]). If samples and targets have distinct or irrelevant
distributions, the unsupervised learning of $\mathbf{Y}$ based on $\mathbf{X}$
is ill-posed.

  * (2)

In supervised learning, the information of learning target $\mathbf{Y}$ is
known to the neural network as supervision [10]. In an ideal situation where
$\gamma$ is a well-defined bijective mapping from $\mathbf{X}$ to $\mathbf{Y}$
(i.e., samples and targets are perfectly paired), learning the distribution of
$\mathbf{Y}$ is principally equivalent to learning the distribution of
$\mathbf{X}$ (e.g., consider the separable case where $\mathbf{X}$ can be
subdivided into disjoint convex sets in Euclidean space according to the label
information in $\mathbf{Y}$ [11, 12]). In more realistic situations where
samples and targets are not perfectly paired, learning the distribution of
$\mathbf{Y}$ is non-trivial and not necessarily consistent with representing
the distribution of $\mathbf{X}$ (e.g., consider the case where noisy labels
exit [13, 14]).

###  I.2 How can neural networks become optimal channels

How can neural networks become optimal channels favorable for deep learning
tasks? This is the central question concerned in our research.

Mathematically, learning tasks are unified by a Markov chain of data
processing in information theory [15] (see Fig. 1 for illustration)

$\displaystyle\mathbf{Y}\xrightarrow{\gamma^{-1}}\mathbf{X}\xrightarrow{\phi}\phi\left(\mathbf{X}\right).$
(1)

Note that the above Markov chain is different from
$\mathbf{X}\xrightarrow{\phi}\phi\left(\mathbf{X}\right)\xrightarrow{\gamma_{\phi}}\mathbf{Y}$,
the Markov chain of hidden variable models [16]. This is because the joint
distribution between samples and targets is given as _a priori_ rather than
something adjustable in deep learning (e.g., neural networks can not modify
task designs). The Markov property in Eq. (1) implies zero conditional mutual
information values
$\mathcal{I}\left(\mathbf{Y};\phi\left(\mathbf{X}\right)|\mathbf{X}\right)=\mathcal{I}\left(\phi\left(\mathbf{X}\right);\mathbf{Y}|\mathbf{X}\right)=0$
[15], leading to a generalized version of the data processing inequality [15,
17, 18]

$\displaystyle\mathcal{I}\left(\phi\left(\mathbf{X}\right);\mathbf{Y}\right)\leq\min\left\\{\mathcal{I}\left(\mathbf{X};\mathbf{Y}\right),\mathcal{I}\left(\phi\left(\mathbf{X}\right);\mathbf{X}\right)\right\\},$
(2)

which can be readily derived from
$\mathcal{I}\left(\phi\left(\mathbf{X}\right);\mathbf{Y}\right)+\mathcal{I}\left(\mathbf{Y};\mathbf{X}|\phi\left(\mathbf{X}\right)\right)=\mathcal{I}\left(\mathbf{X};\mathbf{Y}\right)$
and
$\mathcal{I}\left(\phi\left(\mathbf{X}\right);\mathbf{Y}\right)+\mathcal{I}\left(\phi\left(\mathbf{X}\right);\mathbf{Y}|\mathbf{X}\right)=\mathcal{I}\left(\mathbf{X};\phi\left(\mathbf{X}\right)\right)$
because conditional mutual information is non-negative.

Eq. (2) suggests a clear direction to evaluate the optimality of neural
network representation $\phi\left(\mathbf{X}\right)$. Because the mutual
information between samples and targets, denoted by
$\mathcal{I}\left(\mathbf{X};\mathbf{Y}\right)$, is unoptimizable for the
neural network, the optimality of $\phi\left(\mathbf{X}\right)$ is, at least
partly, determined by maximizing
$\mathcal{I}\left(\phi\left(\mathbf{X}\right);\mathbf{X}\right)$, the mutual
information between samples and their represented counterparts.

###  I.3 Previous studies on neural networks as optimal channels

Given the possible direction for neural networks to become optimal channels,
we review previous efforts that devote to realize this condition during
optimization (e.g., training).

In unsupervised learning, mutual information maximization has been
demonstrated as a promising approach in recent works [19, 20, 21, 22].
Although this approach arises from the idea of maximizing
$\mathcal{I}\left(\phi\left(\mathbf{X}\right);\mathbf{X}\right)$ in Eq. (2),
it differs from the direct maximization because estimating the mutual
information between the entire sample set and its neural network
representation is statistically deficient [23] (i.e., distribution-free
estimators (e.g., statistical bounds) of entropy-related quantities have high
sample complexity [24]). In general, what mutual information maximization
approach follows is a kind of multi-view formulation [25]. Given each
canonical sample $X$ of $\mathbf{X}$ (e.g., an image) and its two different
and potentially overlapping observations (e.g., two views of the image),
$X^{i}$ and $X^{j}$, neural networks are optimized for
$\max_{\phi}\mathcal{I}\left(\phi\left(X^{i}\right);\phi\left(X^{j}\right)\right)$.
Such an idea can date back to Refs. [26, 27] and is valid for optimizing
neural networks because
$\mathcal{I}\left(\phi\left(X^{i}\right);\phi\left(X^{j}\right)\right)\leq\mathcal{I}\left(X;\phi\left(X^{i}\right),\phi\left(X^{j}\right)\right)$
[25], which can be generally used as a lower bound of the mutual information
maximization objective
$\mathcal{I}\left(\phi\left(\mathbf{X}\right);\mathbf{X}\right)$ [27, 25].

In supervised learning, although mutual information is not explicitly used as
an optimization objective, it has been discovered as intrinsically related to
deep metric learning [28, 25]. Let us consider $\left(X,Y,Z\right)$, a triplet
where $X$ is an anchor, $Y$ is a positive sample (e.g., belongs to the same
class as $X$), and $Z$ is a negative sample (e.g., belongs to a different
class compared with $X$). In deep metric learning, neural networks are
optimized to learn a representation $\phi$ such that
$d\left(\phi\left(X\right),\phi\left(Y\right)\right)<d\left(\phi\left(X\right),\phi\left(Z\right)\right)$
for any $\left(X,Y,Z\right)$, where $d\left(\cdot,\cdot\right)$ denotes a
distance measure [28, 28, 29]. Meanwhile, InfoNCE, an extensively applied
lower bound of mutual information [19], can be derived to support maximizing
$\mathcal{I}\left(\phi\left(\mathbf{X}\right);\mathbf{X}\right)$ in Eq. (2) if
all negative samples are drawn from the true marginal distribution [23]. As
proven by Ref. [25], InfoNCE can be equivalently reformulated as an
expectation of the multi-class $n$-pair loss [30], which is a standard triplet
loss in deep metric learning. In the case where negative samples are not
independently drawn, InfoNCE is not valid in estimating mutual information
[25].

In sum, existing deep learning approaches, irrespective of being unsupervised
or supervised, primarily focus on driving neural networks toward optimal
channels during optimization. Maximizing
$\mathcal{I}\left(\phi\left(\mathbf{X}\right);\mathbf{X}\right)$ is used as,
or at least coincides with, a part of the optimization objectives of
mainstream frameworks.

Figure 2: Summary of our main framework and contributions. (a) In Sec. II, we
review and unify classic mean-field theories of neural networks, measure
signal moments and correlation, construct the phase space of information
propagation, and define edge of chaos (EOC) as well as dynamic isometry (DI).
Based on the unified framework, we indicate the limitation of class mean-field
approximation in characterizing neural networks as information channels in
Sec. III. This limitation arises from the independent and identical assumption
held by class mean-field approximation, which enables the correlation between
propagated signal and its original form to be larger than $1$ or smaller than
$-1$. (b) In Secs. IV-V, we overcome the limitation of existing framework by
proposing a corrected mean-field approximation theory of neural networks. We
relate our analysis with Gaussian information bottleneck, based on which, we
can derive an important lower bound of mutual information that is analytically
proven as maximized at dynamic isometry. Apart from analytic proofs, our
theory is also computationally validated on real neural networks. In Sec. VI,
we suggest the relation between our objective of maximizing
$\mathcal{I}\left(\phi\left(\mathbf{X}\right);\mathbf{X}\right)$ and driving
neural networks toward optimal channels from the perspective of information
bottleneck theory, revealing the insights of our theory on deep learning.

###  I.4 Open questions on neural networks as optimal channels

Given our review and unified formalization presented above, one may
erroneously expect that the properties of neural networks as optimal channels
have been completely confirmed. This is generally not true because the studies
on neural networks as optimal channels, just like any other deep learning
research, still remain at their very beginnings. Below, we summarize three
critical open questions in this direction:

  * (I)

Is it possible to maximize
$\mathcal{I}\left(\phi\left(\mathbf{X}\right);\mathbf{X}\right)$ during neural
network such that the subsequent training process can be improved or at least
be accelerated?

  * (II)

If the properties of neural networks as optimal channels are determined by the
dynamics of information propagation within neural networks? Is it possible to
bridge between the static framework of Shannon theory [15] and the dynamic
characterization of information propagation via existing statistical physics
theories of neural networks (e.g., mean-field approximation [31, 32] and
neural tangent kernel [33, 32, 34] theories)? If it is not possible, what are
the main limitations of existing theories and how to resolve them?

  * (III)

What is the precise relation between maximizing
$\mathcal{I}\left(\phi\left(\mathbf{X}\right);\mathbf{X}\right)$ during
initialization and driving neural networks toward optimal channels? Can we
further relate this relation with other deep learning theories (e.g.,
information bottleneck [35, 36, 37]) to explore a unified view?

###  I.5 Our framework and contribution

Motivated by these open questions, we attempt to develop general theories to
verify the possibility of initializing neural networks at optimal states for
maximizing $\mathcal{I}\left(\phi\left(\mathbf{X}\right);\mathbf{X}\right)$,
based on which, we aim at exploring the underlying connections between
information theory and statistical physics in deep learning. Specifically, our
framework and contributions are summarized as the following.

In Sec. II, we review and unify existing works about characterizing neural
networks as information channels and analyzing information propagation within
them (see Fig. 2a). To make analytic derivations possible, our work primarily
focuses on infinite-width neural networks, whose formal definitions are
presented in Sec. II.1. Infinite-width neural networks, irrespective of being
optimized in a Bayesian manner [38, 39, 40] or by gradient descent approaches
[33, 41, 42], are favorable in analytic formulations because they become
Gaussian processes defined with specific kernels (e.g., neural tangent kernel
[33, 32, 34]) at the limits. This property holds across different network
architectures (e.g., fully-connected layers [38], convolutional layers [40,
43], residual connections [40], and recurrent networks [44]), enabling our
theory to analyze mainstream deep learning models. Given an infinite-width
neural network, we unify previous studies to formalize how information
propagates within the network and analytically measure diverse properties of
propagated signals (e.g., the second moment) on each layer in Secs. II.1-II.2.
This formalization relates our analysis with the studies on edge of chaos and
dynamic isometry in deep neural networks [45, 46, 47, 48, 49, 50, 51], whose
formal definitions are presented in Sec. II.3. In sum, our unified framework
lays the foundation of our subsequent analysis.

Given the fundamental definitions in Sec. II, we do not limit ourselves to
adapting classic theories completely. On the contrary, we show that the
independent and identical assumption held by the classic mean-field
approximation of infinite-width neural networks is invalid in analyzing neural
networks as information channels in Sec. III (see Fig. 2a). Specifically, the
independent and identical assumption is proven to imply an Gaussian
distribution of the correlation between propagated signal and its original
form in inputs in Sec. III.1. In other words, the classic mean-field
approximation creates a non-zero probability for the correlation to be larger
than $1$ or smaller than $-1$, which contradicts the definition of correlation
(i.e., a correlation must belongs to the interval of $\left[-1,1\right]$).
Although a correlation quantity can be empirically treated as a constant and
may still be valid (e.g., when the constant is located within
$\left[-1,1\right]$) if it follows a Gaussian distribution with an
infinitesimal variance approaching to 0 or a strictly zero variance, we prove
that the correlation measured under the classic mean-field framework
dissatisfies these two conditions in Sec. III.2. Therefore, the classic mean-
field approximation is invalid in correlation measurement even from an
empirical perspective.

To resolve the limitation of classic mean-field approximation suggested in
Sec. III, we develop a new mean-field-like theory with corrected independent
and identical assumption in Sec. IV.1 (see Fig. 2b). Different from the
classic one, our corrected mean-field approximation does not implies a
Gaussian distribution of correlation with non-zero constant variance because
it excludes the independent and identical assumption on input $\mathbf{X}$.
Certainly, the loose constraints held by the corrected mean-field
approximation propose challenges to analytic characterization of information
channel properties. To overcome these challenges, we introduce Gaussian
information bottleneck into our analysis to relate optimizing
$\mathcal{I}\left(\phi\left(\mathbf{X}\right);\mathbf{X}\right)$ with
maximizing a certain lower bound of mutual information in Sec. IV.2. In the
phase space of information propagation, the lower bound is analytically proven
as maximized at dynamic isometry point, a case where each layer serves as a
random mapping with norm-preserving property during information transmission,
in Secs. IV.2-Sec. IV.3. In other situations where dynamic isometry is absent,
neural networks become highly noisy channels with high information dissipation
rates. In Sec. V, our theory is computationally validated on real neural
networks. Although our theory is initially developed on infinite-width neural
networks, we demonstrate its general applicability on real deep learning
models by showing consistency between our theoretical predictions and
empirical observations on finite-width neural networks with diverse settings
(e.g., different layer widths, layer quantities, and activation functions).

In Sec. V, we relate our theories with information bottleneck theory [35, 36,
37], a special case of rate distortion theory [52] and sufficient statistics
theory [53], to present a unified discussion on the role of maximizing
$\mathcal{I}\left(\phi\left(\mathbf{X}\right);\mathbf{X}\right)$ in driving
neural networks toward optimal channels in deep learning (see Fig. 2b). We
show that supervised learning and unsupervised learning share similar
optimization objectives in terms of information bottleneck. When neural
networks are trained with random data shuffling tricks [54, 55], we
theoretically suggests the possibility that maximizing
$\mathcal{I}\left(\phi\left(\mathbf{X}\right);\mathbf{X}\right)$ in Eq. (2)
serves as a conditional mechanism for neural networks to become optimal in
both supervised learning and unsupervised learning. Our discussions suggest
the potential insights of our work on deep learning and statistical physics.

##  II Infinite-width neural networks as information channels

In this section, we introduce a framework to characterize infinite-width
neural networks as information channels by summarizing and reformulating
existing studies on mean-field behaviours and dynamic isometry of neural
networks [56, 57, 58, 47, 46, 49, 50, 51].

###  II.1 Mean-field behaviours of infinite-width neural networks

Let us consider an arbitrary deep neural network with multiple layers. The
dynamics of cross-layer information propagation (i.e., a signal propagates
from the $l$-th layer to the $\left(l+1\right)$-th layer) is characterized as

$\displaystyle\mathbf{X}^{\left(l+1\right)}=\mathbf{W}^{\left(l+1\right)}\psi\left(\mathbf{X}^{\left(l\right)}\right)+\varepsilon^{\left(l+1\right)},$
(3)

where
$\mathbf{X}^{\left(l\right)}=\left(\mathbf{X}_{1}^{\left(l\right)},\ldots,\mathbf{X}_{N_{l}}^{\left(l\right)}\right)\in\mathbb{R}^{N_{l}}$
denotes the vector of pre-activation signals in the $l$-th layer, parameter
$N_{l}\in\mathbb{N}^{+}$ denotes the width of the $l$-th layer, mapping
$\psi\left(\cdot\right)$ denotes a non-linear activation function, matrix
$\mathbf{W}^{\left(l+1\right)}\in\mathbb{R}^{N_{l+1}}\times\mathbb{R}^{N_{l}}$
defines the weights of all connections between the $l$-th layer and the
$\left(l+1\right)$-th layer, and
$\varepsilon^{\left(l+1\right)}=\left(\varepsilon^{\left(l+1\right)}_{1},\ldots,\varepsilon_{N_{l+1}}^{\left(l+1\right)}\right)\in\mathbb{R}^{N_{l+1}}$
denotes the associated residuals. In common cases, each residual
$\varepsilon^{\left(l+1\right)}_{i}$ is frequently assumed as a Gaussian
variable [56, 57]. Please see Fig. 3a for illustrations.

To offer a clear vision, we begin with formalizing the classic mean-field
approximation of the above neural network before we analyze its limitations in
Sec. III.1. Under the independent and identical assumption of
$\mathbf{W}^{\left(l+1\right)}$ and $\mathbf{X}^{\left(l\right)}$ (e.g., each
element $\mathbf{W}^{\left(l+1\right)}_{i,j}$ in
$\mathbf{W}^{\left(l+1\right)}$ is independently and identically distributed),
we can derive

$\displaystyle\mathbf{X}^{\left(l+1\right)}_{i}=\Big{\langle}\mathbf{W}^{\left(l+1\right)}_{i},\psi\left(\mathbf{X}^{\left(l\right)}\right)\Big{\rangle}+\varepsilon^{\left(l+1\right)}_{i}\xrightarrow{d}\mathcal{N}\left(\mu_{i},\sigma_{i}^{2}\right)$
(4)

as $N_{l+1}\rightarrow\infty$, where $\mathbf{X}^{\left(l+1\right)}_{i}$ and
$\mathbf{W}^{\left(l+1\right)}_{i}$ respectively denote the $i$-th rows of
$\mathbf{X}^{\left(l+1\right)}$ and $\mathbf{W}^{\left(l+1\right)}$ for any
$i\in\\{1,\ldots,N_{l+1}\\}$, notion $\langle\cdot,\cdot\rangle$ defines the
inner product, and $\mathcal{N}\left(\mu_{i},\sigma_{i}\right)$ defines a
specific Gaussian variable.

In general, Eq. (4) approximates each signal propagating in the infinite-width
neural network as a certain Gaussian variable under the central limit theorem,
enabling us to study an ensemble of random neural networks associated with the
original neural network [58] (see Fig. 3a). This idea has been demonstrated as
effective in characterizing the mean-field behaviours of two-layer [56] and
multi-layer [57] neural networks.

###  II.2 Information propagation in infinite-width neural networks

As we have explained, our motivation to consider infinite-width neural
networks is to analytically study information propagation dynamics within
them. Because the signal propagating across layers in an infinite-width neural
network has become a Gaussian variable, we can capture most of its dynamics by
studying the first two moments of it. In our research, we primarily focus on
the second moment since it has been demonstrated as relevant with the
expressivity of neural networks (i.e., the second moment of a signal coincides
with the length of its internal Riemannian manifold in downstream layers)
[58], enabling us to analyze the order-to-chaos expressivity phase transition
[58].

Figure 3: Conceptual illustrations of the classic mean-field approximation of
neural networks. (a) Key settings of classic mean-field approximation, such as
infinite-width condition as well as independent and identical assumption, are
summarized. (b-c) Illustrations of the scalar and vector metrics used in
characterizing information propagation processes are presented. (d) Main steps
of the derivation processes of edge of chaos (EOC) and dynamic isometry (ID)
in the phase space of information propagation are shown.

For convenience, we consider the case studied by Ref. [58], where each
$\mathbf{W}^{\left(l\right)}_{ij}\;\text{i.i.d.}\sim\mathcal{N}\left(0,\frac{\sigma_{w}^{2}}{N_{l-1}}\right)$
and each
$\varepsilon^{\left(l\right)}_{i}\;\text{i.i.d.}\sim\mathcal{N}\left(0,\sigma_{b}^{2}\right)$
(note that “i.i.d.” stands for being independently and identically
distributed). In the infinite-width limit (i.e., $N_{l}\rightarrow\infty$),
the second moment (i.e., variance) of signals in the $l$-layer, denoted by
$\left(\sigma^{\left(l\right)}\right)^{2}$, can be calculated as

$\displaystyle\left(\sigma^{\left(l\right)}\right)^{2}$
$\displaystyle=\frac{1}{N_{l}}\sum_{i=1}^{N_{l}}\left(\mathbf{X}^{\left(l\right)}_{i}\right)^{2},$
(5)
$\displaystyle=\mathbb{E}\left(\Big{\langle}\mathbf{W}^{\left(l\right)}_{i},\psi\left(\mathbf{X}^{\left(l-1\right)}\right)\Big{\rangle}+\varepsilon^{\left(l\right)}_{i}\right)^{2},$
(6)

where $\mathbb{E}\left(\cdot\right)$ denotes the expectation (see Fig. 3b).
Note that the second moment in Eq. (5) reduces to the second origin moment
because $\mathbb{E}\left(\mathbf{X}^{\left(l\right)}_{i}\right)=0$ holds for
each $\mathbf{X}^{\left(l\right)}_{i}$. Eq. (6) can be further reformulated as

$\displaystyle\left(\sigma^{\left(l\right)}\right)^{2}$
$\displaystyle=\mathbb{E}\left[\sum_{j=1}^{N_{l-1}}\mathbf{W}^{\left(l\right)}_{ij}\psi\left(\mathbf{X}^{\left(l-1\right)}_{j}\right)+\varepsilon^{\left(l\right)}_{i}\right]^{2},$
(7)
$\displaystyle=\sum_{j=1}^{N_{l-1}}\mathbb{E}\left[\left(\mathbf{W}^{\left(l\right)}_{ij}\right)^{2}\right]\mathbb{E}\left[\psi\left(\mathbf{X}^{\left(l-1\right)}_{j}\right)^{2}\right]+\mathbb{E}\left[\left(\varepsilon^{\left(l\right)}_{i}\right)^{2}\right],$
(8)
$\displaystyle=\sigma_{w}^{2}\int_{\mathbb{R}}\psi\left(\sqrt{\sigma^{\left(l-1\right)}}x\right)^{2}\mathsf{D}x+\sigma_{b}^{2},$
(9)

where Eq. (8) is derived from the fact that
$\mathbb{E}\left(\mathbf{W}^{\left(l\right)}_{ij}\mathbf{W}^{\left(l\right)}_{ik}\right)=\mathbb{E}\left(\mathbf{W}^{\left(l\right)}_{ij}\right)\mathbb{E}\left(\mathbf{W}^{\left(l\right)}_{ik}\right)$
for any $j\neq k$ under the independent and identical assumption and Eq. (9)
is derived using
$\mathbb{E}\left[\sum_{j=1}^{N_{l-1}}\left(\mathbf{W}^{\left(l\right)}_{ij}\right)^{2}\right]=\sigma_{w}^{2}$
(i.e., weights are independently and identically distributed). In Eq. (9),
notion
$\mathsf{D}x=\frac{\mathsf{d}x}{\sqrt{2\pi}}\exp\left(-\frac{x^{2}}{2}\right)$
is a standard Gaussian measure, where $x\in\mathbb{R}$. In general, Eq. (9)
defines an iterative dynamic process of the second moment of signals in each
layer [58], whose initial condition is

$\displaystyle\left(\sigma^{\left(1\right)}\right)^{2}=\frac{\sigma_{w}^{2}}{N_{1}}\langle\mathbf{X},\mathbf{X}\rangle+\sigma_{b}^{2},$
(10)

where $\mathbf{X}$ is the input vector to the neural network, which is assumed
to obey the independent and identical assumption, i.e., each
$\mathbf{X}_{i}\;\text{i.i.d.}\sim\text{a certain distribution}$, to ensure
that $\mathbf{X}_{i}^{\left(l\right)}\;\text{i.i.d.}\sim\text{a certain
distribution}$ in Eq. (4). One can see Ref. [58] for more analyses of Eq. (9).

Apart from the second moment of a single signal, we can also consider the
relation between two propagating signals. For two input signals,
$\mathbf{X}^{p}$ and $\mathbf{X}^{q}$, correlated with each other because they
propagate within the same neural network, we can mark their representations in
the $l$-layer as $\mathbf{X}^{\left(l\right),p}$ and
$\mathbf{X}^{\left(l\right),q}$ to measure their second moment as (see Fig. 3b
for illustration)

$\displaystyle\sigma^{\left(l\right),pq}=$
$\displaystyle\frac{1}{N_{l}}\sum_{i=1}^{N_{l}}\mathbf{X}^{\left(l\right),p}_{i}\mathbf{X}^{\left(l\right),q}_{i},$
(11) $\displaystyle=$
$\displaystyle\mathbb{E}\Bigg{[}\left(\sum_{j=1}^{N_{l-1}}\mathbf{W}^{\left(l\right)}_{ij}\psi\left(\mathbf{X}^{\left(l-1\right),p}_{j}\right)+\varepsilon^{\left(l\right)}_{i}\right)\left(\sum_{j=1}^{N_{l-1}}\mathbf{W}^{\left(l\right)}_{ij}\psi\left(\mathbf{X}^{\left(l-1\right),q}_{j}\right)+\varepsilon^{\left(l\right)}_{i}\right)\Bigg{]},$
(12) $\displaystyle=$
$\displaystyle\sum_{j=1}^{N_{l-1}}\mathbb{E}\left[\left(\mathbf{W}^{\left(l\right)}_{ij}\right)^{2}\right]\mathbb{E}\left[\psi\left(\mathbf{X}^{\left(l-1\right),p}_{j}\right)\psi\left(\mathbf{X}^{\left(l-1\right),q}_{j}\right)\right]+\mathbb{E}\left[\left(\varepsilon^{\left(l\right)}_{i}\right)^{2}\right],$
(13) $\displaystyle=$
$\displaystyle\sigma_{w}^{2}\int_{\mathbb{R}}\int_{\mathbb{R}}\psi\left(\sigma^{\left(l-1\right),p}x\right)\psi\left[\sigma^{\left(l-1\right),q}\left(\mathcal{C}_{pq}^{\left(l-1\right)}x+\sqrt{1-\left(\mathcal{C}_{pq}^{\left(l-1\right)}\right)^{2}}y\right)\right]\mathsf{D}x\mathsf{D}y+\sigma_{b}^{2},$
(14)

where $\sigma^{\left(l-1\right),p}$ and $\sigma^{\left(l-1\right),q}$ denotes
the standard deviations of $\mathbf{X}^{p}$ and $\mathbf{X}^{q}$ in the
$\left(l-1\right)$-layer defined following Eq. (9). Variables $x$ and $y$ are
independent standard Gaussian variables. Notion
$\mathcal{C}_{pq}^{\left(l-1\right)}$ denotes the correlation between
$\mathbf{X}^{p}$ and $\mathbf{X}^{q}$ in the $\left(l-1\right)$-layer. Please
note that subscript $i$ in Eqs. (11-13) can be eventually dropped in Eq. (14)
because there is independent and identical assumption on the components of
$\mathbf{X}^{p}$ and $\mathbf{X}^{q}$.

###  II.3 Phase space of information propagation

Let us contextualize the above mathematical definitions with physics
backgrounds. While studying information propagation, we expect to understand
how the global extrinsic curvature of latent Riemannian geometry in inputs
(i.e., the relation between two input signals $\mathbf{X}^{p}$ and
$\mathbf{X}^{q}$), a key factor underlying the expressivity of neural networks
[58], evolves across layers.

We wonder if the difference between $\mathbf{X}^{p}$ and $\mathbf{X}^{q}$ will
be principally maintained, enlarged, or reduced during information
propagation. To answer this question, we first explore the stable fixed point
of $\sigma^{\left(l\right)}$ in Eq. (9) as a function of
$\left(\sigma_{w},\sigma_{b}\right)\in\mathbb{R}\times\left(0,\infty\right)$
because the length of each propagating signal in the downstream layer will
rapidly converges to this stable fixed point. After confirming the stable
fixed point, denoted by $\sigma^{*}\left(\sigma_{w},\sigma_{b}\right)$, we set
$\sigma^{\left(l\right),p}=\sigma^{\left(l\right),q}=\sigma^{*}\left(\sigma_{w},\sigma_{b}\right)$
in Eq. (14) and divide Eq. (14) by
$\sigma^{*}\left(\sigma_{w},\sigma_{b}\right)$ to obtain the iterative
dynamics of the correlation between $\mathbf{X}^{p}$ and $\mathbf{X}^{q}$ (see
Fig. 3b)

$\displaystyle\mathcal{C}_{pq}^{\left(l\right)}=$
$\displaystyle\frac{1}{\sigma^{*}\left(\sigma_{w},\sigma_{b}\right)^{2}}\left\\{\sigma_{w}^{2}\int_{\mathbb{R}}\int_{\mathbb{R}}\psi\left(\sigma^{*}\left(\sigma_{w},\sigma_{b}\right)x\right)\psi\left[\sigma^{*}\left(\sigma_{w},\sigma_{b}\right)\left(\mathcal{C}_{pq}^{\left(l-1\right)}x+\sqrt{1-\left(\mathcal{C}_{pq}^{\left(l-1\right)}\right)^{2}}y\right)\right]\mathsf{D}x\mathsf{D}y+\sigma_{b}^{2}\right\\},$
(15)

whose fixed point can be readily found as $\mathcal{C}_{pq}^{*}=1$ after
direct calculation. The stability of this fixed point, however, can not be
directly confirmed. Therefore, we need to analyze
$\zeta\left(\sigma_{w},\sigma_{b}\right)$, the slope of
$\mathcal{C}_{pq}^{\left(l\right)}$ around $\mathcal{C}_{pq}^{*}=1$ given a
setting $\left(\sigma_{w},\sigma_{b}\right)$ (see Fig. 3d)

$\displaystyle\zeta\left(\sigma_{w},\sigma_{b}\right)=$
$\displaystyle\frac{\partial\mathcal{C}_{pq}^{\left(l\right)}}{\partial\mathcal{C}_{pq}^{\left(l-1\right)}}\Bigg{|}_{\mathcal{C}_{pq}^{\left(l\right)}=1},$
(16) $\displaystyle=$
$\displaystyle\sigma_{w}^{2}\int_{\mathbb{R}}\psi^{\prime}\left(\sigma^{*}\left(\sigma_{w},\sigma_{b}\right)x\right)^{2}\mathsf{D}x,$
(17) $\displaystyle=$
$\displaystyle\frac{1}{N_{l}}\mathbb{E}\left\\{\operatorname{tr}\left[\left(\mathbf{D}^{\left(l\right)}\mathbf{W}^{\left(l\right)}\right)^{T}\mathbf{D}^{\left(l\right)}\mathbf{W}^{\left(l\right)}\right]\right\\}.$
(18)

In Eq. (17), notion $\psi^{\prime}\left(\cdot\right)$ denotes the derivative
function of $\psi\left(\cdot\right)$. In Eq. (18), notion
$\mathbf{D}^{\left(l\right)}$ is a diagonal matrix
$\mathbf{D}^{\left(l\right)}=\operatorname{diag}\left(\left[\psi^{\prime}\left(\mathbf{X}^{\left(l\right)}_{1}\right),\ldots,\psi^{\prime}\left(\mathbf{X}^{\left(l\right)}_{N_{l}}\right)\right]\right)$
such that
$\mathbf{J}^{\left(l\right)}=\mathbf{D}^{\left(l\right)}\mathbf{W}^{\left(l\right)}$
can be understood as the input-output Jacobian matrix of the $l$-th layer [47,
46]. Notion $\operatorname{tr}\left(\cdot\right)$ denotes matrix trace. The
expectation is calculated by averaging across all possible configurations of
$\mathbf{D}^{\left(l\right)}\mathbf{W}^{\left(l\right)}$ in Eq. (18). As
suggested by Eq. (18), parameter $\zeta\left(\sigma_{w},\sigma_{b}\right)$ can
be understood as a stretch factor because any random perturbation $\eta$ in
the $\left(l-1\right)$-th layer, $\mathbf{X}^{\left(l\right)}+\nu$, implies a
subsequent perturbation in the $l$-th layer,
$\mathbf{X}^{\left(l+1\right)}+\mathbf{J}^{\left(l-1\right)}\nu$, with a
stretch effect measured by
$\zeta\left(\sigma_{w},\sigma_{b}\right)=\mathbb{E}\left(\|\mathbf{J}^{\left(l-1\right)}\nu\|_{2}^{2}/\|\nu\|_{2}^{2}\right)$
[58]. The effect corresponds to growth if
$\zeta\left(\sigma_{w},\sigma_{b}\right)>1$ and corresponds to shrinkage if
$\zeta\left(\sigma_{w},\sigma_{b}\right)<1$ (see Fig. 3d).

The fixed point $\mathcal{C}_{pq}^{*}=1$ is stable when
$\zeta\left(\sigma_{w},\sigma_{b}\right)<1$ while it is unstable when
$\zeta\left(\sigma_{w},\sigma_{b}\right)>1$ [58, 47, 46]. In the case where
$\mathcal{C}_{pq}^{*}=1$ is stable, all possible relations between
$\mathbf{X}^{p}$ and $\mathbf{X}^{q}$ eventually converge to a strong
correlation (i.e., $\mathbf{X}^{\left(l\right),p}$ and
$\mathbf{X}^{\left(l\right),q}$ become increasingly similar as $l$ increases).
In the case where $\mathcal{C}_{pq}^{*}=1$ is not stable, $\mathbf{X}^{p}$ and
$\mathbf{X}^{q}$ become increasingly separable as they propagate.
Consequently, the condition with $\zeta\left(\sigma_{w},\sigma_{b}\right)=1$
naturally defines a boundary separating between two phases on the plane of
$\left(\sigma_{w},\sigma_{b}\right)$. The first phase, corresponding to the
case where signals become separable during information propagation (i.e.,
$\zeta\left(\sigma_{w},\sigma_{b}\right)>1$), is referred to as the disordered
phase. The second phase, corresponding to the case where signals converge to
correlated states (i.e., $\zeta\left(\sigma_{w},\sigma_{b}\right)<1$), is the
ordered phase. Please see Fig. 3d for illustrations of the phase space. As
suggested by Refs. [46, 49], the ordered and disordered phases correspond to
vanishing and exploding gradient problems, respectively.

Apart from defining a phase transition boundary between ordered and disordered
phases, we can further consider the dynamic isometry condition, a special
point on this boundary (see Fig. 3d). Specifically, we can understand
$\zeta\left(\sigma_{w},\sigma_{b}\right)$ in Eq. (18) as the second moment of
the singular values of
$\mathbf{J}^{\left(l\right)}=\mathbf{D}^{\left(l\right)}\mathbf{W}^{\left(l\right)}$
or, equivalently, the first moment of the singular values of
$\mathbf{H}^{\left(l\right)}=\left(\mathbf{J}^{\left(l\right)}\right)^{T}\mathbf{J}^{\left(l\right)}$

$\displaystyle\zeta\left(\sigma_{w},\sigma_{b}\right)=\frac{1}{N_{l}}\sum_{i=1}^{N_{l}}\theta_{i}^{2}=\frac{1}{N_{l}}\sum_{i=1}^{N_{l}}\lambda_{i},$
(19)

where $\left[\theta_{1},\ldots,\theta_{N_{l}}\right]$ and
$\left[\lambda_{1},\ldots,\lambda_{N_{l}}\right]$ are the singular values of
$\mathbf{J}^{\left(l\right)}$ and $\mathbf{H}^{\left(l\right)}$, respectively.
Dynamic isometry is defined as a case where the singular values of
$\mathbf{H}^{\left(l\right)}$ not only have a first moment of $1$ but also
satisfy $\lambda_{i}=1$ for each $i\in\\{1,\ldots,N_{l}\\}$. To reach the
dynamic isometry condition, we can consider a situation where the second
moment of the singular values of $\mathbf{H}^{\left(l\right)}$ approaches to
$0$ while the first moment equals $1$. As suggested by Ref. [49], free
probability theory [59] can be applied to derive the probability distribution
of the singular values of $\mathbf{H}^{\left(l\right)}$ to realize this
objective. Detailed calculations of dynamic isometry point
$\left(\sigma_{w}^{\diamond},\sigma_{b}^{\diamond}\right)$ across different
activation functions $\psi\left(\cdot\right)$ or network architectures have
been provided by Refs. [49, 50, 51, 47] and we summarize the general method in
Appendix A. Based on the method, it has been demonstrated that orthogonal
weights, i.e.,
$\left(\mathbf{W}^{\left(l\right)}\right)^{T}\mathbf{W}^{\left(l\right)}=\mathbf{I}$
for each $l$ (notion $\mathbf{I}$ denotes the identity matrix), and a non-
ReLU-type activation function $\psi\left(\cdot\right)$, i.e.,
$\psi\left(r\right)\neq\max\left(0,r\right)$ for each $r\in\mathbb{R}$, can
achieve dynamical isometry in neural networks [47, 49].

In sum, we have reviewed and unified existing studies on mean-field behaviours
and dynamic isometry of neural networks [56, 57, 58, 47, 46, 49, 50, 51] to
present a general framework for analyzing information propagation in infinite-
width neural networks. This framework supports us to rethink the limitation of
existing theories and explore undiscovered laws governing neural networks.

##  III On the limitation of classic mean-field approximation

In this section, we suggest the limitation of classic mean-field approximation
in characterizing neural networks as information channels.

###  III.1 Rethinking the classic mean-field approximation: a strict
perspective

Let us rethink the validity of the classic mean-field approximation summarized
in Sec. II in defining the channel capacity of neural networks. At the first
glance, the rethinking seems to be unnecessary because Sec. II has suggested
how neural networks serve as channels where information propagates across
layers in a mean-field manner. However, as we suggested below, the independent
and identical assumption in classic mean-field theory may imply unexpected
errors in correlation measurements.

Figure 4: Conceptual illustrations of the limitation of classic mean-field
approximation in correlation measurement. (a) In real cases, the $i$-th
component of an input signal, $\mathbf{X}_{i}$, and the $j$-th component of
propagated signal in the $l$-th layer, $\mathbf{X}^{\left(l\right)}_{j}$
should always be located within $\left[-1,1\right]$. However, the independent
and identical assumption held by classic mean-field assumption implies a non-
zero probability for the measured correlation,
$\mathcal{C}_{ij}^{\left(0,l\right)}$, to be larger than $1$ or small than
$-1$ because $\mathcal{C}_{ij}^{\left(0,l\right)}$ follows a specific Gaussian
distribution. (b) Although a Gaussian variable with an infinitesimal variance
approaching to $0$ or a strictly zero variance can be empirically treated as a
constant and may be valid to serve as a correlation (e.g., when the constant
is located within $\left[-1,1\right]$), the measured correlation
$\mathcal{C}_{ij}^{\left(0,l\right)}$ is proven to dissatisfy both cases and
has a constant non-zero variance. Consequently, it is invalid even from an
empirical perspective.

Correlation $\mathcal{C}$, irrespective of being measured between any pair of
variables, should be located within the interval of $\left[-1,1\right]$ (see
Fig. 4a). An approximation framework is invalid if it enables a correlation to
be larger than $1$ or smaller than $-1$ with a non-zero probability (see Fig.
4a). Different from the correlation between two inputs, $\mathbf{X}^{p}$ and
$\mathbf{X}^{q}$, in the $l$-th layer defined by Eq. (15), here we consider
$\mathcal{C}_{ij}^{\left(0,l\right)}$, the correlation between the $i$-th
component of an input signal, $\mathbf{X}_{i}$, and the $j$-th component of
propagated signal in the $l$-th layer, $\mathbf{X}^{\left(l\right)}_{j}$, in
Eq. (20). Please see Fig. 3b for illustration. This correlation reflects how
an input evolves during its propagation from the $1$-st layer to the $l$-th
layer. Specifically, we have

$\displaystyle\mathcal{C}_{ij}^{\left(0,l\right)}$
$\displaystyle=\frac{\mathbb{E}\left[\left(\mathbf{X}^{\left(l\right)}_{j}-\mathbb{E}\left(\mathbf{X}^{\left(l\right)}_{j}\right)\right)\Big{(}\mathbf{X}_{i}-\mathbb{E}\left(\mathbf{X}_{i}\right)\Big{)}\right]}{\sigma^{\left(l\right)}_{j}\sigma^{\left(0\right)}_{i}},$
(20)
$\displaystyle=\frac{1}{\sigma^{\left(l\right)}_{j}\sigma^{\left(0\right)}_{i}}\sum_{k=1}^{N_{l-1}}\left\\{\mathbf{W}^{\left(l\right)}_{jk}\int_{\mathbb{R}}\int_{\mathbb{R}}\psi\left(\sigma^{\left(l-1\right)}_{k}x_{k}\right)\sigma^{\left(0\right)}_{i}\left[\mathcal{C}_{ik}^{\left(0,l-1\right)}x_{k}+\sqrt{1-\left(\mathcal{C}_{ik}^{\left(0,l-1\right)}\right)^{2}}y_{k}\right]\mathsf{D}x_{k}\mathsf{D}y_{k}\right\\},$
(21)
$\displaystyle=\frac{1}{\sigma^{\left(l\right)}_{j}}\sum_{k=1}^{N_{l-1}}\mathbf{W}^{\left(l\right)}_{jk}\mathcal{C}_{ik}^{\left(0,l-1\right)}\left(\int_{\mathbb{R}}\psi\left(\sigma^{\left(l-1\right)}_{k}z_{k}\right)z_{k}\mathsf{D}z_{k}\right),$
(22)

where $\sigma^{\left(0\right)}_{i}$ denotes the second moment of
$\mathbf{X}_{i}$, the $i$-th component of input signal $\mathbf{X}$, and
$\sigma^{\left(l\right)}_{j}$ denotes the second moment of
$\mathbf{X}_{j}^{\left(l\right)}$, the $j$-th component of propagated signal
$\mathbf{X}^{\left(l\right)}$ (see Fig. 3b for illustration). In Eq. (21),
every $\mathsf{D}x_{k}$ and $\mathsf{D}y_{k}$ are standard Gaussian measures.
Because these measures are identically and independently distributed, their
integration over $\mathbb{R}$ can be uniformly represented by the integration
of a standard Gaussian measure, $\mathsf{D}z_{k}$, over $\mathbb{R}$ in Eq.
(22).

To verify the possibility for $\mathcal{C}_{ij}^{\left(0,l\right)}$ to be
larger than $1$ or smaller than $-1$, we can analyze its support set (i.e., if
$\left[-1,1\right]$ is a proper subset of the support of
$\mathcal{C}_{ij}^{\left(0,l\right)}$, then
$\mathcal{C}_{ij}^{\left(0,l\right)}$ has a non-zero probability to reach an
invalid value). As shown below, our analysis is implemented based on two main
steps.

First, the second moment term in Eq. (22) is formally measured as (see Fig. 3b
for illustration)

$\displaystyle\left(\sigma^{\left(l\right)}_{j}\right)^{2}=$
$\displaystyle\mathbb{E}\left[\left(\sum_{k=1}^{N_{l-1}}\mathbf{W}^{\left(l\right)}_{jk}\psi\left(\mathbf{X}_{k}^{\left(l-1\right)}\right)+\varepsilon_{j}^{\left(l\right)}\right)^{2}\right]$
$\displaystyle-\mathbb{E}\left[\sum_{k=1}^{N_{l-1}}\mathbf{W}^{\left(l\right)}_{jk}\psi\left(\mathbf{X}_{k}^{\left(l-1\right)}\right)+\varepsilon_{j}^{\left(l\right)}\right]^{2},$
(23)

which can be reformulated as

$\displaystyle\left(\sigma^{\left(l\right)}_{j}\right)^{2}=$
$\displaystyle\left(\varepsilon_{j}^{\left(l\right)}\right)^{2}+\sum_{k=1}^{N_{l-1}}\left(\mathbf{W}^{\left(l\right)}_{jk}\right)^{2}\mathbb{E}\left[\psi\left(\mathbf{X}_{k}^{\left(l-1\right)}\right)^{2}\right]-\left(\varepsilon_{j}^{\left(l\right)}\right)^{2},$
(24) $\displaystyle=$
$\displaystyle\sum_{k=1}^{N_{l-1}}\left(\mathbf{W}^{\left(l\right)}_{jk}\right)^{2}\mathbb{E}\left[\psi\left(\mathbf{X}_{k}^{\left(l-1\right)}\right)^{2}\right].$
(25)

The reformulation holds because of
$\mathbb{E}\left[\psi\left(\mathbf{X}_{k}^{\left(l-1\right)}\right)\right]=0$
(hint: all common non-ReLU-type activation functions in deep learning are odd
functions while the probability density of $\mathbf{X}_{k}^{\left(l-1\right)}$
is an even function) and
$\mathbb{E}\left[\psi\left(\mathbf{X}_{a}^{\left(l-1\right)}\right)\psi\left(\mathbf{X}_{b}^{\left(l-1\right)}\right)\right]=\delta_{a,b}$
where $\delta_{\cdot,\cdot}$ denotes the Kronecker delta function (hint: the
independent and identical assumption). It is trivial that Eq. (25) is
equivalent to

$\displaystyle\left(\sigma^{\left(l\right)}_{j}\right)^{2}=$
$\displaystyle\alpha\sum_{k=1}^{N_{l-1}}\left(\mathbf{W}^{\left(l\right)}_{jk}\right)^{2}\sim\frac{\sigma_{w}^{2}}{N_{l-1}}\chi^{2}\left(N_{l-1}\right),$
(26)

where we define
$\alpha=\mathbb{E}\left[\psi\left(\mathbf{X}_{k}^{\left(l-1\right)}\right)^{2}\right]$
for each $k$ under the independent and identical assumption (i.e., $\alpha$ is
same for every $k$ in the $\left(l-1\right)$-th layer), notion
$\chi^{2}\left(\cdot\right)$ denotes the Chi-square random variable (see Fig.
4a). Eq. (26) is derived from the independent and identical assumption that
$\mathbf{W}^{\left(l\right)}_{jk}\;\text{i.i.d.}\sim\mathcal{N}\left(0,\frac{\sigma_{w}^{2}}{N_{l-1}}\right)$.
In the infinite-width limit, we discover that the variance of such a Chi-
square variable vanishes

$\displaystyle\lim_{N_{l-1}\rightarrow\infty}2N_{l-1}\frac{\sigma_{w}^{4}}{N_{l-1}}=0.$
(27)

In other words, variable $\sigma^{\left(l\right)}_{j}$ can be reasonably
treated as a constant in an infinite-width neural network (Fig. 4a).

Second, we can define
$\beta=\int_{\mathbb{R}}\psi\left(\sigma^{\left(l-1\right)}_{k}z_{k}\right)z_{k}\mathsf{D}z_{k}$
for each $k$ under the independent and identical assumption (i.e., $\beta$ is
same for every $k$ in the $\left(l-1\right)$-th layer), which supports us to
rewrite Eq. (22) as

$\displaystyle\mathcal{C}_{ij}^{\left(0,l\right)}$
$\displaystyle=\frac{\beta}{\sigma^{\left(l\right)}_{j}}\sum_{k=1}^{N_{l-1}}\mathbf{W}^{\left(l\right)}_{jk}\mathcal{C}_{ik}^{\left(0,l-1\right)}\xrightarrow{d}\frac{\beta}{\sigma^{\left(l\right)}_{j}}\mathcal{N}\left(\mu,\sigma\right),$
(28)

where the right part is derived based on the central limit theorem (i.e.,
$\sum_{k=1}^{N_{l-1}}\mathbf{W}^{\left(l\right)}_{jk}\mathcal{C}_{ik}^{\left(0,l-1\right)}\xrightarrow{d}\mathcal{N}\left(\mu,\sigma\right)$).
As suggested by Eq. (28), $\mathcal{C}_{ij}^{\left(0,l\right)}$ is a Gaussian
random variable. Therefore, interval $\left[-1,1\right]$ is always a proper
subset of the support of $\mathcal{C}_{ij}^{\left(0,l\right)}$, suggesting the
limitation of the independent and identical assumption held by classic mean-
field approximation (see Fig. 4a for illustration).

###  III.2 Rethinking the classic mean-field approximation: an empirical
perspective

Certainly, one can still treat the independent and identical assumption as
reasonably valid if the variance of $\mathcal{C}_{ij}^{\left(0,l\right)}$
equals $0$ or becomes an infinitesimal approaching to $0$. In these cases,
correlation $\mathcal{C}_{ij}^{\left(0,l\right)}$ can be generally analyzed as
a constant from an empirical perspective (see Fig. 4b).

However, as we suggest below, correlation
$\mathcal{C}_{ij}^{\left(0,l\right)}$ intrinsically features a non-zero and
finite (i.e., not being infinitesimal) variance in deep neural networks.
Before our formal explanations, we first note that any input $\mathbf{X}$
(e.g., data) of deep neural networks should have a finite dimension even
though we apply classic mean-field approximation to consider infinite-width
neural networks. This is because the dimensionality of $\mathbf{X}$ is
determined by the learning task itself as _a priori_ and should not be
tampered. Given this property, let us consider a case where each component of
$\mathbf{X}$ is independently and identically distributed (i.e., uncorrelated)

$\displaystyle\mathcal{C}\left(\mathbf{X}_{i},\mathbf{X}_{j}\right):=\mathcal{C}_{ij}^{\left(0,0\right)}=\delta_{i,j}.$
(29)

By simple calculation based on Eqs. (20-22) and Eq. (29), we can derive the
correlation between the $i$-th component of an input signal, $\mathbf{X}_{i}$,
and the $j$-th component of propagated signal in the $1$-th layer

$\displaystyle\mathcal{C}_{ij}^{\left(0,1\right)}$
$\displaystyle=\frac{1}{\sigma^{\left(1\right)}_{j}}\mathbf{W}^{\left(1\right)}_{ij}\left(\int_{\mathbb{R}}\psi\left(\sigma_{i}z_{i}\right)z_{i}\mathsf{D}z_{i}\right),$
(30)

where $\sigma^{\left(1\right)}_{j}$ can be further calculated following Eq.
(25)

$\displaystyle\mathcal{C}_{ij}^{\left(0,1\right)}$
$\displaystyle=\frac{\mathbf{W}^{\left(1\right)}_{ij}\left(\int_{\mathbb{R}}\psi\left(\sigma_{i}z_{i}\right)z_{i}\mathsf{D}z_{i}\right)}{\sum_{k=1}^{N_{0}}\left(\mathbf{W}^{\left(1\right)}_{jk}\right)^{2}\mathbb{E}\left[\psi\left(\mathbf{X}_{k}\right)^{2}\right]}.$
(31)

Notion $N_{0}$ measures the dimension of input $\mathbf{X}$. Based on Eq. (31)
and Eqs. (20-22), it is trivial to derive the following variance items

$\displaystyle\operatorname{\mathbb{V}}\left(\mathcal{C}_{ij}^{\left(0,1\right)}\right)$
$\displaystyle=\frac{1}{\mathbb{E}\left[\psi\left(\mathbf{X}_{k}\right)^{2}\right]^{2}}\operatorname{\mathbb{V}}\left(\frac{\mathbf{W}^{\left(1\right)}_{ij}}{\sum_{k=1}^{N_{0}}\left(\mathbf{W}^{\left(1\right)}_{jk}\right)^{2}}\right)$
$\displaystyle\times\left(\int_{\mathbb{R}}\psi\left(\sigma_{i}z_{i}\right)z_{i}\mathsf{D}z_{i}\right)^{2}$
(32)

and

$\displaystyle\operatorname{\mathbb{V}}\left(\mathcal{C}_{ij}^{\left(0,l\right)}\right)$
$\displaystyle=\frac{\sigma_{w}^{2}\left(\int_{\mathbb{R}}\psi\left(\sigma_{i}z_{i}\right)z_{i}\mathsf{D}z_{i}\right)^{2}}{\left(\sigma^{\left(l\right)}_{j}\right)^{2}}\operatorname{\mathbb{V}}\left(\mathcal{C}_{ij}^{\left(0,l-1\right)}\right),$
(33)

where we momentarily use $\operatorname{\mathbb{V}}\left(\cdot\right)$ to
denote the variance to avoid confusions on mathematical symbols.

Let us primarily prove that
$\operatorname{\mathbb{V}}\left(\mathcal{C}_{ij}^{\left(0,l\right)}\right)$ in
Eq. (33) can not be an infinitesimal approaching to $0$ (i.e.,
$\operatorname{\mathbb{V}}\left(\mathcal{C}_{ij}^{\left(0,l\right)}\right)$ is
a finite real number). Because $N_{0}\in\mathbb{N}^{+}$ is finite, we know
that each
$\mathbf{W}^{\left(1\right)}_{ij}\;\text{i.i.d.}\sim\mathcal{N}\left(0,\frac{\sigma_{w}^{2}}{N_{0}}\right)$
has a finite variance, which further implies that
$\operatorname{\mathbb{V}}\left(\frac{\mathbf{W}^{\left(1\right)}_{ij}}{\sum_{k=1}^{N_{0}}\left(\mathbf{W}^{\left(1\right)}_{jk}\right)^{2}}\right)=\mathbb{E}\left(\frac{\left(\mathbf{W}^{\left(1\right)}_{ij}\right)^{2}}{\left(\sum_{k=1}^{N_{0}}\left(\mathbf{W}^{\left(1\right)}_{jk}\right)^{2}\right)^{2}}\right)$
is finite. Because the rest part of terms in Eq. (32) equal certain finite
real numbers, a finite value of
$\operatorname{\mathbb{V}}\left(\frac{\mathbf{W}^{\left(1\right)}_{ij}}{\sum_{k=1}^{N_{0}}\left(\mathbf{W}^{\left(1\right)}_{jk}\right)^{2}}\right)$
makes $\mathcal{C}_{ij}^{\left(0,1\right)}$ intrinsically have a finite
variance. According to Eq. (33), this property further makes
$\operatorname{\mathbb{V}}\left(\mathcal{C}_{ij}^{\left(0,l\right)}\right)$
finite for any $l\in\mathbb{N}^{+}$. Please see Fig. 4b for a summary.

Given that
$\operatorname{\mathbb{V}}\left(\mathcal{C}_{ij}^{\left(0,l\right)}\right)$ in
Eq. (33) is a finite number, we turn to proving
$\operatorname{\mathbb{V}}\left(\mathcal{C}_{ij}^{\left(0,l\right)}\right)\neq
0$. The proof can be readily derived from the following facts. First, term
$\frac{1}{\mathbb{E}\left[\psi\left(\mathbf{X}_{k}\right)^{2}\right]^{2}}$ in
Eq. (32) can not be $0$ because the squared output of an appropriate non-
linear activation function $\psi\left(\cdot\right)$ can not be infinite
(otherwise deep neural networks inevitably involve with numerical issues).
Second, term
$\int_{\mathbb{R}}\psi\left(\sigma_{i}z_{i}\right)z_{i}\mathsf{D}z_{i}$ can
not be $0$ because the integral of the product of $\psi\left(\cdot\right)$ and
$z_{i}$, two odd functions, over $\mathbb{R}$ must be non-zero. Third, term
$\frac{\mathbf{W}^{\left(1\right)}_{ij}}{\sum_{k=1}^{N_{0}}\left(\mathbf{W}^{\left(1\right)}_{jk}\right)^{2}}$
can not have a zero variance otherwise all weights in the $1$-st layer become
the same and immutable. Taken together, we know that
$\operatorname{\mathbb{V}}\left(\mathcal{C}_{ij}^{\left(0,1\right)}\right)\neq
0$ always holds in Eq. (32). Based on the iterative dynamics defined in Eq.
(33), it is not difficult to see that
$\operatorname{\mathbb{V}}\left(\mathcal{C}_{ij}^{\left(0,l\right)}\right)\neq
0$ holds for any $l\in\mathbb{N}^{+}$ because the coefficient term of
$\operatorname{\mathbb{V}}\left(\mathcal{C}_{ij}^{\left(0,l-1\right)}\right)$
can be trivially proven as non-zero following a similar idea. Here we no-
longer repeatedly elaborate these details. One can see Fig. 4b for a summary.

In sum, even from an empirical perspective, the independent and identical
assumption is invalid in correlation measurement because
$\mathcal{C}_{ij}^{\left(0,l\right)}$ is a Gaussian variable with non-zero and
finite variance, whose support may include illogical values (i.e., smaller
than $-1$ or larger than $1$)

##  IV Mutual information maximization at dynamic isometry

In this section, we present our theory on the possibility for neural networks
to be initialized toward optimal information channels. To overcome the
limitation of classic mean-field approximation, we propose a corrected mean-
field approximation framework that does not imply a Gaussian distribution of
correlation with non-zero and finite variance. Then, we introduce Gaussian
information bottleneck [60, 61] into our analysis, which relates our objective
of optimizing $\mathcal{I}\left(\phi\left(\mathbf{X}\right);\mathbf{X}\right)$
with maximizing a lower bound of mutual information and supports analytic
derivations. Based on it, we analytically prove that mutual information
between input and propagated signals is maximized at dynamic isometry.

###  IV.1 Corrected mean-field approximation and Gaussian information
bottleneck

As suggested in Sec. III.1, the key limitation of classic mean-field
approximation arises from the independent and identical assumption applied on
all variables without constraints. Although some variables, such as weight and
bias, can be assumed as independently and identically distributed (i.e., each
$\mathbf{W}^{\left(l\right)}_{ij}\;\text{i.i.d.}\sim\mathcal{N}\left(0,\frac{\sigma_{w}^{2}}{N_{l-1}}\right)$
and each
$\varepsilon^{\left(l\right)}_{i}\;\text{i.i.d.}\sim\mathcal{N}\left(0,\sigma_{b}^{2}\right)$),
it is less reasonable to apply independent and identical assumption on the
components of input $\mathbf{X}$ (i.e., each
$\mathbf{X}_{i}\;\text{i.i.d.}\sim\text{a certain distribution}$). This is
because the joint distribution of the components of $\mathbf{X}$ has been
defined by the learning task as _a priori_ and should not be modified.
Meanwhile, this unreasonable independent and identical assumption also makes
$\mathbf{X}_{i}^{\left(l\right)}\;\text{i.i.d.}\sim\text{a certain
distribution}$ for each $l\in\mathbb{N}^{+}$, which eventually leads to the
central limit theorem in Eq. (28) and implies a Gaussian distribution of
correlation $\mathcal{C}_{ij}^{\left(0,l\right)}$.

Given the above analysis, a natural idea for developing a corrected mean-field
approximation is to exclude the independent and identical assumption on the
components of $\mathbf{X}$ (see Fig. 5a). In the corrected approximation,
there is no constraint on the joint distribution of $\mathbf{X}_{i}$.
Therefore, $\mathbf{X}_{i}^{\left(l\right)}$ in the propagated signal may not
be independently and identically distributed as well. A direct consequence of
this correction lies in that the central limit theorem in Eq. (28) no longer
holds and the empirical distribution of $\mathcal{C}_{ij}^{\left(0,l\right)}$
does not converges to a Gaussian distribution (see Fig. 5a for illustration).
Consequently, the corrected mean-field approximation does not suffer from the
limitation of the classic one while characterizing neural networks as
information channels. Certainly, this correction also proposes critical
challenges to analytic derivations because the distributions of
$\mathbf{X}_{i}^{\left(l\right)}$ and
$\mathcal{C}_{ij}^{\left(0,l\right)}=\frac{\beta}{\sigma^{\left(l\right)}_{j}}\sum_{k=1}^{N_{l-1}}\mathbf{W}^{\left(l\right)}_{jk}\mathcal{C}_{ik}^{\left(0,l-1\right)}$
lack close-form expressions in most general cases.

To create a possibility for analytic derivations, we suggest to include
Gaussian information bottleneck [60, 61] into our analysis. In general, we can
consider a transform that maps $\mathbf{X}$ to an arbitrary Gaussian variable
$\mathbf{Z}$ (there is no constraint on the first and second moments of
$\mathbf{Z}$). With an appropriate transform, we can principally treat
$\mathbf{Z}$ as the “Gaussian part” of $\mathbf{X}$. As suggested by Ref.
[60], optimizing the information bottleneck or mutual information associated
with $\mathbf{Z}$ will also reflect the corresponding optimization associated
with $\mathbf{X}$ (see Fig. 5b). The benefit of such a connection lies in that
Gaussian information bottleneck and Gaussian mutual information have analytic
expressions and clear mathematical properties to support our analysis [60,
61]. Following the idea of Gaussian information bottleneck [60, 61], we
suggest to empirically consider a Gaussian counterpart $\mathbf{Z}$ of a given
$\mathbf{X}$ where $\mathbf{Z}$ is derived following the approach introduced
in Ref. [60]. The algorithm proposed by Ref. [60] ensures that the derived
$\mathbf{Z}$ is a “Gaussian part” of $\mathbf{X}$ with
$\mathcal{I}\left(\phi\left(\mathbf{Z}\right);\mathbf{Z}\right)\leq\mathcal{I}\left(\phi\left(\mathbf{X}\right);\mathbf{X}\right)$.
There is no other constraint on $\mathbf{Z}$. Meanwhile, signal
$\mathbf{Z}^{\left(l\right)}$ is not required to maintain its Gaussian
properties during information propagation in a neural network (i.e.,
$\phi\left(\mathbf{Z}\right)$, the output of a neural network, can be an
arbitrary variable).

In our work, we do not need to explicitly consider the actual form of
$\mathbf{Z}$ derived following Ref. [60]. This is because the following
inequality holds irrespective of what detailed properties that $\mathbf{Z}$
features as long as $\mathbf{Z}$ is a Gaussian variable (see Fig. 5b)

$\displaystyle\mathcal{I}\left(\mathbf{Z};\mathbf{Z}^{\left(l\right)}\right)$
$\displaystyle\geq\mathcal{I}_{\mathcal{N}}\left(\mathbf{Z};\mathbf{Z}^{\left(l\right)}\right),$
(34)
$\displaystyle=\frac{1}{2}\log\left(\frac{\operatorname{det}\left(\mathbf{\Sigma}^{\left(0\right)}\right)\operatorname{det}\left(\mathbf{\Sigma}^{\left(l\right)}\right)}{\operatorname{det}\left(\mathbf{\Sigma}^{\left(0,l\right)}\right)}\right),$
(35)

where $\Sigma^{\left(0\right)}$ and $\Sigma^{\left(l\right)}$ denote the
covariance matrix of $\mathbf{Z}$ and $\mathbf{Z}^{\left(l\right)}$,
respectively (i.e., each element in the matrix denotes the covariance between
a pair of components of signals). Matrix $\Sigma^{\left(0,l\right)}$ is the
covariance matrix measured between the components of $\mathbf{Z}$ and
$\mathbf{Z}^{\left(l\right)}$, which is a direct generalization of
$\Sigma^{\left(0\right)}$ and $\Sigma^{\left(l\right)}$. Please see Fig. 3(c)
for illustrations. Notion
$\mathcal{I}_{\mathcal{N}}\left(\mathbf{Z};\mathbf{Z}^{\left(l\right)}\right)$
denotes a special case of mutual information where $\mathbf{Z}$ and
$\mathbf{Z}^{\left(l\right)}$ are jointly Gaussian while they maintain the
original moment properties (i.e., expectation and covariance remain the same).
Please see detailed proofs of Eqs. (34-35) in Appendix A. Meanwhile, one can
find an equivalent version of Eqs. (34-35) in Ref. [62].

Let us reversely think about the above process. We can analyze the case where
$\mathbf{Z}$ is an arbitrary Gaussian variable to measure mutual information
$\mathcal{I}\left(\phi\left(\mathbf{Z}\right);\mathbf{Z}\right)$. There always
exists a certain input $\mathbf{X}$ that satisfies
$\mathcal{I}\left(\phi\left(\mathbf{Z}\right);\mathbf{Z}\right)\leq\mathcal{I}\left(\phi\left(\mathbf{X}\right);\mathbf{X}\right)$
if we reversely solve the transform problem in Ref. [60]. As suggested by Eqs.
(34-35), we can primarily focus on
$\mathcal{I}_{\mathcal{N}}\left(\mathbf{Z};\mathbf{Z}^{\left(l\right)}\right)$,
a lower bound of
$\mathcal{I}\left(\phi\left(\mathbf{Z}\right);\mathbf{Z}\right)$ in analysis
because it has a closed-form expression. Once
$\mathcal{I}\left(\phi\left(\mathbf{Z}\right);\mathbf{Z}\right)$ is maximized
under specific condition, mutual information terms
$\mathcal{I}\left(\phi\left(\mathbf{Z}\right);\mathbf{Z}\right)$ and
$\mathcal{I}\left(\phi\left(\mathbf{X}\right);\mathbf{X}\right)$ are both
maximized because of their lower bound relations (see Fig. 5b for
illustration).

###  IV.2 Simultaneous maximization of mutual information and $\beta$

Figure 5: Conceptual illustrations of the corrected mean-field approximation
and the maximization of a lower bound of mutual information in the phase space
of information propagation. (a) The corrected mean-field approximation differs
from the classic one by excluding the independent and identical assumption on
the components of input signal $\mathbf{X}$. This difference makes central
limit theorem do not hold and imply a non-Gaussian distribution of correlation
$\mathcal{C}_{ij}^{\left(0,l\right)}$. (b) Although the corrected mean-field
approximation overcomes the limitation of the classic one, its loose
constraints make the analytic characterization of information channel
properties nearly impossible. To create the possibility of analytic
derivations, Gaussian information bottleneck is introduced into our analysis
to relate the arbitrarily distributed input signal $\mathbf{X}$ with a
multivariate Gaussian variable $\mathbf{Z}$. Maximizing
$\mathcal{I}_{\mathcal{N}}$, the lower bound of
$\mathcal{I}\left(\mathbf{Z};\mathbf{Z}^{\left(l\right)}\right)$, is closely
related to optimizing
$\mathcal{I}\left(\mathbf{X};\mathbf{X}^{\left(l\right)}\right)$ according to
Gaussian information bottleneck. Our theory analytically proves that
$\mathcal{I}_{\mathcal{N}}$ is maximized at dynamic isometry, which is
validated by computational experiments on real neural networks as well.

Given the importance of the lower bound of mutual information defined in Eqs.
(34-35), let us begin our formal analysis on it. To relate our calculations of
Eq. (35) with the phase space of information propagation, we can first
consider the mathematical connection between covariance matrix and correlation
matrix. In the limit of infinite-width, we can treat the variance of each
component $\mathbf{Z}^{\left(l\right)}_{i}$ of the propagated signal
$\mathbf{Z}^{\left(l\right)}$ as generally similar (i.e., the fluctuation of
variance is sufficiently small compared with network width such that
$\sigma^{\left(l\right)}_{k}$ principally maintains the same across different
$k$ in the $l$-th layer). This simplification supports us to follow the idea
underlying Eq. (22) to derive a correlation matrix

$\displaystyle\mathbf{C}^{\left(0,l\right)}=\frac{\beta}{\sqrt{N_{l-1}}}\mathbf{C}^{\left(0,l-1\right)}\left(\mathbf{W}^{\left(l\right)}/\sigma_{w}\right)^{2},$
(36)

where the $\left(i,j\right)$-th element in the matrix measures the correlation
between $\mathbf{Z}_{i}$ and $\mathbf{Z}^{\left(l\right)}_{j}$ (see Fig. 3c).
Term $\beta$ remains same as its definition in Sec III.1

$\displaystyle\beta$
$\displaystyle=\sigma_{w}\int_{\mathbb{R}}\frac{\psi\left(\sigma^{\left(l-1\right)}_{k}z_{k}\right)}{\sigma^{\left(l\right)}_{j}}z_{k}\mathsf{D}z_{k},$
(37)
$\displaystyle=\sigma_{w}\int_{\mathbb{R}}\frac{\psi\left(\sigma^{*}\left(\sigma_{w},\sigma_{b}\right)z\right)}{\sigma^{*}\left(\sigma_{w},\sigma_{b}\right)}z\mathsf{D}z,$
(38)

where Eq. (38) is derived by replacing $\sigma_{j}^{\left(l\right)}$ and
$\sigma_{k}^{\left(l-1\right)}$ with the stable fixed point
$\sigma^{*}\left(\sigma_{w},\sigma_{b}\right)$ mentioned in Sec. II.3. As
suggested later in Sec. V, this replacement is reasonable because the
convergence rates of $\sigma_{j}^{\left(l\right)}$ and
$\sigma_{k}^{\left(l-1\right)}$ to
$\sigma^{*}\left(\sigma_{w},\sigma_{b}\right)$ are high. Please note that that
subscript $k$ in Eq. (37) can be eventually dropped in Eq. (38) because
$\beta$ is same across different $k$ in the $\left(l-1\right)$-th layer. Based
on Eq. (36) and Eq. (38), we can derive the recursion equation of
$\mathbf{C}^{\left(0,l\right)}$

$\displaystyle\mathbf{C}^{\left(0,l\right)}=\beta^{l}\mathbf{C}^{\left(0,0\right)}\prod_{i=1}^{l}\frac{\left(\mathbf{W}^{\left(i\right)}/\sigma_{w}\right)^{T}}{\sqrt{N_{i-1}}}.$
(39)

Our next step is to relate Eq. (39) with the lower bound of mutual information
in Eq. (35). Applying the property of the determinant of block matrix, we can
reformulate term
$\operatorname{det}\left(\mathbf{\Sigma}^{\left(0,l\right)}\right)$ in Eq.
(35) as

$\displaystyle\operatorname{det}\left(\mathbf{\Sigma}^{\left(0,l\right)}\right)=\operatorname{det}\left(\mathbf{\Sigma}^{\left(l\right)}\right)$
$\displaystyle\times\operatorname{det}\left(\mathbf{\Sigma}^{\left(0\right)}-\mathbf{D}^{\left(0\right)}\mathbf{C}^{\left(0,l\right)}\mathbf{D}^{\left(l\right)}\left(\mathbf{\Sigma}^{\left(l\right)}\right)^{-1}\mathbf{D}^{\left(l\right)}\left(\mathbf{C}^{\left(0,l\right)}\right)^{T}\mathbf{D}^{\left(0\right)}\right),$
(40)

where
$D^{\left(l\right)}=\operatorname{diag}\left(\mathbf{\Sigma}^{\left(l\right)}\right)^{\frac{1}{2}}$
and
$D^{\left(0\right)}=\operatorname{diag}\left(\mathbf{\Sigma}^{\left(0\right)}\right)^{\frac{1}{2}}$,
i.e., their diagonal elements are the standard deviations of signals. Please
note that $\mathbf{Z}$ and $\mathbf{Z}^{\left(l\right)}$ are jointly Gaussian
while calculating
$\mathcal{I}_{\mathcal{N}}\left(\mathbf{Z};\mathbf{Z}^{\left(l\right)}\right)$.
Therefore, their covariance matrices, $\mathbf{\Sigma}^{\left(0\right)}$ and
$\mathbf{\Sigma}^{\left(l\right)}$, are invertible in Eq. (40). Then, we can
notice that
$\mathbf{\Sigma}^{\left(0\right)}=\mathbf{D}^{(0)}\mathbf{C}^{\left(0,0\right)}\mathbf{D}^{(0)}$
and
$\mathbf{\Sigma}^{\left(l\right)}=\mathbf{D}^{(l)}\mathbf{C}^{\left(l,l\right)}\mathbf{D}^{(l)}$.
Therefore, we can reformulate the lower bound of mutual information as

$\displaystyle\mathcal{I}_{\mathcal{N}}\left(\mathbf{Z};\mathbf{Z}^{\left(l\right)}\right)$
$\displaystyle=$
$\displaystyle\frac{1}{2}\log\left(\frac{\operatorname{det}\left(\mathbf{\Sigma}^{\left(0\right)}\right)\operatorname{det}\left(\mathbf{\Sigma}^{\left(l\right)}\right)}{\operatorname{det}\left(\mathbf{\Sigma}^{\left(0,l\right)}\right)}\right),$
(41) $\displaystyle=$
$\displaystyle\frac{1}{2}\log\left(\frac{\operatorname{det}\left(\mathbf{\Sigma}^{\left(0\right)}\right)}{\operatorname{det}\left(\mathbf{\Sigma}^{\left(0\right)}-\mathbf{D}^{\left(0\right)}\mathbf{C}^{\left(0,l\right)}\left(\mathbf{C}^{\left(l,l\right)}\right)^{-1}\left(\mathbf{C}^{\left(0,l\right)}\right)^{T}\mathbf{D}^{\left(0\right)}\right)}\right),$
(42) $\displaystyle=$
$\displaystyle\frac{1}{2}\log\left(\frac{\operatorname{det}\left(\mathbf{C}^{\left(0,0\right)}\right)}{\operatorname{det}\left(\mathbf{C}^{\left(0,0\right)}-\mathbf{C}^{\left(0,l\right)}\left(\mathbf{C}^{\left(l,l\right)}\right)^{-1}\left(\mathbf{C}^{\left(0,l\right)}\right)^{T}\right)}\right).$
(43)

In Eq. (42), we have replaced
$\mathbf{D}^{\left(l\right)}\left(\mathbf{\Sigma}^{\left(l\right)}\right)^{-1}\mathbf{D}^{\left(l\right)}$
with $\left(\mathbf{C}^{\left(l,l\right)}\right)^{-1}$ for simplification. Eq.
(43) is obtained by dividing the numerator and denominator of Eq. (42) by
$\left(\operatorname{det}\left(\mathbf{D}^{\left(0\right)}\right)\right)^{2}$
simultaneously. In Eq. (43), matrix $\mathbf{C}^{\left(l,l\right)}$ is a
constant matrix across different layers when signals arrive at their stable
states shown in Eq. (38). Matrix $\mathbf{C}^{\left(0,0\right)}$ is fully
determined by input $\mathbf{Z}$ (see Fig. 6a).

After substituting Eq. (39) into Eq. (43) and dividing the numerator and
denominator of the derived result by $\mathbf{C}^{\left(0,0\right)}$, we can
obtain the following equation

$\displaystyle\mathcal{I}_{\mathcal{N}}\left(\mathbf{X};\mathbf{X}^{\left(l\right)}\right)=\frac{1}{2}\log\left(\frac{1}{\operatorname{det}\left(\mathbf{I}-\beta^{2l}\prod_{i=1}^{l}\frac{\left(\mathbf{W}^{\left(i\right)}/\sigma_{w}\right)^{T}}{\sqrt{N_{i-1}}}\left(\mathbf{C}^{\left(l,l\right)}\right)^{-1}\prod_{i=l}^{1}\frac{\mathbf{W}^{\left(i\right)}/\sigma_{w}}{\sqrt{N_{i-1}}}\mathbf{C}^{\left(0,0\right)}\right)}\right).$
(44)

Please note that we have replaced $\left(\mathbf{C}^{(0,0)}\right)^{T}$ with
$\mathbf{C}^{(0,0)}$ to improve the readability of Eq. (44) because
$\mathbf{C}^{(0,0)}$ is a symmetric matrix. In Eq. (44), term
$\prod_{i=1}^{l}\frac{\mathbf{W}^{\left(i\right)}/\sigma_{w}}{\sqrt{N_{i-1}}}$
is completely determined by the type of weight distribution used in neural
network initialization. Meanwhile, matrices $\mathbf{C}^{\left(0,0\right)}$
and $\mathbf{C}^{\left(l,l\right)}$ have been suggested as fully
deterministic. Therefore, term $\beta=\beta\left(\sigma_{w},\sigma_{b}\right)$
is the only one non-trivial variable in Eq. (44) remaining for analysis.

Our analysis guides us to focus on the possibility that
$\mathcal{I}_{\mathcal{N}}\left(\mathbf{Z};\mathbf{Z}^{\left(l\right)}\right)$
can monotonically increase with
$\beta=\beta\left(\sigma_{w},\sigma_{b}\right)$ in Eq. (44) (see Fig. 6a for
illustration). Below, we present our detailed derivations. For convenience, we
define

$\displaystyle\widetilde{\mathbf{M}}=\prod_{i=1}^{l}\frac{\left(\mathbf{W}^{\left(i\right)}/\sigma_{w}\right)^{T}}{\sqrt{N_{i-1}}}\left(\mathbf{C}^{\left(l,l\right)}\right)^{-1}\prod_{i=l}^{1}\frac{\mathbf{W}^{\left(i\right)}/\sigma_{w}}{\sqrt{N_{i-1}}}.$
(45)

as a shorthand. We notice that we can equivalently verify whether
$\operatorname{det}\left(\mathbf{I}-\beta^{2l}\widetilde{\mathbf{M}}\mathbf{C}^{\left(0,0\right)}\right)$
monotonically decreases with $\beta$ since
$\mathcal{I}_{\mathcal{N}}\left(\mathbf{X};\mathbf{X}^{\left(l\right)}\right)$
monotonically decreases with
$\operatorname{det}\left(\mathbf{I}-\beta^{2l}\widetilde{\mathbf{M}}\mathbf{C}^{\left(0,0\right)}\right)$.
Because $\mathbf{C}^{\left(0,0\right)}$ is a positive definite matrix, we can
apply Cholesky factorization on it, i.e.,
$\mathbf{C}^{\left(0,0\right)}=LL^{T}$ where $L$ is a lower triangular matrix
whose diagonal elements are positive. Then, we have

$\displaystyle\operatorname{det}\left(\mathbf{I}-\beta^{2l}\widetilde{\mathbf{M}}\mathbf{C}^{\left(0,0\right)}\right)$
$\displaystyle=$
$\displaystyle\operatorname{det}\left(\mathbf{I}-\beta^{2l}\widetilde{\mathbf{M}}LL^{T}\right),$
(46) $\displaystyle=$
$\displaystyle\operatorname{det}\left(L^{T}\left(L^{T}\right)^{-1}-\beta^{2l}L^{T}\widetilde{\mathbf{M}}LL^{T}\left(L^{T}\right)^{-1}\right),$
(47) $\displaystyle=$
$\displaystyle\operatorname{det}\left(\mathbf{I}-\beta^{2l}L^{T}\widetilde{\mathbf{M}}L\right).$
(48)

Given that $\widetilde{\mathbf{M}}$ is a positive semi-definite matrix, we
know that $L^{T}\widetilde{\mathbf{M}}L$ is also positive semi-definite. Let
$\\{\omega_{i}|i=1,\ldots,\operatorname{dim}\left(\widetilde{\mathbf{M}}\right),\;\omega_{i}>0\\}$
be a set of eigenvalues of matrix $L^{T}\widetilde{\mathbf{M}}L$, we have

$\displaystyle\operatorname{det}\left(\mathbf{I}-\beta^{2l}\widetilde{\mathbf{M}}\mathbf{C}^{\left(0,0\right)}\right)=\prod_{i=1}^{\operatorname{dim}\left(\widetilde{\mathbf{M}}\right)}(1-\beta^{2l}\omega_{i}).$
(49)

Let us assume that the range of $\beta^{2}$ is
$\left[0,\widehat{\beta}^{2}\right)$, where $0$ stands for zero correlation.
To ensure the non-negativity of
$\mathcal{I}_{\mathcal{N}}\left(\mathbf{X};\mathbf{X}^{\left(l\right)}\right)$
(i.e., mutual information can not be negative), we know that
$\operatorname{det}\left(\mathbf{I}-\beta^{2l}\widetilde{\mathbf{M}}\mathbf{C}^{\left(0,0\right)}\right)$,
a continuous function with respect to $\beta^{2l}$, should be in an interval
of $\left(0,1\right]$. According to Eq. (49), this non-negativity requires
that $\omega_{i}\in\left(0,\widehat{\beta}^{-2l}\right)$. Meanwhile, to ensure
that
$\operatorname{det}\left(\mathbf{I}-\beta^{2l}\widetilde{\mathbf{M}}\mathbf{C}^{\left(0,0\right)}\right)$
can monotonically decreases with
$\beta=\beta\left(\sigma_{w},\sigma_{b}\right)$, we can derive
$\omega_{i}\in\left(0,\widehat{\beta}^{-2l}\right)$ based on the continuity
and non-negativity. We can immediately find that these two requirements of
$\omega_{i}$ are consistent with each other. Therefore, we can prove that
$\mathcal{I}_{\mathcal{N}}\left(\mathbf{Z};\mathbf{Z}^{\left(l\right)}\right)$
monotonically increases with $\beta\left(\sigma_{w},\sigma_{b}\right)$ (see
Fig. 6a for a summary).

Given a specific weight distribution and a certain network depth defined for
neural network initialization, terms $\mathbf{C}^{\left(0,0\right)}$,
$\prod_{i=1}^{l}\frac{\mathbf{W}^{\left(i\right)}/\sigma_{w}}{\sqrt{N_{i-1}}}$,
and $\mathbf{C}^{\left(l,l\right)}$ are principally fixed. Therefore,
initializing neural networks for mutual information maximization essentially
requires us to maximize $\beta\left(\sigma_{w},\sigma_{b}\right)$, whose
condition is analyzed in Sec. IV.3.

Figure 6: Conceptual illustrations of the maximization of the lower bound of
mutual information at dynamic isometry. (a) Because
$\mathcal{I}_{\mathcal{N}}\left(\mathbf{Z};\mathbf{Z}^{\left(l\right)}\right)$,
the lower bound of mutual information, monotonically increases with $\beta\leq
1$, maximizing
$\mathcal{I}_{\mathcal{N}}\left(\mathbf{Z};\mathbf{Z}^{\left(l\right)}\right)$
is equivalent to maximizing $\beta\leq 1$, whose condition is analytically
derived as $\sigma^{*}\rightarrow 0$. Meanwhile, the condition of dynamic
isometry can also be proven as $\sigma^{*}\rightarrow 0$ based on the free
probability theory. (b) Consequently, there are equivalent relations among
dynamic isometry, lower bound maximization, and $\beta=1$. As shown in the
analytically calculated plane of $\beta\left(\sigma_{w},\sigma_{b}\right)$
(the edge of chaos is represented by a purple line while dynamic isometry
point is marked by a star), there is $\beta=1$ at dynamic isometry (DI),
suggesting the maximization of
$\mathcal{I}_{\mathcal{N}}\left(\mathbf{Z};\mathbf{Z}^{\left(l\right)}\right)$.
Here the plane of $\beta\left(\sigma_{w},\sigma_{b}\right)$ is calculated
based on a commonly used non-linear activation function,
$\tanh\left(\cdot\right)$, in deep learning.

###  IV.3 Mutual information maximization at dynamic isometry

As we have proved in Sec. IV.2, the maximization of the lower bound of mutual
information is equivalent to the maximization of
$\beta\left(\sigma_{w},\sigma_{b}\right)$. Below, we prove that
$\beta\left(\sigma_{w},\sigma_{b}\right)$ is maximized at dynamic isometry.
For convenience, we use $\beta$, $\zeta$, and $\sigma^{*}$ as the shorthands
of $\beta\left(\sigma_{w},\sigma_{b}\right)$,
$\zeta\left(\sigma_{w},\sigma_{b}\right)$, and
$\sigma^{*}\left(\sigma_{w},\sigma_{b}\right)$ in our derivations.

Reformulating the one-dimensional integral as two-dimensional integral by
Cauchy–Schwarz inequality, we can derive

$\displaystyle\left(\int_{\mathbb{R}}\psi\left(\sigma^{*}z\right)z\mathsf{D}z\right)^{2}$
$\displaystyle=\int_{\mathbb{R}}\int_{\mathbb{R}}\psi\left(\sigma^{*}x\right)\psi\left(\sigma^{*}y\right)xy\mathsf{D}x\mathsf{D}y,$
(50)
$\displaystyle\leq\int_{\mathbb{R}}\int_{\mathbb{R}}\psi\left(\sigma^{*}x\right)^{2}y^{2}\mathsf{D}x\mathsf{D}y,$
(51)

where equality holds only if $\psi\left(\sigma^{*}x\right)^{2}\propto x^{2}$.
Because $\sigma^{*}$ satisfies

$\displaystyle\left(\sigma^{*}\right)^{2}=\sigma_{w}^{2}\int_{\mathbb{R}}\psi\left(\sigma^{*}z\right)^{2}\mathsf{D}z+\sigma_{b}^{2},$
(52)

we can combine Eq. (52) with Eq. (38) and Eq. (51) to prove $\beta<1$ when
$\sigma^{*}\neq 0$

$\displaystyle\left(\int_{\mathbb{R}}\psi\left(\sigma^{*}z\right)z\mathsf{D}z\right)^{2}$
$\displaystyle<\int_{\mathbb{R}}\int_{\mathbb{R}}\psi\left(\sigma^{*}x\right)^{2}y^{2}\mathsf{D}x\mathsf{D}y,$
(53)
$\displaystyle\sigma_{w}^{2}\left(\int_{\mathbb{R}}\psi\left(\sigma^{*}z\right)z\mathsf{D}z\right)^{2}$
$\displaystyle<\left(\sigma^{*}\right)^{2},$ (54) $\displaystyle\beta^{2}$
$\displaystyle<1.$ (55)

As for the situation where $\sigma^{*}\rightarrow 0$, we can prove that
$\mathcal{I}_{\mathcal{N}}\left(\mathbf{X};\mathbf{X}^{\left(l\right)}\right)$
is maximized at dynamic isometry by proving $\beta=1$ under the corresponding
condition. As _a priori_ knowledge, we can know
$\lim_{z\rightarrow+\infty}\psi\left(z\right)=c\in\mathbb{R}$ (i.e., the
integral result is a constant) and
$|\psi^{\prime}\left(0\right)|\geq|\psi^{\prime}\left(z\right)|$ because
$\psi\left(\cdot\right)$, an activation function of neural networks, is
usually an odd function that is convex in $\left[0,\infty\right]$ and
satisfies $\lim_{z\rightarrow+\infty}\psi\left(z\right)=c\in\mathbb{R}$. These
properties support us to derive the following proof. First, based on the
convex property of $\psi\left(\cdot\right)$ and
$1=\int_{\mathbb{R}}z^{2}\mathsf{D}z$ (i.e., $\mathsf{D}z$ is a Gaussian
measure), we can derive

$\displaystyle\left(\sigma^{\left(l\right)}\right)^{2}=\int_{\mathbb{R}}\left(\sigma^{\left(l\right)}\right)^{2}z^{2}\mathsf{D}z\geq\sigma_{w}^{2}\int_{\mathbb{R}}\psi\left(\sigma^{\left(l\right)}z\right)^{2}\mathsf{D}z$
(56)

when $\sigma_{w}\leq\frac{1}{\psi^{\prime}\left(0\right)}$. Here the equality
holds only if $\sigma^{\left(l\right)}=0$. Based on Eq. (56), we can see that
any $\left(\sigma^{\left(l\right)}\right)^{2}$ will decrease until it arrives
at $\left(\sigma^{*}\right)^{2}=0$ when
$\sigma_{w}\leq\frac{1}{\psi^{\prime}\left(0\right)}$ and $\sigma_{b}=0$. In
other words, point $\sigma^{*}\rightarrow 0$ is a stable fixed point only if
$\sigma_{w}\leq\frac{1}{\psi^{\prime}\left(0\right)}$ and $\sigma_{b}=0$.
Second, given the condition for $\sigma^{*}\rightarrow 0$ to become a stable
fixed point, we can further prove that $\beta=1$ may emerge under this
condition. Our derivation utilizes an equality obtained in Appendix B.

$\displaystyle\psi^{\prime}\left(0\right)=\lim_{\sigma^{*}\rightarrow
0}\int_{\mathbb{R}}\frac{\psi\left(\sigma^{*}z\right)}{\sigma^{*}}z\mathsf{D}z.$
(57)

Substituting Eq. (57) into Eq. (38), we can see the desired combination of
$\left(\sigma_{w},\sigma_{b}\right)$ for $\beta=1$ when $\sigma^{*}\rightarrow
0$

$\displaystyle\sigma_{w}$
$\displaystyle=\frac{\beta}{\lim_{\sigma^{*}\rightarrow
0}\int_{\mathbb{R}}\frac{\psi\left(\sigma^{*}z\right)}{\sigma^{*}\left(\sigma_{w},\sigma_{b}\right)}z\mathsf{D}z}=\frac{1}{\psi^{\prime}\left(0\right)},$
(58) $\displaystyle\sigma_{b}$ $\displaystyle=0.$ (59)

To this point, we have derived Eqs. (58-59) as the sufficient and necessary
condition for $\beta$ to take its maximum, i.e.
$\mathcal{I}_{\mathcal{N}}\left(\mathbf{X};\mathbf{X}^{\left(l\right)}\right)$
takes its maximum, which will be shown as exactly the condition of dynamical
isometry in our subsequent analysis.

Although previous studies have studied the condition of dynamical isometry in
linear [63] and non-linear neural networks [47, 49] as suggested in Sec. II.3,
it remains unclear if it is possible to relate the condition of dynamical
isometry with our theory. Below, we present our detailed derivations of
dynamic isometry point
$\left(\sigma_{w}^{\diamond},\sigma_{b}^{\diamond}\right)$ based on free
probability theory [59] to suggest such a possibility.

The key idea in our derviations is to take advantage of the property of
$\mathcal{S}$-transform concerning matrix multiplication [64, 65]

$\displaystyle\mathcal{S}_{UV}\left(z\right)=\mathcal{S}_{U}\left(z\right)\mathcal{S}_{V}\left(z\right),$
(60)

where $U$ and $V$ are two freely independent random matrices. Because the
Jacobian matrix of the neural network can be defined as

$\displaystyle\mathbf{J}=\frac{\partial\mathbf{X}^{\left(l\right)}}{\partial\mathbf{X}}=\prod_{i=1}^{l}\mathbf{\widetilde{D}}^{\left(i\right)}\mathbf{W}^{\left(i\right)},$
(61)

we can derive the $\mathcal{S}$-transform of the Jacobian matrix

$\displaystyle\mathcal{S}_{\mathbf{J}\mathbf{J}^{T}}\left(z\right)=\mathcal{S}^{L}_{\left(\mathbf{\widetilde{D}}^{\left(i\right)}\right)^{2}}\mathcal{S}^{L}_{\left(\mathbf{W}^{\left(i\right)}\right)^{2}\mathbf{W}^{\left(i\right)}}\left(z\right),$
(62)

where $\mathbf{\widetilde{D}}^{\left(i\right)}$ is a diagonal matrix whose
diagonal elements are
$\mathbf{\widetilde{D}}^{(i)}_{jj}=\psi^{\prime}\left(\mathbf{X}^{\left(i\right)}_{j}\right)$
(please note that the definition is different from the matrix
$\mathbf{{D}}^{(i)}$ analyzed before). In the derivation of Eq. (62), we have
applied that
$\mathbf{\widetilde{D}}^{\left(a\right)}=\mathbf{\widetilde{D}}^{\left(b\right)}$
and $\mathbf{W}^{\left(a\right)}=\mathbf{W}^{\left(b\right)}$ for any pair of
$\left(a,b\right)$. This property holds because every layer in the neural
network shares the same network initialization settings and signals in every
layer share the same marginal distribution if they are at the stable fixed
point.

Following the idea in Ref. [49], we can derive an implicit equation for
eigenvalue spectrum of
$\mathbf{H}^{\left(l\right)}=\left(\mathbf{J}^{\left(l\right)}\right)^{T}\mathbf{J}^{\left(l\right)}$
based on Eq. (62)

$\displaystyle M_{\mathbf{H}^{\left(l\right)}}(z)=$ $\displaystyle
M_{\mathbf{\widetilde{D}}^{2}}\left(z^{\frac{1}{l}}\mathcal{S}_{\mathbf{W}^{T}\mathbf{W}}\left(M_{\mathbf{H}^{\left(l\right)}}(z)\right)\left(1+\frac{1}{M_{\mathbf{H}^{\left(l\right)}}(z)}\right)^{1-\frac{1}{l}}\right),$
(63)

where $M$ is a moment generating function. After expanding Eq. (63) in the
powers of $z^{-1}$, the expression of the first two moments of the eigenvalue
spectrum of $\mathbf{H}^{\left(l\right)}$ can be obtained as

$\displaystyle m_{1}=\left(\sigma_{w}^{2}\tilde{\mu}_{1}\right)^{l},$ (64)
$\displaystyle
m_{2}=\left(\sigma_{w}^{2}\tilde{\mu}_{1}\right)^{2l}\left(\frac{\tilde{\mu}_{2}}{\tilde{\mu}_{1}^{2}}+\frac{1}{l}-1-s_{1}\right)l,$
(65)

where we define
$\tilde{\mu}_{k}=\int\psi^{\prime}\left(\sigma^{*}z\right)^{2k}\mathsf{D}z$.
Meanwhile, one can notice that $m_{1}$ is exactly equivalent to $\zeta$
defined in Eqs. (16-17). For scaled orthogonal initialization (i.e., weight
matrices are initialized as orthogonal random matrices), we have $s_{1}=0$.
For scaled Gaussian initialization (i.e., weight matrices are initialized as
Gaussian random matrices), we have $s_{1}=-1$. More calculation details of
$s_{1}$ can be seen in Ref. [49].

As suggested in Sec. II.3, dynamic isometry requires that the first moment of
the eigenvalue spectrum of $\mathbf{H}^{\left(l\right)}$ equals $1$ while the
second moment approaches to $0$. In deep neural networks, we reasonably relax
the restriction on the second moment and require that the second moment does
not increases with network depth $l$ significantly. Applying the Cauchy-
Schwarz inequality, we can derive

$\displaystyle\tilde{\mu}_{1}^{2}=\left(\int\psi^{\prime}\left(\sigma^{*}z\right)^{2}\mathsf{D}z\right)^{2}\leq\int\psi^{\prime}\left(\sigma^{*}z\right)^{4}\mathsf{D}z=\tilde{\mu}_{2},$
(66)

where the equality holds (i.e., $\tilde{\mu}_{1}^{2}=\tilde{\mu}_{2}$) if

$\displaystyle\psi^{\prime}\left(\sigma^{*}z\right)^{2}\propto\psi^{\prime}\left(\sigma^{*}z\right)^{4}\Leftrightarrow\sigma^{*}\rightarrow
0.$ (67)

In the case where $m_{1}=1$ and $\tilde{\mu}_{1}^{2}=\tilde{\mu}_{2}$, we can
readily obtain $m_{2}=1$ (i.e., the second moment is a small constant) for
orthogonal initialization and $m_{2}=l+1$ (i.e., the second moment increases
with network depth linearly instead of exhibiting explosive growth) for
Gaussian initialization. Moreover, we can know that $\sigma^{*}\rightarrow 0$
enables $m_{1}=1$ (or equivalently $\zeta=1$) to imply
$\sigma_{w}=\frac{1}{\psi^{\prime}\left(0\right)}$, which is exactly the
condition of the maximization of $\beta$ (or the maximization of the lower
bound of mutual information) defined in Eqs. (58-59) (see Fig. 6a for a
summary).

In sum, we have proven that the maximization of the lower bound of mutual
information and dynamic isometry shares the same condition (i.e., Eqs.
(58-59)). In other words, mutual information
$\mathcal{I}\left(\mathbf{Z};\mathbf{Z}^{\left(l\right)}\right)$ and its lower
bound
$\mathcal{I}_{\mathcal{N}}\left(\mathbf{Z};\mathbf{Z}^{\left(l\right)}\right)$
are maximized at dynamic isometry (see Fig. 6b). Because $\mathbf{Z}$ is a
“Gaussian part” of $\mathbf{X}$ with
$\mathcal{I}\left(\mathbf{Z}^{\left(l\right)};\mathbf{Z}\right)\leq\mathcal{I}\left(\mathbf{X}^{\left(l\right)};\mathbf{X}\right)$
(or equivalently
$\mathcal{I}\left(\phi\left(\mathbf{Z}\right);\mathbf{Z}\right)\leq\mathcal{I}\left(\phi\left(\mathbf{X}\right);\mathbf{X}\right)$)
as suggested in Sec. IV.1, we know that
$\mathcal{I}\left(\mathbf{X}^{\left(l\right)};\mathbf{X}\right)$ is maximized
at dynamic isometry in more general cases.

Figure 7: The plane of the lower bound of mutual information computationally
derived on real neural networks. Experiment settings, such as width, depth,
input type, and initialization type, are presented along corresponding planes.
Same as Fig. 6b, the edge of chaos (EOC) is marked by a purple line while
dynamic isometry is marked by a star.

##  V Experimental validations

In previous sections, we have present our main theory on the equivalence of
dynamic isometry and mutual information maximization, which is developed on
infinite-width neural networks. One may reasonably question whether our theory
is valid on real finite-width neural networks in deep learning. Below, we
validate our theory on real neural networks with various settings (e.g.,
different widths, depths, inputs, and initialization conditions).

In Fig. 7, we implement our experiments on finite-width neural networks with a
widely applied non-linear activation function, $\tanh\left(\cdot\right)$. To
comprehensively verify the robustness of our theory against finite size
effects, we design these neural networks with different widths and depths. To
suggest the applicability of our theory on more general cases where input
$\mathbf{Z}$ may not necessarily propagate at its stable state, we distinguish
between stable inputs (i.e., propagating at the stable state as our
theoretical derivations require) and unit variance inputs (i.e., with a unit
covariance matrix that have not been considered in our previous derivations).
To show the capacity of our theory to characterize orthogonal and Gaussian
initialization, we conduct experiments under both initialization conditions.
In our experiments, we measure the lower bound of mutual information between
input $\mathbf{Z}$ and output $\phi\left(\mathbf{Z}\right)$ (i.e.,
$\mathbf{Z}^{\left(l\right)}$ when $l$ stands for the last layer of the neural
network) based on Eq. (35), in which $\mathbf{\Sigma}^{\left(0\right)}$,
$\mathbf{\Sigma}^{\left(l\right)}$, and $\mathbf{\Sigma}^{\left(0,l\right)}$
are computationally estimated from the data. Given a combination of width,
depth, input type, and initialization type, we repeat our measurement under
each condition of $\left(\sigma_{w},\sigma_{b}\right)$ to obtain a plane of
the lower bound of mutual information. To offer a clear vision, we also
illustrate the measured lower bound of mutual information along the edge of
chaos given orthogonal or Gaussian initialization and stable inputs as two
instances (see Fig. 8). As shown in Fig. 7 and Fig. 8, our experiment results
are consistent with theoretical prediction that the lower bound of mutual
information is maximized at the dynamic isometry point
$\left(\sigma_{w}^{\diamond},\sigma_{b}^{\diamond}\right)=\left(1,0\right)$
(note that this point is confirmed by whether $\sigma^{*}\rightarrow 0$).
Moreover, the distribution of the lower bound of mutual information
corroborates the distribution of $\beta$ analytically calculated in Fig. 6b,
where both $\beta$ and the lower bound of mutual information on the edge of
chaos increase as the condition moves towards dynamic isometry.

In sum, we have observed consistency between our theory and experiment
results, which suggesting the applicability of our theory on real neural
networks in deep learning.

Figure 8: Two instances of the computationally measured lower bound of mutual
information on the edge of chaos. Experiment settings, such as width, depth,
input type, and initialization type, are presented. One can see that
$\mathcal{I}_{\mathcal{N}}$, the lower bound of mutual information, is
maximized near $\sigma_{b}=0$, which corresponds to the dynamic isometry
point.

##  VI Analysis with information bottleneck

To this point, we have analytically developed and computationally validated
our theory about mutual information maximization at dynamic isometry. As an
interdisciplinary research, our work not only focuses on the statistical
physics of neural networks but also aims at offering insights on deep learning
techniques. When
$\mathcal{I}_{\mathcal{N}}\left(\mathbf{Z};\mathbf{Z}^{\left(l\right)}\right)$,
the lower bound of mutual information
$\mathcal{I}\left(\mathbf{Z};\mathbf{Z}^{\left(l\right)}\right)$ is maximized
because the neural network is initialized at dynamic isometry, we can know
$\mathcal{I}\left(\mathbf{X};\mathbf{X}^{\left(l\right)}\right)$ is also
maximized according to
$\mathcal{I}\left(\mathbf{Z}^{\left(l\right)};\mathbf{Z}\right)\leq\mathcal{I}\left(\mathbf{X}^{\left(l\right)};\mathbf{X}\right)$
in Sec. IV.1. However, we have suggested that maximizing
$\mathcal{I}\left(\mathbf{X};\mathbf{X}^{\left(l\right)}\right)$ is not
equivalent to making the neural network an optimal channel in deep learning.
Below, we attempt to present a comprehensive analysis on the precise relation
between maximizing
$\mathcal{I}\left(\mathbf{X};\mathbf{X}^{\left(l\right)}\right)$ and driving
neural networks toward optimal channels.

Given the difference between unsupervised and supervised learning, we
subdivide our analysis into two cases:

  * (1)

In supervised learning, both sample, $\mathbf{X}$, and target, $\mathbf{Y}$,
are accessible for the neural network. Therefore, the optimization objective
reduces to the classic information bottleneck [35, 36, 37], which is actually
a special case of rate distortion theory [52] and sufficient statistics theory
[53]

$\displaystyle\max_{\phi}\mathcal{L}_{s}\left(\phi\right):=\underbrace{\mathcal{I}\left(\phi\left(\mathbf{X}\right);\mathbf{Y}\right)}_{\text{Encoding}}\underbrace{-\tau\mathcal{I}\left(\phi\left(\mathbf{X}\right);\mathbf{X}\right)}_{\text{Compression}}.$
(68)

In general, Eq. (68) defines an objective that the neural network maximizes
$\mathcal{I}\left(\phi\left(\mathbf{X}\right);\mathbf{Y}\right)$, the capacity
to learn $\mathbf{Y}$, during encoding and minimizes the complexity of
representation,
$\mathcal{I}\left(\phi\left(\mathbf{X}\right);\mathbf{X}\right)$, during
compression. During initialization, we suggest to maximize
$\mathcal{I}\left(\phi\left(\mathbf{X}\right);\mathbf{X}\right)$ by
initializing the neural network at dynamic isometry because we know
$\mathcal{I}\left(\phi\left(\mathbf{X}\right);\mathbf{Y}\right)\leq\mathcal{I}\left(\phi\left(\mathbf{X}\right);\mathbf{X}\right)$
according to the Markov chain in Eq. (1). This approach can avoid that
$\mathcal{I}\left(\phi\left(\mathbf{X}\right);\mathbf{Y}\right)$ is bounded by
a small value of
$\mathcal{I}\left(\phi\left(\mathbf{X}\right);\mathbf{X}\right)$ and can not
be thoroughly optimized during encoding. During compression, we suggest to
minimize $\mathcal{I}\left(\phi\left(\mathbf{X}\right);\mathbf{X}\right)$
while controlling the loss of
$\mathcal{I}\left(\phi\left(\mathbf{X}\right);\mathbf{Y}\right)$ (i.e.,
ensuring
$\Delta|\mathcal{I}\left(\phi\left(\mathbf{X}\right);\mathbf{Y}\right)|<\kappa$
where $\kappa\rightarrow 0$) to avoid neural network over-fitting.

  * (2)

In unsupervised learning, the only information accessible to the neural
network is sample $\mathbf{X}$. Therefore, the optimization of neural network
towards optimal information channel can be implemented following

$\displaystyle\max_{\phi}\mathcal{L}_{u}\left(\phi\right):=\underbrace{\mathcal{I}\left(\phi\left(\mathbf{X}\right);\mathbf{X}\right)}_{\text{Encoding}}\underbrace{-\tau\mathcal{I}\left(\phi\left(\mathbf{X}\right);\mathbf{A}\right)}_{\text{Compression}},$
(69)

where $\mathbf{A}\subset\mathbb{N}^{+}$ denotes the index set of samples.
Parameter $\tau\in\left(0,\infty\right)$ denotes a Lagrange multiplier. The
objective in Eq. (69) requires the neural network to maximize the encoded
information in its representation,
$\mathcal{I}\left(\phi\left(\mathbf{X}\right);\mathbf{X}\right)$, during
encoding and reduce the dependence of neural network representation on sample
index, $\mathcal{I}\left(\phi\left(\mathbf{X}\right);\mathbf{A}\right)$,
during compression. This optimization enables the neural network to learn
$\mathbf{Y}=\gamma\left(\mathbf{X}\right)$ under the assumption that sample
distribution matches target distribution (e.g., mapping $\gamma$ is invertible
such that $\mathcal{I}\left(\phi\left(\mathbf{X}\right);\mathbf{Y}\right)$ can
be indirectly optimized through maximizing
$\mathcal{I}\left(\phi\left(\mathbf{X}\right);\mathbf{X}\right)$). An
intrinsic difference between Eq. (69) and classic information bottleneck [35,
36, 37] lies that there is no strict Markov chain among $\mathbf{A}$,
$\mathbf{X}$, $\phi\left(\mathbf{X}\right)$, and $\mathbf{Y}$ because the
definition of index set $\mathbf{A}$ is rather flexible in practice. During
initialization, we suggest to maximize
$\mathcal{I}\left(\phi\left(\mathbf{X}\right);\mathbf{X}\right)$ by
initializing the neural network at dynamic isometry. During encoding, we
suggest to continue to maximize
$\mathcal{I}\left(\phi\left(\mathbf{X}\right);\mathbf{X}\right)$ and apply
random data shuffling, a standard trick in real training processes [54, 55],
to make the neural network learn samples rather than over-fit sample index.
During compression, we suggest to minimize
$\mathcal{I}\left(\phi\left(\mathbf{X}\right);\mathbf{A}\right)$ and ensure
$\Delta|\mathcal{I}\left(\phi\left(\mathbf{X}\right);\mathbf{X}\right)|<\kappa$
for $\kappa\rightarrow 0$.

In sum, although neural network initialization can not completely determine
the performance in subsequent learning tasks, maximizing
$\mathcal{I}\left(\phi\left(\mathbf{X}\right);\mathbf{X}\right)$ based on
initialization at dynamic isometry is beneficial for neural network
optimization during the encoding phase. The above analysis may be closely
related to the empirically observed benefits of dynamic isometry for neural
network training (e.g., for convolutional neural networks [51] and recurrent
neural networks [66]).

##  VII Conclusion

In this research, we have explored a frequently neglected possibility that
neural networks can be initialized toward optimal information channels in deep
learning.

Compared with some prior studies on related topics (e.g., studies on mean-
field dynamics [56, 57, 58, 47, 46, 49, 50, 51] and mutual information
maximization [62] in neural networks), our work may contribute to both physics
and deep learning in the following aspects (see a summary in Fig. 2). First,
we present a unified framework to summarize existing works concerning the
classic mean-filed approximation of information propagation in neural
networks. Second, we indicate the limitation of classic mean-field
approximation in characterizing neural networks as information channels (i.e.,
the implied unreasonable distribution of the correlation measured between
inputs and propagated signals). Third, we propose a corrected mean-field
approximation of infinite-width neural networks to overcome the limitation of
the classic one. Based on the proposed approximation framework and the
mechanism underlying Gaussian information bottleneck, we analytically prove
that neural networks can realize mutual information maximization between
inputs and outputs when they are initialized at dynamic isometry, a case where
neural networks serve as norm-preserving random mappings during information
propagation. Although initially proposed for infinite-width neural networks,
our theory is successfully validated on real finite-width neural networks.
Fourth, we have explored an in-depth analysis on the relation between the
mutual information maximization emerged at dynamic isometry and driving neural
networks towards optimal channels in deep learning tasks.

As a preliminary research, there are diverse intriguing details in our work
remained for future exploration (e.g., the accurate quantification of finite-
size effects on our theory). The suggested connection between statistical
physics and deep learning may be considered as a starting point for more
comprehensive interdisciplinary studies bridging between these two fields.

## Acknowledgements

Authors appreciate Wenqing Wei, who studies at the School of science and
engineer, Chinese University of Hong kong (Shenzhen), for the inspiration on
the proof of that
$\mathcal{I}_{\mathcal{N}}\left(\mathbf{X};\mathbf{X}^{\left(l\right)}\right)$
monotonically increases with $\beta(\sigma_{w},\sigma_{b})$. Sirui Huang, who
studies at the Qiuzhen College, Tsinghua University, is acknowledged for
applying the dominated convergence theorem to derive Eq. (82). Congwen Zhang,
who studies at the ZhiLi College, Tsinghua University, is appreciated for
proof reading.

This project is supported by the Artificial and General Intelligence Research
Program of Guo Qiang Research Institute at Tsinghua University (2020GQG1017)
as well as the Tsinghua University Initiative Scientific Research Program.

##  Appendix A Calculation of dynamic isometry

In this section, we present our proof of Eqs. (34-35). One can also see
similar derivations in Ref. [62].

Let us consider $f(\mathbf{Z})$, the probability density function of an
arbitrary variable, and $g(\mathbf{Z})$, the probability density function of a
multivariate Gaussian variable

$\displaystyle g\left(\mathbf{Z}\right)=$
$\displaystyle\frac{1}{\sqrt{(2\pi)^{n}\operatorname{det}\left(\mathbf{\Sigma}\right)}}$
$\displaystyle\times\exp\left[-\frac{1}{2}\left(\mathbf{Z}-\mathbb{E}\left(\mathbf{Z}\right)\right)^{T}\mathbf{\Sigma}^{(-1)}\left(\mathbf{Z}-\mathbb{E}\left(\mathbf{Z}\right)\right)\right],$
(70)

where $\mathbf{\Sigma}$ and $\mathbb{E}\left(\mathbf{Z}\right)$ denote the
variance matrix and the mean vector shared by $f(\mathbf{Z})$ and
$g(\mathbf{Z})$. Then, we can derive

$\displaystyle\int
f(\mathbf{Z})\log\left[g(\mathbf{Z})\right]\mathrm{d}\mathbf{Z}$
$\displaystyle=$
$\displaystyle-\frac{1}{2}\log\left[(2\pi)^{n}\operatorname{det}\left(\mathbf{\Sigma}\right)\right]-\frac{1}{2}$
$\displaystyle\times\int
f(\mathbf{Z})\log\left[(\mathbf{Z}-\mathbb{E}\left(\mathbf{Z}\right))^{T}\mathbf{\Sigma}^{-1}(\mathbf{Z}-\mathbb{E}\left(\mathbf{Z}\right))\right]\mathrm{d}\mathbf{Z},$
(71) $\displaystyle=$
$\displaystyle-\frac{1}{2}\log\left[(2\pi)^{n}\operatorname{det}\left(\mathbf{\Sigma}\right)\right]-\frac{\operatorname{tr}\left(\mathbf{\Sigma}^{-1}\mathbf{\Sigma}\right)}{2},$
(72) $\displaystyle=$ $\displaystyle\int
g(\mathbf{Z})\log\left[g(\mathbf{X})\right]\mathrm{d}\mathbf{Z}.$ (73)

After replacing $g\left(\mathbf{Z}\right)$ and $f\left(\mathbf{Z}\right)$ by
$g\left(\mathbf{Z};\mathbf{Z}^{\left(l\right)}\right)$ and
$f\left(\mathbf{Z};\mathbf{Z}^{\left(l\right)}\right)$, constraining the
marginal distribution as
$f_{\mathbf{Z}}(\mathbf{Z})=g_{\mathbf{Z}}(\mathbf{Z})$, and using
$f\left(\mathbf{Z};\mathbf{Z}^{\left(l\right)}\right)$ as the joint
distribution of $\left(\mathbf{Z};\mathbf{Z}^{\left(l\right)}\right)$, we can
obtain

$\displaystyle\mathcal{I}\left(\mathbf{Z};\mathbf{Z}^{\left(l\right)}\right)-\mathcal{I}_{\mathcal{N}}\left(\mathbf{Z};\mathbf{Z}^{\left(l\right)}\right)$
$\displaystyle=$ $\displaystyle\int
f\left(\mathbf{Z};\mathbf{Z}^{\left(l\right)}\right)\log\left[\frac{f\left(\mathbf{Z};\mathbf{Z}^{\left(l\right)}\right)}{f_{\mathbf{Z}}\left(\mathbf{Z}\right)f_{\mathbf{Z}^{\left(l\right)}}\left(\mathbf{Z}^{\left(l\right)}\right)}\right]\mathrm{d}\mathbf{Z}\mathrm{d}\mathbf{Z}^{\left(l\right)}$
$\displaystyle-\int
g\left(\mathbf{Z};\mathbf{Z}^{\left(l\right)}\right)\log\left[\frac{g\left(\mathbf{Z};\mathbf{Z}^{\left(l\right)}\right)}{g_{\mathbf{Z}}\left(\mathbf{Z}\right)g_{\mathbf{Z}^{\left(l\right)}}\left(\mathbf{Z}^{\left(l\right)}\right)}\right]\mathrm{d}\mathbf{Z}\mathrm{d}\mathbf{Z}^{\left(l\right)},$
(74) $\displaystyle=$ $\displaystyle-\int
f_{\mathbf{Z}^{\left(l\right)}}\left(\mathbf{Z}^{\left(l\right)}\right)\log\left[f_{\mathbf{Z}^{\left(l\right)}}\left(\mathbf{Z}^{\left(l\right)}\right)\right]\mathrm{d}\mathbf{Z}^{\left(l\right)}$
$\displaystyle+\int
g_{\mathbf{Z}^{\left(l\right)}}\left(\mathbf{Z}^{\left(l\right)}\right)\log\left[g_{\mathbf{Z}^{\left(l\right)}}\left(\mathbf{Z}^{\left(l\right)}\right)\right]\mathrm{d}\mathbf{Z}^{\left(l\right)}$
$\displaystyle+\int
f\left(\mathbf{Z};\mathbf{Z}^{\left(l\right)}\right)\log\left[f\left(\mathbf{Z};\mathbf{Z}^{\left(l\right)}\right)\right]\mathrm{d}\mathbf{Z}\mathrm{d}\mathbf{Z}^{\left(l\right)}$
$\displaystyle-\int
g\left(\mathbf{Z};\mathbf{Z}^{\left(l\right)}\right)\log\left[g\left(\mathbf{Z};\mathbf{Z}^{\left(l\right)}\right)\right]\mathrm{d}\mathbf{Z}\mathrm{d}\mathbf{Z}^{\left(l\right)},$
(75) $\displaystyle=$ $\displaystyle\int
f_{\mathbf{Z}^{\left(l\right)}}\left(\mathbf{Z}^{\left(l\right)}\right)\log\left[\frac{g_{\mathbf{Z}^{\left(l\right)}}\left(\mathbf{Z}^{\left(l\right)}\right)}{f_{\mathbf{Z}^{\left(l\right)}}\left(\mathbf{Z}^{\left(l\right)}\right)}\right]\mathrm{d}\mathbf{Z}^{\left(l\right)}$
$\displaystyle-\int
f\left(\mathbf{Z};\mathbf{Z}^{\left(l\right)}\right)\log\left[\frac{g\left(\mathbf{Z};\mathbf{Z}^{\left(l\right)}\right)}{f\left(\mathbf{Z};\mathbf{Z}^{\left(l\right)}\right)}\right]\mathrm{d}\mathbf{Z}\mathrm{d}\mathbf{Z}^{\left(l\right)},$
(76) $\displaystyle=$ $\displaystyle\int
f\left(\mathbf{Z};\mathbf{Z}^{\left(l\right)}\right)\log\left[\frac{g_{\mathbf{Z}^{\left(l\right)}}\left(\mathbf{Z}^{\left(l\right)}\right)f\left(\mathbf{Z};\mathbf{Z}^{\left(l\right)}\right)}{f_{\mathbf{Z}^{\left(l\right)}}\left(\mathbf{Z}^{\left(l\right)}\right)g\left(\mathbf{Z};\mathbf{Z}^{\left(l\right)}\right)}\right]\mathrm{d}\mathbf{Z}\mathrm{d}\mathbf{Z}^{\left(l\right)},$
(77) $\displaystyle=$ $\displaystyle\int
f_{\mathbf{Z}^{\left(l\right)}}\left(\mathbf{Z}^{\left(l\right)}\right)D_{KL}\left(f_{\mathbf{Z}|\mathbf{Z}^{\left(l\right)}}||g_{\mathbf{Z}|\mathbf{Z}^{\left(l\right)}}\right)\mathrm{d}\mathbf{Z}^{\left(l\right)},$
(78) $\displaystyle\geq$ $\displaystyle 0,$ (79)

where $f_{\mathbf{Z}}(\mathbf{Z})=g_{\mathbf{Z}}(\mathbf{Z})$ is used in Eq.
(75) and Eq. (73) is used to derive Eq.(76). Based on Eq.(79), Eqs. (34-35) in
the main text can be proven.

##  Appendix B Necessary derivations of Eq. (57)

In this section, we present our derivations of Eq. (57). Let us reformulate
the right side of Eq. (57) as

$\displaystyle\int_{\mathbb{R}}\frac{\psi\left(\sigma^{*}z\right)}{\sigma^{*}}z\mathsf{D}z=\int_{\mathbb{R}}\frac{\psi\left(\sigma^{*}z\right)-\psi\left(0\right)}{\sigma^{*}z}z^{2}\mathsf{D}z,$
(80)

where we have used the fact that $\psi\left(0\right)=0$. Because
$\psi\left(\cdot\right)$ is a odd function that is convex in $[0,+\infty]$, we
can know $|\psi^{\prime}\left(0\right)|>|\psi^{\prime}\left(z\right)|$ for any
$z\neq 0$. Then, we have

$\displaystyle\lim_{\sigma^{*}\rightarrow
0}\frac{\psi\left(\sigma^{*}z\right)-\psi\left(0\right)}{\sigma^{*}z}z^{2}\mathsf{D}z=\psi^{\prime}\left(0\right).$
(81)

Based on the dominated convergence theorem [67], we can obtain

$\displaystyle\lim_{\sigma^{*}\rightarrow
0}\int_{\mathbb{R}}\frac{\psi\left(\sigma^{*}z\right)-\psi\left(0\right)}{\sigma^{*}z}z^{2}\mathsf{D}z=\psi^{\prime}\left(0\right)\int_{\mathbb{R}}z^{2}\mathsf{D}z=\psi^{\prime}\left(0\right)$
(82)

based on Eq. (81), which finishes our derivations on Eq. (57) in the main
text.

## References

  * Bengio  _et al._ [2013] Y. Bengio, A. Courville, and P. Vincent, Representation learning: A review and new perspectives, IEEE transactions on pattern analysis and machine intelligence 35, 1798 (2013).
  * Lesort  _et al._ [2018] T. Lesort, N. Díaz-Rodríguez, J.-F. Goudou, and D. Filliat, State representation learning for control: An overview, Neural Networks 108, 379 (2018).
  * Hamilton  _et al._ [2017] W. L. Hamilton, R. Ying, and J. Leskovec, Representation learning on graphs: Methods and applications, arXiv preprint arXiv:1709.05584 (2017).
  * Schölkopf  _et al._ [2021] B. Schölkopf, F. Locatello, S. Bauer, N. R. Ke, N. Kalchbrenner, A. Goyal, and Y. Bengio, Toward causal representation learning, Proceedings of the IEEE 109, 612 (2021).
  * Dike  _et al._ [2018] H. U. Dike, Y. Zhou, K. K. Deveerasetty, and Q. Wu, Unsupervised learning based on artificial neural network: A review, in  _2018 IEEE International Conference on Cyborg and Bionic Systems (CBS)_ (IEEE, 2018) pp. 322–327.
  * Xu and Wunsch [2005] R. Xu and D. Wunsch, Survey of clustering algorithms, IEEE Transactions on neural networks 16, 645 (2005).
  * Xu and Tian [2015] D. Xu and Y. Tian, A comprehensive survey of clustering algorithms, Annals of Data Science 2, 165 (2015).
  * Radford  _et al._ [2015] A. Radford, L. Metz, and S. Chintala, Unsupervised representation learning with deep convolutional generative adversarial networks, arXiv preprint arXiv:1511.06434 (2015).
  * Arora  _et al._ [2019] S. Arora, H. Khandeparkar, M. Khodak, O. Plevrakis, and N. Saunshi, A theoretical analysis of contrastive unsupervised representation learning, arXiv preprint arXiv:1902.09229 (2019).
  * Kotsiantis  _et al._ [2007] S. B. Kotsiantis, I. Zaharakis, P. Pintelas, _et al._ , Supervised machine learning: A review of classification techniques, Emerging artificial intelligence applications in computer engineering 160, 3 (2007).
  * Hastie  _et al._ [2009] T. Hastie, R. Tibshirani, J. H. Friedman, and J. H. Friedman,  _The elements of statistical learning: data mining, inference, and prediction_ , Vol. 2 (Springer, 2009).
  * Deisenroth  _et al._ [2020] M. P. Deisenroth, A. A. Faisal, and C. S. Ong,  _Mathematics for machine learning_ (Cambridge University Press, 2020).
  * Song  _et al._ [2022] H. Song, M. Kim, D. Park, Y. Shin, and J.-G. Lee, Learning from noisy labels with deep neural networks: A survey, IEEE Transactions on Neural Networks and Learning Systems (2022).
  * Frénay and Verleysen [2013] B. Frénay and M. Verleysen, Classification in the presence of label noise: a survey, IEEE transactions on neural networks and learning systems 25, 845 (2013).
  * Cover [1999] T. M. Cover,  _Elements of information theory_ (John Wiley & Sons, 1999).
  * Shamir  _et al._ [2010] O. Shamir, S. Sabato, and N. Tishby, Learning and generalization with the information bottleneck, Theoretical Computer Science 411, 2696 (2010).
  * Kang and Ulukus [2010] W. Kang and S. Ulukus, A new data processing inequality and its applications in distributed source and channel coding, IEEE Transactions on Information Theory 57, 56 (2010).
  * Zhou  _et al._ [2021] C. Zhou, Q. Zhuang, M. Mattina, and P. N. Whatmough, Strong data processing inequality in neural networks with noisy neurons and its implications, in  _2021 IEEE International Symposium on Information Theory (ISIT)_ (IEEE, 2021) pp. 1170–1175.
  * Oord  _et al._ [2018] A. v. d. Oord, Y. Li, and O. Vinyals, Representation learning with contrastive predictive coding, arXiv preprint arXiv:1807.03748 (2018).
  * Henaff [2020] O. Henaff, Data-efficient image recognition with contrastive predictive coding, in  _International conference on machine learning_ (PMLR, 2020) pp. 4182–4192.
  * Tian  _et al._ [2020] Y. Tian, D. Krishnan, and P. Isola, Contrastive multiview coding, in  _European conference on computer vision_ (Springer, 2020) pp. 776–794.
  * Hjelm  _et al._ [2018] R. D. Hjelm, A. Fedorov, S. Lavoie-Marchildon, K. Grewal, P. Bachman, A. Trischler, and Y. Bengio, Learning deep representations by mutual information estimation and maximization, arXiv preprint arXiv:1808.06670 (2018).
  * Poole  _et al._ [2019] B. Poole, S. Ozair, A. Van Den Oord, A. Alemi, and G. Tucker, On variational bounds of mutual information, in  _International Conference on Machine Learning_ (PMLR, 2019) pp. 5171–5180.
  * McAllester and Stratos [2020] D. McAllester and K. Stratos, Formal limitations on the measurement of mutual information, in  _International Conference on Artificial Intelligence and Statistics_ (PMLR, 2020) pp. 875–884.
  * Tschannen  _et al._ [2019] M. Tschannen, J. Djolonga, P. K. Rubenstein, S. Gelly, and M. Lucic, On mutual information maximization for representation learning, arXiv preprint arXiv:1907.13625 (2019).
  * Becker and Hinton [1992] S. Becker and G. E. Hinton, Self-organizing neural network that discovers surfaces in random-dot stereograms, Nature 355, 161 (1992).
  * Linsker [1988] R. Linsker, Self-organization in a perceptual network, Computer 21, 105 (1988).
  * Ge [2018] W. Ge, Deep metric learning with hierarchical triplet loss, in  _Proceedings of the European Conference on Computer Vision (ECCV)_ (2018) pp. 269–285.
  * Yu and Tao [2019] B. Yu and D. Tao, Deep metric learning with tuplet margin loss, in  _Proceedings of the IEEE/CVF International Conference on Computer Vision_ (2019) pp. 6490–6499.
  * Sohn [2016] K. Sohn, Improved deep metric learning with multi-class n-pair loss objective, Advances in neural information processing systems 29 (2016).
  * Li and Sompolinsky [2021] Q. Li and H. Sompolinsky, Statistical mechanics of deep linear neural networks: The backpropagating kernel renormalization, Physical Review X 11, 031059 (2021).
  * Bahri  _et al._ [2020] Y. Bahri, J. Kadmon, J. Pennington, S. S. Schoenholz, J. Sohl-Dickstein, and S. Ganguli, Statistical mechanics of deep learning, Annual Review of Condensed Matter Physics 11 (2020).
  * Jacot  _et al._ [2018] A. Jacot, F. Gabriel, and C. Hongler, Neural tangent kernel: Convergence and generalization in neural networks, Advances in neural information processing systems 31 (2018).
  * Golikov  _et al._ [2022] E. Golikov, E. Pokonechnyy, and V. Korviakov, Neural tangent kernel: A survey, arXiv preprint arXiv:2208.13614 (2022).
  * Tishby  _et al._ [2000] N. Tishby, F. C. Pereira, and W. Bialek, The information bottleneck method, arXiv preprint physics/0004057 (2000).
  * Alemi  _et al._ [2016] A. A. Alemi, I. Fischer, J. V. Dillon, and K. Murphy, Deep variational information bottleneck, arXiv preprint arXiv:1612.00410 (2016).
  * Higgins  _et al._ [2017] I. Higgins, L. Matthey, A. Pal, C. Burgess, X. Glorot, M. Botvinick, S. Mohamed, and A. Lerchner, beta-VAE: Learning basic visual concepts with a constrained variational framework, in  _International Conference on Learning Representations_ (2017).
  * Lee  _et al._ [2017] J. Lee, Y. Bahri, R. Novak, S. S. Schoenholz, J. Pennington, and J. Sohl-Dickstein, Deep neural networks as gaussian processes, arXiv preprint arXiv:1711.00165 (2017).
  * Matthews  _et al._ [2018] A. G. d. G. Matthews, M. Rowland, J. Hron, R. E. Turner, and Z. Ghahramani, Gaussian process behaviour in wide deep neural networks, arXiv preprint arXiv:1804.11271 (2018).
  * Garriga-Alonso  _et al._ [2018] A. Garriga-Alonso, C. E. Rasmussen, and L. Aitchison, Deep convolutional networks as shallow gaussian processes, arXiv preprint arXiv:1808.05587 (2018).
  * Lee  _et al._ [2019] J. Lee, L. Xiao, S. Schoenholz, Y. Bahri, R. Novak, J. Sohl-Dickstein, and J. Pennington, Wide neural networks of any depth evolve as linear models under gradient descent, Advances in neural information processing systems 32 (2019).
  * Chizat  _et al._ [2019] L. Chizat, E. Oyallon, and F. Bach, On lazy training in differentiable programming, Advances in Neural Information Processing Systems 32 (2019).
  * Novak  _et al._ [2018] R. Novak, L. Xiao, J. Lee, Y. Bahri, G. Yang, J. Hron, D. A. Abolafia, J. Pennington, and J. Sohl-Dickstein, Bayesian deep convolutional networks with many channels are gaussian processes, arXiv preprint arXiv:1810.05148 (2018).
  * Yang [2019] G. Yang, Wide feedforward or recurrent neural networks of any architecture are gaussian processes, Advances in Neural Information Processing Systems 32 (2019).
  * Sirignano and Spiliopoulos [2022] J. Sirignano and K. Spiliopoulos, Mean field analysis of deep neural networks, Mathematics of Operations Research 47, 120 (2022).
  * Schoenholz  _et al._ [2016] S. S. Schoenholz, J. Gilmer, S. Ganguli, and J. Sohl-Dickstein, Deep information propagation, arXiv preprint arXiv:1611.01232 (2016).
  * Pennington  _et al._ [2017] J. Pennington, S. Schoenholz, and S. Ganguli, Resurrecting the sigmoid in deep learning through dynamical isometry: theory and practice, Advances in neural information processing systems 30 (2017).
  * Yang and Schoenholz [2017] G. Yang and S. Schoenholz, Mean field residual networks: On the edge of chaos, Advances in neural information processing systems 30 (2017).
  * Pennington  _et al._ [2018] J. Pennington, S. Schoenholz, and S. Ganguli, The emergence of spectral universality in deep networks, in  _International Conference on Artificial Intelligence and Statistics_ (PMLR, 2018) pp. 1924–1932.
  * Chen  _et al._ [2018] M. Chen, J. Pennington, and S. Schoenholz, Dynamical isometry and a mean field theory of rnns: Gating enables signal propagation in recurrent neural networks, in  _International Conference on Machine Learning_ (PMLR, 2018) pp. 873–882.
  * Xiao  _et al._ [2018] L. Xiao, Y. Bahri, J. Sohl-Dickstein, S. Schoenholz, and J. Pennington, Dynamical isometry and a mean field theory of cnns: How to train 10,000-layer vanilla convolutional neural networks, in  _International Conference on Machine Learning_ (PMLR, 2018) pp. 5393–5402.
  * Berger [2003] T. Berger, Rate-distortion theory, Wiley Encyclopedia of Telecommunications (2003).
  * Kleven [2021] H. J. Kleven, Sufficient statistics revisited, Annual Review of Economics 13, 515 (2021).
  * Nguyen  _et al._ [2022] T. T. Nguyen, F. Trahay, J. Domke, A. Drozd, E. Vatai, J. Liao, M. Wahib, and B. Gerofi, Why globally re-shuffle? revisiting data shuffling in large scale deep learning, in  _2022 IEEE International Parallel and Distributed Processing Symposium (IPDPS)_ (IEEE, 2022) pp. 1085–1096.
  * Summers and Dinneen [2021] C. Summers and M. J. Dinneen, Nondeterminism and instability in neural network optimization, in  _International Conference on Machine Learning_ (PMLR, 2021) pp. 9913–9922.
  * Mei  _et al._ [2019] S. Mei, T. Misiakiewicz, and A. Montanari, Mean-field theory of two-layers neural networks: dimension-free bounds and kernel limit, in  _Conference on Learning Theory_ (PMLR, 2019) pp. 2388–2464.
  * Nguyen [2019] P.-M. Nguyen, Mean field limit of the learning dynamics of multilayer neural networks, arXiv preprint arXiv:1902.02880 (2019).
  * Poole  _et al._ [2016] B. Poole, S. Lahiri, M. Raghu, J. Sohl-Dickstein, and S. Ganguli, Exponential expressivity in deep neural networks through transient chaos, Advances in neural information processing systems 29 (2016).
  * Mingo and Speicher [2017] J. A. Mingo and R. Speicher,  _Free probability and random matrices_ , Vol. 35 (Springer, 2017).
  * Painsky and Tishby [2017] A. Painsky and N. Tishby, Gaussian lower bound for the information bottleneck limit., J. Mach. Learn. Res. 18, 213 (2017).
  * Chechik  _et al._ [2003] G. Chechik, A. Globerson, N. Tishby, and Y. Weiss, Information bottleneck for gaussian variables, Advances in Neural Information Processing Systems 16 (2003).
  * Ughi [2022] G. Ughi,  _Studies on neural networks: Information propagation at initialisation and robustness to adversarial examples_ , Ph.D. thesis, University of Oxford (2022).
  * Saxe  _et al._ [2013] A. M. Saxe, J. L. McClelland, and S. Ganguli, Exact solutions to the nonlinear dynamics of learning in deep linear neural networks, arXiv preprint arXiv:1312.6120 (2013).
  * Speicher [1994] R. Speicher, Multiplicative functions on the lattice of non-crossing partitions and free convolution, Mathematische Annalen 298, 611 (1994).
  * Voiculescu  _et al._ [1992] D. V. Voiculescu, K. J. Dykema, and A. Nica,  _Free random variables_ , 1 (American Mathematical, 1992).
  * Gilboa  _et al._ [2019] D. Gilboa, B. Chang, M. Chen, G. Yang, S. S. Schoenholz, E. H. Chi, and J. Pennington, Dynamical isometry and a mean field theory of lstms and grus, arXiv preprint arXiv:1901.08987 (2019).
  * Stein and Shakarchi [2009] E. M. Stein and R. Shakarchi,  _Real analysis: measure theory, integration, and Hilbert spaces_ (Princeton University Press, 2009).