ALBERT: A Lite BERT for Self-supervised Learning of Language Representation
"ALBERT: A Lite BERT for Self-supervised Learning of Language Representations" is a paper by Zhenzhong Lan and others from Google Research. Published in 2019, the paper presents ALBERT, a new model architecture that is a lighter and more memory-efficient variant of BERT, a transformer-based model for NLP tasks.
ALBERT introduces two key changes to the BERT architecture to reduce the model size and increase training speed without compromising performance:
-
Factorized embedding parameterization: This is a method that separates the size of the hidden layers from the size of the vocabulary embeddings. This is done by using two separate matrices for the embedding layer: one for the token embeddings and one for the hidden layers. The token embeddings are projected to the dimension of the hidden layers through a linear transformation. This way, the model can have a large vocabulary of embeddings with a relatively small embedding size, and then project these embeddings into the larger dimension required by the hidden layers. The model parameters are thus factorized into two smaller matrices, reducing the total number of parameters.
-
Cross-layer parameter sharing: This is a technique that shares parameters across the multiple layers in the transformer. In BERT, each transformer layer has separate parameters, while in ALBERT, all layers share the same parameters. This reduces the model size and also the number of computations required during training. There are different ways to implement cross-layer parameter sharing, such as sharing all parameters, sharing feed-forward network parameters, or sharing attention parameters.
The structure of the ALBERT transformer layer is similar to that of the original Transformer model. Given the output of the layer (l) as (H^l), the output of the self-attention sub-layer as (A^l), and the output of the feed-forward network as (F^l), the output of each layer is computed as:
[ H^l = \text{LayerNorm}(H^{l-1} + A^l) ]
[ F^l = \text{LayerNorm}(H^l + \text{FFN}(H^l)) ]
But remember that in ALBERT, the parameters for computing (A^l) and (F^l) are shared across all layers.
ALBERT uses the same training objectives as BERT: masked language modeling and next sentence prediction. However, the authors later introduced a sentence-order prediction task to replace the next sentence prediction task, as it was found to be more effective.
The main findings of the paper are that ALBERT performs comparably to much larger BERT models while being significantly smaller and faster to train. This is largely due to the two main architectural changes: factorized embedding parameterization and cross-layer parameter sharing.
The implications of the ALBERT paper are significant. It introduces a new approach to building transformer models that are more memory-efficient and faster to train, making them more accessible for researchers and developers with limited resources. It also contributes to our understanding of how to design effective architectures for large-scale language models. However, as with all models, it's important to consider the limitations and specific requirements of the task at hand when deciding whether to use ALBERT. For example, while ALBERT is very powerful, it may be overkill for simple tasks or tasks where the training data is very different from the pretraining data.
Additionally, while ALBERT reduces the model size, the computational resources required to train ALBERT are still substantial. This highlights the ongoing challenge in the field of AI research of balancing performance with resource efficiency and accessibility. Finally, as with all powerful AI models, it's important to consider the ethical implications of its use. For example, the ability of ALBERT to generate fluent text could be misused to produce misleading or harmful content. This underscores the importance of using such models responsibly and in a way that benefits society.
The ALBERT model has been widely adopted in the NLP community and has inspired further research into efficient model architectures and the use of parameter sharing in deep learning. This work demonstrates that it's possible to achieve high performance on NLP tasks with models that are significantly smaller and more efficient than previous state-of-the-art models. However, as with all AI research, it's important to continue pushing the boundaries of what's possible while also considering the broader implications of the technology.
Architecture
The main innovation in ALBERT's model architecture is the factorization of the embedding matrix into two smaller matrices.
In the original BERT model, the token embeddings are of the same size as the hidden states. So, if we have a vocabulary size (V), a hidden size (H), and (L) layers, the size of the token embedding matrix is (V \times H). This means that the number of parameters in the embedding layer scales with the size of the hidden layers, which can be quite large.
In ALBERT, the authors propose to separate the size of the token embeddings from the size of the hidden layers. They introduce an additional projection layer that maps the token embeddings to the size of the hidden layers. So, instead of having a single (V \times H) matrix, they have a (V \times E) matrix for the token embeddings, and an (E \times H) matrix for the projection, where (E) is the embedding size and is typically much smaller than (H).
So, the token embeddings are first looked up in the (V \times E) matrix, resulting in a (N \times E) matrix for a sequence of (N) tokens. This is then projected to the (H)-dimensional space using the (E \times H) matrix, resulting in a (N \times H) matrix that can be fed into the transformer layers.
This factorization of the embedding matrix reduces the number of parameters and makes the model more efficient. The smaller (E) dimension also means that the model can handle a larger vocabulary without significantly increasing the number of parameters.
This architectural change is one of the key reasons why ALBERT is more parameter-efficient than BERT. The other main change in ALBERT is the use of cross-layer parameter sharing, which further reduces the number of parameters and increases training efficiency.
If (E) is less than (H), and the projection is linear, how is any information being added?
That's a great question. When the embedding size (E) is less than the hidden size (H), and a linear projection is used to project the embeddings into the larger hidden space, no new information is technically added to the model. The projection is essentially stretching the (E)-dimensional embeddings into a (H)-dimensional space.
The purpose of this factorized embedding parameterization is not to add new information, but to reduce the number of parameters in the model. By having a smaller (E), the model can handle a larger vocabulary without significantly increasing the number of parameters. This makes the model more efficient and faster to train.
It's worth noting that while no new information is added by the projection, the higher-dimensional space may allow the model to learn more complex representations in the subsequent layers {at the cost of more paramters in the model from the (E \times H) matrix}. The transformer layers can potentially learn to use the additional dimensions to capture more complex patterns in the data.
However, it's also possible that some information could be lost or distorted during the projection, especially if (H) is much larger than (E). This is a trade-off that the designers of ALBERT chose to make in order to achieve a smaller and more efficient model. The impressive performance of ALBERT on a range of NLP tasks suggests that this trade-off was successful, but it's possible that for certain tasks or datasets, a different trade-off might be more effective.
Tags: BERT, ALBERT, 2019