Can We Losslessly Scale Up a Transformer Model? (Part 1)
Looking at the title, readers might be puzzled — isn't everyone trying to shrink large models these days? Why are you trying to scale up a small model? Here's the background: it's generally true that bigger models with more data tend to perform better, but under limited compute, pretraining a large model from scratch takes far too long. If you also need several rounds of hyperparameter tuning on top of that, months can go by.
That's when a "poor man's trick" comes to mind (rich folks can ignore this): what if we first train a small model with the same number of layers, then scale it up and keep training? That way, once the small model has finished pretraining, its scaled-up weights can serve as a very strong initialization for the large model, which means fewer training steps are needed in the large-model stage, shortening the overall training time.
So, can a small model be losslessly scaled up into a large model? This post analyzes this question from a theoretical standpoint.
What "lossless" means
Some readers might think: of course this is possible — a large model's fitting capacity is certainly greater than a small model's. That's true, and from the perspective of fitting capacity alone, this is definitely achievable. But that's not the whole story of what we mean by "lossless scale-up" here. more
Taking BERT as an example, since pretraining is essentially an MLM task, "lossless scale-up" here means:
Can we transform a small model into a large model through some transformation, such that the output remains exactly unchanged?
By "transformation" here we mean a deterministic transformation applied to the weights, rather than continued training via gradient descent. By "output remains exactly unchanged" we mean that for the same input, the small model and the large model produce exactly identical predictions — that is, although the two models look different on the surface, they are mathematically the exact same function. This is why we call it "lossless scale-up." Since it's lossless, we can at least guarantee that the large model is no worse than the small model, so continuing to pretrain from this point should theoretically yield a positive gain. Whether this "small-then-large" pretraining scheme actually matches or beats training the large model from scratch from the start is a question that needs to be settled experimentally, and is not the focus of this post.
Intuitively, this kind of scale-up doesn't seem too hard — operations like "repeating" or "zero-padding" naturally scale up model weights. Indeed, this is the direction we explore. The difficulty lies in carefully analyzing what happens to every module of the model after scaling, to make sure the final result is truly lossless.
Attempts
Below we take "scaling a BERT model up by 2×" as our working example, to figure out the final form of the transformation. Here, "scaling up" refers only to enlarging the hidden vector dimension, without changing the number of layers or the number of heads in multi-head attention.
Embedding
First, the input layer is the embedding layer, so we need to tackle the embedding layer's scale-up first. This is also the simplest part: we simply double the dimensionality of every token's vector. The main choices are "repeating" and "zero-padding":
\begin{equation}\begin{array}{ll} \text{repeat:} & [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}\end{equation}
Both are candidate schemes, but intuitively, zero-padding introduces too many zeros, leading to excessive sparsity and repeated values, which hurts weight diversity. So we opt for the repeating scheme. That said, even within "repeating," there is more than one way to do it — for instance, $[x_1,x_2,x_3,x_4,x_1,x_2,x_3,x_4]$ is another possible scheme, but the later analysis of the attention layer will show that this latter scheme is not viable.
Beyond that, we generally also want the transformation to be orthogonal, which tends to maximize the model's stability. Specifically, the most basic property of an orthogonal transformation is that it preserves vector norms, so we adjust the final repeating transformation to:
\begin{equation}\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}\end{equation}
or written more compactly as $\tilde{x}_i = x_{\lceil i/2\rceil} / \sqrt{2}$, where $\lceil \cdot\rceil$ is the ceiling function. We call this "repeat, then divide by $\sqrt{2}$."
LayerNorm
The layer right after the embedding layer is LayerNorm. Before the transformation, LayerNorm computes:
\begin{equation}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}\end{equation}
After the transformation, we have:
\begin{equation}\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}\end{equation}
That is to say, the "subtract the mean, divide by the standard deviation" step automatically cancels out the factor $1/\sqrt{2}$, and the result is simply the direct repetition of the pre-scale-up result. If we also transform the parameter vector $\beta,\gamma$ according to formula $\eqref{eq:vt}$, the result becomes $\tilde{y}_i = y_{\lceil i/2\rceil} / \sqrt{2}$, consistent with the embedding layer's transformation result. Our goal is precisely to make the "net transformation" of every layer the same simple transformation: "repeat, then divide by $\sqrt{2}$."
FeedForward
In principle, we should analyze the attention layer next, but the FeedForward layer is relatively simpler, and analyzing it first will also help us understand the attention layer's transformation later. So let's 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:
\begin{equation} y_j = \mathcal{A}\left(\sum_{i=1}^d x_i w_{i,j} + b_j\right)\end{equation}
Here $\mathcal{A}(\cdot)$ is the activation function. Building on our previous experience, we try the following transformation:
\begin{equation}\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}\end{equation}
That is, we transform $b_j$ according to formula $\eqref{eq:vt}$, and for $w_{i,j}$ we try the following form of transformation:
\begin{equation}\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} \\ \hdashline 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} \\ \hdashline \vdots & \vdots & \vdots & \vdots & \ddots & \vdots & \vdots\\ \hdashline 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)\end{equation}
Here $D$ is the output dimension size; we assume that when the model is scaled up 2×, $D$ is also doubled. It's not hard to see that this transformation amounts to applying the transformation $\eqref{eq:vt}$ to both the row and column directions of the weight matrix $w_{i,j}$. In this case,
\begin{equation}\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}\end{equation}
This shows that the transformation $\eqref{eq:wt}$ achieves what we want for a plain linear layer — the scaled-up result is exactly "repeat, then divide by $\sqrt{2}$." However, this is not enough on its own, because the fully-connected layer also has an activation function $\mathcal{A}(\cdot)$, and the issue now is that $\mathcal{A}(x/\sqrt{2})$ is not necessarily equal to $\mathcal{A}(x)/\sqrt{2}$. If they're not equal, we can't make the overall transformation equivalent to "repeat, then divide by $\sqrt{2}$."
In fact, the GeLU activation used by BERT does not satisfy this identity; a linear activation (i.e., no activation function) obviously does satisfy it, and a common nonlinear activation function that does satisfy this identity is ReLU (and also LeakyReLU). So one straightforward fix is to switch the FeedForward layer's activation to ReLU. In fact, this is already a common choice in pretrained models — Baidu's ERNIE and Google's T5 both use ReLU as the FeedForward activation function.
So does that mean FeedForward layers with non-ReLU activations, like BERT's, are simply out of luck? Not quite, because the FeedForward layer is a composition of two fully-connected layers. We just need to divide by $\sqrt{2}$ one fewer time when transforming the first fully-connected layer, and divide by an extra $\sqrt{2}$ when transforming the second. Specifically, the first fully-connected weight becomes:
\begin{equation} \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}\end{equation}
In this case we get:
\begin{equation}\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) \end{equation}
The result here is a direct repetition of the original result, without dividing by $\sqrt{2}$. Given this, it suffices for the subsequent fully-connected layer to divide by one extra $\sqrt{2}$, i.e., the weight of the following fully-connected layer is transformed as:
\begin{equation} \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}\end{equation}
With this, the overall effect of the FeedForward layer becomes equivalent to "repeat, then divide by $\sqrt{2}$."
Attention
Now we come to the hardest nut to crack — the transformation of the attention layer. The attention layer first applies three linear layers to transform each input vector into $q,k,v$:
\begin{equation} 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)} \end{equation}
Based on our earlier analysis of the FeedForward layer, we know that in order to achieve the "repeat, then divide by $\sqrt{2}$" effect for all of $q,k,v$, we simply need to follow transformation $\eqref{eq:wt}$. But the attention layer isn't a plain fully-connected layer — after transforming, we need to check whether the attention matrix stays unchanged. Let's compute the inner product:
\begin{equation}\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\end{equation}
where $d'$ is the corresponding head_size. This result tells us that the above transformation preserves the inner product, and so should also preserve the attention matrix. But there's a trap here! For a model like T5, whose inner products aren't rescaled afterward, we're indeed done at this point. However, for a model like BERT, the inner product is followed by division by $\sqrt{d'}$ before the softmax. Once the model is scaled up, dividing by $\sqrt{d'}$ becomes dividing by $\sqrt{2d'}$, so preserving the inner product alone is no longer enough to keep the attention matrix unchanged. We also need to multiply the weights going into $q,k$ by an extra factor of $\sqrt[4]{2}$, so the final transformation should be:
\begin{equation}\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}\end{equation}
With this transformation, the attention matrix stays unchanged, and since $\tilde{v}_i = v_{\lceil i/2\rceil} / \sqrt{2}$, the final output is also $\tilde{o}_i = o_{\lceil i/2\rceil} / \sqrt{2}$.
The above analysis only covers a single attention head, but of course attention has multiple heads, and the outputs of the multiple heads are concatenated and then passed through a fully-connected layer. Since all heads are treated equally and independently, the conclusion above still basically holds, and the final fully-connected layer just needs to be transformed according to formula $\eqref{eq:wt}$ to preserve attention's transformation behavior. However, one effect of having multiple heads is that when we do the "repeating," it must be done locally, within each head.
Specifically, when implementing multi-head attention, we don't literally perform separate fully-connected operations for each head — instead we do one big fully-connected operation and then reshape. This lets us compare the reshape results of two different repeating schemes:
$$\begin{array}{c:c} [x_1,x_2,x_3,x_4,x_5,x_6] & [x_1,x_2,x_3,x_4,x_5,x_6] \\ \downarrow & \downarrow \\ [x_1,x_1,x_2,x_2,x_3,x_3,x_4,x_4,x_5,x_5,x_6,x_6] & [x_1,x_2,x_3,x_4,x_5,x_6,x_1,x_2,x_3,x_4,x_5,x_6] \\ \downarrow & \downarrow \\ \begin{pmatrix}x_1,x_1,x_2,x_2 \\ x_3,x_3,x_4,x_4 \\ x_5,x_5,x_6,x_6\end{pmatrix} & \begin{pmatrix}x_1,x_2,x_3,x_4 \\ x_5,x_6,x_1,x_2 \\ x_3,x_4,x_5,x_6\end{pmatrix} \\ \end{array}$$
Note that before scaling up, the reshape result is $\begin{pmatrix}x_1,x_2 \\ x_3,x_4 \\ x_5,x_6\end{pmatrix}$, so comparing the reshape results of the two different repeating schemes, we find that the second scheme's reshaped result is completely scrambled — it's not equivalent to repeating within each head separately. So we're forced to choose the first repeating scheme.
Output probability distribution
With the above analysis, we can scale the entire encoder up 2× while achieving the "repeat, then divide by $\sqrt{2}$" effect. What remains is the output part, i.e., converting the encoder's output vector into a probability distribution over tokens. There are a few cases to consider here.
For models like GPT and T5, the logits for the probability distribution are computed by multiplying the encoder output directly by the transpose of the embedding matrix (possibly with an added bias). Since the embedding matrix already embodies the "repeat, then divide by $\sqrt{2}$" operation, and the encoder's output is also "repeat, then divide by $\sqrt{2}$," the two combine and exactly cancel out. So from the perspective of the probability distribution, the output is completely unchanged.
BERT, however, has an extra fully-connected layer: it first applies a GeLU-activated fully-connected layer, and only then multiplies by the transpose of the embedding matrix and adds a bias term to get the logits. As discussed in the "FeedForward" section, a fully-connected layer with non-ReLU activation cannot achieve the "repeat, then divide by $\sqrt{2}$" effect — it can only achieve plain "repetition" via the transformation $\eqref{eq:wt-2}$. So in order to still get the "divide by $\sqrt{2}$" effect, the LayerNorm right after it needs to divide by an extra factor of $\sqrt{2}$ during its transformation.
Of course, if the activation is ReLU, then transforming according to formula $\eqref{eq:wt}$ leaves everything completely unchanged. Also, in cases like mT5, where the final logits transformation matrix isn't shared with the embedding layer, the output can also be kept completely unchanged simply by adjusting this final transformation matrix.
RoPE positional encoding
All the analysis above only applies to the case where the individual neurons are uncorrelated — that is, where any two components $x_i,x_j$ of a vector have no particular relationship to each other. But if the model uses "rotary position embedding (RoPE)," this assumption no longer holds, because RoPE operates on pairs of components — $[x_1,x_2]$ forms one group, $[x_3,x_4]$ forms another, and so on.
If we still repeat according to the earlier formula $\eqref{eq:vt}$, then after the transformation the groupings would become $[x_1,x_1]$ as one group, $[x_2,x_2]$ as another, and so on — no longer consistent with the original groupings, which would introduce significant errors. In this case, the repeating should also be done in pairs:
\begin{equation}\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:vt-2}\end{equation}
Of course, since the standard RoPE has no trainable weights and varies according to a fixed scheme, even repeating this way cannot fully guarantee an identical result. In other words, if RoPE is used, lossless scale-up generally cannot be fully achieved. Still, empirical tests show that scaling up this way does cause some performance loss for the corresponding RoFormer model, but not much, and it can be quickly recovered by continued training.
Conclusion
We can now confirm that for BERT, if the nonlinear activation is ReLU, then BERT can be losslessly scaled up directly; if the nonlinear activation is not ReLU, we can still achieve scale-up that is lossless in terms of MLM accuracy (in fact, with more fine-grained adjustments, fully lossless scale-up is also achievable, but the transformation for each layer ends up somewhat inconsistent and less elegant). For models like GPT and T5, regardless of the activation function used (including the GLU activation used by mT5, which can also be handled with appropriate customization), lossless scale-up can in fact always be achieved.
The transformations for scaling BERT's weights up by 2× are summarized below:
$$\begin{array}{l|l} \hline \text{Embedding} & \tilde{x}_i = \frac{1}{\sqrt{2}} x_{\lceil i/2\rceil} \\ \hline \text{LayerNorm} & \tilde{\beta}_i = \frac{1}{\sqrt{2}} \beta_{\lceil i/2\rceil},\quad \tilde{\gamma}_i = \frac{1}{\sqrt{2}} \gamma_{\lceil i/2\rceil} \\ \hline \text{Attention} & \begin{array}{l} \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)} \\ \tilde{w}_{i,j}^{(o)}=\frac{1}{2}w_{\lceil i/2\rceil,\lceil j/2\rceil}^{(o)},\quad \tilde{b}_j^{(o)}=\frac{1}{\sqrt{2}}b_{\lceil j/2\rceil}^{(o)} \end{array} \\ \hline \text{FeedForward} & \begin{array}{l} \tilde{w}_{i,j}^{(1)}=\frac{1}{\sqrt{2}}w_{\lceil i/2\rceil,\lceil j/2\rceil}^{(1)},\quad \tilde{b}_j^{(1)}=b_{\lceil j/2\rceil}^{(1)} \\ \tilde{w}_{i,j}^{(2)}=\frac{1}{2\sqrt{2}}w_{\lceil i/2\rceil,\lceil j/2\rceil}^{(2)},\quad \tilde{b}_j=\frac{1}{2}b_{\lceil j/2\rceil}^{(2)} \end{array} \\ \hline \text{output probability distribution} & \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} \\ \hline \end{array}$$
For other, slightly different models, one can perform a similar analysis following the same line of thinking. If RoPE is used, just switch the repeating scheme to formula $\eqref{eq:vt-2}$; if scaling up by a factor of $k$, just replace most instances of 2 in the table with $k$. In short, if attention has no scaling factor (no dividing by $\sqrt{d'}$) and the FeedForward activation is ReLU (or LeakyReLU), then the transformation for scaling up by a factor of $k$ is the simplest: just have every dimension of the weights "repeat $k$ times and divide by $\sqrt{k}$."
Summary
This post analyzed, from a mathematical perspective, the possibility of directly scaling up a Transformer model, and arrived at several viable transformations, establishing the feasibility of losslessly scaling up a Transformer model. This provides a point of reference for realizing progressive training of large models.
Translated automatically with claude-sonnet-5; all equations are reproduced verbatim from the source. Copyright remains with the original author.