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: Authors offer an alternative for masked LM pretraining that's more sample-efficient called replaced token detection. Their method basically replaces certain input tokens with alternatives which are sampled from a generator and train a discriminative model to determine whether its generated or real. The work shows empirical success getting better results than GPT with a fraction of the compute on GLUE and others.