ELECTRA: Pre-training Text Encoders as Discriminators Rather Than Generators

1 INTRODUCTION . Current state-of-the-art representation learning methods for language can be viewed as learning denoising autoencoders ( Vincent et al. , 2008 ) . They select a small subset of the unlabeled input sequence ( typically 15 % ) , mask the identities of those tokens ( e.g. , BERT ; Devlin et al . ( 2019 ) ) or attention to those tokens ( e.g. , XLNet ; Yang et al . ( 2019 ) ) , and then train the network to recover the original input . While more effective than conventional language-model pre-training due to learning bidirectional representations , these masked language modeling ( MLM ) approaches incur a substantial compute cost because the network only learns from 15 % of the tokens per example . As an alternative , we propose replaced token detection , a pre-training task in which the model learns to distinguish real input tokens from plausible but synthetically generated replacements . Instead of masking , our method corrupts the input by replacing some tokens with samples from a proposal distribution , which is typically the output of a small masked language model . This corruption procedure solves a mismatch in BERT ( although not in XLNet ) where the network sees artificial [ MASK ] tokens during pre-training but not when being fine-tuned on downstream tasks . We then pre-train the network as a discriminator that predicts for every token whether it is an original or a replacement . In contrast , MLM trains the network as a generator that predicts the original identities of the corrupted tokens . A key advantage of our discriminative task is that the model learns from all input tokens instead of just the small masked-out subset , making it more computationally efficient . Although our approach is reminiscent of training the discriminator of a GAN , our method is not adversarial in that the generator producing corrupted tokens is trained with maximum likelihood due to the difficulty of applying GANs to text ( Caccia et al. , 2018 ) . We call our approach ELECTRA1 for “ Efficiently Learning an Encoder that Classifies Token Replacements Accurately. ” As in prior work , we apply it to pre-train Transformer text encoders ( Vaswani et al. , 2017 ) that can be fine-tuned on downstream tasks . Through a series of ablations , we show that learning from all input positions causes ELECTRA to train much faster than BERT . We also show ELECTRA achieves higher accuracy on downstream tasks when fully trained . Most current pre-training methods require large amounts of compute to be effective , raising concerns about their cost and accessibility . Since pre-training with more compute almost always results in better downstream accuracies , we argue an important consideration for pre-training methods should be compute efficiency as well as absolute downstream performance . From this viewpoint , we train ELECTRA models of various sizes and evaluate their downstream performance vs. their compute requirement . In particular , we run experiments on the GLUE natural language understanding benchmark ( Wang et al. , 2019 ) and SQuAD question answering benchmark ( Rajpurkar et al. , 2016 ) . ELECTRA substantially outperforms MLM-based methods such as BERT and XLNet given the same model size , data , and compute ( see Figure 1 ) . For example , we build an ELECTRA-Small model that can be trained on 1 GPU in 4 days.2 ELECTRA-Small outperforms a comparably small BERT model by 5 points on GLUE , and even outperforms the much larger GPT model ( Radford et al. , 2018 ) . Our approach also works well at large scale , where we train an ELECTRA-Large model that performs comparably to RoBERTa ( Liu et al. , 2019 ) and XLNet ( Yang et al. , 2019 ) , despite having fewer parameters and using 1/4 of the compute for training . Training ELECTRA-Large further results in an even stronger model that outperforms ALBERT ( Lan et al. , 2019 ) on GLUE and sets a new state-of-the-art for SQuAD 2.0 . Taken together , our results indicate that the discriminative task of distinguishing real data from challenging negative samples is more compute-efficient and parameter-efficient than existing generative approaches for language representation learning . 2 METHOD . We first describe the replaced token detection pre-training task ; see Figure 2 for an overview . We suggest and evaluate several modeling improvements for this method in Section 3.2 . 1Code and pre-trained weights will be released at https : //github.com/google-research/ electra 2It has 1/20th the parameters and requires 1/135th the pre-training compute of BERT-Large . Published as a conference paper at ICLR 2020 artist artist [ MASK ] v artist sample Our approach trains two neural networks , a generator G and a discriminator D. Each one primarily consists of an encoder ( e.g. , a Transformer network ) that maps a sequence on input tokens x = [ x1 , ... , xn ] into a sequence of contextualized vector representations h ( x ) = [ h1 , ... , hn ] . For a given position t , ( in our case only positions where xt = [ MASK ] ) , the generator outputs a probability for generating a particular token xt with a softmax layer : pG ( xt|x ) = exp ( e ( xt ) ThG ( x ) t ) / ∑ x′ exp ( e ( x′ ) ThG ( x ) t ) where e denotes token embeddings . For a given position t , the discriminator predicts whether the token xt is “ real , ” i.e. , that it comes from the data rather than the generator distribution , with a sigmoid output layer : D ( x , t ) = sigmoid ( wThD ( x ) t ) The generator is trained to perform masked language modeling ( MLM ) . Given an input x = [ x1 , x2 , ... , xn ] , MLM first select a random set of positions ( integers between 1 and n ) to mask out m = [ m1 , ... , mk ] .3 The tokens in the selected positions are replaced with a [ MASK ] token : we denote this as xmasked = REPLACE ( x , m , [ MASK ] ) . The generator then learns to predict the original identities of the masked-out tokens . The discriminator is trained to distinguish tokens in the data from tokens that have been replaced by generator samples . More specifically , we create a corrupted example xcorrupt by replacing the masked-out tokens with generator samples and train the discriminator to predict which tokens in xcorrupt match the original input x . Formally , model inputs are constructed according to mi ∼ unif { 1 , n } for i = 1 to k xmasked = REPLACE ( x , m , [ MASK ] ) x̂i ∼ pG ( xi|xmasked ) for i ∈m xcorrupt = REPLACE ( x , m , x̂ ) and the loss functions are LMLM ( x , θG ) = E ( ∑ i∈m − log pG ( xi|xmasked ) ) LDisc ( x , θD ) = E ( n∑ t=1 −1 ( xcorruptt = xt ) logD ( xcorrupt , t ) − 1 ( x corrupt t 6= xt ) log ( 1−D ( xcorrupt , t ) ) ) Although similar to the training objective of a GAN , there are several key differences . First , if the generator happens to generate the correct token , that token is considered “ real ” instead of “ fake ” ; we found this formulation to moderately improve results on downstream tasks . More importantly , the generator is trained with maximum likelihood rather than being trained adversarially to fool the discriminator . Adversarially training the generator is challenging because it is impossible to backpropagate through sampling from the generator . Although we experimented circumventing this issue 3Typically k = d0.15ne , i.e. , 15 % of the tokens are masked out . by using reinforcement learning to train the generator ( see Appendix F ) , this performed worse than maximum-likelihood training . Lastly , we do not supply the generator with a noise vector as input , as is typical with a GAN . We minimize the combined loss min θG , θD ∑ x∈X LMLM ( x , θG ) + λLDisc ( x , θD ) over a large corpus X of raw text . We approximate the expectations in the losses with a single sample . We don ’ t back-propagate the discriminator loss through the generator ( indeed , we can ’ t because of the sampling step ) . After pre-training , we throw out the generator and fine-tune the discriminator on downstream tasks . 3 EXPERIMENTS . 3.1 EXPERIMENTAL SETUP . We evaluate on the General Language Understanding Evaluation ( GLUE ) benchmark ( Wang et al. , 2019 ) and Stanford Question Answering ( SQuAD ) dataset ( Rajpurkar et al. , 2016 ) . GLUE contains a variety of tasks covering textual entailment ( RTE and MNLI ) question-answer entailment ( QNLI ) , paraphrase ( MRPC ) , question paraphrase ( QQP ) , textual similarity ( STS ) , sentiment ( SST ) , and linguistic acceptability ( CoLA ) . See Appendix C for more details on the GLUE tasks . Our evaluation metrics are Spearman correlation for STS , Matthews correlation for CoLA , and accuracy for the other GLUE tasks ; we generally report the average score over all tasks . For SQuAD , we evaluate on versions 1.1 , in which models select the span of text answering a question , and 2.0 , in which some questions are unanswerable by the passage . We use the standard evaluation metrics of Exact-Match ( EM ) and F1 scores . For most experiments we pre-train on the same data as BERT , which consists of 3.3 Billion tokens from Wikipedia and BooksCorpus ( Zhu et al. , 2015 ) . However , for our Large model we pre-trained on the data used for XLNet ( Yang et al. , 2019 ) , which extends the BERT dataset to 33B tokens by including data from ClueWeb ( Callan et al. , 2009 ) , CommonCrawl , and Gigaword ( Parker et al. , 2011 ) . All of the pre-training and evaluation is on English data , although we think it would be interesting to apply our methods to multilingual data in the future . Our model architecture and most hyperparameters are the same as BERT ’ s . For fine-tuning on GLUE , we add simple linear classifiers on top of ELECTRA . For SQuAD , we add the questionanswering module from XLNet on top of ELECTRA , which is slightly more sophisticated than BERT ’ s in that it jointly rather than independently predicts the start and end positions and has a “ answerability ” classifier added for SQuAD 2.0 . Some of our evaluation datasets are small , which means accuracies of fine-tuned models can vary substantially depending on the random seed . We therefore report the median of 10 fine-tuning runs from the same pre-trained checkpoint for each result . Unless stated otherwise , results are on the dev set . See the appendix for further training details and hyperparameter values . 3.2 MODEL EXTENSIONS . We improve our method by proposing and evaluating several extensions to the model . Unless stated otherwise , these experiments use the same model size and training data as BERT-Base . Weight Sharing We propose improving the efficiency of the pre-training by sharing weights between the generator and discriminator . If the generator and discriminator are the same size , all of the transformer weights can be tied . However , we found it to be more efficient to have a small generator , in which case we only share the embeddings ( both the token and positional embeddings ) of the generator and discriminator . In this case we use embeddings the size of the discriminator ’ s hidden states.4 The “ input ” and “ output ” token embeddings of the generator are always tied as in BERT . We compare the weight tying strategies when the generator is the same size as the discriminator . We train these models for 500k steps . GLUE scores are 83.6 for no weight tying , 84.3 for tying token embeddings , and 84.4 for tying all weights . We hypothesize that ELECTRA benefits from 4We add linear layers to the generator to project the embeddings into generator-hidden-sized representations . tied token embeddings because masked language modeling is particularly effective at learning these representations : while the discriminator only updates tokens that are present in the input or are sampled by the generator , the generator ’ s softmax over the vocabulary densely updates all token embeddings . On the other hand , tying all encoder weights caused little improvement while incurring the significant disadvantage of requiring the generator and discriminator to be the same size . Based on these findings , we use tied embeddings for further experiments in this paper . Smaller Generators If the generator and discriminator are the same size , training ELECTRA would take around twice as much compute per step as training only with masked language modeling . We suggest using a smaller generator to reduce this factor . Specifically , we make models smaller by decreasing the layer sizes while keeping the other hyperparameters constant . We also explore using an extremely simple “ unigram ” generator that samples fake tokens according their frequency in the train corpus . GLUE scores for differently-sized generators and discriminators are shown in the left of Figure 3 . All models are trained for 500k steps , which puts the smaller generators at a disadvantage in terms of compute because they require less compute per training step . Nevertheless , we find that models work best with generators 1/4-1/2 the size of the discriminator . We speculate that having too strong of a generator may pose a too-challenging task for the discriminator , preventing it from learning as effectively . In particular , the discriminator may have to use many of its parameters modeling the generator rather than the actual data distribution . Further experiments in this paper use the best generator size found for the given discriminator size . Training Algorithms Lastly , we explore other training algorithms for ELECTRA , although these did not end up improving results . The proposed training objective jointly trains the generator and discriminator . We experiment with instead using the following two-stage training procedure : 1 . Train only the generator with LMLM for n steps . 2 . Initialize the weights of the discriminator with the weights of the generator . Then train the discriminator with LDisc for n steps , keeping the generator ’ s weights frozen . Note that the weight initialization in this procedure requires having the same size for the generator and discriminator . We found that without the weight initialization the discriminator would sometimes fail to learn at all beyond the majority class , perhaps because the generator started so far ahead of the discriminator . Joint training on the other hand naturally provides a curriculum for the discriminator where the generator starts off weak but gets better throughout training . We also explored training the generator adversarially as in a GAN , using reinforcement learning to accommodate the discrete operations of sampling from the generator . See Appendix F for details . Results are shown in the right of Figure 3 . During two-stage training , downstream task performance notably improves after the switch from the generative to the discriminative objective , but does not end up outscoring joint training . Although still outperforming BERT , we found adversarial training to underperform maximum-likelihood training . Further analysis suggests the gap is caused by two problems with adversarial training . First , the adversarial generator is simply worse at masked language modeling ; it achieves 58 % accuracy at masked language modeling compared to 65 % accuracy for an MLE-trained one . We believe the worse accuracy is mainly due to the poor sample efficiency of reinforcement learning when working in the large action space of generating text . Secondly , the adversarially trained generator produces a low-entropy output distribution where most of the probability mass is on a single token , which means there is not much diversity in the generator samples . Both of these problems have been observed in GANs for text in prior work ( Caccia et al. , 2018 ) .

Summary: The paper proposed a novel sample-efficient pretraining task. One inefficiency of BERT is that only 15% tokens are used for training in each example. The paper introduced a generator+discriminator framework to optimize the utility of training examples. The generator task is the MLM which predicts the masked word. The author adds a discriminator to further learn from the example by classifying each word to be either generated or original. In this way, more words can be used. This method looks as only adding the discrimination task after BERT pretraining task. But, the authors later show that the best GLUE scores can be obtained only when both generator and discriminator are co-trained. Moreover, the adversarial ELECTRA perform worse. All these observations are interesting. It will be helpful if the authors provide more empirical analysis why the adversarial ELECTRA perform worse or failed. Is it because the GAN is hard to train or the adversarial task doesn't fit the pretraining?