[Paper Review] ELECTRA: Pre-training Text Encoders as Discriminators Rather Than Generators

Kevin B. Clark, Minh-Thang Luong+2 more|arXiv (Cornell University)|Mar 23, 2020

Topic Modeling48 references541 citations

TL;DR

ELECTRA introduces replaced token detection, a discriminative pre-training task where a generator creates plausible token replacements and a discriminator learns to detect which tokens were replaced. This yields higher downstream performance with substantially less compute than MLM-based methods like BERT.

ABSTRACT

Masked language modeling (MLM) pre-training methods such as BERT corrupt the input by replacing some tokens with [MASK] and then train a model to reconstruct the original tokens. While they produce good results when transferred to downstream NLP tasks, they generally require large amounts of compute to be effective. As an alternative, we propose a more sample-efficient pre-training task called replaced token detection. Instead of masking the input, our approach corrupts it by replacing some tokens with plausible alternatives sampled from a small generator network. Then, instead of training a model that predicts the original identities of the corrupted tokens, we train a discriminative model that predicts whether each token in the corrupted input was replaced by a generator sample or not. Thorough experiments demonstrate this new pre-training task is more efficient than MLM because the task is defined over all input tokens rather than just the small subset that was masked out. As a result, the contextual representations learned by our approach substantially outperform the ones learned by BERT given the same model size, data, and compute. The gains are particularly strong for small models; for example, we train a model on one GPU for 4 days that outperforms GPT (trained using 30x more compute) on the GLUE natural language understanding benchmark. Our approach also works well at scale, where it performs comparably to RoBERTa and XLNet while using less than 1/4 of their compute and outperforms them when using the same amount of compute.

Motivation & Objective

• Improve pre-training efficiency and performance for Transformer encoders compared to masked language modeling (MLM) like BERT.
• Develop a discriminative pre-training task that uses replacements sampled from a generator rather than masking tokens.
• Enable learning from all input tokens rather than only the masked subset to accelerate convergence and improve representations.
• Demonstrate scalability and efficiency across small and large model regimes on GLUE and SQuAD benchmarks.

Proposed method

• Propose two-network pre-training with a generator G and a discriminator D, both based on Transformer encoders.
• Corrupt input by replacing a subset of tokens with samples from G, forming a corrupted sequence.
• Train D to predict for every token whether it is original or a generator replacement (replaced token detection).
• Train G with max-likelihood masked language modeling to generate plausible replacements (not adversarially).
• Use a combined objective: L = E[MLM loss of G] + lambda * E[Disc loss of D], where Disc loss is a binary classification over each token in the corrupted sequence.
• Explore weight sharing between G and D (embeddings shared, sometimes entire weights tied) and different generator sizes to balance compute and performance.
• Evaluate on GLUE and SQuAD, comparing ELECTRA with BERT, XLNet, RoBERTa, and GPT under similar compute and data regimes.

Experimental results

Research questions

• RQ1 Does learning from all input tokens via replaced token detection improve efficiency and performance relative to traditional MLM pre-training?
• RQ2 How do generator size, weight sharing strategies, and training algorithms affect ELECTRA’s sample efficiency and downstream performance?
• RQ3 Can ELECTRA achieve comparable or superior results to state-of-the-art models (RoBERTa, XLNet) with less pre-training compute?
• RQ4 How does ELECTRA perform in small-model regimes and on SQuAD 2.0’s answerability task?

Key findings

• ELECTRA substantially outperforms MLM-based methods (e.g., BERT) given the same model size, data, and compute on GLUE and SQuAD.
• ELECTRA-Small trained on 1 GPU in 4 days outperforms GPT and is competitive with larger models, while requiring far less compute and with far fewer parameters.
• In large-scale settings, ELECTRA-Large achieves performance comparable to RoBERTa and XLNet with under 1/4 the pre-training compute and outperforms them when using similar compute.
• Learning from all input tokens (discriminator objective) is a major contributor to efficiency and performance gains; training with a smaller generator relative to the discriminator further improves results.
• Two-stage and adversarial training variants did not surpass joint ML objective; maximum-likelihood training of the generator yielded better downstream results.
• Across model sizes, ELECTRA’s gains are more pronounced as the model size decreases, indicating improved parameter efficiency.