In articles such as "Understanding Model Parameter Initialization Strategies from a Geometric Perspective" and "A Brief Discussion on Initialization, Parameterization, and Normalization of Transformer", we have discussed model initialization methods. The general idea is: if an n \times n square matrix is initialized with independent and identically distributed (i.i.d.) values with a mean of 0 and a variance of 1/n, it approximates an orthogonal matrix, ensuring that the second moment (or variance) of the data remains roughly constant during propagation.
But what if it is an m \times n non-square matrix? The common approach (Xavier initialization) is to consider both forward and backward propagation, thus using an i.i.d. initialization with a mean of 0 and a variance of 2/(m+n). However, this "averaging" is somewhat heuristic. This article explores whether there are better averaging schemes.
Basic Review
Xavier initialization considers the following fully connected layer (assuming m input nodes and n output nodes): y_j = b_j + \sum_i x_i w_{i,j} where b_j is generally initialized to 0, and the mean of w_{i,j} is also generally 0. As calculated in "A Brief Discussion on Initialization, Parameterization, and Normalization of Transformer", we have: \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] Therefore, to maintain the second moment, we set the initialization variance of w_{i,j} to 1/m (when the mean is 0, the variance equals the second moment).
However, this derivation only considers forward propagation. We also need the model to have reasonable gradients, which requires stability during backpropagation. Let the loss function of the model be l. According to the chain rule, we have: \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} Note that the summation is now over j, and the dimension of the sum is n. Under the same assumptions, we have: \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] Thus, to keep the second moment of backpropagation constant, we set the initialization variance of w_{i,j} to 1/n.
One is 1/m and the other is 1/n. When m \neq n, there is a conflict. Since both are equally important, Xavier initialization simply averages the two dimensions and uses 2/(m+n) as the variance for initialization.
Geometric Mean
Now let us consider two composite fully connected layers (temporarily ignoring the bias terms): y = xW_1 W_2 where x \in \mathbb{R}^m, W_1 \in \mathbb{R}^{m \times n}, W_2 \in \mathbb{R}^{n \times m}. That is, the input is m-dimensional, transformed to n dimensions, and then transformed back to m dimensions. Similar operations occur, for example, in the FFN layer of BERT (though the FFN layer includes an activation function in the middle).
Based on the stability of forward propagation, we should initialize W_1 with a variance of 1/m and W_2 with a variance of 1/n. However, if we require W_1 and W_2 to be initialized with the same variance, then clearly, to keep the variance of x and y the same, both W_1 and W_2 need to be initialized with a distribution having a variance of 1/\sqrt{mn}. The result is the same when considering backpropagation.
In this way, we derive a new dimension averaging strategy: the geometric mean \sqrt{mn}. Through this strategy, we can ensure that in a multi-layer composite network, if the input and output dimensions remain the same, the variance remains constant (regardless of forward or backward propagation). If we used the arithmetic mean (m+n)/2, and assuming m < n, then since (m+n)^2/4 \geq mn, the variance would shrink during forward/backward propagation.
Quadratic Mean
Another perspective is to treat this as a dual minimization problem: suppose the chosen variance is t. In forward propagation, we want (mt-1)^2 to be as small as possible, and in backpropagation, we want (nt-1)^2 to be as small as possible. Therefore, we consider: (mt-1)^2 + (nt-1)^2 The above expression reaches its minimum when t = (m+n)/(m^2+n^2). This yields an averaging scheme based on a quadratic fraction: (m^2+n^2)/(m+n).
It is easy to prove that: \frac{m^2+n^2}{m+n} \geq \frac{m+n}{2} \geq \sqrt{mn} From the derivation process, the quadratic mean on the left aims to keep the variance of each step of forward and backward propagation as constant as possible; thus, it can be considered a local optimum. The geometric mean on the right aims to keep the variance of the "initial input" and "final output" as constant as possible; thus, it can be considered a global optimum in some sense. The arithmetic mean in the middle is a solution that lies between the global and local optima.
From this perspective, it seems that the "heuristic" arithmetic mean of Xavier initialization might be a reasonable "Middle Way" choice?
Summary
This article briefly reflects on dimension averaging schemes for non-square matrices in initialization methods. For a long time, it seems that few have questioned the default arithmetic mean. I have derived the possibilities of different averaging strategies from two different perspectives. As for which strategy is better, I have not conducted detailed experiments; interested readers may try them out themselves. Of course, it is also possible that under current optimization strategies, the default initialization scheme already works well enough, making detailed tuning unnecessary.
Original Address: https://kexue.fm/archives/8725
For more details regarding reposting, please refer to: "Scientific Space FAQ"