ALBERT: A Lite BERT for Self-supervised Learning of Language Representations

1 INTRODUCTION . Full network pre-training ( Dai & Le , 2015 ; Radford et al. , 2018 ; Devlin et al. , 2019 ; Howard & Ruder , 2018 ) has led to a series of breakthroughs in language representation learning . Many nontrivial NLP tasks , including those that have limited training data , have greatly benefited from these pre-trained models . One of the most compelling signs of these breakthroughs is the evolution of machine performance on a reading comprehension task designed for middle and high-school English exams in China , the RACE test ( Lai et al. , 2017 ) : the paper that originally describes the task and formulates the modeling challenge reports then state-of-the-art machine accuracy at 44.1 % ; the latest published result reports their model performance at 83.2 % ( Liu et al. , 2019 ) ; the work we present here pushes it even higher to 89.4 % , a stunning 45.3 % improvement that is mainly attributable to our current ability to build high-performance pretrained language representations . Evidence from these improvements reveals that a large network is of crucial importance for achieving state-of-the-art performance ( Devlin et al. , 2019 ; Radford et al. , 2019 ) . It has become common practice to pre-train large models and distill them down to smaller ones ( Sun et al. , 2019 ; Turc et al. , 2019 ) for real applications . Given the importance of model size , we ask : Is having better NLP models as easy as having larger models ? An obstacle to answering this question is the memory limitations of available hardware . Given that current state-of-the-art models often have hundreds of millions or even billions of parameters , it is easy to hit these limitations as we try to scale our models . Training speed can also be significantly hampered in distributed training , as the communication overhead is directly proportional to the number of parameters in the model . Existing solutions to the aforementioned problems include model parallelization ( Shazeer et al. , 2018 ; Shoeybi et al. , 2019 ) and clever memory management ( Chen et al. , 2016 ; Gomez et al. , 2017 ) . ∗Work done as an intern at Google Research , driving data processing and downstream task evaluations . These solutions address the memory limitation problem , but not the communication overhead . In this paper , we address all of the aforementioned problems , by designing A Lite BERT ( ALBERT ) architecture that has significantly fewer parameters than a traditional BERT architecture . ALBERT incorporates two parameter reduction techniques that lift the major obstacles in scaling pre-trained models . The first one is a factorized embedding parameterization . By decomposing the large vocabulary embedding matrix into two small matrices , we separate the size of the hidden layers from the size of vocabulary embedding . This separation makes it easier to grow the hidden size without significantly increasing the parameter size of the vocabulary embeddings . The second technique is cross-layer parameter sharing . This technique prevents the parameter from growing with the depth of the network . Both techniques significantly reduce the number of parameters for BERT without seriously hurting performance , thus improving parameter-efficiency . An ALBERT configuration similar to BERT-large has 18x fewer parameters and can be trained about 1.7x faster . The parameter reduction techniques also act as a form of regularization that stabilizes the training and helps with generalization . To further improve the performance of ALBERT , we also introduce a self-supervised loss for sentence-order prediction ( SOP ) . SOP primary focuses on inter-sentence coherence and is designed to address the ineffectiveness ( Yang et al. , 2019 ; Liu et al. , 2019 ) of the next sentence prediction ( NSP ) loss proposed in the original BERT . As a result of these design decisions , we are able to scale up to much larger ALBERT configurations that still have fewer parameters than BERT-large but achieve significantly better performance . We establish new state-of-the-art results on the well-known GLUE , SQuAD , and RACE benchmarks for natural language understanding . Specifically , we push the RACE accuracy to 89.4 % , the GLUE benchmark to 89.4 , and the F1 score of SQuAD 2.0 to 92.2 . 2 RELATED WORK . 2.1 SCALING UP REPRESENTATION LEARNING FOR NATURAL LANGUAGE . Learning representations of natural language has been shown to be useful for a wide range of NLP tasks and has been widely adopted ( Mikolov et al. , 2013 ; Le & Mikolov , 2014 ; Dai & Le , 2015 ; Peters et al. , 2018 ; Devlin et al. , 2019 ; Radford et al. , 2018 ; 2019 ) . One of the most significant changes in the last two years is the shift from pre-training word embeddings , whether standard ( Mikolov et al. , 2013 ; Pennington et al. , 2014 ) or contextualized ( McCann et al. , 2017 ; Peters et al. , 2018 ) , to full-network pre-training followed by task-specific fine-tuning ( Dai & Le , 2015 ; Radford et al. , 2018 ; Devlin et al. , 2019 ) . In this line of work , it is often shown that larger model size improves performance . For example , Devlin et al . ( 2019 ) show that across three selected natural language understanding tasks , using larger hidden size , more hidden layers , and more attention heads always leads to better performance . However , they stop at a hidden size of 1024 , presumably because of the model size and computation cost problems . It is difficult to experiment with large models due to computational constraints , especially in terms of GPU/TPU memory limitations . Given that current state-of-the-art models often have hundreds of millions or even billions of parameters , we can easily hit memory limits . To address this issue , Chen et al . ( 2016 ) propose a method called gradient checkpointing to reduce the memory requirement to be sublinear at the cost of an extra forward pass . Gomez et al . ( 2017 ) propose a way to reconstruct each layer ’ s activations from the next layer so that they do not need to store the intermediate activations . Both methods reduce the memory consumption at the cost of speed . Raffel et al . ( 2019 ) proposed to use model parallelization to train a giant model . In contrast , our parameter-reduction techniques reduce memory consumption and increase training speed . 2.2 CROSS-LAYER PARAMETER SHARING . The idea of sharing parameters across layers has been previously explored with the Transformer architecture ( Vaswani et al. , 2017 ) , but this prior work has focused on training for standard encoderdecoder tasks rather than the pretraining/finetuning setting . Different from our observations , Dehghani et al . ( 2018 ) show that networks with cross-layer parameter sharing ( Universal Transformer , UT ) get better performance on language modeling and subject-verb agreement than the standard transformer . Very recently , Bai et al . ( 2019 ) propose a Deep Equilibrium Model ( DQE ) for transformer networks and show that DQE can reach an equilibrium point for which the input embedding and the output embedding of a certain layer stay the same . Our observations show that our embeddings are oscillating rather than converging . Hao et al . ( 2019 ) combine a parameter-sharing transformer with the standard one , which further increases the number of parameters of the standard transformer . 2.3 SENTENCE ORDERING OBJECTIVES . ALBERT uses a pretraining loss based on predicting the ordering of two consecutive segments of text . Several researchers have experimented with pretraining objectives that similarly relate to discourse coherence . Coherence and cohesion in discourse have been widely studied and many phenomena have been identified that connect neighboring text segments ( Hobbs , 1979 ; Halliday & Hasan , 1976 ; Grosz et al. , 1995 ) . Most objectives found effective in practice are quite simple . Skipthought ( Kiros et al. , 2015 ) and FastSent ( Hill et al. , 2016 ) sentence embeddings are learned by using an encoding of a sentence to predict words in neighboring sentences . Other objectives for sentence embedding learning include predicting future sentences rather than only neighbors ( Gan et al. , 2017 ) and predicting explicit discourse markers ( Jernite et al. , 2017 ; Nie et al. , 2019 ) . Our loss is most similar to the sentence ordering objective of Jernite et al . ( 2017 ) , where sentence embeddings are learned in order to determine the ordering of two consecutive sentences . Unlike most of the above work , however , our loss is defined on textual segments rather than sentences . BERT ( Devlin et al. , 2019 ) uses a loss based on predicting whether the second segment in a pair has been swapped with a segment from another document . We compare to this loss in our experiments and find that sentence ordering is a more challenging pretraining task and more useful for certain downstream tasks . Concurrently to our work , Wang et al . ( 2019 ) also try to predict the order of two consecutive segments of text , but they combine it with the original next sentence prediction in a three-way classification task rather than empirically comparing the two . 3 THE ELEMENTS OF ALBERT . In this section , we present the design decisions for ALBERT and provide quantified comparisons against corresponding configurations of the original BERT architecture ( Devlin et al. , 2019 ) . 3.1 MODEL ARCHITECTURE CHOICES . The backbone of the ALBERT architecture is similar to BERT in that it uses a transformer encoder ( Vaswani et al. , 2017 ) with GELU nonlinearities ( Hendrycks & Gimpel , 2016 ) . We follow the BERT notation conventions and denote the vocabulary embedding size as E , the number of encoder layers as L , and the hidden size as H . Following Devlin et al . ( 2019 ) , we set the feed-forward/filter size to be 4H and the number of attention heads to be H/64 . There are three main contributions that ALBERT makes over the design choices of BERT . Factorized embedding parameterization . In BERT , as well as subsequent modeling improvements such as XLNet ( Yang et al. , 2019 ) and RoBERTa ( Liu et al. , 2019 ) , the WordPiece embedding size E is tied with the hidden layer size H , i.e. , E ≡ H . This decision appears suboptimal for both modeling and practical reasons , as follows . From a modeling perspective , WordPiece embeddings are meant to learn context-independent representations , whereas hidden-layer embeddings are meant to learn context-dependent representations . As experiments with context length indicate ( Liu et al. , 2019 ) , the power of BERT-like representations comes from the use of context to provide the signal for learning such context-dependent representations . As such , untying the WordPiece embedding size E from the hidden layer size H allows us to make a more efficient usage of the total model parameters as informed by modeling needs , which dictate that H E. From a practical perspective , natural language processing usually require the vocabulary size V to be large.1 If E ≡ H , then increasing H increases the size of the embedding matrix , which has size 1Similar to BERT , all the experiments in this paper use a vocabulary size V of 30,000 . V ×E . This can easily result in a model with billions of parameters , most of which are only updated sparsely during training . Therefore , for ALBERT we use a factorization of the embedding parameters , decomposing them into two smaller matrices . Instead of projecting the one-hot vectors directly into the hidden space of size H , we first project them into a lower dimensional embedding space of size E , and then project it to the hidden space . By using this decomposition , we reduce the embedding parameters from O ( V × H ) to O ( V × E + E × H ) . This parameter reduction is significant when H E. We choose to use the same E for all word pieces because they are much more evenly distributed across documents compared to whole-word embedding , where having different embedding size ( Grave et al . ( 2017 ) ; Baevski & Auli ( 2018 ) ; Dai et al . ( 2019 ) ) for different words is important . Cross-layer parameter sharing . For ALBERT , we propose cross-layer parameter sharing as another way to improve parameter efficiency . There are multiple ways to share parameters , e.g. , only sharing feed-forward network ( FFN ) parameters across layers , or only sharing attention parameters . The default decision for ALBERT is to share all parameters across layers . All our experiments use this default decision unless otherwise specified . We compare this design decision against other strategies in our experiments in Sec . 4.5 . Similar strategies have been explored by Dehghani et al . ( 2018 ) ( Universal Transformer , UT ) and Bai et al . ( 2019 ) ( Deep Equilibrium Models , DQE ) for Transformer networks . Different from our observations , Dehghani et al . ( 2018 ) show that UT outperforms a vanilla Transformer . Bai et al . ( 2019 ) show that their DQEs reach an equilibrium point for which the input and output embedding of a certain layer stay the same . Our measurement on the L2 distances and cosine similarity show that our embeddings are oscillating rather than converging . Figure 1 shows the L2 distances and cosine similarity of the input and output embeddings for each layer , using BERT-large and ALBERT-large configurations ( see Table 1 ) . We observe that the transitions from layer to layer are much smoother for ALBERT than for BERT . These results show that weight-sharing has an effect on stabilizing network parameters . Although there is a drop for both metrics compared to BERT , they nevertheless do not converge to 0 even after 24 layers . This shows that the solution space for ALBERT parameters is very different from the one found by DQE . Inter-sentence coherence loss . In addition to the masked language modeling ( MLM ) loss ( Devlin et al. , 2019 ) , BERT uses an additional loss called next-sentence prediction ( NSP ) . NSP is a binary classification loss for predicting whether two segments appear consecutively in the original text , as follows : positive examples are created by taking consecutive segments from the training corpus ; negative examples are created by pairing segments from different documents ; positive and negative examples are sampled with equal probability . The NSP objective was designed to improve performance on downstream tasks , such as natural language inference , that require reasoning about the relationship between sentence pairs . However , subsequent studies ( Yang et al. , 2019 ; Liu et al. , 2019 ) found NSP ’ s impact unreliable and decided to eliminate it , a decision supported by an improvement in downstream task performance across several tasks . We conjecture that the main reason behind NSP ’ s ineffectiveness is its lack of difficulty as a task , as compared to MLM . As formulated , NSP conflates topic prediction and coherence prediction in a single task2 . However , topic prediction is easier to learn compared to coherence prediction , and also overlaps more with what is learned using the MLM loss . We maintain that inter-sentence modeling is an important aspect of language understanding , but we propose a loss based primarily on coherence . That is , for ALBERT , we use a sentence-order prediction ( SOP ) loss , which avoids topic prediction and instead focuses on modeling inter-sentence coherence . The SOP loss uses as positive examples the same technique as BERT ( two consecutive segments from the same document ) , and as negative examples the same two consecutive segments but with their order swapped . This forces the model to learn finer-grained distinctions about discourse-level coherence properties . As we show in Sec . 4.6 , it turns out that NSP can not solve the SOP task at all ( i.e. , it ends up learning the easier topic-prediction signal , and performs at randombaseline level on the SOP task ) , while SOP can solve the NSP task to a reasonable degree , presumably based on analyzing misaligned coherence cues . As a result , ALBERT models consistently improve downstream task performance for multi-sentence encoding tasks .

Summary: This paper investigates improving upon BERT by reducing complexity in terms of free parameters and memory footprint as well as computation steps. They propose 2 strategies for doing this: 1) Splitting the embedding matrix into two smaller matrices (going from V x A to V x B + B x A where B <<<< A); 2) layer-wise parameter sharing. They also utilize sentence order prediction to help with training. These coupled with a bunch of other choices such as using the lamb optimizer, certain hyperparameters etc help show dramatic empirical gains across the board on a wide variety of NLP/NLU tasks.