# On the impact of activation and normalization in obtaining isometric
embeddings at initialization

Amir Joudaki
BMI, ETH Zürich
<EMAIL_ADDRESS>
&Hadi Daneshmand
MIT LIDS
<EMAIL_ADDRESS>
&Francis Bach
INRIA-ENS-PSL Paris
<EMAIL_ADDRESS>

###### Abstract

In this paper, we explore the structure of the penultimate Gram matrix in deep
neural networks, which contains the pairwise inner products of outputs
corresponding to a batch of inputs. In several architectures it has been
observed that this Gram matrix becomes degenerate with depth at
initialization, which dramatically slows training. Normalization layers, such
as batch or layer normalization, play a pivotal role in preventing the rank
collapse issue. Despite promising advances, the existing theoretical results
(i) do not extend to layer normalization, which is widely used in
transformers, (ii) can not characterize the bias of normalization
quantitatively at finite depth.

To bridge this gap, we provide a proof that layer normalization, in
conjunction with activation layers, biases the Gram matrix of a multilayer
perceptron towards isometry at an exponential rate with depth at
initialization. We quantify this rate using the Hermite expansion of the
activation function, highlighting the importance of higher order ($\geq 2$)
Hermite coefficients in the bias towards isometry.

##  1 Introduction

The Gram matrix of the penultimate layer in deep neural networks (DNNs) is a
key object of study in recent theoretical works on neural network dynamics [1;
2; 3; 4; 5]. This matrix measures the similarity between the outputs for a
batch of inputs. The Gram matrix can become degenerate as the network depth
increases at initialization for various neural architectures such as
multilayer perceptrons (MLPs) [6; 7], convolutional networks [8], and
transformers [9]. The degenerate Gram matrix can slow down the training
process and impact generalization for DNNs [10; 2; 11]. To prevent this
problem, several techniques have been proposed to modify the neural
architecture, such as residual connections [9], shaping activations [12; 13],
and normalization layers [1; 7].

Mean field theory has been extensively used to characterize the effect of
neural architectures on the Gram matrix in the limit of infinitely deep and
wide neural networks. In this setting, the Gram matrix can be expressed as a
fixed point of a recurrence equation that depends on the neural architecture
[14; 1; 2]. A fixed-point analysis can provide insights about the the
structure of Gram matrices in deep neural networks; thereby shedding light on
the degeneracy of Gram matrices in networks without normalization [14], and
how to avoid this degeneracy with batch normalization [1].

One of the main difficulties in studying infinitely deep networks is that the
limiting Gram matrix is determined by a difference equation that has multiple
fixed-points which can be degenerate or non-degenerate [1]. The network
structure determines the stability of these fixed points in their local
neighbourhood. For neural networks with batch normalization, [1] prove the
non-degenerate fixed points are unstable. However, this stability only holds
in a small neighborhood, and does not guarantee that the limiting Gram matrix
will converge to a non-degenerate fixed-point for random inputs.

To go beyond a local analysis, Yang et al. [1] prove the global convergence of
the Gram matrix for non-degenerate inputs and networks with linear activation
and batch normalization layers. They further show empirically that for a much
wider set of activations and a wide class of inputs, only well-conditioned
Gram matrices are reached. Based on these observations, Yang et al. [1]
postulate that such global convergence holds for non-linear activations under
mild assumptions on the input distribution.

In the present work, we prove the global convergence of Gram matrices towards
isometry for networks with non-linear activations and layer normalization [15]
in the regime of infinite width. We provide quantitative rates for this
convergence towards isometric Gram matrices in terms of the Hermite expansion
of the activation. Our analysis reveals the striking isometry bias of
normalization layers, including layer (and batch) normalization, that holds
even for networks with a finite width independent from other components.

##  2 Preliminaries

#### Notation.

Let $\langle x,y\rangle$ be the inner product of vectors $x$ and $y,$ and
$\|x\|^{2}=\langle x,x\rangle$ the squared Euclidean norm of $x$. For a matrix
$X,$ we write $X_{i\cdot}$ and $X_{\cdot i}$ for the $i$-th row and column of
$X,$ respectively. We use ${W\sim N(\mu,\sigma^{2})^{m\times n}}$ to indicate
that $W$ is an $m\times n$ Gaussian matrix with i.i.d. elements from
$N(\mu,\sigma^{2}).$ We denote by $\bm{0}_{n}$ the zero vector of size $n.$
Given vector $x\in\mathbb{R}^{n},$ $\overline{x}$ denotes the arithmetic mean
of $\frac{1}{n}\sum_{i=1}^{n}x_{i}.$ Lastly, $I_{n}$ is the identity matrix of
size $n.$

#### Normalization layers.

Normalization layers are widely used in deep neural networks (DNNs) to improve
their training stability and speed. Among them, batch normalization [16] and
layer normalization [15] are two prominent examples that have been
successfully used to train deep models [17; 18]. Batch normalization ensures
each feature within a layer across a mini-batch has zero mean and unit
variance, while layer normalization centers and divides the output of each
layer by its standard deviation. Inspired by these two main types of
normalization, we will consider layer normalization
$\text{LN}\mathrel{\mathop{\ordinarycolon}}\mathbb{R}^{d}\to\mathbb{R}^{d}$
and batch normalization
$\text{BN}\mathrel{\mathop{\ordinarycolon}}\mathbb{R}^{d\times
n}\to\mathbb{R}^{d\times n},$ defined in Table 2.1.

Width | $d\in\mathbb{N}$ | Batch size | $n\in\mathbb{N}$
---|---|---|---
Depth | $L\in\mathbb{N}$ | Input | $x\in\mathbb{R}^{d}$
Input batch | $X\in\mathbb{R}^{d\times n}$ | Gaussian weights | $W^{1},\dots,W^{L}\sim N(0,1)^{d\times d}$
Activation | $\sigma\mathrel{\mathop{\ordinarycolon}}\mathbb{R}\to\mathbb{R}$ | Centering | $x-\overline{x}$
Layer Norm | ${\textsc{LN}}(x)=\frac{x}{\sqrt{\frac{1}{d}\sum_{i}^{d}x_{i}^{2}}}$ | Batch Norm | ${\textsc{BN}}(X)_{ij}=\frac{X_{ij}}{\sqrt{\frac{1}{n}\sum_{k}^{d}X_{ik}^{2}}}$
Isometry | $\mathcal{I}(M)\mathrel{\mathop{\ordinarycolon}}=\frac{\det(M)^{1/n}}{\frac{1}{n}\mathrm{tr}(M)}\in[0,1]$ | Isometry gap | $-\log\mathcal{I}(M)\in[0,\infty]$
Table 2.1:  Building blocks we consider in this work. Isometry is defined for
PSD matrices.

Note that in Table 2.1, we distinguish between the centering and normalization
steps in our definitions for layer (batch) normalization. Yet, we note that we
will not depart from the standard MLP architectures since we include centering
before normalization as ${\textsc{LN}}(x-\overline{x})$ used in the original
definition [15]. This separation helps us to decouple the effect of projection
onto the $\sqrt{d}$-sphere from the centering. While the influence of the
normalization step is crucial in reaching isometry in a deterministic fashion,
we will see that the effects of mean reduction is essential when combined with
non-linear activations and random weights. Similarly, we split batch
normalization into centering and normalization steps in our definitions.

#### MLP setup.

The subject of our analysis is an MLP with constant width $d$ across the
layers and $L$ layers, which takes input $x\in\mathbb{R}^{d}$, and maps it to
output $x^{L}\in\mathbb{R}^{d},$ with hidden representations as

$\displaystyle\begin{cases}x^{\ell+1}=\frac{1}{\sqrt{d}}{\textsc{LN}}(h_{\ell}-\overline{h}_{\ell}),\quad
h_{\ell}=\sigma(W^{\ell}x^{\ell}),&\ell=0,\dots,L-1\\\
x^{0}\mathrel{\mathop{\ordinarycolon}}=\frac{1}{\sqrt{d}}{\textsc{LN}}(x),&\text{input.
}\end{cases}$ (1)

Compared to Ba et al. [15], the order of activation and layer normalization
presented here is different. However, recently it has been proposed that
applying layer normalization before fully connected and then activation, as
prescribed in equation above, improve the gradients and training of
transformer architecture [19].

###  2.1 Gram matrices and isometry

Given $n$ data points $\\{x_{i}\\}_{i\leq n}\in\mathbb{R}^{d}$, the Gram
matrix $G^{\ell}$ of the feature vectors
$x_{1}^{\ell},\dots,x_{n}^{\ell}\in\mathbb{R}^{d}$ at layer $\ell$ of the
network is defined as

$\displaystyle G^{\ell}\mathrel{\mathop{\ordinarycolon}}=\left[\langle
x^{\ell}_{i},x_{j}^{\ell}\rangle\right]_{i,j\leq n},$
$\displaystyle\ell=1,\dots,L.$ (2)

Intuitively, an isometric Gram matrix implies that the network preserves the
distances and angles between the input data points after mapping them to the
feature space. Isometry of Gram matrices can be quantified using the
eigenvalues of $G^{\ell}.$ One possible way to formulate isometry is to use
the ratio of the volume and scale of the parallelepiped spanned by the feature
vectors $x_{1}^{\ell},\dots,x^{\ell}_{n}$. For example, consider two points on
a plane $x_{1},x_{2}\in\mathbb{R}^{2}$ with lengths $a=|x_{1}|,b=|x_{2}|$ and
angle $\theta=\angle(x_{1},x_{2})$. The ratio is given by
$ab\sin(\theta)/(a^{2}+b^{2})$, which is maximized when $a=b$ and
$\theta=\pi/2$. This is shown for $n=2$ and $n=3$ feature vectors in Figure
2.1.

$\frac{\|x_{1}\|\|x_{2}\|\sin\theta}{\frac{1}{2}(\|x_{1}\|^{2}+\|x_{2}\|^{2})}$$\frac{\text{Vol}(\\{x_{i}\\}_{i\leq
n})^{2/n}}{\frac{1}{n}\sum_{i}^{n}\|x_{i}\|^{2}}=\frac{\det(X^{\top}X)^{1/n}}{\frac{1}{n}\mathrm{tr}(X^{\top}X)}$$\vec{x_{1}}$$\vec{x_{2}}$$\theta$$\vec{x_{1}}$$\vec{x_{2}}$$\theta$$\vec{x_{1}}$$\vec{x_{2}}$$\theta$$\vec{x_{1}}$$\vec{x_{2}}$$\theta$$\vec{x_{1}}$$\vec{x_{2}}$$\vec{x_{3}}$$\vec{x_{1}}$$\vec{x_{2}}$$\vec{x_{3}}$$\vec{x_{1}}$$\vec{x_{2}}$$\vec{x_{3}}$$\vec{x_{1}}$$\vec{x_{2}}$$\vec{x_{3}}$
Figure 2.1: A geometric interpretation of isometry: higher volume corresponds
to higher isometry.

Inspired by this intuition, we can define isometry of the Gram matrix.

######  Definition 1.

Let $M$ be an $n\times n$ positive semi-definite matrix. We define the
_isometry $\mathcal{I}(M)$_ of $M$ as the ratio of its normalized determinant
to its normalized trace:

$\displaystyle\mathcal{I}(M)\mathrel{\mathop{\ordinarycolon}}=\frac{\det(M)^{1/n}}{\frac{1}{n}\mathrm{tr}(M)}.$
(3)

The function $\mathcal{I}(M)$ defined in equation (3) quantifies how well $M$
approximates the identity matrix $I_{n}$. We can easily check that
$\mathcal{I}(M)$ has some desirable properties (see Lemma B.1 for formal
statements and proofs):

  * •

Scale-invariance: Multiplying $M$ by a constant does not affect
$\mathcal{I}(M)$.

  * •

Isometry-preserving: $\mathcal{I}(M)$ ranges between $0$ and $1$, with $0$ and
$1$ corresponding to degenerate and identity matrices respectively.

  * •

Isometry gap: $-\log\mathcal{I}(M)$ lies between $0$ and $\infty,$ with $0$
and $\infty$ indicating identity and degenerate matrices respectively.

These properties suggest that $\mathcal{I}(M)$ is a suitable function for
measuring how close a matrix is to being an isometry, i.e., a transformation
that preserves distances between metric spaces. Moreover, there is a clear
link between isometry and normalization, which we will explore in the next
section. We will often measure how far a matrix is from isometry by its
negative log $-\log\mathcal{I}(M),$ which we will call _isometry gap_.

##  3 Isometry bias of normalization

This notion of isometry has a remarkable property: if we normalize each point
by its Euclidean norm then the isometry of their associated Gram matrix does
not decrease. We formalize this property in the following theorem.

######  Theorem 3.1.

Given $n$ samples $\\{x_{i}\\}_{i\leq
n}\subset\mathbb{R}^{d}\setminus\\{\bm{0}_{d}\\}$, and their projection onto
the unit sphere
$\widetilde{x}_{i}\mathrel{\mathop{\ordinarycolon}}=x_{i}/\|x_{i}\|,$ and
their respective Gram matrices $G=\left[\langle
x_{i},x_{j}\rangle\right]_{i,j\leq n}$ and
$\widetilde{G}=\left[\langle\widetilde{x_{i}},\widetilde{x_{j}}\rangle\right]_{i,j\leq
n}.$ The isometry of Gram matrices obeys

$\displaystyle\mathcal{I}\left(\widetilde{G}\right)\geq\mathcal{I}(G)\left(1+\frac{\frac{1}{n}\sum_{i}^{n}(a_{i}-\bar{a})^{2}}{\bar{a}^{2}}\right),$
$\displaystyle\text{where
}a_{i}\mathrel{\mathop{\ordinarycolon}}=\|x_{i}\|,\bar{a}\mathrel{\mathop{\ordinarycolon}}=\frac{1}{n}\sum_{i}^{n}a_{i}.$
(4)

Theorem 3.1 shows shows a subtle property of normalization: Because the terms
$(a_{i}-\overline{a})^{2}$ are always non-negative, the left hand side is
always greater than or equal to $\mathcal{I}(G).$ It further quantifies the
improvement in isometry as a function of fluctuations in norms: the terms
$\overline{a}$ and $\frac{1}{n}\sum_{i}^{n}(a_{i}-\bar{a})^{2}$ correspond to
the sample mean and variance of $a_{1},\dots,a_{n}.$ Thus, the more diverse
the norms, the larger the improvement in isometry.

###  3.1 Implications for normalization layers

Theorem 3.1 allows us to highlight the isometry bias of layer and batch
normalization.

######  Corollary 3.2.

Consider $n$ vectors before and after layer-normalization $\\{x_{i}\\}_{i\leq
n}\subset\mathbb{R}^{d}\setminus\\{\bm{0}_{d}\\}$ and
$\\{\widetilde{x}_{i}\\}_{i\leq
n},\widetilde{x_{i}}\mathrel{\mathop{\ordinarycolon}}={\textsc{LN}}(x_{i}).$
Define their respective Gram matrices
$G\mathrel{\mathop{\ordinarycolon}}=[\langle x_{i},x_{j}\rangle]_{i,j\leq n},$
and
$\widetilde{G}\mathrel{\mathop{\ordinarycolon}}=[\langle\widetilde{x}_{i},\widetilde{x}_{j}\rangle]_{i,j\leq
n}.$ We have:

$\displaystyle\mathcal{I}\left(\widetilde{G}\right)\geq\mathcal{I}(G)\left(1+\frac{\frac{1}{n}\sum_{i}^{n}(a_{i}-\bar{a})^{2}}{\bar{a}^{2}}\right),$
$\displaystyle\text{where
}a_{i}\mathrel{\mathop{\ordinarycolon}}=\|x_{i}\|,\bar{a}\mathrel{\mathop{\ordinarycolon}}=\frac{1}{n}\sum_{i}^{n}a_{i}.$

Observe that we can view the layer normalization step in Table 2.1 as a
projection onto the $\sqrt{d}$-sphere, which is equivalent to the unit-norm
projection in Theorem 3.1 up to a constant scale factor. Since isometry is
scale-invariant, this implies that corollary 3.2 follows directly from Theorem
3.1.

Moreover, corollary 3.2 shows that the isometry of layer normalization is
deterministic and does not rely on the random weights. This means that layer
normalization always preserves or enhances the isometry of representations,
even during training. We provide empirical evidence for this in discussion.

Despite the seemingly vast differences between layer normalization and batch
normalization, the following corollary shows an intimate link between them
through the prism of isometry.

######  Corollary 3.3.

Given $n$ samples in a mini-batch before $X\in\mathbb{R}^{d\times n},$ and
after normalization $\widetilde{X}={\textsc{BN}}(X)$ and define covariance
matrices $C\mathrel{\mathop{\ordinarycolon}}=XX^{\top}$ and
$\widetilde{C}\mathrel{\mathop{\ordinarycolon}}=XX^{\top}.$ We have:

$\displaystyle\mathcal{I}\left(\widetilde{C}\right)\geq\mathcal{I}(C)\left(1+\frac{\frac{1}{d}\sum_{i}^{d}(a_{i}-\bar{a})^{2}}{\bar{a}^{2}}\right),$
$\displaystyle\text{where
}a_{i}\mathrel{\mathop{\ordinarycolon}}=\|X_{i\cdot}\|,\bar{a}\mathrel{\mathop{\ordinarycolon}}=\frac{1}{n}\sum_{i=1}^{d}a_{i}.$

Gram matrices of networks with batch normalization have been subject of many
previous studies at network initialization: it has been postulated that BN
prevents rank collapse issue [7] and that it orthogonalizes the
representations [20], and that it imposes isometry [1]. The isometry results
implied by Corollaries 3.3 give a geometric interpretation of all these
findings: it is straightforward to verify that maximum isometry is achieved
with identity Gram matrix, which implies orthogonalization of these results.
Rather strikingly, while all the results stated before have been of
probabilistic nature, the improved isometry of Gram matrix implied by
Corollary 3.3 holds deterministically.

###  3.2 Isometry across the layers: batch normalization vs. layer
normalization

In this subsection, we investigate how normalization affects the isometry of
the representations. We use the ratio of the variance and the square of the
mean of the norms as a measure of normalization bias, The normalization bias
can be written as:

$\displaystyle\text{normalization
bias}=\frac{\text{var}(\text{norms})}{\text{mean}^{2}(\text{norms})},$ (5)

where for layer normalization, we compute this ratio for different samples
$\|x_{i}^{\ell}\|$, and for batch normalization, we compute it for different
features within a batch. We plot this ratio for MLPs with random weights with
or without layer or batch normalization in Figure 3.1. We observe that batch
normalization has a non-zero normalization bias that persists with depth,
while layer normalization has a variable normalization bias that depends on
the activation function. This implies that layer normalization cannot prevent
the rank collapse of the Gram matrices as depth increases, as seen in Figure
3.2. To understand this phenomenon better, we analyze how activation functions
and layer normalization interact.

Figure 3.1: Isometry bias of normalization layer vs. batch normalization:
quantifying the normalization bias in layer and batch normalization, as
defined in equation (5). Figure 3.2: Isometry bias of normalization: MLP with
ReLU (blue) and identity activation (orange) without normalization (left
column), with layer normalization (middle column), and with batch
normalization (right column).

##  4 Achieving Isometry through Nonlinear Activation Functions

Neural networks typically employ a nonlinear function
$\sigma\mathrel{\mathop{\ordinarycolon}}\mathbb{R}\to\mathbb{R}$, with the
choice of function depending on the network architecture. For instance,
rectified linear units (ReLU) are commonly used in convolutional networks for
image processing, while various activations are utilized in large language
models. The choice of activation can significantly influence the optimization
of neural network architectures [21]. Motivated by this impact, we examine the
interplay between layer normalization and nonlinear activation functions.
However, nonlinear activations can make it more challenging to theoretically
understand the dynamics of representations through layers. One approach to
address this issue involves employing Hermite polynomials, which serve as a
functional basis for the Gaussian kernel. They have been successfully used to
analyze the kernel function associated with neural networks at initialization
[22] and to derive neural tangent kernel derivations of infinitely wide
networks [23]. Leveraging the Hermite expansion of the activation, we
characterize the isometry bias introduced by activation functions.

###  4.1 Hermite expansion of activation functions

Hermite polynomials possess the properties of completeness and orthogonality
under the Gaussian weight kernel. This means that any function in the space of
square-integrable functions with respect to the Gaussian kernel, satisfying
the condition $\int_{-\infty}^{\infty}\sigma(x)^{2}e^{-x^{2}/2}dx<\infty$, can
be expressed as a linear combination of Hermite polynomials.

We use the _normalized_ Hermite polynomials and their coefficients.

######  Definition 2.

Define normalized Hermite polynomials $\mathit{He}_{k}(x)$ as

$\displaystyle\mathit{He}_{k}(x)\mathrel{\mathop{\ordinarycolon}}=(k!)^{-1/2}(-1)^{k}e^{x^{2}/2}\frac{d^{k}}{dx^{k}}e^{-x^{2}/2},$

and for a given scalar function $\sigma$, define its normalized coefficients
${c_{k}}$ as (See section A for more details):

$\displaystyle
c_{k}\mathrel{\mathop{\ordinarycolon}}=\int_{x}\sigma(x)\mathit{He}_{k}(x)e^{-x^{2}/2}dx.$

To quantify the influence of activations on the isometry, we define the notion
of non-linearity strength for activations.

######  Definition 3 (Non-linearity strength).

Given activation $\sigma$ with Hermite expansion $\\{c_{k}\\}_{k\geq 0},$
define its _non-linearity strength $\beta_{0}$_ as:

$\displaystyle\beta_{0}\mathrel{\mathop{\ordinarycolon}}=2-\frac{c_{1}^{2}}{\sum_{k=1}^{\infty}c_{k}^{2}}$
(6)

Table 4.1 represents the non-linearity strength of various activation in
closed form.

| $\mathit{He}_{1}(x)$ | $\mathit{He}_{2}(x)$ | Sine | Exponential | Step | ReLU
---|---|---|---|---|---|---
Formula | $x$ | $x^{2}-1$ | $\sin(x)$ | $\exp(x)$ | $\bm{1}[x>0]$ | $\max(x,0)$
$\beta_{0}$ | 1 | 2 | $2-\frac{2e}{e^{2}-1}$ | $2-\frac{1}{e-1}$ | $2-\frac{2}{\pi}$ | $\frac{3\pi-4}{2\pi-2}$
Table 4.1: Different activation functions and their non-linearity strength
$\beta_{0}$ (equation (6)).

###  4.2 Mean Field Isometry for MLPs

In this section, we analyze how the isometry of input Gram matrix $G^{0}$
changes as it passes through the network layers, following the definitions in
equations (1) and (2). We use the mean field dynamic of Gram matrices, which
has been previously studied [1; 14; 24], to understand the effects of
normalization and nonlinear activations on the isometry. The mean field
dynamics of Gram matrices is given by

$\displaystyle G^{\ell+1}_{*}=\mathbb{E}_{h\sim
N(0,G^{\ell}_{*})}\left[\phi(\sigma(h))\phi(\sigma(h))^{\top}\right],$
$\displaystyle\text{ where
}[\phi(a)]_{i}\mathrel{\mathop{\ordinarycolon}}=(a_{i}-\mathbb{E}a_{i})/\sqrt{\operatorname*{Var}{a_{i}}}.$
(7)

This equation represents the expected Gram matrix of layer $\ell+1,$ given the
Gram matrix of the previous layer is $G^{\ell}_{*},$ with $\phi$ taking the
role of layer normalization operator in the mean field regime (see section A
for more details). The sequence $G^{\ell}_{*}$ approximates the dynamics of
$G^{\ell},$ and this correspondence becomes exact for infinitely wide MLPs. In
the rest of this section, we analyze the above dynamical system, and we will
use the notation subscript $G^{\ell}_{*}$ to make this distinction explicit.

Inspired by the isometry bias of normalization layers, we analyze how the
isometry of $G^{\ell}_{*}$ changes with depth through the lens of isometry.
Interestingly, the negative log of isometry $-\log\mathcal{I}(G^{\ell}_{*}),$
can serve as a Lyapunov function for the above dynamics. The following Theorem
proves the isometry gap may decay with depth at an exponential rate which
depends on the non-linearity strength of the activation.

######  Theorem 4.1.

Let $\sigma$ be an activation function with a Hermite expansion and a non-
linearity strength $\beta_{0},$ (see equation (6)). Given non-degenerate input
Gram matrix $G^{0}_{*},$ for sufficiently deep layer
$\ell\gtrsim\beta_{0}^{-1}(-n\log\mathcal{I}(G_{*}^{0})+\log(4n))$, we have

$\displaystyle-\log\mathcal{I}(G^{\ell}_{*})\leq\exp(-\ell\log\beta_{0}-n\log\mathcal{I}(G^{0}_{*})+\log(4n)).$
(8)

Theorem 4.1 confirms the main claim of this paper: MLPs with normalization
layers and non-linear activations can achieve isometric embedding by
increasing depth at initialization. Thus, this theorem sheds light on the role
of layer normalization in training DNNs, since isometry of Gram matrices
influences the training performance [2; 8; 11].

Remarkably, the condition on $G^{0}_{*}$ being non-degenerate is essential to
establish non-vacuous bounds for isometry. For example, if the input batch
contains a duplicated sample as input, Gram matrices across all layers remain
degenerate. Theorem 4.1 states that as long as the input Gram matrix is non-
degenerate, the embedding becomes isometric with depth at an exponential rate.

The decay rate towards isometry can be quantified in terms of the non-
linearity strength of the activation function, i.e., $\beta_{0},$ which is
computed in closed form for various activations in Table 4.1. Figure 4.1
compares the theoretical rate with those observed in practice for three
activations, showing that the established rates from mean field analysis are
predictive in practice for a wide MLP.

Figure 4.1: Isometry bias of activation depth $\ell$. Solid blue traces show
isometry of MLP with $n=10,$ $d=1000,$ and various activations $\sigma$. Solid
lines show average of $\\#10$ independent runs. The dashed traces are
theoretical upper bounds given in Theorem 4.1.

#### The role of non-linear terms in achieving isometry.

According to Theorem 4.1, the isometry gap depends on the non-linearity
strength $\beta_{0},$ which is determined by the non-linear coefficients
$c_{k}$ of the Hermite expansion of the activation function. To illustrate
this dependence, we compare the isometry gap of LN-MLP with different
activation functions: $\sigma(x)=\mathit{He}_{k}(x)$ for $k\in\\{1,2,3\\}$.
Figure 4.2 shows the decay rate for these Hermite bases. We observe the
higher-order polynomials obtains a faster rate which validate the role of non-
linearity strength in Theorem 4.1. Moreover, the linear activation
$\mathit{He}_{1}(x)=x$ does not lead to the isometry. This suggests that non-
linear terms in the Hermite expansion are crucial for achieving good geometric
properties of MLP and imposing a stronger isometry bias.

Figure 4.2: Isometry with various Hermite polynomials as activation functions
in MLP with batch and layer normalization. Batch size $n=10$, width $d=1000.$

##  5 Discussion

In this study, we explored the influence of layer normalization and nonlinear
activation functions on the isometry of MLP representations. Our findings open
up several avenues for future research.

#### Normalization bias during training.

Interestingly, the normalization bias of Theorem 3.1 is deterministic and
remains valid irrespective of the assumptions in our mean-field analysis.
Figure 5.1 reveals that this bias not only persists but also increases during
training with SGD on the CIFAR10 dataset [25]. Investigating the role of this
isometry bias during and after training presents an intriguing research
direction.

Figure 5.1: Normalization bias, as defined in equation (5) (y-axis) of Gram
matrices of MLP representations versus layer depth (x-axis), during 3 epochs
of SGD with learning rate $0.001$ (each epoch marked in a different color) on
CIFAR10. Model parameters: ReLU activation, batch size $n=256$, width
$d=1000$, depth $L=30$. The solid lines represent the average over five
independent runs.

#### Beyond normalization.

Our mean-field theory reveals a significant insight: when applied post-
activation, layer normalization inherently biases and scales the layers
towards isometry. Specifically, our results establish the necessary conditions
for isometry with layer normalization. However, it is worth investigating
whether similar outcomes can be achieved without layer normalization. Our
empirical observations suggest that certain activations, such as ReLU, require
layer normalization to attain isometry. In contrast, other activations, which
can be considered as “self-normalizing” (e.g., SELU [21] and hyperbolic
tangent), can achieve isometry with only offset and scale adjustments (see
Figure 5.2).

Figure 5.2: Exploring the potential of achieving isometry with self-normalized
activations. Batch size $n=10$, width $d=1000.$

#### Ordering of normalization and activation layers.

Theorem 4.1 highlights that the ordering of activation and normalization
layers has a critical impact on the isometry. Figure 5.3 demonstrates that a
different ordering can lead to a non-isotropic Gram matrix. Remarkably, the
structure analyzed in this paper is used in transformers [26].

Figure 5.3: Impact of the order of activation and normalization layers. Batch
size $n=10$, width $d=1000.$

#### Conclusion.

In conclusion, our study highlights the importance of layer normalization and
nonlinear activation functions in the isometry of representations in deep
random neural networks. Our findings not only contribute to a better
understanding of the interplay between various components of neural networks
but also suggest potential avenues for future research, such as investigating
the role of isometry bias during training and exploring the possibility of
self-normalizing activations.

## References

  * Yang et al. [2019] Greg Yang, Jeffrey Pennington, Vinay Rao, Jascha Sohl-Dickstein, and Samuel S. Schoenholz.  A mean field theory of batch normalization.  In _International Conference on Learning Representations_ , 2019.
  * Pennington et al. [2018] Jeffrey Pennington, Samuel Schoenholz, and Surya Ganguli.  The emergence of spectral universality in deep networks.  In _International Conference on Artificial Intelligence and Statistics_ , pages 1924–1932, 2018.
  * de G. Matthews et al. [2018] Alexander G. de G. Matthews, Jiri Hron, Mark Rowland, Richard E. Turner, and Zoubin Ghahramani.  Gaussian process behaviour in wide deep neural networks.  In _International Conference on Learning Representations_ , 2018.
  * Li et al. [2022] Mufan Bill Li, Mihai Nica, and Daniel M. Roy.  The neural covariance sde: Shaped infinite depth-and-width networks at initialization.  _Advances in Neural Information Processing Systems_ , 2022.
  * Jacot et al. [2018] Arthur Jacot, Franck Gabriel, and Clément Hongler.  Neural tangent kernel: Convergence and generalization in neural networks.  _Advances in Neural Information Processing Systems_ , 31, 2018.
  * Saxe et al. [2013] Andrew M. Saxe, James L. McClelland, and Surya Ganguli.  Exact solutions to the nonlinear dynamics of learning in deep linear neural networks.  _arXiv preprint arXiv:1312.6120_ , 2013.
  * Daneshmand et al. [2020] Hadi Daneshmand, Jonas Kohler, Francis Bach, Thomas Hofmann, and Aurelien Lucchi.  Batch normalization provably avoids ranks collapse for randomly initialised deep networks.  _Advances in Neural Information Processing Systems_ , 33:18387–18398, 2020.
  * Bjorck et al. [2018] Nils Bjorck, Carla P. Gomes, Bart Selman, and Kilian Q. Weinberger.  Understanding batch normalization.  _Advances in Neural Information Processing Systems_ , 31, 2018.
  * Dong et al. [2021] Yihe Dong, Jean-Baptiste Cordonnier, and Andreas Loukas.  Attention is not all you need: Pure attention loses rank doubly exponentially with depth.  In _International Conference on Machine Learning_ , pages 2793–2803, 2021.
  * Noci et al. [2022] Lorenzo Noci, Sotiris Anagnostidis, Luca Biggio, Antonio Orvieto, Sidak Pal Singh, and Aurelien Lucchi.  Signal propagation in transformers: Theoretical perspectives and the role of rank collapse.  _arXiv preprint arXiv:2206.03126_ , 2022.
  * Xiao et al. [2018] Lechao Xiao, Yasaman Bahri, Jascha Sohl-Dickstein, Samuel Schoenholz, and Jeffrey 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_ , pages 5393–5402, 2018.
  * Martens et al. [2021] James Martens, Andy Ballard, Guillaume Desjardins, Grzegorz Swirszcz, Valentin Dalibard, Jascha Sohl-Dickstein, and Samuel S. Schoenholz.  Rapid training of deep neural networks without skip connections or normalization layers using deep kernel shaping.  _arXiv preprint arXiv:2110.01765_ , 2021.
  * Zhang et al. [2022] Guodong Zhang, Aleksandar Botev, and James Martens.  Deep learning without shortcuts: Shaping the kernel with tailored rectifiers.  _arXiv preprint arXiv:2203.08120_ , 2022.
  * Schoenholz et al. [2017] Samuel S. Schoenholz, Justin Gilmer, Surya Ganguli, and Jascha Sohl-Dickstein.  Deep information propagation.  In _International Conference on Learning Representations_ , 2017.
  * Ba et al. [2016] Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E. Hinton.  Layer normalization.  _arXiv preprint arXiv:1607.06450_ , 2016.
  * Ioffe and Szegedy [2015] Sergey Ioffe and Christian Szegedy.  Batch normalization: Accelerating deep network training by reducing internal covariate shift.  In _International Conference on Machine Learning_ , pages 448–456. pmlr, 2015.
  * He et al. [2016] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun.  Deep residual learning for image recognition.  In _Proceedings of the IEEE conference on computer vision and pattern recognition_ , pages 770–778, 2016.
  * Devlin et al. [2018] Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova.  Bert: Pre-training of deep bidirectional transformers for language understanding.  _NAACL_ , 2018.
  * Xiong et al. [2020] Ruibin Xiong, Yunchang Yang, Di He, Kai Zheng, Shuxin Zheng, Chen Xing, Huishuai Zhang, Yanyan Lan, Liwei Wang, and Tieyan Liu.  On layer normalization in the transformer architecture.  In _International Conference on Machine Learning_ , pages 10524–10533, 2020.
  * Daneshmand et al. [2021] Hadi Daneshmand, Amir Joudaki, and Francis Bach.  Batch normalization orthogonalizes representations in deep random networks.  _Advances in Neural Information Processing Systems_ , 34:4896–4906, 2021.
  * Klambauer et al. [2017] Günter Klambauer, Thomas Unterthiner, Andreas Mayr, and Sepp Hochreiter.  Self-normalizing neural networks.  _Advances in Neural Information Processing Systems_ , 30, 2017.
  * Daniely et al. [2016] Amit Daniely, Roy Frostig, and Yoram Singer.  Toward deeper understanding of neural networks: The power of initialization and a dual view on expressivity.  _Advances in Neural Information Processing Systems_ , 29, 2016.
  * Yang [2019] Greg Yang.  Scaling limits of wide neural networks with weight sharing: Gaussian process behavior, gradient independence, and neural tangent kernel derivation.  _arXiv preprint arXiv:1902.04760_ , 2019.
  * Poole et al. [2016] Ben Poole, Subhaneil Lahiri, Maithra Raghu, Jascha Sohl-Dickstein, and Surya Ganguli.  Exponential expressivity in deep neural networks through transient chaos.  _Advances in Neural Information Processing Systems_ , 29, 2016.
  * Krizhevsky and Hinton [2009] Alex Krizhevsky and Geoffrey Hinton.  Learning multiple layers of features from tiny images.  Technical report, University of Toronto, 2009.
  * Vaswani et al. [2017] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin.  Attention is all you need.  _Advances in Neural Information Processing Systems_ , 30, 2017.
  * Gershgorin [1931] Semyon Aranovich Gershgorin.  Uber die abgrenzung der eigenwerte einer matrix.  (6):749–754, 1931.
  * Mehler [1866] F Gustav Mehler.  Ueber die entwicklung einer function von beliebig vielen variablen nach laplaceschen functionen höherer ordnung.  _Journal für die Reine und Angewandte Mathematik (in German)_ , 1866\.
  * Li et al. [2021] Mufan Li, Mihai Nica, and Dan Roy.  The future is log-gaussian: Resnets and their infinite-depth-and-width limit at initialization.  _Advances in Neural Information Processing Systems_ , 34:7852–7864, 2021.
  * Joudaki et al. [2023] Amir Joudaki, Hadi Daneshmand, and Francis Bach.  On bridging the gap between mean field and finite width in deep random neural networks with batch normalization.  _International Conference on Machine Learning_ , 2023.
  * Glorot and Bengio [2010] Xavier Glorot and Yoshua Bengio.  Understanding the difficulty of training deep feedforward neural networks.  In _Proceedings of the thirteenth international conference on artificial intelligence and statistics_ , pages 249–256. JMLR Workshop and Conference Proceedings, 2010.

## Appendix outline

The appendix is partitioned into four main components, each serving its
purpose as described:

  1. 1.

Section A details the proof of Theorem 4.1 alongside numerical confirmation of
significant steps:

     * •

Section A.1 provides an elaborate review of the mean-field Gram dynamics.

     * •

Section A.2 presents the Lyapunov function $\gamma$, and establishes that this
function provides both upper and lower bounds for the isometry, thereby
implying that geometric contraction of $\gamma(G^{\ell})$ indicates a
geometric contraction of isometry gap $-\log\mathcal{I}(G^{\ell})$.

     * •

Section A.3 proves that $\gamma(G^{\ell})$ exhibits an exponential contraction
in depth with rate $\beta_{0}$.

  2. 2.

Section B demonstrates the proof of Theorem 3.1.

  3. 3.

Section C empirically elucidates some key insights drawn from our mean-field
theory:

     * •

Section C.1 underscores the role of centering in layer normalization to obtain
isometry.

     * •

Section C.2 experimentally probes the role of projection of each layer to
$\sqrt{d}$ sphere.

     * •

Section C.3 investigates the stabilizing impact of normalization on mean-field
predictions.

     * •

Section C.4 explores the effect of gain on isometry, and links the rate to the
associated non-linearity strength.

  4. 4.

Lastly, Section D gives additional details concerning the experiments reported
in the main text and appendix.

##  Appendix A Proof of Theorem 4.1

###  A.1 Mean-field Gram Dynamics

Recall the mean-field Gram dynamics stated in equation (7):

$\displaystyle G^{\ell+1}_{*}=\mathbb{E}_{h\sim
N(0,G^{\ell}_{*})}\left[\phi(\sigma(h))\phi(\sigma(h))^{\top}\right],$
$\displaystyle\text{ where
}[\phi(a)]_{i}\mathrel{\mathop{\ordinarycolon}}=(a_{i}-\mathbb{E}a_{i})/\sqrt{\operatorname*{Var}{a_{i}}},$
(9)

Assuming that inputs are encoded as columns of $X$, we can restate the MLP
dynamics as follows

$\displaystyle
X^{0}\mathrel{\mathop{\ordinarycolon}}=\frac{1}{\sqrt{d}}{\textsc{LN}}(X-\overline{X}),$
inputs (10) $\displaystyle
H^{\ell}\mathrel{\mathop{\ordinarycolon}}=W^{\ell}X^{\ell},W^{\ell}\sim
N(0,1)^{d\times d},$ preactivation (11) $\displaystyle
A^{\ell}\mathrel{\mathop{\ordinarycolon}}=\sigma(H^{\ell}),$ activations (12)
$\displaystyle
X^{\ell+1}=\frac{1}{\sqrt{d}}{\textsc{LN}}(A^{\ell}-\overline{A^{\ell}})$
$\displaystyle\text{centering \& normalization},$ (13)

where centering and layer normalization are applied column-wise, as defined in
the main text. Observe that Gram matrix of representations can be written as

$\displaystyle G^{\ell+1}_{d}$ $\displaystyle=X^{{\ell+1}^{\top}}X^{\ell+1}$
(14)
$\displaystyle=\frac{1}{d}\sum_{k=1}^{d}\left(\frac{A_{k1}^{\ell}-\mu_{1}}{s_{1}},\dots,\frac{A_{kn}^{\ell}-\mu_{n}}{s_{n}}\right)^{\otimes
2},$
$\displaystyle\mu_{i}\mathrel{\mathop{\ordinarycolon}}=\frac{1}{d}\sum_{k=1}^{d}(A^{\ell}_{ki}),s_{i}\mathrel{\mathop{\ordinarycolon}}=\sqrt{\frac{1}{d}\sum_{k=1}^{d}(A^{\ell}_{ki}-\mu_{i})^{2}},$
(15)

where $\otimes$ denotes Hadamard product, and subscript d emphasises the
dependence of Gram on width $d.$ Note that conditioned on the previous layer,
rows of $H^{\ell}$ and $A^{\ell}$ are i.i.d. , because of independence of rows
of $W^{\ell}.$ Thus, by law of large numbers, in the infinitely wide network
regime, $\mu_{i}$ and $s_{i}$ will converge to the expected mean and variance
respectively $\lim_{d\to\infty}\mu_{i}=\mathbb{E}A^{\ell}_{1i},$ and
$\lim_{d\to\infty}s_{i}=\sqrt{\text{Var}(A^{\ell}_{1i})},$ for all
$i=1,\dots,n.$ By construction of $\phi,$ in the infinitely wide regime, we
can rewrite Gram dynamics as
$\lim_{d\to\infty}G^{\ell}_{d}=\frac{1}{d}\sum_{k=1}^{d}\phi(A^{\ell}_{i\cdot})^{\top}\phi(A^{\ell}_{i\cdot}).$
We can invoke the fact that rows of $A^{\ell}$ are i.i.d. to conclude that
$G_{d}$ is the sample Gram matrix that converges to its expectation

$\displaystyle\lim_{d\to\infty}G^{\ell}_{d}=\mathbb{E}\phi(A^{\ell}_{1\cdot})\phi(A^{\ell}_{1\cdot})^{\top}=\mathbb{E}_{h\sim
N(0,G^{\ell})}\phi(\sigma(h))\phi(\sigma(h))^{\top}=\mathrel{\mathop{\ordinarycolon}}G^{\ell}_{*},$
(16)

where ∗ denotes the mean-field regime $d\to\infty.$ This concludes the
connection between the mean-field Gram dynamics and infinitely wide Gram
dynamics.

###  A.2 Introducing a potential

Here we will introduce a Lyapunov function that enables us to precisely
quantify the isometry of activations in deep networks:

######  Definition 4.

Given a positive semidefinite matrix $G\in\mathbb{R}^{n\times n}$, we define
$\gamma\mathrel{\mathop{\ordinarycolon}}\mathbb{R}^{n\times
n}\to\mathbb{R}^{\geq 0}$ as:

$\displaystyle\gamma(G)\mathrel{\mathop{\ordinarycolon}}=\max_{i\neq
j}\frac{|\widetilde{G}_{ij}|}{1-|\widetilde{G}_{ij}|},$
$\displaystyle\widetilde{G}_{ij}\mathrel{\mathop{\ordinarycolon}}=G_{ij}/\sqrt{G_{ii}G_{jj}}.$
(17)

Remarkably, $\gamma$ obey exhibits an geometric contraction under one MLP
layer update, which is stated in the following theorem:

######  Theorem A.1.

Let $G$ be PSD matrix with unit diagonals $G_{ii}=1,i=1,\dots,n.$ It holds:

$\displaystyle\gamma\left(\mathbb{E}_{h\sim
N(0,G)}\left[\phi(\sigma(h))\phi(\sigma(h))^{\top}\right]\right)\leq\gamma(G)/\beta_{0},$
$\displaystyle\text{ where
}[\phi(a)]_{i}\mathrel{\mathop{\ordinarycolon}}=(a_{i}-\mathbb{E}a_{i})/\sqrt{\operatorname*{Var}{a_{i}}}.$
(18)

Thus, we may apply Theorem A.1 iteratively to prove that in the mean-field,
the Lyapunov function $\gamma(G^{\ell})$ decays at an exponential rate
$\beta_{0}.$ A straightforward induction over layers leads to a decay rate in
$\gamma$, which is presented in the next corollary.

######  Corollary A.2.

If activation $\sigma$ has Hermite expansion with non-linearity strength
$\beta_{0}$, the mean-field Gram matrices $G^{\ell}_{*}$, obey Lyapunov of
these Gram matrices decays at an exponential rate $\beta_{0}$:

$\displaystyle\gamma(G^{\ell}_{*})\leq\gamma(G^{0}_{*})\beta_{0}^{-\ell},$
(19)

where $G^{0}_{*}$ denotes the input Gram matrix.

In Figure A.1, we experimentally validated the above equation for MLPs with a
finite width and activations $\\{\text{tanh ,relu, sigmoid}\\}$ where we
observe that the mean field analysis well predicates the decay in $\gamma$.

Figure A.1: $\gamma(G^{\ell})$ vs depth $\ell,$ for MLP with $n=10,$
$d=10000,$ various activations $\sigma$. Solid lines shows average of $10$
independent runs. The dashed traces the theoretical upper bounds given in
Corollary A.2.

Interestingly, we can connect the Lyapunov $\gamma$ to the isometry, by
proving an upper and lower based on the determinant $G$ based on $\gamma(G)$,
when $G$ is PSD and has unit diagonals:

######  Lemma A.3.

For PSD matrix $G$ with unit diagonals holds:

$\displaystyle\left(1-(n-1)\max_{i\neq j}|G_{ij}|\right)^{n}\leq\det(G)\leq
1-\max_{i\neq j}G_{ij}^{2},$ (20)

where the lower bound holds if $(n-1)|G_{ij}|\leq 1.$

Now, we are ready to prove the main theorem.

######  Theorem A.4 (Restated Theorem 4.1).

Let $\sigma$ be an activation function with a Hermite expansion and a non-
linearity strength $\beta_{0},$ (see equation (6)). Given non-degenerate input
Gram matrix $G^{0}_{*},$ then for sufficiently large layer
$\ell\gtrsim\beta_{0}^{-1}(-n\log\mathcal{I}(G_{*}^{0})+\log(4n))$, we have

$\displaystyle-\log\mathcal{I}(G^{\ell}_{*})\leq\exp(-\ell\log\beta_{0}-n\log\mathcal{I}(G^{0}_{*})+\log(4n)).$
(21)

###### Proof of Theorem 4.1.

First, note that we have the following transformation

$\displaystyle t=\frac{x}{1-x}\implies x=\frac{t}{t+1}\implies\max_{i\neq
j}|G_{ij}|=\frac{\gamma(G)}{1+\gamma(G)},$ (22)

we have

$\displaystyle\det(G_{*}^{\ell})$ $\displaystyle\geq(1-(n-1)\max_{i\neq
j}|G_{{ij}_{*}}^{\ell}|)^{n}$ Lemma A.3 (23) $\displaystyle\det(G_{*}^{\ell})$
$\displaystyle=\left(1-(n-1)(\gamma(G^{\ell}_{*})/(1+\gamma(G^{\ell}_{*})))\right)^{n}$
Invoke (22) (24)
$\displaystyle\geq\left(1-(n-1)\gamma(G^{\ell}_{*})\right)^{n}$ (25)
$\displaystyle\implies-\frac{1}{n}\log\det(G_{*}^{\ell})$
$\displaystyle\leq-\log(1-(n-1)\gamma(G_{*}^{\ell}))$ (26) $\displaystyle\leq
2n\gamma(G_{*}^{\ell})$ for $\gamma(G^{\ell}_{*})<1/(2n-2)$ (27)
$\displaystyle\leq 2n\gamma(G_{0})\beta_{0}^{-\ell}.$ (28)

By upper bound of Lemma A.3 we have $\max_{i\neq
j}|G_{ij}|\leq\sqrt{1-\det(G)}\leq 1-\det(G)/2,$ implying $\gamma(G_{0})\leq
2\det(G_{0})^{-1},$ which allows us to

$\displaystyle-\log\mathcal{I}(G_{*}^{\ell})$ $\displaystyle\leq
4n\det(G_{0})^{-1}\beta_{0}^{-\ell}$ (29)
$\displaystyle\leq\exp(-\ell\log\beta_{0}-n\log\mathcal{I}(G_{0})+\log(4n))$
(30)

which concludes the proof. ∎

###### Proof of Lemma A.3.

_Lower bound._ The lower bound is a result of Gershgorin circle theorem
[Gershgorin, 1931], which implies that every eigenvalue must be within the
disc $[1-(n-1)\max_{i\neq j}|G_{ij}|,1+(n-1)\max_{i\neq j}|G_{ij}|].$ Thus,
the determinant is lower-bounded by $(1-(n-1)\max_{i\neq j}|G_{ij}|)^{n}.$

_Upper bound._ Since $G$ is PSD, we can write $G=X^{\top}X,$ which implies
that columns of $X$ are unit norm $\|x_{i}\|=1,$ and $G$ encodes the angles
between them $\cos\angle(x_{i},x_{j})=\langle x_{i},x_{j}\rangle=G_{ij}.$
Furthermore, we have $\det(G)=\text{vol}(X)^{2},$ where the volume refers to
the parallelepiped spanned by the columns of $X.$ With this formulation, we
write volume recursively as volume spanned by columns $1$ up to $n-1,$ times
the projection distance of $x_{n}$ from their span

$\displaystyle\text{vol}(x_{1},\dots,x_{n})=\text{dist}(x_{n},\text{span}(x_{1},\dots,x_{n-1}))\text{vol}(x_{1},\dots,x_{n-1}),$

where span refers to the space of all linear combinations of these vectors.
Note that projection distance of $x_{i}$ onto $x_{j}$ can be written as
$\sin\angle(x_{i},x_{j})=\sqrt{1-\langle
x_{i},x_{j}\rangle^{2}}=\sqrt{1-G_{ij}^{2}}.$ Since the projection distance
onto the linear span cannot be greater than projection distance onto a single
vector, which is bounded by $\sqrt{1-\max_{i\neq j}G_{ij}^{2}}$. This
concludes the upper bound that $\det(G)\leq 1-\max_{i\neq j}G_{ij}^{2}$. ∎

###  A.3 Dual activation and proof of Thm. A.1

According to Daniely et al. [2016], we leverage the notion dual activation
associated with the activation function $\sigma,$ which is defined in
following.

######  Definition 5.

Given activation
$\sigma\mathrel{\mathop{\ordinarycolon}}\mathbb{R}\to\mathbb{R}$ that is
square integrable with respect to the Gaussian kernel, define its dual
activation
$\hat{\sigma}\mathrel{\mathop{\ordinarycolon}}\mathbb{R}\to\mathbb{R},$ and
its mean reduced dual
$\bar{\sigma}\mathrel{\mathop{\ordinarycolon}}\mathbb{R}\to\mathbb{R}$ as:

$\displaystyle\hat{\sigma}(\rho)\mathrel{\mathop{\ordinarycolon}}=\mathbb{E}\sigma(x)\sigma(y)$
$\displaystyle\bar{\sigma}(\rho)\mathrel{\mathop{\ordinarycolon}}=\operatorname{\mathbb{E}}\left[{(\sigma(x)-\mathbb{E}\sigma(x))(\sigma(y)-\mathbb{E}\sigma(y))}\right]$
(31)

The following lemma connects the dual activation and its mean-reduced version
with the Hermite expansion of $\sigma$:

######  Lemma A.5.

Given two standard Gaussian variables $X,Y\sim N(0,1)$ with covariance
$\mathbb{E}XY=\rho,$ and activation $\sigma$ with normalized Hermite
coefficients ${c_{k}}_{k},$ we have

$\displaystyle\hat{\sigma}(\rho)=\sum_{k=0}^{\infty}c_{k}^{2}\rho^{k},$
$\displaystyle\bar{\sigma}(\rho)=\sum_{k=1}^{\infty}c_{k}^{2}\rho^{k}.$ (32)

Finally, we have the tools to prove the first main theorem.

###### Proof of Theorem A.1.

Let $a=\sigma(h)$ for $h\sim N(0,G).$ Note that by definition of the dual-
activation we have $\mathbb{E}aa^{\top}=[\hat{\sigma}(G_{ij})]_{i,j\leq n}.$
Since we assumed $G$ has unit diagonals, we have $h_{i}\sim N(0,1),$ which
implies that $\mathbb{E}a_{i}=\mathbb{E}_{h_{i}\sim
N(0,1)}\sigma(h_{i})=c_{0},$ implying that
$\mathbb{E}(a_{i}-\mathbb{E}a_{i})(a_{j}-\mathbb{E}a_{j})^{\top}=\hat{\sigma}(G_{ij})-c_{0}^{2}=\bar{\sigma}(G_{ij}).$
Furthermore, the variance can be driven as
$\operatorname*{Var}(a_{i})=\mathbb{E}a_{i}^{2}-(\mathbb{E}a_{i})^{2}=\sum_{k=0}^{\infty}c_{k}^{2}=\sum_{k=1}^{\infty}c_{k}^{2}=\bar{\sigma}(1)$
for all $i=1,\dots,n.$ Thus, we have
$\mathbb{E}\phi(a_{i})\phi(a_{j})=\bar{\sigma}(G_{ij})/\bar{\sigma}(1).$ In
the matrix form we have

$\displaystyle\gamma\left(\mathbb{E}_{h\sim
N(0,G)}\left[\phi(\sigma(h))\phi(\sigma(h))^{\top}\right]\right)=\gamma\left([\bar{\sigma}(G_{ij})/\bar{\sigma}(1)]_{i,j\leq
n}\right)$ (33)

The remainder proof relies on the following contractive property of Gram
matrix potential:

######  Lemma A.6.

Consider activation $\sigma,$ with normalized Hermite coefficients
$\\{c_{k}\\}_{k\geq 0}.$ For all $\rho\in(0,1),$ the mean-reduced dual
activation $\bar{\sigma}$ obeys

$\displaystyle\frac{|\bar{\sigma}(\rho)|/\bar{\sigma}(1)}{1-|\bar{\sigma}(\rho)|/\bar{\sigma}(1)}\leq\beta_{0}^{-1}\frac{|\rho|}{1-|\rho|},$
(34)

which the right hand-side is strictly larger if some nonlinear coefficient is
nonzero $c_{k}\neq 0$ for some $k\geq 2.$

Thus we can apply Lemma A.6 on each element $i\neq j$ to conclude that

$\displaystyle\frac{|\bar{\sigma}(G_{ij})|/\bar{\sigma}(1)}{1-|\bar{\sigma}(G_{ij})|/\bar{\sigma}(1)}\leq\frac{|G_{ij}|}{1-|G_{ij}|}\beta_{0}^{-1}$
$\displaystyle\text{ Lemma~{}\ref{lem:hermite_nonlinearity} for all }i\neq j$
(35) $\displaystyle\leq\beta^{-1}_{0}\max_{i\neq
j}\frac{|G_{ij}|}{1-|G_{ij}|}$ (36) $\displaystyle=\beta_{0}^{-1}\gamma(G).$
(37)

since the inequality holds for any value of $i\neq j,$ we can take the maximum
over $i\neq jj$ to write:

$\displaystyle\gamma\left(\bar{\sigma}(G)/\bar{\sigma}(1)\right)=\max_{i\neq
j}\frac{|\bar{\sigma}(G_{ij})|/\bar{\sigma}(1)}{1-|\bar{\sigma}(G_{ij})|/\bar{\sigma}(1)}\leq\beta_{0}^{-1}\gamma(G),$
(39)

which concludes the proof. ∎

###### Proof of Lemma A.6.

Note the ratio is invariant to scaling of $\sum_{k=1}^{\infty}c_{k}^{2}$.
Hence, we assume $\sum_{k=1}^{\infty}c_{k}^{2}=1$ without loss of generality.
With this simplification, we have $\bar{\sigma}(1)=1.$ For the positive range
$\rho\in[0,1]$ we have

$\displaystyle(\frac{\bar{\sigma}(\rho)}{1-\bar{\sigma}(\rho)})(\frac{\rho}{1-\rho})^{-1}$
$\displaystyle=\frac{\rho^{-1}\sum_{k=1}^{\infty}c_{k}^{2}\rho^{k}}{(1-\rho)^{-1}\sum_{k=1}^{\infty}c_{k}^{2}(1-\rho^{k})}$
(40)
$\displaystyle=\frac{\sum_{k=1}^{\infty}c_{k}^{2}\rho^{k-1}}{\sum_{k=1}^{\infty}c_{k}^{2}(\frac{1-\rho^{k}}{1-\rho})}$
(41)
$\displaystyle=\frac{\sum_{k=1}^{\infty}c_{k}^{2}\rho^{k-1}}{\sum_{k=1}^{\infty}c_{k}^{2}\left(\sum_{i=0}^{k-1}\rho^{i}\right)}$
(42)
$\displaystyle\leq\frac{\sum_{k=1}^{\infty}c_{k}^{2}\rho^{k-1}}{(\sum_{k=1}^{\infty}c_{k}^{2}\rho^{k-1})+\sum_{k=2}^{\infty}c_{k}^{2}}$
(43)
$\displaystyle\leq\max_{\rho\in[0,1]}\frac{\sum_{k=1}^{\infty}c_{k}^{2}\rho^{k-1}}{(\sum_{k=1}^{\infty}c_{k}^{2}\rho^{k-1})+\sum_{k=2}^{\infty}c_{k}^{2}}$
(44)
$\displaystyle=\frac{\sum_{k=1}^{\infty}c_{k}^{2}}{\sum_{k=1}^{\infty}c_{k}^{2}+\sum_{k=1}^{\infty}c_{k}^{2}-c_{1}^{2}}$
(45) $\displaystyle=\frac{1}{\beta_{0}}.$ (46)

Thus for $\rho\in[0,1]$ we have

$\displaystyle\frac{\bar{\sigma}(\rho)}{1-\bar{\sigma}(\rho)}\leq\frac{\rho}{1-\rho}\beta_{0}^{-1}.$
(47)

By Jensen inequality for convex function $x\mapsto|x|$ we have

$\displaystyle|\bar{\sigma}(\rho)|=|\sum_{k=1}^{\infty}c_{k}^{2}\rho^{k}|$
$\displaystyle\leq\sum_{k=1}^{\infty}c_{k}^{2}|\rho|^{k}=\bar{\sigma}(|\rho|)$
(48)

Because $x\mapsto x/(1-x)$ is monotonically increasing for $x\in[0,1]$, we
have

$\displaystyle\frac{|\bar{\sigma}(\rho)|}{1-|\bar{\sigma}(\rho)|}\leq\frac{\bar{\sigma}(|\rho|)}{1-\bar{\sigma}(|\rho|)}\leq\frac{|\rho|}{1-|\rho|}\beta_{0}^{-1}.$
(49)

Where we invoked the inequality that was proven for $\rho\in[0,1],$ because
$|\rho|\in[0,1].$ ∎

The following lemma, which is a consequence of Mehler’s formula, is at the
hart of proof of Lemma A.5:

######  Lemma A.7 (Consequence of Mehler’s kernel).

If $X,Y\sim N(0,1)$ with covariance $EXY=\rho$ we have

$E_{X,Y}\mathit{He}_{n}(X)\mathit{He}_{k}(Y)=\rho^{n}\delta_{nk}$

where $\delta_{nk}$ is the Dirac delta.

###### Proof of Lemma A.5.

Let $\sigma(x)=\sum_{k=0}c_{k}\mathit{He}_{k}(x),$ denote the Hermite
expansion of $\sigma.$ Thus, we have

$\displaystyle\hat{\sigma}(\rho)$
$\displaystyle\mathrel{\mathop{\ordinarycolon}}=\mathbb{E}_{X,Y}\sigma(X)\sigma(Y)$
(50)
$\displaystyle=\sum_{n,k=0}^{\infty}c_{n}c_{k}\mathbb{E}_{X,Y}\mathit{He}_{n}(X)\mathit{He}_{k}(Y)$
(51) $\displaystyle=\sum_{k=0}^{\infty}c_{k}^{2}\rho^{k},$ (52)

where in the last line we applied result of Lemma A.7. For the mean-reduced
dual kernel $\bar{\sigma},$ observe that $\mathbb{E}_{X\sim
N(0,1)}\sigma(X)=c_{0}.$ Thus, the reduction of mean will cancell the $k=0$
term, which concludes the proof that
$\bar{\sigma}(\rho)=\sum_{k=1}c_{k}^{2}\rho^{k}.$ ∎

###### Proof of Lemma A.7.

The property can be deduced from Mehler’s formula Mehler [1866]. The formula
states that

$\displaystyle\frac{1}{\sqrt{1-\rho^{2}}}\exp\left(-\frac{\rho^{2}(x^{2}+y^{2})-2xy\rho}{2(1-\rho^{2})}\right)$
(53) $\displaystyle=\sum_{m=0}^{\infty}\mathit{He}_{m}(x)\mathit{He}_{m}(y),$
(54)

where the $m!$ factor difference is due to the definition of Hermite
polynomials with an additional $1/\sqrt{m!}$ compared to the one used in
Mehler’s kernel. Observe that the left hand side is equal to
$p(x,y)/p(x)p(y)$, where $p(x,y)$ is the joint PDF of $(X,Y)$, and $p(x),p(y)$
are PDF of $X$ and $Y$ respectively. Therefore, we can take the expectation
using the expansion

$\displaystyle\mathbb{E}_{X,Y}\left[\mathit{He}_{n}(X)\mathit{He}_{k}(Y)\right]$
$\displaystyle=\int\mathit{He}_{n}(x)\mathit{He}_{k}(y)p(x,y)dxdy$ (55)
$\displaystyle=\sum_{m=0}^{\infty}\rho^{m}\int\mathit{He}_{n}(x)\mathit{He}_{k}(y)\mathit{He}_{m}(x)\mathit{He}_{m}(y)dp(x)dp(y)$
(56) $\displaystyle=\sum_{m=0}^{\infty}\rho^{m}\mathbb{E}_{X\sim
N(0,1)}\left[\mathit{He}_{n}(X)\mathit{He}_{m}(X)\right]\mathbb{E}_{Y\sim
N(0,1)}\left[\mathit{He}_{k}(Y)\mathit{He}_{m}(Y)\right]$ (57)
$\displaystyle=\rho^{n}\delta_{nk}$ (58)

where in the last line we used the orthogonality property $E_{X\sim
N(0,1)}H_{k}(x)H_{n}(X)=\delta_{nk}$. ∎

##  Appendix B Proof of Theorem 3.1 and properties of isometry

#### Basic properties of isometry

It is straightforward to check isometry obeys the following basic isometry-
preserving properties:

######  Lemma B.1.

For PSD matrix $M,$ the isometry defined in (3) obeys the following
properties: 1) scale-invariance $\mathcal{I}(cM)=\mathcal{I}(M),$ 2) only
takes value in the unit range $\mathcal{I}(M)\in[0,1]$ 3) it takes its maximum
value if and only if $M$ is identity $\mathcal{I}(M)=1\iff M=I_{n},$ and 3)
takes minimum value if and only if $M$ is degenerate $\mathcal{I}(M)=0.$

###### Proof of Lemma B.1.

The scale-invariance is trivially true as scaling $M$ by any constant will
scale $\det(M)^{1/n}$ and $\mathrm{tr}(M)$ by the same amount. The proof of
other properties is a straightforward consequence of writing the isometry in
terms of the eigenvalues
$\mathcal{I}(M)=(\prod_{i}\lambda_{i})^{1/n}/(\frac{1}{n}\sum_{i}\lambda_{i}),$
where $\lambda_{i}$’s are eigenvalues of $M.$ By arithmetic vs geometric mean
inequality over the eigenvalues we have
$(\prod_{i}\lambda_{i})^{1/n}\leq\frac{1}{n}\sum_{i}\lambda_{i}),$ which
proves that $\mathcal{I}(M)\in[0,1].$ Furthermore, the inequality is tight iff
the values are all equal $\lambda_{1}=\dots=\lambda_{n},$ which holds only for
identity $M=I_{n}$. Finally, isometry is zero iff at least one eigenvalue is
zero, which is the case for degenerate matrix $M.$ ∎

###### Proof of Theorem 3.1.

Define
$D\mathrel{\mathop{\ordinarycolon}}=\mathrm{diag}(a_{1}/\sqrt{d},\dots,a_{n}/\sqrt{d}).$
Observe that $C=D\widetilde{G}D,$ implying
$\det(G)=\det(\widetilde{G})\det(D)^{2}.$ Because $\widetilde{x_{i}}$’s have
norm $\sqrt{d},$ diagonals of Gram after normalization are constant
$\widetilde{G}_{ii}=d,$ implying $\frac{1}{n}\mathrm{tr}(\widetilde{G})=d.$ We
have

$\displaystyle\frac{\mathcal{I}(\widetilde{G})}{\mathcal{I}(G)}$
$\displaystyle=\frac{\frac{1}{n}\mathrm{tr}(G)}{\frac{1}{n}\mathrm{tr}(\widetilde{G})}\frac{\det(\widetilde{G})^{1/n}}{\det(G)^{1/n}}$
(59)
$\displaystyle=\frac{\frac{1}{n}\sum_{i}^{n}a_{i}^{2}}{d}\frac{\det(\widetilde{G})^{1/n}}{\det(\widetilde{G})^{1/n}(d^{-n}\prod_{i}^{n}a_{i}^{2})^{1/n}}$
(60)
$\displaystyle=\frac{(\frac{1}{n}\sum_{i}a_{i})^{2}}{(\prod_{i}^{n}a_{i})^{2/n}}\frac{\frac{1}{d}\sum_{i}^{n}a_{i}^{2}}{(\frac{1}{n}\sum_{i}a_{i})^{2}}$
(61)
$\displaystyle=1+\frac{\frac{1}{n}\sum_{i}^{n}(a_{i}-\bar{a})^{2}}{\bar{a}^{2}},$
$\displaystyle\bar{a}\mathrel{\mathop{\ordinarycolon}}=\frac{1}{n}\sum_{i}^{n}a_{i}$
(62)

∎

##  Appendix C Implications of our mean-field theory

Our Theorem 4.1 proof underscores the pivotal role of two elements of layer
normalization in controlling the Hermite expansion: (i) centering and (ii)
projection onto the $\sqrt{d}$-sphere. We shall delve into the significance of
centering in Section C.1, the projection’s role in Section C.2, and the
stabilizing effects of projection in Section C.3.

#### Experimental Setup.

In this section, $a^{\ell}$ refers to the post-activation vector
$a^{\ell}\mathrel{\mathop{\ordinarycolon}}=\sigma(W^{\ell}x^{\ell}),W^{\ell}\sim
N(0,1/d)^{d\times d}.$ Compared to equation (1), the weight matrix has
absorbed the factor $1/\sqrt{d},$ and thus its elements are drawn from
$N(0,1/d)$ instead of $N(0,1).$ Note that the offset $c_{0}=\mathbb{E}_{z\sim
N(0,1)}\sigma(z)$, and $\bar{\sigma}(1)$ and dual-activation
$\bar{\sigma}(1)=\mathbb{E}_{z\sim
N(0,1)}(\sigma(z)-c_{0})^{2}=\sum_{k=1}^{\infty}c_{k}^{2},$ are absolute
constants given the activation. We shall examine four different configurations
emerging from our mean-field theory to underscore its implications:

###  C.1 Importance of centering in obtaining isometry

The proof of Theorem 4.1 necessitates centering in layer normalization to
ensure that the $c_{0}$ term in the Hermite expansion disappears. Figure C.1
demonstrates through numerical simulations that this condition is crucial for
achieving isometry. Interestingly, activations with $c_{0}=0$, such as
hyperbolic tangent, do not show any differential effect with or without
centering. Conversely, for activations with a substantial offset term, like
sigmoid, ReLU, or exponential function, the difference is markedly pronounced.

Figure C.1: The contribution of centering in achieving isometry. Batch size
$n=10$, width $d=1000.$ No centering (blue):
$x^{\ell+1}={\textsc{LN}}(a^{\ell})$. Layer centering (orange):
$x^{\ell+1}={\textsc{LN}}(a^{\ell}-\overline{a^{\ell}})$

To further corroborate the hypothesis that centering solely ensures the
disappearance of the $c_{0}$ term, we can explicitly remove the offset term
$c_{0}$ from activation as $\sigma(x)-c_{0}.$ Strikingly, our experiments
presented in Figure C.2 indicate that this $c_{0}$-centering is equally or
more efficient than mean-reduction in achieving isometry. This can be
rationalized by the fact that while centering in layer norm converges to
$c_{0}$ in expectation, for finite $d=1000,$ the mean value is merely an
approximation of $c_{0}.$ This supports our hypothesis that the disappearance
of $c_{0}$ is imperative for achieving isometry, and provides fresh insights
into the role of centering in layer normalization.

Figure C.2: Layer vs mean-field centering in obtaining isometry. Batch size
$n=10$, width $d=1000.$
Layer centering (orange):
$x^{\ell+1}={\textsc{LN}}(a^{\ell}-\overline{a^{\ell}}).$ Mean-field centering
(blue): $x^{\ell+1}={\textsc{LN}}(a^{\ell}-c_{0})$.

###  C.2 The importance of projection in achieving isometry

Our mean-field theory emphasizes the role of projection onto the
$\sqrt{d}$-sphere in addition to centering to achieve isometry. Figure C.3
underlines the importance of this projection for different activations. It’s
evident that projection plays a pivotal role in achieving isometry,
particularly significant for ReLU and sigmoid.

Figure C.3: Importance of projection in obtaining isometry, batch size $n=10$,
width $d=1000.$ No projection (blue): $x^{\ell+1}=(a^{\ell}-c_{0})$. Layer
projection (orange): $x^{\ell+1}={\textsc{LN}}(a^{\ell}-c_{0})$.

According to our mean-field analysis, the factor
$\frac{1}{d}\sum_{i=1}^{d}(a^{\ell}_{i}-\overline{a^{\ell}})^{2}$ in layer
normalization converges to variance $\text{var}_{z\sim
N(0,1)}(\sigma(z))=\bar{\sigma}(1)=\sum_{k=1}^{\infty}c_{k}^{2}.$ Thus, as the
width increases, the layer normalization operator
${\textsc{LN}}(a^{\ell}-\overline{a^{\ell}})$ will converge to
$(a^{\ell}-c_{0})/\sqrt{\bar{\sigma}(1)}.$ Figure C.4 demonstrates that the
constant scaling $1/\sqrt{\bar{\sigma}(1)},$ achieves comparable isometry to
layer projection for hyperbolic tangent and sigmoid, and enhances the isometry
for ReLU. This further strengthens the insight from our theory that one
significant function of projection is to ensure that Hermite coefficients are
normalized.

Figure C.4: Layer vs mean-field projections. Batch size $n=10$, width
$d=1000.$
Layer projection (blue): $x^{\ell+1}={\textsc{LN}}(a^{\ell}-c_{0})$. Mean-
field projection (orange):
$x^{\ell+1}=\frac{a^{\ell}-c_{0}}{\sqrt{\bar{\sigma}(1)}}.$

###  C.3 Normalization’s role in stabilizing mean-field accuracy through depth

Numerous theoretical studies conjecture that mean field predictions may not be
reliable for considerably deep neural networks [Li et al., 2021, Joudaki et
al., 2023]. Mean field analysis encounters $O(1/\sqrt{\text{width}})$ error in
each layer when the network width is finite. This error may accumulate with
depth, making mean field predictions increasingly inaccurate with an increase
in depth. However, Figure C.5 illustrates that layer normalization controls
this error accumulation through depth. This might be attributable to the
isometry bias induced by normalization, as proven in Theorem 3.1. Similarly,
batch normalization also prevents error propagation with depth by imposing the
same isometry [Joudaki et al., 2023]. This observation warrants future
research on the essential role normalization plays in ensuring the accuracy of
mean field predictions.

Figure C.5: Role of normalization in stabilizing mean-field accuracy for ReLU,
batch size $n=10.$
Mean-field normalization (blue):
$x^{\ell+1}=\frac{a^{\ell}-c_{0}}{\sqrt{\bar{\sigma}(1)}}$. Layer
normalization (orange): $x^{\ell+1}={\textsc{LN}}(a^{\ell}-c_{0})$

###  C.4 Quantifying the influence of gain on isometry through non-linearity
strength

The concept of gain in neural networks is vital and closely connected with the
weights initialization. A neural network with properly initialized weights can
learn faster, have a lesser chance of getting stuck at sub-optimal solutions,
and provide better generalization. The impact of gain can be visualized
through the lens of weight initialization strategies such as Xavier
normalization [Glorot and Bengio, 2010], which has shown significant
effectiveness in optimizing neural networks. These initialization strategies
apply a gain value to the weights, which is a scaling factor, to ensure a good
signal flow through many layers during the forward and backward passes. The
gain value essentially determines the variance of the weights in the
initialization stage.

As an extension to our prior investigations, we delve into understanding the
influence of gain on isometry, predominantly through our calculated metric,
the non-linearity strength, denoted as $\beta_{0}$, as a function of gain
$\alpha.$ For certain instances, such as ReLU, sine, and exponential
activations, we are capable of deriving $\beta_{0}(\alpha)$ in a closed form.
Table C.1 presents a few of these cases.

$\sigma$ | $\exp(\alpha x)$ | $\sin(\alpha x)$ | $\max(\alpha x,0)$
---|---|---|---
$\beta_{0}$ | $\frac{2-2e^{\alpha^{2}}+\alpha^{2}}{1-e^{\alpha^{2}}}$ | $\frac{2(-1+e^{2\alpha^{2}}-e^{\alpha^{2}}\alpha^{2})}{-1+e^{2\alpha^{2}}}$ | $\frac{4-3\pi}{2-2\pi}$
Table C.1: Relationship between non-linearity strength and gain.

For a more extensive selection of activation functions, we have numerically
computed the non-linearity strength as a function of gain $\alpha$, as
visualized in Figure C.6. Leveraging Theorem 4.1, these computed values
provide an estimation for the isometry rate for various activations. This
correlation proves to be remarkably predictive, as shown in Figure C.7.
Remarkably, in the case of ReLU activation, its unique characteristics lead to
both our closed form $\beta_{0}$ (refer Table C.1) and rate towards isometry
(refer Figure C.7) remaining consistent across different values of $\alpha$.

Figure C.6: Analysis of gain through the lens of non-linearity strength.
Figure C.7: Influence of gain on the attainment of isometry. Figure C.8:
Isometry vs depth for hyperbolic tangent activation. The traces depict
activations with different gains $\alpha=0.5,1,2,5.$ The dash traces are the
lines are parallel to $\beta_{0}^{-\ell}$. Figure C.9: Isometry vs depth for
mean-field centering and normalization to Xavier initialization.

#### Comparison to Xavier gain for initialization

Inspired by the results so far, we can compare mean-field centering and
normalization to Xavier gain for activations. Figure C.9 demonstrates that all
mean-field based gains improve the isometry when compared with Xavier
initialization. However, this is markedly stronger for ReLU and leaky ReLU. We
can explain this starker contrast by the fact that both activations have a
significant offset term $c_{0},$ which is not corrected by the Xavier
initialization. We can explain this difference by conceptually comparing the
two methods as shaping the activation function as follows:

mean-field correction
$\displaystyle\frac{\sigma(x)-c_{0}}{\sqrt{\bar{\sigma}(1)}}$ (63) Xavier
$\displaystyle\frac{\sigma(x)}{\sqrt{\hat{\sigma}(1)}}.$ (64)

Recall that $\hat{\sigma}(1)=\bar{\sigma}(1)-c_{0}^{2}.$ Thus, for activations
that we have $c_{0}\neq 0,$ such as ReLU and leaky ReLU, the difference mean-
field approach with centering achieves a higher isometry compared to Xavier.

##  Appendix D Details about empirical validations

#### Codes and reproducibility.

We implemented our experiments in Python using the PyTorch framework
111https://pytorch.org.. All the figures are reproducible with the code
attached in the supplementary.

#### Hardware.

All experiments, except for Figure 5.1 were run on a single AMD Ryzen 9 3900X
12-Core CPU, which takes about 5 minutes overall. For Figure 5.1, we trained
our models on a single NVIDIA GeForce RTX 3090 GPU, which takes about 3
minutes.

#### Figures.

The solid lines in all plots represent the average performance over multiple
independent runs, and the shaded regions indicate the confidence intervals.
Unless stated otherwise, each average is computed over #10 independent runs.

#### Training procedure for Figure 5.1

In each epoch, the isometry gap and normalization bias (equation (5)) are
computed per each batch, and then averaged over the entire test set. The epoch
$i$ corresponds to the network at after $i$ steps of training on the training
set of CIFAR10 (epoch $0$ means network is at initialization).