Rethinking the Dimension-Averaging Strategy for Non-Square Matrices in Initialization

初始化方法中非方阵的维度平均策略思考 苏剑林 · 2021-10-18 · Machine Learning model optimization gradient

In posts such as Understanding Model Parameter Initialization from a Geometric Perspective] and A Brief Discussion on Initialization, Parameterization, and Normalization in Transformers], we have discussed model initialization methods. The basic idea is: if a $n\times n$ square matrix is initialized i.i.d. with mean 0 and variance $1/n$, then it approximates an orthogonal matrix, so that the second moment (or variance) of the data stays roughly unchanged as it propagates through.

What about a non-square matrix of shape $m\times n$? The common approach (Xavier initialization) is to jointly consider forward and backward propagation, and thus initialize i.i.d. with mean 0 and variance $2/(m+n)$. But this averaging is largely a "rule of thumb," and in this post we'll explore whether there's a better averaging scheme.

A Quick Recap

Xavier initialization considers a fully-connected layer as follows (let the number of input nodes be $m$ and the number of output nodes be $n$):

\begin{equation} y_j = b_j + \sum_i x_i w_{i,j}\end{equation}more

Here $b_j$ is generally initialized to 0, and the initialization mean of $w_{i,j}$ is also usually 0. In A Brief Discussion on Initialization, Parameterization, and Normalization in Transformers] we already computed

\begin{equation} \mathbb{E}[y_j^2] = \sum_{i} \mathbb{E}[x_i^2] \mathbb{E}[w_{i,j}^2]= m\mathbb{E}[x_i^2]\mathbb{E}[w_{i,j}^2]\end{equation}

So, in order to keep the second moment unchanged, we set the initialization variance of $w_{i,j}$ to $1/m$ (since the mean is 0, the variance equals the second moment).

However, this derivation only considers the forward pass. We also need the model to have reasonable gradients, meaning the backward pass should remain stable as well. Suppose the model's loss function is $l$; by the chain rule we have

\begin{equation}\frac{\partial l}{\partial x_i} = \sum_j \frac{\partial l}{\partial y_j} \frac{\partial y_j}{\partial x_i}=\sum_j \frac{\partial l}{\partial y_j} w_{i,j}\end{equation}

Note that here the sum is over $j$, with the summation dimension being $n$, so under the same assumptions we get

\begin{equation} \mathbb{E}\left[\left(\frac{\partial l}{\partial x_i}\right)^2\right] = \sum_{j} \mathbb{E}\left[\left(\frac{\partial l}{\partial y_j}\right)^2\right] \mathbb{E}[w_{i,j}^2]= n \mathbb{E}\left[\left(\frac{\partial l}{\partial y_j}\right)^2\right]\mathbb{E}[w_{i,j}^2]\end{equation}

So, to keep the second moment unchanged during backpropagation, we should set the initialization variance of $w_{i,j}$ to $1/n$.

One gives $1/m$, the other gives $1/n$, and these conflict when $m\neq n$. But both are equally important, so Xavier initialization simply averages the two dimensions and initializes with variance $2/(m+n)$.

Geometric Mean

Now let's consider two fully-connected layers composed together (ignoring bias terms for now):

\begin{equation} y = xW_1 W_2 \end{equation}

where $x\in\mathbb{R}^m,W_1\in\mathbb{R}^{m\times n},W_2\in\mathbb{R}^{n\times m}$, meaning the input has dimension $m$, is transformed to dimension $n$, and then transformed back to dimension $m$. A similar operation appears, for instance, in BERT's FFN layer (though the FFN layer has an extra activation function in the middle).

Based on the stability of the forward pass, we should initialize $W_1$ with variance $1/m$ and $W_2$ with variance $1/n$. But what if we require $W_1$ and $W_2$ to be initialized with the same variance? Then, clearly, in order to keep the variance of $x,y$ unchanged, $W_1,W_2$ both need to be initialized with variance $1/\sqrt{mn}$. If we consider the backward pass as well, we arrive at the same result.

This gives us a new dimension-averaging strategy: the geometric mean, $\sqrt{mn}$. With this averaging strategy, when multiple layers are composed and the input/output dimensions are unchanged, the variance stays constant (regardless of whether we're looking at the forward or backward pass). By contrast, with the algebraic mean $(m+n)/2$, if we assume $m < n$, then according to $(m+n)^2/4\geq mn$, the variance would shrink during forward/backward propagation.

Quadratic Mean

Another angle to think about this is as a dual minimization problem: suppose we choose variance $t$. During the forward pass we want $(mt-1)^2$ to be as small as possible, and during the backward pass we want $(nt-1)^2$ to be as small as possible. So, taking both into account:

\begin{equation}(mt-1)^2 + (nt-1)^2 \end{equation}

This expression attains its minimum when $t=(m+n)/(m^2+n^2)$, which gives us an averaging scheme based on a "quadratic fraction": $(m^2+n^2)/(m+n)$.

It's easy to prove that:

\begin{equation}\frac{m^2+n^2}{m+n} \geq \frac{m+n}{2}\geq \sqrt{mn}\end{equation}

From the derivation, the quadratic mean on the left aims to keep the variance nearly unchanged at every single step of forward and backward propagation, so we can regard it as a local optimum. The geometric mean on the right, on the other hand, aims to keep the variance of the "very first input" and the "very last output" nearly unchanged, so in some sense it can be regarded as a global optimum. The algebraic mean in the middle is a solution sitting somewhere between the global and local optima.

Seen this way, perhaps Xavier initialization's "rule of thumb" algebraic mean isn't such a bad choice after all — a kind of golden mean.

Summary

This post has been a brief reflection on dimension-averaging schemes for non-square matrices in initialization methods. It seems that people have generally taken the default algebraic mean for granted without much question, but here I've derived the possibility of different averaging strategies from two different angles. As for which averaging strategy works better in practice, I haven't run careful experiments myself — readers who are interested are welcome to try it out. Of course, it's also possible that, given the many optimization tricks in use today, the default initialization scheme already works well enough that there's no real need to fine-tune it further.