Can We Losslessly Enlarge a Transformer Model? (Part 1)

Embedding

First, the input layer is the Embedding layer, so we must first solve the enlargement of the Embedding layer. This is one of the simplest parts, which involves enlarging the vector dimension of each token by 2x. There are two main operations: "repetition" and "zero-padding": \begin{array}{ll} \text{Repetition:} & [x_1, x_2, x_3, x_4] \to [x_1, x_1, x_2, x_2, x_3, x_3, x_4, x_4] \\ \text{Zero-padding:} & [x_1, x_2, x_3, x_4] \to [x_1, x_2, x_3, x_4, 0, 0, 0, 0] \end{array} Both schemes are candidates, but intuitively, zero-padding introduces too many zeros, leading to excessive sparsity and too many repetitions of the same value, which is detrimental to weight diversity. Therefore, we choose the repetition scheme. However, even within repetition, there is more than one way; for example, [x_1, x_2, x_3, x_4, x_1, x_2, x_3, x_4] is also a scheme. But subsequent analysis of the Attention layer shows that the latter scheme is undesirable.

In addition, we usually hope the transformation is orthogonal, which generally ensures the stability of the model to the greatest extent. Specifically, the most basic property of an orthogonal transformation is that it does not change the norm of the vector. Therefore, we adjust the final repetition transformation to: \begin{pmatrix}x_1 \\ x_2 \\ \vdots \\ x_d\end{pmatrix} \quad \to \quad \begin{pmatrix}\tilde{x}_1 \\ \tilde{x}_2 \\ \tilde{x}_3 \\ \tilde{x}_4 \\ \vdots \\ \tilde{x}_{2d-1} \\ \tilde{x}_{2d}\end{pmatrix} = \frac{1}{\sqrt{2}}\begin{pmatrix}x_1 \\ x_1 \\ x_2 \\ x_2 \\ \vdots \\ x_d \\ x_d \end{pmatrix} \label{eq:vt} Or simply denoted as \tilde{x}_i = x_{\lceil i/2\rceil} / \sqrt{2}, where \lceil \cdot \rceil is the ceiling operation. We call this "repeat then divide by \sqrt{2}."

LayerNorm

The next layer after Embedding is LayerNorm. Before the transformation, the LayerNorm operation is: y_i = \frac{x_i - \mu}{\sigma} \times \gamma_i + \beta_i \quad \mu = \frac{1}{d}\sum_{i=1}^d x_i \quad \sigma = \sqrt{\frac{1}{d}\sum_{i=1}^d (x_i-\mu)^2} After the transformation, we have: \begin{aligned} &\tilde{\mu} = \frac{1}{2d}\sum_{i=1}^{2d} \tilde{x}_i = \frac{1}{d}\sum_{i=1}^{d} \frac{x_i}{\sqrt{2}} = \frac{\mu}{\sqrt{2}} \\ &\tilde{\sigma} = \sqrt{\frac{1}{2d}\sum_{i=1}^{2d} (\tilde{x}_i-\tilde{\mu})^2} = \sqrt{\frac{1}{d}\sum_{i=1}^{d} \left(\frac{x_i}{\sqrt{2}}-\frac{\mu}{\sqrt{2}}\right)^2} = \frac{\sigma}{\sqrt{2}} \\ &\frac{\tilde{x}_i-\tilde{\mu}}{\tilde{\sigma}} = \frac{x_{\lceil i/2\rceil} / \sqrt{2} - \mu/\sqrt{2}}{\sigma/\sqrt{2}} = \frac{x_{\lceil i/2\rceil} - \mu}{\sigma} \end{aligned} This means that the "subtract mean and divide by standard deviation" step automatically cancels out the 1/\sqrt{2} factor. The result is a direct repetition of the pre-enlargement result. If we also transform the parameter vectors \beta, \gamma according to formula [eq:vt], the result will be \tilde{y}_i = y_{\lceil i/2\rceil} / \sqrt{2}, which is consistent with the transformation result of the Embedding layer. Our goal is to ensure as much as possible that the "net transformation" of each layer is the same simple transformation: "repeat then divide by \sqrt{2}."

FeedForward

Following the sequence, we should analyze the Attention layer next, but the FeedForward layer is relatively simpler, and its analysis results will help in understanding the Attention layer’s transformation. Therefore, we consider the FeedForward layer first.

The FeedForward layer is just a composition of two fully connected layers, so we only need to analyze a single fully connected layer: y_j = \mathcal{A}\left(\sum_{i=1}^d x_i w_{i,j} + b_j\right) where \mathcal{A}(\cdot) is the activation function. Based on previous experience, we try the following transformation: \tilde{w}_{i,j}=\frac{1}{2}w_{\lceil i/2\rceil,\lceil j/2\rceil},\quad \tilde{b}_j=\frac{1}{\sqrt{2}}b_{\lceil j/2\rceil} \label{eq:wt} That is, b_j is transformed according to formula [eq:vt], while for w_{i,j}, we attempt the following form: \begin{pmatrix}w_{1,1} & w_{1,2} & \cdots & w_{1,D} \\ w_{2,1} & w_{2,2} & \cdots & w_{2,D} \\ \vdots & \vdots & \ddots & \vdots \\ w_{d,1} & w_{d,2} & \cdots & w_{d,D}\end{pmatrix} \quad \to \quad \frac{1}{2}\left(\begin{array}{cc|cc|c|cc} w_{1,1} & w_{1,1} & w_{1,2} & w_{1,2} & \cdots & w_{1,D} & w_{1,D} \\ w_{1,1} & w_{1,1} & w_{1,2} & w_{1,2} & \cdots & w_{1,D} & w_{1,D} \\ \hline w_{2,1} & w_{2,1} & w_{2,2} & w_{2,2} & \cdots & w_{2,D} & w_{2,D} \\ w_{2,1} & w_{2,1} & w_{2,2} & w_{2,2} & \cdots & w_{2,D} & w_{2,D} \\ \hline \vdots & \vdots & \vdots & \vdots & \ddots & \vdots & \vdots\\ \hline w_{d,1} & w_{d,1} & w_{d,2} & w_{d,2} & \cdots & w_{d,D} & w_{d,D} \\ w_{d,1} & w_{d,1} & w_{d,2} & w_{d,2} & \cdots & w_{d,D} & w_{d,D}\end{array}\right) where D is the output dimension. We assume D also doubles. It is easy to see that this transformation applies [eq:vt] to both rows and columns of w_{i,j}. Then: \begin{aligned} \sum_{i=1}^{2d} \tilde{x}_i \tilde{w}_{i,j} + \tilde{b}_j =&\, 2\sum_{i=1}^d \frac{x_i}{\sqrt{2}} \frac{w_{i,\lceil j/2\rceil}}{2} + \frac{b_{\lceil j/2\rceil}}{\sqrt{2}} \\ =&\, \frac{1}{\sqrt{2}}\left(\sum_{i=1}^d x_i w_{i,\lceil j/2\rceil} + b_{\lceil j/2\rceil}\right) \end{aligned} This shows that transformation [eq:wt] for a linear layer satisfies our ideal goal—the enlarged result is "repeat then divide by \sqrt{2}." However, this is not enough because the layer has an activation function \mathcal{A}(\cdot). The problem is that \mathcal{A}(x/\sqrt{2}) may not equal \mathcal{A}(x)/\sqrt{2}.

In fact, the GeLU activation function used in BERT does not satisfy this identity. Linear activation functions obviously do, and a common non-linear activation function that satisfies this is the ReLU (including LeakyReLU) function. Thus, a direct solution is to use ReLU in the FeedForward layer. This is already a common choice for pre-trained models like Baidu’s Ernie and Google’s T5.

What about models like BERT with non-ReLU activations? We can still manage because the FeedForward layer is a composite of two layers. We can divide by one less \sqrt{2} in the first layer and one more in the second. Specifically, for the first layer: \tilde{w}_{i,j}=\frac{1}{\sqrt{2}}w_{\lceil i/2\rceil,\lceil j/2\rceil},\quad \tilde{b}_j=b_{\lceil j/2\rceil} \label{eq:wt-2} Then: \mathcal{A}\left(\sum_{i=1}^{2d} \tilde{x}_i \tilde{w}_{i,j} + \tilde{b}_j\right) = \mathcal{A}\left(\sum_{i=1}^d x_i w_{i,\lceil j/2\rceil} + b_{\lceil j/2\rceil}\right) The result is a direct repetition without the 1/\sqrt{2} factor. We then compensate in the second fully connected layer by dividing by an extra \sqrt{2}: \tilde{w}_{i,j}=\frac{1}{2\sqrt{2}}w_{\lceil i/2\rceil,\lceil j/2\rceil},\quad \tilde{b}_j=\frac{1}{2}b_{\lceil j/2\rceil} This makes the entire FeedForward layer equivalent to "repeat then divide by \sqrt{2}."

Attention

Now for the "hardest nut to crack"—the Attention layer. It first transforms input vectors into q, k, v: q_j = \sum_{i=1}^d x_i w_{i,j}^{(q)} + b_j^{(q)}, \quad k_j = \sum_{i=1}^d x_i w_{i,j}^{(k)} + b_j^{(k)}, \quad v_j = \sum_{i=1}^d x_i w_{i,j}^{(v)} + b_j^{(v)} To make q, k, v follow "repeat then divide by \sqrt{2}," we use transformation [eq:wt]. We then check the Attention matrix (inner product): \sum_{i=1}^{2d'} \tilde{q}_i \tilde{k}_i = 2\sum_{i=1}^{d'} \frac{q_i}{\sqrt{2}}\frac{k_i}{\sqrt{2}} = \sum_{i=1}^{d'} q_i k_i where d' is the head size. This keeps the inner product unchanged. However, there is a trap! Models like BERT divide the inner product by \sqrt{d'} before Softmax. After enlargement, d' becomes 2d', so we divide by \sqrt{2d'}. To keep the Attention matrix unchanged, we must scale q and k weights by an additional \sqrt[4]{2} each. The final transformations are: \begin{aligned} &\tilde{w}_{i,j}^{(q)}=\frac{\sqrt[4]{2}}{2}w_{\lceil i/2\rceil,\lceil j/2\rceil}^{(q)},\quad \tilde{b}_j^{(q)}=\frac{\sqrt[4]{2}}{\sqrt{2}}b_{\lceil j/2\rceil}^{(q)}\\ &\tilde{w}_{i,j}^{(k)}=\frac{\sqrt[4]{2}}{2}w_{\lceil i/2\rceil,\lceil j/2\rceil}^{(k)},\quad \tilde{b}_j^{(k)}=\frac{\sqrt[4]{2}}{\sqrt{2}}b_{\lceil j/2\rceil}^{(k)}\\ &\tilde{w}_{i,j}^{(v)}=\frac{1}{2}w_{\lceil i/2\rceil,\lceil j/2\rceil}^{(v)},\quad \tilde{b}_j^{(v)}=\frac{1}{\sqrt{2}}b_{\lceil j/2\rceil}^{(v)} \end{aligned} This keeps the Attention matrix unchanged, and since \tilde{v}_i = v_{\lceil i/2\rceil} / \sqrt{2}, the output is also \tilde{o}_i = o_{\lceil i/2\rceil} / \sqrt{2}.

For multi-head attention, the outputs are concatenated and passed through a final linear layer. Since heads are independent, the logic holds, but we must repeat "locally" within each head. If we repeat globally, the reshape operation for multi-head attention will scramble the dimensions. Thus, we must use the interleaved repetition scheme: [x_1, x_1, x_2, x_2, \dots].

Output Probability Distribution

Finally, the Encoder output is converted to token probabilities. For GPT or T5, which multiply the output by the transpose of the Embedding matrix, the 1/\sqrt{2} factors from both cancel out (due to the doubling of the inner product sum), keeping the logits unchanged.

For BERT, there is an additional GeLU fully connected layer before the Embedding transpose. As discussed, non-ReLU layers require transformation [eq:wt-2] to achieve "repetition," and the subsequent LayerNorm must then divide by an extra \sqrt{2} to maintain the "divide by \sqrt{2}" effect.

RoPE Position Encoding

The previous analysis assumes neurons are independent. However, Rotary Position Embedding (RoPE) groups components in pairs ([x_1, x_2], [x_3, x_4], \dots). Standard repetition [x_1, x_1, x_2, x_2] breaks these pairs. Instead, we must repeat in pairs: \begin{array}{c} [x_1,x_2,x_3,x_4,\cdots,x_{d-1},x_d] \\ \downarrow\\ \frac{1}{\sqrt{2}}[x_1,x_2,x_1,x_2,x_3,x_4,x_3,x_4,\cdots,x_{d-1},x_d,x_{d-1},x_d] \end{array} \label{eq:rope-rep} While RoPE is not trainable and its progression is fixed, this grouped repetition minimizes performance loss, which can be recovered quickly through further training.