Discarding Constraints, Enhancing Models: Improving ALBERT Performance with One Line of Code

Low-rank Decomposition

First is the Embedding layer. Taking the Chinese version of BERT-base as an example, the total number of tokens is approximately 20,000, and the embedding dimension is 768. Therefore, the total number of parameters in the Embedding layer is about 15 million, accounting for roughly 1/6 of the total parameters. The first part ALBERT targets is the Embedding layer. It reduces the Embedding dimension to 128 and then transforms it back to 768 through a 128 \times 768 matrix. This reduces the Embedding layer parameters to 1/6 of the original size. This is known as low-rank decomposition.

Parameter Sharing

Next is the Transformer section. In Transformer-based architectures like BERT, the core consists of modules composed of self-attention, layer normalization, and fully connected layers (1D convolutions with a kernel size of 1). We call these “Transformer blocks.” A BERT model is a stack of multiple Transformer blocks. The figure on the left below is a schematic of BERT-base, which stacks 12 Transformer blocks.

Note that in BERT’s design, the input and output shapes of each Transformer block are the same, which means it is reasonable to use the output of the current block as the input for the same block. This tells us that the same block can be reused iteratively instead of using a new block for every layer. ALBERT employs the simplest and most direct solution: all layers share the same Transformer block (as shown in the figure on the right above)! Consequently, in ALBERT-base, the parameters for the Transformer blocks are reduced to 1/12 of those in BERT-base.

Brief Commentary

In addition to the two points above, a significant difference between ALBERT and BERT is that during the pre-training phase, the NSP (Next Sentence Prediction) task was replaced by the SOP (Sentence-Order Prediction) task. However, this is not part of the model architecture and is not the focus of this article; readers can find relevant materials on their own.

Overall, ALBERT is a model designed to reduce parameter counts, with the hope that this reduction acts as a form of regularization, thereby lowering the risk of overfitting and improving final performance. But did the results meet the authors’ expectations? Looking at its “track record,” ALBERT refreshed the GLUE leaderboard with its largest version upon release, so it seems to have met expectations. However, ALBERT is not always ideal, nor is it the “small model” we might imagine.

For any model, we are primarily concerned with two indicators: speed and performance. From the two diagrams of BERT and ALBERT above, we can see that during the inference phase (forward propagation), ALBERT is essentially no different from BERT. Therefore, under the same specifications (e.g., both being the “base” version), the inference speeds of ALBERT and BERT are the same. Strictly speaking, ALBERT might even be slightly slower because of the additional matrix operation in the Embedding section. In other words, ALBERT does not bring an improvement in inference speed!

So, between ALBERT and BERT of the same specification, which performs better? The ALBERT paper already provides the answer: for versions up to “large,” ALBERT performs worse than BERT. Only at the “xlarge” and “xxlarge” versions does ALBERT begin to outperform BERT. However, RoBERTa-style pre-training has mitigated BERT’s shortcomings, so the only ALBERT version that can be said to consistently outperform BERT/RoBERTa is “xxlarge.” Yet, ALBERT-xxlarge is such a massive model that it is difficult for most people to run.

Thus, it can basically be said: (under the premise that most people can run it) at the same inference speed, ALBERT performs worse; at the same performance level, ALBERT is slower.

What about the training phase? One point not mentioned yet is that ALBERT’s parameter-sharing design has a strong regularizing effect, so ALBERT removes dropout. Parameter sharing and the removal of dropout do save some GPU memory and increase training speed, but my evaluation suggests the improvement is only around 10%–20%. That is to say, even if ALBERT’s parameters are reduced to 1/10 or less of BERT’s, it does not mean its memory usage drops to 1/10, nor does it mean training speed increases tenfold. On the contrary, it only offers a slight improvement.

Performance Evaluation

Without further ado, let the results speak. Four tasks were selected for testing. To ensure reproducibility, each experiment was run three times, and the table shows the average of the three results. The “unshared” version refers to the model with the parameter-sharing constraint removed. The “Training Speed” column refers to the training time per epoch on a single TITAN RTX card, provided for relative comparison.

The experiments were conducted using bert4keras. For the unshared version, one only needs to load the ALBERT weights during build_transformer_model and set model=’albert_unshared’. This is the “one line of code” mentioned in the title.

First is a relatively simple text sentiment classification task.

\begin{array}{c|c|c|c|c} \hline \text{Model} & \text{Valid Set} & \text{Speed} & \text{Metric after 1st Epoch} & \text{Test Set} \\ \hline \text{small\_unshared} & 94.66\% & 38s & 90.75\% & 94.35\% \\ \text{small} & 94.57\% & 33s & 91.02\% & 94.52\% \\ \hline \text{tiny\_unshared} & 94.02\% & 23s & 88.09\% & 94.13\% \\ \text{tiny} & 94.14\% & 20s & 90.18\% & 93.78\% \\ \hline \end{array}

After removing parameter sharing, the training time increased slightly, as expected. As for model performance, results were mixed. Considering that the accuracy for this task is already quite high, it might not highlight the gap between models, so we continue testing with more complex tasks.

Next, we try CLUE’s IFLYTEK long text classification. The results are as follows:

\begin{array}{c|c|c|c} \hline \text{Model} & \text{Dev Set} & \text{Speed} & \text{Metric after 1st Epoch} \\ \hline \text{small\_unshared} & 57.73\% & 27s & 49.35\% \\ \text{small} & 57.14\% & 24s & 48.21\% \\ \hline \text{tiny\_unshared} & 55.91\% & 16s & 47.89\% \\ \text{tiny} & 56.42\% & 14s & 43.84\% \\ \hline \end{array}

Here, the advantage of the unshared version begins to show, primarily reflected in faster overall convergence (see the metric after the first epoch). The best performance of the “small” version is significantly better. The “tiny” version’s best performance was slightly worse, but with fine-tuning of the learning rate, the tiny_unshared version can actually outperform the tiny version (but that introduces too many variables; the table shows results with strictly controlled variables).

Then we try a more comprehensive task: Information Extraction. The results are as follows:

\begin{array}{c|c|c|c} \hline \text{Model} & \text{Dev Set} & \text{Speed} & \text{Metric after 1st Epoch} \\ \hline \text{small\_unshared} & 77.89\% & 375s & 61.11\% \\ \text{small} & 77.69\% & 335s & 46.58\% \\ \hline \text{tiny\_unshared} & 76.44\% & 235s & 49.74\% \\ \text{tiny} & 75.94\% & 215s & 31.66\% \\ \hline \end{array}

As can be seen, in more comprehensive and complex tasks, the unshared version of the model consistently outperforms the original model of the same scale.

The last one is Reading Comprehension QA using seq2seq. The results are as follows:

\begin{array}{c|c|c|c} \hline \text{Model} & \text{Dev Set} & \text{Speed} & \text{Metric after 1st Epoch} \\ \hline \text{small\_unshared} & 68.80\% & 607s & 57.02\% \\ \text{small} & 66.66\% & 582s & 50.93\% \\ \hline \text{tiny\_unshared} & 66.15\% & 455s & 48.64\% \\ \text{tiny} & 63.41\% & 443s & 37.47\% \\ \hline \end{array}

The main purpose of this task is to test the model’s text generation capability. It can be seen that in this task, the unshared version significantly outperforms the original model, and even the unshared tiny version approaches the performance of the original small model.

Analysis and Reflection

The models above were experiments on ALBERT tiny/small. In fact, experiments were also conducted on the “base” version, and the conclusions were basically consistent with the tiny and small versions. However, because experiments on the base version (and naturally the large and xlarge versions) take too long, they were not completed in full (nor repeated three times), so they are not included here. But overall, the results for the tiny/small versions are representative.

The above experiments indicate that ALBERT with parameter sharing removed performs at a level that is equal to or better than the original ALBERT in downstream tasks. This suggests that for many NLP tasks, parameter sharing might not be a good constraint. Readers might wonder: “Why does ALBERT with parameter sharing start to outperform BERT without parameter sharing when the model scale reaches xlarge or xxlarge?” I will attempt to provide an explanation.

Theoretically, BERT’s methods for preventing overfitting include dropout and weight decay. ALBERT also uses weight decay, but dropout is absent in ALBERT. Therefore, we can think from the perspective of dropout. Many experiments have shown that dropout is indeed an effective strategy for reducing overfitting risk, but existing experimental models are generally far smaller than BERT-xlarge or BERT-xxlarge. Thus, the effectiveness of dropout in ultra-large models remains debatable. In fact, dropout suffers from inconsistency between training and inference—that is, “strictly speaking, the training model and the prediction model are not the same model.” I personally feel that as models become larger and deeper, this inconsistency is further amplified. Therefore, I believe dropout is not an effective method for preventing overfitting in ultra-large models. What if we just remove dropout in BERT? That doesn’t work well either, because without dropout, there are few means to suppress BERT’s overfitting; with such a large parameter count, overfitting would be severe. ALBERT removes dropout and introduces implicit regularization through parameter sharing, allowing the model to grow larger and deeper without degrading, and even performing better.

Conversely, for ALBERT’s parameter sharing to perform well, the model must be sufficiently large and deep. Therefore, if we are using the base, small, or even tiny versions, we should not use parameter sharing. For small models, parameter sharing is an unnecessary restriction on the model’s representational capacity. In these cases, removing parameter sharing leads to better performance.