Understanding Model Parameter Initialization Strategies from a Geometric Perspective

Orthogonality at Your Fingertips

Some time ago, I wrote "Distribution of the Angle Between Two Random Vectors in n-Dimensional Space", one of the corollaries of which is:

Corollary 1: Any two random vectors in a high-dimensional space are almost always orthogonal.

In fact, Corollary 1 is the starting point for the entire geometric perspective of this article! A further corollary is:

Corollary 2: If n^2 numbers are randomly selected from \mathcal{N}(0, 1/n) to form an n\times n matrix, this matrix is approximately an orthogonal matrix. The larger n is, the better the approximation.

Skeptical readers can verify this numerically:

import numpy as np

n = 100
W = np.random.randn(n, n) / np.sqrt(n)
X = np.dot(W.T, W)  # Matrix multiplied by its own transpose
print(X)  # See if it is close to the identity matrix
print(np.square(X - np.eye(n)).mean())  # Calculate the MSE with the identity matrix

I believe that for most readers, seeing Corollary 2 for the first time will be more or less surprising. An orthogonal matrix is a matrix that satisfies \boldsymbol{W}^{\top}\boldsymbol{W}=\boldsymbol{I}, which means its inverse is equal to its transpose. Generally, the difficulty of solving for the inverse and the transpose of a matrix is vastly different, so we tend to feel that "inverse = transpose" should be a very strict condition. However, Corollary 2 tells us that a randomly sampled matrix is already close to being an orthogonal matrix, which is admittedly somewhat counter-intuitive. When I first realized this, I was also quite surprised.

It’s Actually Not That Hard to Understand

However, once we get used to the fact in Corollary 1 that "any two random vectors in high-dimensional space are almost always orthogonal," we can indeed understand and derive this result quickly. For a quick derivation, we can first consider the standard normal distribution \mathcal{N}(0,1). Note that Corollary 1 requires the sampling direction to be uniform, and the standard normal distribution satisfies this requirement. Sampling an n\times n matrix from \mathcal{N}(0,1) can be seen as sampling n vectors of n dimensions. Since these n vectors are all random vectors, they are naturally close to being pairwise orthogonal.

Of course, pairwise orthogonality does not yet make an orthogonal matrix, because an orthogonal matrix also requires each vector to have a norm of 1. We have \mathbb{E}_{x\sim \mathcal{N}(0,1)}\left[x^2\right]=1, which means the norm of an n-dimensional vector sampled from \mathcal{N}(0,1) is approximately \sqrt{n}. Therefore, to approach orthogonality, each element needs to be divided by \sqrt{n}, which is equivalent to changing the sampling variance from 1 to 1/n.

Furthermore, the sampling distribution does not necessarily have to be a normal distribution; for example, a uniform distribution U\left[-\sqrt{3/n}, \sqrt{3/n}\right] also works. In fact, we have:

Corollary 3: An n\times n matrix independently and repeatedly sampled from any distribution p(x) with a mean of 0 and a variance of 1/n is close to an orthogonal matrix.

We can understand Corollary 3 from a more mathematical perspective: assume \boldsymbol{x}=(x_1,x_2,\dots,x_n) and \boldsymbol{y}=(y_1,y_2,\dots,y_n) are both sampled from p(x). Then we have: \begin{aligned} \langle \boldsymbol{x}, \boldsymbol{y}\rangle =&\, n\times \frac{1}{n}\sum_{k=1}^n x_k y_k\\ \approx&\, n\times \mathbb{E}_{x\sim p(x),y\sim p(x)}[xy]\\ =&\, n\times \mathbb{E}_{x\sim p(x)}[x]\times \mathbb{E}_{y\sim p(x)}[y]\\ =&\,0 \end{aligned} And: \begin{aligned} \Vert\boldsymbol{x}\Vert^2 =&\, n\times \frac{1}{n}\sum_{k=1}^n x_k^2\\ \approx&\, n\times \mathbb{E}_{x\sim p(x)}\left[x^2\right]\\ =&\, n\times \left(\mu^2 + \sigma^2\right)\\ =&\,1 \end{aligned} Thus, any two vectors are approximately orthonormal, and therefore the sampled matrix is also close to an orthogonal matrix.

Now We Can Talk About Initialization

We have discussed orthogonal matrices at length, but essentially, it was all to pave the way for understanding the geometric meaning of initialization methods. If readers still remember linear algebra, they should recall that the important significance of an orthogonal matrix is that it preserves the norm of a vector during transformation. Expressed mathematically, if \boldsymbol{W}\in \mathbb{R}^{n\times n} is an orthogonal matrix and \boldsymbol{x}\in\mathbb{R}^n is any vector, then the norm of \boldsymbol{x} is equal to the norm of \boldsymbol{W}\boldsymbol{x}: \Vert\boldsymbol{W}\boldsymbol{x}\Vert^2 = \boldsymbol{x}^{\top}\boldsymbol{W}^{\top}\boldsymbol{W}\boldsymbol{x}=\boldsymbol{x}^{\top}\boldsymbol{x}=\Vert\boldsymbol{x}\Vert^2 Consider a fully connected layer: \boldsymbol{y}=\boldsymbol{W}\boldsymbol{x} + \boldsymbol{b} A deep learning model is essentially a nesting of fully connected layers. Therefore, to prevent the final output of the model from "exploding" or "vanishing" during the initialization phase, one idea is to let the model maintain the norm during initialization.

A natural initialization strategy formed by this idea is to "initialize \boldsymbol{b} with all zeros and initialize \boldsymbol{W} with a random orthogonal matrix." Corollary 2 has already told us that an n\times n matrix sampled from \mathcal{N}(0, 1/n) is already close to an orthogonal matrix, so we can sample from \mathcal{N}(0, 1/n) to initialize \boldsymbol{W}. This is the Xavier initialization strategy, also called Glorot initialization in some frameworks because the author is Xavier Glorot. Furthermore, the sampling distribution is not necessarily \mathcal{N}(0, 1/n); as Corollary 3 stated, you can sample from any distribution with a mean of 0 and a variance of 1/n.

The above discussion assumes the input and output dimensions are both n. What if the input is n-dimensional and the output is m-dimensional? In this case, \boldsymbol{W}\in\mathbb{R}^{m\times n}, and the condition to keep the norm of \boldsymbol{W}\boldsymbol{x} unchanged is still \boldsymbol{W}^{\top}\boldsymbol{W}=\boldsymbol{I}. However, when m < n, this is impossible; when m \geq n, it is possible, and based on similar derivations as before, we can obtain:

Corollary 4: When m \geq n, an m\times n matrix independently and repeatedly sampled from any distribution p(x) with a mean of 0 and a variance of 1/m approximately satisfies \boldsymbol{W}^{\top}\boldsymbol{W}=\boldsymbol{I}.

So, if m > n, you only need to change the variance of the sampling distribution to 1/m. As for m < n, although there is no direct derivation, this practice can still be followed, as a reasonable strategy should be universal. Note that this modification is slightly different from the original design of Xavier initialization; it is a dual version of "LeCun initialization" (where the variance is 1/n), while the variance of Xavier initialization is 2/(m+n), which is an intuitive approach that averages forward and backward propagation. Here, we mainly consider forward propagation.

Some readers might still wonder: you only considered scenarios without activation functions. Even if the norm of \boldsymbol{y} is the same as \boldsymbol{x}, it will change after \boldsymbol{y} passes through an activation function. This is indeed the case, and at this point, one must analyze specific problems. For example, \tanh(x) \approx x when x is small, so Xavier initialization can be considered directly applicable to \tanh activation. For \text{relu}, one can assume that about half of the elements in \text{relu}(\boldsymbol{y}) will be set to zero, so the norm becomes approximately 1/\sqrt{2} of the original. To keep the norm unchanged, one can multiply \boldsymbol{W} by \sqrt{2}, meaning the initialization variance changes from 1/m to 2/m. This is the initialization strategy proposed by Kaiming He for \text{relu}.

Of course, in practice, it is difficult to adjust the variance perfectly for every activation function. Therefore, a more general approach is to add an operation similar to Layer Normalization directly after the activation function to explicitly restore the norm. This is where various Normalization techniques come into play (feel free to read my previous work "What Role Does BN Actually Play? An Analysis from Scratch").

A Brief Summary

This article primarily derives the conclusion that "any n\times n matrix with mean 0 and variance 1/n is close to an orthogonal matrix" from the premise that "any two random vectors in high-dimensional space are almost always orthogonal," thereby providing a geometric perspective on related initialization strategies. I believe this geometric perspective is more intuitive and easier to understand than a purely statistical one.

Reprinting please include the original address: https://kexue.fm/archives/7180

For more detailed reprinting matters, please refer to: "Scientific Space FAQ"