Understanding Model Parameter Initialization from a Geometric Perspective

从几何视角来理解模型参数的初始化策略 苏剑林 · 2020-01-16 · Mathematics model probability geometry optimization

For complex models, parameter initialization is especially important. A poor initialization is often not just a matter of degraded model performance—it can mean the model simply fails to train at all or never converges. A common adaptive initialization strategy in deep learning is Xavier initialization, which constructs the initial weights by sampling randomly from a normal distribution $\mathcal{N}\left(0,\frac{2}{fan_{in} + fan_{out}}\right)$, where $fan_{in}$ is the input dimension and $fan_{out}$ is the output dimension. Other initialization strategies are basically similar, differing only in their underlying assumptions, which leads to slightly different final forms.

The standard derivation of initialization strategies is based on probability and statistics: roughly, one assumes that the input data has mean 0 and variance 1, expects the output data to also have mean 0 and variance 1, and then derives the mean and variance conditions that the initial transformation must satisfy. There's nothing wrong with this process in principle, but to me it still doesn't feel intuitive enough, and the derivation relies on a fair number of assumptions. This post aims to understand model initialization from a geometric perspective instead, offering a more intuitive derivation.

Orthogonality, Almost for Free

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

Corollary 1: In high-dimensional space, any two random vectors are almost always perpendicular to each other.

In fact, Corollary 1 is precisely the starting point of the entire geometric perspective of this post! A further corollary of it is:

Corollary 2: If we randomly draw $n^2$ numbers from $\mathcal{N}(0, 1/n)$ to form a $n\times n$ matrix, this matrix is approximately orthogonal, and 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)  # 矩阵乘以自身的转置
print(X)  # 看看是否接近单位阵
print(np.square(X - np.eye(n)).mean())  # 计算与单位阵的mse

I suspect most readers, on first seeing Corollary 2, find it somewhat surprising. An orthogonal matrix is one satisfying $\boldsymbol{W}^{\top}\boldsymbol{W}=\boldsymbol{I}$, i.e., its inverse equals its transpose. For a general matrix, computing the inverse is vastly harder than computing the transpose, so intuitively "inverse = transpose" feels like it should be a very stringent condition. Yet Corollary 2 tells us that a randomly sampled matrix is already close to orthogonal—which does feel a bit counterintuitive. I remember being quite surprised myself when I first realized this.

It's Actually Not That Hard to Understand

However, once we get used to the fact stated in Corollary 1—that in high-dimensional space any two random vectors are almost perpendicular—we can quickly understand and derive this result. For a quick derivation, let's first consider the standard normal distribution $\mathcal{N}(0,1)$, noting that Corollary 1 requires the sampled directions to be uniform, and the standard normal distribution indeed satisfies this. If we sample a $n\times n$ matrix from $\mathcal{N}(0,1)$, we can view it as $n$ vectors of dimension $n$. Since these $n$ vectors are all random vectors, they are naturally close to orthogonal to one another.

Of course, pairwise orthogonality alone doesn't make an orthogonal matrix, because an orthogonal matrix also requires each vector to have unit norm, and we have $\mathbb{E}_{x\sim \mathcal{N}(0,1)}\left[x^2\right]=1$, which means the norm of a $n$-dimensional vector sampled from $\mathcal{N}(0,1)$ is approximately $\sqrt{n}$. So to get closer to orthogonality, we also need to divide each element by $\sqrt{n}$, which is equivalent to changing the sampling variance from 1 to $1/n$.

Moreover, the sampling distribution doesn't have to be normal—a uniform distribution $U\left[-\sqrt{3/n}, \sqrt{3/n}\right]$ works just as well. In fact, we have:

Corollary 3: A $n\times n$ matrix obtained by independently and repeatedly sampling from any distribution $p(x)$ with mean 0 and variance $1/n$ is approximately orthogonal.

We can understand Corollary 3 from a more mathematical angle: suppose $\boldsymbol{x}=(x_1,x_2,\dots,x_n),\boldsymbol{y}=(y_1,y_2,\dots,y_n)$ are all sampled from $p(x)$, then we have

\begin{equation}\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}\end{equation}

and

\begin{equation}\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}\end{equation}

so any two vectors are approximately orthonormal, and hence the sampled matrix is approximately orthogonal.

Now We Can Talk About Initialization

After all this discussion of orthogonal matrices, it's really just laying the groundwork for understanding the geometric meaning of initialization methods. If readers still remember their linear algebra, they'll recall that the key significance of an orthogonal matrix is that it preserves vector norms under transformation. In mathematical terms, if $\boldsymbol{W}\in \mathbb{R}^{n\times n}$ is an orthogonal matrix and $\boldsymbol{x}\in\mathbb{R}^n$ is an arbitrary vector, then the norm of $\boldsymbol{x}$ equals the norm of $\boldsymbol{W}\boldsymbol{x}$:

\begin{equation}\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\end{equation}

Consider a fully connected layer:

\begin{equation}\boldsymbol{y}=\boldsymbol{W}\boldsymbol{x} + \boldsymbol{b}\end{equation}

Deep learning models are fundamentally nested compositions of fully connected layers, so to prevent the model's final output from "exploding" or "collapsing" right at initialization, one idea is to have the model preserve vector norms at initialization.

This idea naturally leads to an initialization strategy: "initialize $\boldsymbol{b}$ to all zeros, and initialize $\boldsymbol{W}$ as a random orthogonal matrix." And Corollary 2 already tells us that a $n\times n$ matrix sampled from $\mathcal{N}(0, 1/n)$ is already close to orthogonal, so we can initialize $\boldsymbol{W}$ by sampling from $\mathcal{N}(0, 1/n)$. This is exactly Xavier initialization—some frameworks call it Glorot initialization, since the author's name is Xavier Glorot. Also, the sampling distribution doesn't have to be $\mathcal{N}(0, 1/n)$; as Corollary 3 states, you can sample from any distribution with mean 0 and variance $1/n$.

The discussion above assumed the input and output dimensions are both $n$. What if the input has dimension $n$ and the output has dimension $m$? In that case $\boldsymbol{W}\in\mathbb{R}^{m\times n}$, and the condition for preserving the norm of $\boldsymbol{W}\boldsymbol{x}$ is still $\boldsymbol{W}^{\top}\boldsymbol{W}=\boldsymbol{I}$. However, when $m < n$, this is impossible; when $m \geq n$, it can hold, and following a similar derivation as before, we get:

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

So, if $m > n$, we just need to change the sampling variance to $1/m$. As for the case $m < n$, although there's no direct derivation, we can still adopt the same approach—after all, a reasonable strategy ought to be universal. Note that this adjustment differs somewhat from the original design of Xavier initialization: it's actually the dual version of "LeCun initialization" (whose variance is $1/n$), whereas Xavier initialization uses a variance of $2/(m+n)$, which averages the intuitions from forward propagation and backpropagation. Here we're mainly considering forward propagation.

Some readers might object: you've only considered the case without an activation function. Even if the norm of $\boldsymbol{y}$ matches that of $\boldsymbol{x}$, once $\boldsymbol{y}$ passes through an activation function, things change. This is indeed the case, and here one really has to analyze it problem by problem. For instance, $\tanh(x)$ satisfies $\tanh(x)\approx x$ when $x$ is small, so Xavier initialization can be considered directly applicable to $\tanh$ activation. Another example: for $\text{relu}$, we can assume roughly half the elements of $\text{relu}(\boldsymbol{y})$ get zeroed out, so the norm becomes roughly $1/\sqrt{2}$ times the original—and to keep the norm unchanged, we can scale $\boldsymbol{W}$ by $\sqrt{2}$, i.e., change the initialization variance from $1/m$ to $2/m$. This is precisely the initialization strategy proposed by Kaiming He for $\text{relu}$.

Of course, it's genuinely hard to fine-tune the variance adjustment for every single activation function, so a more general approach is to simply add an operation similar to Layer Normalization right after the activation function, explicitly restoring the norm. This is where various Normalization tricks come into play~ (Feel free to continue reading an earlier post, What Does BatchNorm Actually Do? A Speculative Analysis.)

A Brief Summary

This post has mainly derived, from the fact that "in high-dimensional space, any two random vectors are almost always perpendicular," the conclusion that "any $n\times n$ matrix with mean 0 and variance $1/n$ is approximately orthogonal," and from there offered a geometric perspective on related initialization strategies. I daresay this geometric viewpoint is more intuitive and easier to grasp than the purely statistical one.